refactoring canceling tasks to helpermethod

Author: [email protected]

rsocket/helpers.py

@@ -1,5 +1,6 @@
 import asyncio
 import struct
+from asyncio import Task
 from typing import TypeVar
 from contextlib import contextmanager
 from typing import Any
@@ -11,6 +12,7 @@
 from rsocket.exceptions import RSocketMimetypeTooLong
 from rsocket.exceptions import RSocketTransportError
 from rsocket.frame import Frame
+from rsocket.logger import logger
 from rsocket.payload import Payload
 
 _default = object()
@@ -111,3 +113,15 @@ def parse_well_known_encoding(buffer: bytes, encoding_name_provider: Callable[[W
         offset = 1 + real_mime_type_length
 
     return metadata_encoding, offset
+
+
+async def cancel_if_task_exists(task: Optional[Task]):
+    if task is not None:
+        task.cancel()
+
+        try:
+            await task
+        except asyncio.CancelledError:
+            logger().debug('Asyncio task cancellation error: %s', str(task))
+        except RuntimeError:
+            logger().error('Runtime error', exc_info=True)

rsocket/rsocket_base.py

@@ -29,7 +29,7 @@
 from rsocket.handlers.request_response_responder import RequestResponseResponder
 from rsocket.handlers.request_stream_requester import RequestStreamRequester
 from rsocket.handlers.request_stream_responder import RequestStreamResponder
-from rsocket.helpers import payload_from_frame, async_noop
+from rsocket.helpers import payload_from_frame, async_noop, cancel_if_task_exists
 from rsocket.lease import DefinedLease, NullLease, Lease
 from rsocket.logger import logger
 from rsocket.payload import Payload
@@ -413,8 +413,8 @@ async def close(self):
         logger().debug('%s: Closing', self._log_identifier())
 
         self._is_closing = True
-        await self._cancel_if_task_exists(self._sender_task)
-        await self._cancel_if_task_exists(self._receiver_task)
+        await cancel_if_task_exists(self._sender_task)
+        await cancel_if_task_exists(self._receiver_task)
 
         await self._close_transport()
 
@@ -430,16 +430,6 @@ async def _close_transport(self):
                     logger().debug('Transport already closed or failed to close: %s', str(exception))
                     pass
 
-    async def _cancel_if_task_exists(self, task):
-        if task is not None:
-            task.cancel()
-            try:
-                await task
-            except asyncio.CancelledError:
-                logger().debug('%s: Asyncio task cancellation error: %s', self._log_identifier(), str(task))
-            except RuntimeError:
-                logger().error('Runtime error', exc_info=True)
-
     async def __aenter__(self) -> 'RSocketBase':
         return self
 

rsocket/rsocket_client.py

@@ -7,7 +7,7 @@
 from reactivestreams.publisher import Publisher
 from rsocket.exceptions import RSocketNoAvailableTransport
 from rsocket.extensions.mimetypes import WellKnownMimeTypes
-from rsocket.helpers import create_future
+from rsocket.helpers import create_future, cancel_if_task_exists
 from rsocket.logger import logger
 from rsocket.payload import Payload
 from rsocket.request_handler import BaseRequestHandler
@@ -93,7 +93,7 @@ async def close(self):
 
     async def _close(self, reconnect=False):
         if not reconnect:
-            await self._cancel_if_task_exists(self._reconnect_task)
+            await cancel_if_task_exists(self._reconnect_task)
         else:
             logger().debug('%s: Closing before reconnect', self._log_identifier())
 
@@ -141,7 +141,7 @@ def _before_sender(self):
         self._keepalive_task = self._start_task_if_not_closing(self._keepalive_send_task)
 
     async def _finally_sender(self):
-        await self._cancel_if_task_exists(self._keepalive_task)
+        await cancel_if_task_exists(self._keepalive_task)
 
     def _update_last_keepalive(self):
         self._last_server_keepalive = datetime.now()
@@ -171,4 +171,4 @@ async def _receiver_listen(self):
         try:
             await super()._receiver_listen()
         finally:
-            await self._cancel_if_task_exists(keepalive_timeout_task)
+            await cancel_if_task_exists(keepalive_timeout_task)

rsocket/transports/abstract_messaging.py

@@ -12,8 +12,6 @@ def __init__(self):
     async def next_frame_generator(self):
         frame = await self._incoming_frame_queue.get()
 
-        self._incoming_frame_queue.task_done()
-
         if isinstance(frame, Exception):
             raise frame
 

rsocket/transports/aioquic_transport.py

@@ -7,7 +7,7 @@
 
 from rsocket.exceptions import RSocketTransportError
 from rsocket.frame import Frame
-from rsocket.helpers import wrap_transport_exception
+from rsocket.helpers import wrap_transport_exception, cancel_if_task_exists
 from rsocket.logger import logger
 from rsocket.rsocket_server import RSocketServer
 from rsocket.transports.abstract_messaging import AbstractMessagingTransport
@@ -95,16 +95,16 @@ async def incoming_data_listener(self):
 
                 if isinstance(data, Exception):
                     self._incoming_frame_queue.put_nowait(data)
+                    return
                 else:
                     async for frame in self._frame_parser.receive_data(data, 0):
                         self._incoming_frame_queue.put_nowait(frame)
 
-                self._incoming_bytes_queue.task_done()
         except asyncio.CancelledError:
             logger().debug('Asyncio task canceled: incoming_data_listener')
         except Exception:
             self._incoming_frame_queue.put_nowait(RSocketTransportError())
 
     async def close(self):
-        self._listener.cancel()
+        await cancel_if_task_exists(self._listener)
         self._quic_protocol.close()