add graceful shutdown for sessions

Author: Sam Daniel Thangarajan

src/nasdaq_protocols/common/session.py

@@ -107,13 +107,17 @@ 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)
     _task: asyncio.Task = attrs.field(init=False, default=None)
     _stopped: bool = attrs.field(init=False, default=False)
 
     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
+
     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:
@@ -133,8 +137,14 @@ def is_stopped(self):
 
     async def _process(self):
         while not self._stopped:
-            if len(self._buffer) > 0:
+            len_before = 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
+
             await asyncio.sleep(0.0001)
 
     async def _process_1(self):
@@ -200,6 +210,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)
@@ -255,6 +266,13 @@ def initiate_close(self) -> None:
             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()
+
     async def close(self):
         """
         Close the session, the session cannot be used after this call.
@@ -336,7 +354,10 @@ def connection_lost(self, exc):
         :meta private:
         """
         self.log.debug('%s> connection lost', self.session_id)
-        self.initiate_close()
+        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}')
 
     async def on_message(self, msg):
         await self._msg_queue.put(msg)

tests/test_common_asyncsession.py

@@ -13,7 +13,9 @@
 @attrs.define(auto_attribs=True)
 class SampleTestReader(common.Reader):
     def deserialize(self) -> Any:
-        return self._buffer.decode('ascii'), False, False
+        result = self._buffer.decode('ascii')
+        self._buffer = bytearray()
+        return result, False, False
 
     @staticmethod
     def create(session_id, on_msg, on_close):
@@ -39,6 +41,7 @@ async def send_heartbeat(self):
         pass
 
     async def on_close(self):
+        LOG.warning("Session %s closed", self.session_id)
         self.closed.set()
 
 
@@ -103,3 +106,31 @@ async def test__asyncsession__missed_remote_heartbeats__session_is_closed(mock_s
 
     # test client session is closed
     assert client_session.is_closed()
+
+
+async def test__asyncsession__stream_after_connect_and_close__client_session_able_to_read_all_messages(mock_server_session):
+    event_loop = asyncio.get_running_loop()
+    port, server_session = mock_server_session
+
+    # server sends 100 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)
+
+    session_ = SampleTestClientSession(session_id=common.SessionId())
+    _, session_ = 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:
+        msg = await asyncio.wait_for(session_.received.get(), 1)
+        if 'end-of-stream' in msg:
+            found_end_of_stream = True
+
+    await session_.closed.wait()
+
+    assert found_end_of_stream
+    assert session_.closed.is_set()
+    assert session_.is_closed()