FIX: Fix subscribe in Live callback

Author: nmacholl

CHANGELOG.md

@@ -5,6 +5,9 @@
 #### Enhancements
 - Increase `Live` session connection & authentication timeouts
 
+#### Bug fixes
+- Fixed an issue where calling `Live.subscribe` from a `Live` client callback would cause a deadlock
+
 ## 0.31.0 - 2024-03-05
 
 #### Enhancements

databento/live/session.py

@@ -434,15 +434,12 @@ def subscribe(
                     loop=self._loop,
                 )
 
-        asyncio.run_coroutine_threadsafe(
-            self._subscribe_task(
+            self._protocol.subscribe(
                 schema=schema,
                 symbols=symbols,
                 stype_in=stype_in,
                 start=start,
-            ),
-            loop=self._loop,
-        ).result()
+            )
 
     def resume_reading(self) -> None:
         """
@@ -565,21 +562,3 @@ async def _connect_task(
         )
 
         return transport, protocol
-
-    async def _subscribe_task(
-        self,
-        schema: Schema | str,
-        symbols: Iterable[str | int] | str | int = ALL_SYMBOLS,
-        stype_in: SType | str = SType.RAW_SYMBOL,
-        start: str | int | None = None,
-    ) -> None:
-        with self._lock:
-            if self._protocol is None:
-                return
-
-            self._protocol.subscribe(
-                schema=schema,
-                symbols=symbols,
-                stype_in=stype_in,
-                start=start,
-            )

tests/mock_live_server.py

@@ -16,6 +16,7 @@
 from functools import singledispatchmethod
 from io import BytesIO
 from os import PathLike
+from typing import Any
 from typing import Callable
 from typing import NewType
 from typing import TypeVar
@@ -32,9 +33,7 @@
 from databento.live.gateway import SessionStart
 from databento.live.gateway import SubscriptionRequest
 from databento.live.gateway import parse_gateway_message
-from databento_dbn import Metadata
 from databento_dbn import Schema
-from databento_dbn import SType
 
 
 LIVE_SERVER_VERSION: str = "1.0.0"
@@ -100,7 +99,8 @@ def __init__(
         self._version: str = version
         self._is_authenticated: bool = False
         self._is_streaming: bool = False
-        self._repeater_tasks: set[asyncio.Task[None]] = set()
+        self._subscriptions: list[SubscriptionRequest] = []
+        self._tasks: set[asyncio.Task[None]] = set()
 
         self._dbn_path = pathlib.Path(dbn_path)
         self._user_api_keys = user_api_keys
@@ -155,6 +155,18 @@ def is_streaming(self) -> bool:
         """
         return self._is_streaming
 
+    @property
+    def dataset_path(self) -> pathlib.Path:
+        """
+        The path to the DBN files for serving.
+
+        Returns
+        -------
+        Path
+
+        """
+        return self._dbn_path / (self._dataset or "")
+
     @property
     def mode(self) -> MockLiveMode:
         """
@@ -205,6 +217,18 @@ def session_id(self) -> str:
         """
         return str(hash(self))
 
