Ignore all methods except Close and Close-OK when channel/connection is closing

Author: matusvalo

amqp/channel.py

@@ -31,6 +31,10 @@
 consumer_tag=%r exchange=%r routing_key=%r.\
 """
 
+IGNORED_METHOD_DURING_CHANNEL_CLOSE = """\
+Received method %s during closing channel. This method will be ignored\
+"""
+
 
 class VDeprecationWarning(DeprecationWarning):
     pass
@@ -108,6 +112,7 @@ def __init__(self, connection,
 
         super(Channel, self).__init__(connection, channel_id)
 
+        self.is_closing = False
         self.is_open = False
         self.active = True  # Flow control
         self.returned_messages = Queue()
@@ -159,6 +164,21 @@ def collect(self):
         self.events.clear()
         self.no_ack_consumers.clear()
 
+    def dispatch_method(self, method_sig, payload, content):
+        if self.is_closing and method_sig not in (
+            spec.Channel.Close, spec.Channel.CloseOk
+        ):
+            # When channel.close() was called we must ignore all methods except
+            # Channel.close and Channel.CloseOk
+            AMQP_LOGGER.warning(
+                IGNORED_METHOD_DURING_CHANNEL_CLOSE, method_sig
+            )
+            return
+        else:
+            return super(Channel, self).dispatch_method(
+                method_sig, payload, content
+            )
+
     def _do_revive(self):
         self.is_open = False
         self.open()

amqp/connection.py

@@ -39,6 +39,10 @@ class SSLError(Exception):  # noqa
 Start from server, version: %d.%d, properties: %s, mechanisms: %s, locales: %s
 """.strip()
 
+IGNORED_METHOD_DURING_CONN_CLOSE = """\
+Received method %s during closing connection. This method will be ignored\
+"""
+
 __all__ = ['Connection']
 
 AMQP_LOGGER = logging.getLogger('amqp')
@@ -233,6 +237,7 @@ def __init__(self, host='localhost:5672', userid='guest', password='guest',
         self.frame_handler_cls = frame_handler
         self.frame_writer_cls = frame_writer
 
+        self.is_closing = False
         self._handshake_complete = False
 
         self.channels = {}
@@ -292,6 +297,21 @@ def _setup_listeners(self):
             spec.Connection.CloseOk: self._on_close_ok,
         })
 
+    def dispatch_method(self, method_sig, payload, content):
+        if self.is_closing and method_sig not in (
+            spec.Connection.Close, spec.Connection.CloseOk
+        ):
+            # When Connection.close() was called we must ignore
+            # all methods except Connection.close and Connection.CloseOk
+            AMQP_LOGGER.warning(
+                IGNORED_METHOD_DURING_CONN_CLOSE, method_sig
+            )
+            return
+        else:
+            return super(Connection, self).dispatch_method(
+                method_sig, payload, content
+            )
+
     def connect(self, callback=None):
         # Let the transport.py module setup the actual
         # socket connection to the broker.
@@ -576,10 +596,11 @@ def close(self, reply_code=0, reply_text='', method_sig=(0, 0),
                 wait=spec.Connection.CloseOk,
             )
         except (OSError, IOError, SSLError):
-            self.is_closing = False
             # close connection
             self.collect()
             raise
