183 lines
6.7 KiB
Python
183 lines
6.7 KiB
Python
"""Round-robin admission by user, with bounded pending work."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import threading
|
|
import time
|
|
from collections import Counter, deque
|
|
from collections.abc import Callable
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class QueueFullError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class FairAsyncGate:
|
|
def __init__(self, capacity: int, per_user: int, queue_limit: int = 256):
|
|
self.capacity, self.per_user, self.queue_limit = capacity, min(capacity, per_user), queue_limit
|
|
self._active = Counter()
|
|
self._queues: dict[str, deque] = {}
|
|
self._order = deque()
|
|
|
|
def _dispatch(self):
|
|
while sum(self._active.values()) < self.capacity:
|
|
for _ in range(len(self._order)):
|
|
key = self._order.popleft()
|
|
if self._active[key] < self.per_user:
|
|
future = self._queues[key].popleft()
|
|
if self._queues[key]:
|
|
self._order.append(key)
|
|
else:
|
|
del self._queues[key]
|
|
self._active[key] += 1
|
|
future.set_result(None)
|
|
break
|
|
self._order.append(key)
|
|
else:
|
|
return
|
|
|
|
def _release(self, key):
|
|
self._active[key] -= 1
|
|
if not self._active[key]:
|
|
del self._active[key]
|
|
self._dispatch()
|
|
|
|
@asynccontextmanager
|
|
async def slot(self, key: str):
|
|
if len(self._queues.get(key, ())) >= max(1, self.queue_limit // max(2, self.capacity)) or sum(len(queue) for queue in self._queues.values()) >= self.queue_limit:
|
|
raise QueueFullError("Background queue is full; retry later")
|
|
future = asyncio.get_running_loop().create_future()
|
|
if key not in self._queues:
|
|
self._queues[key] = deque()
|
|
self._order.append(key)
|
|
self._queues[key].append(future)
|
|
self._dispatch()
|
|
try:
|
|
await asyncio.shield(future)
|
|
except BaseException:
|
|
if future.done():
|
|
self._release(key)
|
|
else:
|
|
self._queues[key].remove(future)
|
|
if not self._queues[key]:
|
|
del self._queues[key]
|
|
self._order.remove(key)
|
|
future.cancel()
|
|
self._dispatch()
|
|
raise
|
|
try:
|
|
yield
|
|
finally:
|
|
self._release(key)
|
|
|
|
|
|
@dataclass
|
|
class _Job:
|
|
key: str
|
|
run: Callable[[float], None]
|
|
deadline: float
|
|
cancelled: Callable[[], bool]
|
|
expired: Callable[[bool], None]
|
|
|
|
|
|
class FairBackgroundScheduler:
|
|
"""Only admitted work enters the executor, so queued jobs occupy no worker.
|
|
|
|
Deadlines include queue time. Expiry/cancellation never releases the slot of
|
|
a running job; its execution wrapper must first stop the actual work.
|
|
"""
|
|
|
|
def __init__(self, capacity=3, per_user=2, queue_limit=256):
|
|
self.capacity, self.per_user, self.queue_limit = capacity, min(capacity, per_user), queue_limit
|
|
self._condition = threading.Condition()
|
|
self._queues: dict[str, deque[_Job]] = {}
|
|
self._order = deque()
|
|
self._active = Counter()
|
|
self._closed = False
|
|
self._pool = ThreadPoolExecutor(max_workers=capacity, thread_name_prefix="subagent-fair")
|
|
self._dispatcher = threading.Thread(target=self._dispatch, name="subagent-admission", daemon=True)
|
|
self._dispatcher.start()
|
|
|
|
def submit(self, key, run, *, timeout, cancelled, expired):
|
|
with self._condition:
|
|
if self._closed or len(self._queues.get(key, ())) >= max(1, self.queue_limit // max(2, self.capacity)) or sum(map(len, self._queues.values())) >= self.queue_limit:
|
|
raise QueueFullError("Subagent queue is full; retry later")
|
|
if key not in self._queues:
|
|
self._queues[key] = deque()
|
|
self._order.append(key)
|
|
self._queues[key].append(_Job(key, run, time.monotonic() + timeout, cancelled, expired))
|
|
self._condition.notify_all()
|
|
|
|
def _execute(self, job):
|
|
try:
|
|
remaining = job.deadline - time.monotonic()
|
|
if job.cancelled() or remaining <= 0:
|
|
self._expire(job, job.cancelled())
|
|
else:
|
|
job.run(remaining)
|
|
finally:
|
|
with self._condition:
|
|
self._active[job.key] -= 1
|
|
if not self._active[job.key]:
|
|
del self._active[job.key]
|
|
self._condition.notify_all()
|
|
|
|
@staticmethod
|
|
def _expire(job, cancelled):
|
|
try:
|
|
job.expired(cancelled)
|
|
except Exception:
|
|
logger.exception("Background expiry callback failed for user %s", job.key)
|
|
|
|
def _dispatch(self):
|
|
with self._condition:
|
|
while not self._closed:
|
|
for key in list(self._order):
|
|
retained = deque()
|
|
for job in self._queues[key]:
|
|
if job.cancelled() or job.deadline <= time.monotonic():
|
|
self._expire(job, job.cancelled())
|
|
else:
|
|
retained.append(job)
|
|
if retained:
|
|
self._queues[key] = retained
|
|
else:
|
|
del self._queues[key]
|
|
self._order.remove(key)
|
|
admitted = False
|
|
if sum(self._active.values()) < self.capacity:
|
|
for _ in range(len(self._order)):
|
|
key = self._order.popleft()
|
|
if self._active[key] < self.per_user:
|
|
job = self._queues[key].popleft()
|
|
if self._queues[key]:
|
|
self._order.append(key)
|
|
else:
|
|
del self._queues[key]
|
|
self._active[key] += 1
|
|
self._pool.submit(self._execute, job)
|
|
admitted = True
|
|
break
|
|
self._order.append(key)
|
|
if not admitted:
|
|
self._condition.wait(timeout=0.05)
|
|
|
|
def shutdown(self, wait=True):
|
|
with self._condition:
|
|
self._closed = True
|
|
for queue in self._queues.values():
|
|
for job in queue:
|
|
self._expire(job, True)
|
|
self._queues.clear()
|
|
self._order.clear()
|
|
self._condition.notify_all()
|
|
self._dispatcher.join()
|
|
self._pool.shutdown(wait=wait)
|