reconnecting managements, publishers and consumers

Author: DanielePalaia

examples/reconnection/reconnection_example.py

@@ -129,7 +129,7 @@ def main() -> None:
                     publisher.publish(Message(body="test"))
             except ConnectionClosed:
                 print("publisher closing exception, resubmitting")
-                publisher = connection.publisher(addr)
+                # publisher = connection.publisher(addr)
                 continue
 
         print("closing publisher")
@@ -153,9 +153,6 @@ def main() -> None:
         except KeyboardInterrupt:
             pass
         except ConnectionClosed:
-            consumer = connection.consumer(
-                addr_queue, message_handler=MyMessageHandler()
-            )
             continue
         except Exception as e:
             print("consumer exited for exception " + str(e))
@@ -165,7 +162,6 @@ def main() -> None:
     print("cleanup")
     consumer.close()
 
-    management = connection.management()
     print("unbind")
     management.unbind(bind_name)
 

rabbitmq_amqp_python_client/connection.py

@@ -62,7 +62,7 @@ def __init__(
              uri: Single node connection URI
              uris: List of URIs for multi-node setup
              ssl_context: SSL configuration for secure connections
-             on_disconnection_handler: Callback for handling disconnection events
+             reconnect: Ability to automatically reconnect in case of disconnections from the server
 
         Raises:
              ValueError: If neither uri nor uris is provided
@@ -76,10 +76,10 @@ def __init__(
         self._addr: Optional[str] = uri
         self._addrs: Optional[list[str]] = uris
         self._conn: BlockingConnection
-        self._management: Management
         self._conf_ssl_context: Union[
             PosixSslConfigurationContext, WinSslConfigurationContext, None
         ] = ssl_context
+        self._managements: list[Management] = []
         self._reconnect = reconnect
         self._ssl_domain = None
         self._connections = []  # type: ignore
@@ -158,7 +158,7 @@ def dial(self) -> None:
                 on_disconnection_handler=self._on_disconnection,
             )
 
-        self._open()
+        # self._open()
         logger.debug("Connection to the server established")
 
     def _win_store_to_cert(
@@ -185,7 +185,12 @@ def management(self) -> Management:
         Returns:
             Management: The management interface for performing administrative tasks
         """
-        return self._management
+        if len(self._managements) == 0:
+            management = Management(self._conn)
+            management.open()
+            self._managements.append(management)
+
+        return self._managements[0]
 
     # closes the connection to the AMQP 1.0 server.
     def close(self) -> None:
@@ -196,9 +201,9 @@ def close(self) -> None:
         """
         logger.debug("Closing connection")
         try:
-            for publisher in self._publishers:
+            for publisher in self._publishers[:]:
                 publisher.close()
-            for consumer in self._consumers:
+            for consumer in self._consumers[:]:
                 consumer.close()
             self._conn.close()
         except Exception as e:
@@ -228,8 +233,9 @@ def publisher(self, destination: str = "") -> Publisher:
                     "destination address must start with /queues or /exchanges"
                 )
         publisher = Publisher(self._conn, destination)
+        publisher._set_publishers_list(self._publishers)
         self._publishers.append(publisher)
-        return self._publishers[self._publishers.index(publisher)]
+        return publisher
 
     def consumer(
         self,
@@ -265,26 +271,36 @@ def consumer(
 
     def _on_disconnection(self) -> None:
 
-        print("disconnected")
-
         if self in self._connections:
             self._connections.remove(self)
 
-        print("reconnecting")
         self._conn = BlockingConnection(
             url=self._addr,
             urls=self._addrs,
             ssl_domain=self._ssl_domain,
             on_disconnection_handler=self._on_disconnection,
         )
-        self._open()
+
         self._connections.append(self)
 
-        for index, publisher in enumerate(self._publishers):
-            # publisher = self._publishers.pop(index)
-            # address = publisher.address
-            self._publishers.remove(publisher)
-            # self._publishers.insert(index, Publisher(self._conn, address))
+        for i, management in enumerate(self._managements):
+            # Update the broken connection and sender in the management
+            self._managements[i]._update_connection(self._conn)
+
+        for i, publisher in enumerate(self._publishers):
+            # Update the broken connection and sender in the publisher
+            self._publishers[i]._update_connection(self._conn)
 
         for i, consumer in enumerate(self._consumers):
-            self._consumers.remove(consumer)
+            # Update the broken connection and sender in the consumer
+            self._consumers[i]._update_connection(self._conn)
+
+    @property
+    def active_producers(self) -> int:
+        """Returns the number of active publishers"""
+        return len(self._publishers)
+
+    @property
+    def active_consumers(self) -> int:
+        """Returns the number of active consumers"""
+        return len(self._consumers)

rabbitmq_amqp_python_client/consumer.py

@@ -65,6 +65,10 @@ def _open(self) -> None:
             logger.debug("Creating Sender")
             self._receiver = self._create_receiver(self._addr)
 
+    def _update_connection(self, conn: BlockingConnection) -> None:
+        self._conn = conn
+        self._receiver = self._create_receiver(self._addr)
+
     def _set_consumers_list(self, consumers: []) -> None:  # type: ignore
         self._consumers = consumers
 

rabbitmq_amqp_python_client/management.py

@@ -53,6 +53,13 @@ def __init__(self, conn: BlockingConnection):
         self._receiver: Optional[BlockingReceiver] = None
         self._conn = conn
 
+    def _update_connection(self, conn: BlockingConnection) -> None:
+        self._conn = conn
+        self._sender = self._create_sender(CommonValues.management_node_address.value)
+        self._receiver = self._create_receiver(
+            CommonValues.management_node_address.value,
+        )
+
     def open(self) -> None:
         """
         Open the management connection by creating sender and receiver.

rabbitmq_amqp_python_client/publisher.py

@@ -46,6 +46,10 @@ def __init__(self, conn: BlockingConnection, addr: str = ""):
         self._publishers: list[Publisher] = []
         self._open()
 
+    def _update_connection(self, conn: BlockingConnection) -> None:
+        self._conn = conn
+        self._sender = self._create_sender(self._addr)
+
     def _open(self) -> None:
         if self._sender is None:
             logger.debug("Creating Sender")

tests/test_connection.py

@@ -81,9 +81,11 @@ def test_connection_reconnection() -> None:
     # delay
     time.sleep(5)
     # simulate a disconnection
-    delete_all_connections()
     # raise a reconnection
     management = connection.management()
+
+    delete_all_connections()
+
     stream_name = "test_stream_info_with_validation"
     queue_specification = StreamSpecification(
         name=stream_name,
@@ -95,7 +97,6 @@ def test_connection_reconnection() -> None:
         disconnected = True
 
     # check that we reconnected
-    management = connection.management()
     management.declare_queue(queue_specification)
     management.delete_queue(stream_name)
     management.close()

tests/test_publisher.py

@@ -270,13 +270,11 @@ def test_disconnection_reconnection() -> None:
     # delay
     time.sleep(5)
     messages_to_publish = 10000
-    queue_name = "test-queue"
+    queue_name = "test-queue-reconnection"
     management = connection_test.management()
 
     management.declare_queue(QuorumQueueSpecification(name=queue_name))
 
-    management.close()
-
     publisher = connection_test.publisher(
         destination=AddressHelper.queue_address(queue_name)
     )
@@ -291,9 +289,6 @@ def test_disconnection_reconnection() -> None:
 
             except ConnectionClosed:
                 disconnected = True
-                publisher = connection_test.publisher(
-                    destination=AddressHelper.queue_address(queue_name)
-                )
                 continue
 
             except Exception:
@@ -303,14 +298,6 @@ def test_disconnection_reconnection() -> None:
 
     publisher.close()
 
-    # cleanup, we need to create a new connection as the previous one
-    # was closed by the test
-
-    connection_test = environment.connection()
-    connection_test.dial()
-
-    management = connection_test.management()
-
     # purge the queue and check number of published messages
     message_purged = management.purge_queue(queue_name)
 
@@ -335,14 +322,10 @@ def test_queue_info_for_stream_with_validations(connection: Connection) -> None:
     management = connection.management()
     management.declare_queue(queue_specification)
 
-    print("before creating publisher")
-
     publisher = connection.publisher(
         destination=AddressHelper.queue_address(stream_name)
     )
 
-    print("after creating publisher")
-
     for i in range(messages_to_send):
 
         publisher.publish(Message(body="test"))
@@ -392,3 +375,55 @@ def test_publish_per_message_exchange(connection: Connection) -> None:
     assert accepted_2 is True
     assert purged_messages_queue == 2
     assert raised is False
+
+
+def test_multiple_publishers(environment: Environment) -> None:
+
+    stream_name = "test_multiple_publisher_1"
+    stream_name_2 = "test_multiple_publisher_2"
+    connection = environment.connection()
+    connection.dial()
+
+    stream_specification = StreamSpecification(
+        name=stream_name,
+    )
+    management = connection.management()
+    management.declare_queue(stream_specification)
+
+    stream_specification = StreamSpecification(
+        name=stream_name_2,
+    )
+    management.declare_queue(stream_specification)
+
+    destination = AddressHelper.queue_address(stream_name)
+    destination_2 = AddressHelper.queue_address(stream_name_2)
+    connection.publisher(destination)
+
+    assert connection.active_producers == 1
+
+    publisher_2 = connection.publisher(destination_2)
+
+    assert connection.active_producers == 2
+
+    publisher_2.close()
+
+    assert connection.active_producers == 1
+
+    connection.publisher(destination_2)
+
+    assert connection.active_producers == 2
+
+    connection.close()
+
+    assert connection.active_producers == 0
+
+    # cleanup
+    connection = environment.connection()
+    connection.dial()
+    management = connection.management()
+
+    management.delete_queue(stream_name)
+
+    management.delete_queue(stream_name_2)
+
+    management.close()