MOD: Live client iteration queue

Author: nmacholl

CHANGELOG.md

@@ -2,8 +2,23 @@
 
 ## 0.28.0 - TBD
 
-#### Breaking change
+#### Enhancements
+- Substantially increased iteration queue size
+- Added methods `DBNQueue.enable` and `DBNQueue.disable` for controlling queue consumption
+- Added method `DBNQueue.is_enabled` to signal the queue can accept records
+- Added method `DBNQueue.is_full` to signal the queue has reached capacity
+- Added enabled checks to `DBNQueue.put` and `DBNQueue.put_nowait`
+
+#### Breaking changes
 - Iterating a `Live` client after the streaming session has started will now raise a `ValueError`. Calling `Live.start` is not necessary when iterating the `Live` client
+- Moved constant `databento.live.client.DEFAULT_QUEUE_SIZE` to `databento.live.session.DBN_QUEUE_CAPACITY`
+- Removed `maxsize` parameter from `DBNQueue` constructor, `None` is now used instead
+- Removed property `DBNQueue.enabled`, use `DBNQueue.is_enabled` instead
+- Removed method `DBNQueue.is_half_full`, use `DBNQueue.is_full` instead
+
+#### Bug fixes
+- Fixed an issue where DBN records could be dropped while iterating
+- Fixed an issue where async iteration would block the event loop
 
 ## 0.27.0 - 2024-01-23
 

databento/live/client.py

@@ -9,7 +9,7 @@
 from collections.abc import Iterable
 from concurrent import futures
 from os import PathLike
-from typing import IO, Final
+from typing import IO
 
 import databento_dbn
 from databento_dbn import Schema
@@ -33,7 +33,6 @@
 
 
 logger = logging.getLogger(__name__)
-DEFAULT_QUEUE_SIZE: Final = 2048
 
 
 class Live:
