102 lines
3.0 KiB
Python
102 lines
3.0 KiB
Python
"""Slot-scoped opaque strategy-instance state."""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from typing import Optional, Protocol
|
|
from uuid import UUID
|
|
|
|
_STATE_CAP_BYTES = int(os.environ.get("TREX_STRATEGY_STATE_CAP_BYTES", "4096"))
|
|
|
|
REASON_VERSION_CONFLICT = "version_conflict"
|
|
REASON_SLOT_INACTIVE = "slot_inactive"
|
|
REASON_STATE_TOO_LARGE = "state_too_large"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StrategyStateEntry:
|
|
"""One opaque entry in a slot's strategy_instance_id map."""
|
|
|
|
version: int
|
|
data: str
|
|
updated_at: datetime
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SetResult:
|
|
ok: bool
|
|
version: Optional[int] = None
|
|
reason: Optional[str] = None
|
|
|
|
|
|
def _byte_len(data: str) -> int:
|
|
return len(data.encode("utf-8"))
|
|
|
|
|
|
def _cas(
|
|
existing: Optional[StrategyStateEntry], *, expected_version: int, data: str, slot_active: bool
|
|
) -> SetResult:
|
|
if not slot_active:
|
|
return SetResult(ok=False, reason=REASON_SLOT_INACTIVE)
|
|
if _byte_len(data) > _STATE_CAP_BYTES:
|
|
return SetResult(ok=False, reason=REASON_STATE_TOO_LARGE)
|
|
stored_version = existing.version if existing is not None else 0
|
|
if expected_version != stored_version:
|
|
return SetResult(ok=False, reason=REASON_VERSION_CONFLICT, version=stored_version)
|
|
return SetResult(ok=True, version=stored_version + 1)
|
|
|
|
|
|
class StrategyStateStore(Protocol):
|
|
def set(
|
|
self,
|
|
slot_id: UUID,
|
|
strategy_instance_id: str,
|
|
*,
|
|
expected_version: int,
|
|
data: str,
|
|
slot_active: bool,
|
|
) -> SetResult: ...
|
|
|
|
def get_for_slots(self, slot_ids: list[UUID]) -> dict[UUID, dict[str, StrategyStateEntry]]: ...
|
|
|
|
|
|
class InMemoryStrategyStateStore:
|
|
"""Process-local slot-scoped state-map store."""
|
|
|
|
def __init__(self) -> None:
|
|
self._store: dict[UUID, dict[str, StrategyStateEntry]] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def set(
|
|
self,
|
|
slot_id: UUID,
|
|
strategy_instance_id: str,
|
|
*,
|
|
expected_version: int,
|
|
data: str,
|
|
slot_active: bool,
|
|
) -> SetResult:
|
|
with self._lock:
|
|
entries = self._store.setdefault(slot_id, {})
|
|
existing = entries.get(strategy_instance_id)
|
|
result = _cas(
|
|
existing, expected_version=expected_version, data=data, slot_active=slot_active
|
|
)
|
|
if not result.ok:
|
|
if not entries:
|
|
self._store.pop(slot_id, None)
|
|
return result
|
|
entries[strategy_instance_id] = StrategyStateEntry(
|
|
version=result.version, # type: ignore[arg-type]
|
|
data=data,
|
|
updated_at=datetime.now(tz=timezone.utc),
|
|
)
|
|
return result
|
|
|
|
def get_for_slots(self, slot_ids: list[UUID]) -> dict[UUID, dict[str, StrategyStateEntry]]:
|
|
wanted = set(slot_ids)
|
|
with self._lock:
|
|
return {sid: dict(entries) for sid, entries in self._store.items() if sid in wanted}
|