From 6ae9a26b39db86dab3c7a069c02138ff75689c37 Mon Sep 17 00:00:00 2001 From: Simon Fishel Date: Fri, 2 Oct 2026 13:35:28 -0700 Subject: [PATCH 1/5] feat: probe queries that go quiet for a lost terminal event (WBC-1179) The driver trusts a single pushed state_updated per query. When that frame is lost on a healthy socket, the query waits forever (2026-09-17 staging MCP canary: the session logged success, the driver never saw it). Add a bounded watchdog on the reader thread. Any inbound event for an execution, progress included, resets that query's clock. After N seconds of silence it re-sends retrieve_results as a probe, backing off N, 2N, 4N, then 8N. Replies to a retrieve reset the clock but not the backoff, so a long, silent query isn't re-asked every N seconds. The check runs on every reader-loop iteration (gated to at most once per min(1s, N)), so steady traffic for one query can't starve another's watchdog. When the watchdog is on, recv() is capped at that check interval, so a read_timeout of None or a long one can't park the reader past a due probe. Probes never wait for the send lock: if a user thread is mid-send, maybe stalled, the probe is skipped until the next check instead of stopping recv() for every query on the connection. Probed states: EXECUTION_REQUESTED, PENDING, RUNNING. A query whose results were already requested gets one probe, after 4N: the session re-serializes and re-sends the full result for every retrieve, so duplicates are costly. Duplicate replies are dropped by the existing claim in complete_query. Add the session's PENDING state, so a probe reply can never fail a healthy query as an unknown state. sql-session's execution cache is a small LRU that can evict a query that is still running; its terminal state_updated still arrives. So a probe's "Execution not found" doesn't fail the query: it stops probing it and keeps waiting. After the normal retrieve, not-found still fails it. New additive kwarg stale_query_probe_seconds (default 30.0) on connect(), connect_direct() and Connection. None disables the watchdog; an invalid value disables it with a warning and never raises. --- README.md | 5 + tests/test_disconnect.py | 5 +- tests/test_driver.py | 17 ++ tests/test_watchdog.py | 503 +++++++++++++++++++++++++++++++++++++ wherobots/db/connection.py | 259 +++++++++++++++++-- wherobots/db/constants.py | 2 + wherobots/db/driver.py | 15 ++ wherobots/db/types.py | 3 + 8 files changed, 791 insertions(+), 18 deletions(-) create mode 100644 tests/test_watchdog.py diff --git a/README.md b/README.md index 07edb56..d99cfb0 100644 --- a/README.md +++ b/README.md @@ -277,3 +277,8 @@ users may find useful: caller needs to interrupt it (e.g. on client disconnect or timeout). The event is checked before each HTTP request, between retry attempts, and before the WebSocket handshake. +* `stale_query_probe_seconds`: how long a running query may go without + any message from the SQL session before the driver asks the session + for its state again, in case the query's completion event was lost. + Repeat probes back off to at most 8 times this interval. Defaults to + 30 seconds; pass `None` to disable probing. diff --git a/tests/test_disconnect.py b/tests/test_disconnect.py index 5f420a0..da70d8c 100644 --- a/tests/test_disconnect.py +++ b/tests/test_disconnect.py @@ -582,9 +582,10 @@ def slow_close(): def test_slow_handshake_does_not_consume_reader_join(): # The handshake uses most of close_timeout. A long read_timeout keeps the # reader parked in recv(), so only the end-of-stream delivered 0.3s after - # the handshake wakes it; close() must still be waiting then. + # the handshake wakes it; close() must still be waiting then. The watchdog + # would cap recv() at its 1s check interval, so it's off here. ws = Transport() - conn = Connection(ws, read_timeout=10) + conn = Connection(ws, read_timeout=10, stale_query_probe_seconds=None) def slow_close(): time.sleep(0.9) diff --git a/tests/test_driver.py b/tests/test_driver.py index 6c9cb7d..3295cd0 100644 --- a/tests/test_driver.py +++ b/tests/test_driver.py @@ -75,6 +75,23 @@ def test_session_id_is_derived_from_status_url(self, mock_post, mock_get, suffix ) assert kwargs["session_id"] == "test-id" + @pytest.mark.parametrize("probe_seconds", [None, 5.0]) + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_stale_query_probe_seconds_is_forwarded( + self, mock_post, mock_get, probe_seconds + ): + _, kwargs = _run_connect_full( + mock_post, mock_get, stale_query_probe_seconds=probe_seconds + ) + assert kwargs["stale_query_probe_seconds"] == probe_seconds + + @patch("wherobots.db.driver.requests.get") + @patch("wherobots.db.driver.requests.post") + def test_stale_query_probes_default_to_30_seconds(self, mock_post, mock_get): + _, kwargs = _run_connect_full(mock_post, mock_get) + assert kwargs["stale_query_probe_seconds"] == 30.0 + @patch("wherobots.db.driver.requests.get") @patch("wherobots.db.driver.requests.post") def test_omitted_region_runtime_not_sent(self, mock_post, mock_get): diff --git a/tests/test_watchdog.py b/tests/test_watchdog.py new file mode 100644 index 0000000..b783366 --- /dev/null +++ b/tests/test_watchdog.py @@ -0,0 +1,503 @@ +"""The stale-query watchdog re-asks the session when a terminal event is lost.""" + +import json +import logging +import queue +import threading +import time +from unittest.mock import MagicMock, patch + +import cbor2 +import pytest +from websockets.exceptions import ConnectionClosedOK + +from wherobots.db.connection import Connection +from wherobots.db.driver import connect_direct +from wherobots.db.errors import OperationalError +from wherobots.db.models import Store +from wherobots.db.types import ExecutionState, ResultsFormat, StorageFormat + +PROBE = 0.05 +"""Probe interval used throughout: small, so the tests stay fast.""" + +SLOW_PROBE = 0.2 +"""For assertions that something does *not* happen within N: a larger N, so a +briefly stalled test thread on a loaded runner can't produce a false probe.""" + +NOT_FOUND = "Execution not found" + +READ_TIMEOUT = 0.01 + +STORE = Store.for_download(format=StorageFormat.PARQUET) + +RESULT_URI = "https://presigned.example.com/results.parquet" + + +class Session: + """A fake SQL session transport. + + Records every request with its send time, and lets a test answer probes + (``retrieve_results`` requests) through ``on_probe``. ``before_send`` runs + before a request is recorded, so a test can stall a send. + """ + + def __init__(self, on_probe=None): + self.incoming = queue.Queue() + self.sent = [] + self.on_probe = on_probe + self.before_send = None + self.close_timeout = 1.0 + self.socket = MagicMock() + self.socket.shutdown.side_effect = lambda how: self.close() + + def recv(self, timeout): + try: + value = self.incoming.get(timeout=timeout) + except queue.Empty: + raise TimeoutError from None + if isinstance(value, Exception): + raise value + if isinstance(value, bytes): + return value + return json.dumps(value) + + def send(self, value): + message = json.loads(value) + if self.before_send is not None: + self.before_send(message) + self.sent.append((time.monotonic(), message)) + if message["kind"] == "retrieve_results" and self.on_probe is not None: + self.on_probe(self, message) + + def close(self): + self.incoming.put(ConnectionClosedOK(None, None)) + + def put(self, **message): + self.incoming.put(message) + + def requests(self, kind, execution_id=None): + return [ + (sent_at, message) + for sent_at, message in list(self.sent) + if message["kind"] == kind + and (execution_id is None or message["execution_id"] == execution_id) + ] + + def probe_times(self, execution_id=None): + return [t for t, _ in self.requests("retrieve_results", execution_id)] + + +def wait_until(predicate, timeout=3.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.005) + return predicate() + + +def connect(session, **kwargs): + kwargs.setdefault("stale_query_probe_seconds", PROBE) + return Connection(session, read_timeout=READ_TIMEOUT, **kwargs) + + +def execute(conn, session, sql="SELECT 1", store=None): + """Start a query that the session reports as running, then goes quiet.""" + cursor = conn.cursor() + cursor.execute(sql, store=store) + execution_id = session.requests("execute_sql")[-1][1]["execution_id"] + session.put(kind="state_updated", execution_id=execution_id, state="running") + return cursor, execution_id + + +def results_of(cursor): + return cursor._Cursor__queue + + +def tracked(conn, execution_id): + return conn._Connection__queries.get(execution_id) + + +def reply(session, execution_id, state="succeeded", **fields): + session.put( + kind="execution_result", + execution_id=execution_id, + state=state, + results=None, + **fields, + ) + + +def test_disabled_watchdog_never_probes(): + session = Session() + conn = connect(session, stale_query_probe_seconds=None) + try: + cursor, execution_id = execute(conn, session, store=STORE) + time.sleep(PROBE * 8) + assert session.requests("retrieve_results") == [] + assert tracked(conn, execution_id) is not None + finally: + conn.close() + + +def test_steady_traffic_for_one_query_does_not_starve_anothers_watchdog(): + # read_timeout is well above the progress cadence, so recv() never times + # out while B is streaming. + session = Session() + conn = Connection(session, read_timeout=0.5, stale_query_probe_seconds=SLOW_PROBE) + stop = threading.Event() + try: + _, quiet = execute(conn, session) + _, busy = execute(conn, session) + + def stream_progress(): + while not stop.is_set(): + session.put(kind="execution_progress", execution_id=busy) + time.sleep(0.005) + + streamer = threading.Thread(target=stream_progress, daemon=True) + streamer.start() + assert wait_until(lambda: session.probe_times(quiet), timeout=2) + # Progress events are activity: the busy query is never probed. + assert session.probe_times(busy) == [] + finally: + stop.set() + conn.close() + + +def test_successive_probes_back_off_to_a_cap(): + session = Session( + on_probe=lambda s, m: reply(s, m["execution_id"], state="running") + ) + conn = connect(session) + try: + _, execution_id = execute(conn, session) + registered = session.requests("execute_sql")[0][0] + assert wait_until(lambda: len(session.probe_times()) >= 5, timeout=10) + times = session.probe_times()[:5] + assert times[0] - registered >= PROBE + gaps = [later - earlier for earlier, later in zip(times, times[1:])] + # N, then 2N, 4N, 8N, capped at 8N; never a tight loop. The upper + # bound (next doubling) still tells each step, and the cap, apart. + for gap, expected in zip(gaps, [2, 4, 8, 8]): + assert PROBE * expected <= gap < PROBE * expected * 2, gaps + finally: + conn.close() + + +def test_probe_recovers_non_store_query_and_requests_configured_format(): + rows = b'[{"x": 1}, {"x": 2}]' + + def on_probe(s, m): + s.incoming.put( + cbor2.dumps( + { + "kind": "execution_result", + "execution_id": m["execution_id"], + "state": "succeeded", + "results": {"result_bytes": rows, "format": "json"}, + } + ) + ) + + session = Session(on_probe=on_probe) + conn = connect(session, results_format=ResultsFormat.JSON) + try: + cursor, execution_id = execute(conn, session) + result = results_of(cursor).get(timeout=3) + assert result.error is None + assert result.results == [{"x": 1}, {"x": 2}] + (_, probe), *_ = session.requests("retrieve_results") + assert probe == { + "kind": "retrieve_results", + "execution_id": execution_id, + "format": "json", + } + finally: + conn.close() + + +def test_lost_retrieve_results_reply_is_re_requested(): + asked = [] + + def on_probe(s, m): + # Lose the reply to the driver's own request; answer the probe. + asked.append(m) + if len(asked) > 1: + s.incoming.put( + cbor2.dumps( + { + "kind": "execution_result", + "execution_id": m["execution_id"], + "state": "succeeded", + "results": {"result_bytes": b"[]", "format": "json"}, + } + ) + ) + + session = Session(on_probe=on_probe) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session) + session.put(kind="state_updated", execution_id=execution_id, state="succeeded") + result = results_of(cursor).get(timeout=3) + assert result.error is None + assert result.results == [] + assert len(asked) == 2 + finally: + conn.close() + + +@pytest.mark.parametrize("state", ["running", "pending"]) +def test_non_terminal_probe_reply_keeps_waiting_and_probing(state): + session = Session(on_probe=lambda s, m: reply(s, m["execution_id"], state=state)) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session) + assert wait_until(lambda: len(session.probe_times()) >= 2) + query = tracked(conn, execution_id) + assert query is not None + assert query.state == ExecutionState[state.upper()] + assert results_of(cursor).empty() + # The probe replies aren't events of their own: backoff keeps growing. + first, second = session.probe_times()[:2] + assert second - first >= PROBE * 2 + finally: + conn.close() + + +def test_genuine_events_postpone_probes_and_reset_backoff(): + session = Session( + on_probe=lambda s, m: reply(s, m["execution_id"], state="running") + ) + conn = connect(session, stale_query_probe_seconds=SLOW_PROBE) + try: + _, execution_id = execute(conn, session) + assert wait_until(lambda: len(session.probe_times()) >= 2) + query = tracked(conn, execution_id) + assert wait_until(lambda: query.watch.probes_sent >= 2) + # Progress is activity: no probes while it flows, and backoff resets. + deadline = time.monotonic() + SLOW_PROBE * 4 + probes_before = len(session.probe_times()) + while time.monotonic() < deadline: + session.put(kind="execution_progress", execution_id=execution_id) + time.sleep(SLOW_PROBE / 5) + assert wait_until(lambda: query.watch.probes_sent == 0, timeout=1) + assert len(session.probe_times()) == probes_before + finally: + conn.close() + + +def test_probe_not_found_keeps_waiting_for_an_evicted_query(caplog): + # The session's small LRU cache can evict a query that is still running. + # Its terminal state_updated still arrives (from the future's callback), + # but a probe gets "Execution not found". That must not fail the query. + session = Session( + on_probe=lambda s, m: s.put( + kind="error", execution_id=m["execution_id"], message=NOT_FOUND + ) + ) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session, store=STORE) + assert wait_until(lambda: "no longer tracks execution" in caplog.text) + time.sleep(PROBE * 8) + assert results_of(cursor).empty() + assert tracked(conn, execution_id) is not None + # Not probeable any more: no further probes. + assert len(session.probe_times()) == 1 + session.put( + kind="state_updated", + execution_id=execution_id, + state="succeeded", + result_uri=RESULT_URI, + size=3, + ) + result = results_of(cursor).get(timeout=3) + assert result.error is None + assert result.store_result.result_uri == RESULT_URI + finally: + conn.close() + + +def test_progress_while_probe_outstanding_does_not_unmask_not_found(caplog): + # A progress event between the probe and its reply must not make the + # not-found reply look like an answer to a normal request. + def on_probe(s, m): + s.put(kind="execution_progress", execution_id=m["execution_id"]) + s.put(kind="error", execution_id=m["execution_id"], message=NOT_FOUND) + + session = Session(on_probe=on_probe) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session) + assert wait_until(lambda: "no longer tracks execution" in caplog.text) + assert results_of(cursor).empty() + assert tracked(conn, execution_id) is not None + finally: + conn.close() + + +def test_not_found_for_the_normal_results_request_still_fails(): + # After the genuine state_updated: succeeded, there is no other terminal + # event to wait for: a not-found reply fails the query, as without the + # watchdog. + session = Session( + on_probe=lambda s, m: s.put( + kind="error", execution_id=m["execution_id"], message=NOT_FOUND + ) + ) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session) + session.put(kind="state_updated", execution_id=execution_id, state="succeeded") + error = results_of(cursor).get(timeout=3).error + assert isinstance(error, OperationalError) + assert NOT_FOUND in str(error) + assert tracked(conn, execution_id) is None + finally: + conn.close() + + +def test_not_found_after_results_requested_probe_still_fails(): + # The normal retrieve's reply is lost, and the one probe that follows it + # gets not-found: the terminal event has already arrived, so fail. + asked = [] + + def on_probe(s, m): + asked.append(m) + if len(asked) > 1: + s.put(kind="error", execution_id=m["execution_id"], message=NOT_FOUND) + + session = Session(on_probe=on_probe) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session) + session.put(kind="state_updated", execution_id=execution_id, state="succeeded") + error = results_of(cursor).get(timeout=3).error + assert isinstance(error, OperationalError) + assert NOT_FOUND in str(error) + assert len(asked) == 2 + finally: + conn.close() + + +def test_results_requested_gets_one_probe_no_earlier_than_4n(): + session = Session() + conn = connect(session) + try: + _, execution_id = execute(conn, session) + session.put(kind="state_updated", execution_id=execution_id, state="succeeded") + assert wait_until(lambda: len(session.probe_times()) >= 2) + requested, probed = session.probe_times() + assert probed - requested >= PROBE * 4 + # Never a second probe: duplicate retrieves re-send the full result. + time.sleep(PROBE * 16) + assert len(session.probe_times()) == 2 + finally: + conn.close() + + +def test_stalled_user_send_does_not_hold_up_other_results(): + session = Session() + conn = connect(session) + sending, release = threading.Event(), threading.Event() + + def stall_new_queries(message): + if message["kind"] == "execute_sql" and message["statement"] == "SLOW": + sending.set() + release.wait(timeout=5) + + try: + _, quiet = execute(conn, session) + cursor, ready = execute(conn, session, sql="SELECT 2") + session.before_send = stall_new_queries + sender = threading.Thread(target=conn.cursor().execute, args=("SLOW",)) + sender.start() + assert sending.wait(timeout=1) + # Let a probe for the quiet query come due while the send is stalled. + time.sleep(PROBE * 3) + started = time.monotonic() + session.put( + kind="state_updated", + execution_id=ready, + state="cancelled", + ) + result = results_of(cursor).get(timeout=3) + delay = time.monotonic() - started + release.set() + sender.join(timeout=3) + assert result.error is None + assert delay < 1.0, delay + # The skipped probe goes out once the send lock is free. + assert wait_until(lambda: session.probe_times(quiet)) + finally: + release.set() + conn.close() + + +def test_recovery_is_logged_only_for_a_successful_probe_reply(caplog): + caplog.set_level(logging.INFO) + session = Session( + on_probe=lambda s, m: s.incoming.put( + cbor2.dumps( + { + "kind": "execution_result", + "execution_id": m["execution_id"], + "state": "succeeded", + "results": {"result_bytes": b"not json", "format": "json"}, + } + ) + ) + ) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session) + assert isinstance(results_of(cursor).get(timeout=3).error, OperationalError) + assert "recovered" not in caplog.text + finally: + conn.close() + + +def test_watchdog_runs_when_read_timeout_is_none(): + # recv() with no timeout would block forever on a quiet connection, which + # is exactly the lost-terminal-event shape. + session = Session( + on_probe=lambda s, m: reply(s, m["execution_id"], state="running") + ) + conn = Connection(session, read_timeout=None, stale_query_probe_seconds=PROBE) + try: + _, execution_id = execute(conn, session) + assert wait_until(lambda: session.probe_times(execution_id)) + finally: + conn.close() + + +@pytest.mark.parametrize( + "value", + [0, -1, -0.5, float("nan"), float("inf"), float("-inf"), True, "30", object()], +) +def test_invalid_probe_interval_disables_watchdog(value, caplog): + session = Session() + conn = connect(session, stale_query_probe_seconds=value) + try: + assert "stale-query probes disabled" in caplog.text + assert conn._Connection__probe_after is None + _, execution_id = execute(conn, session) + time.sleep(PROBE * 3) + assert session.requests("retrieve_results") == [] + assert tracked(conn, execution_id) is not None + finally: + conn.close() + + +def test_connect_direct_forwards_probe_interval(): + session = Session() + target = "wherobots.db.driver.websockets.sync.client.connect" + with patch(target, return_value=session): + conn = connect_direct("wss://compute/sql", stale_query_probe_seconds=PROBE) + try: + _, execution_id = execute(conn, session) + assert wait_until(lambda: session.probe_times(execution_id)) + finally: + conn.close() diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 5e28b7b..c3cc164 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -3,8 +3,9 @@ import math import textwrap import threading +import time import uuid -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, Callable, Dict from wherobots.db.redaction import get_statement_type, redact_sql @@ -15,7 +16,7 @@ import websockets.exceptions import websockets.sync.client -from .constants import DEFAULT_READ_TIMEOUT_SECONDS +from .constants import DEFAULT_READ_TIMEOUT_SECONDS, DEFAULT_STALE_QUERY_PROBE_SECONDS from ._transport import abort_connection from .cursor import Cursor from .errors import NotSupportedError, OperationalError @@ -69,6 +70,83 @@ def _handshake_bound(ws: Any) -> float | None: return None +_STALE_CHECK_MAX_INTERVAL_SECONDS = 1.0 +"""Upper bound on how often the reader thread scans for stale queries.""" + +_MAX_PROBE_BACKOFF_DOUBLINGS = 3 +"""Probe spacing doubles from N up to 8N (2**3 * N), then stays there.""" + +_RESULTS_PROBE_AFTER_FACTOR = 4 +"""A query whose results were already requested gets one probe, after 4N. + +Its terminal event has arrived; only the retrieve_results reply is missing. +The session re-serializes and re-sends the whole result (up to 100 MiB) for +every retrieve, synchronously on its event loop, so duplicates are costly. A +lost reply is also the rarer case: retrieve replies are sent inline on the +session's main loop, not pushed from a callback like state_updated. +""" + +_EXECUTION_NOT_FOUND = "Execution not found" +"""sql-session's error for a retrieve_results it has no cache entry for. + +Its execution cache is a small LRU that can evict a query that is still +running, whose terminal state_updated still arrives (it's pushed from the +query's own done-callback). So for a probe, this means "can't answer", not +"failed". +""" + +_PROBED_STATES = frozenset( + { + ExecutionState.EXECUTION_REQUESTED, + ExecutionState.PENDING, + ExecutionState.RUNNING, + ExecutionState.RESULTS_REQUESTED, + } +) +"""States in which the driver is waiting on a session event it can re-ask for.""" + + +def _probe_interval(value: Any) -> float | None: + """The validated stale-query probe interval, or ``None`` to disable it. + + Like ``_handshake_bound``, returns ``None`` for a non-numeric (including + ``bool``), non-finite or non-positive value. Never raises ``Exception``: + ``connect()`` callers pass this through, and a bad value must not cost + them their connection. + """ + try: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + value = float(value) + if not math.isfinite(value) or value <= 0: + return None + return value + except Exception: + return None + + +@dataclass +class _Watch: + """Stale-query watchdog bookkeeping for one query. + + Only the reader thread touches it once the query is registered; ``__send`` + stamps ``last_activity`` under ``__lock`` just before registration. + """ + + last_activity: float = field(default_factory=time.monotonic) + """``time.monotonic()`` of the last inbound event (or probe) for the query.""" + probes_sent: int = 0 + """Probes sent since the last event that wasn't a reply to a retrieve.""" + probe_outstanding: bool = False + """A probe was sent and no reply to a retrieve has arrived since.""" + probeable: bool = True + """False once the session has said it can't answer probes for the query.""" + results_requested: bool = False + """The normal (non-probe) retrieve_results was sent.""" + results_probed: bool = False + """The one probe allowed after ``results_requested`` was sent.""" + + @dataclass class Query: sql: str @@ -76,6 +154,7 @@ class Query: state: ExecutionState handler: Callable[[Any], None] store: Store | None = None + watch: _Watch = field(default_factory=_Watch) class Connection: @@ -101,13 +180,29 @@ def __init__( data_compression: DataCompression | None = None, geometry_representation: GeometryRepresentation | None = None, session_id: str | None = None, + stale_query_probe_seconds: float | None = DEFAULT_STALE_QUERY_PROBE_SECONDS, ): self.__ws = ws - self.__read_timeout = read_timeout self.__results_format = results_format self.__data_compression = data_compression self.__geometry_representation = geometry_representation self.__progress_handler: ProgressHandler | None = None + self.__probe_after = _probe_interval(stale_query_probe_seconds) + if self.__probe_after is None and stale_query_probe_seconds is not None: + logging.warning( + "Invalid stale_query_probe_seconds (%s); stale-query probes disabled.", + type(stale_query_probe_seconds).__name__, + ) + self.__next_stale_check = 0.0 + # recv() must wake up for the watchdog's checks, also when the caller + # asked for no read timeout (websockets accepts None) or a long one. + self.__recv_timeout = read_timeout + if self.__probe_after is not None: + check = min(_STALE_CHECK_MAX_INTERVAL_SECONDS, self.__probe_after) + if read_timeout is None or ( + isinstance(read_timeout, (int, float)) and read_timeout > check + ): + self.__recv_timeout = check self.__session_id = session_id self.__lock = threading.Lock() @@ -200,6 +295,10 @@ def __receive_loop(self) -> None: while True: if self.__shutdown_done.is_set(): return + # Checked on every iteration, not only when recv() times out: on a + # busy connection, steady traffic for one query would otherwise + # starve another query's watchdog. + self.__probe_stale_queries() try: self.__listen() except TimeoutError: @@ -214,6 +313,67 @@ def __receive_loop(self) -> None: except Exception as e: logging.exception("Error handling message from SQL session", exc_info=e) + def __probe_stale_queries(self) -> None: + """Re-ask the session about queries that have gone quiet. + + Defence in depth for a lost terminal event: the driver otherwise trusts + a single pushed ``state_updated`` and would wait forever if it never + arrives. A probe is a ``retrieve_results`` request; the session answers + with the query's current state (and its results or result location + once it has succeeded), or with an error for an unknown execution. + + Runs on the reader thread, which owns every ``Query``'s watchdog fields. + Never raises: a watchdog failure must not disturb the receive loop. + """ + probe_after = self.__probe_after + if probe_after is None: + return + try: + now = time.monotonic() + if now < self.__next_stale_check: + return + self.__next_stale_check = now + min( + _STALE_CHECK_MAX_INTERVAL_SECONDS, probe_after + ) + with self.__lock: + if self.__closed: + return + candidates = list(self.__queries.values()) + for query in candidates: + watch = query.watch + if not watch.probeable or query.state not in _PROBED_STATES: + continue + if watch.results_requested: + if watch.results_probed: + continue + wait = probe_after * _RESULTS_PROBE_AFTER_FACTOR + else: + doublings = min(watch.probes_sent, _MAX_PROBE_BACKOFF_DOUBLINGS) + wait = probe_after * 2**doublings + silence = now - watch.last_activity + if silence < wait: + continue + if not self.__request_results(query.execution_id, probe=True): + # A user thread is sending, possibly stalled. Don't wait + # behind it: that would stop recv() for every query. The + # lock is busy for the rest too; retry at the next check. + break + watch.probes_sent += 1 + watch.probe_outstanding = True + watch.results_probed = watch.results_requested + # Space the next probe from this one, not from the last event. + watch.last_activity = now + logging.info( + "No events for query %s in %.1fs while %s; probed the SQL " + "session for its state (probe %d).", + query.execution_id, + silence, + query.state, + watch.probes_sent, + ) + except Exception: + logging.exception("Stale-query probe failed") + def __connection_error(self, execution_id: str) -> OperationalError: message = ( f"SQL connection lost (session={self.__session_id or 'unknown'}, " @@ -299,6 +459,25 @@ def __listen(self) -> None: # Invalid event. return + query = self.__queries.get(execution_id) + # Whether this event answers a probe (possibly along with the normal + # retrieve). Read before the bookkeeping below resets it. + probe_reply = False + if query is not None: + watch = query.watch + # Any inbound event, progress included, shows the query is alive. + watch.last_activity = time.monotonic() + if kind == EventKind.EXECUTION_RESULT or kind == EventKind.ERROR: + # A reply to a retrieve_results, which may be our own probe: it + # resets the timer but not the backoff, so a long-running + # silent query isn't probed every N seconds forever. + probe_reply = watch.probe_outstanding + watch.probe_outstanding = False + else: + watch.probes_sent = 0 + if kind == EventKind.STATE_UPDATED: + watch.probe_outstanding = False + # Progress events are independent of the query state machine and don't # require a tracked query — the handler is connection-level. if kind == EventKind.EXECUTION_PROGRESS: @@ -318,7 +497,6 @@ def __listen(self) -> None: logging.exception("Progress handler raised an exception") return - query = self.__queries.get(execution_id) if not query: logging.warning( "Received %s event for unknown execution ID %s", kind, execution_id @@ -332,6 +510,18 @@ def complete_query(result: ExecutionResult) -> None: with self.__lock: claimed = self.__queries.pop(execution_id, None) if claimed is not None: + if probe_reply and result.error is None: + if claimed.watch.results_requested: + # The normal retrieve's reply may have won the race. + logging.info( + "Query %s results arrived after a stale-query probe.", + execution_id, + ) + else: + logging.info( + "Query %s recovered from a stale-query probe reply.", + execution_id, + ) claimed.handler(result) def fail_query(action: str, error: Exception) -> None: @@ -428,8 +618,25 @@ def fail_query(action: str, error: Exception) -> None: # details. pass elif kind == EventKind.ERROR: - query.state = ExecutionState.FAILED error = message.get("message") + if ( + probe_reply + and error == _EXECUTION_NOT_FOUND + and not query.watch.results_requested + and query.state != ExecutionState.FAILED + ): + # The session evicted the query from its cache, but it may well + # still be running; its terminal event will still arrive. After + # the normal retrieve, though, there's nothing left to wait for. + query.watch.probeable = False + logging.warning( + "SQL session no longer tracks execution %s (likely evicted " + "from its result cache); cannot probe it, waiting for its " + "terminal event.", + execution_id, + ) + return + query.state = ExecutionState.FAILED complete_query(ExecutionResult(error=OperationalError(error))) else: logging.warning("Received unknown %s event!", kind) @@ -456,7 +663,13 @@ def _handle_results(self, execution_id: str, results: Dict[str, Any]) -> Any: else: raise NotSupportedError("Unsupported results format") - def __send(self, message: Dict[str, Any], query: Query | None = None) -> None: + def __send( + self, + message: Dict[str, Any], + query: Query | None = None, + blocking: bool = True, + ) -> bool: + """Send a request; return ``False`` only if non-blocking and busy.""" # Serialization and redaction are local work. Fail before registration, # without poisoning unrelated cursors or misreporting a transport loss. request = json.dumps(message) @@ -465,16 +678,20 @@ def __send(self, message: Dict[str, Any], query: Query | None = None) -> None: # unguarded call would redact on every request even with DEBUG off. if logging.getLogger().isEnabledFor(logging.DEBUG): logging.debug("Request: %s", self.__redacted_request(message)) - with self.__send_lock: + if not self.__send_lock.acquire(blocking=blocking): + return False + try: with self.__lock: if self.__closed: if query is not None: raise self.__connection_error(query.execution_id) - return + return True if query is not None: + # The watchdog's clock starts at registration. + query.watch.last_activity = time.monotonic() self.__queries[query.execution_id] = query elif message.get("execution_id") not in self.__queries: - return + return True try: self.__ws.send(request) except (websockets.exceptions.ConnectionClosed, OSError): @@ -486,8 +703,11 @@ def __send(self, message: Dict[str, Any], query: Query | None = None) -> None: self.__queries.pop(query.execution_id, None) raise else: - return + return True + finally: + self.__send_lock.release() self.__fail_pending() + return True @staticmethod def __redacted_request(message: Dict[str, Any]) -> str: @@ -504,7 +724,7 @@ def __redacted_request(message: Dict[str, Any]) -> str: def __recv(self) -> Dict[str, Any]: try: - frame = self.__ws.recv(timeout=self.__read_timeout) + frame = self.__ws.recv(timeout=self.__recv_timeout) except TimeoutError: raise # Idle polls are expected, not terminal I/O failures. except OSError as e: @@ -559,10 +779,14 @@ def __execute_sql( ) return execution_id - def __request_results(self, execution_id: str) -> None: + def __request_results(self, execution_id: str, probe: bool = False) -> bool: + """Send retrieve_results; ``False`` if a probe found the sender busy. + + A probe never waits for ``__send_lock`` (see ``__probe_stale_queries``). + """ query = self.__queries.get(execution_id) if not query: - return + return True request = { "kind": RequestKind.RETRIEVE_RESULTS.value, @@ -575,9 +799,12 @@ def __request_results(self, execution_id: str) -> None: if self.__geometry_representation: request["geometry"] = self.__geometry_representation.value - query.state = ExecutionState.RESULTS_REQUESTED - logging.info("Requesting results from %s ...", execution_id) - self.__send(request) + if not probe: + # A probe only asks; the reply carries the query's actual state. + query.state = ExecutionState.RESULTS_REQUESTED + query.watch.results_requested = True + logging.info("Requesting results from %s ...", execution_id) + return self.__send(request, blocking=not probe) def __cancel_query(self, execution_id: str) -> None: """Cancels the query with the given execution ID.""" diff --git a/wherobots/db/constants.py b/wherobots/db/constants.py index 0b409cd..965bd59 100644 --- a/wherobots/db/constants.py +++ b/wherobots/db/constants.py @@ -17,6 +17,8 @@ # Bound on the WebSocket close handshake performed by a graceful close(). DEFAULT_CLOSE_TIMEOUT_SECONDS: float = 1.0 DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS: float = 900 +# Silence, per query, before the driver re-asks the session for its state. +DEFAULT_STALE_QUERY_PROBE_SECONDS: float = 30.0 MAX_MESSAGE_SIZE: int = 100 * 2**20 # 100MiB PROTOCOL_VERSION: Version = Version("1.0.0") diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index 932cff6..399bced 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -25,6 +25,7 @@ DEFAULT_READ_TIMEOUT_SECONDS, DEFAULT_SESSION_TYPE, DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS, + DEFAULT_STALE_QUERY_PROBE_SECONDS, MAX_MESSAGE_SIZE, PARAM_STYLE, PROTOCOL_VERSION, @@ -125,6 +126,7 @@ def connect( geometry_representation: Union[GeometryRepresentation, None] = None, cancel_event: Union[threading.Event, None] = None, extra_headers: Union[Dict[str, str], None] = None, + stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS, ) -> Connection: """Create a connection to a Wherobots SQL session. @@ -146,6 +148,12 @@ def connect( own ``client=dbapi;ver=`` hop to the right of it. This header is advisory only and never affects authentication; ``extra_headers`` cannot be used to override the ``Authorization``/``X-API-Key`` headers. + :param stale_query_probe_seconds: How long a running query may go without + any event from the SQL session before the driver re-asks the session + for its state, in case its completion event was lost. Later probes + back off to at most 8 times this interval. ``None`` disables probing; + an invalid value (non-numeric, non-finite, or not positive) also + disables it. """ if not token and not api_key: raise ValueError("At least one of `token` or `api_key` is required") @@ -271,6 +279,7 @@ def get_session_uri() -> str: session_id=urllib.parse.urlparse(session_id_url) .path.rstrip("/") .rsplit("/", 1)[-1], + stale_query_probe_seconds=stale_query_probe_seconds, ) @@ -299,7 +308,12 @@ def connect_direct( geometry_representation: Union[GeometryRepresentation, None] = None, cancel_event: Union[threading.Event, None] = None, session_id: str | None = None, + stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS, ) -> Connection: + """Connect to a SQL session's WebSocket endpoint directly. + + :param stale_query_probe_seconds: See :func:`connect`. + """ uri_with_protocol = f"{uri}/{protocol}" ssl_context = ssl.create_default_context() ssl_context.load_verify_locations(certifi.where()) @@ -344,4 +358,5 @@ def ws_connect() -> websockets.sync.client.ClientConnection: data_compression=data_compression, geometry_representation=geometry_representation, session_id=session_id, + stale_query_probe_seconds=stale_query_probe_seconds, ) diff --git a/wherobots/db/types.py b/wherobots/db/types.py index 4c0745e..a1db45f 100644 --- a/wherobots/db/types.py +++ b/wherobots/db/types.py @@ -9,6 +9,9 @@ class ExecutionState(LowercaseStrEnum): EXECUTION_REQUESTED = auto() "Execution of a query has been requested by the driver." + PENDING = auto() + "The SQL session has accepted the query but not started running it yet." + RUNNING = auto() "The SQL session has reported the query is running." From 470f6919ed56b0ab7595ab0508d023fd3bec7997 Mon Sep 17 00:00:00 2001 From: Simon Fishel Date: Fri, 2 Oct 2026 13:35:29 -0700 Subject: [PATCH 2/5] feat: recover a store query's result location from a probe reply (WBC-1179) A store-only execution has no cached Arrow table, so the session answers retrieve_results for it with succeeded and no results. Sessions with the probe contract (sql-session step 16) also echo the top-level result_uri and size, as state_updated carries them. On an execution_result with a result_uri, complete with that StoreResult, as the state_updated path does. The normal path never requests results for a store-configured query, so such a reply can only answer a probe. Succeeded without a result_uri on a store query is an older session or an empty store result, which can't be told apart. Log a warning and keep waiting rather than complete empty, which would silently drop an old session's real result, and stop probing it. Known limitation: an empty store result whose terminal event was lost waits as it does today. --- tests/test_watchdog.py | 85 ++++++++++++++++++++++++++++++++++++++ wherobots/db/connection.py | 54 +++++++++++++++++------- 2 files changed, 125 insertions(+), 14 deletions(-) diff --git a/tests/test_watchdog.py b/tests/test_watchdog.py index b783366..e218ffc 100644 --- a/tests/test_watchdog.py +++ b/tests/test_watchdog.py @@ -128,6 +128,91 @@ def reply(session, execution_id, state="succeeded", **fields): ) +def test_probe_reply_with_result_uri_completes_store_query(caplog): + caplog.set_level(logging.INFO) + session = Session( + on_probe=lambda s, m: reply( + s, m["execution_id"], result_uri=RESULT_URI, size=42 + ) + ) + conn = connect(session) + try: + # The terminal state_updated is never delivered. + cursor, execution_id = execute(conn, session, store=STORE) + result = results_of(cursor).get(timeout=3) + assert result.error is None + assert result.store_result.result_uri == RESULT_URI + assert result.store_result.size == 42 + assert tracked(conn, execution_id) is None + assert len(session.probe_times()) == 1 + assert f"No events for query {execution_id}" in caplog.text + assert f"Query {execution_id} recovered" in caplog.text + finally: + conn.close() + + +def test_store_query_succeeded_without_result_uri_keeps_waiting(caplog): + caplog.set_level(logging.INFO) + session = Session(on_probe=lambda s, m: reply(s, m["execution_id"])) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session, store=STORE) + assert wait_until(lambda: "no result location" in caplog.text) + warnings = [r for r in caplog.records if "no result location" in r.message] + assert warnings[0].levelno == logging.WARNING + # Not completed empty, and not re-probed: the answer wouldn't change. + time.sleep(PROBE * 8) + assert results_of(cursor).empty() + assert tracked(conn, execution_id).state == ExecutionState.SUCCEEDED + assert len(session.probe_times()) == 1 + # The genuine terminal event still completes it. + session.put( + kind="state_updated", + execution_id=execution_id, + state="succeeded", + result_uri=RESULT_URI, + size=7, + ) + assert results_of(cursor).get(timeout=3).store_result.result_uri == RESULT_URI + finally: + conn.close() + + +@pytest.mark.parametrize("late_event_first", [True, False]) +def test_late_state_updated_racing_probe_reply_delivers_once(late_event_first): + def on_probe(s, m): + late = dict( + kind="state_updated", + execution_id=m["execution_id"], + state="succeeded", + result_uri=RESULT_URI, + size=1, + ) + probe_reply = dict( + kind="execution_result", + execution_id=m["execution_id"], + state="succeeded", + results=None, + result_uri=RESULT_URI, + size=1, + ) + for message in (late, probe_reply) if late_event_first else (probe_reply, late): + s.incoming.put(message) + + session = Session(on_probe=on_probe) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session, store=STORE) + assert results_of(cursor).get(timeout=3).store_result.result_uri == RESULT_URI + assert wait_until(session.incoming.empty) + time.sleep(PROBE * 4) + assert results_of(cursor).empty() + assert tracked(conn, execution_id) is None + assert len(session.probe_times()) == 1 + finally: + conn.close() + + def test_disabled_watchdog_never_probes(): session = Session() conn = connect(session, stale_query_probe_seconds=None) diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index c3cc164..62ca4bf 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -538,6 +538,18 @@ def fail_query(action: str, error: Exception) -> None: logging.error("%s", message) complete_query(ExecutionResult(error=OperationalError(message))) + def complete_with_store_result(result_uri: str) -> None: + # Results are stored in cloud storage + store_result = StoreResult(result_uri=result_uri, size=message.get("size")) + logging.info( + "Query %s results stored at: %s (size: %s)", + execution_id, + result_uri, + store_result.size, + ) + query.state = ExecutionState.COMPLETED + complete_query(ExecutionResult(store_result=store_result)) + # Incoming state transitions are handled here. if kind == EventKind.STATE_UPDATED or kind == EventKind.EXECUTION_RESULT: try: @@ -553,19 +565,7 @@ def fail_query(action: str, error: Exception) -> None: if kind == EventKind.STATE_UPDATED: result_uri = message.get("result_uri") if result_uri: - # Results are stored in cloud storage - store_result = StoreResult( - result_uri=result_uri, - size=message.get("size"), - ) - logging.info( - "Query %s results stored at: %s (size: %s)", - execution_id, - result_uri, - store_result.size, - ) - query.state = ExecutionState.COMPLETED - complete_query(ExecutionResult(store_result=store_result)) + complete_with_store_result(result_uri) return if query.store is not None: @@ -587,8 +587,34 @@ def fail_query(action: str, error: Exception) -> None: fail_query("request", error) return - # Otherwise, process the results from the execution_result event. + # Otherwise, this execution_result answers a retrieve_results. + # For a store-configured query, that can only be a stale-query + # probe: the normal path never requests results for one. + # Sessions with the probe contract echo the store location. + result_uri = message.get("result_uri") + if result_uri: + complete_with_store_result(result_uri) + return + results = message.get("results") + if query.store is not None and (results is None or results == {}): + # Succeeded without a result location: either an older + # session that doesn't echo it, or an empty store result + # (no store_path). The two are indistinguishable, and + # completing empty would silently drop an old session's + # real result. Keep waiting for the genuine state_updated, + # exactly as without the watchdog, and don't probe again: + # the answer won't change. Known limitation: an empty store + # result whose terminal event was lost can't be recovered. + query.watch.probeable = False + logging.warning( + "Query %s succeeded, but the SQL session's reply has no " + "result location (older SQL session, or an empty store " + "result); cannot recover it, still waiting.", + execution_id, + ) + return + if results is None or results == {}: logging.warning("Got no results back from %s.", execution_id) query.state = ExecutionState.COMPLETED From 104ee7e7187303e69d147f6a616402d96d735324 Mon Sep 17 00:00:00 2001 From: Simon Fishel Date: Fri, 9 Oct 2026 11:54:17 -0700 Subject: [PATCH 3/5] fix: log the true silence and the invalid value in watchdog messages (WBC-1179) The probe log reported silence since the last probe, not the last event: each probe resets last_activity to space the next one, so once probes go unanswered (the case worth diagnosing) it under-reported, e.g. 60s for a real 90s. Track the last inbound event separately in _Watch.last_event (set on registration and inbound messages only) and log that, noting when the previous probe went unanswered. Probe timing is unchanged. The invalid stale_query_probe_seconds warning logged only the type name, so 0 read as "(int)". Log a truncated repr, falling back to the type name if repr raises, since this runs in the constructor. Reword the eviction-guard comment: the terminal event still arrives only if the query is still running; if it was already lost, nothing more can be done. --- tests/test_watchdog.py | 62 ++++++++++++++++++++++++++++++++++++++ wherobots/db/connection.py | 59 ++++++++++++++++++++++++++++++------ 2 files changed, 111 insertions(+), 10 deletions(-) diff --git a/tests/test_watchdog.py b/tests/test_watchdog.py index e218ffc..34341c3 100644 --- a/tests/test_watchdog.py +++ b/tests/test_watchdog.py @@ -3,6 +3,7 @@ import json import logging import queue +import re import threading import time from unittest.mock import MagicMock, patch @@ -250,6 +251,33 @@ def stream_progress(): conn.close() +def test_unanswered_probe_log_reports_silence_since_the_last_event(caplog): + caplog.set_level(logging.INFO) + session = Session() # Never answers probes. + conn = connect(session, stale_query_probe_seconds=SLOW_PROBE) + try: + _, execution_id = execute(conn, session) + assert wait_until(lambda: len(session.probe_times()) >= 2, timeout=5) + last_event = tracked(conn, execution_id).watch.last_event + second_probe_at = session.probe_times()[1] + finally: + conn.close() + logs = [ + r.getMessage() for r in caplog.records if "No events for query" in r.message + ] + assert len(logs) >= 2, logs + first, second = logs[:2] + assert "(probe 1)" in first + assert "(probe 2, previous probe unanswered)" in second + logged = float(re.search(r" in ([0-9.]+)s while ", second).group(1)) + # Probe 1 goes out N after the last event, probe 2 a further 2N later: the + # silence is >= 3N, not the 2N since probe 1. + assert logged >= SLOW_PROBE * 3 - 0.05, logged + # Loose on purpose: the >= 3N check above is what catches a regression; + # this only confirms the figure tracks the measured silence. + assert abs(logged - (second_probe_at - last_event)) < SLOW_PROBE, logged + + def test_successive_probes_back_off_to_a_cap(): session = Session( on_probe=lambda s, m: reply(s, m["execution_id"], state="running") @@ -576,6 +604,40 @@ def test_invalid_probe_interval_disables_watchdog(value, caplog): conn.close() +class _BadRepr: + def __repr__(self): + raise RuntimeError("no repr") + + +@pytest.mark.parametrize( + "value, expected", + [(0, "(0)"), (-5, "(-5)"), (_BadRepr(), "()")], +) +def test_invalid_probe_interval_warning_names_the_value(value, expected, caplog): + conn = connect(Session(), stale_query_probe_seconds=value) + try: + assert conn._Connection__probe_after is None + warnings = [r.getMessage() for r in caplog.records if "disabled" in r.message] + assert warnings == [ + f"Invalid stale_query_probe_seconds {expected}; " + "stale-query probes disabled." + ] + finally: + conn.close() + + +def test_invalid_probe_interval_warning_truncates_a_huge_value(caplog): + conn = connect(Session(), stale_query_probe_seconds=10**400) + try: + assert conn._Connection__probe_after is None + (warning,) = [r.getMessage() for r in caplog.records if "disabled" in r.message] + quoted = re.search(r"\((.*)\);", warning).group(1) + assert quoted.startswith("1000") and quoted.endswith("...") + assert len(quoted) <= 80 + finally: + conn.close() + + def test_connect_direct_forwards_probe_interval(): session = Session() target = "wherobots.db.driver.websockets.sync.client.connect" diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 62ca4bf..2ca562c 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -125,16 +125,46 @@ def _probe_interval(value: Any) -> float | None: return None +_INVALID_VALUE_REPR_CHARS = 80 +"""Longest repr of an invalid setting that a warning quotes.""" + + +def _describe_invalid(value: Any) -> str: + """A bounded ``repr`` of ``value`` for a warning. Never raises ``Exception``. + + Truncated so a huge value (``10**400``) can't flood the log; falls back to + the type name if ``repr`` fails, since this runs in the constructor. + """ + try: + text = repr(value) + except Exception: + try: + return f"" + except Exception: + return "" + if len(text) > _INVALID_VALUE_REPR_CHARS: + text = text[: _INVALID_VALUE_REPR_CHARS - 3] + "..." + return text + + @dataclass class _Watch: """Stale-query watchdog bookkeeping for one query. Only the reader thread touches it once the query is registered; ``__send`` - stamps ``last_activity`` under ``__lock`` just before registration. + stamps ``last_activity`` and ``last_event`` under ``__lock`` just before + registration. """ last_activity: float = field(default_factory=time.monotonic) - """``time.monotonic()`` of the last inbound event (or probe) for the query.""" + """``time.monotonic()`` of the last inbound event (or probe) for the query. + + Spaces the probes; see ``last_event`` for the true silence.""" + last_event: float = field(default_factory=time.monotonic) + """``time.monotonic()`` of the last inbound event (or registration). + + Unlike ``last_activity``, probes don't update it, so the probe log reports + how long the query has really been silent.""" probes_sent: int = 0 """Probes sent since the last event that wasn't a reply to a retrieve.""" probe_outstanding: bool = False @@ -191,7 +221,7 @@ def __init__( if self.__probe_after is None and stale_query_probe_seconds is not None: logging.warning( "Invalid stale_query_probe_seconds (%s); stale-query probes disabled.", - type(stale_query_probe_seconds).__name__, + _describe_invalid(stale_query_probe_seconds), ) self.__next_stale_check = 0.0 # recv() must wake up for the watchdog's checks, also when the caller @@ -358,6 +388,10 @@ def __probe_stale_queries(self) -> None: # behind it: that would stop recv() for every query. The # lock is busy for the rest too; retry at the next check. break + # Replies to a retrieve and state_updated clear the flag; + # progress doesn't. Still set means the last probe went + # unanswered. + unanswered = watch.probe_outstanding watch.probes_sent += 1 watch.probe_outstanding = True watch.results_probed = watch.results_requested @@ -365,11 +399,12 @@ def __probe_stale_queries(self) -> None: watch.last_activity = now logging.info( "No events for query %s in %.1fs while %s; probed the SQL " - "session for its state (probe %d).", + "session for its state (probe %d%s).", query.execution_id, - silence, + now - watch.last_event, query.state, watch.probes_sent, + ", previous probe unanswered" if unanswered else "", ) except Exception: logging.exception("Stale-query probe failed") @@ -466,7 +501,7 @@ def __listen(self) -> None: if query is not None: watch = query.watch # Any inbound event, progress included, shows the query is alive. - watch.last_activity = time.monotonic() + watch.last_activity = watch.last_event = time.monotonic() if kind == EventKind.EXECUTION_RESULT or kind == EventKind.ERROR: # A reply to a retrieve_results, which may be our own probe: it # resets the timer but not the backoff, so a long-running @@ -651,9 +686,12 @@ def complete_with_store_result(result_uri: str) -> None: and not query.watch.results_requested and query.state != ExecutionState.FAILED ): - # The session evicted the query from its cache, but it may well - # still be running; its terminal event will still arrive. After - # the normal retrieve, though, there's nothing left to wait for. + # The session evicted the query from its cache. If the query is + # still running, its terminal event will still arrive; if that + # event was already lost, nothing more can be done (still + # running and finished-with-a-lost-event look the same from + # here), so keep waiting either way. After the normal + # retrieve, though, there's nothing left to wait for. query.watch.probeable = False logging.warning( "SQL session no longer tracks execution %s (likely evicted " @@ -714,7 +752,8 @@ def __send( return True if query is not None: # The watchdog's clock starts at registration. - query.watch.last_activity = time.monotonic() + watch = query.watch + watch.last_activity = watch.last_event = time.monotonic() self.__queries[query.execution_id] = query elif message.get("execution_id") not in self.__queries: return True From a0840447ab9f6fae628adf412ef88028845a738f Mon Sep 17 00:00:00 2001 From: Simon Fishel Date: Fri, 9 Oct 2026 12:11:23 -0700 Subject: [PATCH 4/5] fix: keep probe attribution across non-terminal state updates (WBC-1179) A probe for an evicted queued or running query can be followed by a genuine state_updated: running before the probe's "Execution not found" reply. Clearing probe_outstanding on that update made the reply look like an answer to a normal retrieve, so it bypassed the eviction guard and failed a query that could still complete. Only replies to a retrieve (execution_result/error) clear the flag now. The not-found cases that must still fail are gated by results_requested, which the normal retrieve sets after a terminal success. --- tests/test_watchdog.py | 34 ++++++++++++++++++++++++++++++++++ wherobots/db/connection.py | 17 +++++++++++------ 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/tests/test_watchdog.py b/tests/test_watchdog.py index 34341c3..527af2b 100644 --- a/tests/test_watchdog.py +++ b/tests/test_watchdog.py @@ -451,6 +451,40 @@ def on_probe(s, m): conn.close() +def test_running_state_update_while_probe_outstanding_does_not_unmask_not_found( + caplog, +): + # An evicted query that is still running: a genuine state_updated: running + # can arrive between the probe and its not-found reply. The reply is still + # the probe's and must not fail the query. + def on_probe(s, m): + s.put(kind="state_updated", execution_id=m["execution_id"], state="running") + s.put(kind="error", execution_id=m["execution_id"], message=NOT_FOUND) + + session = Session(on_probe=on_probe) + conn = connect(session) + try: + cursor, execution_id = execute(conn, session, store=STORE) + assert wait_until(lambda: "no longer tracks execution" in caplog.text) + time.sleep(PROBE * 8) + assert results_of(cursor).empty() + assert tracked(conn, execution_id) is not None + # Not probeable any more: no further probes. + assert len(session.probe_times()) == 1 + session.put( + kind="state_updated", + execution_id=execution_id, + state="succeeded", + result_uri=RESULT_URI, + size=3, + ) + result = results_of(cursor).get(timeout=3) + assert result.error is None + assert result.store_result.result_uri == RESULT_URI + finally: + conn.close() + + def test_not_found_for_the_normal_results_request_still_fails(): # After the genuine state_updated: succeeded, there is no other terminal # event to wait for: a not-found reply fails the query, as without the diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 2ca562c..712c7d0 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -168,7 +168,10 @@ class _Watch: probes_sent: int = 0 """Probes sent since the last event that wasn't a reply to a retrieve.""" probe_outstanding: bool = False - """A probe was sent and no reply to a retrieve has arrived since.""" + """A probe was sent and no reply to a retrieve has arrived since. + + Other events, a non-terminal ``state_updated`` included, leave it set: they + don't answer the probe, whose reply may still be on its way.""" probeable: bool = True """False once the session has said it can't answer probes for the query.""" results_requested: bool = False @@ -388,9 +391,9 @@ def __probe_stale_queries(self) -> None: # behind it: that would stop recv() for every query. The # lock is busy for the rest too; retry at the next check. break - # Replies to a retrieve and state_updated clear the flag; - # progress doesn't. Still set means the last probe went - # unanswered. + # Only a reply to a retrieve clears the flag; other events, + # state_updated included, don't. Still set means the last + # probe went unanswered. unanswered = watch.probe_outstanding watch.probes_sent += 1 watch.probe_outstanding = True @@ -509,9 +512,11 @@ def __listen(self) -> None: probe_reply = watch.probe_outstanding watch.probe_outstanding = False else: + # A genuine event resets the backoff, but leaves an outstanding + # probe outstanding: a non-terminal state_updated can arrive + # between a probe and its reply, which must still be + # recognised as the probe's (see the not-found guard below). watch.probes_sent = 0 - if kind == EventKind.STATE_UPDATED: - watch.probe_outstanding = False # Progress events are independent of the query state machine and don't # require a tracked query — the handler is connection-level. From 9fd773181e298146cb8e8bde578cf49baf2c68f4 Mon Sep 17 00:00:00 2001 From: Simon Fishel Date: Fri, 9 Oct 2026 12:14:02 -0700 Subject: [PATCH 5/5] fix: don't let a probe block on a full socket (WBC-1179) A probe runs on the reader thread. Its non-blocking acquire of the send lock only avoids waiting behind another sender; its own ws.send() can still block when the socket's send buffer is full, e.g. a peer that keeps sending but stops reading. The reader then stops processing every query, and keepalive can't help: the stalled sendall() holds websockets' protocol mutex, which pings need too. On the non-blocking path, check the socket is writable first and treat a full one like a busy lock: skip, and retry at the next check. The check uses poll (or a selector), so fds above FD_SETSIZE work, and never raises: anything unusable, including test fakes, counts as writable. --- tests/test_watchdog.py | 124 +++++++++++++++++++++++++++++++++++++ wherobots/db/_transport.py | 33 ++++++++++ wherobots/db/connection.py | 27 ++++++-- 3 files changed, 178 insertions(+), 6 deletions(-) diff --git a/tests/test_watchdog.py b/tests/test_watchdog.py index 527af2b..43d58b9 100644 --- a/tests/test_watchdog.py +++ b/tests/test_watchdog.py @@ -2,8 +2,10 @@ import json import logging +import os import queue import re +import socket import threading import time from unittest.mock import MagicMock, patch @@ -12,6 +14,7 @@ import pytest from websockets.exceptions import ConnectionClosedOK +from wherobots.db._transport import socket_writable from wherobots.db.connection import Connection from wherobots.db.driver import connect_direct from wherobots.db.errors import OperationalError @@ -583,6 +586,127 @@ def stall_new_queries(message): conn.close() +def fill_send_buffer(sock): + """Write to ``sock`` until its send buffer is full; restore its timeout.""" + timeout = sock.gettimeout() + sock.setblocking(False) + try: + # Large writes first, then single bytes: some kernels accept a last + # partial write. + for chunk in (b"x" * 65536, b"x"): + try: + while True: + sock.send(chunk) + except BlockingIOError: + pass + finally: + sock.settimeout(timeout) + + +class SocketSession(Session): + """A fake session whose requests are written to a real socket. + + Like websockets' sync client, ``send`` blocks while the socket's send + buffer is full (bounded here, so a regression fails instead of hanging). + """ + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.socket, self.peer = socket.socketpair() + self.socket.settimeout(5) + self.before_send = lambda message: self.socket.sendall( + json.dumps(message).encode() + ) + + def fill(self): + """Fill the send buffer, as a peer that stops reading would.""" + fill_send_buffer(self.socket) + + def drain(self): + self.peer.setblocking(False) + try: + while self.peer.recv(1 << 20): + pass + except BlockingIOError: + pass + + def dispose(self): + self.socket.close() + self.peer.close() + + +def test_probe_skips_a_full_socket_instead_of_stalling_the_reader(): + session = SocketSession() + conn = connect(session) + try: + _, quiet = execute(conn, session) + cursor, ready = execute(conn, session, sql="SELECT 2") + session.fill() + assert not socket_writable(session) + # Let a probe for the quiet query come due while the socket is full. + time.sleep(PROBE * 3) + started = time.monotonic() + session.put(kind="state_updated", execution_id=ready, state="cancelled") + result = results_of(cursor).get(timeout=3) + delay = time.monotonic() - started + assert result.error is None + assert delay < 1.0, delay + assert session.probe_times(quiet) == [] + assert tracked(conn, quiet) is not None + # The skipped probe goes out once the peer reads again. + session.drain() + assert wait_until(lambda: session.probe_times(quiet)) + finally: + conn.close() + session.dispose() + + +def test_socket_writable_proceeds_without_a_real_socket(): + # A failure must not disable probing: anything unusable counts as writable. + assert socket_writable(Session()) # MagicMock socket + assert socket_writable(object()) # No socket at all + + class Broken: + @property + def socket(self): + raise RuntimeError("no socket") + + assert socket_writable(Broken()) + closed = MagicMock() + closed.socket, peer = socket.socketpair() + closed.socket.close() + peer.close() + assert socket_writable(closed) + + +def test_socket_writable_handles_fds_above_fd_setsize(): + # select.select() rejects fds >= FD_SETSIZE (1024); a server holding many + # connections reaches them. + resource = pytest.importorskip("resource") # POSIX-only + target = 2048 + soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) + if soft <= target: + if hard != resource.RLIM_INFINITY and hard <= target: + pytest.skip("RLIMIT_NOFILE too low for a high fd") + resource.setrlimit(resource.RLIMIT_NOFILE, (target + 1, hard)) + a, b = socket.socketpair() + high = None + try: + os.dup2(a.fileno(), target) + high = socket.socket(fileno=target) + ws = MagicMock() + ws.socket = high + assert socket_writable(ws) + fill_send_buffer(high) + assert not socket_writable(ws) + finally: + if high is not None: + high.close() + a.close() + b.close() + resource.setrlimit(resource.RLIMIT_NOFILE, (soft, hard)) + + def test_recovery_is_logged_only_for_a_successful_probe_reply(caplog): caplog.set_level(logging.INFO) session = Session( diff --git a/wherobots/db/_transport.py b/wherobots/db/_transport.py index 8f7989c..e9017c5 100644 --- a/wherobots/db/_transport.py +++ b/wherobots/db/_transport.py @@ -1,6 +1,8 @@ """Transport termination independent of the WebSocket protocol's send lock.""" import errno +import select +import selectors import socket from websockets.sync.client import ClientConnection @@ -22,3 +24,34 @@ def abort_connection(ws: ClientConnection) -> None: finally: # Even if shutdown fails, prevent any later send on this socket object. ws.socket.close() + + +def socket_writable(ws: ClientConnection) -> bool: + """Whether a small write to the transport's socket won't block right now. + + For a send that must not stall the reader thread: a full send buffer + (a peer that keeps sending but stops reading) would park sendall(), which + also holds the library's protocol mutex, so even keepalive pings stall. + + Never raises. Returns ``True`` (proceed, as without this check) on any + failure or for a transport without a real socket, such as a test fake: + ``False`` would silently disable the caller's sends. Uses ``poll`` (or a + selector), not ``select.select``, which fails for fds above FD_SETSIZE. + """ + try: + sock = getattr(ws, "socket", None) + if not isinstance(sock, socket.socket): + return True + fd = sock.fileno() + if fd < 0: + return True # Closed: let the send itself report it. + if hasattr(select, "poll"): + poller = select.poll() + poller.register(fd, select.POLLOUT) + # POLLERR/POLLHUP count as writable: the send then fails fast. + return bool(poller.poll(0)) + with selectors.DefaultSelector() as selector: + selector.register(fd, selectors.EVENT_WRITE) + return bool(selector.select(0)) + except Exception: + return True diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index 712c7d0..5d53298 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -17,7 +17,7 @@ import websockets.sync.client from .constants import DEFAULT_READ_TIMEOUT_SECONDS, DEFAULT_STALE_QUERY_PROBE_SECONDS -from ._transport import abort_connection +from ._transport import abort_connection, socket_writable from .cursor import Cursor from .errors import NotSupportedError, OperationalError from .models import ExecutionResult, ProgressInfo, Store, StoreResult @@ -387,9 +387,10 @@ def __probe_stale_queries(self) -> None: if silence < wait: continue if not self.__request_results(query.execution_id, probe=True): - # A user thread is sending, possibly stalled. Don't wait - # behind it: that would stop recv() for every query. The - # lock is busy for the rest too; retry at the next check. + # A user thread is sending, possibly stalled, or the + # socket can't take a write. Don't wait: that would stop + # recv() for every query. Either holds for the rest too; + # retry at the next check. break # Only a reply to a retrieve clears the flag; other events, # state_updated included, don't. Still set means the last @@ -738,7 +739,11 @@ def __send( query: Query | None = None, blocking: bool = True, ) -> bool: - """Send a request; return ``False`` only if non-blocking and busy.""" + """Send a request; return ``False`` only if non-blocking and busy. + + Busy means another sender holds the send lock, or the socket's send + buffer is full. + """ # Serialization and redaction are local work. Fail before registration, # without poisoning unrelated cursors or misreporting a transport loss. request = json.dumps(message) @@ -762,6 +767,15 @@ def __send( self.__queries[query.execution_id] = query elif message.get("execution_id") not in self.__queries: return True + # A non-blocking send runs on the reader thread (a probe). Not + # waiting for the lock isn't enough: with the socket's send buffer + # full (a peer that keeps sending but stops reading), ws.send() + # itself would block, stopping recv() for every query, and would + # hold the library's protocol mutex, which keepalive pings need + # too. Treat it like a busy lock: skip, and retry at the next + # check. Checked outside __lock: it's a syscall on the transport. + if not blocking and not socket_writable(self.__ws): + return False try: self.__ws.send(request) except (websockets.exceptions.ConnectionClosed, OSError): @@ -852,7 +866,8 @@ def __execute_sql( def __request_results(self, execution_id: str, probe: bool = False) -> bool: """Send retrieve_results; ``False`` if a probe found the sender busy. - A probe never waits for ``__send_lock`` (see ``__probe_stale_queries``). + A probe never waits for ``__send_lock`` or a full socket (see + ``__probe_stale_queries`` and ``__send``). """ query = self.__queries.get(execution_id) if not query: