trex-core/trex_core/application/strategy_state.py

"""Slot strategy-state store (dcy-028 / REQ-079, BAT-042).

A strategy daemon (e.g. trex-tsl) accumulates opaque state across a slot's
occupancy — for a Trailing Stop Loss the peak price since entry. `evaluate()`
must stay pure, so the state cannot live in the daemon (lost on restart) nor in
an external store (orphan-cleanup problem). Since the state's lifecycle matches
the slot occupancy exactly, trex-core stores it server-side keyed by
`(slot_id, strategy_id)`, round-trips it on the slots snapshot (read), accepts
version-CAS writes via a dedicated RPC, and deletes it when the slot closes.

trex-core never interprets `data` (opaque UTF-8 string). It validates only
(1) valid UTF-8 and (2) a byte size cap — encoding (JSON/base64) is the
strategy's business (dcy-028 OQ-1).

Occupancy guard (dcy-028 R-1 / REQ-081): a permanent slot is reused, so two
occupancies of the same slot both start at version 0. A write computed against
occupancy X must not land on a later occupancy Y. The slot carries a monotonic
`occupancy_seq`; the CAS rejects a write whose `expected_occupancy_seq` differs
from the slot's current value (`occupancy_changed`). The version CAS alone
cannot catch this at the first write (0 vs 0).

Single-writer assumption: one daemon operates a given slot·strategy. The CAS is
for retry / duplicate / occupancy-turnover safety, not multi-writer concurrency.
"""
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

# dcy-028 OQ-1: hard ceiling for the opaque `data` blob, measured in UTF-8
# bytes. Operators target 1~2KB for the known TSL use case; 4KB is the cap.
_STATE_CAP_BYTES = int(os.environ.get("TREX_STRATEGY_STATE_CAP_BYTES", "4096"))

# Reject reason tokens (dcy-028 / REQ-080). The full enumerated set; consumers
# fail-closed on anything unrecognized.
REASON_OCCUPANCY_CHANGED = "occupancy_changed"
REASON_VERSION_CONFLICT = "version_conflict"
REASON_SLOT_INACTIVE = "slot_inactive"
REASON_STATE_TOO_LARGE = "state_too_large"


@dataclass(frozen=True)
class StrategyStateEntry:
    """One stored opaque state value for a (slot_id, strategy_id)."""

    strategy_id: str
    occupancy_seq: int
    version: int
    data: str
    updated_at: datetime


@dataclass(frozen=True)
class SetResult:
    """Outcome of a CAS set. ok=True carries the new version; else `reason`."""

    ok: bool
    version: Optional[int] = None
    reason: Optional[str] = None
    occupancy_seq: Optional[int] = None  # current slot occupancy_seq on conflict


def _byte_len(data: str) -> int:
    return len(data.encode("utf-8"))


class StrategyStateStore(Protocol):
    def set(
        self,
        slot_id: UUID,
        strategy_id: str,
        *,
        expected_occupancy_seq: int,
        expected_version: int,
        data: str,
        current_occupancy_seq: int,
        slot_active: bool,
    ) -> SetResult:
        """occupancy_seq + version CAS upsert of opaque `data` (dcy-028 FR-2)."""
        ...

    def get_for_slots(
        self, slot_ids: list[UUID], current_occupancy_seq: dict[UUID, int]
    ) -> dict[UUID, dict[str, StrategyStateEntry]]:
        """Batched read for the snapshot (C-1). Returns slot_id -> {strategy_id: entry}.

        Only entries whose stored occupancy_seq matches the slot's current
        occupancy_seq are returned (stale entries from a prior occupancy are
        invisible — defense in depth alongside close-delete, dcy-028 FR-3)."""
        ...

    def delete_all(self, slot_id: UUID) -> None:
        """Delete all strategy_state for a slot (close hook, dcy-028 C-3)."""
        ...


def _cas(
    existing: Optional[StrategyStateEntry],
    *,
    expected_occupancy_seq: int,
    expected_version: int,
    data: str,
    current_occupancy_seq: int,
    slot_active: bool,
) -> SetResult:
    """Pure CAS decision shared by all store implementations (dcy-028 FR-2)."""
    if not slot_active:
        return SetResult(ok=False, reason=REASON_SLOT_INACTIVE)
    if expected_occupancy_seq != current_occupancy_seq:
        return SetResult(
            ok=False, reason=REASON_OCCUPANCY_CHANGED, occupancy_seq=current_occupancy_seq
        )
    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 InMemoryStrategyStateStore:
    """Process-local strategy-state store (paper / tests / explicit double)."""

    def __init__(self) -> None:
        # (slot_id, strategy_id) -> StrategyStateEntry
        self._store: dict[tuple[UUID, str], StrategyStateEntry] = {}
        self._lock = threading.Lock()

    def set(
        self,
        slot_id: UUID,
        strategy_id: str,
        *,
        expected_occupancy_seq: int,
        expected_version: int,
        data: str,
        current_occupancy_seq: int,
        slot_active: bool,
    ) -> SetResult:
        with self._lock:
            existing = self._store.get((slot_id, strategy_id))
            result = _cas(
                existing,
                expected_occupancy_seq=expected_occupancy_seq,
                expected_version=expected_version,
                data=data,
                current_occupancy_seq=current_occupancy_seq,
                slot_active=slot_active,
            )
            if not result.ok:
                return result
            self._store[(slot_id, strategy_id)] = StrategyStateEntry(
                strategy_id=strategy_id,
                occupancy_seq=current_occupancy_seq,
                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], current_occupancy_seq: dict[UUID, int]
    ) -> dict[UUID, dict[str, StrategyStateEntry]]:
        wanted = set(slot_ids)
        out: dict[UUID, dict[str, StrategyStateEntry]] = {}
        with self._lock:
            for (sid, strat), entry in self._store.items():
                if sid not in wanted:
                    continue
                # stale-occupancy entries are invisible (dcy-028 FR-3 defense)
                if entry.occupancy_seq != current_occupancy_seq.get(sid):
                    continue
                out.setdefault(sid, {})[strat] = entry
        return out

    def delete_all(self, slot_id: UUID) -> None:
        with self._lock:
            for key in [k for k in self._store if k[0] == slot_id]:
                del self._store[key]