"""Bounded concurrency for heavy Jyotish API compute endpoints. Request threads that run rectification scans or consultation workflows share one process-wide semaphore sized for the 2 vCPU production host. A saturated request waits up to 8 seconds when fewer than 4 threads are already waiting. If that queue is full or the wait expires, the caller still gets HTTP 429 and Retry-After. A thread is never parked without a deadline. Health checks and other light routes must not call this gate. """ from __future__ import annotations import logging import os import threading import time DEFAULT_CONCURRENCY = 2 DEFAULT_RETRY_AFTER_SECONDS = 2 DEFAULT_QUEUE_WAIT_SECONDS = 8.0 DEFAULT_MAX_WAITERS = 4 ENV_CONCURRENCY = "JYOTISH_HEAVY_COMPUTE_CONCURRENCY" ENV_RETRY_AFTER = "JYOTISH_HEAVY_COMPUTE_RETRY_AFTER_SECONDS" ENV_GATE_CHART_OVERVIEW = "JYOTISH_GATE_CHART_OVERVIEW" _LOG = logging.getLogger("jyotish.heavy_compute") HEAVY_COMPUTE_PATHS = frozenset( { "/api/rectification/sensitivity_scan", "/api/active_rectification_events", "/api/active_rectification_events_v4", "/api/rectification/v5/candidate-features", "/api/rectification/v5/score", "/api/rectification/v5/varga-segments", "/api/rectification/v5/diagnostics", "/api/rectification/v5/vedastro-validate", "/api/rectification/v5/range_reading", "/api/rectification/v5/block_scan", "/api/dynamic_rectification_opportunities", "/api/dynamic_rectification_score", "/api/high_rigor_workflow", "/api/consultation_workflow", "/api/professional_reading", "/api/professional_report_reference", "/api/vedastro/range_scan", "/api/vedastro_gateway/run", "/api/thematic_report", "/api/qizheng", "/api/western", "/api/ephemeris_events", } ) _state_lock = threading.Lock() _semaphore: threading.BoundedSemaphore | None = None _retry_after_seconds = DEFAULT_RETRY_AFTER_SECONDS _limit = DEFAULT_CONCURRENCY _wait_seconds = DEFAULT_QUEUE_WAIT_SECONDS _max_waiters = DEFAULT_MAX_WAITERS _waiters = 0 _holders: list[dict] = [] _holder_seq = 0 class _HeldSlot: """One acquired heavy-compute permit, paired with the path that holds it.""" def __init__(self, semaphore: threading.BoundedSemaphore, holder_id: int) -> None: self._semaphore = semaphore self._holder_id = holder_id self._released = False def release(self) -> None: with _state_lock: if self._released: return self._released = True _holders[:] = [item for item in _holders if item["id"] != self._holder_id] self._semaphore.release() class HeavyComputeBusy(RuntimeError): """The short queue could not take this request; callers return HTTP 429.""" error_code = "ERR_COMPUTE_BUSY" def __init__(self, retry_after_seconds: int) -> None: super().__init__( "Heavy compute capacity is saturated; retry after the Retry-After delay." ) self.retry_after_seconds = max(int(retry_after_seconds), 1) def _parse_positive_int(raw: str | None, default: int) -> int: try: value = int(str(raw or "").strip()) except (TypeError, ValueError): return default return value if value >= 1 else default def is_heavy_compute_path(path: str) -> bool: if path in HEAVY_COMPUTE_PATHS: return True return path.startswith("/api/vedastro_gateway/jobs/") and path.endswith("/run") def reset_heavy_compute_gate( *, limit: int | None = None, retry_after_seconds: int | None = None, wait_seconds: float | None = None, max_waiters: int | None = None, ) -> None: """Rebuild the process-wide semaphore. Tests must call this after env changes.""" global _semaphore, _retry_after_seconds, _limit, _holders, _wait_seconds, _max_waiters, _waiters resolved_limit = ( limit if limit is not None else _parse_positive_int(os.environ.get(ENV_CONCURRENCY), DEFAULT_CONCURRENCY) ) resolved_retry = ( retry_after_seconds if retry_after_seconds is not None else _parse_positive_int(os.environ.get(ENV_RETRY_AFTER), DEFAULT_RETRY_AFTER_SECONDS) ) resolved_wait = DEFAULT_QUEUE_WAIT_SECONDS if wait_seconds is None else max(0.0, float(wait_seconds)) resolved_waiters = DEFAULT_MAX_WAITERS if max_waiters is None else max(0, int(max_waiters)) with _state_lock: _retry_after_seconds = resolved_retry _limit = resolved_limit _wait_seconds = resolved_wait _max_waiters = resolved_waiters _waiters = 0 _holders = [] _semaphore = threading.BoundedSemaphore(resolved_limit) def _ensure_locked() -> tuple[threading.BoundedSemaphore, int]: global _semaphore, _retry_after_seconds, _limit if _semaphore is None: _limit = _parse_positive_int(os.environ.get(ENV_CONCURRENCY), DEFAULT_CONCURRENCY) _retry_after_seconds = _parse_positive_int( os.environ.get(ENV_RETRY_AFTER), DEFAULT_RETRY_AFTER_SECONDS ) _semaphore = threading.BoundedSemaphore(_limit) return _semaphore, _retry_after_seconds def _holder_view(now: float) -> list[dict]: return [ { "path": item["path"], "held_seconds": round(max(now - item["started"], 0.0), 3), } for item in _holders ] def _remember_locked(path: str, semaphore: threading.BoundedSemaphore) -> _HeldSlot: """Record a permit that this thread already holds. Caller owns ``_state_lock``.""" global _holder_seq _holder_seq += 1 holder_id = _holder_seq _holders.append({"id": holder_id, "path": path, "started": time.monotonic()}) return _HeldSlot(semaphore, holder_id) def _reject_busy(path: str, retry_after: int) -> None: with _state_lock: snapshot = _holder_view(time.monotonic()) _LOG.warning( "heavy_compute_busy rejected_path=%s holders=%s", path, snapshot, ) raise HeavyComputeBusy(retry_after) def _acquire_recorded(path: str) -> _HeldSlot: """Take one permit. Wait briefly when full; still raise once the queue cannot.""" global _waiters with _state_lock: semaphore, retry_after = _ensure_locked() if semaphore.acquire(blocking=False): return _remember_locked(path, semaphore) if _waiters >= _max_waiters: queued = False else: _waiters += 1 queued = True wait_seconds = _wait_seconds if not queued: _reject_busy(path, retry_after) acquired = False try: acquired = bool(semaphore.acquire(blocking=True, timeout=wait_seconds)) finally: with _state_lock: _waiters -= 1 if not acquired: _reject_busy(path, retry_after) with _state_lock: return _remember_locked(path, semaphore) def heavy_compute_waiter_count() -> int: """Threads parked inside the short queue. Tests use this to avoid races.""" with _state_lock: return _waiters def acquire_heavy_compute_slot(path: str) -> _HeldSlot | None: """Acquire a slot for a gated path. Light paths return None.""" if not is_heavy_compute_path(path): return None return _acquire_recorded(path) def _chart_overview_gate_enabled() -> bool: return os.environ.get(ENV_GATE_CHART_OVERVIEW, "").strip() == "1" def acquire_heavy_compute_slot_for_request(path: str, body: dict | None) -> _HeldSlot | None: """Path gate, plus the chart-overview switch and the warm-consultation exception. ``/api/chart`` stays out of ``HEAVY_COMPUTE_PATHS``. It takes a slot only when ``JYOTISH_GATE_CHART_OVERVIEW=1``, the body does not skip the VedAstro overview, and the body is not an async enqueue. Default is off. """ payload = body if isinstance(body, dict) else {} if path == "/api/consultation_workflow" and payload.get("consult_card_warm_only") is True: return None if path == "/api/chart" and _chart_overview_gate_enabled(): if payload.get("async") or payload.get("enqueue"): return None if payload.get("skip_vedastro_main_entry_overview"): return None return _acquire_recorded(path) return acquire_heavy_compute_slot(path) def heavy_compute_status() -> dict: """Process-local counts only. No request body and no user identity.""" with _state_lock: _ensure_locked() return {"limit": _limit, "in_use": len(_holders)} def release_heavy_compute_slot(slot: _HeldSlot | None) -> None: if slot is not None: slot.release()