add graceful shutdown for sessions

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
 
@@ -108,17 +109,36 @@ class Reader(Stoppable):
     on_msg_coro: OnMsgCoro = attrs.field(validator=Validators.not_none())
     on_close_coro: OnCloseCoro = attrs.field(validator=Validators.not_none())
     _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}')
 
+    @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:
@@ -133,8 +153,13 @@ def is_stopped(self):
 
     async def _process(self):
         while not self._stopped:
-            if len(self._buffer) > 0:
+            len_before = len_after = len(self._buffer)
+            if len_before > 0:
                 await self._process_1()
+                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)
 
     async def _process_1(self):
@@ -200,6 +225,7 @@ class AsyncSession(asyncio.Protocol, abc.ABC, Generic[T]):
     on_msg_coro: OnMsgCoro = attrs.field(kw_only=True, default=None)
     on_close_coro: OnCloseCoro = attrs.field(kw_only=True, default=None)
     dispatch_on_connect: bool = attrs.field(kw_only=True, default=True)
+    graceful_shutdown: bool = attrs.field(kw_only=True, default=True)
     _reader: Reader = attrs.field(init=False, default=None)
     _transport: asyncio.Transport = attrs.field(init=False, default=None)
     _closed: bool = attrs.field(init=False, default=False)
@@ -242,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
@@ -253,22 +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 close(self):
+        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, 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()
 
@@ -336,7 +378,7 @@ def connection_lost(self, exc):
         :meta private:
         """
         self.log.debug('%s> connection lost', self.session_id)
-        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

@@ -58,6 +58,9 @@ async def reader__stop__reader_is_stopped(**kwargs):
 
     assert handler.closed.is_set()
 
+    # Double stop should be no-op
+    await common.stop_task(reader)
+
 
 @reader_test
 async def reader__one_msg_per_packet__msg_is_read(**kwargs):
@@ -148,6 +151,51 @@ 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)
+
+    await reader.stop()
+
+
+@reader_test
+async def reader__buffer_until_drained__nested_call_raises_error(**kwargs):
+    handler, reader, input_factory, output_factory = all_test_params(**kwargs)
+
+    async with reader.buffer_until_drained():
+        with pytest.raises(RuntimeError):
+            async with reader.buffer_until_drained():
+                pass
+
+    await reader.stop()
+
+
 @pytest.fixture(scope='function', params=READER_TESTS)
 async def reader_clientapp_common_tests(request, handler):
     async def _test(reader_factory, input_factory, output_factory):