ADD: Actively monitor for hung Live sessions

Author: nmacholl

CHANGELOG.md

@@ -5,6 +5,8 @@
 #### Enhancements
 - Added a property `Live.subscription_requests` which returns a list of tuples containing every `SubscriptionRequest` for the live session
 - Changed the return value of `Live.subscribe()` to `int`, the value of the subscription ID, which can be used to index into the `Live.subscription_requests` property
+- Added feature to automatically monitor for hung connections in the `Live` client
+  - Hung connections will be disconnected client side with a `BentoError`
 
 #### Breaking changes
 - Several log messages have been reformatted to improve clarity and reduce redundancy, especially at debug levels

databento/live/session.py

@@ -4,6 +4,7 @@
 import dataclasses
 import itertools
 import logging
+import math
 import queue
 import struct
 import threading
@@ -35,6 +36,7 @@
 CONNECT_TIMEOUT_SECONDS: Final = 10.0
 DBN_QUEUE_CAPACITY: Final = 2**20
 DEFAULT_REMOTE_PORT: Final = 13000
+CLIENT_TIMEOUT_MARGIN_SECONDS: Final = 10
 
 
 class DBNQueue(queue.SimpleQueue):  # type: ignore [type-arg]
@@ -212,6 +214,7 @@ def __init__(
         self._user_callbacks = user_callbacks
         self._user_streams = user_streams
         self._last_ts_event: int | None = None
+        self._last_msg_loop_time: float = math.inf
 
     def received_metadata(self, metadata: databento_dbn.Metadata) -> None:
         if self._metadata:
@@ -235,6 +238,7 @@ def received_record(self, record: DBNRecord) -> None:
         if self._dbn_queue.is_enabled():
             self._queue_for_iteration(record)
         self._last_ts_event = record.ts_event
+        self._last_msg_loop_time = self._loop.time()
 
         return super().received_record(record)
 
@@ -324,7 +328,7 @@ def __init__(
 
         self._api_key = api_key
         self._ts_out = ts_out
-        self._heartbeat_interval_s = heartbeat_interval_s
+        self._heartbeat_interval_s = heartbeat_interval_s or 30
 
         self._protocol: _SessionProtocol | None = None
         self._transport: asyncio.Transport | None = None
@@ -333,6 +337,7 @@ def __init__(
         self._subscriptions: list[tuple[SubscriptionRequest, ...]] = []
         self._reconnect_policy = ReconnectPolicy(reconnect_policy)
         self._reconnect_task: asyncio.Task[None] | None = None
+        self._heartbeat_monitor_task: asyncio.Task[None] | None = None
 
         self._dataset = ""
 
@@ -436,8 +441,6 @@ def stop(self) -> None:
         with self._lock:
             if self._transport is None:
                 return
-            if self._protocol is not None:
-                self._protocol.disconnected.add_done_callback(lambda _: self._cleanup())
             self._loop.call_soon_threadsafe(self._transport.close)
 
     def start(self) -> None:
@@ -454,6 +457,9 @@ def start(self) -> None:
             if self._protocol is None:
                 raise ValueError("session is not connected")
             self._protocol.start()
+            self._heartbeat_monitor_task = self._loop.create_task(
+                self._heartbeat_monitor(),
+            )
 
     def subscribe(
         self,
@@ -553,6 +559,8 @@ def _cleanup(self) -> None:
             if stream.is_managed:
                 stream.close()
 
+        if self._heartbeat_monitor_task is not None:
+            self._heartbeat_monitor_task.cancel()
         self._user_callbacks.clear()
         self._user_streams.clear()
         self._user_reconnect_callbacks.clear()
@@ -647,6 +655,21 @@ async def _connect_task(
 
         return transport, protocol
 
+    async def _heartbeat_monitor(self) -> None:
+        while not self._protocol.disconnected.done():
+            await asyncio.sleep(1)
+            gap = self._loop.time() - self._protocol._last_msg_loop_time
+            if gap > (self._heartbeat_interval_s + CLIENT_TIMEOUT_MARGIN_SECONDS):
+                logger.error(
+                    "disconnecting client due to timeout, no data received for %d second(s)",
+                    int(gap),
+                )
+                self._protocol.disconnected.set_exception(
+                    BentoError(
+                        f"Gateway timeout: {gap:.0f} second(s) since last message",
+                    ),
+                )
+
     async def _reconnect(self) -> None:
         while True:
             try:

tests/test_live_client.py

@@ -4,7 +4,6 @@
 
 from __future__ import annotations
 
-import asyncio
 import pathlib
 import platform
 import random
@@ -279,7 +278,7 @@ async def test_live_client_reuse(
     live_client.stop()
     assert live_client.session_id == first_session_id
 
-    await asyncio.sleep(1)
+    await live_client.wait_for_close()
 
     live_client.subscribe(
         dataset=Dataset.GLBX_MDP3,

tests/test_live_client_reconnect.py

@@ -65,9 +65,9 @@ async def test_reconnect_before_start(
     reconnect_policy: ReconnectPolicy = ReconnectPolicy.RECONNECT,
 ) -> None:
     """
-    Test that a reconnect policy of "reconnect_do_not_replay" reconnects a
-    client but does not send the session start command if the session was not
-    streaming previously.
+    Test that a reconnect policy of "reconnect" reconnects a client but does
+    not send the session start command if the session was not streaming
+    previously.
     """
     # Arrange
     live_client = client.Live(
@@ -138,8 +138,8 @@ async def test_reconnect_subscriptions(
     reconnect_policy: ReconnectPolicy = ReconnectPolicy.RECONNECT,
 ) -> None:
     """
-    Test that a reconnect policy of "reconnect_do_not_replay" re-sends the
-    subscription requests with a start of `None`.
+    Test that a reconnect policy of "reconnect" re-sends the subscription
+    requests with a start of `None`.
     """
     # Arrange
     live_client = client.Live(
@@ -192,8 +192,8 @@ async def test_reconnect_callback(
     reconnect_policy: ReconnectPolicy = ReconnectPolicy.RECONNECT,
 ) -> None:
     """
-    Test that a reconnect policy of "reconnect_do_not_replay" will cause a user
-    supplied reconnection callback to be executed when a reconnection occurs.
+    Test that a reconnect policy of "reconnect" will cause a user supplied
+    reconnection callback to be executed when a reconnection occurs.
     """
     # Arrange
     live_client = client.Live(