execution_queue.py

import asyncio
import os
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Optional


@dataclass
class QueueMetrics:
    total_enqueued: int = 0
    total_completed: int = 0
    total_rejected: int = 0
    total_errors: int = 0
    active_executions: int = 0
    queue_depth: int = 0
    avg_wait_ms: float = 0.0
    avg_exec_ms: float = 0.0
    _wait_samples: list = field(default_factory=list)
    _exec_samples: list = field(default_factory=list)

    def record_wait(self, ms: float):
        self._wait_samples.append(ms)
        if len(self._wait_samples) > 1000:
            self._wait_samples = self._wait_samples[-500:]
        self.avg_wait_ms = sum(self._wait_samples) / len(self._wait_samples)

    def record_exec(self, ms: float):
        self._exec_samples.append(ms)
        if len(self._exec_samples) > 1000:
            self._exec_samples = self._exec_samples[-500:]
        self.avg_exec_ms = sum(self._exec_samples) / len(self._exec_samples)

    def snapshot(self) -> dict:
        return {
            "total_enqueued": self.total_enqueued,
            "total_completed": self.total_completed,
            "total_rejected": self.total_rejected,
            "total_errors": self.total_errors,
            "active_executions": self.active_executions,
            "queue_depth": self.queue_depth,
            "avg_wait_ms": round(self.avg_wait_ms, 2),
            "avg_exec_ms": round(self.avg_exec_ms, 2),
        }


class ExecutionQueue:
    def __init__(
        self,
        max_concurrent: int = 0,
        max_queue_depth: int = 0,
        max_per_user: int = 0,
    ):
        self.max_concurrent = max_concurrent or int(os.getenv("RUNTIME_MAX_CONCURRENT", "10"))
        self.max_queue_depth = max_queue_depth or int(os.getenv("RUNTIME_MAX_QUEUE_DEPTH", "50"))
        self.max_per_user = max_per_user or int(os.getenv("RUNTIME_MAX_PER_USER", "3"))

        self._semaphore = asyncio.Semaphore(self.max_concurrent)
        self._user_counts: dict[str, int] = defaultdict(int)
        self._pending: int = 0
        self.metrics = QueueMetrics()

    def _check_admission(self, user_id: Optional[str]) -> Optional[str]:
        if self._pending >= self.max_queue_depth:
            return f"queue full ({self._pending}/{self.max_queue_depth})"
        if user_id and self._user_counts[user_id] >= self.max_per_user:
            return f"per-user limit reached ({self._user_counts[user_id]}/{self.max_per_user})"
        return None

    async def execute(self, user_id: Optional[str], coro):
        rejection = self._check_admission(user_id)
        if rejection:
            self.metrics.total_rejected += 1
            raise QueueFullError(rejection)

        self._pending += 1
        self.metrics.total_enqueued += 1
        self.metrics.queue_depth = self._pending
        if user_id:
            self._user_counts[user_id] += 1

        enqueue_time = time.monotonic()
        try:
            await self._semaphore.acquire()
            wait_ms = (time.monotonic() - enqueue_time) * 1000
            self.metrics.record_wait(wait_ms)
            self.metrics.active_executions += 1

            exec_start = time.monotonic()
            try:
                result = await coro
                self.metrics.total_completed += 1
                return result
            except Exception:
                self.metrics.total_errors += 1
                raise
            finally:
                exec_ms = (time.monotonic() - exec_start) * 1000
                self.metrics.record_exec(exec_ms)
                self.metrics.active_executions -= 1
        finally:
            self._semaphore.release()
            self._pending -= 1
            self.metrics.queue_depth = self._pending
            if user_id:
                self._user_counts[user_id] -= 1
                if self._user_counts[user_id] <= 0:
                    del self._user_counts[user_id]


class QueueFullError(Exception):
    pass