managing disconnections

DanielePalaia · DanielePalaia · commit 524b78a371a9 · 2025-01-28T09:55:37.000+01:00
diff --git a/examples/getting_started/reconnection_example.py b/examples/getting_started/reconnection_example.py
@@ -0,0 +1,214 @@
+# type: ignore
+
+
+import time
+
+from rabbitmq_amqp_python_client import (
+    AddressHelper,
+    AMQPMessagingHandler,
+    BindingSpecification,
+    Connection,
+    ConnectionClosed,
+    Event,
+    ExchangeSpecification,
+    Message,
+    QuorumQueueSpecification,
+)
+
+connection = None
+management = None
+publisher = None
+consumer = None
+
+
+def on_disconnected():
+
+    print("disconnected")
+    exchange_name = "test-exchange"
+    queue_name = "example-queue"
+    routing_key = "routing-key"
+
+    global connection
+    global management
+    global publisher
+    global consumer
+
+    addr = AddressHelper.exchange_address(exchange_name, routing_key)
+    addr_queue = AddressHelper.queue_address(queue_name)
+
+    connection = create_connection()
+    if management is not None:
+        management = connection.management()
+    if publisher is not None:
+        publisher = connection.publisher(addr)
+    if consumer is not None:
+        consumer = connection.consumer(addr_queue, handler=MyMessageHandler())
+
+
+class MyMessageHandler(AMQPMessagingHandler):
+
+    def __init__(self):
+        super().__init__()
+        self._count = 0
+
+    def on_message(self, event: Event):
+        print("received message: " + str(event.message.annotations))
+
+        # accepting
+        self.delivery_context.accept(event)
+
+        # in case of rejection (+eventually deadlettering)
+        # self.delivery_context.discard(event)
+
+        # in case of requeuing
+        # self.delivery_context.requeue(event)
+
+        # annotations = {}
+        # annotations[symbol('x-opt-string')] = 'x-test1'
+        # in case of requeuing with annotations added
+        # self.delivery_context.requeue_with_annotations(event, annotations)
+
+        # in case of rejection with annotations added
+        # self.delivery_context.discard_with_annotations(event)
+
+        print("count " + str(self._count))
+
+        self._count = self._count + 1
+
+        if self._count == 100:
+            print("closing receiver")
+            # if you want you can add cleanup operations here
+            # event.receiver.close()
+            # event.connection.close()
+
+    def on_connection_closed(self, event: Event):
+        # if you want you can add cleanup operations here
+        print("connection closed")
+
+    def on_link_closed(self, event: Event) -> None:
+        # if you want you can add cleanup operations here
+        print("link closed")
+
+
+def create_connection() -> Connection:
+    connection = Connection(
+        "amqp://guest:guest@localhost:5672/", on_disconnection_handler=on_disconnected
+    )
+    connection.dial()
+
+    return connection
+
+
+def main() -> None:
+
+    exchange_name = "test-exchange"
+    queue_name = "example-queue"
+    routing_key = "routing-key"
+    messages_to_publish = 50000
+
+    global connection
+    global management
+    global publisher
+    global consumer
+
+    print("connection to amqp server")
+    if connection is None:
+        connection = create_connection()
+
+    if management is None:
+        management = connection.management()
+
+    print("declaring exchange and queue")
+    management.declare_exchange(ExchangeSpecification(name=exchange_name, arguments={}))
+
+    management.declare_queue(
+        QuorumQueueSpecification(name=queue_name)
+        # QuorumQueueSpecification(name=queue_name, dead_letter_exchange="dead-letter")
+    )
+
+    print("binding queue to exchange")
+    bind_name = management.bind(
+        BindingSpecification(
+            source_exchange=exchange_name,
+            destination_queue=queue_name,
+            binding_key=routing_key,
+        )
+    )
+
+    addr = AddressHelper.exchange_address(exchange_name, routing_key)
+
+    addr_queue = AddressHelper.queue_address(queue_name)
+
+    print("create a publisher and publish a test message")
+    if publisher is None:
+        publisher = connection.publisher(addr)
+
+    print("purging the queue")
+    messages_purged = management.purge_queue(queue_name)
+
+    print("messages purged: " + str(messages_purged))
+    # management.close()
+
+    # publish 10 messages
+    while True:
+        for i in range(messages_to_publish):
+
+            if i % 1000 == 0:
+                print("publishing")
+            try:
+                publisher.publish(Message(body="test"))
+            except ConnectionClosed:
+                print("publisher closing exception, resubmitting")
+                continue
+
+        print("closing")
+        try:
+            publisher.close()
+        except ConnectionClosed:
+            print("publisher closing exception, resubmitting")
+            continue
+        break
+
+    print(
+        "create a consumer and consume the test message - press control + c to terminate to consume"
+    )
+    if consumer is None:
+        consumer = connection.consumer(addr_queue, handler=MyMessageHandler())
+
+    while True:
+        try:
+            consumer.run()
+        except KeyboardInterrupt:
+            pass
+        except ConnectionClosed:
+            time.sleep(1)
+            continue
+        except Exception as e:
+            print("consumer exited for exception " + str(e))
+
+        break
+
+    print("cleanup")
+    consumer.close()
+    # once we finish consuming if we close the connection we need to create a new one
+    # connection = create_connection()
+    # management = connection.management()
+
+    print("unbind")
+    management.unbind(bind_name)
+
+    print("delete queue")
+    management.delete_queue(queue_name)
+
+    print("delete exchange")
+    management.delete_exchange(exchange_name)
+
+    print("closing connections")
+    management.close()
+    print("after management closing")
+    connection.close()
+    print("after connection closing")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/rabbitmq_amqp_python_client/__init__.py b/rabbitmq_amqp_python_client/__init__.py
@@ -16,6 +16,7 @@
 from .qpid.proton._delivery import Delivery
 from .qpid.proton._events import Event
 from .qpid.proton._message import Message
