@@ -159,6 +159,7 @@ def __init__(
159159 self ._evictions = 0
160160 self ._bytes_evicted = 0
161161 self ._put_failures = 0
162+ self ._reservations : dict [str , int ] = {}
162163 self ._lock = threading .RLock ()
163164
164165 def put (self , block : CacheBlock ) -> bool :
@@ -263,6 +264,80 @@ def fetch(self, lease_id: str, *, now: float | None = None) -> tuple[CacheBlock,
263264 self ._bytes_served += lease .transfer_bytes
264265 return tuple (blocks )
265266
267+ def reserve (self , reservation_id : str , byte_count : int ) -> None :
268+ """Reserve cache capacity before an expensive Prefill job starts."""
269+ if not reservation_id or byte_count <= 0 :
270+ raise ValueError ("reservation id and byte count must be positive" )
271+ with self ._lock :
272+ if reservation_id in self ._reservations :
273+ raise ValueError ("duplicate cache reservation" )
274+ requested = int (byte_count )
275+ reserved = sum (self ._reservations .values ()) + requested
276+ if reserved > self .max_bytes :
277+ raise ValueError (
278+ f"snapshot reservation { requested }
279+ f"cache budget { self .max_bytes - sum (self ._reservations .values ())} ,
280+ )
281+ self ._expire_leases (time .time ())
282+ self ._evict_to_limit (self .max_bytes - reserved )
283+ if self ._bytes_used > self .max_bytes - reserved :
284+ raise ValueError ("cache capacity is pinned by active leases" )
285+ self ._reservations [reservation_id ] = requested
286+
287+ def release_reservation (self , reservation_id : str ) -> None :
288+ with self ._lock :
289+ self ._reservations .pop (reservation_id , None )
290+
291+ def publish_and_lease (
292+ self ,
293+ blocks : Sequence [CacheBlock ],
294+ block_hashes : Sequence [bytes ],
295+ * ,
296+ reservation_id : str ,
297+ lease_seconds : float = DEFAULT_LEASE_SECONDS ,
298+ ) -> PrefixLease :
299+ """Atomically publish and pin the final computed snapshot."""
300+ if not blocks or len (blocks ) != len (block_hashes ):
301+ raise ValueError ("one computed snapshot is required per block hash" )
302+ if lease_seconds <= 0 :
303+ raise ValueError ("lease_seconds must be > 0" )
304+ final = blocks [- 1 ]
305+ if final .block_hash != bytes (block_hashes [- 1 ]):
306+ raise ValueError ("final snapshot hash does not match request" )
307+ now = time .time ()
308+ with self ._lock :
309+ reserved = self ._reservations .get (reservation_id )
310+ if reserved is None :
311+ raise ValueError ("unknown cache reservation" )
312+ if final .nbytes > reserved :
313+ raise ValueError (
314+ f"final snapshot { final .nbytes } { reserved } ,
315+ )
316+ self ._expire_leases (now )
317+ self ._evict_to_limit (self .max_bytes - final .nbytes )
318+ if self ._bytes_used > self .max_bytes - final .nbytes :
319+ raise ValueError ("cache capacity is pinned by active leases" )
320+ self ._put_locked (final )
321+ lease_id = secrets .token_urlsafe (18 )
322+ lease = PrefixLease (
323+ lease_id = lease_id ,
324+ block_hashes = (final .block_hash ,),
325+ hit_block_count = len (block_hashes ),
326+ hit_token_count = final .token_count ,
327+ transfer_bytes = final .nbytes ,
328+ cache_epoch = self ._epoch ,
329+ expires_at_unix = now + lease_seconds ,
330+ payload_sha256 = final .payload_sha256 ,
331+ )
332+ self ._leases [lease_id ] = lease
333+ del self ._reservations [reservation_id ]
334+ # Preserve longest useful boundaries when spare capacity remains.
335+ for block in reversed (blocks [:- 1 ]):
336+ if block .nbytes + self ._bytes_used > self .max_bytes :
337+ continue
338+ self ._put_locked (block )
339+ return lease
340+
266341 def stats (self ) -> CacheStats :
267342 with self ._lock :
268343 return CacheStats (
@@ -285,9 +360,12 @@ def resize(self, max_bytes: int) -> bool:
285360 raise ValueError ("max_bytes must be > 0" )
286361 with self ._lock :
287362 previous = self .max_bytes
363+ reserved = sum (self ._reservations .values ())
364+ if int (max_bytes ) < reserved :
365+ return False
288366 self .max_bytes = int (max_bytes )
289- self ._evict_to_budget ( )
290- if self ._bytes_used > self .max_bytes :
367+ self ._evict_to_limit ( self . max_bytes - reserved )
368+ if self ._bytes_used > self .max_bytes - reserved :
291369 self .max_bytes = max (previous , self ._bytes_used )
292370 return False
293371 return True
@@ -322,8 +400,11 @@ def _pinned_hashes(self) -> set[bytes]:
322400 }
323401
324402 def _evict_to_budget (self ) -> None :
403+ self ._evict_to_limit (self .max_bytes )
404+
405+ def _evict_to_limit (self , limit : int ) -> None :
325406 pinned = self ._pinned_hashes ()
326- while self ._bytes_used > self . max_bytes and self ._blocks :
407+ while self ._bytes_used > limit and self ._blocks :
327408 victim = next ((h for h in self ._blocks if h not in pinned ), None )
328409 if victim is None :
329410 break
@@ -333,6 +414,19 @@ def _evict_to_budget(self) -> None:
333414 self ._bytes_evicted += block .nbytes
334415 self ._epoch += 1
335416
417+ def _put_locked (self , block : CacheBlock ) -> bool :
418+ existing = self ._blocks .get (block .block_hash )
419+ if existing is not None :
420+ if existing .payload_sha256 != block .payload_sha256 :
421+ self ._put_failures += 1
422+ raise ValueError ("content-address collision with different payload" )
423+ self ._blocks .move_to_end (block .block_hash )
424+ return False
425+ self ._blocks [block .block_hash ] = block
426+ self ._bytes_used += block .nbytes
427+ self ._epoch += 1
428+ return True
429+
336430
337431def total_payload_bytes (blocks : Iterable [CacheBlock ]) -> int :
338432 return sum (block .nbytes for block in blocks )
0 commit comments