Jyotisha/scripts/api_heavy_compute_gate.py

"""Fail-fast bounded concurrency for heavy Jyotish API compute endpoints.

Request threads that run rectification scans or high-rigor workflows share one
process-wide semaphore sized for the 2 vCPU production host. Saturated requests
return immediately; they are not queued. 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
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
_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):
    """No heavy-compute slot is free; callers must fail fast with 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,
) -> None:
    """Rebuild the process-wide semaphore. Tests must call this after env changes."""
    global _semaphore, _retry_after_seconds, _limit, _holders
    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)
    )
    with _state_lock:
        _retry_after_seconds = resolved_retry
        _limit = resolved_limit
        _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 _acquire_recorded(path: str) -> _HeldSlot:
    """Take one permit and remember which path holds it. Fail fast when full."""
    global _holder_seq
    with _state_lock:
        semaphore, retry_after = _ensure_locked()
        if not semaphore.acquire(blocking=False):
            snapshot = _holder_view(time.monotonic())
            failed = True
            slot = None
        else:
            _holder_seq += 1
            holder_id = _holder_seq
            _holders.append({"id": holder_id, "path": path, "started": time.monotonic()})
            failed = False
            snapshot = []
            slot = _HeldSlot(semaphore, holder_id)
    if failed:
        _LOG.warning(
            "heavy_compute_busy rejected_path=%s holders=%s",
            path,
            snapshot,
        )
        raise HeavyComputeBusy(retry_after)
    assert slot is not None
    return slot


def acquire_heavy_compute_slot(path: str) -> _HeldSlot | None:
    """Acquire a slot for a gated path. Light paths return None. Fail-fast on saturation."""
    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()