Coverage for src/lilbee/server/chat_dispatch/concurrency.py: 100%
82 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Process-wide chat-generation admission gate shared by every chat route."""
3from __future__ import annotations
5import asyncio
6from collections import deque
7from functools import lru_cache
9# Default upper bound for waiting on a free chat slot before surfacing a real
10# "busy" response. Tuned to absorb the typical opencode retry storm (one user
11# turn can take 30-60s on slow local models) while still bounding the worst case
12# for a genuinely-stuck backend.
13DEFAULT_BUSY_WAIT_S = 60.0
16class ChatBusyError(Exception):
17 """Raised when every chat slot is in use after waiting."""
20class ChatGate:
21 """Admit up to the backend's live slot capacity; the rest wait FIFO, then 429.
23 Capacity is read at admission time, not fixed when the gate is created, so a
24 model swap or provider change that changes the slot count takes effect at
25 once, and a single in-process model (capacity 1) is never oversubscribed.
26 The fleet reports its ``--parallel`` slot count, so its continuous-batching
27 slots are actually used instead of one request at a time. The latest count
28 any caller reports is remembered, so a capacity increase also admits the
29 waiters that are already queued rather than leaving them to time out
30 against a backend that now has room.
32 Hand-rolled on purpose, and the reason is worth keeping because the obvious
33 replacement almost fits. ``anyio.CapacityLimiter`` has a runtime-settable
34 ``total_tokens`` that would make a capacity increase wake waiters for free,
35 but it binds each token to the borrowing task and raises ``RuntimeError``
36 when another task releases it. Slots here are released from wherever cleanup
37 happens first, including a response's after-send hook, which is not the task
38 that acquired. ``asyncio.Semaphore`` allows the cross-task release but has no
39 live capacity. Neither fits until the streaming cleanup paths are collapsed
40 onto one task.
42 Slots are *handed to* waiters rather than merely signalled. Waking a waiter
43 reserves its slot in the same synchronous step, and a request arriving while
44 a waiter is queued joins the back of the queue instead of testing the
45 counter. Without both, a newcomer could take the slot in the window between
46 the wake-up and the woken waiter resuming, and the waiter would rejoin at
47 the tail: the requests that had waited longest would be overtaken
48 repeatedly and would be the ones to time out.
49 """
51 def __init__(self) -> None:
52 self._in_flight = 0
53 # Last slot count reported by a caller; the gate has no other view of it.
54 self._ceiling = 1
55 self._waiters: deque[asyncio.Future[None]] = deque()
57 @property
58 def in_flight(self) -> int:
59 """Chat generations currently admitted."""
60 return self._in_flight
62 async def acquire(self, capacity: int, timeout: float) -> None:
63 """Reserve a slot, waiting up to *timeout*; raise ChatBusyError if full."""
64 loop = asyncio.get_running_loop()
65 self._observe_capacity(max(1, capacity))
66 # Admit straight away only with room AND nobody ahead in the queue.
67 if not self._waiters and self._in_flight < self._ceiling:
68 self._in_flight += 1
69 return
70 if timeout <= 0:
71 raise ChatBusyError(_busy_message(timeout))
72 waiter: asyncio.Future[None] = loop.create_future()
73 self._waiters.append(waiter)
74 try:
75 # asyncio.timeout, not wait_for: 3.11's wait_for swallows a task
76 # cancellation that races a completed waiter (fixed in 3.12).
77 async with asyncio.timeout(timeout):
78 await waiter
79 # Returning means a slot was reserved for us by whoever woke us;
80 # there is nothing left to claim.
81 except TimeoutError as exc:
82 self._abandon(waiter)
83 raise ChatBusyError(_busy_message(timeout)) from exc
84 except asyncio.CancelledError:
85 self._abandon(waiter)
86 raise
87 finally:
88 if waiter in self._waiters:
89 self._waiters.remove(waiter)
91 async def release(self) -> None:
92 """Free an acquired slot and admit whoever it makes room for.
94 Contains no awaits: the decrement and wake-ups run synchronously on the
95 event loop, so a cancellation delivered to the caller (for example a
96 client disconnect tearing down a streaming response) can never abort
97 the release halfway and leak the slot.
98 """
99 if self._in_flight > 0:
100 self._in_flight -= 1
101 self._admit_waiters()
103 def _observe_capacity(self, ceiling: int) -> None:
104 """Record the caller's live slot count and admit anyone it now fits."""
105 self._ceiling = ceiling
106 self._admit_waiters()
108 def _admit_waiters(self) -> None:
109 """Hand a reserved slot to each queued waiter that now fits."""
110 while self._in_flight < self._ceiling:
111 waiter = self._pop_live_waiter()
112 if waiter is None:
113 return
114 # Reserve before waking: the slot is the waiter's from this moment,
115 # so nothing entering acquire() in between can take it.
116 self._in_flight += 1
117 waiter.set_result(None)
119 def _pop_live_waiter(self) -> asyncio.Future[None] | None:
120 """Remove and return the oldest waiter still able to take a slot."""
121 while self._waiters:
122 waiter = self._waiters.popleft()
123 if not waiter.done():
124 return waiter
125 return None
127 def _abandon(self, waiter: asyncio.Future[None]) -> None:
128 """Give back a slot that was reserved for a waiter that then bailed out."""
129 if waiter.done() and not waiter.cancelled():
130 if self._in_flight > 0:
131 self._in_flight -= 1
132 self._admit_waiters()
135def _busy_message(timeout: float) -> str:
136 rendered = f"{timeout:.1f}s" if timeout < 1 else f"{timeout:.0f}s"
137 return f"Chat backend busy: all slots in use after {rendered}. Retry shortly."
140@lru_cache(maxsize=1)
141def chat_gate() -> ChatGate:
142 """Return the process-wide chat admission gate (created lazily)."""
143 return ChatGate()
146async def acquire_chat_slot_or_busy(capacity: int, timeout: float | None = None) -> None:
147 """Reserve one of *capacity* chat slots, waiting up to *timeout*.
149 Raises :class:`ChatBusyError` only when the wait times out. Returns with the
150 slot held; callers must ``await release_chat_slot()`` in a ``finally``.
151 *timeout* defaults to :data:`DEFAULT_BUSY_WAIT_S`.
152 """
153 effective_timeout = DEFAULT_BUSY_WAIT_S if timeout is None else timeout
154 await chat_gate().acquire(capacity, effective_timeout)
157async def release_chat_slot() -> None:
158 """Release a chat slot reserved by :func:`acquire_chat_slot_or_busy`."""
159 await chat_gate().release()
162class ChatSlotGuard:
163 """Releases one acquired chat slot at most once across multiple cleanup paths.
165 A streaming route acquires its slot before the SSE generator runs; a client
166 that disconnects before the generator's first iteration means the generator
167 body (and its ``finally``) never executes. Every cleanup path (generator
168 ``finally``, response after-send hook, explicit ``aclose``) releases through
169 the same guard, so whichever fires first frees the slot and the rest no-op.
170 """
172 def __init__(self) -> None:
173 self._released = False
175 @property
176 def released(self) -> bool:
177 """True once the slot has been freed."""
178 return self._released
180 async def release(self) -> None:
181 """Free the slot on first call; later calls are no-ops."""
182 if self._released:
183 return
184 self._released = True
185 # ChatGate.release never yields to the event loop, so no cancellation
186 # can land between flipping the flag and the slot actually freeing.
187 await release_chat_slot()