trex-core/trex_core/application/balance_health.py

"""Broker-balance worker health tracking with bounded backoff."""
from __future__ import annotations

import math
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from threading import Lock
from typing import Callable


@dataclass(frozen=True)
class BalanceWorkerHealth:
    worker: str
    last_success_at: datetime | None
    consecutive_failures: int
    failure_code: str | None
    failure_at: datetime | None
    next_retry_at: datetime | None

    def as_dict(self, *, now: datetime, stale_after_sec: float) -> dict[str, object]:
        stale = False
        if self.last_success_at is not None and stale_after_sec >= 0:
            stale = (now - self.last_success_at).total_seconds() > stale_after_sec
        state = "ok"
        if self.last_success_at is None:
            state = "unavailable"
        elif self.consecutive_failures > 0:
            state = "unavailable"
        elif stale:
            state = "stale"
        return {
            "worker": self.worker,
            "state": state,
            "last_success_at": (
                self.last_success_at.astimezone(UTC).isoformat()
                if self.last_success_at is not None
                else None
            ),
            "consecutive_failures": self.consecutive_failures,
            "failure_code": self.failure_code,
            "failure_at": (
                self.failure_at.astimezone(UTC).isoformat()
                if self.failure_at is not None
                else None
            ),
            "next_retry_at": (
                self.next_retry_at.astimezone(UTC).isoformat()
                if self.next_retry_at is not None
                else None
            ),
            "stale": stale,
        }


class BalanceHealthTracker:
    def __init__(
        self,
        *,
        clock: Callable[[], datetime] | None = None,
        base_backoff_sec: float = 15.0,
        max_backoff_sec: float = 300.0,
    ) -> None:
        self._clock = clock or (lambda: datetime.now(tz=UTC))
        self._base_backoff_sec = max(1.0, float(base_backoff_sec))
        self._max_backoff_sec = max(self._base_backoff_sec, float(max_backoff_sec))
        self._lock = Lock()
        self._states: dict[str, BalanceWorkerHealth] = {}
        if self._base_backoff_sec >= self._max_backoff_sec:
            self._max_backoff_exponent = 0
        else:
            ratio = self._max_backoff_sec / self._base_backoff_sec
            self._max_backoff_exponent = int(math.floor(math.log2(ratio)))

    def _now(self) -> datetime:
        return self._clock()

    def _get(self, worker: str) -> BalanceWorkerHealth:
        return self._states.get(
            worker,
            BalanceWorkerHealth(
                worker=worker,
                last_success_at=None,
                consecutive_failures=0,
                failure_code=None,
                failure_at=None,
                next_retry_at=None,
            ),
        )

    def record_success(self, worker: str, *, captured_at: datetime | None = None) -> str | None:
        with self._lock:
            prev = self._get(worker)
            recovered = prev.consecutive_failures > 0
            now = self._now()
            self._states[worker] = BalanceWorkerHealth(
                worker=worker,
                last_success_at=captured_at or now,
                consecutive_failures=0,
                failure_code=None,
                failure_at=None,
                next_retry_at=None,
            )
        return "recovered" if recovered else None

    def record_failure(self, worker: str, *, code: str) -> tuple[str | None, float]:
        with self._lock:
            prev = self._get(worker)
            now = self._now()
            failures = prev.consecutive_failures + 1
            exponent = max(0, failures - 1)
            if exponent > self._max_backoff_exponent:
                backoff = self._max_backoff_sec
            else:
                backoff = min(
                    self._base_backoff_sec * (2.0 ** exponent),
                    self._max_backoff_sec,
                )
            self._states[worker] = BalanceWorkerHealth(
                worker=worker,
                last_success_at=prev.last_success_at,
                consecutive_failures=failures,
                failure_code=code,
                failure_at=now,
                next_retry_at=now + timedelta(seconds=backoff),
            )
            lifecycle: str | None = None
            if prev.consecutive_failures == 0 or prev.failure_code != code:
                lifecycle = "detected"
            elif failures in {3, 6, 12}:
                lifecycle = "ongoing"
        return lifecycle, backoff

    def ready(self, worker: str) -> tuple[bool, float]:
        with self._lock:
            state = self._get(worker)
            now = self._now()
            retry_at = state.next_retry_at
            if retry_at is None:
                return True, 0.0
            remain = (retry_at - now).total_seconds()
            if remain <= 0:
                return True, 0.0
            return False, remain

    def status(self, worker: str, *, stale_after_sec: float) -> dict[str, object]:
        with self._lock:
            state = self._get(worker)
            now = self._now()
        return state.as_dict(now=now, stale_after_sec=stale_after_sec)