From 0e2ec7f8b1c766a7a1ba53c4f0c0b478de91570c Mon Sep 17 00:00:00 2001 From: Ryan Avery Date: Mon, 21 Sep 2026 13:37:25 -0700 Subject: [PATCH] feat(dbapi): log bounded session diagnostics after connection loss --- tests/test_session_diagnostics.py | 278 ++++++++++++++++++++++++++++++ wherobots/db/_diagnostics.py | 125 ++++++++++++++ wherobots/db/connection.py | 13 +- wherobots/db/driver.py | 20 +++ 4 files changed, 434 insertions(+), 2 deletions(-) create mode 100644 tests/test_session_diagnostics.py create mode 100644 wherobots/db/_diagnostics.py diff --git a/tests/test_session_diagnostics.py b/tests/test_session_diagnostics.py new file mode 100644 index 0000000..041a419 --- /dev/null +++ b/tests/test_session_diagnostics.py @@ -0,0 +1,278 @@ +"""Optional diagnostics are bounded and never delay a cursor's terminal result.""" +import json +import threading +from unittest.mock import MagicMock, patch + +import pytest +from websockets.exceptions import ConnectionClosedError + +from wherobots.db import _diagnostics +from wherobots.db._diagnostics import SessionDiagnostics +from wherobots.db.connection import Connection +from wherobots.db.driver import connect_direct +from wherobots.db.errors import OperationalError +from test_disconnect import Transport + + +def response(payload, status=200): + result = MagicMock(status_code=status) + result.__enter__.return_value = result + result.iter_content.return_value = [json.dumps(payload).encode()] + return result + + +@pytest.fixture +def workers(monkeypatch): + # Keep real concurrency; remove only jitter and record workers for cleanup. + threads = [] + thread_type = threading.Thread + + def thread(*args, **kwargs): + result = thread_type(*args, **kwargs) + threads.append(result) + return result + + monkeypatch.setattr(_diagnostics.threading, "Thread", thread) + monkeypatch.setattr(_diagnostics.time, "sleep", lambda _: None) + yield threads + for worker in threads: + worker.join(timeout=3) + assert not worker.is_alive() + + +@pytest.mark.parametrize( + "payload,expected", + [ + ( + {"status": "FAILED", "firstFailure": {"message": "Evicted"}}, + ("FAILED", "Evicted"), + ), + ({"status": "READY"}, ("READY", None)), + ({"firstFailure": {"message": "Evicted"}}, (None, "Evicted")), + ({}, None), + (None, None), + ([], None), + ({"status": 1, "firstFailure": {"message": 5}}, None), + ({"firstFailure": "invalid"}, None), + ], +) +def test_status_and_failure_parsing(payload, expected): + with patch.object( + _diagnostics.requests, "get", return_value=response(payload) + ) as get: + assert ( + SessionDiagnostics._fetch( + "https://api/session/1", {"Authorization": "Bearer test"} + ) + == expected + ) + assert get.call_args.kwargs == { + "headers": {"Authorization": "Bearer test"}, + "timeout": 1.0, + "allow_redirects": False, + "stream": True, + } + + +def test_message_and_status_lengths_are_bounded(): + payload = {"status": "s" * 200, "firstFailure": {"message": "m" * 5000}} + with patch.object(_diagnostics.requests, "get", return_value=response(payload)): + status, message = SessionDiagnostics._fetch("https://api/session/1", {}) + assert status == "s" * 128 + assert message == "m" * 4096 + + +def test_streaming_response_size_is_bounded(): + result = response({}) + chunks_consumed = [] + + def chunks(chunk_size): + for i in range(1000): + chunks_consumed.append(i) + yield b"x" * chunk_size + + result.iter_content.side_effect = chunks + with patch.object(_diagnostics.requests, "get", return_value=result): + assert SessionDiagnostics._fetch("https://api/session/1", {}) is None + assert len(chunks_consumed) == 9 + result.__exit__.assert_called_once() + + +@pytest.mark.parametrize("status", [302, 404, 503]) +def test_http_errors_are_logged_without_parsing(status, caplog): + result = response({}, status) + with caplog.at_level("DEBUG"), patch.object( + _diagnostics.requests, "get", return_value=result + ): + assert SessionDiagnostics._fetch("https://api/session/1", {}) is None + result.iter_content.assert_not_called() + assert f"HTTP status: {status}" in caplog.text + + +def test_inflight_and_recent_requests_are_shared(workers, caplog): + service = SessionDiagnostics() + entered = threading.Event() + release = threading.Event() + + def fetch(*args): + entered.set() + assert release.wait(timeout=3) + return "FAILED", "Evicted" + + with caplog.at_level("WARNING"), patch.object( + service, "_fetch", side_effect=fetch + ) as get: + try: + service.request("https://api/session/1", {"Authorization": "secret"}, "1") + assert entered.wait(timeout=1) + for _ in range(50): + service.request( + "https://api/session/1", {"Authorization": "secret"}, "1" + ) + assert len(workers) == 1 + release.set() + workers[0].join(timeout=2) + service.request("https://api/session/1", {"Authorization": "secret"}, "1") + get.assert_called_once() + finally: + release.set() + assert "status='FAILED'" in caplog.text + assert "firstFailure='Evicted'" in caplog.text + assert "secret" not in caplog.text + + +def test_authentication_and_endpoint_contexts_are_not_shared(workers): + service = SessionDiagnostics() + with patch.object(service, "_fetch", return_value=None) as get: + for url, token in [ + ("https://one/session/1", "a"), + ("https://one/session/1", "b"), + ("https://two/session/1", "a"), + ]: + service.request(url, {"Authorization": token}, "1") + workers[-1].join(timeout=1) + assert get.call_count == 3 + assert all( + isinstance(key[1], bytes) and len(key[1]) == 32 for key in service._entries + ) + + +def test_worker_budget_bounds_stalled_lookups(workers): + service = SessionDiagnostics() + release = threading.Event() + + def fetch(*args): + assert release.wait(timeout=3) + return None + + with patch.object(service, "_fetch", side_effect=fetch): + try: + for i in range(100): + service.request(f"https://api/session/{i}", {}, str(i)) + assert len(workers) == service._active == 4 + assert len(service._entries) == 4 + finally: + release.set() + for worker in workers: + worker.join(timeout=2) + assert service._active == 0 + + +def test_cache_expires_and_has_a_fixed_size(workers, monkeypatch): + service = SessionDiagnostics() + monkeypatch.setattr(_diagnostics, "_MAX_ENTRIES", 2) + with patch.object(service, "_fetch", return_value=None) as get: + for i in range(3): + service.request(f"https://api/session/{i}", {}, str(i)) + workers[-1].join(timeout=1) + assert len(service._entries) == 2 + assert get.call_count == 3 + for key in service._entries: + service._entries[key] = 0.0 + service.request("https://api/session/2", {}, "2") + workers[-1].join(timeout=1) + assert get.call_count == 4 + + +def test_worker_exception_is_debug_logged_without_credentials(workers, caplog): + service = SessionDiagnostics() + with caplog.at_level("DEBUG"), patch.object( + service, "_fetch", side_effect=ValueError("secret token") + ): + service.request("https://api/session/1", {}, "1") + workers[0].join(timeout=1) + assert "lookup failed (ValueError)" in caplog.text + assert "secret token" not in caplog.text + assert service._active == 0 + + +def test_thread_start_failure_is_debug_logged_and_releases_slot(caplog): + service = SessionDiagnostics() + with caplog.at_level("DEBUG"), patch.object( + threading.Thread, "start", side_effect=RuntimeError + ): + service.request("https://api/session/1", {}, "1") + assert "Could not start session diagnostic thread" in caplog.text + assert service._active == 0 + + +@pytest.mark.parametrize("failure", ["send", "receive"]) +def test_stalled_lookup_does_not_delay_failure_delivery(failure, workers): + service = SessionDiagnostics() + entered = threading.Event() + release = threading.Event() + + def fetch(*args): + entered.set() + assert release.wait(timeout=3) + return None + + ws = Transport() + conn = Connection( + ws, on_connection_lost=lambda: service.request("https://api/session/1", {}, "1") + ) + with patch.object(service, "_fetch", side_effect=fetch): + try: + if failure == "send": + ws.send = MagicMock(side_effect=ConnectionClosedError(None, None)) + cursor = conn.cursor() + cursor.execute("SELECT 1") + if failure == "receive": + ws.incoming.put(ConnectionClosedError(None, None)) + assert isinstance( + cursor._Cursor__queue.get(timeout=1).error, OperationalError + ) + assert entered.wait(timeout=1) + assert not release.is_set() # Error preceded lookup completion. + finally: + release.set() + conn.close() + + +def test_explicit_close_does_not_request_diagnostics(): + callback = MagicMock() + conn = Connection(Transport(), on_connection_lost=callback) + cursor = conn.cursor() + cursor.execute("SELECT 1") + conn.close() + callback.assert_not_called() + + +@pytest.mark.parametrize("suffix", ["", "/", "/?ignored=1"]) +def test_direct_status_url_derivation_and_header_snapshot(suffix): + headers = {"Authorization": "original"} + url = "https://api/session/session-1" + suffix + with patch( + "wherobots.db.driver.websockets.sync.client.connect", return_value=Transport() + ), patch.object(_diagnostics.session_diagnostics, "request") as request: + conn = connect_direct( + "wss://compute/sql", session_status_url=url, headers=headers + ) + headers["Authorization"] = "changed" + cursor = conn.cursor() + cursor.execute("SELECT 1") + conn._Connection__ws.incoming.put(ConnectionClosedError(None, None)) + conn._Connection__thread.join(timeout=2) + request.assert_called_once_with(url, {"Authorization": "original"}, "session-1") + with pytest.raises(OperationalError, match="session=session-1"): + cursor.fetchall() diff --git a/wherobots/db/_diagnostics.py b/wherobots/db/_diagnostics.py new file mode 100644 index 0000000..28ba023 --- /dev/null +++ b/wherobots/db/_diagnostics.py @@ -0,0 +1,125 @@ +"""Bounded, best-effort session diagnostics after cursor failure delivery.""" + +import hashlib +import json +import logging +import random +import threading +import time +from collections import OrderedDict + +import requests + + +_logger = logging.getLogger(__name__) +_MAX_WORKERS = 4 +_MAX_ENTRIES = 128 +_CACHE_SECONDS = 30.0 +_MAX_RESPONSE_BYTES = 65536 + + +class SessionDiagnostics: + """Share in-flight and recent lookups by URL and authentication context. + + A hung DNS/TLS/HTTP operation retains its slot, so even failures that ignore + socket timeouts cannot cause unbounded threads. At capacity we skip optional + diagnostics. There are no retries, and no caller waits for a lookup. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._entries: OrderedDict[tuple[str, bytes], float | None] = OrderedDict() + self._active = 0 + + def request( + self, url: str, headers: dict[str, str] | None, session_id: str + ) -> None: + headers = dict(headers or {}) + # Distinct credentials must never share diagnostic responses. Keep only + # a digest in the short-lived registry, never credentials in log output. + context = hashlib.sha256(json.dumps(sorted(headers.items())).encode()).digest() + key = (url, context) + now = time.monotonic() + with self._lock: + for expired, deadline in list(self._entries.items()): + if deadline is not None and deadline <= now: + del self._entries[expired] + if key in self._entries: + return + if self._active >= _MAX_WORKERS: + _logger.debug("Session diagnostics at capacity; skipping lookup") + return + if len(self._entries) >= _MAX_ENTRIES: + # Never evict a running request: it must remain deduplicated. + oldest = next( + k for k, deadline in self._entries.items() if deadline is not None + ) + del self._entries[oldest] + self._entries[key] = None + self._active += 1 + + def lookup() -> None: + try: + # Spread correlated failures across processes without delaying + # execute(), fetch(), or connection shutdown. + time.sleep(random.uniform(0.0, 0.5)) + details = self._fetch(url, headers) + if details is not None: + status, message = details + _logger.warning( + "SQL session diagnostic (session=%s): status=%r; firstFailure=%r", + session_id, + status, + message, + ) + except Exception as exc: + # Request exceptions may contain URLs/credentials. Log the type, + # not arbitrary exception text or the authenticated request. + _logger.debug( + "Session diagnostic lookup failed (%s)", type(exc).__name__ + ) + finally: + with self._lock: + self._active -= 1 + self._entries[key] = time.monotonic() + _CACHE_SECONDS + + try: + threading.Thread( + target=lookup, daemon=True, name="wherobots-session-diagnostic" + ).start() + except RuntimeError: + with self._lock: + self._active -= 1 + self._entries[key] = time.monotonic() + _CACHE_SECONDS + _logger.debug("Could not start session diagnostic thread") + + @staticmethod + def _fetch( + url: str, headers: dict[str, str] + ) -> tuple[str | None, str | None] | None: + with requests.get( + url, headers=headers, timeout=1.0, allow_redirects=False, stream=True + ) as response: + if response.status_code != 200: + _logger.debug( + "Session diagnostic HTTP status: %s", response.status_code + ) + return None + data = bytearray() + for chunk in response.iter_content(chunk_size=8192): + if len(data) + len(chunk) > _MAX_RESPONSE_BYTES: + _logger.debug("Session diagnostic response exceeds byte limit") + return None + data.extend(chunk) + payload = json.loads(data) + if not isinstance(payload, dict): + return None + status = payload.get("status") + status = status[:128] if isinstance(status, str) else None + failure = payload.get("firstFailure") + message = failure.get("message") if isinstance(failure, dict) else None + message = message[:4096] if isinstance(message, str) else None + return (status, message) if status is not None or message is not None else None + + +session_diagnostics = SessionDiagnostics() diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py index e4fb00c..2e3cc20 100644 --- a/wherobots/db/connection.py +++ b/wherobots/db/connection.py @@ -70,6 +70,7 @@ def __init__( data_compression: DataCompression | None = None, geometry_representation: GeometryRepresentation | None = None, session_id: str | None = None, + on_connection_lost: Callable[[], None] | None = None, ): self.__ws = ws self.__read_timeout = read_timeout @@ -79,6 +80,9 @@ def __init__( self.__progress_handler: ProgressHandler | None = None self.__session_id = session_id + # Internal notification hook: schedule optional diagnostics, never wait + # for remote results here. All cursor failures are delivered first. + self.__on_connection_lost = on_connection_lost self.__lock = threading.Lock() self.__send_lock = threading.Lock() self.__shutdown_done = threading.Event() @@ -103,7 +107,7 @@ def close(self) -> None: decoder or callback can outlive the bounded reader join. """ deadline = time.monotonic() + 1.0 - self.__fail_pending() + self.__fail_pending(notify=False) # A handler may close its own connection during terminal delivery. if self.__shutdown_owner == threading.get_ident(): return @@ -170,7 +174,7 @@ def __connection_error(self, execution_id: str) -> OperationalError: ) return OperationalError(message) - def __fail_pending(self) -> None: + def __fail_pending(self, notify: bool = True) -> None: # Stop admission first. Do not wait for __send_lock: its owner may be # blocked in network I/O. __closed means closing until shutdown_done. with self.__lock: @@ -201,6 +205,11 @@ def __fail_pending(self) -> None: finally: self.__shutdown_owner = None self.__shutdown_done.set() + if notify and pending and self.__on_connection_lost is not None: + try: + self.__on_connection_lost() + except Exception: + logging.debug("Could not schedule session diagnostics") def __listen(self) -> None: """Waits for the next message from the SQL session and processes it. diff --git a/wherobots/db/driver.py b/wherobots/db/driver.py index 63d1b24..3ef119d 100644 --- a/wherobots/db/driver.py +++ b/wherobots/db/driver.py @@ -267,6 +267,7 @@ def get_session_uri() -> str: data_compression=data_compression, geometry_representation=geometry_representation, cancel_event=cancel_event, + session_status_url=session_id_url, session_id=urllib.parse.urlparse(session_id_url) .path.rstrip("/") .rsplit("/", 1)[-1], @@ -298,6 +299,7 @@ def connect_direct( geometry_representation: Union[GeometryRepresentation, None] = None, cancel_event: Union[threading.Event, None] = None, session_id: str | None = None, + session_status_url: str | None = None, ) -> Connection: uri_with_protocol = f"{uri}/{protocol}" ssl_context = ssl.create_default_context() @@ -335,6 +337,23 @@ def ws_connect() -> websockets.sync.client.ClientConnection: except Exception as e: raise InterfaceError("Failed to connect to SQL session!") from e + on_connection_lost = None + if session_status_url is not None: + from ._diagnostics import session_diagnostics + + if session_id is None: + session_id = ( + urllib.parse.urlparse(session_status_url) + .path.rstrip("/") + .rsplit("/", 1)[-1] + ) + diagnostic_headers = dict(headers or {}) + + def on_connection_lost() -> None: + session_diagnostics.request( + session_status_url, diagnostic_headers, session_id or "unknown" + ) + return Connection( ws, read_timeout=read_timeout, @@ -342,4 +361,5 @@ def ws_connect() -> websockets.sync.client.ClientConnection: data_compression=data_compression, geometry_representation=geometry_representation, session_id=session_id, + on_connection_lost=on_connection_lost, )