Artur Mukhamadiev
5aad7a0b45
In-memory AdmissionController (idle/reserved/running, FIFO tickets keyed by
the session cookie, reservation and heartbeat expiry, promotion after
cleanup, EMA wait estimate) behind POST /api/v1/queue/join,
GET /api/v1/queue/status and DELETE /api/v1/queue/ticket. POST /api/v1/query
claims the session's reservation first and joins implicitly when idle.
Settings: CONFLUENCE_WEB_QUEUE_{RESERVATION_SECONDS,HEARTBEAT_SECONDS,MAX_LENGTH}.
296 lines
11 KiB
Python
296 lines
11 KiB
Python
"""In-memory admission queue for the single-runner query gate.
|
|
|
|
Implements the state machine and estimate from docs/QUEUE_SPECIFICATION.md:
|
|
one runner slot that is idle, reserved for a single session, or running for
|
|
a single session, with a FIFO queue of waiting tickets behind it. All public
|
|
methods are synchronous and perform no suspension, so a single call is
|
|
atomic under asyncio's cooperative scheduling -- the same reasoning
|
|
backend.runner.QueryRunner relies on for its own gate.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import secrets
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Callable, Dict, List, Optional
|
|
|
|
from backend.errors import QueueFullError, TicketNotFoundError
|
|
|
|
# Estimate smoothing factor and history window (spec section 4/5 constants).
|
|
ESTIMATE_ALPHA = 0.3
|
|
ESTIMATE_HISTORY = 20
|
|
|
|
# Runs cancelled by the client before this many seconds do not update the estimate.
|
|
CANCELLED_RUN_FLOOR_SECONDS = 5.0
|
|
|
|
TicketState = str # "queued" | "ready" | "running"
|
|
|
|
|
|
@dataclass
|
|
class Ticket:
|
|
ticket_id: str
|
|
session_id: str
|
|
joined_at: float
|
|
last_seen_at: float
|
|
state: TicketState
|
|
reserved_at: Optional[float] = None
|
|
expires_at: Optional[float] = None
|
|
|
|
|
|
def _generate_ticket_id() -> str:
|
|
# Opaque, case-sensitive, [A-Za-z0-9_-] per CONTRACTS section 1.
|
|
return "q_" + secrets.token_urlsafe(12)
|
|
|
|
|
|
class AdmissionController:
|
|
"""FIFO admission queue in front of the single query runner slot."""
|
|
|
|
def __init__(
|
|
self,
|
|
reservation_seconds: float,
|
|
heartbeat_seconds: float,
|
|
max_queue_length: int,
|
|
clock: Callable[[], float] = time.monotonic,
|
|
) -> None:
|
|
self.reservation_seconds = reservation_seconds
|
|
self.heartbeat_seconds = heartbeat_seconds
|
|
self.max_queue_length = max_queue_length
|
|
self.clock = clock
|
|
|
|
self._tickets: Dict[str, Ticket] = {}
|
|
self._queue: List[str] = [] # session_ids in "queued" state, FIFO
|
|
self._reserved_session: Optional[str] = None
|
|
self._running_session: Optional[str] = None
|
|
self._running_started_at: Optional[float] = None
|
|
|
|
self._run_durations: List[float] = []
|
|
self._mean_duration: Optional[float] = None
|
|
|
|
# -- Public API -----------------------------------------------------
|
|
|
|
def join(self, session_id: str) -> dict:
|
|
"""Idempotent join. Returns the (possibly pre-existing) ticket view."""
|
|
self._expire()
|
|
existing = self._tickets.get(session_id)
|
|
if existing is not None:
|
|
return self._view(existing)
|
|
|
|
now = self.clock()
|
|
if self._reserved_session is None and self._running_session is None:
|
|
# Runner is idle: book the place directly.
|
|
ticket = Ticket(
|
|
ticket_id=_generate_ticket_id(),
|
|
session_id=session_id,
|
|
joined_at=now,
|
|
last_seen_at=now,
|
|
state="ready",
|
|
reserved_at=now,
|
|
expires_at=now + self.reservation_seconds,
|
|
)
|
|
self._tickets[session_id] = ticket
|
|
self._reserved_session = session_id
|
|
return self._view(ticket)
|
|
|
|
if len(self._queue) >= self.max_queue_length:
|
|
raise QueueFullError("The queue is full; please try again later")
|
|
|
|
ticket = Ticket(
|
|
ticket_id=_generate_ticket_id(),
|
|
session_id=session_id,
|
|
joined_at=now,
|
|
last_seen_at=now,
|
|
state="queued",
|
|
)
|
|
self._tickets[session_id] = ticket
|
|
self._queue.append(session_id)
|
|
return self._view(ticket)
|
|
|
|
def status(self, session_id: str) -> dict:
|
|
"""Heartbeat poll. Refreshes last_seen_at; raises TicketNotFoundError if absent."""
|
|
self._expire()
|
|
ticket = self._tickets.get(session_id)
|
|
if ticket is None:
|
|
raise TicketNotFoundError("No active queue ticket for this session")
|
|
ticket.last_seen_at = self.clock()
|
|
return self._view(ticket)
|
|
|
|
def leave(self, session_id: str) -> None:
|
|
"""Drop the session's ticket, if any. Promotes the next ticket if it held the reservation."""
|
|
self._expire()
|
|
ticket = self._tickets.pop(session_id, None)
|
|
if ticket is None:
|
|
return
|
|
if session_id in self._queue:
|
|
self._queue.remove(session_id)
|
|
if self._reserved_session == session_id:
|
|
self._reserved_session = None
|
|
self._promote()
|
|
|
|
def claim(self, session_id: str) -> bool:
|
|
"""Attempt to move a query request's session into the running slot.
|
|
|
|
Returns True if the session may proceed (claimed or implicitly joined),
|
|
False if it must be refused with 409 busy. Does not disturb the queue
|
|
on refusal.
|
|
"""
|
|
now = self.clock()
|
|
|
|
# The reserved session's own claim must see 409 once its window has
|
|
# passed, even when this very call is what discovers the expiry --
|
|
# it must not fall through to idle/implicit-join for itself.
|
|
if self._reserved_session == session_id:
|
|
ticket = self._tickets.get(session_id)
|
|
expired = ticket is None or (
|
|
ticket.expires_at is not None and now >= ticket.expires_at
|
|
)
|
|
if expired:
|
|
self._tickets.pop(session_id, None)
|
|
self._reserved_session = None
|
|
self._promote()
|
|
return False
|
|
ticket.state = "running"
|
|
ticket.reserved_at = None
|
|
ticket.expires_at = None
|
|
self._reserved_session = None
|
|
self._running_session = session_id
|
|
self._running_started_at = now
|
|
return True
|
|
|
|
# Not the current holder: evaluate lazy expiries (a stale reservation
|
|
# held by someone else, or stale queued heartbeats) before deciding.
|
|
self._expire()
|
|
|
|
if self._running_session is not None:
|
|
return False
|
|
if self._reserved_session is not None:
|
|
return False
|
|
|
|
# Idle: implicit join, reserve, and claim in one step.
|
|
ticket = Ticket(
|
|
ticket_id=_generate_ticket_id(),
|
|
session_id=session_id,
|
|
joined_at=now,
|
|
last_seen_at=now,
|
|
state="running",
|
|
)
|
|
self._tickets[session_id] = ticket
|
|
self._running_session = session_id
|
|
self._running_started_at = now
|
|
return True
|
|
|
|
def run_end(self, session_id: str, duration_seconds: float, cancelled: bool = False) -> None:
|
|
"""Record run completion (success, failure, timeout, cancellation, disconnect).
|
|
|
|
Must be called after the runner's own cleanup has finished so the next
|
|
holder's reservation window never overlaps residual cleanup.
|
|
"""
|
|
self._tickets.pop(session_id, None)
|
|
if self._running_session == session_id:
|
|
self._running_session = None
|
|
self._running_started_at = None
|
|
|
|
if not (cancelled and duration_seconds < CANCELLED_RUN_FLOOR_SECONDS):
|
|
self._record_duration(duration_seconds)
|
|
|
|
self._promote()
|
|
|
|
def expire(self) -> None:
|
|
"""Public hook for periodic maintenance; evaluates lazy expiries eagerly."""
|
|
self._expire()
|
|
|
|
# -- Internal ---------------------------------------------------------
|
|
|
|
def _record_duration(self, duration_seconds: float) -> None:
|
|
self._run_durations.append(max(0.0, duration_seconds))
|
|
if len(self._run_durations) > ESTIMATE_HISTORY:
|
|
self._run_durations = self._run_durations[-ESTIMATE_HISTORY:]
|
|
mean = self._run_durations[0]
|
|
for d in self._run_durations[1:]:
|
|
mean = ESTIMATE_ALPHA * d + (1 - ESTIMATE_ALPHA) * mean
|
|
self._mean_duration = mean
|
|
|
|
def _promote(self) -> None:
|
|
if self._reserved_session is not None or self._running_session is not None:
|
|
return
|
|
if not self._queue:
|
|
return
|
|
session_id = self._queue.pop(0)
|
|
ticket = self._tickets.get(session_id)
|
|
if ticket is None:
|
|
# Should not happen; keep state consistent and try the next one.
|
|
self._promote()
|
|
return
|
|
now = self.clock()
|
|
ticket.state = "ready"
|
|
ticket.reserved_at = now
|
|
ticket.expires_at = now + self.reservation_seconds
|
|
self._reserved_session = session_id
|
|
|
|
def _expire(self) -> None:
|
|
now = self.clock()
|
|
|
|
# Reservation expiry.
|
|
if self._reserved_session is not None:
|
|
ticket = self._tickets.get(self._reserved_session)
|
|
if ticket is None or (ticket.expires_at is not None and now >= ticket.expires_at):
|
|
if ticket is not None:
|
|
self._tickets.pop(self._reserved_session, None)
|
|
self._reserved_session = None
|
|
self._promote()
|
|
|
|
# Heartbeat expiry for queued tickets only.
|
|
stale = [
|
|
sid
|
|
for sid in self._queue
|
|
if (ticket := self._tickets.get(sid)) is not None
|
|
and now - ticket.last_seen_at > self.heartbeat_seconds
|
|
]
|
|
for sid in stale:
|
|
self._tickets.pop(sid, None)
|
|
if sid in self._queue:
|
|
self._queue.remove(sid)
|
|
|
|
def _position(self, session_id: str, ticket: Ticket) -> int:
|
|
if ticket.state == "running":
|
|
return 0
|
|
front: List[str] = []
|
|
if self._reserved_session is not None:
|
|
front.append(self._reserved_session)
|
|
front.extend(self._queue)
|
|
try:
|
|
return front.index(session_id) + 1
|
|
except ValueError:
|
|
return 0
|
|
|
|
def _eta_seconds(self, ahead: int) -> Optional[float]:
|
|
if self._mean_duration is None:
|
|
return None
|
|
if self._running_session is not None:
|
|
started_at = self._running_started_at if self._running_started_at is not None else self.clock()
|
|
elapsed = self.clock() - started_at
|
|
return max(0.0, self._mean_duration * (ahead + 1) - elapsed)
|
|
return max(0.0, self._mean_duration * (ahead + 1))
|
|
|
|
def _view(self, ticket: Ticket) -> dict:
|
|
position = self._position(ticket.session_id, ticket)
|
|
ahead = max(0, position - 1)
|
|
eta = self._eta_seconds(ahead)
|
|
runner = "running" if self._running_session is not None else "reserved"
|
|
|
|
reservation_expires_in: Optional[float] = None
|
|
if ticket.state == "ready" and ticket.expires_at is not None:
|
|
reservation_expires_in = max(0.0, ticket.expires_at - self.clock())
|
|
|
|
return {
|
|
"ticket_id": ticket.ticket_id,
|
|
"status": ticket.state,
|
|
"position": position,
|
|
"ahead": ahead,
|
|
"eta_seconds": round(eta) if eta is not None else None,
|
|
"reservation_expires_in_seconds": (
|
|
round(reservation_expires_in) if reservation_expires_in is not None else None
|
|
),
|
|
"runner": runner,
|
|
}
|