diff --git a/docs/docker.md b/docs/docker.md index cf1b6d51..eade60fc 100644 --- a/docs/docker.md +++ b/docs/docker.md @@ -84,6 +84,7 @@ Auth is always enabled. `TURNSTONE_JWT_SECRET` is required. |----------|---------|-------------| | `TURNSTONE_DB_BACKEND` | `sqlite` | Storage backend: `sqlite` or `postgresql` | | `TURNSTONE_DB_URL` | — | Database URL (e.g. `postgresql+psycopg://user:pass@postgres:5432/turnstone`). For SQLite, defaults to `/data/.turnstone.db` | +| `TURNSTONE_DB_LISTEN_URL` | (falls back to `TURNSTONE_DB_URL`) | Direct-to-PostgreSQL URL for the console's dedicated `LISTEN` connection. Set this when `TURNSTONE_DB_URL` points at PgBouncer in transaction pooling mode — LISTEN is session state and the transaction-pooled connection can't hold it. See [pgbouncer.md](pgbouncer.md). | | `TURNSTONE_DB_POOL_SIZE` | `2` | PostgreSQL connection pool size per process (default: 2 base + 3 overflow = 5 max) | | `POSTGRES_USER` | `turnstone` | PostgreSQL container username (used in default `TURNSTONE_DB_URL` for cluster/channel) | | `POSTGRES_PASSWORD` | — | PostgreSQL container password (required for production and cluster profiles) | diff --git a/docs/pgbouncer.md b/docs/pgbouncer.md index 486b05b2..0e079a73 100644 --- a/docs/pgbouncer.md +++ b/docs/pgbouncer.md @@ -199,4 +199,28 @@ does not support prepared statements. Turnstone's SQLAlchemy layer does not use server-side prepared statements by default, so this is not an issue. +**LISTEN / NOTIFY not supported in transaction mode** — PgBouncer's +transaction pooling assigns a real server connection only for the +duration of each transaction, then returns it to the pool. PostgreSQL +`LISTEN` is session state — a transaction-pooled client can't hold the +multi-statement session a long-lived `LISTEN` needs. The console's +`NotifyDispatcher` (reactive node discovery via the `services` channel) +therefore opens a **dedicated, direct-to-Postgres** connection that +bypasses PgBouncer. + +Configure via `config.toml` `[database] listen_url` (preferred — +co-located with the main `url`) or the `TURNSTONE_DB_LISTEN_URL` env var +(config.toml wins when both are set). Defaults to the main DB URL when +unset. + +| Setting | Behaviour | +|---|---| +| unset | Listener uses `TURNSTONE_DB_URL` as-is. Fine when PgBouncer is in **session** mode, or when there's no pooler in front of Postgres. With transaction-mode PgBouncer the listener's `LISTEN` will fail and the dispatcher retries with exponential backoff (1 s → 30 s cap) without ever succeeding. Reactive NOTIFY-driven node discovery is silently lost; the cluster collector's 60 s `_discovery_loop` is the only remaining backstop. | +| set to direct-to-PG URL (e.g. `postgresql://…/turnstone`) | Listener bypasses PgBouncer for its one dedicated connection. Reactive discovery latency drops from up-to-60 s to ~500 ms. The rest of the storage layer continues to go through PgBouncer in transaction mode. | + +Set this whenever PgBouncer is in transaction mode (the recommended +setting per this doc). The override only adds one long-lived PG +connection per console process — sized into the cluster's +`max_connections` budget alongside the pool. + See also: [Docker deployment](docker.md) · [Security](security.md) diff --git a/tests/test_console.py b/tests/test_console.py index 54c6ed0c..e0597077 100644 --- a/tests/test_console.py +++ b/tests/test_console.py @@ -3,6 +3,7 @@ import asyncio import json import queue +from typing import Any from unittest.mock import MagicMock import pytest @@ -149,6 +150,78 @@ class TestCollectorDiscovery: assert c._nodes["node-a"].started == 1234567890.0 +class TestCollectorNotifyWireIn: + """NotifyDispatcher-driven discovery — reactive node visibility.""" + + def test_start_subscribes_to_services_channel(self): + # Stub dispatcher records subscriptions without spawning threads. + class _StubDispatcher: + def __init__(self): + self.subscriptions: list[tuple[str, Any]] = [] + + def subscribe(self, channel, handler): + self.subscriptions.append((channel, handler)) + return lambda: None + + stub = _StubDispatcher() + storage = MockStorage() + c = ClusterCollector( + storage=storage, + discovery_interval=999, + notify_dispatcher=stub, + ) + try: + c.start() + assert len(stub.subscriptions) == 1 + channel, handler = stub.subscriptions[0] + assert channel == "services" + assert handler == c._on_services_notify + finally: + c.stop() + + def test_no_dispatcher_means_no_subscribe(self): + # Collector without a dispatcher (single-node / SQLite dev) just + # falls back to the 60 s discovery-loop polling — no error. + c = _make_collector(MockStorage()) + try: + c.start() + assert c._notify_unsubscribe is None + finally: + c.stop() + + def test_on_notify_runs_discovery(self): + # Construct a synthetic Notify and invoke the handler directly — + # asserts the wire-in delegates back to ``_discover_nodes``. + from turnstone.core.storage._notify import Notify + + storage = MockStorage() + c = _make_collector(storage) + c._running = True # bypass start() so we don't spawn threads + q: queue.Queue[dict[str, Any]] = queue.Queue() + c.register_listener(q) + + storage.services = [ + {"service_id": "node-z", "url": "http://z:8080", "metadata": "{}"}, + ] + c._on_services_notify(Notify(channel="services", payload="{}", pid=0)) + + event = q.get_nowait() + assert event["type"] == "node_joined" + assert event["node_id"] == "node-z" + + def test_on_notify_when_not_running_is_noop(self): + # If a stray notify arrives after stop, the handler doesn't run + # discovery on a half-torn-down collector. + from turnstone.core.storage._notify import Notify + + storage = MockStorage() + storage.services = [{"service_id": "node-y", "url": "http://y:8080", "metadata": "{}"}] + c = _make_collector(storage) + # _running stays False (never called start()). + c._on_services_notify(Notify(channel="services", payload="{}", pid=0)) + assert c.get_overview()["nodes"] == 0 + + class TestCollectorSnapshot: """Applying node_snapshot SSE events.""" diff --git a/tests/test_notify_dispatcher.py b/tests/test_notify_dispatcher.py new file mode 100644 index 00000000..59c41c3a --- /dev/null +++ b/tests/test_notify_dispatcher.py @@ -0,0 +1,386 @@ +"""Tests for the console-side ``NotifyDispatcher``. + +Exercises the dispatcher against the SQLite synthetic-sweep path so the +suite runs without a Postgres dependency. The PG path is shaped the +same way (same handler invocation semantics) — the only difference is +the underlying stream's wake-up source, which is covered separately in +``test_storage_notify.py::TestPostgresNotify``. +""" + +from __future__ import annotations + +import threading +import time + +import pytest + + +@pytest.fixture +def dispatcher_factory(storage): + """Yield a factory that constructs + tracks dispatchers for teardown.""" + from turnstone.console.notify_dispatcher import NotifyDispatcher + + created: list[NotifyDispatcher] = [] + + def _make(*, channels: list[str]) -> NotifyDispatcher: + d = NotifyDispatcher(storage, channels=channels) + created.append(d) + return d + + yield _make + + for d in created: + d.stop(timeout=2.0) + + +def _wait_for(predicate, deadline_sec: float = 3.0) -> bool: + """Poll ``predicate`` until True or timeout. Returns bool.""" + deadline = time.monotonic() + deadline_sec + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.02) + return False + + +class TestSubscribe: + def test_subscribe_registers_handler(self, dispatcher_factory, storage): + d = dispatcher_factory(channels=["alpha"]) + seen: list = [] + d.subscribe("alpha", lambda n: seen.append(n)) + d.start() + # Fire a notify via the storage layer — dispatcher delivers to handler. + storage.notify("alpha", "hello") + assert _wait_for(lambda: any(n.payload == "hello" for n in seen)) + + def test_subscribe_undeclared_channel_raises(self, dispatcher_factory): + d = dispatcher_factory(channels=["alpha"]) + with pytest.raises(ValueError, match="not declared"): + d.subscribe("beta", lambda n: None) + + def test_subscribe_returns_unsubscribe_callable(self, dispatcher_factory, storage): + d = dispatcher_factory(channels=["alpha"]) + seen: list = [] + unsub = d.subscribe("alpha", lambda n: seen.append(n)) + d.start() + storage.notify("alpha", "first") + assert _wait_for(lambda: any(n.payload == "first" for n in seen)) + unsub() + # After unsubscribe, the handler no longer fires. Drain old hits + # so the next notify-vs-handler-count check is unambiguous. + seen.clear() + storage.notify("alpha", "second") + # Give the dispatcher a beat to deliver if it were going to. + time.sleep(0.2) + assert not any(n.payload == "second" for n in seen) + + def test_construction_requires_at_least_one_channel(self, storage): + from turnstone.console.notify_dispatcher import NotifyDispatcher + + with pytest.raises(ValueError, match="at least one"): + NotifyDispatcher(storage, channels=[]) + + def test_duplicate_channels_deduplicated(self, dispatcher_factory): + d = dispatcher_factory(channels=["alpha", "alpha", "beta"]) + assert d.channels == ["alpha", "beta"] + + +class TestDispatch: + def test_multiple_handlers_each_invoked(self, dispatcher_factory, storage): + d = dispatcher_factory(channels=["alpha"]) + seen_a: list = [] + seen_b: list = [] + d.subscribe("alpha", lambda n: seen_a.append(n)) + d.subscribe("alpha", lambda n: seen_b.append(n)) + d.start() + storage.notify("alpha", "shared") + assert _wait_for(lambda: seen_a and seen_b) + assert seen_a[0].payload == "shared" + assert seen_b[0].payload == "shared" + + def test_handler_exception_does_not_break_dispatch(self, dispatcher_factory, storage): + d = dispatcher_factory(channels=["alpha"]) + survived: list = [] + + def _broken(_n): + msg = "boom" + raise RuntimeError(msg) + + d.subscribe("alpha", _broken) + d.subscribe("alpha", lambda n: survived.append(n)) + d.start() + storage.notify("alpha", "after_broken") + # The second handler runs even though the first raised. + assert _wait_for(lambda: any(n.payload == "after_broken" for n in survived)) + + def test_dispatch_filters_by_channel(self, dispatcher_factory, storage): + d = dispatcher_factory(channels=["alpha", "beta"]) + seen_a: list = [] + seen_b: list = [] + d.subscribe("alpha", lambda n: seen_a.append(n)) + d.subscribe("beta", lambda n: seen_b.append(n)) + d.start() + storage.notify("alpha", "for_a") + storage.notify("beta", "for_b") + assert _wait_for(lambda: seen_a and seen_b) + assert all(n.payload == "for_a" for n in seen_a) + assert all(n.payload == "for_b" for n in seen_b) + + +class TestReconnect: + """Reconnect + synthetic ``reconcile`` notify on stream-open success. + + Uses a stub storage that owns its own listen stream so the test can + drive a controlled stream-error sequence — the SQLite path can't + raise :class:`NotifyConnectionError`, and the PG path requires a + real database outage to exercise this code, neither of which fits a + unit test. The dispatcher's threading and reconcile-pending logic + are storage-agnostic — the dispatcher sees the same + :class:`NotifyStream` Protocol regardless of backend. + """ + + def test_reconcile_fires_after_reopen_not_before(self): + from turnstone.console.notify_dispatcher import NotifyDispatcher + from turnstone.core.storage._notify import Notify, NotifyConnectionError + + # State machine: open -> first poll raises NotifyConnectionError + # -> dispatcher waits backoff then reopens -> second open's first + # poll blocks forever (test stops the dispatcher before then). + # The fix: synthetic reconcile fires AFTER the second open + # succeeds, not after the first open fails. + sequence: list[str] = [] + reopen_event = threading.Event() + + class _StubStream: + def __init__(self, fail_first_poll: bool): + self._fail = fail_first_poll + self._closed = False + + def poll(self, _timeout): + if self._closed: + return [] + if self._fail: + self._fail = False + sequence.append("poll_raises") + msg = "fake-disconnect" + raise NotifyConnectionError(msg) + sequence.append("poll_returns") + # Block until close to simulate a quiet steady-state. + time.sleep(0.5) + return [] + + def close(self): + self._closed = True + + class _StubStorage: + def __init__(self): + self._open_count = 0 + + def listen(self, _channels): + import contextlib as _contextlib + + @_contextlib.contextmanager + def _cm(): + self._open_count += 1 + sequence.append(f"open_{self._open_count}") + if self._open_count == 2: + reopen_event.set() + stream = _StubStream(fail_first_poll=(self._open_count == 1)) + try: + yield stream + finally: + stream.close() + + return _cm() + + # Speed up backoff so the reopen happens promptly in the test. + import turnstone.console.notify_dispatcher as nd_mod + + original_backoff = nd_mod._RECONNECT_BACKOFF_INITIAL + nd_mod._RECONNECT_BACKOFF_INITIAL = 0.05 + try: + d = NotifyDispatcher(_StubStorage(), channels=["alpha"]) + got: list[Notify] = [] + d.subscribe("alpha", lambda n: got.append(n)) + d.start() + try: + # Wait for the second open (post-reconnect). + assert reopen_event.wait(3.0), "dispatcher did not reopen after disconnect" + # Reconcile should be delivered shortly after the reopen. + deadline = time.monotonic() + 2.0 + while time.monotonic() < deadline: + if any(n.payload == "reconcile" for n in got): + break + time.sleep(0.02) + assert any(n.payload == "reconcile" for n in got), ( + f"no reconcile delivered; sequence={sequence}, got={got}" + ) + # The reconcile must NOT fire before the second open — + # if it did, the index of 'open_2' in sequence would + # come after any reconcile-emitting work. Check ordering: + # 'open_1' < 'poll_raises' < 'open_2' (synthesize happens + # inside the with-block of the SECOND open). + ix_open_1 = sequence.index("open_1") + ix_raises = sequence.index("poll_raises") + ix_open_2 = sequence.index("open_2") + assert ix_open_1 < ix_raises < ix_open_2 + finally: + d.stop(timeout=2.0) + finally: + nd_mod._RECONNECT_BACKOFF_INITIAL = original_backoff + + def test_generic_exception_path_also_synthesizes_reconcile(self): + """Exceptions thrown during ``listen()`` (not via stream.poll) still trigger reconcile. + + Models the ``psycopg.connect()`` / initial ``LISTEN`` failure + shape, which doesn't go through the stream's exception + translator and would hit the generic ``except Exception`` + branch. Pre-fix, that branch emitted no reconcile. + """ + from turnstone.console.notify_dispatcher import NotifyDispatcher + + reopen_event = threading.Event() + + class _StubStream: + def __init__(self): + self._closed = False + + def poll(self, _timeout): + if self._closed: + return [] + time.sleep(0.5) + return [] + + def close(self): + self._closed = True + + class _StubStorage: + def __init__(self): + self._open_count = 0 + + def listen(self, _channels): + import contextlib as _contextlib + + self._open_count += 1 + if self._open_count == 1: + # First open raises a generic exception (e.g. + # ``psycopg.OperationalError`` from a failed connect) + # — landing in the dispatcher's generic except branch. + msg = "fake-connect-failure" + raise RuntimeError(msg) + + @_contextlib.contextmanager + def _cm(): + reopen_event.set() + stream = _StubStream() + try: + yield stream + finally: + stream.close() + + return _cm() + + import turnstone.console.notify_dispatcher as nd_mod + + original_backoff = nd_mod._RECONNECT_BACKOFF_INITIAL + nd_mod._RECONNECT_BACKOFF_INITIAL = 0.05 + try: + d = NotifyDispatcher(_StubStorage(), channels=["alpha"]) + got: list = [] + d.subscribe("alpha", lambda n: got.append(n)) + d.start() + try: + assert reopen_event.wait(3.0), "dispatcher did not reopen after generic exception" + deadline = time.monotonic() + 2.0 + while time.monotonic() < deadline: + if any(n.payload == "reconcile" for n in got): + break + time.sleep(0.02) + assert any(n.payload == "reconcile" for n in got), ( + "no reconcile delivered after generic-exception recovery" + ) + finally: + d.stop(timeout=2.0) + finally: + nd_mod._RECONNECT_BACKOFF_INITIAL = original_backoff + + +class TestCoalescing: + """Same-channel burst collapses to one handler invocation per batch.""" + + def test_burst_coalesces_to_one_handler_call_per_channel(self, dispatcher_factory, storage): + d = dispatcher_factory(channels=["alpha"]) + invocations: list = [] + # Slow handler to ensure all bursts queue up before the first + # call returns — gives the dispatch loop time to coalesce. + coalesce_gate = threading.Event() + + def _slow_handler(n): + invocations.append(n) + coalesce_gate.wait(0.05) + + d.subscribe("alpha", _slow_handler) + d.start() + # Burst of 10 notifies on the same channel — should coalesce + # down to many fewer handler invocations. + for i in range(10): + storage.notify("alpha", str(i)) + # Wait until the dispatch settles (handler is called at least once + # and the queue empties). + deadline = time.monotonic() + 2.0 + while time.monotonic() < deadline: + if invocations and d._dispatch_queue.empty(): + time.sleep(0.1) # allow any final coalesced call to land + break + time.sleep(0.02) + coalesce_gate.set() + # At least one handler call; well fewer than 10 (coalescing + # collapsed the burst). Exact count depends on timing — typical + # is 1-2 invocations per burst on a fast machine. + assert invocations, "handler never fired" + assert len(invocations) < 10, ( + f"expected coalescing to collapse burst of 10; got {len(invocations)} invocations" + ) + + +class TestLifecycle: + def test_start_is_idempotent(self, dispatcher_factory): + d = dispatcher_factory(channels=["alpha"]) + d.start() + d.start() # No-op, no thread doubling + # Single listener + single dispatch thread are spawned regardless. + # Inspect by name so we don't depend on the exact thread count of + # the test runner. + listener_threads = [ + t for t in threading.enumerate() if t.name == "notify-dispatcher-listener" + ] + dispatch_threads = [ + t for t in threading.enumerate() if t.name == "notify-dispatcher-dispatch" + ] + assert len(listener_threads) == 1 + assert len(dispatch_threads) == 1 + + def test_stop_is_idempotent(self, dispatcher_factory): + d = dispatcher_factory(channels=["alpha"]) + d.start() + d.stop(timeout=2.0) + d.stop(timeout=2.0) # No-op, no error + + def test_stop_without_start_is_noop(self, dispatcher_factory): + d = dispatcher_factory(channels=["alpha"]) + d.stop(timeout=1.0) # No-op, no thread to join + + def test_stop_joins_threads(self, dispatcher_factory): + d = dispatcher_factory(channels=["alpha"]) + d.start() + # Capture thread references then stop and assert they exited. + threads_before = [ + t + for t in threading.enumerate() + if t.name in {"notify-dispatcher-listener", "notify-dispatcher-dispatch"} + ] + assert threads_before + d.stop(timeout=3.0) + time.sleep(0.05) + for t in threads_before: + assert not t.is_alive(), f"{t.name} still alive after stop" diff --git a/tests/test_storage_notify.py b/tests/test_storage_notify.py new file mode 100644 index 00000000..a59332cd --- /dev/null +++ b/tests/test_storage_notify.py @@ -0,0 +1,212 @@ +"""Tests for the storage layer's cross-process ``notify`` / ``listen`` API. + +Covers SQLite (synthetic-sweep + in-process fan-out) and PostgreSQL +(real ``LISTEN``/``NOTIFY``). The PG-only cases are gated on the +``--storage-backend=postgresql`` flag so they no-op on default CI runs. +""" + +from __future__ import annotations + +import threading +import time + +import pytest + + +def _drain_until(stream, predicate, deadline_sec: float = 5.0): + """Poll ``stream`` until ``predicate`` matches one of the drained notifies. + + Returns the matching notify or raises ``TimeoutError``. Tests use + this so timing flakes against the bounded-blocking ``poll`` shape + don't masquerade as logic bugs. + """ + deadline = time.monotonic() + deadline_sec + while time.monotonic() < deadline: + remaining = max(0.05, deadline - time.monotonic()) + for n in stream.poll(min(0.5, remaining)): + if predicate(n): + return n + msg = "no matching notify drained before deadline" + raise TimeoutError(msg) + + +class TestSqliteNotify: + """SQLite path: in-process fan-out + synthetic sweep.""" + + def test_notify_no_listeners_is_noop(self, storage): + # No exception, no side effect — safe to always call from dispatch. + storage.notify("services", '{"op": "INSERT"}') + + def test_notify_delivers_to_in_process_listener(self, storage): + with storage.listen(["services"]) as stream: + storage.notify("services", '{"op": "INSERT"}') + got = _drain_until(stream, lambda n: n.payload == '{"op": "INSERT"}') + assert got.channel == "services" + assert got.pid == 0 + + def test_notify_filters_by_channel(self, storage): + with storage.listen(["services"]) as stream: + storage.notify("other_channel", "ignored") + storage.notify("services", "wanted") + got = _drain_until(stream, lambda n: True) + assert got.payload == "wanted" + + def test_multiple_listeners_each_get_event(self, storage): + # Two streams open on the same channel; each gets its own copy. + with storage.listen(["services"]) as s1, storage.listen(["services"]) as s2: + storage.notify("services", "broadcast") + got1 = _drain_until(s1, lambda n: True) + got2 = _drain_until(s2, lambda n: True) + assert got1.payload == "broadcast" + assert got2.payload == "broadcast" + + def test_close_stops_stream(self, storage): + with storage.listen(["services"]) as stream: + pass + # After context exit, the stream is closed; poll returns [] without + # blocking. A second close() is idempotent. + assert stream.poll(0.05) == [] + stream.close() + + def test_synthetic_sweep_emits_after_interval(self, storage): + # Force a short sweep interval via direct attribute override — + # the production default (``_SQLITE_NOTIFY_SWEEP_INTERVAL``) is + # tuned for an idle dev backstop and is way too long for a test. + with storage.listen(["services"]) as stream: + stream._sweep_interval = 0.1 + # First poll: not yet at the interval, so likely empty. + stream.poll(0.05) + # Wait past the interval, then poll again — should emit a + # synthetic-sweep notify per declared channel. + time.sleep(0.15) + got = _drain_until(stream, lambda n: n.payload == "sweep") + assert got.channel == "services" + assert got.payload == "sweep" + + def test_empty_channel_list_yields_empty_stream(self, storage): + with storage.listen([]) as stream: + # No channels — poll returns [] regardless of how long we wait. + assert stream.poll(0.05) == [] + + +# --------------------------------------------------------------------------- +# PostgreSQL path — gated on --storage-backend=postgresql. +# --------------------------------------------------------------------------- + + +@pytest.fixture +def _is_postgres(storage): + """Skip the wrapped test when the active backend isn't Postgres.""" + if storage.__class__.__name__ != "PostgreSQLBackend": + pytest.skip("PostgreSQL-specific test") + return True + + +class TestPostgresNotify: + def test_round_trip(self, storage, _is_postgres): + # Open a listener, fire a notify on a regular pooled connection, + # drain the listener within a reasonable bound (PG NOTIFY is + # typically sub-100ms on a local socket). + with storage.listen(["pytest_round_trip"]) as stream: + # Tiny sleep so the LISTEN settles before the NOTIFY fires — + # otherwise the notify can arrive on the connection before + # the LISTEN is registered (race only visible in tests). + time.sleep(0.05) + storage.notify("pytest_round_trip", '{"hello": "world"}') + got = _drain_until(stream, lambda n: True, deadline_sec=3.0) + assert got.channel == "pytest_round_trip" + assert got.payload == '{"hello": "world"}' + assert got.pid > 0 + + def test_concurrent_notifies_all_arrive(self, storage, _is_postgres): + with storage.listen(["pytest_concurrent"]) as stream: + time.sleep(0.05) + for i in range(5): + storage.notify("pytest_concurrent", str(i)) + seen: set[str] = set() + deadline = time.monotonic() + 3.0 + while len(seen) < 5 and time.monotonic() < deadline: + for n in stream.poll(0.2): + seen.add(n.payload) + assert seen == {"0", "1", "2", "3", "4"} + + def test_close_aborts_blocked_poll(self, storage, _is_postgres): + # poll() should return promptly once close() runs on another thread. + with storage.listen(["pytest_close"]) as stream: + done = threading.Event() + result: list[list] = [] + + def _poll_long(): + result.append(stream.poll(5.0)) + done.set() + + t = threading.Thread(target=_poll_long, daemon=True) + t.start() + time.sleep(0.1) + stream.close() + assert done.wait(2.0), "close() did not unblock poll()" + # No notify arrived, so the polled batch is empty — but the + # poll loop must have exited well under the 5 s timeout. + assert result == [[]] + + +class TestServicesTriggerFilter: + """Migration 053's trigger: fires on real changes, quiet on heartbeats. + + PG-only — the SQLite path has no trigger and is covered by + :class:`TestSqliteNotify`. Verifies the in-trigger ``IS NOT DISTINCT + FROM`` filter — a heartbeat-only UPDATE (same url + same metadata, + only ``last_heartbeat`` changed) must NOT emit a NOTIFY, since + ``register_service`` runs the same UPSERT on every 30 s tick × N + nodes and the channel would otherwise flood. + """ + + def test_insert_fires_notify(self, storage, _is_postgres): + with storage.listen(["services"]) as stream: + time.sleep(0.05) + storage.register_service("server", "pytest-trigger-node", "http://127.0.0.1:1") + got = _drain_until(stream, lambda n: True, deadline_sec=3.0) + assert got.channel == "services" + assert '"op": "INSERT"' in got.payload or "INSERT" in got.payload + # Cleanup so concurrent suites don't pick up the row. + storage.deregister_service("server", "pytest-trigger-node") + + def test_delete_fires_notify(self, storage, _is_postgres): + storage.register_service("server", "pytest-trigger-node-del", "http://127.0.0.1:2") + with storage.listen(["services"]) as stream: + time.sleep(0.05) + storage.deregister_service("server", "pytest-trigger-node-del") + got = _drain_until(stream, lambda n: True, deadline_sec=3.0) + assert "DELETE" in got.payload + + def test_url_change_update_fires_notify(self, storage, _is_postgres): + storage.register_service("server", "pytest-trigger-node-url", "http://127.0.0.1:3") + with storage.listen(["services"]) as stream: + time.sleep(0.05) + # UPSERT with different url — UPDATE path with url diff, + # trigger must fire. + storage.register_service("server", "pytest-trigger-node-url", "http://127.0.0.1:9") + got = _drain_until(stream, lambda n: True, deadline_sec=3.0) + assert "UPDATE" in got.payload + storage.deregister_service("server", "pytest-trigger-node-url") + + def test_heartbeat_only_update_is_quiet(self, storage, _is_postgres): + # Open the LISTEN session FIRST so PG delivers the INSERT NOTIFY + # to this connection — pg_notify routes only to sessions that + # have LISTENed at COMMIT time, so an INSERT committed before the + # listen opens would be lost and the drain would time out instead + # of exercising the heartbeat-quiet check below. + with storage.listen(["services"]) as stream: + time.sleep(0.05) + storage.register_service("server", "pytest-trigger-node-hb", "http://127.0.0.1:4") + # Drain the INSERT notify so subsequent polls see only what + # heartbeats emit (if anything). + _drain_until(stream, lambda n: True, deadline_sec=2.0) + # Now fire a heartbeat tick — same url + same metadata, + # only last_heartbeat updates. Trigger must NOT emit. + storage.heartbeat_service("server", "pytest-trigger-node-hb") + # Poll long enough that any spurious notify would have + # arrived; the channel must stay silent. + spurious = stream.poll(0.5) + assert spurious == [], f"heartbeat-only update emitted unexpected notify: {spurious}" + storage.deregister_service("server", "pytest-trigger-node-hb") diff --git a/turnstone.example.toml b/turnstone.example.toml index d3072914..3b3d9aa1 100644 --- a/turnstone.example.toml +++ b/turnstone.example.toml @@ -62,6 +62,13 @@ [database] # url = "" # postgres://user:pass@host/db or /path/to.db # env: TURNSTONE_DB_URL +# listen_url = "" # direct-to-postgres URL for the console's + # dedicated LISTEN connection. Set this when + # `url` points at pgbouncer in transaction + # pooling mode (LISTEN holds session state and + # is incompatible with transaction pooling — + # see docs/pgbouncer.md). Defaults to `url` + # when unset. env: TURNSTONE_DB_LISTEN_URL # SSL params (passed through to SQLAlchemy connection): # sslmode = "prefer" # disable, allow, prefer, require, verify-ca, verify-full # sslrootcert = "" # path to CA cert for verify-ca/verify-full diff --git a/turnstone/console/collector.py b/turnstone/console/collector.py index 781ae07c..8cd6c082 100644 --- a/turnstone/console/collector.py +++ b/turnstone/console/collector.py @@ -25,9 +25,13 @@ import httpx_sse from turnstone.core.workstream import WorkstreamKind if TYPE_CHECKING: + from collections.abc import Callable + from turnstone.console.metrics import ConsoleMetrics + from turnstone.console.notify_dispatcher import NotifyDispatcher from turnstone.console.router import ConsoleRouter from turnstone.core.auth import ServiceTokenManager + from turnstone.core.storage._notify import Notify from turnstone.core.storage._protocol import StorageBackend log = logging.getLogger("turnstone.console.collector") @@ -73,6 +77,7 @@ class ClusterCollector: tls_cert: tuple[str, str] | None = None, router: ConsoleRouter | None = None, console_metrics: ConsoleMetrics | None = None, + notify_dispatcher: NotifyDispatcher | None = None, ): self._storage = storage self._discovery_interval = discovery_interval @@ -82,6 +87,8 @@ class ClusterCollector: self._console_metrics = console_metrics self._tls_verify = tls_verify self._tls_cert = tls_cert + self._notify_dispatcher = notify_dispatcher + self._notify_unsubscribe: Callable[[], None] | None = None self._lock = threading.Lock() self._nodes: dict[str, NodeSnapshot] = {} @@ -128,6 +135,15 @@ class ClusterCollector: def start(self) -> None: """Start background threads.""" self._running = True + # Subscribe to the ``services`` channel for reactive node discovery. + # NOTIFY-driven wake-ups bring new-node visibility from up-to-60 s + # (next discovery tick) down to ~500 ms on Postgres; the 60 s + # discovery loop still runs as the backstop for crash-shaped node + # loss (NOTIFY only fires on actual writes, not on crash exits). + if self._notify_dispatcher is not None: + self._notify_unsubscribe = self._notify_dispatcher.subscribe( + "services", self._on_services_notify + ) for target, name in [ (self._discovery_loop, "console-discovery"), (self._sse_manager_thread, "console-sse"), @@ -145,6 +161,10 @@ class ClusterCollector: its ``finally`` cleanup (cancel tasks, close AsyncClient). """ self._running = False + if self._notify_unsubscribe is not None: + with contextlib.suppress(Exception): + self._notify_unsubscribe() + self._notify_unsubscribe = None # Request cancellation of all SSE tasks so they don't block the # manager's cleanup. The manager coroutine exits when _running is # False and handles remaining task cancellation in its finally block. @@ -156,6 +176,26 @@ class ClusterCollector: t.join(timeout=5) log.info("ClusterCollector stopped") + def _on_services_notify(self, notify: Notify) -> None: + """Run a discovery tick when the ``services`` channel fires. + + The dispatcher delivers both real Postgres notifications and + synthetic ``reconcile`` wake-ups after a reconnect — both shape + the same way: re-read ``services`` and diff against in-memory + state. Re-uses :meth:`_discover_nodes` so the timer-driven + backstop and the NOTIFY-driven fast-path share one code path. + """ + from turnstone.core.storage._registry import StorageUnavailableError + + if not self._running: + return + try: + self._discover_nodes() + except StorageUnavailableError: + pass # already logged by storage layer + except Exception: + log.exception("Node discovery error (notify-driven)") + def _fanout(self, event: dict[str, Any]) -> None: """Copy an event to all registered SSE listener queues.""" with self._listeners_lock: diff --git a/turnstone/console/notify_dispatcher.py b/turnstone/console/notify_dispatcher.py new file mode 100644 index 00000000..3a713393 --- /dev/null +++ b/turnstone/console/notify_dispatcher.py @@ -0,0 +1,326 @@ +"""Console-side multiplexer for PostgreSQL ``LISTEN``/``NOTIFY`` events. + +Holds a single dedicated listen connection (via :meth:`StorageBackend.listen`), +drains it on a listener thread, and fans notifications out to per-channel +handlers on a dedicated dispatch thread so a slow handler doesn't back up +the connection. + +Consumers register at construction time by passing their channel in +:attr:`channels`, then call :meth:`subscribe` to attach a handler. +Registering an undeclared channel raises — the construction list is the +single source of truth so wire-in is explicit (each future consumer +touches the dispatcher construction call site at +``turnstone/console/server.py::main`` to add its channel). + +On connection loss the listener wakes its handlers with a synthetic +``Notify(channel, payload="reconcile", pid=0)`` so every consumer +re-reads the underlying rows; their normal "reconcile on any wake-up" +code path covers both real notifications and reconnect recovery +identically. +""" + +from __future__ import annotations + +import contextlib +import queue +import threading +import time +from typing import TYPE_CHECKING + +from turnstone.core.log import get_logger +from turnstone.core.storage._notify import Notify, NotifyConnectionError + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + + from turnstone.core.storage._protocol import StorageBackend + + +log = get_logger(__name__) + + +# Backoff (seconds) between reconnect attempts after :class:`NotifyConnectionError`. +# Doubles each failure, capped at the max — long enough that a Postgres outage +# doesn't burn CPU on reconnect spins, short enough that recovery is fast. +_RECONNECT_BACKOFF_INITIAL: float = 1.0 +_RECONNECT_BACKOFF_MAX: float = 30.0 + +# Poll cadence on the listener thread. Short enough that ``stop`` lands +# promptly without joining a long-blocked notifies() call; long enough +# that we don't burn CPU on empty polls. +_LISTENER_POLL_TIMEOUT: float = 1.0 + +# Cap on the inter-thread dispatch queue. Drops oldest if a slow handler +# falls behind (logs once per drop bucket). Sized larger than the expected +# steady-state notification rate (services trigger fires only on +# register/restart/deregister — order of hundreds per hour at the 100-node +# design ceiling). +_DISPATCH_QUEUE_MAX: int = 1024 + + +class NotifyDispatcher: + """Holds the dedicated listen connection and fans events to handlers. + + Lifecycle: construct with the declared channel list, attach + handlers via :meth:`subscribe`, then call :meth:`start`. :meth:`stop` + closes the connection and joins the worker threads. Idempotent in + both directions so console teardown can call stop unconditionally. + """ + + def __init__(self, storage: StorageBackend, channels: Iterable[str]) -> None: + ch_list = [str(c) for c in channels if c] + if not ch_list: + msg = "NotifyDispatcher requires at least one declared channel" + raise ValueError(msg) + self._storage = storage + self._channels: list[str] = list(dict.fromkeys(ch_list)) # de-dupe, preserve order + self._handlers: dict[str, list[Callable[[Notify], None]]] = { + ch: [] for ch in self._channels + } + self._handlers_lock = threading.Lock() + self._lifecycle_lock = threading.Lock() + self._started = False + self._stopping = threading.Event() + self._listener_thread: threading.Thread | None = None + self._dispatch_thread: threading.Thread | None = None + self._dispatch_queue: queue.Queue[Notify | None] = queue.Queue(maxsize=_DISPATCH_QUEUE_MAX) + self._drop_count = 0 + + @property + def channels(self) -> list[str]: + """Snapshot copy of declared channels.""" + return list(self._channels) + + def subscribe(self, channel: str, handler: Callable[[Notify], None]) -> Callable[[], None]: + """Attach ``handler`` to ``channel``; return an unsubscribe callable. + + Safe to call before or after :meth:`start`. Raises if the + channel was not declared at construction time — the channel + list is fixed so the dispatcher knows up-front which LISTENs + to issue (consumers added in follow-up PRs touch the + construction call site). + """ + if channel not in self._handlers: + msg = ( + f"channel {channel!r} not declared at construction; " + f"declared channels: {sorted(self._handlers)}" + ) + raise ValueError(msg) + with self._handlers_lock: + self._handlers[channel].append(handler) + + def _unsubscribe() -> None: + with self._handlers_lock, contextlib.suppress(ValueError): + self._handlers[channel].remove(handler) + + return _unsubscribe + + def start(self) -> None: + """Open the listen stream and start the listener + dispatch threads. + + Idempotent — repeat calls log a debug line and return without + spawning a second listener. + """ + with self._lifecycle_lock: + if self._started: + log.debug("notify_dispatcher.start_noop_already_started") + return + self._started = True + self._stopping.clear() + self._listener_thread = threading.Thread( + target=self._listener_loop, + name="notify-dispatcher-listener", + daemon=True, + ) + self._dispatch_thread = threading.Thread( + target=self._dispatch_loop, + name="notify-dispatcher-dispatch", + daemon=True, + ) + self._listener_thread.start() + self._dispatch_thread.start() + log.info( + "notify_dispatcher.started", + channels=self._channels, + ) + + def stop(self, timeout: float = 5.0) -> None: + """Signal shutdown and join the worker threads. + + Idempotent — safe to call multiple times. Workers exit on the + next iteration of their poll loops; :meth:`stop` blocks up to + ``timeout`` seconds per thread before giving up (the threads are + daemons so the process can exit regardless). + """ + with self._lifecycle_lock: + if not self._started: + return + self._stopping.set() + listener = self._listener_thread + dispatcher = self._dispatch_thread + # Sentinel wakes the dispatch loop out of queue.get(). + with contextlib.suppress(queue.Full): + self._dispatch_queue.put_nowait(None) + if listener is not None: + listener.join(timeout=timeout) + if dispatcher is not None: + dispatcher.join(timeout=timeout) + with self._lifecycle_lock: + self._listener_thread = None + self._dispatch_thread = None + self._started = False + log.info("notify_dispatcher.stopped") + + # ------------------------------------------------------------------ + # Internal threading + # ------------------------------------------------------------------ + + def _listener_loop(self) -> None: + """Drain the storage stream onto the dispatch queue, reconnecting on loss. + + After any disconnect — whether surfaced through the stream's + :class:`NotifyConnectionError` (post-open ``poll`` failure) or + through the generic exception path (``psycopg.connect`` / + initial ``LISTEN`` execute failures during reopen, which are + NOT wrapped by the stream) — the loop sets a ``reconcile_pending`` + flag, waits the backoff, then enqueues one synthetic ``reconcile`` + notify per channel ONLY after the next stream successfully + reopens. Handlers see the synthetic notify and re-read the + relevant rows on the same code path they use for any real event, + closing the missed-notification window regardless of which + exception type caused the disconnect. + """ + backoff = _RECONNECT_BACKOFF_INITIAL + reconcile_pending = False + while not self._stopping.is_set(): + try: + with self._storage.listen(self._channels) as stream: + log.debug( + "notify_dispatcher.stream_open", + channels=self._channels, + ) + # Stream is open — reset backoff for the next outage + # and flush any pending reconcile so consumers see a + # wake-up against a now-live DB. + backoff = _RECONNECT_BACKOFF_INITIAL + if reconcile_pending: + self._synthesize_reconcile() + reconcile_pending = False + while not self._stopping.is_set(): + batch = stream.poll(_LISTENER_POLL_TIMEOUT) + for n in batch: + self._enqueue(n) + except NotifyConnectionError as exc: + if self._stopping.is_set(): + return + log.warning( + "notify_dispatcher.connection_lost", + error=str(exc), + backoff_seconds=backoff, + ) + reconcile_pending = True + if self._stopping.wait(backoff): + return + backoff = min(backoff * 2.0, _RECONNECT_BACKOFF_MAX) + except Exception: + if self._stopping.is_set(): + return + log.exception("notify_dispatcher.listener_unexpected_error") + reconcile_pending = True + if self._stopping.wait(backoff): + return + backoff = min(backoff * 2.0, _RECONNECT_BACKOFF_MAX) + log.debug("notify_dispatcher.listener_exiting") + + def _synthesize_reconcile(self) -> None: + """Push one synthetic ``reconcile`` notify per channel on reconnect. + + Reconcile-on-wake is the same logic handlers run for any real + notification, so a single synthetic event per channel covers + any notifications missed during the connection-loss window. + """ + for ch in self._channels: + self._enqueue(Notify(channel=ch, payload="reconcile", pid=0)) + + def _enqueue(self, notify: Notify) -> None: + """Put a notify on the dispatch queue, dropping oldest on overflow.""" + try: + self._dispatch_queue.put_nowait(notify) + except queue.Full: + # Drop oldest to make room — a slow handler shouldn't be able + # to silently block the listener thread. Log once per power + # of two so a sustained backpressure problem shows up + # in logs without flooding. + self._drop_count += 1 + if self._drop_count & (self._drop_count - 1) == 0: + log.warning( + "notify_dispatcher.dispatch_queue_full_dropping_oldest", + drops_total=self._drop_count, + channel=notify.channel, + ) + with contextlib.suppress(queue.Empty): + self._dispatch_queue.get_nowait() + with contextlib.suppress(queue.Full): + self._dispatch_queue.put_nowait(notify) + + def _dispatch_loop(self) -> None: + """Pull notifies off the queue and invoke handlers per channel. + + Notifies queued on the same channel coalesce per dispatch batch: + after blocking ``get()`` returns one notify, the loop drains + whatever else is already queued and collapses to one + ``per-channel`` notify before invoking handlers. The payload is + signal-only by design (handlers reconcile by re-reading the + underlying rows), so N same-channel notifies have the same + observable effect as one — coalescing turns an N-node deploy + burst into a single ``_discover_nodes`` per channel instead of N. + + Each handler runs under exception suppression so one buggy + consumer can't take down the dispatch thread. + """ + while not self._stopping.is_set(): + try: + first = self._dispatch_queue.get(timeout=_LISTENER_POLL_TIMEOUT) + except queue.Empty: + continue + if first is None: + # Sentinel from :meth:`stop`. + return + # Coalesce by channel: keep the most recent payload per + # channel from this drain batch. Drops a stop sentinel + # silently — the next loop iteration will see _stopping set + # and exit anyway, so we don't need to re-queue the sentinel. + per_channel: dict[str, Notify] = {first.channel: first} + stop_seen = False + while True: + try: + nxt = self._dispatch_queue.get_nowait() + except queue.Empty: + break + if nxt is None: + stop_seen = True + continue + per_channel[nxt.channel] = nxt + for notify in per_channel.values(): + with self._handlers_lock: + handlers = list(self._handlers.get(notify.channel, ())) + for handler in handlers: + t0 = time.monotonic() + try: + handler(notify) + except Exception: + log.exception( + "notify_dispatcher.handler_failed", + channel=notify.channel, + ) + else: + elapsed_ms = (time.monotonic() - t0) * 1000.0 + if elapsed_ms > 100.0: + log.debug( + "notify_dispatcher.handler_slow", + channel=notify.channel, + elapsed_ms=round(elapsed_ms, 1), + ) + if stop_seen: + return + log.debug("notify_dispatcher.dispatch_exiting") diff --git a/turnstone/console/server.py b/turnstone/console/server.py index c5f5f265..2df534b1 100644 --- a/turnstone/console/server.py +++ b/turnstone/console/server.py @@ -4770,6 +4770,12 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]: await close_oidc_state(app.state) app.state.collector.stop() + # Stop the dispatcher after the collector — collector.stop() drops its + # subscription, so the dispatcher's dispatch thread won't fire into a + # half-torn-down collector during shutdown. + notify_dispatcher = getattr(app.state, "notify_dispatcher", None) + if notify_dispatcher is not None: + notify_dispatcher.stop() audit_exec_shutdown = getattr(app.state, "audit_executor", None) if audit_exec_shutdown is not None: _set_audit_executor(None) @@ -11882,6 +11888,7 @@ def create_app( console_url: str = "", router: ConsoleRouter | None = None, console_metrics: ConsoleMetrics | None = None, + notify_dispatcher: Any = None, ) -> Starlette: """Build the Starlette ASGI application for the console dashboard.""" _spec = build_console_spec() @@ -12508,6 +12515,7 @@ def create_app( lifespan=_lifespan, ) app.state.collector = collector + app.state.notify_dispatcher = notify_dispatcher app.state.jwt_secret = jwt_secret app.state.auth_storage = auth_storage app.state.proxy_token_mgr = proxy_token_mgr @@ -12603,7 +12611,7 @@ def main() -> None: from turnstone.core.config import add_config_arg, apply_config add_config_arg(parser) - apply_config(parser, ["console", "auth"]) + apply_config(parser, ["console", "auth", "database"]) args = parser.parse_args() from turnstone.core.log import configure_logging_from_args @@ -12622,6 +12630,13 @@ def main() -> None: db_backend = os.environ.get("TURNSTONE_DB_BACKEND", "sqlite") db_url = os.environ.get("TURNSTONE_DB_URL", "") db_path = os.environ.get("TURNSTONE_DB_PATH", "") + # Optional dedicated LISTEN URL — config.toml ``[database] listen_url`` + # (lifted onto args by ``apply_config``) wins over env, and an empty + # value falls through to the main DB URL inside the storage layer. + # Only used by the ``NotifyDispatcher``; ignored on SQLite. + db_listen_url = getattr(args, "db_listen_url", None) or os.environ.get( + "TURNSTONE_DB_LISTEN_URL", "" + ) auth_storage = init_storage( db_backend, path=db_path, @@ -12630,6 +12645,7 @@ def main() -> None: sslrootcert=os.environ.get("TURNSTONE_DB_SSLROOTCERT", ""), sslcert=os.environ.get("TURNSTONE_DB_SSLCERT", ""), sslkey=os.environ.get("TURNSTONE_DB_SSLKEY", ""), + listen_url=db_listen_url, ) except Exception: log.info("Console storage not available — admin API disabled, JWT-only auth") @@ -12659,11 +12675,21 @@ def main() -> None: router = ConsoleRouter(storage=auth_storage) console_metrics = ConsoleMetrics() + # NotifyDispatcher multiplexes the dedicated LISTEN connection for all + # console-side consumers. Currently one channel: ``services`` for + # reactive node discovery. Followup PRs (ConfigStore live reload, + # scheduler immediate dispatch) add additional channels here. + from turnstone.console.notify_dispatcher import NotifyDispatcher + + notify_dispatcher = NotifyDispatcher(auth_storage, channels=["services"]) + notify_dispatcher.start() + collector = ClusterCollector( storage=auth_storage, token_manager=collector_token_mgr, router=router, console_metrics=console_metrics, + notify_dispatcher=notify_dispatcher, ) collector.start() @@ -12755,6 +12781,7 @@ def main() -> None: console_url=console_url, router=router, console_metrics=console_metrics, + notify_dispatcher=notify_dispatcher, ) log.info("Console starting on %s", console_url) diff --git a/turnstone/core/config.py b/turnstone/core/config.py index 8eb6d4f0..60ca523b 100644 --- a/turnstone/core/config.py +++ b/turnstone/core/config.py @@ -143,6 +143,7 @@ _CONFIG_MAP: dict[str, dict[str, str]] = { "sslrootcert": "db_sslrootcert", "sslcert": "db_sslcert", "sslkey": "db_sslkey", + "listen_url": "db_listen_url", }, "judge": { "enabled": "judge_enabled", diff --git a/turnstone/core/storage/_notify.py b/turnstone/core/storage/_notify.py new file mode 100644 index 00000000..0896f266 --- /dev/null +++ b/turnstone/core/storage/_notify.py @@ -0,0 +1,62 @@ +"""Cross-process notification primitive shared by all storage backends. + +Provides a uniform ``notify`` / ``listen`` shape over PostgreSQL's +``LISTEN`` / ``NOTIFY`` and a SQLite synthetic-sweep fallback. + +Consumers subscribe to one or more channels, drain a +:class:`NotifyStream` via :meth:`NotifyStream.poll`, and reconcile by +re-reading the relevant rows on every wake-up. Payloads are signal-only +(<= 8 KiB on Postgres) — full event content is delivered by SSE or +in-process callbacks elsewhere; this primitive is the "go re-read these +rows" wake-up channel, nothing more. + +The PostgreSQL implementation requires a session-mode connection +(``pgbouncer`` in transaction mode is incompatible with LISTEN). See +the ``listen`` docs on each backend for the deployment-config detail. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + + +@dataclass(frozen=True) +class Notify: + """One notification draining out of a :class:`NotifyStream`.""" + + channel: str + payload: str + pid: int + + +class NotifyConnectionError(Exception): + """Raised when a :class:`NotifyStream`'s underlying connection drops. + + Consumers handle this by closing the stream, reconciling against the + relevant table (re-reading whatever rows the channel describes), and + reopening with a fresh :meth:`StorageBackend.listen` call. + """ + + +class NotifyStream(Protocol): + """Bounded-blocking pull interface for cross-process notifications. + + Returned by :meth:`StorageBackend.listen` as a context manager; the + consumer drains via :meth:`poll` in a loop, typically with a short + timeout so the loop can also observe a shutdown flag. + """ + + def poll(self, timeout: float) -> list[Notify]: + """Wait up to ``timeout`` seconds for notifications. + + Returns the list of notifications received during the wait + (possibly empty on timeout). Raises :class:`NotifyConnectionError` + if the underlying connection was dropped — the caller reconciles + and re-listens. + """ + ... + + def close(self) -> None: + """Stop the stream; subsequent :meth:`poll` calls return ``[]``.""" + ... diff --git a/turnstone/core/storage/_postgresql.py b/turnstone/core/storage/_postgresql.py index 8b2c15d5..23b59101 100644 --- a/turnstone/core/storage/_postgresql.py +++ b/turnstone/core/storage/_postgresql.py @@ -3,13 +3,16 @@ from __future__ import annotations import contextlib +import os import threading import time from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Iterable, Iterator + + from turnstone.core.storage._notify import Notify, NotifyStream import sqlalchemy as sa @@ -120,11 +123,105 @@ def _escape_ilike(s: str) -> str: return s.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") +def _resolve_pg_listen_url(override: str, sqlalchemy_url: str) -> str: + """Resolve the URL used by the dedicated LISTEN connection. + + Precedence: + + 1. ``override`` — explicitly passed via :class:`PostgreSQLBackend` + constructor, typically wired from ``[database] listen_url`` in + ``config.toml`` or ``--db-listen-url`` on the CLI. + 2. ``TURNSTONE_DB_LISTEN_URL`` environment variable. + 3. The engine's main DB URL. + + The override is for deployments where the regular ``TURNSTONE_DB_URL`` + points at a ``pgbouncer`` running in transaction pooling mode (the + project default per ``docs/pgbouncer.md``). LISTEN holds session + state and is incompatible with transaction pooling; the dispatcher + needs to bypass pgbouncer for that single connection. When neither + override is set and ``TURNSTONE_DB_URL`` already points at Postgres + directly (no pooler in between), the fallback uses the engine URL + as-is. + + The SQLAlchemy ``+psycopg`` driver suffix is stripped so the URL is + consumable by ``psycopg.connect`` directly. + """ + raw = override.strip() or os.environ.get("TURNSTONE_DB_LISTEN_URL", "").strip() + raw = raw or sqlalchemy_url + return raw.replace("postgresql+psycopg://", "postgresql://", 1) + + +class _PostgreSQLNotifyStream: + """PostgreSQL ``listen`` stream — drains ``conn.notifies`` per poll. + + Owns a dedicated psycopg autocommit connection. Each :meth:`poll` + waits up to ``timeout`` seconds for notifications and returns them + as a list — empty on timeout, raises :class:`NotifyConnectionError` + on connection loss (caller reconciles + re-listens). + + Closing the stream from another thread is the supported abort path: + ``close`` calls ``conn.close()``, which causes the in-flight + :meth:`poll` to wake (the next call returns ``[]`` because + ``_closed`` is set). + """ + + def __init__(self, conn: Any, channels: list[str]) -> None: + self._conn = conn + self._channels = list(channels) + self._closed = False + self._close_lock = threading.Lock() + + def poll(self, timeout: float) -> list[Notify]: + from turnstone.core.storage._notify import Notify, NotifyConnectionError + + if self._closed: + return [] + out: list[Notify] = [] + try: + # psycopg3 generator yields whatever's available within the + # window, then stops — bounded blocking semantics. Per-call + # generator (not a long-lived one) so close() can abort by + # closing the connection without leaving a half-consumed + # generator behind. + for n in self._conn.notifies(timeout=max(0.0, timeout)): + out.append(Notify(channel=n.channel, payload=n.payload, pid=n.pid)) + except Exception as exc: + if self._closed: + # Graceful close-from-another-thread surfaced as an + # operational error inside notifies() — swallow it, + # let the caller observe close via the next poll + # returning ``[]``. + return out + raise NotifyConnectionError(str(exc)) from exc + return out + + def close(self) -> None: + with self._close_lock: + if self._closed: + return + self._closed = True + conn = self._conn + # Best-effort UNLISTEN + close. An already-broken connection + # raises here; the consumer's reconciliation logic will catch + # the underlying ``NotifyConnectionError`` on the next poll if + # any waiter is still blocked. + with contextlib.suppress(Exception): + conn.execute("UNLISTEN *") + with contextlib.suppress(Exception): + conn.close() + + class PostgreSQLBackend: """PostgreSQL implementation of the StorageBackend protocol.""" def __init__( - self, url: str, pool_size: int = 2, max_overflow: int = 3, *, create_tables: bool = True + self, + url: str, + pool_size: int = 2, + max_overflow: int = 3, + *, + create_tables: bool = True, + listen_url: str = "", ) -> None: self._engine = sa.create_engine( url, @@ -134,6 +231,12 @@ class PostgreSQLBackend: ) self._db_unavailable = False self._db_unavailable_lock = threading.Lock() + # Operator override for the dedicated LISTEN connection's URL. + # Empty string means "fall back through env var, then the main + # engine URL" — see :func:`_resolve_pg_listen_url` for the full + # precedence rules. Threaded through ``init_storage`` from + # ``config.toml [database] listen_url`` / ``--db-listen-url``. + self._listen_url_override = listen_url if create_tables: metadata.create_all(self._engine) @@ -1863,6 +1966,61 @@ class PostgreSQLBackend: conn.commit() return result.rowcount > 0 + # -- Cross-process notifications ------------------------------------------- + + def notify(self, channel: str, payload: str = "") -> None: + """Broadcast a wake-up via ``pg_notify`` on a pooled connection. + + ``channel`` and ``payload`` are bound as parameters so this is + safe to call with operator-supplied strings without quoting + gymnastics. Postgres caps the payload at 8 KiB — keep payloads + signal-only (a JSON id list, an op name) and let consumers + re-read the underlying rows on wake-up. + """ + with self._conn() as conn: + conn.execute( + sa.text("SELECT pg_notify(:channel, :payload)"), + {"channel": channel, "payload": payload}, + ) + conn.commit() + + @contextlib.contextmanager + def listen(self, channels: Iterable[str]) -> Iterator[NotifyStream]: + """Subscribe to channels on a dedicated session-mode connection. + + Opens a fresh ``psycopg`` connection in autocommit mode (the + SQLAlchemy pool is incompatible with LISTEN — it recycles + connections back into a pool that may be transaction-pooled by + pgbouncer). Channel names are interpolated via + ``psycopg.sql.Identifier`` so caller-supplied channel strings + can't inject SQL. + + ``TURNSTONE_DB_LISTEN_URL`` overrides the engine URL — see + :func:`_resolve_pg_listen_url` for the bypass-URL rationale. + + Yields a :class:`_PostgreSQLNotifyStream`; the connection is + closed on context exit. + """ + import psycopg + from psycopg import sql + + ch_list = [str(c) for c in channels if c] + sqlalchemy_url = self._engine.url.render_as_string(hide_password=False) + listen_url = _resolve_pg_listen_url(self._listen_url_override, sqlalchemy_url) + conn = psycopg.connect(listen_url, autocommit=True) + stream: _PostgreSQLNotifyStream | None = None + try: + for ch in ch_list: + conn.execute(sql.SQL("LISTEN {}").format(sql.Identifier(ch))) + stream = _PostgreSQLNotifyStream(conn, ch_list) + yield stream + finally: + if stream is not None: + stream.close() + else: + with contextlib.suppress(Exception): + conn.close() + # -- Node metadata --------------------------------------------------------- def get_node_metadata(self, node_id: str) -> list[dict[str, Any]]: diff --git a/turnstone/core/storage/_protocol.py b/turnstone/core/storage/_protocol.py index f330b88c..dfef0147 100644 --- a/turnstone/core/storage/_protocol.py +++ b/turnstone/core/storage/_protocol.py @@ -5,8 +5,10 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any, Protocol, TypedDict, runtime_checkable if TYPE_CHECKING: + from collections.abc import Iterable from contextlib import AbstractContextManager + from turnstone.core.storage._notify import NotifyStream from turnstone.core.workstream import WorkstreamKind @@ -1020,6 +1022,33 @@ class StorageBackend(Protocol): """Remove a service registration. Returns True if existed.""" ... + # -- Cross-process notifications ------------------------------------------- + + def notify(self, channel: str, payload: str = "") -> None: + """Broadcast a wake-up on ``channel`` to any listening process. + + Payloads are signal-only — a JSON-encoded string identifying + which rows to re-read, capped well below Postgres's 8 KiB + ``NOTIFY`` payload limit. Full event content is NOT delivered + this way; consumers reconcile by reading the relevant table on + wake-up. Safe to call from any thread. + """ + ... + + def listen(self, channels: Iterable[str]) -> AbstractContextManager[NotifyStream]: + """Subscribe to one or more channels for cross-process wake-ups. + + Returns a context manager wrapping a :class:`NotifyStream` the + caller drains via :meth:`NotifyStream.poll`. PostgreSQL holds a + dedicated session-mode connection for the lifetime of the + context (incompatible with ``pgbouncer`` transaction pooling — + see :class:`PostgreSQLBackend.listen` for the bypass-URL config). + SQLite emits a synthetic-sweep wake on its own cadence (see + ``_SQLITE_NOTIFY_SWEEP_INTERVAL``) per subscribed channel so + consumer code is identical across backends. + """ + ... + # -- Node metadata --------------------------------------------------------- def get_node_metadata(self, node_id: str) -> list[dict[str, Any]]: diff --git a/turnstone/core/storage/_registry.py b/turnstone/core/storage/_registry.py index d1ad9ce8..56066464 100644 --- a/turnstone/core/storage/_registry.py +++ b/turnstone/core/storage/_registry.py @@ -34,6 +34,7 @@ def init_storage( sslrootcert: str = "", sslcert: str = "", sslkey: str = "", + listen_url: str = "", ) -> StorageBackend: """Initialize the storage backend singleton. @@ -43,6 +44,13 @@ def init_storage( url: PostgreSQL connection URL (e.g. postgresql+psycopg://user:pass@host/db) pool_size: Connection pool size (PostgreSQL only) run_migrations: Whether to run Alembic migrations on init + listen_url: Optional dedicated PostgreSQL URL for the dispatcher's + ``LISTEN`` connection. Required only when ``url`` points at a + ``pgbouncer`` running in transaction pooling mode (LISTEN + holds session state and is incompatible with transaction + pooling — see ``docs/pgbouncer.md``). Empty string means + "fall back through ``TURNSTONE_DB_LISTEN_URL`` env var, then + the main ``url``." Ignored on SQLite. """ global _storage @@ -84,7 +92,12 @@ def init_storage( sep = "&" if "?" in url else "?" url += sep + urlencode(ssl_params) - _storage = PostgreSQLBackend(url, pool_size=pool_size, create_tables=create_tables) + _storage = PostgreSQLBackend( + url, + pool_size=pool_size, + create_tables=create_tables, + listen_url=listen_url, + ) log.info("Storage initialized: PostgreSQL") else: diff --git a/turnstone/core/storage/_sqlite.py b/turnstone/core/storage/_sqlite.py index 3396fa16..d81f02c0 100644 --- a/turnstone/core/storage/_sqlite.py +++ b/turnstone/core/storage/_sqlite.py @@ -3,14 +3,18 @@ from __future__ import annotations import contextlib +import queue import threading +import time from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any import sqlalchemy as sa if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Iterable, Iterator + + from turnstone.core.storage._notify import Notify, NotifyStream from turnstone.core.log import get_logger from turnstone.core.storage._protocol import ( @@ -129,6 +133,90 @@ def _fts5_query(query: str) -> str: return " ".join(safe) +# Synthetic-sweep cadence for the SQLite ``listen`` fallback. SQLite is +# the dev-only path where reactive latency isn't load-bearing — a single +# console process, no cross-process notify semantics to recover from. +# 300 s sits comfortably above the existing per-consumer timers +# (cluster collector's 60 s ``discovery_interval``, any future +# ConfigStore/scheduler reload cadences) so the sweep is a true backstop +# rather than a duplicate tick. Future consumers that need tighter +# SQLite-mode reactive latency should pass a custom interval through +# :meth:`SQLiteBackend.listen` rather than lowering this default. +_SQLITE_NOTIFY_SWEEP_INTERVAL: float = 300.0 + + +class _SQLiteNotifyStream: + """SQLite ``listen`` stream — synthetic sweep + in-process fan-out. + + Each poll either drains queued in-process notifies (delivered by a + same-process :meth:`SQLiteBackend.notify` call) or emits one + synthetic ``Notify(channel, payload="sweep", pid=0)`` per subscribed + channel once :attr:`_sweep_interval` has elapsed since the previous + sweep, whichever happens first. Consumers handle both shapes the + same way: re-read the relevant rows on every wake-up. + """ + + def __init__( + self, + backend: SQLiteBackend, + channels: list[str], + sweep_interval: float, + ) -> None: + self._backend = backend + self._channels = list(channels) + self._sweep_interval = sweep_interval + self._queue: queue.Queue[Any] = queue.Queue() + self._closed = False + self._last_sweep = time.monotonic() + if self._channels: + backend._notify_register(self._channels, self._queue) + + def poll(self, timeout: float) -> list[Notify]: + from turnstone.core.storage._notify import Notify + + if self._closed: + return [] + deadline = time.monotonic() + max(0.0, timeout) + # Emit a synthetic-sweep tick on the first poll where the sweep + # interval has elapsed. Single tick per channel per interval — + # PG-equivalent "one wake-up per change" semantics, not a burst. + now = time.monotonic() + if self._channels and now - self._last_sweep >= self._sweep_interval: + self._last_sweep = now + for ch in self._channels: + with contextlib.suppress(Exception): + self._queue.put_nowait(Notify(channel=ch, payload="sweep", pid=0)) + out: list[Notify] = [] + try: + while True: + if self._closed: + break + if out: + # Drain everything already queued without further + # blocking — produces "one poll returns the burst" + # semantics so the consumer reconciles once per wake. + item = self._queue.get_nowait() + else: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + item = self._queue.get(timeout=remaining) + out.append(item) + except queue.Empty: + # End-of-drain: the blocking get hit its deadline OR a + # get_nowait found the queue empty. Either way we return + # whatever was already collected. + pass + return out + + def close(self) -> None: + if self._closed: + return + self._closed = True + if self._channels: + self._backend._notify_unregister(self._channels, self._queue) + + class SQLiteBackend: """SQLite implementation of the StorageBackend protocol.""" @@ -155,6 +243,15 @@ class SQLiteBackend: self._fts5_available = False self._db_unavailable = False self._db_unavailable_lock = threading.Lock() + # In-process notify fan-out: channel name -> list of stream queues. + # SQLite has no cross-process LISTEN/NOTIFY, so notifications are + # delivered synchronously to any open ``listen`` stream in the same + # process. Streams register on open + unregister on close; the + # synthetic-sweep timer below covers consumers that need a periodic + # wake regardless of producer activity (matching the PG-side + # discovery-loop cadence). + self._notify_lock = threading.Lock() + self._notify_subs: dict[str, list[queue.Queue[Any]]] = {} if create_tables: self._init_schema() @@ -2009,6 +2106,75 @@ class SQLiteBackend: conn.commit() return result.rowcount > 0 + # -- Cross-process notifications ------------------------------------------- + + def notify(self, channel: str, payload: str = "") -> None: + """In-process broadcast — SQLite has no cross-process channel. + + SQLite deployments are single-process by design (no shared backend + across nodes); the storage layer delivers to any ``listen`` stream + open in the same process. Cross-process consumers wouldn't be + served regardless — the synthetic-sweep wake-up in :meth:`listen` + is the parity fallback so consumer code stays backend-agnostic. + """ + from turnstone.core.storage._notify import Notify + + with self._notify_lock: + subs = list(self._notify_subs.get(channel, ())) + for q in subs: + with contextlib.suppress(Exception): + q.put(Notify(channel=channel, payload=payload, pid=0)) + + @contextlib.contextmanager + def listen( + self, + channels: Iterable[str], + *, + sweep_interval: float = _SQLITE_NOTIFY_SWEEP_INTERVAL, + ) -> Iterator[NotifyStream]: + """Subscribe to channels — synthetic-sweep + in-process fan-out. + + The returned stream wakes every ``sweep_interval`` seconds with + one ``Notify(channel, payload="sweep", pid=0)`` per subscribed + channel; the default (:data:`_SQLITE_NOTIFY_SWEEP_INTERVAL`) + suits a dev backstop with a 60 s consumer-side timer. Callers + that need a tighter cadence (e.g. a future consumer without its + own polling timer) pass a smaller value here. In-process + :meth:`notify` calls deliver immediately on top of the sweep. + Either path produces a wake-up; consumers reconcile by re-reading + the relevant rows. + + Channel names are de-duplicated so callers passing the same name + twice don't double-deliver each notify to a single stream. + """ + # de-dupe + preserve insertion order — passing the same channel + # twice would otherwise register the stream's queue against that + # channel twice and deliver each notify multiple times. + ch_list = list(dict.fromkeys(str(c) for c in channels if c)) + stream = _SQLiteNotifyStream(self, ch_list, sweep_interval=sweep_interval) + try: + yield stream + finally: + stream.close() + + def _notify_register(self, channels: list[str], q: queue.Queue[Any]) -> None: + """Subscribe a stream's queue to in-process notifies on ``channels``.""" + with self._notify_lock: + for ch in channels: + self._notify_subs.setdefault(ch, []).append(q) + + def _notify_unregister(self, channels: list[str], q: queue.Queue[Any]) -> None: + """Detach a stream's queue from in-process notifies on ``channels``.""" + with self._notify_lock: + for ch in channels: + subs = self._notify_subs.get(ch) + if subs is None: + continue + with contextlib.suppress(ValueError): + subs.remove(q) + if not subs: + self._notify_subs.pop(ch, None) + # -- Node metadata --------------------------------------------------------- def get_node_metadata(self, node_id: str) -> list[dict[str, Any]]: diff --git a/turnstone/core/storage/migrations/versions/053_services_notify_trigger.py b/turnstone/core/storage/migrations/versions/053_services_notify_trigger.py new file mode 100644 index 00000000..c3644d3f --- /dev/null +++ b/turnstone/core/storage/migrations/versions/053_services_notify_trigger.py @@ -0,0 +1,95 @@ +"""Trigger ``pg_notify('services', ...)`` on service registry changes. + +The console-side :class:`NotifyDispatcher` (`turnstone/console/notify_dispatcher.py`) +holds a dedicated LISTEN connection and fans channel events out to handlers. +This migration installs the producer side for the ``services`` channel — +the cluster collector subscribes so new-node discovery is reactive instead +of polling every 60 s. + +The trigger filters heartbeat-only UPDATEs in-trigger (same url + same +metadata, only ``last_heartbeat`` changed): ``register_service`` is an +UPSERT, so a node restart that changes url/metadata still fires; a plain +heartbeat tick stays quiet to avoid flooding the channel on every +30 s × N-nodes cluster tick. Channel payload is a small JSON object — +service_type, service_id, op — well below PG's 8 KiB NOTIFY limit; the +handler reconciles by re-reading ``services`` rather than relying on +the payload content. + +SQLite is a no-op for this migration — the SQLite backend's in-process +:meth:`notify` doesn't go through a trigger, and the synthetic-sweep +fallback in :meth:`listen` covers consumer parity. + +What this trigger does NOT cover: crashed-node detection. A node that +dies without running its deregister handshake leaves a stale row that +ages out via the existing 120 s heartbeat-expiry filter. The 60 s +discovery loop in the collector keeps running as the backstop for +crash-shaped node loss. + +Revision ID: 053 +Revises: 052 +Create Date: 2026-05-10 +""" + +import sqlalchemy as sa +from alembic import op + +revision = "053" +down_revision = "052" +branch_labels = None +depends_on = None + + +_TRIGGER_FN_NAME = "turnstone_notify_services" +_TRIGGER_NAME = "services_notify" + + +_TRIGGER_FN_SQL = f""" +CREATE OR REPLACE FUNCTION {_TRIGGER_FN_NAME}() RETURNS trigger AS $$ +BEGIN + -- Skip heartbeat-only UPDATEs: same url and metadata, only + -- ``last_heartbeat`` changed. ``register_service`` is an UPSERT + -- (on_conflict_do_update), so node restarts that change url or + -- metadata MUST still fire — only no-op heartbeat ticks stay + -- quiet. IS NOT DISTINCT FROM treats NULLs as equal so a row + -- with NULL metadata before/after doesn't trip the diff. + IF TG_OP = 'UPDATE' + AND OLD.url IS NOT DISTINCT FROM NEW.url + AND OLD.metadata IS NOT DISTINCT FROM NEW.metadata THEN + RETURN NULL; + END IF; + + PERFORM pg_notify( + 'services', + json_build_object( + 'service_type', COALESCE(NEW.service_type, OLD.service_type), + 'service_id', COALESCE(NEW.service_id, OLD.service_id), + 'op', TG_OP + )::text + ); + RETURN NULL; +END; +$$ LANGUAGE plpgsql; +""" + + +_TRIGGER_SQL = f""" +CREATE TRIGGER {_TRIGGER_NAME} +AFTER INSERT OR UPDATE OR DELETE ON services +FOR EACH ROW EXECUTE FUNCTION {_TRIGGER_FN_NAME}(); +""" + + +def upgrade() -> None: + bind = op.get_bind() + if bind.dialect.name != "postgresql": + return + op.execute(sa.text(_TRIGGER_FN_SQL)) + op.execute(sa.text(_TRIGGER_SQL)) + + +def downgrade() -> None: + bind = op.get_bind() + if bind.dialect.name != "postgresql": + return + op.execute(sa.text(f"DROP TRIGGER IF EXISTS {_TRIGGER_NAME} ON services")) + op.execute(sa.text(f"DROP FUNCTION IF EXISTS {_TRIGGER_FN_NAME}()"))