more work in progress

Author: Sam Daniel Thangarajan

src/nasdaq_protocols/common/message_queue.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 from contextlib import asynccontextmanager
 from itertools import count
 from typing import Any, Callable, Awaitable
@@ -26,19 +27,25 @@ class DispatchableMessageQueue(Stoppable):
     on_msg_coro: DispatcherCoro = None
     _closed: bool = attrs.field(init=False, default=False)
     _msg_queue: asyncio.Queue = attrs.field(init=False, default=None)
+    _buffer_msg_queue: asyncio.Queue | None = attrs.field(init=False, default=None)
     _recv_task: asyncio.Task = attrs.field(init=False, default=None)
     _dispatcher_task: asyncio.Task = attrs.field(init=False, default=None)
 
     def __attrs_post_init__(self):
         self._msg_queue = asyncio.Queue()
         self.start_dispatching(self.on_msg_coro)
 
+    def __len__(self) -> int:
+        """Return the number of entries in the queue."""
+        return self._msg_queue.qsize()
+
     async def put(self, msg: Any) -> None:
         """
         put an entry into the queue.
         :param msg: Any
         """
-        await self._msg_queue.put(msg)
+        queue = self._buffer_msg_queue if self._buffer_msg_queue else self._msg_queue
+        await queue.put(msg)
 
     async def get(self):
         """get an entry from the queue.
@@ -59,7 +66,8 @@ def put_nowait(self, msg: Any):
         put an entry into the queue.
         :param msg: Any
         """
-        self._msg_queue.put_nowait(msg)
+        queue = self._buffer_msg_queue if self._buffer_msg_queue else self._msg_queue
+        queue.put_nowait(msg)
 
     def get_nowait(self) -> Any | None:
         """
@@ -118,6 +126,27 @@ def start_dispatching(self, on_msg_coro: DispatcherCoro) -> None:
             self._dispatcher_task = asyncio.create_task(self._start_dispatching(), name=f'{self.session_id}-dispatcher')
             self.log.debug('%s> queue dispatcher started.', self.session_id)
 
+    @contextlib.asynccontextmanager
+    async def buffer_until_drained(self, discard_buffer: bool = False):
+        """Async context manager that waits until the buffer is drained."""
+        if self._buffer_msg_queue:
+            raise StateError('Already blocking new messages, cannot nest block_until_empty')
+
+        self._buffer_msg_queue = asyncio.Queue()
+
+        try:
+            while not self._msg_queue.empty():
+                await asyncio.sleep(0.0001)
+            yield
+        finally:
+            if not discard_buffer:
+                if self.is_dispatching():
+                    async with self.pause_dispatching():
+                        self._msg_queue = self._buffer_msg_queue
+                else:
+                    self._msg_queue = self._buffer_msg_queue
+            self._buffer_msg_queue = None
+
     async def stop(self) -> None:
         """
         Stop the queue.
@@ -137,8 +166,7 @@ async def _start_dispatching(self):
         counter = count(1)
         while True:
             try:
-                msg = await self._msg_queue.get()
-                await self.on_msg_coro(msg)
+                await self.on_msg_coro(await self._msg_queue.get())
                 self.log.debug('%s> dispatched message %s', self.session_id, next(counter))
             except asyncio.CancelledError:
                 break

src/nasdaq_protocols/common/session.py

@@ -1,5 +1,6 @@
 import abc
 import asyncio
+import contextlib
 from typing import Any, Callable, Coroutine, Generic, Type, TypeVar
 from itertools import count
 
@@ -107,22 +108,37 @@ class Reader(Stoppable):
     session_id: Any = attrs.field(validator=Validators.not_none())
     on_msg_coro: OnMsgCoro = attrs.field(validator=Validators.not_none())
     on_close_coro: OnCloseCoro = attrs.field(validator=Validators.not_none())
-    _close_when_drained: bool = attrs.field(init=False, default=False)
     _buffer: bytearray = attrs.field(init=False, factory=bytearray)
+    _drain_buffer: bytearray | None = attrs.field(init=False, default=None)
     _task: asyncio.Task = attrs.field(init=False, default=None)
     _stopped: bool = attrs.field(init=False, default=False)
+    _drain_mode: asyncio.Event | None = attrs.field(init=False, default=None)
 
     def __attrs_post_init__(self):
         self._task = asyncio.create_task(self._process(), name=f'reader:{self.session_id}')
 
-    def close_when_drained(self):
-        self._close_when_drained = True
+    @contextlib.asynccontextmanager
+    async def buffer_until_drained(self, discard_buffer: bool = False):
+        """Async context manager that waits until the buffer is drained."""
+        if self._drain_mode:
+            raise StateError('Already draining, cannot nest buffer_until_drained')
+        self._drain_mode = asyncio.Event()
+        self._drain_buffer = bytearray()
+        try:
+            await self._drain_mode.wait()
+            yield
+        finally:
+            self._drain_mode = None
+            if not discard_buffer:
+                self._buffer.extend(self._drain_buffer)
+            self._drain_buffer = None
 
     def on_data(self, data: bytes):
         self.log.debug('%s> on_data: existing = %s, received = %s', self.session_id, self._buffer, data)
         if len(data) == 0:
             return
