Merge pull request #164 from mosquito/feafure/dataclasses

Conection stuck fixes

Author: mosquito

aiormq/abc.py

@@ -1,11 +1,15 @@
 import asyncio
+import dataclasses
+import io
+import logging
 from abc import ABC, abstractmethod, abstractproperty
 from types import TracebackType
 from typing import (
-    Any, Awaitable, Callable, Coroutine, Dict, Iterable, NamedTuple, Optional,
-    Set, Tuple, Type, Union,
+    Any, Awaitable, Callable, Coroutine, Dict, Iterable, Optional, Set, Tuple,
+    Type, Union,
 )
 
+import pamqp
 from pamqp import commands as spec
 from pamqp.base import Frame
 from pamqp.body import ContentBody
@@ -22,24 +26,24 @@
 
 # noinspection PyShadowingNames
 class TaskWrapper:
-    __slots__ = "exception", "task"
+    __slots__ = "_exception", "task"
 
-    exception: Union[BaseException, Type[BaseException]]
+    _exception: Union[BaseException, Type[BaseException]]
     task: asyncio.Task
 
     def __init__(self, task: asyncio.Task):
         self.task = task
-        self.exception = asyncio.CancelledError
+        self._exception = asyncio.CancelledError
 
     def throw(self, exception: ExceptionType) -> None:
-        self.exception = exception
+        self._exception = exception
         self.task.cancel()
 
     async def __inner(self) -> Any:
         try:
             return await self.task
         except asyncio.CancelledError as e:
-            raise self.exception from e
+            raise self._exception from e
 
     def __await__(self, *args: Any, **kwargs: Any) -> Any:
         return self.__inner().__await__()
@@ -59,7 +63,8 @@ def __repr__(self) -> str:
 GetResultType = Union[Basic.GetEmpty, Basic.GetOk]
 
 
-class DeliveredMessage(NamedTuple):
+@dataclasses.dataclass(frozen=True)
+class DeliveredMessage:
     delivery: Union[spec.Basic.Deliver, spec.Basic.Return, GetResultType]
     header: ContentHeader
     body: bytes
@@ -137,7 +142,8 @@ def message_count(self) -> Optional[int]:
 ]
 
 
-class SSLCerts(NamedTuple):
+@dataclasses.dataclass(frozen=True)
+class SSLCerts:
     cert: Optional[str]
     key: Optional[str]
     capath: Optional[str]
@@ -146,7 +152,8 @@ class SSLCerts(NamedTuple):
     verify: bool
 
 
-class FrameReceived(NamedTuple):
+@dataclasses.dataclass(frozen=True)
+class FrameReceived:
     channel: int
     frame: str
 
@@ -182,11 +189,56 @@ class FrameReceived(NamedTuple):
 ]
 
 
-class ChannelFrame(NamedTuple):
-    channel_number: int
-    frames: Iterable[Union[FrameType, Heartbeat, ContentBody]]
+@dataclasses.dataclass(frozen=True)
+class ChannelFrame:
+    payload: bytes
+    should_close: bool
     drain_future: Optional[asyncio.Future] = None
 
+    def drain(self) -> None:
+        if not self.should_drain:
+            return
+
+        if self.drain_future is not None:
+            self.drain_future.set_result(None)
+
+    @property
+    def should_drain(self) -> bool:
+        return self.drain_future is not None and not self.drain_future.done()
+
+    @classmethod
+    def marshall(
+        cls, channel_number: int,
+        frames: Iterable[Union[FrameType, Heartbeat, ContentBody]],
+        drain_future: Optional[asyncio.Future] = None,
+    ) -> "ChannelFrame":
+        should_close = False
+
+        with io.BytesIO() as fp:
+            for frame in frames:
+                if should_close:
+                    logger = logging.getLogger(
+                        "aiormq.connection"
+                    ).getChild(
+                        "marshall"
+                    )
+
+                    logger.warning(
+                        "It looks like you are going to send a frame %r after "
+                        "the connection is closed, it's pointless, "
+                        "the frame is dropped.", frame,
+                    )
+                    continue
+                if isinstance(frame, spec.Connection.CloseOk):
+                    should_close = True
+                fp.write(pamqp.frame.marshal(frame, channel_number))
+
+            return cls(
+                payload=fp.getvalue(),
+                drain_future=drain_future,
+                should_close=should_close,
+            )
+
 
 class AbstractFutureStore:
     futures: Set[Union[asyncio.Future, TaskType]]