+    @property
+    def subscriptions(self) -> tuple[SubscriptionRequest, ...]:
+        """
+        The received subscriptions.
+
+        Returns
+        -------
+        tuple[SubscriptionRequest, ...]
+
+        """
+        return tuple(self._subscriptions)
+
     @property
     def version(self) -> str:
         """
@@ -353,11 +377,11 @@ def _(self, message: AuthenticationRequest) -> None:
         logger.info("received CRAM response: %s", message.auth)
         if self.is_authenticated:
             logger.error("authentication request sent when already authenticated")
-            self.__transport.close()
+            self.__transport.write_eof()
             return
         if self.is_streaming:
             logger.error("authentication request sent while streaming")
-            self.__transport.close()
+            self.__transport.write_eof()
             return
 
         _, bucket_id = message.auth.split("-")
@@ -400,54 +424,49 @@ def _(self, message: SubscriptionRequest) -> None:
         logger.info("received subscription request: %s", str(message).strip())
         if not self.is_authenticated:
             logger.error("subscription request sent while unauthenticated")
-            self.__transport.close()
+            self.__transport.write_eof()
 
-        if self.is_streaming:
-            logger.error("subscription request sent while streaming")
-            self.__transport.close()
+        self._subscriptions.append(message)
 
-        self._schemas.append(Schema(message.schema))
+        if self.is_streaming:
+            self.create_server_task(message)
 
     @handle_client_message.register(SessionStart)
     def _(self, message: SessionStart) -> None:
         logger.info("received session start request: %s", str(message).strip())
         self._is_streaming = True
 
-        dataset_path = self._dbn_path / (self._dataset or "")
+        for sub in self.subscriptions:
+            self.create_server_task(sub)
+
+    def create_server_task(self, message: SubscriptionRequest) -> None:
         if self.mode is MockLiveMode.REPLAY:
-            for schema in self._schemas:
-                for test_data_path in dataset_path.glob(f"*{schema}.dbn.zst"):
-                    decompressor = zstandard.ZstdDecompressor().stream_reader(
-                        test_data_path.read_bytes(),
-                    )
-                    logger.info(
-                        "streaming %s for %s schema",
-                        test_data_path.name,
-                        schema,
-                    )
-                    self.__transport.write(decompressor.readall())
+            task = asyncio.create_task(self.replay_task(schema=Schema(message.schema)))
+        else:
+            task = asyncio.create_task(self.repeater_task(schema=Schema(message.schema)))
 
-            logger.info(
-                "data streaming for %d schema(s) completed",
-                len(self._schemas),
-            )
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.remove)
+        task.add_done_callback(self.check_done)
 
+    def check_done(self, _: Any) -> None:
+        if not self._tasks:
+            logger.info("streaming tasks completed")
             self.__transport.write_eof()
-            self.__transport.close()
 
-        elif self.mode is MockLiveMode.REPEAT:
-            metadata = Metadata("UNIT.TEST", 0, SType.RAW_SYMBOL, [], [], [], [])
-            self.__transport.write(bytes(metadata))
-
-            loop = asyncio.get_event_loop()
-            for schema in self._schemas:
-                task = loop.create_task(self.repeater(schema))
-                self._repeater_tasks.add(task)
-                task.add_done_callback(self._repeater_tasks.remove)
-        else:
-            raise ValueError(f"unsupported mode {MockLiveMode.REPEAT}")
+    async def replay_task(self, schema: Schema) -> None:
+        for test_data_path in self.dataset_path.glob(f"*{schema}.dbn.zst"):
+            decompressor = zstandard.ZstdDecompressor().stream_reader(
+                test_data_path.read_bytes(),
+            )
+            logger.info(
+                "streaming %s for %s schema",
+                test_data_path.name,
+                schema,
+            )
+            self.__transport.write(decompressor.readall())
 
-    async def repeater(self, schema: Schema) -> None:
+    async def repeater_task(self, schema: Schema) -> None:
         struct = SCHEMA_STRUCT_MAP[schema]
         repeated = bytes(struct(*[0] * 12))  # for now we only support MBP_1
 

tests/test_live_client.py

@@ -503,6 +503,51 @@ async def test_live_subscribe_large_symbol_list(
     assert reconstructed == large_symbol_list
 
 
+async def test_live_subscribe_from_callback(
+    live_client: client.Live,
+    mock_live_server: MockLiveServer,
+) -> None:
+    """
+    Test that `Live.subscribe` can be called from a callback.
+    """
+    # Arrange
+    live_client.subscribe(
+        dataset=Dataset.GLBX_MDP3,
+        schema=Schema.OHLCV_1H,
+        stype_in=SType.RAW_SYMBOL,
+        symbols="TEST0",
+    )
+
+    def cb_sub(_: DBNRecord) -> None:
+        live_client.subscribe(
+            dataset=Dataset.GLBX_MDP3,
+            schema=Schema.MBO,
+            stype_in=SType.RAW_SYMBOL,
+            symbols="TEST1",
+        )
+
+    live_client.add_callback(cb_sub)
+
+    # Act
+    first_sub = mock_live_server.get_message_of_type(
+        gateway.SubscriptionRequest,
+        timeout=1,
+    )
+
+    live_client.start()
+
+    second_sub = mock_live_server.get_message_of_type(
+        gateway.SubscriptionRequest,
+        timeout=1,
+    )
+
+    await live_client.wait_for_close()
+
+    # Assert
+    assert first_sub.symbols == "TEST0"
+    assert second_sub.symbols == "TEST1"
+
+
 @pytest.mark.usefixtures("mock_live_server")
 def test_live_stop(
     live_client: client.Live,