106 lines
3.5 KiB
Python
106 lines
3.5 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Bounded polling lifecycle and progress health for the AI usage exporter."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
|
||
|
|
POLL_STARTUP_GRACE_SECONDS = 45
|
||
|
|
POLL_PROGRESS_BUDGET_SECONDS = 120
|
||
|
|
|
||
|
|
|
||
|
|
class PollingEngine:
|
||
|
|
"""Run isolated provider polls and track bounded forward progress."""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
collector: Any,
|
||
|
|
*,
|
||
|
|
interval: float,
|
||
|
|
startup_grace: float = POLL_STARTUP_GRACE_SECONDS,
|
||
|
|
progress_timeout: float | None = None,
|
||
|
|
clock: Any = time.monotonic,
|
||
|
|
) -> None:
|
||
|
|
self.collector = collector
|
||
|
|
self.interval = interval
|
||
|
|
self.startup_grace = startup_grace
|
||
|
|
self.progress_timeout = (
|
||
|
|
interval + POLL_PROGRESS_BUDGET_SECONDS
|
||
|
|
if progress_timeout is None
|
||
|
|
else progress_timeout
|
||
|
|
)
|
||
|
|
self._clock = clock
|
||
|
|
self._started_at = clock()
|
||
|
|
self._last_progress: float | None = None
|
||
|
|
self._progress_lock = threading.Lock()
|
||
|
|
self._thread: threading.Thread | None = None
|
||
|
|
|
||
|
|
def _mark_progress(self) -> None:
|
||
|
|
"""Record completion of one bounded provider attempt."""
|
||
|
|
with self._progress_lock:
|
||
|
|
self._last_progress = self._clock()
|
||
|
|
|
||
|
|
def poll_once(self) -> None:
|
||
|
|
"""Refresh every provider even if another provider fails unexpectedly."""
|
||
|
|
for provider in ("openai", "anthropic"):
|
||
|
|
started = time.time()
|
||
|
|
monotonic_started = time.monotonic()
|
||
|
|
try:
|
||
|
|
self.collector.refresh_provider(provider)
|
||
|
|
except Exception as error:
|
||
|
|
print(
|
||
|
|
f"{provider} quota collection isolated: {type(error).__name__}",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
self.collector.record_failure(
|
||
|
|
provider,
|
||
|
|
started=started,
|
||
|
|
monotonic_started=monotonic_started,
|
||
|
|
)
|
||
|
|
except Exception as record_error:
|
||
|
|
print(
|
||
|
|
f"{provider} quota failure accounting deferred: "
|
||
|
|
f"{type(record_error).__name__}",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
finally:
|
||
|
|
self._mark_progress()
|
||
|
|
|
||
|
|
def run(self) -> None:
|
||
|
|
"""Poll forever without allowing a cycle-level exception to stop the thread."""
|
||
|
|
while True:
|
||
|
|
try:
|
||
|
|
self.poll_once()
|
||
|
|
except Exception as error:
|
||
|
|
print(
|
||
|
|
f"quota polling cycle deferred: {type(error).__name__}", flush=True
|
||
|
|
)
|
||
|
|
time.sleep(self.interval)
|
||
|
|
|
||
|
|
def start(self) -> None:
|
||
|
|
"""Start the daemon poller exactly once."""
|
||
|
|
if self._thread is not None:
|
||
|
|
return
|
||
|
|
self._thread = threading.Thread(
|
||
|
|
target=self.run,
|
||
|
|
name="ai-usage-poller",
|
||
|
|
daemon=True,
|
||
|
|
)
|
||
|
|
self._thread.start()
|
||
|
|
|
||
|
|
def is_healthy(self) -> bool:
|
||
|
|
"""Report thread liveness and progress, independent of provider success."""
|
||
|
|
thread = self._thread
|
||
|
|
if thread is None or not thread.is_alive():
|
||
|
|
return False
|
||
|
|
now = self._clock()
|
||
|
|
with self._progress_lock:
|
||
|
|
last_progress = self._last_progress
|
||
|
|
if last_progress is None:
|
||
|
|
return now - self._started_at <= self.startup_grace
|
||
|
|
return now - last_progress <= self.progress_timeout
|