aiormq/channel.py

@@ -2,12 +2,14 @@
 import io
 import logging
 from collections import OrderedDict
-from contextlib import suppress
 from functools import partial
 from io import BytesIO
 from random import getrandbits
 from types import MappingProxyType
-from typing import Any, Dict, List, Mapping, Optional, Set, Type, Union
+from typing import (
+    Any, Awaitable, Callable, Dict, List, Mapping, Optional, Set, Tuple, Type,
+    Union,
+)
 from uuid import UUID
 
 import pamqp.frame
@@ -166,7 +168,7 @@ async def rpc(
             try:
                 await countdown(
                     self.write_queue.put(
-                        ChannelFrame(
+                        ChannelFrame.marshall(
                             channel_number=self.number,
                             frames=[frame],
                         ),
@@ -195,7 +197,7 @@ async def rpc(
 
                 self.__close_event.set()
                 self.write_queue.put_nowait(
-                    ChannelFrame(
+                    ChannelFrame.marshall(
                         channel_number=self.number,
                         frames=[
                             spec.Channel.Close(
@@ -268,7 +270,7 @@ async def __get_content_header(self) -> ContentHeader:
 
         return frame
 
-    async def _on_deliver(self, frame: spec.Basic.Deliver) -> None:
+    async def _on_deliver_frame(self, frame: spec.Basic.Deliver) -> None:
         header: ContentHeader = await self.__get_content_header()
         message = await self._read_content(frame, header)
 
@@ -281,7 +283,7 @@ async def _on_deliver(self, frame: spec.Basic.Deliver) -> None:
             # noinspection PyAsyncCall
             self.create_task(consumer(message))
 
-    async def _on_get(
+    async def _on_get_frame(
         self, frame: Union[spec.Basic.GetOk, spec.Basic.GetEmpty],
     ) -> None:
         message = None
@@ -308,7 +310,7 @@ async def _on_get(
         getter.set_result((frame, message))
         return
 
-    async def _on_return(self, frame: spec.Basic.Return) -> None:
+    async def _on_return_frame(self, frame: spec.Basic.Return) -> None:
         header: ContentHeader = await self.__get_content_header()
         message = await self._read_content(frame, header)
         message_id = message.header.properties.message_id
@@ -366,7 +368,7 @@ def _confirm_delivery(
             DeliveryError(None, frame),
         )  # pragma: nocover
 
-    async def _on_confirm(self, frame: ConfirmationFrameType) -> None:
+    async def _on_confirm_frame(self, frame: ConfirmationFrameType) -> None:
         if not self.publisher_confirms:  # pragma: nocover
             return
 
@@ -386,61 +388,59 @@ async def _on_confirm(self, frame: ConfirmationFrameType) -> None:
         else:
             self._confirm_delivery(frame.delivery_tag, frame)
 
+    async def _on_cancel_frame(
+        self,
+        frame: Union[spec.Basic.CancelOk, spec.Basic.Cancel],
+    ) -> None:
+        if frame.consumer_tag is not None:
+            self.consumers.pop(frame.consumer_tag, None)
+
+    async def _on_close_frame(self, frame: spec.Channel.Close) -> None:
+        exc: BaseException = exception_by_code(frame)
+        self.write_queue.put_nowait(
+            ChannelFrame.marshall(
+                channel_number=self.number,
+                frames=[spec.Channel.CloseOk()],
+            ),
+        )
+        self.connection.channels.pop(self.number, None)
+        raise exc
+
+    async def _on_close_ok_frame(self, _: spec.Channel.CloseOk) -> None:
+        self.connection.channels.pop(self.number, None)
+        raise ChannelClosed()
+
     async def _reader(self) -> None:
-        while True:
-            try:
+        hooks: Mapping[Any, Tuple[bool, Callable[[Any], Awaitable[None]]]]
+
+        hooks = {
+            spec.Basic.Deliver: (False, self._on_deliver_frame),
+            spec.Basic.GetOk: (True, self._on_get_frame),
+            spec.Basic.GetEmpty: (True, self._on_get_frame),
+            spec.Basic.Return: (False, self._on_return_frame),
+            spec.Basic.Cancel: (False, self._on_cancel_frame),
+            spec.Basic.CancelOk: (True, self._on_cancel_frame),
+            spec.Channel.Close: (False, self._on_close_frame),
+            spec.Channel.CloseOk: (False, self._on_close_ok_frame),
+            spec.Basic.Ack: (False, self._on_confirm_frame),
+            spec.Basic.Nack: (False, self._on_confirm_frame),
+        }
+
+        try:
+            while True:
                 frame = await self._get_frame()
+                should_add_to_rpc, hook = hooks.get(type(frame), (True, None))
 
-                if isinstance(frame, spec.Basic.Deliver):
-                    with suppress(Exception):
-                        await self._on_deliver(frame)
-                    continue
-                elif isinstance(
-                    frame, (spec.Basic.GetOk, spec.Basic.GetEmpty),
-                ):
-                    with suppress(Exception):
-                        await self._on_get(frame)
-                elif isinstance(frame, spec.Basic.Return):
-                    with suppress(Exception):
-                        await self._on_return(frame)
-                    continue
-                elif isinstance(frame, spec.Basic.Cancel):
-                    if frame.consumer_tag is None:
-                        continue
-                    self.consumers.pop(frame.consumer_tag, None)
-                    continue
-                elif isinstance(frame, spec.Basic.CancelOk):
-                    if frame.consumer_tag is not None:
-                        self.consumers.pop(frame.consumer_tag, None)
-                elif isinstance(frame, (spec.Basic.Ack, spec.Basic.Nack)):
-                    with suppress(Exception):
-                        await self._on_confirm(frame)
-                    continue
-                elif isinstance(frame, spec.Channel.Close):
-                    exc: BaseException = exception_by_code(frame)
-                    self.write_queue.put_nowait(
-                        ChannelFrame(
-                            channel_number=self.number,
-                            frames=[spec.Channel.CloseOk()],
-                        ),
-                    )
-
-                    self.connection.channels.pop(self.number, None)
-                    await self._cancel_tasks(exc)
-                    return
-                elif isinstance(frame, spec.Channel.CloseOk):
-                    if self.__close_event.is_set():
-                        await self._cancel_tasks(asyncio.TimeoutError())
-                        self.connection.channels.pop(self.number, None)
+                if hook is not None:
+                    try:
+                        await hook(frame)
+                    except asyncio.CancelledError as e:
+                        await self._cancel_tasks(e)
                         return
-
-                await self.rpc_frames.put(frame)
-            except asyncio.CancelledError:
-                return
-            except Exception as e:  # pragma: nocover
-                log.debug("Channel reader exception %r", exc_info=e)
-                await self._cancel_tasks(e)
-                raise
+                if should_add_to_rpc:
+                    await self.rpc_frames.put(frame)
+        except Exception as e:
+            await self._cancel_tasks(e)
 
     @task
     async def _on_close(self, exc: Optional[ExceptionType] = None) -> None:
@@ -454,6 +454,7 @@ async def _on_close(self, exc: Optional[ExceptionType] = None) -> None:
                 timeout=self.connection.connection_tune.heartbeat or None,
             )
         self.connection.channels.pop(self.number, None)
+        self.__close_event.set()
 
     async def basic_get(
         self, queue: str = "", no_ack: bool = False,
@@ -525,7 +526,7 @@ async def basic_ack(
         drain_future = self.create_future() if wait else None
 
         await self.write_queue.put(
-            ChannelFrame(
+            ChannelFrame.marshall(
                 frames=[
                     spec.Basic.Ack(
                         delivery_tag=delivery_tag,
@@ -553,7 +554,7 @@ async def basic_nack(
         drain_future = self.create_future() if wait else None
 
         await self.write_queue.put(
-            ChannelFrame(
+            ChannelFrame.marshall(
                 frames=[
                     spec.Basic.Nack(
                         delivery_tag=delivery_tag,
@@ -574,7 +575,7 @@ async def basic_reject(
     ) -> None:
         drain_future = self.create_future()
         await self.write_queue.put(
-            ChannelFrame(
+            ChannelFrame.marshall(
                 channel_number=self.number,
                 frames=[
                     spec.Basic.Reject(
@@ -664,7 +665,7 @@ async def basic_publish(
 
             await countdown(
                 self.write_queue.put(
-                    ChannelFrame(
+                    ChannelFrame.marshall(
                         frames=body_frames,
                         channel_number=self.number,
                         drain_future=drain_future,