+        finally:
+            self.is_closing = False
 
     def _on_close(self, reply_code, reply_text, class_id, method_id):
         """Request a connection close.

t/integration/test_integration.py

@@ -413,6 +413,74 @@ def test_connection_methods(self, method, callback):
                 conn.drain_events(0)
                 callback_mock.assert_called_once()
 
+    def test_channel_ignore_methods_during_close(self):
+        # Test checking that py-amqp will discard any received methods
+        # except Close and Close-OK after sending Channel.Close method
+        # to server.
+        frame_writer_cls_mock = Mock()
+        conn = Connection(frame_writer=frame_writer_cls_mock)
+        consumer_tag = 'amq.ctag-PCmzXGkhCw_v0Zq7jXyvkg'
+        with patch.object(conn, 'Transport') as transport_mock:
+            handshake(conn, transport_mock)
+
+            channel_id = 1
+            transport_mock().read_frame.side_effect = [
+                # Inject Open Handshake
+                build_frame_type_1(
+                    spec.Channel.OpenOk,
+                    channel=channel_id,
+                    args=(1, False),
+                    arg_format='Lb'
+                ),
+                # Inject basic-deliver response
+                build_frame_type_1(
+                    spec.Basic.Deliver,
+                    channel=1,
+                    arg_format='sLbss',
+                    args=(
+                        # consumer-tag, delivery-tag, redelivered,
+                        consumer_tag, 1, False,
+                        # exchange-name, routing-key
+                        'foo_exchange', 'routing-key'
+                    )
+                ),
+                build_frame_type_2(
+                    channel=1,
+                    body_len=12,
+                    properties=b'0\x00\x00\x00\x00\x00\x01'
+                ),
+                build_frame_type_3(
+                    channel=1,
+                    body=b'Hello World!'
+                ),
+                # Inject close method
+                build_frame_type_1(
+                    spec.Channel.CloseOk,
+                    channel=channel_id
+                ),
+            ]
+
+            frame_writer_mock = frame_writer_cls_mock()
+            frame_writer_mock.reset_mock()
+
+            with patch('amqp.Channel._on_basic_deliver') as on_deliver_mock:
+                ch = conn.channel(channel_id=channel_id)
+                ch.close()
+                on_deliver_mock.assert_not_called()
+            frame_writer_mock.assert_has_calls(
+                [
+                    call(
+                        1, 1, spec.Channel.Open, dumps('s', ('',)),
+                        None
+                    ),
+                    call(
+                        1, 1, spec.Channel.Close, dumps('BsBB', (0, '', 0, 0)),
+                        None
+                    )
+                ]
+            )
+            assert ch.is_open is False
+
     def test_channel_open_close(self):
         # Test checking opening and closing channel
         frame_writer_cls_mock = Mock()

t/unit/test_channel.py

@@ -7,7 +7,7 @@
 
 from amqp import spec
 from amqp.basic_message import Message
-from amqp.channel import Channel
+from amqp.channel import Channel, IGNORED_METHOD_DURING_CHANNEL_CLOSE
 from amqp.exceptions import (ConsumerCancelled, MessageNacked, NotFound,
                              RecoverableConnectionError)
 from amqp.platform import pack
@@ -26,6 +26,9 @@ def setup_conn(self):
         self.c = Channel(self.conn, 1)
         self.c.send_method = Mock(name='send_method')
 
+    def test_init_is_closing(self):
+        assert self.c.is_closing is False
+
     def test_init_confirm_enabled(self):
         self.conn.confirm_publish = True
         c = Channel(self.conn, 2)
@@ -88,6 +91,7 @@ def test_close(self):
             (30, 'text', spec.Queue.Declare[0], spec.Queue.Declare[1]),
             wait=spec.Channel.CloseOk,
         )
+        assert self.c.is_closing is False
         assert self.c.connection is None
 
     def test_on_close(self):
@@ -598,3 +602,29 @@ def test_on_basic_nack(self):
         self.c.events['basic_nack'].add(callback)
         self.c._on_basic_nack(123, True)
         callback.assert_called_with(123, True)
+
+    @pytest.mark.parametrize(
+        "method",
+        (spec.Channel.Close, spec.Channel.CloseOk, spec.Basic.Deliver)
+    )
+    def test_dispatch_method_close(self, method):
+        with patch('amqp.channel.AbstractChannel.dispatch_method') as dm_mock:
+            self.c.dispatch_method(method, 'aa', 'bb')
+            dm_mock.assert_called_once_with(method, 'aa', 'bb')
+
+    @pytest.mark.parametrize(
+        "method",
+        (spec.Channel.Close, spec.Channel.CloseOk, spec.Basic.Deliver)
+    )
+    def test_dispatch_method_closing_channel(self, method, caplog):
+        self.c.is_closing = True
+        with patch('amqp.channel.AbstractChannel.dispatch_method') as dm_mock:
+            self.c.dispatch_method(method, 'aa', 'bb')
+            if method in (spec.Channel.Close, spec.Channel.CloseOk):
+                dm_mock.assert_called_once_with(method, 'aa', 'bb')
+            else:
+                dm_mock.assert_not_called()
+                assert caplog.records[0].msg == \
+                    IGNORED_METHOD_DURING_CHANNEL_CLOSE
+                assert caplog.records[0].args[0] == method
+                assert caplog.records[0].levelname == 'WARNING'

t/unit/test_connection.py

@@ -6,12 +6,12 @@
 import pytest
 
 from amqp import Connection, spec
-from amqp.connection import SSLError
+from amqp.connection import SSLError, IGNORED_METHOD_DURING_CONN_CLOSE
 from amqp.exceptions import ConnectionError, NotFound, ResourceError
 from amqp.five import items
 from amqp.sasl import AMQPLAIN, EXTERNAL, GSSAPI, PLAIN, SASL
 from amqp.transport import TCPTransport
-from case import ContextMock, Mock, call, patch
+from case import ContextMock, Mock, call, patch, sentinel
 
 
 class test_Connection:
@@ -491,3 +491,47 @@ def test_heartbeat_tick(self):
     def test_server_capabilities(self):
         self.conn.server_properties['capabilities'] = {'foo': 1}
         assert self.conn.server_capabilities == {'foo': 1}
+
+    @pytest.mark.parametrize(
+        "method",
+        (
+            spec.Connection.Close,
+            spec.Connection.CloseOk,
+            spec.Connection.TuneOk
+        )
+    )
+    def test_dispatch_method(self, method):
+        mocked_method = 'amqp.connection.AbstractChannel.dispatch_method'
+        with patch(mocked_method) as dm_mock:
+            self.conn.dispatch_method(
+                method, sentinel.PAYLOAD, sentinel.CONTENT
+            )
+            dm_mock.assert_called_once_with(
+                method, sentinel.PAYLOAD, sentinel.CONTENT
+            )
+
+    @pytest.mark.parametrize(
+        "method",
+        (
+            spec.Connection.Close,
+            spec.Connection.CloseOk,
+            spec.Connection.TuneOk
+        )
+    )
+    def test_dispatch_method_closing_connection(self, method, caplog):
+        mocked_method = 'amqp.connection.AbstractChannel.dispatch_method'
+        self.conn.is_closing = True
+        with patch(mocked_method) as dm_mock:
+            self.conn.dispatch_method(
+                method, sentinel.PAYLOAD, sentinel.CONTENT
+            )
+            if method in (spec.Connection.Close, spec.Connection.CloseOk):
+                dm_mock.assert_called_once_with(
+                    method, sentinel.PAYLOAD, sentinel.CONTENT
+                )
+            else:
+                dm_mock.assert_not_called()
+                assert caplog.records[0].msg == \
+                    IGNORED_METHOD_DURING_CONN_CLOSE
+                assert caplog.records[0].args[0] == method
+                assert caplog.records[0].levelname == 'WARNING'