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

Author: matusvalo

amqp/channel.py

@@ -108,6 +108,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 +160,18 @@ 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
+            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

@@ -233,6 +233,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 +293,18 @@ 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
+            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 +589,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()