From 470a6af6a99fc76c20d94dc26c1958f3434b2b6d Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Fri, 24 Apr 2026 17:29:38 -0700 Subject: [PATCH] feat(core): add state_writer for buffered set_state persistence MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ``turnstone.core.state_writer.StateWriter`` buffers non-terminal ``update_workstream_state`` writes (last state per ws_id wins) and flushes them on a ~1s cadence (configurable). Terminal ERROR transitions and close()'s 'closed' write bypass the buffer. Bounded buffer (``max_buffer=10000`` default) evicts the oldest ws_id on insertion overflow — protects against unbounded growth when storage is unreachable. ``discard(ws_id)`` drops any pending buffered transition AND waits on a flush_lock for any in-flight write to complete; this is the close-path hook that preserves the bug-3 invariant (a closed row can't be resurrected by a buffered transient writing AFTER close's sync 'closed'). 13 unit tests cover coalescing, flush_now, bounded buffer, the discard / in-flight-flush wait, lifecycle (start/shutdown idempotence), wake-on-record latency, and resilience to storage errors poisoning subsequent flushes. --- tests/test_state_writer.py | 296 +++++++++++++++++++++++++++++++++ turnstone/core/state_writer.py | 205 +++++++++++++++++++++++ 2 files changed, 501 insertions(+) create mode 100644 tests/test_state_writer.py create mode 100644 turnstone/core/state_writer.py diff --git a/tests/test_state_writer.py b/tests/test_state_writer.py new file mode 100644 index 00000000..72bef787 --- /dev/null +++ b/tests/test_state_writer.py @@ -0,0 +1,296 @@ +"""Unit tests for ``turnstone.core.state_writer``. + +Tests cover the contract callers depend on: + +* Buffered transitions coalesce per ws_id (last state wins). +* ``flush_now=True`` bypasses the buffer (used for terminal ERROR + transitions and any other write that must be durable on return). +* ``discard`` drops pending and waits for any in-flight flush to + complete (the bug-3 invariant — close()'s sync ``closed`` write must + not be overtaken by a buffered transient). +* Bounded buffer evicts oldest under capacity pressure. +* DB error during flush doesn't poison the loop; subsequent flushes + still run. +* Shutdown drains any pending entries synchronously. +""" + +from __future__ import annotations + +import threading +import time + +from turnstone.core.state_writer import StateWriter + + +class _FakeStorage: + """Records update_workstream_state calls. Optionally raises or pauses.""" + + def __init__(self, *, raises: BaseException | None = None) -> None: + self.calls: list[tuple[str, str]] = [] + self.raises = raises + self._call_lock = threading.Lock() + # Optional gate to pin a write inside update_workstream_state + # so the test can race ``discard`` against an in-flight flush. + self.write_gate: threading.Event | None = None + # Set by the writer thread once it enters update_workstream_state. + self.write_started = threading.Event() + + def update_workstream_state(self, ws_id: str, state: str) -> None: + if self.write_gate is not None: + self.write_started.set() + self.write_gate.wait(timeout=2.0) + with self._call_lock: + self.calls.append((ws_id, state)) + if self.raises is not None: + raise self.raises + + +def _drain(writer: StateWriter) -> None: + """Trigger a single flush synchronously.""" + writer._flush_once() + + +# --------------------------------------------------------------------------- +# Coalescing + flush +# --------------------------------------------------------------------------- + + +def test_buffered_transitions_coalesce_per_ws_id() -> None: + storage = _FakeStorage() + writer = StateWriter(storage) + + writer.record("ws-1", "thinking") + writer.record("ws-1", "running") + writer.record("ws-1", "idle") + writer.record("ws-2", "thinking") + + _drain(writer) + # Only the latest state per ws_id should land. + assert sorted(storage.calls) == sorted([("ws-1", "idle"), ("ws-2", "thinking")]) + + +def test_flush_now_bypasses_buffer_and_writes_sync() -> None: + storage = _FakeStorage() + writer = StateWriter(storage) + + # Pre-buffer something for a different ws_id to prove the sync + # path doesn't drain the whole buffer. + writer.record("ws-other", "running") + + writer.record("ws-err", "error", flush_now=True) + # ws-err landed sync, ws-other still buffered. + assert ("ws-err", "error") in storage.calls + assert ("ws-other", "running") not in storage.calls + + _drain(writer) + assert ("ws-other", "running") in storage.calls + + +def test_flush_now_swallows_storage_error() -> None: + storage = _FakeStorage(raises=RuntimeError("db down")) + writer = StateWriter(storage) + # Should not raise — set_state path can't recover from a storage + # write failure mid-transition. + writer.record("ws-1", "error", flush_now=True) + + +# --------------------------------------------------------------------------- +# Bounded buffer +# --------------------------------------------------------------------------- + + +def test_bounded_buffer_evicts_oldest_on_capacity() -> None: + storage = _FakeStorage() + writer = StateWriter(storage, max_buffer=3) + + writer.record("ws-1", "running") + writer.record("ws-2", "running") + writer.record("ws-3", "running") + # ws-4 forces eviction of ws-1 (oldest). + writer.record("ws-4", "running") + + _drain(writer) + landed = {ws_id for ws_id, _ in storage.calls} + assert "ws-1" not in landed + assert {"ws-2", "ws-3", "ws-4"} <= landed + + +def test_bounded_buffer_update_existing_does_not_evict() -> None: + storage = _FakeStorage() + writer = StateWriter(storage, max_buffer=2) + + writer.record("ws-1", "running") + writer.record("ws-2", "running") + # Update existing — must not evict. + writer.record("ws-1", "idle") + + _drain(writer) + landed = dict(storage.calls) + assert landed["ws-1"] == "idle" + assert landed["ws-2"] == "running" + + +# --------------------------------------------------------------------------- +# Resilience +# --------------------------------------------------------------------------- + + +def test_storage_error_does_not_poison_subsequent_flushes() -> None: + storage = _FakeStorage(raises=RuntimeError("db blip")) + errors: list[Exception] = [] + writer = StateWriter(storage, on_flush_error=errors.append) + + writer.record("ws-1", "running") + _drain(writer) + # Error was surfaced via callback. + assert len(errors) == 1 + + # Storage recovers; next flush succeeds. + storage.raises = None + writer.record("ws-2", "idle") + _drain(writer) + assert ("ws-2", "idle") in storage.calls + + +# --------------------------------------------------------------------------- +# discard / close-race +# --------------------------------------------------------------------------- + + +def test_discard_drops_pending_buffered_state() -> None: + storage = _FakeStorage() + writer = StateWriter(storage) + + writer.record("ws-close", "running") + writer.discard("ws-close") + _drain(writer) + assert storage.calls == [] + + +def test_discard_waits_for_in_flight_flush_to_complete() -> None: + """The bug-3 invariant: ``close()`` calls ``discard`` BEFORE its + sync ``state='closed'`` write. If a flusher was mid-write for the + same ws_id, the flusher's write must complete BEFORE + ``discard`` returns — so ``close()``'s sync write strictly + follows the flusher's transient write, leaving 'closed' as the + final state. (If discard returned early, close's 'closed' write + could be overwritten by the flusher's late 'running' write.) + """ + storage = _FakeStorage() + storage.write_gate = threading.Event() + writer = StateWriter(storage) + + writer.record("ws-A", "running") + + # Kick off a flush in a background thread; it will block inside + # update_workstream_state on storage.write_gate. + flush_done = threading.Event() + + def _flush_in_bg() -> None: + writer._flush_once() + flush_done.set() + + flusher = threading.Thread(target=_flush_in_bg, daemon=True) + flusher.start() + assert storage.write_started.wait(timeout=1.0) + assert flush_done.is_set() is False # writer is pinned + + # Call discard concurrently — it must NOT return until the flush + # completes. + discard_done = threading.Event() + + def _discard_in_bg() -> None: + writer.discard("ws-A") + discard_done.set() + + discarder = threading.Thread(target=_discard_in_bg, daemon=True) + discarder.start() + # discard should be blocked on flush_lock. + time.sleep(0.05) + assert discard_done.is_set() is False, "discard returned before flusher released the write" + + # Release the writer; both threads should complete now. + storage.write_gate.set() + flusher.join(timeout=2.0) + discarder.join(timeout=2.0) + assert flush_done.is_set() + assert discard_done.is_set() + # The flusher's write went through. + assert ("ws-A", "running") in storage.calls + + +# --------------------------------------------------------------------------- +# Lifecycle +# --------------------------------------------------------------------------- + + +def test_start_starts_flusher_and_buffered_writes_land() -> None: + storage = _FakeStorage() + writer = StateWriter(storage, flush_interval=0.05) + writer.start() + try: + writer.record("ws-1", "running") + # Wait up to 1s for the flusher to drain. + for _ in range(20): + if storage.calls: + break + time.sleep(0.05) + assert ("ws-1", "running") in storage.calls + finally: + writer.shutdown(timeout=2.0) + + +def test_shutdown_drains_pending_synchronously() -> None: + storage = _FakeStorage() + # Long flush interval so no automatic drain happens. + writer = StateWriter(storage, flush_interval=60.0) + writer.start() + try: + writer.record("ws-1", "running") + writer.record("ws-2", "thinking") + finally: + writer.shutdown(timeout=2.0) + landed = {ws_id for ws_id, _ in storage.calls} + assert {"ws-1", "ws-2"} <= landed + + +def test_start_is_idempotent() -> None: + storage = _FakeStorage() + writer = StateWriter(storage, flush_interval=0.05) + writer.start() + first_thread = writer._thread + writer.start() + assert writer._thread is first_thread + writer.shutdown(timeout=2.0) + + +def test_shutdown_is_idempotent() -> None: + storage = _FakeStorage() + writer = StateWriter(storage, flush_interval=0.05) + writer.start() + writer.shutdown(timeout=2.0) + # Second shutdown is a no-op, must not raise. + writer.shutdown(timeout=2.0) + + +# --------------------------------------------------------------------------- +# Wake-on-record +# --------------------------------------------------------------------------- + + +def test_record_wakes_flusher_immediately() -> None: + """Single transitions get persisted within ~one round-trip rather + than waiting up to flush_interval seconds.""" + storage = _FakeStorage() + # Long interval — only the wake event should drive the flush. + writer = StateWriter(storage, flush_interval=10.0) + writer.start() + try: + writer.record("ws-1", "running") + for _ in range(30): + if storage.calls: + break + time.sleep(0.02) + assert ("ws-1", "running") in storage.calls + finally: + writer.shutdown(timeout=2.0) diff --git a/turnstone/core/state_writer.py b/turnstone/core/state_writer.py new file mode 100644 index 00000000..f51d0820 --- /dev/null +++ b/turnstone/core/state_writer.py @@ -0,0 +1,205 @@ +"""Buffered workstream-state persistence. + +``SessionManager.set_state`` previously held ``ws._lock`` across a +synchronous Postgres ``UPDATE`` for every ``thinking → running → idle +→ attention`` transition — multiple writes per turn, with +per-workstream observers serialising behind each round-trip. This +module replaces that with a write-behind buffer: + +* Non-terminal transitions buffer in a per-ws_id dict (last state wins + per ws_id — coalesced). +* A daemon flusher drains the buffer to ``storage.update_workstream_state`` + every ``flush_interval`` seconds (default 1.0s; loop wakes early on + ``record``). +* Terminal transitions (``ERROR``) and ``close()`` bypass the buffer + via ``record(..., flush_now=True)`` / ``discard(ws_id)`` — those + paths must be durable before observers see the transition. +* Bounded buffer (``max_buffer``): when full, the oldest ws_id's + pending state is evicted on insertion of a new ws_id. All entries + are non-terminal (terminals bypass), so eviction is safe. + +The bug-3 invariant the close path must keep holding (a closed ws +row can't be resurrected by a buffered transient state writing +"running" after close's sync "closed" write): + +1. ``close()`` acquires ``ws._lock`` and sets ``ws._closed = True``. +2. ``close()`` calls :meth:`StateWriter.discard` to drop any pending + buffered transition for the ws_id AND wait for any in-progress + flush to complete (so a flusher mid-write can't sneak through + AFTER ``close()``'s sync write). +3. ``close()`` writes ``state='closed'`` synchronously to storage. +4. Any later ``set_state`` for this ws_id sees ``ws._closed=True`` + under ``ws._lock`` and short-circuits — never reaches + :meth:`StateWriter.record`. +""" + +from __future__ import annotations + +import contextlib +import threading +from typing import TYPE_CHECKING, Any + +from turnstone.core.log import get_logger + +if TYPE_CHECKING: + from collections.abc import Callable + +log = get_logger(__name__) + + +class StateWriter: + """Buffered ``update_workstream_state`` writer. + + Construct once per process; pass to :class:`SessionManager`. + Lifecycle managed by the host's ASGI lifespan: call + :meth:`start` on startup, :meth:`shutdown` on teardown. + """ + + def __init__( + self, + storage: Any, + *, + flush_interval: float = 1.0, + max_buffer: int = 10_000, + on_flush_error: Callable[[Exception], None] | None = None, + ) -> None: + self._storage = storage + self._flush_interval = flush_interval + self._max_buffer = max_buffer + self._on_flush_error = on_flush_error + # ws_id → state.value. Python dict preserves insertion order, so + # iterating the buffer yields oldest-first for FIFO eviction. + self._buffer: dict[str, str] = {} + self._lock = threading.Lock() + # Held by the flusher while it's iterating + writing the + # snapshotted batch. ``discard`` waits on it so close() can + # ensure no stray write follows its sync ``state='closed'``. + self._flush_lock = threading.Lock() + self._wake = threading.Event() + self._stop = threading.Event() + self._thread: threading.Thread | None = None + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def record(self, ws_id: str, state: str, *, flush_now: bool = False) -> None: + """Buffer (or sync-write) a state transition. + + ``flush_now=True`` writes synchronously and bypasses the + buffer — used for ERROR transitions where durability matters + before any observer sees the state. Errors are logged and + swallowed to match the prior ``set_state`` behaviour (which + wrapped its DB call in a try/except for the same reason). + """ + if flush_now: + try: + self._storage.update_workstream_state(ws_id, state) + except Exception as exc: + log.debug( + "state_writer.flush_now_failed ws=%s", + ws_id[:8], + exc_info=True, + ) + self._notify_error(exc) + return + with self._lock: + # Bounded buffer. If a new ws_id arrives at capacity, drop + # the oldest pending entry. Updates to an existing key + # don't grow the buffer. + if ws_id not in self._buffer and len(self._buffer) >= self._max_buffer: + evict_id = next(iter(self._buffer)) + self._buffer.pop(evict_id) + log.warning( + "state_writer.buffer_full evicted=%s — DB unreachable?", + evict_id[:8], + ) + self._buffer[ws_id] = state + # Wake the flusher so a single transition gets persisted within + # ~one round-trip rather than waiting up to flush_interval. + # Coalescing across bursts still happens because the flusher + # snapshots the buffer atomically. + self._wake.set() + + def discard(self, ws_id: str) -> None: + """Drop any pending buffered state for ``ws_id`` and wait for any + in-progress flush to complete. + + Called by ``SessionManager.close`` (and ``close_idle``) under + ``ws._lock`` after ``ws._closed=True`` and BEFORE the sync + ``state='closed'`` write. After this returns, no buffered or + in-flight write for ``ws_id`` can land in storage AFTER the + caller's sync ``closed`` write. + """ + with self._lock: + self._buffer.pop(ws_id, None) + # If a flusher is currently writing, wait for it to finish. + # The flusher snapshots the buffer under self._lock then writes + # under self._flush_lock, so any write of ``ws_id`` already + # in-flight will complete before this returns. + with self._flush_lock: + pass + + def start(self) -> None: + """Start the background flusher thread. Idempotent.""" + if self._thread is not None and self._thread.is_alive(): + return + self._stop.clear() + self._wake.clear() + self._thread = threading.Thread( + target=self._loop, + name="state-writer-flush", + daemon=True, + ) + self._thread.start() + + def shutdown(self, *, timeout: float = 5.0) -> None: + """Stop the flusher and drain any pending writes synchronously. + + Idempotent — safe to call multiple times. Best-effort drain + even if the flusher thread doesn't exit cleanly. + """ + self._stop.set() + self._wake.set() + if self._thread is not None: + self._thread.join(timeout=timeout) + self._thread = None + # Final synchronous drain. The flusher may have exited mid-loop + # without picking up the last record(s); make sure they land. + self._flush_once() + + # ------------------------------------------------------------------ + # Flusher internals + # ------------------------------------------------------------------ + + def _loop(self) -> None: + while not self._stop.is_set(): + self._wake.wait(timeout=self._flush_interval) + self._wake.clear() + if self._stop.is_set(): + break + self._flush_once() + + def _flush_once(self) -> None: + with self._lock: + if not self._buffer: + return + pending = self._buffer + self._buffer = {} + with self._flush_lock: + for ws_id, state in pending.items(): + try: + self._storage.update_workstream_state(ws_id, state) + except Exception as exc: + log.debug( + "state_writer.flush_failed ws=%s", + ws_id[:8], + exc_info=True, + ) + self._notify_error(exc) + + def _notify_error(self, exc: Exception) -> None: + if self._on_flush_error is None: + return + with contextlib.suppress(Exception): + self._on_flush_error(exc)