supporting reconnection of streams by last consumed offset

Author: DanielePalaia

examples/streams/example_with_streams.py

@@ -4,6 +4,7 @@
     AddressHelper,
     AMQPMessagingHandler,
     Connection,
+    ConnectionClosed,
     Environment,
     Event,
     Message,
@@ -12,7 +13,7 @@
     StreamSpecification,
 )
 
-MESSAGES_TO_PUBLISH = 100
+MESSAGES_TO_PUBLISH = 1
 
 
 class MyMessageHandler(AMQPMessagingHandler):
@@ -21,7 +22,7 @@ def __init__(self):
         super().__init__()
         self._count = 0
 
-    def on_message(self, event: Event):
+    def on_amqp_message(self, event: Event):
         # just messages with banana filters get received
         print(
             "received message from stream: "
@@ -86,7 +87,7 @@ def main() -> None:
     queue_name = "example-queue"
 
     print("connection to amqp server")
-    environment = Environment("amqp://guest:guest@localhost:5672/")
+    environment = Environment("amqp://guest:guest@localhost:5672/", reconnect=True)
     connection = create_connection(environment)
 
     management = connection.management()
@@ -134,14 +135,22 @@ def main() -> None:
 
     publisher.close()
 
-    try:
-        consumer.run()
-    except KeyboardInterrupt:
-        pass
+    while True:
+        try:
+            consumer.run()
+        except KeyboardInterrupt:
+            pass
+        except ConnectionClosed:
+            print("connection closed")
+            continue
+        except Exception as e:
+            print("consumer exited for exception " + str(e))
+
+        break
 
     #
     print("delete queue")
-    management.delete_queue(queue_name)
+    # management.delete_queue(queue_name)
 
     print("closing connections")
     management.close()

rabbitmq_amqp_python_client/amqp_consumer_handler.py

@@ -21,10 +21,18 @@ 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()
+        self._offset = 0
 
     def on_amqp_message(self, event: Event) -> None:
         pass
 
     def on_message(self, event: Event) -> None:
         print("first level callback")
+        if "x-stream-offset" in event.message.annotations:
+            print("setting offset")
+            self._offset = int(event.message.annotations["x-stream-offset"])
         self.on_amqp_message(event)
+
+    @property
+    def offset(self) -> int:
+        return self._offset

rabbitmq_amqp_python_client/connection.py

@@ -292,6 +292,7 @@ def _on_disconnection(self) -> None:
             self._publishers[i]._update_connection(self._conn)
 
         for i, consumer in enumerate(self._consumers):
+            print("reconnecting consumer")
             # Update the broken connection and sender in the consumer
             self._consumers[i]._update_connection(self._conn)
 

rabbitmq_amqp_python_client/consumer.py

@@ -1,12 +1,12 @@
 import logging
 from typing import Literal, Optional, Union, cast
 
+from .amqp_consumer_handler import AMQPMessagingHandler
 from .entities import StreamOptions
 from .options import (
     ReceiverOptionUnsettled,
     ReceiverOptionUnsettledWithFilters,
 )
-from .qpid.proton._handlers import MessagingHandler
 from .qpid.proton._message import Message
 from .qpid.proton.utils import (
     BlockingConnection,
@@ -37,7 +37,7 @@ def __init__(
         self,
         conn: BlockingConnection,
         addr: str,
-        handler: Optional[MessagingHandler] = None,
+        handler: Optional[AMQPMessagingHandler] = None,
         stream_options: Optional[StreamOptions] = None,
         credit: Optional[int] = None,
     ):
@@ -67,7 +67,23 @@ def _open(self) -> None:
 
     def _update_connection(self, conn: BlockingConnection) -> None:
         self._conn = conn
-        self._receiver = self._create_receiver(self._addr)
+        if self._stream_options is None:
+            print("creating new receiver without stream")
+            self._receiver = self._conn.create_receiver(
+                self._addr,
+                options=ReceiverOptionUnsettled(self._addr),
+                handler=self._handler,
+            )
+        else:
+            print("creating new stream receiver")
+            self._stream_options.offset(self._handler.offset - 1)  # type: ignore
+            self._receiver = self._conn.create_receiver(
+                self._addr,
+                options=ReceiverOptionUnsettledWithFilters(
+                    self._addr, self._stream_options
+                ),
+                handler=self._handler,
+            )
 
     def _set_consumers_list(self, consumers: []) -> None:  # type: ignore
         self._consumers = consumers
@@ -149,3 +165,7 @@ def _create_receiver(self, addr: str) -> BlockingReceiver:
     def address(self) -> str:
         """Get the current publisher address."""
         return self._addr
+
+    @property
+    def handler(self) -> Optional[AMQPMessagingHandler]:
+        return self._handler

tests/conftest.py

@@ -20,6 +20,8 @@
     FriendlyName,
 )
 
+from .http_requests import delete_all_connections
+
 os.chdir(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 
@@ -150,6 +152,22 @@ def on_message(self, event: Event):
             raise ConsumerTestException("consumed")
 
 
+class MyMessageHandlerAcceptStreamOffsetReconnect(AMQPMessagingHandler):
+
+    def __init__(self, starting_offset: Optional[int] = None):
+        super().__init__()
+        self._received = 0
+        self._starting_offset = starting_offset
+
+    def on_message(self, event: Event):
+        if self._received == 5:
+            delete_all_connections()
+        self.delivery_context.accept(event)
+        self._received = self._received + 1
+        if self._received == 10:
+            raise ConsumerTestException("consumed")
+
+
 class MyMessageHandlerNoack(AMQPMessagingHandler):
 
     def __init__(self):

tests/test_streams.py

@@ -10,6 +10,7 @@
 from .conftest import (
     ConsumerTestException,
     MyMessageHandlerAcceptStreamOffset,
+    MyMessageHandlerAcceptStreamOffsetReconnect,
 )
 from .utils import publish_messages
 
@@ -370,3 +371,42 @@ def test_stream_match_unfiltered(
     consumer.close()
 
     management.delete_queue(stream_name)
+
+
+def test_stream_reconnection(connection: Connection, environment: Environment) -> None:
+
+    consumer = None
+    stream_name = "test_stream_info_with_filtering"
+    messages_to_send = 10
+
+    queue_specification = StreamSpecification(
+        name=stream_name,
+    )
+    management = connection.management()
+    management.declare_queue(queue_specification)
+
+    addr_queue = AddressHelper.queue_address(stream_name)
+
+    # consume and then publish
+    try:
+        stream_filter_options = StreamOptions()
+        stream_filter_options.filter_values(["banana"])
+        stream_filter_options.filter_match_unfiltered(True)
+        connection_consumer = environment.connection()
+        connection_consumer.dial()
+        consumer = connection_consumer.consumer(
+            addr_queue,
+            # disconnection and check happens here
+            message_handler=MyMessageHandlerAcceptStreamOffsetReconnect(),
+            stream_filter_options=stream_filter_options,
+        )
+        # send with annotations filter banana
+        publish_messages(connection, messages_to_send, stream_name)
+        consumer.run()
+    # ack to terminate the consumer
+    except ConsumerTestException:
+        pass
+
+    consumer.close()
+
+    management.delete_queue(stream_name)