mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
feat(console): reactive node discovery via PG LISTEN/NOTIFY dispatcher (#505)
* feat(console): reactive node discovery via PG LISTEN/NOTIFY dispatcher Add a console-side `NotifyDispatcher` that holds a dedicated PostgreSQL `LISTEN` connection and fans wake-ups out to per-channel handlers on a separate dispatch thread. Cluster collector subscribes to a new `services` channel and runs node discovery reactively — new-node / graceful-deregister visibility drops from up-to-60 s to ~500 ms on Postgres, with the 60 s discovery loop retained as the backstop for crash-shaped node loss (NOTIFY only fires on real writes). Storage layer gains a uniform `notify` / `listen` API: - PostgreSQL: real `pg_notify` / `LISTEN` on a dedicated session-mode connection that bypasses pgbouncer (mandatory: pgbouncer is required in transaction-pool mode per docs, which is incompatible with LISTEN). - SQLite: in-process fan-out + synthetic-sweep fallback so consumer code is identical across backends. `TURNSTONE_DB_LISTEN_URL` (or `[database] listen_url` in config.toml) points the dispatcher's connection direct-to-Postgres. Defaults to the main DB URL when unset. Migration 053 installs the `services_notify` trigger; it filters heartbeat-only UPDATEs in-trigger so the 30 s × N-nodes heartbeat tick stays quiet, while INSERT, DELETE, and url/metadata-changing UPDATE still fire. Dispatcher detail: - Two threads: listener (drains stream → bounded queue) and dispatch (invokes handlers under exception suppression). Same-channel notifies coalesce per dispatch batch so an N-node deploy burst is one `_discover_nodes` per channel. - Reconnect uses exponential backoff (1 s → 30 s cap). After any successful reopen — whether the prior failure was a stream-poll error or a connect / initial-LISTEN error — one synthetic Notify with payload="reconcile" is enqueued per channel so handlers re-read on the same code path they use for real events. Future consumers (ConfigStore live reload, scheduler immediate dispatch, audit live-tail) plug in by adding their channel to the dispatcher's construction list. Tests: 22 dispatcher tests (incl. reconnect + coalescing under stub storage), 7 SQLite notify-stream tests, 4 PG-gated trigger-filter tests, 4 collector wire-in tests. All pass; ruff + mypy clean. * fix(notify): address Copilot review on #505 - _sqlite.py: SQLiteBackend.listen() now de-dupes channel names via dict.fromkeys before constructing the stream — duplicates would otherwise register the queue twice and double-deliver each notify. - _sqlite.py: SQLiteBackend.listen() gains a keyword-only sweep_interval parameter (defaults to _SQLITE_NOTIFY_SWEEP_INTERVAL) — matches what the comment at the constant already promised, and lets future consumers without their own polling timer pick a tighter cadence without reaching into private stream attributes. - _sqlite.py: documented the `except queue.Empty: pass` end-of-drain termination so it's not mistaken for swallowing an unexpected error. - _postgresql.py: docstring referenced :func:`_pg_listen_url` which was renamed to _resolve_pg_listen_url during PR development. - notify_dispatcher.py: module docstring referenced a non-existent _bootstrap_console_subsystem; wire-in is at console/server.py::main. Refuted (no change, false positives from github-code-quality bot): - 4× "Statement has no effect" on Protocol-method `...` ellipsis bodies (idiomatic Python Protocol declaration, not dead code). - 2× "Mixed import style" in tests — `import ... as nd_mod` is intentional to allow attribute assignment for monkey-patching the module's `_RECONNECT_BACKOFF_INITIAL` constant inside try/finally.
This commit is contained in:
@@ -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) |
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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"
|
||||
@@ -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")
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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")
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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 ``[]``."""
|
||||
...
|
||||
@@ -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]]:
|
||||
|
||||
@@ -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]]:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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]]:
|
||||
|
||||
@@ -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}()"))
|
||||
Reference in New Issue
Block a user