FIX: Fix excessive receive buffer allocation

Author: nmacholl

databento/live/client.py

@@ -145,7 +145,6 @@ def __next__(self) -> DBNRecord:
                     "yielding %s record from next",
                     type(record).__name__,
                 )
-                self._dbn_queue.task_done()
                 return record
             finally:
                 if not self._dbn_queue.half_full() and not self._session.is_reading():

databento/live/protocol.py

@@ -31,7 +31,7 @@
 from databento.live.gateway import SubscriptionRequest
 
 
-MIN_BUFFER_SIZE: int = 64 * 1024  # 64kb
+RECV_BUFFER_SIZE: int = 64 * 2**10  # 64kb
 
 logger = logging.getLogger(__name__)
 
@@ -100,7 +100,7 @@ def __init__(
     ) -> None:
         self.__api_key = api_key
         self.__transport: asyncio.Transport | None = None
-        self.__buffer: bytearray
+        self.__buffer: bytearray = bytearray(RECV_BUFFER_SIZE)
 
         self._dataset = validate_semantic_string(dataset, "dataset")
         self._ts_out = ts_out
@@ -231,7 +231,8 @@ def get_buffer(self, sizehint: int) -> bytearray:
         asycnio.BufferedProtocol.get_buffer
 
         """
-        self.__buffer = bytearray(max(sizehint, MIN_BUFFER_SIZE))
+        if len(self.__buffer) < sizehint:
+            self.__buffer = bytearray(sizehint)
         return self.__buffer
 
     def buffer_updated(self, nbytes: int) -> None:

tests/test_live_client.py

@@ -822,9 +822,9 @@ async def test_live_stream_to_dbn_with_tiny_buffer(
 ) -> None:
     """
     Test that DBN data streamed by the MockLiveServer is properly re-
-    constructed client side when using the small values for MIN_BUFFER_SIZE.
+    constructed client side when using the small values for RECV_BUFFER_SIZE.
     """
-    monkeypatch.setattr(protocol, "MIN_BUFFER_SIZE", buffer_size)
+    monkeypatch.setattr(protocol, "RECV_BUFFER_SIZE", buffer_size)
     output = tmp_path / "output.dbn"
 
     live_client.subscribe(