-        self._buffer.extend(data)
+        buffer = self._drain_buffer if self._drain_buffer is not None else self._buffer
+        buffer.extend(data)
 
     async def stop(self):
         if self._stopped:
@@ -137,13 +153,12 @@ def is_stopped(self):
 
     async def _process(self):
         while not self._stopped:
-            len_before = len(self._buffer)
+            len_before = len_after = len(self._buffer)
             if len_before > 0:
                 await self._process_1()
-            if self._close_when_drained and len(self._buffer) == len_before:
-                self.log.debug('%s> buffer drained, closing reader.', self.session_id)
-                await self.stop()
-                return
+                len_after = len(self._buffer)
+            if self._drain_mode and not self._drain_mode.is_set() and len_after == len_before:
+                self._drain_mode.set()
 
             await asyncio.sleep(0.0001)
 
@@ -253,7 +268,7 @@ def is_closed(self) -> bool:
         """
         return self._closed
 
-    def initiate_close(self) -> None:
+    def initiate_close(self, drain: bool = False) -> None:
         """
         Initiate close of the session.
         An asynchronous task is created which will close the session and all its
@@ -264,29 +279,38 @@ def initiate_close(self) -> None:
         """
         if self._closed or self._closing_task:
             return
-        self._closing_task = asyncio.create_task(self.close(), name=f'asyncsession-close:{self.session_id}')
 
-    async def _graceful_close(self):
-        await stop_task([
-            self._local_hb_monitor,
-            self._remote_hb_monitor,
-        ])
-        self._reader.close_when_drained()
+        if drain:
+            name = f'asyncsession-close-drain:{self.session_id}'
+        else:
+            name = f'asyncsession-close:{self.session_id}'
+
+        self._closing_task = asyncio.create_task(self.close(drain), name=name)
+
+    @contextlib.asynccontextmanager
+    async def buffer_until_drained(self, discard_buffer: bool = False):
+        """Async context manager that waits until both the reader and message queue are drained."""
+        async with self._reader.buffer_until_drained(discard_buffer=discard_buffer):
+            async with self._msg_queue.buffer_until_drained(discard_buffer=discard_buffer):
+                yield
 
-    async def close(self):
+    async def close(self, drain: bool = False):
         """
         Close the session, the session cannot be used after this call.
         """
         if not self._closed:
             self._closed = True
-            await stop_task([
-                self._msg_queue,
-                self._local_hb_monitor,
-                self._remote_hb_monitor,
-                self._reader
-            ])
             if self._transport:
                 self._transport.close()
+
+            await stop_task([self._local_hb_monitor, self._remote_hb_monitor])
+
+            if drain:
+                async with self.buffer_until_drained(discard_buffer=True):
+                    self.log.debug('%s> close: drained session.', self.session_id)
+
+            await stop_task([self._msg_queue,self._reader])
+
             if self.on_close_coro:
                 await self.on_close_coro()
 
@@ -354,12 +378,7 @@ def connection_lost(self, exc):
         :meta private:
         """
         self.log.debug('%s> connection lost', self.session_id)
-        if self.graceful_shutdown:
-            if self._closed or self._closing_task:
-                return
-            self._closing_task = asyncio.create_task(self._graceful_close(), name=f'asyncsession-graceful-close:{self.session_id}')
-        else:
-            self.initiate_close()
+        self.initiate_close(self.graceful_shutdown)
 
     async def on_message(self, msg):
         await self._msg_queue.put(msg)

tests/reader_app_tests.py

@@ -148,6 +148,37 @@ async def reader__empty_data__no_effect(**kwargs):
     await reader.stop()
 
 
+@reader_test
+async def reader__buffer_until_drained__buffers_not_discarded(**kwargs):
+    handler, reader, input_factory, output_factory = all_test_params(**kwargs)
+
+    discard_buffers_modes = [False, True]
+    for discard_buffer in discard_buffers_modes:
+        reader.on_data(input_factory(1))
+        async with reader.buffer_until_drained(discard_buffer=discard_buffer):
+            first = await handler.received_messages.get()
+            assert first == output_factory(1)
+
+            # Send another message while buffering
+            reader.on_data(input_factory(2))
+            await asyncio.sleep(0.1)
+
+            # Ensure no new messages are received while buffering
+            try:
+                second = await asyncio.wait_for(handler.received_messages.get(), timeout=0.5)
+                assert False, f'Unexpected message received: {second}'
+            except asyncio.TimeoutError:
+                pass
+
+        if not discard_buffer:
+            # After exiting the context, the buffered message should be processed
+            second = await asyncio.wait_for(handler.received_messages.get(), timeout=0.5)
+            assert second == output_factory(2)
+        else:
+            with pytest.raises(asyncio.TimeoutError):
+                await asyncio.wait_for(handler.received_messages.get(), timeout=0.5)
+
+
 @pytest.fixture(scope='function', params=READER_TESTS)
 async def reader_clientapp_common_tests(request, handler):
     async def _test(reader_factory, input_factory, output_factory):

