improving disconnection management

Author: DanielePalaia

examples/getting_started/getting_started.py

@@ -23,7 +23,7 @@ def __init__(self):
         super().__init__()
         self._count = 0
 
-    def on_message(self, event: Event):
+    def on_amqp_message(self, event: Event):
         print("received message: " + str(event.message.body))
 
         # accepting
@@ -147,7 +147,7 @@ def main() -> None:
     consumer.close()
     # once we finish consuming if we close the connection we need to create a new one
     # connection = create_connection()
-    # management = connection.management()
+    management = connection.management()
 
     print("unbind")
     management.unbind(bind_name)

examples/reconnection/reconnection_example.py

@@ -1,76 +1,22 @@
 # type: ignore
-
-
-import time
-from dataclasses import dataclass
-from typing import Optional
-
 from rabbitmq_amqp_python_client import (
     AddressHelper,
     AMQPMessagingHandler,
     Connection,
     ConnectionClosed,
-    Consumer,
     Environment,
     Event,
     ExchangeSpecification,
     ExchangeToQueueBindingSpecification,
-    Management,
     Message,
-    Publisher,
     QuorumQueueSpecification,
 )
 
-
 # here we keep track of the objects we need to reconnect
-@dataclass
-class ConnectionConfiguration:
-    connection: Optional[Connection] = None
-    management: Optional[Management] = None
-    publisher: Optional[Publisher] = None
-    consumer: Optional[Consumer] = None
-
+MESSAGES_TO_PUBLISH = 50000
 
-connection_configuration = ConnectionConfiguration()
-MESSAGES_TO_PUBLSH = 50000
 
-
-# disconnection callback
-# here you can cleanup or reconnect
-def on_disconnection():
-
-    print("disconnected")
-    global environment
-    exchange_name = "test-exchange"
-    queue_name = "example-queue"
-    routing_key = "routing-key"
-
-    global connection_configuration
-
-    addr = AddressHelper.exchange_address(exchange_name, routing_key)
-    addr_queue = AddressHelper.queue_address(queue_name)
-
-    if connection_configuration.connection is not None:
-        connection_configuration.connection = create_connection()
-    if connection_configuration.management is not None:
-        connection_configuration.management = (
-            connection_configuration.connection.management()
-        )
-    if connection_configuration.publisher is not None:
-        connection_configuration.publisher = (
-            connection_configuration.connection.publisher(addr)
-        )
-    if connection_configuration.consumer is not None:
-        connection_configuration.consumer = (
-            connection_configuration.connection.consumer(
-                addr_queue, message_handler=MyMessageHandler()
-            )
-        )
-
-
-environment = Environment(
-    uri="amqp://guest:guest@localhost:5672/", on_disconnection_handler=on_disconnection
-)
+environment = Environment(uri="amqp://guest:guest@localhost:5672/", reconnect=True)
 
 
 class MyMessageHandler(AMQPMessagingHandler):
@@ -102,7 +48,7 @@ def on_message(self, event: Event):
 
         self._count = self._count + 1
 
-        if self._count == MESSAGES_TO_PUBLSH:
+        if self._count == MESSAGES_TO_PUBLISH:
             print("closing receiver")
             # if you want you can add cleanup operations here
 
@@ -136,29 +82,22 @@ def main() -> None:
     queue_name = "example-queue"
     routing_key = "routing-key"
 
-    global connection_configuration
-
     print("connection to amqp server")
-    if connection_configuration.connection is None:
-        connection_configuration.connection = create_connection()
-
-    if connection_configuration.management is None:
-        connection_configuration.management = (
-            connection_configuration.connection.management()
-        )
+    connection = create_connection()
+    management = connection.management()
+    publisher = None
+    consumer = None
 
     print("declaring exchange and queue")
-    connection_configuration.management.declare_exchange(
-        ExchangeSpecification(name=exchange_name)
-    )
+    management.declare_exchange(ExchangeSpecification(name=exchange_name))
 