+from .qpid.proton._utils import ConnectionClosed
 from .qpid.proton.handlers import MessagingHandler
 from .queues import (
     ClassicQueueSpecification,
@@ -53,4 +54,5 @@
     "AMQPMessagingHandler",
     "ArgumentOutOfRangeException",
     "Delivery",
+    "ConnectionClosed",
 ]
diff --git a/rabbitmq_amqp_python_client/connection.py b/rabbitmq_amqp_python_client/connection.py
@@ -1,5 +1,5 @@
 import logging
-from typing import Optional
+from typing import Annotated, Callable, Optional, TypeVar
 
 from .address_helper import validate_address
 from .consumer import Consumer
@@ -11,16 +11,22 @@
 
 logger = logging.getLogger(__name__)
 
+MT = TypeVar("MT")
+CB = Annotated[Callable[[MT], None], "Message callback type"]
+
 
 class Connection:
-    def __init__(self, addr: str):
+    def __init__(self, addr: str, on_disconnection_handler: Optional[CB] = None):  # type: ignore
         self._addr: str = addr
         self._conn: BlockingConnection
         self._management: Management
+        self._on_disconnection_handler = on_disconnection_handler
 
     def dial(self) -> None:
         logger.debug("Establishing a connection to the amqp server")
-        self._conn = BlockingConnection(self._addr)
+        self._conn = BlockingConnection(
+            self._addr, on_disconnection_handler=self._on_disconnection_handler
+        )
         self._open()
         logger.debug("Connection to the server established")
 
diff --git a/rabbitmq_amqp_python_client/qpid/proton/_utils.py b/rabbitmq_amqp_python_client/qpid/proton/_utils.py
@@ -71,6 +71,11 @@ class Literal(metaclass=GenericMeta):
     )
     from ._transport import SSLDomain
 
+from typing import Annotated, TypeVar
+
+MT = TypeVar("MT")
+CB = Annotated[Callable[[MT], None], "Message callback type"]
+
 
 class BlockingLink:
     def __init__(
@@ -423,6 +428,7 @@ def __init__(
         heartbeat: Optional[float] = None,
         urls: Optional[List[str]] = None,
         reconnect: Union[None, Literal[False], "Backoff"] = None,
+        on_disconnection_handler: Optional[CB] = None,
         **kwargs
     ) -> None:
         self.disconnected = False
@@ -432,6 +438,7 @@ def __init__(
         self.container.start()
         self.conn = None
         self.closing = False
+        self._on_disconnection_handler = on_disconnection_handler
         # Preserve previous behaviour if neither reconnect nor urls are supplied
         if url is not None and urls is None and reconnect is None:
             reconnect = False
@@ -631,7 +638,10 @@ def on_connection_remote_close(self, event: "Event") -> None:
         """
         Event callback for when the link peer closes the connection.
         """
+        print("connection remote closed")
         if event.connection.state & Endpoint.LOCAL_ACTIVE:
+            if self._on_disconnection_handler is not None:
+                event.container.schedule(0, self._on_disconnection_handler())
             event.connection.close()
             if not self.closing:
                 raise ConnectionClosed(event.connection)
