m4
541 lines
20 KiB
Python
541 lines
20 KiB
Python
"""Local durable outbox and background sender for UsageEvent v1."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import atexit
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import random
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from urllib.parse import urlsplit, urlunsplit
|
|
|
|
import httpx
|
|
|
|
from .schema import UsageEventV1
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _now() -> str:
|
|
return datetime.now(UTC).isoformat().replace("+00:00", "Z")
|
|
|
|
|
|
def _int_env(name: str, default: int, minimum: int = 1) -> int:
|
|
try:
|
|
return max(minimum, int(os.getenv(name, str(default))))
|
|
except ValueError:
|
|
return default
|
|
|
|
|
|
def _float_env(name: str, default: float) -> float:
|
|
try:
|
|
return max(0.05, float(os.getenv(name, str(default))))
|
|
except ValueError:
|
|
return default
|
|
|
|
|
|
def _fsync_dir(path: Path) -> None:
|
|
try:
|
|
descriptor = os.open(path, os.O_RDONLY)
|
|
except OSError:
|
|
return
|
|
try:
|
|
os.fsync(descriptor)
|
|
except OSError:
|
|
pass
|
|
finally:
|
|
os.close(descriptor)
|
|
|
|
|
|
def _endpoint(sink_url: str, suffix: str) -> str:
|
|
parts = urlsplit(sink_url)
|
|
base = parts.path.removesuffix("/api/usage/events")
|
|
return urlunsplit((parts.scheme, parts.netloc, f"{base}{suffix}", "", ""))
|
|
|
|
|
|
class UsageSpool:
|
|
"""Synchronous durable enqueue with an asynchronous at-least-once sender."""
|
|
|
|
def __init__(self) -> None:
|
|
self.sink_url = os.environ["EVOSCIENTIST_USAGE_SINK_URL"]
|
|
self.token = os.environ["EVOSCIENTIST_USAGE_SINK_TOKEN"]
|
|
self.deployment_id = os.environ["EVOSCIENTIST_DEPLOYMENT_ID"]
|
|
self.workspace_id = os.environ["EVOSCIENTIST_WORKSPACE_ID"]
|
|
self.root = Path(os.environ["EVOSCIENTIST_USAGE_SPOOL_DIR"])
|
|
if not self.root.is_absolute():
|
|
raise ValueError("EVOSCIENTIST_USAGE_SPOOL_DIR must be absolute")
|
|
self.tmp = self.root / "tmp"
|
|
self.pending = self.root / "pending"
|
|
self.inflight = self.root / "inflight"
|
|
self.quarantine = self.root / "quarantine"
|
|
for directory in (self.tmp, self.pending, self.inflight, self.quarantine):
|
|
directory.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
self.status_path = self.root / "status.json"
|
|
self.max_files = _int_env("EVOSCIENTIST_USAGE_SPOOL_MAX_FILES", 100_000)
|
|
self.max_bytes = _int_env("EVOSCIENTIST_USAGE_SPOOL_MAX_BYTES", 1_073_741_824)
|
|
self.max_event_bytes = _int_env("EVOSCIENTIST_USAGE_MAX_EVENT_BYTES", 262_144)
|
|
self.lease_seconds = _int_env("EVOSCIENTIST_USAGE_INFLIGHT_LEASE_SECONDS", 120)
|
|
self.heartbeat_interval = _float_env(
|
|
"EVOSCIENTIST_USAGE_HEARTBEAT_INTERVAL_SECONDS", 15
|
|
)
|
|
self.retry_initial = _float_env("EVOSCIENTIST_USAGE_RETRY_INITIAL_SECONDS", 1)
|
|
self.retry_max = _float_env("EVOSCIENTIST_USAGE_RETRY_MAX_SECONDS", 60)
|
|
self.connect_timeout = _float_env(
|
|
"EVOSCIENTIST_USAGE_HTTP_CONNECT_TIMEOUT_SECONDS", 1
|
|
)
|
|
self.request_timeout = _float_env("EVOSCIENTIST_USAGE_HTTP_TIMEOUT_SECONDS", 3)
|
|
self.unsupported_reprobe = _float_env(
|
|
"EVOSCIENTIST_USAGE_UNSUPPORTED_REPROBE_SECONDS", 300
|
|
)
|
|
self.schema_reprobe = _float_env(
|
|
"EVOSCIENTIST_USAGE_SCHEMA_REPROBE_SECONDS", 60
|
|
)
|
|
self.quarantine_retention_days = _int_env(
|
|
"EVOSCIENTIST_USAGE_QUARANTINE_RETENTION_DAYS", 90
|
|
)
|
|
self._state_lock = threading.Lock()
|
|
self._size_lock = threading.Lock()
|
|
self._cached_files = 0
|
|
self._cached_bytes = 0
|
|
self._size_cache_at = 0.0
|
|
self.first_loss_at: str | None = None
|
|
self.degraded_reason: str | None = None
|
|
self.last_error_code: str | None = None
|
|
self._load_status()
|
|
self._clean_expired_quarantine()
|
|
self._clean_stale_local_artifacts()
|
|
self._last_cleanup_at = time.monotonic()
|
|
self._stop = threading.Event()
|
|
self._wake = threading.Event()
|
|
self._thread = threading.Thread(
|
|
target=self._run, name="evoscientist-usage-sender", daemon=True
|
|
)
|
|
self._thread.start()
|
|
|
|
def _load_status(self) -> None:
|
|
try:
|
|
data = json.loads(self.status_path.read_text(encoding="utf-8"))
|
|
self.first_loss_at = data.get("first_loss_at")
|
|
self.degraded_reason = data.get("tracking_degraded_reason")
|
|
except (OSError, ValueError, TypeError):
|
|
pass
|
|
|
|
def _persist_status(self) -> None:
|
|
temporary = self.root / f"status.{os.getpid()}.{uuid.uuid4().hex}.tmp"
|
|
payload = json.dumps(
|
|
{
|
|
"first_loss_at": self.first_loss_at,
|
|
"tracking_degraded_reason": self.degraded_reason,
|
|
},
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
).encode()
|
|
try:
|
|
descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
|
try:
|
|
os.write(descriptor, payload)
|
|
os.fsync(descriptor)
|
|
finally:
|
|
os.close(descriptor)
|
|
os.replace(temporary, self.status_path)
|
|
_fsync_dir(self.root)
|
|
except OSError:
|
|
logger.exception("Could not persist usage tracking degraded status")
|
|
try:
|
|
temporary.unlink(missing_ok=True)
|
|
except OSError:
|
|
pass
|
|
|
|
def mark_degraded(self, reason: str) -> None:
|
|
with self._state_lock:
|
|
if self.first_loss_at is None:
|
|
self.first_loss_at = _now()
|
|
if self.degraded_reason is None:
|
|
self.degraded_reason = reason[:512]
|
|
self._persist_status()
|
|
self._wake.set()
|
|
|
|
def _spool_size(self) -> tuple[int, int]:
|
|
with self._size_lock:
|
|
if time.monotonic() - self._size_cache_at < 1:
|
|
return self._cached_files, self._cached_bytes
|
|
count = 0
|
|
size = 0
|
|
for directory in (self.pending, self.inflight, self.quarantine):
|
|
try:
|
|
for item in os.scandir(directory):
|
|
if item.is_file() and item.name.endswith(".json"):
|
|
count += 1
|
|
try:
|
|
size += item.stat().st_size
|
|
except OSError:
|
|
pass
|
|
except OSError:
|
|
continue
|
|
with self._size_lock:
|
|
self._cached_files = count
|
|
self._cached_bytes = size
|
|
self._size_cache_at = time.monotonic()
|
|
return count, size
|
|
|
|
def _adjust_spool_size(self, count: int, size: int) -> None:
|
|
with self._size_lock:
|
|
self._cached_files = max(0, self._cached_files + count)
|
|
self._cached_bytes = max(0, self._cached_bytes + size)
|
|
|
|
def enqueue(self, event: UsageEventV1) -> None:
|
|
try:
|
|
payload = json.dumps(
|
|
event.model_dump(mode="json"),
|
|
ensure_ascii=False,
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
).encode("utf-8")
|
|
if len(payload) > self.max_event_bytes:
|
|
self.mark_degraded("event_too_large")
|
|
logger.error("Usage event %s exceeds spool event limit", event.event_id)
|
|
return
|
|
count, size = self._spool_size()
|
|
if count >= self.max_files or size + len(payload) > self.max_bytes:
|
|
self.mark_degraded("spool_soft_limit_reached")
|
|
logger.error("Usage spool soft limit reached; event was not persisted")
|
|
return
|
|
key = hashlib.sha256(event.event_id.encode()).hexdigest()
|
|
temporary = self.tmp / f"{key}.{os.getpid()}.{uuid.uuid4().hex}.tmp"
|
|
descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
|
try:
|
|
os.write(descriptor, payload)
|
|
os.fsync(descriptor)
|
|
finally:
|
|
os.close(descriptor)
|
|
lock_path = self.root / f"{key}.lock"
|
|
deadline = time.monotonic() + 0.05
|
|
while True:
|
|
try:
|
|
lock_fd = os.open(
|
|
lock_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600
|
|
)
|
|
break
|
|
except FileExistsError:
|
|
if time.monotonic() < deadline:
|
|
time.sleep(0.001)
|
|
continue
|
|
os.replace(
|
|
temporary,
|
|
self.quarantine
|
|
/ f"{key}.lock-contention.{uuid.uuid4().hex}.json",
|
|
)
|
|
self._adjust_spool_size(1, len(payload))
|
|
_fsync_dir(self.quarantine)
|
|
logger.error(
|
|
"Usage event %s quarantined after spool lock contention",
|
|
event.event_id,
|
|
)
|
|
self.mark_degraded("spool_lock_contention")
|
|
self._wake.set()
|
|
return
|
|
try:
|
|
target = self.pending / f"{key}.json"
|
|
existing = target if target.exists() else self.inflight / target.name
|
|
if existing.exists():
|
|
if existing.read_bytes() != payload:
|
|
os.replace(
|
|
temporary,
|
|
self.quarantine / f"{key}.conflict.{uuid.uuid4().hex}.json",
|
|
)
|
|
self._adjust_spool_size(1, len(payload))
|
|
_fsync_dir(self.quarantine)
|
|
self.mark_degraded("local_event_payload_conflict")
|
|
else:
|
|
temporary.unlink(missing_ok=True)
|
|
else:
|
|
os.replace(temporary, target)
|
|
self._adjust_spool_size(1, len(payload))
|
|
_fsync_dir(self.pending)
|
|
finally:
|
|
os.close(lock_fd)
|
|
lock_path.unlink(missing_ok=True)
|
|
except Exception:
|
|
self.mark_degraded("spool_write_failed")
|
|
logger.exception("Usage capture could not persist a terminal event")
|
|
finally:
|
|
self._wake.set()
|
|
|
|
def _recover_stale_inflight(self) -> None:
|
|
cutoff = time.time() - self.lease_seconds
|
|
try:
|
|
items = list(self.inflight.glob("*.json"))
|
|
except OSError:
|
|
return
|
|
for item in items:
|
|
try:
|
|
if item.stat().st_mtime < cutoff:
|
|
os.replace(item, self.pending / item.name)
|
|
_fsync_dir(self.inflight)
|
|
_fsync_dir(self.pending)
|
|
except (FileNotFoundError, OSError):
|
|
continue
|
|
|
|
def _clean_expired_quarantine(self) -> None:
|
|
cutoff = time.time() - self.quarantine_retention_days * 86_400
|
|
removed = 0
|
|
for item in self.quarantine.glob("*.json"):
|
|
try:
|
|
if item.stat().st_mtime < cutoff:
|
|
item.unlink()
|
|
removed += 1
|
|
except (FileNotFoundError, OSError):
|
|
continue
|
|
if removed:
|
|
logger.warning("Removed %d expired usage quarantine events", removed)
|
|
_fsync_dir(self.quarantine)
|
|
|
|
def _clean_stale_local_artifacts(self) -> None:
|
|
cutoff = time.time() - self.lease_seconds
|
|
candidates = [*self.tmp.glob("*.tmp"), *self.root.glob("*.lock")]
|
|
for item in candidates:
|
|
try:
|
|
if item.stat().st_mtime < cutoff:
|
|
item.unlink()
|
|
except (FileNotFoundError, OSError):
|
|
continue
|
|
|
|
def _counts(self) -> tuple[int, int, int, int]:
|
|
def files(directory: Path) -> list[Path]:
|
|
try:
|
|
return list(directory.glob("*.json"))
|
|
except OSError:
|
|
return []
|
|
|
|
pending = files(self.pending)
|
|
inflight = files(self.inflight)
|
|
quarantine = files(self.quarantine)
|
|
size = 0
|
|
for item in pending + inflight + quarantine:
|
|
try:
|
|
size += item.stat().st_size
|
|
except OSError:
|
|
pass
|
|
return len(pending), len(inflight), len(quarantine), size
|
|
|
|
def _heartbeat(self, client: httpx.Client) -> None:
|
|
pending, inflight, quarantine, size = self._counts()
|
|
with self._state_lock:
|
|
body = {
|
|
"deployment_id": self.deployment_id,
|
|
"workspace_id": self.workspace_id,
|
|
"emitter_version": "2.0",
|
|
"schema_version": 1,
|
|
"sender_status": (
|
|
"degraded"
|
|
if self.first_loss_at or self.degraded_reason
|
|
else "healthy"
|
|
),
|
|
"spool_pending": pending,
|
|
"spool_inflight": inflight,
|
|
"spool_quarantined": quarantine,
|
|
"spool_bytes": size,
|
|
"first_loss_at": self.first_loss_at,
|
|
"tracking_degraded_reason": self.degraded_reason,
|
|
"last_error_code": self.last_error_code,
|
|
"sent_at": _now(),
|
|
}
|
|
try:
|
|
response = client.post(
|
|
_endpoint(self.sink_url, "/api/usage/sources/heartbeat"), json=body
|
|
)
|
|
if response.status_code >= 400:
|
|
self.last_error_code = f"heartbeat_http_{response.status_code}"
|
|
except httpx.HTTPError:
|
|
self.last_error_code = "heartbeat_unreachable"
|
|
|
|
def _capable(self, client: httpx.Client) -> bool:
|
|
try:
|
|
response = client.get(_endpoint(self.sink_url, "/api/usage/capabilities"))
|
|
if response.status_code != 200:
|
|
self.last_error_code = f"capabilities_http_{response.status_code}"
|
|
return False
|
|
versions = response.json().get("supported_schema_versions", [])
|
|
if 1 not in versions:
|
|
self.last_error_code = "schema_incompatible"
|
|
return False
|
|
self.last_error_code = None
|
|
return True
|
|
except (httpx.HTTPError, ValueError, TypeError):
|
|
self.last_error_code = "collector_unreachable"
|
|
return False
|
|
|
|
def _send_one(self, client: httpx.Client) -> bool:
|
|
pending: Path | None = None
|
|
lock_fd: int | None = None
|
|
lock_path: Path | None = None
|
|
try:
|
|
candidates = self.pending.glob("*.json")
|
|
for candidate in candidates:
|
|
candidate_lock = self.root / f"{candidate.stem}.lock"
|
|
try:
|
|
descriptor = os.open(
|
|
candidate_lock,
|
|
os.O_WRONLY | os.O_CREAT | os.O_EXCL,
|
|
0o600,
|
|
)
|
|
except FileExistsError:
|
|
continue
|
|
pending = candidate
|
|
lock_fd = descriptor
|
|
lock_path = candidate_lock
|
|
break
|
|
except OSError:
|
|
return False
|
|
if pending is None:
|
|
return False
|
|
inflight = self.inflight / pending.name
|
|
try:
|
|
os.replace(pending, inflight)
|
|
_fsync_dir(self.pending)
|
|
_fsync_dir(self.inflight)
|
|
except (FileNotFoundError, OSError):
|
|
return True
|
|
finally:
|
|
if lock_fd is not None:
|
|
os.close(lock_fd)
|
|
if lock_path is not None:
|
|
lock_path.unlink(missing_ok=True)
|
|
try:
|
|
response = client.post(
|
|
self.sink_url,
|
|
content=inflight.read_bytes(),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
if response.status_code == 200:
|
|
status = response.json().get("status")
|
|
if status in {"accepted", "duplicate"}:
|
|
event_size = inflight.stat().st_size
|
|
inflight.unlink(missing_ok=True)
|
|
self._adjust_spool_size(-1, -event_size)
|
|
_fsync_dir(self.inflight)
|
|
self.last_error_code = None
|
|
return True
|
|
if response.status_code in {400, 409, 413, 422}:
|
|
os.replace(inflight, self.quarantine / inflight.name)
|
|
_fsync_dir(self.inflight)
|
|
_fsync_dir(self.quarantine)
|
|
self.mark_degraded(f"collector_rejected_event_{response.status_code}")
|
|
self.last_error_code = f"event_http_{response.status_code}"
|
|
return True
|
|
self.last_error_code = f"event_http_{response.status_code}"
|
|
except (httpx.HTTPError, OSError, ValueError, TypeError):
|
|
self.last_error_code = "event_send_failed"
|
|
try:
|
|
os.replace(inflight, self.pending / inflight.name)
|
|
_fsync_dir(self.inflight)
|
|
_fsync_dir(self.pending)
|
|
except (FileNotFoundError, OSError):
|
|
pass
|
|
return False
|
|
|
|
def _run(self) -> None:
|
|
headers = {"Authorization": f"Bearer {self.token}"}
|
|
timeout = httpx.Timeout(self.request_timeout, connect=self.connect_timeout)
|
|
retry = self.retry_initial
|
|
last_heartbeat = 0.0
|
|
collector_capable = False
|
|
next_probe = 0.0
|
|
with httpx.Client(headers=headers, timeout=timeout) as client:
|
|
while not self._stop.is_set():
|
|
self._recover_stale_inflight()
|
|
now = time.monotonic()
|
|
if now - self._last_cleanup_at >= 3_600:
|
|
self._clean_expired_quarantine()
|
|
self._clean_stale_local_artifacts()
|
|
self._last_cleanup_at = now
|
|
if now - last_heartbeat >= self.heartbeat_interval:
|
|
self._heartbeat(client)
|
|
last_heartbeat = now
|
|
if not collector_capable and now >= next_probe:
|
|
collector_capable = self._capable(client)
|
|
if not collector_capable:
|
|
if self.last_error_code == "capabilities_http_404":
|
|
probe_delay = self.unsupported_reprobe
|
|
elif self.last_error_code == "schema_incompatible":
|
|
probe_delay = self.schema_reprobe
|
|
else:
|
|
probe_delay = min(self.retry_max, retry)
|
|
next_probe = time.monotonic() + probe_delay
|
|
if not collector_capable:
|
|
wait_for = max(0.05, next_probe - time.monotonic())
|
|
self._wake.wait(wait_for)
|
|
self._wake.clear()
|
|
retry = min(self.retry_max, retry * 2)
|
|
continue
|
|
progressed = self._send_one(client)
|
|
if progressed:
|
|
retry = self.retry_initial
|
|
continue
|
|
if self.last_error_code in {
|
|
"event_http_401",
|
|
"event_http_403",
|
|
"event_http_404",
|
|
"event_http_426",
|
|
}:
|
|
collector_capable = False
|
|
next_probe = time.monotonic() + (
|
|
self.unsupported_reprobe
|
|
if self.last_error_code == "event_http_404"
|
|
else self.schema_reprobe
|
|
)
|
|
self._wake.wait(
|
|
self.heartbeat_interval
|
|
if not any(self.pending.glob("*.json"))
|
|
else retry * random.uniform(0.8, 1.2)
|
|
)
|
|
self._wake.clear()
|
|
retry = min(self.retry_max, retry * 2)
|
|
|
|
def close(self) -> None:
|
|
self._stop.set()
|
|
self._wake.set()
|
|
if self._thread.is_alive():
|
|
self._thread.join(timeout=1)
|
|
|
|
|
|
_singleton_lock = threading.Lock()
|
|
_singleton: tuple[int, UsageSpool] | None = None
|
|
|
|
|
|
def get_usage_spool() -> UsageSpool:
|
|
global _singleton
|
|
pid = os.getpid()
|
|
with _singleton_lock:
|
|
if _singleton is None or _singleton[0] != pid:
|
|
_singleton = (pid, UsageSpool())
|
|
return _singleton[1]
|
|
|
|
|
|
def mark_tracking_degraded(reason: str) -> None:
|
|
if os.getenv("EVOSCIENTIST_USAGE_TRACKING", "").strip().lower() not in {
|
|
"1",
|
|
"true",
|
|
"yes",
|
|
"on",
|
|
}:
|
|
return
|
|
try:
|
|
get_usage_spool().mark_degraded(reason)
|
|
except Exception:
|
|
logger.exception("Could not mark usage tracking as degraded: %s", reason)
|
|
|
|
|
|
def _close_singleton() -> None:
|
|
if _singleton is not None:
|
|
_singleton[1].close()
|
|
|
|
|
|
atexit.register(_close_singleton)
|