-    connection_configuration.management.declare_queue(
+    management.declare_queue(
         QuorumQueueSpecification(name=queue_name)
         # QuorumQueueSpecification(name=queue_name, dead_letter_exchange="dead-letter")
     )
 
     print("binding queue to exchange")
-    bind_name = connection_configuration.management.bind(
+    bind_name = management.bind(
         ExchangeToQueueBindingSpecification(
             source_exchange=exchange_name,
             destination_queue=queue_name,
@@ -171,34 +110,32 @@ def main() -> None:
     addr_queue = AddressHelper.queue_address(queue_name)
 
     print("create a publisher and publish a test message")
-    if connection_configuration.publisher is None:
-        connection_configuration.publisher = (
-            connection_configuration.connection.publisher(addr)
-        )
+    if publisher is None:
+        publisher = connection.publisher(addr)
 
     print("purging the queue")
-    messages_purged = connection_configuration.management.purge_queue(queue_name)
+    messages_purged = management.purge_queue(queue_name)
 
     print("messages purged: " + str(messages_purged))
-    # management.close()
 
     # publishing messages
     while True:
-        for i in range(MESSAGES_TO_PUBLSH):
+        for i in range(MESSAGES_TO_PUBLISH):
 
             if i % 1000 == 0:
                 print("published 1000 messages...")
             try:
-                if connection_configuration.publisher is not None:
-                    connection_configuration.publisher.publish(Message(body="test"))
+                if publisher is not None:
+                    publisher.publish(Message(body="test"))
             except ConnectionClosed:
                 print("publisher closing exception, resubmitting")
+                publisher = connection.publisher(addr)
                 continue
 
         print("closing publisher")
         try:
-            if connection_configuration.publisher is not None:
-                connection_configuration.publisher.close()
+            if publisher is not None:
+                publisher.close()
         except ConnectionClosed:
             print("publisher closing exception, resubmitting")
             continue
@@ -207,43 +144,39 @@ def main() -> None:
     print(
         "create a consumer and consume the test message - press control + c to terminate to consume"
     )
-    if connection_configuration.consumer is None:
-        connection_configuration.consumer = (
-            connection_configuration.connection.consumer(
-                addr_queue, message_handler=MyMessageHandler()
-            )
-        )
+    if consumer is None:
+        consumer = connection.consumer(addr_queue, message_handler=MyMessageHandler())
 
     while True:
         try:
-            connection_configuration.consumer.run()
+            consumer.run()
         except KeyboardInterrupt:
             pass
         except ConnectionClosed:
-            time.sleep(1)
+            consumer = connection.consumer(
+                addr_queue, message_handler=MyMessageHandler()
+            )
             continue
         except Exception as e:
             print("consumer exited for exception " + str(e))
 
         break
 
     print("cleanup")
-    connection_configuration.consumer.close()
-    # once we finish consuming if we close the connection we need to create a new one
-    # connection = create_connection()
-    # management = connection.management()
+    consumer.close()
 
+    management = connection.management()
     print("unbind")
-    connection_configuration.management.unbind(bind_name)
+    management.unbind(bind_name)
 
     print("delete queue")
-    connection_configuration.management.delete_queue(queue_name)
+    management.delete_queue(queue_name)
 
     print("delete exchange")
-    connection_configuration.management.delete_exchange(exchange_name)
+    management.delete_exchange(exchange_name)
 
     print("closing connections")
-    connection_configuration.management.close()
+    management.close()
     print("after management closing")
     environment.close()
     print("after connection closing")

rabbitmq_amqp_python_client/amqp_consumer_handler.py

@@ -1,4 +1,5 @@
 from .delivery_context import DeliveryContext
+from .qpid.proton._events import Event
 from .qpid.proton.handlers import MessagingHandler
 
 """
@@ -20,3 +21,10 @@ def __init__(self, auto_accept: bool = False, auto_settle: bool = True):
         """
         super().__init__(auto_accept=auto_accept, auto_settle=auto_settle)
         self.delivery_context: DeliveryContext = DeliveryContext()
+
+    def on_amqp_message(self, event: Event) -> None:
+        pass
+
+    def on_message(self, event: Event) -> None:
+        print("first level callback")
+        self.on_amqp_message(event)

rabbitmq_amqp_python_client/connection.py

@@ -35,7 +35,7 @@ def __init__(
         # multi-node mode
         uris: Optional[list[str]] = None,
         ssl_context: Optional[SslConfigurationContext] = None,
-        on_disconnection_handler: Optional[CB] = None,  # type: ignore
+        reconnect: bool = False,
     ):
         """
          Initialize a new Connection instance.
@@ -59,11 +59,13 @@ def __init__(
         self._addrs: Optional[list[str]] = uris
         self._conn: BlockingConnection
         self._management: Management
-        self._on_disconnection_handler = on_disconnection_handler
+        self._reconnect = reconnect
         self._conf_ssl_context: Optional[SslConfigurationContext] = ssl_context
         self._ssl_domain = None
         self._connections = []  # type: ignore
         self._index: int = -1
+        self._publishers: list[Publisher] = []
+        self._consumers: list[Consumer] = []
 
     def _set_environment_connection_list(self, connections: []):  # type: ignore
         self._connections = connections
@@ -91,12 +93,21 @@ def dial(self) -> None:
                         self._conf_ssl_context.client_cert.client_key,
                         self._conf_ssl_context.client_cert.password,
                     )
-        self._conn = BlockingConnection(
-            url=self._addr,
-            urls=self._addrs,
-            ssl_domain=self._ssl_domain,
-            on_disconnection_handler=self._on_disconnection_handler,
-        )
+
+        if self._reconnect is False:
+            self._conn = BlockingConnection(
+                url=self._addr,
+                urls=self._addrs,
+                ssl_domain=self._ssl_domain,
+            )
+        else:
+            self._conn = BlockingConnection(
+                url=self._addr,
+                urls=self._addrs,
+                ssl_domain=self._ssl_domain,
+                on_disconnection_handler=self._on_disconnection,
+            )
+
         self._open()
         logger.debug("Connection to the server established")
 
@@ -122,6 +133,10 @@ def close(self) -> None:
         """
         logger.debug("Closing connection")
         try:
+            for publisher in self._publishers:
+                publisher.close()
+            for consumer in self._consumers:
+                consumer.close()
             self._conn.close()
         except Exception as e:
             logger.error(f"Error closing connection: {e}")
@@ -150,7 +165,8 @@ def publisher(self, destination: str = "") -> Publisher:
                     "destination address must start with /queues or /exchanges"
                 )
         publisher = Publisher(self._conn, destination)
-        return publisher
+        self._publishers.append(publisher)
+        return self._publishers[self._publishers.index(publisher)]
 
     def consumer(
         self,
@@ -181,4 +197,31 @@ def consumer(
         consumer = Consumer(
             self._conn, destination, message_handler, stream_filter_options, credit
         )
+        self._consumers.append(consumer)
         return 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, consumer in enumerate(self._consumers):
+            self._consumers.remove(consumer)