@@ -85,7 +84,7 @@ def __init__(
         self._dataset: Dataset | str = ""
         self._ts_out = ts_out
 
-        self._dbn_queue: DBNQueue = DBNQueue(maxsize=DEFAULT_QUEUE_SIZE)
+        self._dbn_queue: DBNQueue = DBNQueue()
         self._metadata: SessionMetadata = SessionMetadata()
         self._symbology_map: dict[int, str | int] = {}
         self._user_callbacks: dict[RecordCallback, ExceptionCallback | None] = {
@@ -120,10 +119,37 @@ def __aiter__(self) -> Live:
         return iter(self)
 
     async def __anext__(self) -> DBNRecord:
+        if not self._dbn_queue.is_enabled():
+            raise ValueError("iteration has not started")
+
+        loop = asyncio.get_running_loop()
+
         try:
-            return next(self)
-        except StopIteration:
-            raise StopAsyncIteration
+            return self._dbn_queue.get_nowait()
+        except queue.Empty:
+            while True:
+                try:
+                    return await loop.run_in_executor(
+                        None,
+                        self._dbn_queue.get,
+                        True,
+                        0.1,
+                    )
+                except queue.Empty:
+                    if self._session.is_disconnected():
+                        break
+        finally:
+            if not self._dbn_queue.is_full() and not self._session.is_reading():
+                logger.debug(
+                    "resuming reading with %d pending records",
+                    self._dbn_queue.qsize(),
+                )
+                self._session.resume_reading()
+
+        self._dbn_queue.disable()
+        await self.wait_for_close()
+        logger.debug("completed async iteration")
+        raise StopAsyncIteration
 
     def __iter__(self) -> Live:
         logger.debug("starting iteration")
@@ -138,30 +164,28 @@ def __iter__(self) -> Live:
         return self
 
     def __next__(self) -> DBNRecord:
-        if self._dbn_queue is None:
+        if not self._dbn_queue.is_enabled():
             raise ValueError("iteration has not started")
 
-        while not self._session.is_disconnected() or self._dbn_queue.qsize() > 0:
+        while True:
             try:
-                record = self._dbn_queue.get(block=False)
+                record = self._dbn_queue.get(timeout=0.1)
             except queue.Empty:
-                continue
+                if self._session.is_disconnected():
+                    break
             else:
-                logger.debug(
-                    "yielding %s record from next",
-                    type(record).__name__,
-                )
                 return record
             finally:
-                if not self._dbn_queue.half_full() and not self._session.is_reading():
+                if not self._dbn_queue.is_full() and not self._session.is_reading():
                     logger.debug(
                         "resuming reading with %d pending records",
-                        self._dbn_queue._qsize(),
+                        self._dbn_queue.qsize(),
                     )
                     self._session.resume_reading()
 
-        self._dbn_queue._enabled.clear()
+        self._dbn_queue.disable()
         self.block_for_close()
+        logger.debug("completed iteration")
         raise StopIteration
 
     def __repr__(self) -> str:

databento/live/session.py

@@ -25,41 +25,98 @@
 
 AUTH_TIMEOUT_SECONDS: Final = 2.0
 CONNECT_TIMEOUT_SECONDS: Final = 5.0
+DBN_QUEUE_CAPACITY: Final = 2**20
 DEFAULT_REMOTE_PORT: Final = 13000
 
 
-class DBNQueue(queue.Queue):  # type: ignore [type-arg]
+class DBNQueue(queue.SimpleQueue):  # type: ignore [type-arg]
     """
     Queue for DBNRecords that can only be pushed to when enabled.
-
-    Parameters
-    ----------
-    maxsize : int
-        The `maxsize` for the Queue.
-
     """
 
-    def __init__(self, maxsize: int) -> None:
-        super().__init__(maxsize)
+    def __init__(self) -> None:
+        super().__init__()
         self._enabled = threading.Event()
 
-    @property
-    def enabled(self) -> bool:
+    def is_enabled(self) -> bool:
         """
-        True if the Queue will allow pushing.
+        Return True if the Queue will allow pushing; False otherwise.
 
         A queue should only be enabled when it has a consumer.
 
         """
         return self._enabled.is_set()
 
-    def half_full(self) -> bool:
+    def is_full(self) -> bool:
+        """
+        Return True when the queue has reached capacity; False otherwise.
+        """
+        return self.qsize() > DBN_QUEUE_CAPACITY
+
+    def enable(self) -> None:
+        """
+        Enable the DBN queue for pushing.
+        """
+        self._enabled.set()
+
+    def disable(self) -> None:
+        """
+        Disable the DBN queue for pushing.
+        """
+        self._enabled.clear()
+
+    def put(self, item: DBNRecord, block: bool = True, timeout: float | None = None) -> None:
+        """
+        Put an item on the queue if the queue is enabled.
+
+        Parameters
+        ----------
+        item: DBNRecord
+            The DBNRecord to put into the queue
+        block: bool, default True
+            Block if necessary until a free slot is available or the `timeout` is reached
+        timeout: float | None, default None
+            The maximum amount of time to block, when `block` is True, for the queue to become enabled.
+
+        Raises
+        ------
+        BentoError
+            If the queue is not enabled.
+            If the queue is not enabled within `timeout` seconds.
+
+        See Also
+        --------
+        queue.SimpleQueue.put
+
         """
-        Return True when the queue has reached half capacity.
+        if self._enabled.wait(timeout):
+            return super().put(item, block, timeout)
+        if timeout is not None:
+            raise BentoError(f"queue is not enabled after {timeout} second(s)")
+        raise BentoError("queue is not enabled")
+
+    def put_nowait(self, item: DBNRecord) -> None:
         """
-        with self.mutex:
-            return self._qsize() > self.maxsize // 2
+        Put an item on the queue, if the queue is enabled, without blocking.
 
+        Parameters
+        ----------
+        item: DBNRecord
+            The DBNRecord to put into the queue
+
+        Raises
+        ------
+        BentoError
+            If the queue is not enabled.
+
+        See Also
+        --------
+        queue.SimpleQueue.put_nowait
+
+        """
+        if self.is_enabled():
+            return super().put_nowait(item)
+        raise BentoError("queue is not enabled")
 
 @dataclasses.dataclass
 class SessionMetadata:
@@ -173,22 +230,16 @@ def received_record(self, record: DBNRecord) -> None:
                 if exc_callback is not None:
                     exc_callback(exc)
 
-        if self._dbn_queue.enabled:
-            try:
-                self._dbn_queue.put_nowait(record)
-            except queue.Full:
-                logger.error(
-                    "record queue is full; dropped %s record ts_event=%s",
-                    type(record).__name__,
-                    record.ts_event,
+        if self._dbn_queue.is_enabled():
+            self._dbn_queue.put(record)
+
+            # DBNQueue has no max size; so check if it's above capacity, and if so, pause reading
+            if self._dbn_queue.is_full():
+                logger.warning(
+                    "record queue is full; %d record(s) to be processed",
+                    self._dbn_queue.qsize(),
                 )
-            else:
-                if self._dbn_queue.half_full():
-                    logger.warning(
-                        "record queue is full; %d record(s) to be processed",
-                        self._dbn_queue._qsize(),
-                    )
-                    self.transport.pause_reading()
+                self.transport.pause_reading()
 
         return super().received_record(record)
 

tests/test_live_client.py

@@ -837,7 +837,7 @@ async def test_live_async_iteration_backpressure(
     when the queue is depleted when iterating asynchronously.
     """
     # Arrange
-    monkeypatch.setattr(client, "DEFAULT_QUEUE_SIZE", 4)
+    monkeypatch.setattr(session, "DBN_QUEUE_CAPACITY", 2)
 
     live_client = client.Live(
         key=test_api_key,
@@ -878,10 +878,11 @@ async def test_live_async_iteration_dropped(
     test_api_key: str,
 ) -> None:
     """
-    Test that an artificially small queue size will drop messages when full.
+    Test that an artificially small queue size will not drop messages when
+    full.
     """
     # Arrange
-    monkeypatch.setattr(client, "DEFAULT_QUEUE_SIZE", 1)
+    monkeypatch.setattr(session, "DBN_QUEUE_CAPACITY", 1)
 
     live_client = client.Live(
         key=test_api_key,
@@ -911,7 +912,7 @@ async def test_live_async_iteration_dropped(
     records = list(live_it)
 
     # Assert
-    assert len(records) == 1
+    assert len(records) == 4
     assert live_client._dbn_queue.empty()
 
 

tests/test_live_session.py

@@ -0,0 +1,46 @@
+import pytest
+from databento.common.error import BentoError
+from databento.live.session import DBNQueue
+
+
+def test_dbn_queue_put(
+    timeout: float = 0.01,
+) -> None:
+    """
+    Test that DBNQueue.put raises a BentoError if disabled.
+
+    The `timeout` is required, otherwise we will block forever.
+
+    """
+    # Arrange
+    queue = DBNQueue()
+
+    # Act, Assert
+    with pytest.raises(BentoError):
+        queue.put(None, timeout=timeout)
+
+    queue.enable()
+    queue.put(None, timeout=timeout)
+
+    queue.disable()
+    with pytest.raises(BentoError):
+        queue.put(None, timeout=timeout)
+
+
+def test_dbn_queue_put_nowait() -> None:
+    """
+    Test that DBNQueue.put_nowait raises a BentoError if disabled.
+    """
+    # Arrange
+    queue = DBNQueue()
+
+    # Act, Assert
+    with pytest.raises(BentoError):
+        queue.put_nowait(None)
+
+    queue.enable()
+    queue.put_nowait(None)
+
+    queue.disable()
+    with pytest.raises(BentoError):
+        queue.put_nowait(None)