Jyotisha/scripts/api_heavy_compute_gate.py

"""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()