ewilliams28/turnstone
1
0
Fork 0
mirror of https://github.com/turnstonelabs/turnstone.git synced 2026-08-27 06:14:48 -06:00
Code Issues Packages Projects Releases Wiki Activity

feat(core): add state_writer for buffered set_state persistence

Browse Source 
``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.
This commit is contained in:
Patrick Buckley
2026-04-24 17:29:38 -07:00
committed by Patrick Buckley
parent 4e791cfb15
commit 470a6af6a9
2 changed files with 501 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_state_writer.py
+296
View File
@@ -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)
turnstone/core/state_writer.py
+205
View File
@@ -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)