tests/test_common_asyncsession.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import socket
 from typing import Any
 
 import attrs
@@ -12,31 +13,50 @@
 
 @attrs.define(auto_attribs=True)
 class SampleTestReader(common.Reader):
+    separator: bytes = attrs.field(kw_only=True, default=None)
+
     def deserialize(self) -> Any:
-        result = self._buffer.decode('ascii')
-        self._buffer = bytearray()
+        idx, sep_len = len(self._buffer), len(self.separator) if self.separator else 0
+        if self.separator:
+            idx = self._buffer.find(self.separator)
+            if idx == -1:
+                return None, False, False
+
+        result = self._buffer[:idx].decode('ascii')
+        self._buffer = self._buffer[idx + sep_len:]
         return result, False, False
 
     @staticmethod
     def create(session_id, on_msg, on_close):
         return SampleTestReader(session_id, on_msg, on_close)
 
+    @staticmethod
+    def creator(separator: bytes):
+        def _creator(session_id, on_msg, on_close):
+            return SampleTestReader(session_id, on_msg, on_close, separator=separator)
+        return _creator
+
 
 @attrs.define(auto_attribs=True)
 class SampleTestClientSession(common.AsyncSession):
     session_id: Any = attrs.field(validator=common.Validators.not_none())
     reader_factory: common.ReaderFactory = SampleTestReader.create
     received: asyncio.Queue = attrs.field(init=False, factory=asyncio.Queue)
     closed: asyncio.Event = attrs.field(init=False, factory=asyncio.Event)
+    slow_client: bool = attrs.field(kw_only=True, default=False)
 
     def __attrs_post_init__(self):
-        self.on_msg_coro = self.received.put
+        self.on_msg_coro = self.slow_on_msg if self.slow_client else self.received.put
         self.on_close_coro = self.on_close
         super().__attrs_post_init__()
 
     def send_msg(self, data: str):
         self._transport.write(data.encode('ascii'))
 
+    async def slow_on_msg(self, msg):
+        await self.received.put(msg)
+        await asyncio.sleep(0.01)
+
     async def send_heartbeat(self):
         pass
 
@@ -135,32 +155,56 @@ def stream_messages_and_close(session, _):
     assert session_.closed.is_set()
     assert session_.is_closed()
 
-async def test__asyncsession__stream_after_connect_and_close__no_graceful_close_in_client_session(mock_server_session):
+
+@pytest.mark.parametrize('graceful_shutdown', [False, True])
+@pytest.mark.parametrize('num_messages', [1, 10, 100])
+@pytest.mark.parametrize('slow_client', [False, True])
+async def test__asyncsession__graceful_shutdown__enables_client_to_read_all_messages(mock_server_session, graceful_shutdown, num_messages, slow_client):
     event_loop = asyncio.get_running_loop()
     port, server_session = mock_server_session
 
-    # server sends 100000 messages after connection and then closes the session
-    def stream_messages_and_close(session, _):
-        for i in range(100000):
-            session.send(f'msg{i}|'.encode('ascii'))
-        session.send(f'end-of-stream'.encode('ascii'))
-        session.close()
-    server_session.when_connect().do(stream_messages_and_close)
+    # server sends n messages after connection and then closes the session
+    server_session.when_connect().do(get_streamer_close_function(num_messages))
 
-    session_ = SampleTestClientSession(session_id=common.SessionId(), graceful_shutdown=False)
-    _, session_ = await event_loop.create_connection(lambda: session_, '127.0.0.1', port=port)
+    session_ = SampleTestClientSession(
+        session_id=common.SessionId(),
+        reader_factory=SampleTestReader.creator(separator=b'|'),
+        graceful_shutdown=graceful_shutdown,
+        slow_client=slow_client
+    )
+
+    _1, _2 = await event_loop.create_connection(
+        lambda: session_, '127.0.0.1', port=port
+    )
 
     found_end_of_stream = False
-    while not session_.closed.is_set() and not found_end_of_stream:
+    while not found_end_of_stream:
         try:
             msg = await asyncio.wait_for(session_.received.get(), 0.1)
+            if 'end-of-stream' == msg:
+                found_end_of_stream = True
+        except asyncio.TimeoutError:
+            if session_.closed.is_set():
+                break
         except Exception:
             break
-        if 'end-of-stream' in msg:
-            found_end_of_stream = True
 
+    task_id = asyncio.current_task().get_name()
+    LOG.info('Waiting for session to close, from task %s', task_id)
     await session_.closed.wait()
+    LOG.info('Session closed')
 
     assert session_.closed.is_set()
     assert session_.is_closed()
-    assert found_end_of_stream == False
+
+    assert found_end_of_stream == graceful_shutdown, \
+        "Graceful shutdown should allow reading all messages"
+
+
+def get_streamer_close_function(message_count: int):
+    def stream_messages_and_close(session, _):
+        for i in range(message_count):
+            session.send(f'msg{i}|'.encode('ascii'))
+        session.send(f'end-of-stream|'.encode('ascii'))
+        session.close()
+    return stream_messages_and_close