[Backport release-1.19] Return the next message after a streaming subscription reconnects

dapr/aio/clients/grpc/subscription.py

@@ -1,17 +1,52 @@
 import asyncio
+import contextlib
+import logging
+import random
+from typing import Optional
 
 from grpc import StatusCode  # type: ignore[attr-defined]
-from grpc.aio import AioRpcError
+from grpc.aio import EOF, AioRpcError  # type: ignore[attr-defined]
 
 from dapr.aio.clients.health import DaprHealth
 from dapr.clients.grpc._response import TopicEventResponse
 from dapr.common.pubsub.subscription import (
     StreamCancelledError,
     StreamInactiveError,
+    StreamReconnectError,
     SubscriptionMessage,
 )
 from dapr.proto import api_v1, appcallback_v1
 
+logger = logging.getLogger(__name__)
+
+MAX_RECONNECT_ATTEMPTS = 5
+RECONNECT_BACKOFF_INITIAL_SECONDS = 0.2
+RECONNECT_BACKOFF_MAX_SECONDS = 5.0
+_TRANSIENT_STATUS_CODES = (StatusCode.UNAVAILABLE, StatusCode.UNKNOWN, StatusCode.INTERNAL)
+
+
+class _TransientStreamError(Exception):
+    """The stream failed in a way a reconnect can fix."""
+
+
+def _current_task_is_cancelling() -> bool:
+    """Tell whether the running task itself was asked to cancel.
+
+    Task.cancelling() only exists on Python 3.11+. On 3.10 this returns False, so a
+    cancellation that coincides with close() is reported as StreamInactiveError there.
+    """
+    current_task = asyncio.current_task()
+    cancelling = getattr(current_task, 'cancelling', None)
+    if cancelling is None:
+        return False
+    return cancelling() > 0
+
+
+def _reconnect_backoff_seconds(reconnect_attempts: int) -> float:
+    backoff_exponential = RECONNECT_BACKOFF_INITIAL_SECONDS * 2 ** (reconnect_attempts - 1)
+    backoff_capped = min(RECONNECT_BACKOFF_MAX_SECONDS, backoff_exponential)
+    return random.uniform(backoff_capped / 2, backoff_capped)
+
 
 class Subscription:
     def __init__(self, stub, pubsub_name, topic, metadata=None, dead_letter_topic=None):
@@ -21,10 +56,19 @@ def __init__(self, stub, pubsub_name, topic, metadata=None, dead_letter_topic=No
         self._metadata = metadata or {}
         self._dead_letter_topic = dead_letter_topic or ''
         self._stream = None
-        self._send_queue: asyncio.Queue[api_v1.SubscribeTopicEventsRequestAlpha1] = asyncio.Queue()
+        self._send_queue: asyncio.Queue[Optional[api_v1.SubscribeTopicEventsRequestAlpha1]] = (
+            asyncio.Queue()
+        )
         self._stream_active = asyncio.Event()
+        self._closed = asyncio.Event()
 
     async def start(self):
+        if self._closed.is_set():
+            raise StreamInactiveError('Stream is not active')
+        send_queue: asyncio.Queue[Optional[api_v1.SubscribeTopicEventsRequestAlpha1]] = (
+            asyncio.Queue()
+        )
+
         async def outgoing_request_iterator():
             try:
                 initial_request = api_v1.SubscribeTopicEventsRequestAlpha1(
@@ -39,49 +83,108 @@ async def outgoing_request_iterator():
 
                 while self._stream_active.is_set():
                     try:
-                        response = await asyncio.wait_for(self._send_queue.get(), timeout=1.0)
+                        response = await asyncio.wait_for(send_queue.get(), timeout=1.0)
+                        if response is None:
+                            return
                         yield response
                     except asyncio.TimeoutError:
                         continue
             except Exception as e:
                 raise Exception(f'Error while writing to stream: {e}')
 
+        # The stream gets its own send queue, so the request iterator of a stream that already
+        # failed can't take its acks.
+        self._send_queue = send_queue
         self._stream = self._stub.SubscribeTopicEventsAlpha1(outgoing_request_iterator())
         self._stream_active.set()
-        await self._stream.read()  # discard the initial message
+        try:
+            await self._stream.read()  # discard the initial message
+        except (Exception, asyncio.CancelledError):
+            self._stream_active.clear()
+            send_queue.put_nowait(None)
+            if _current_task_is_cancelling():
+                raise
+            if self._closed.is_set():
+                raise StreamInactiveError('Stream is not active')
+            raise
 
     async def reconnect_stream(self):
-        await self.close()
+        await self._close_stream()
         await DaprHealth.wait_for_sidecar()
-        print('Attempting to reconnect...')
+        logger.info('Subscription stream reconnecting...')
         await self.start()
 
     async def next_message(self):
-        if not self._stream_active.is_set():
+        """Get the next message from the stream.
+
+        If the stream fails with a transient error or the server ends it, the stream is
+        reconnected and read again, up to MAX_RECONNECT_ATTEMPTS times per call, with a jittered
+        exponential backoff between reconnects.
+        A reconnect whose new stream fails with a transient error counts as one attempt.
+
+        Returns:
+            SubscriptionMessage: The next message.
+
+        Raises:
+            StreamInactiveError: If the subscription is closed.
+            StreamCancelledError: If the stream was cancelled.
+            StreamReconnectError: If the stream still fails after the reconnects.
+            Exception: On any other error.
+        """
+        reconnect_attempts = 0
+        while True:
+            try:
+                if reconnect_attempts > 0:
+                    await self.reconnect_stream()
+                return await self._read_next_message()
+            except _TransientStreamError as e:
+                stream_error = str(e)
+            except AioRpcError as e:
+                # Raised by reconnect_stream() when the new stream's initial read fails.
+                if e.code() not in _TRANSIENT_STATUS_CODES:
+                    raise
+                stream_error = f'{e.details()} Status Code: {e.code()}'
+
+            if reconnect_attempts >= MAX_RECONNECT_ATTEMPTS:
+                raise StreamReconnectError(
+                    f'Subscription stream still failing after {reconnect_attempts} '
+                    f'reconnect attempts: {stream_error}'
+                )
+            if reconnect_attempts > 0:
+                backoff_seconds = _reconnect_backoff_seconds(reconnect_attempts)
+                with contextlib.suppress(asyncio.TimeoutError):
+                    await asyncio.wait_for(self._closed.wait(), timeout=backoff_seconds)
+                if self._closed.is_set():
+                    raise StreamInactiveError('Stream is not active')
+            reconnect_attempts += 1
+            logger.warning('Subscription stream error: %s — reconnecting', stream_error)
+
+    async def _read_next_message(self) -> SubscriptionMessage:
+        """Read one message from the current stream.
+
+        Raises:
+            _TransientStreamError: If the stream failed in a way a reconnect can fix.
+        """
+        if not self._stream_active.is_set() or self._stream is None:
             raise StreamInactiveError('Stream is not active')
 
+        send_queue = self._send_queue
         try:
-            if self._stream is not None:
-                message = await self._stream.read()
-                if message is None:
-                    return None
-                return SubscriptionMessage(message.event_message)
+            message = await self._stream.read()
         except AioRpcError as e:
-            if e.code() == StatusCode.UNAVAILABLE or e.code() == StatusCode.UNKNOWN:
-                print(
-                    f'gRPC error while reading from stream: {e.details()}, '
-                    f'Status Code: {e.code()}. '
-                    f'Attempting to reconnect...'
-                )
-                await self.reconnect_stream()
-            elif e.code() == StatusCode.CANCELLED:
+            if e.code() in _TRANSIENT_STATUS_CODES:
+                raise _TransientStreamError(f'{e.details()} Status Code: {e.code()}') from e
+            if e.code() == StatusCode.CANCELLED:
                 raise StreamCancelledError('Stream has been cancelled')
-            else:
-                raise Exception(f'gRPC error while reading from subscription stream: {e} ')
+            raise Exception(f'gRPC error while reading from subscription stream: {e} ')
         except Exception as e:
             raise Exception(f'Error while fetching message: {e}')
 
-        return None
+        if message is EOF:
+            raise _TransientStreamError('Stream ended by the server')
+        subscription_message = SubscriptionMessage(message.event_message)
+        subscription_message._origin_send_queue = send_queue
+        return subscription_message
 
     async def respond(self, message, status):
         try:
@@ -92,9 +195,14 @@ async def respond(self, message, status):
             msg = api_v1.SubscribeTopicEventsRequestAlpha1(event_processed=response)
             if not self._stream_active.is_set():
                 raise StreamInactiveError('Stream is not active')
-            await self._send_queue.put(msg)
+            send_queue = self._send_queue
+            origin_send_queue = message._origin_send_queue
+            if origin_send_queue is not None and origin_send_queue is not send_queue:
+                logger.debug('Dropping response to message %s from a replaced stream', message.id())
+                return
+            await send_queue.put(msg)
         except Exception as e:
-            print(f"Can't send message: {e}")
+            logger.warning("Can't send response for message %s: %s", message.id(), e)
 
     async def respond_success(self, message):
         await self.respond(message, TopicEventResponse('success').status)
@@ -106,6 +214,11 @@ async def respond_drop(self, message):
         await self.respond(message, TopicEventResponse('drop').status)
 
     async def close(self):
+        self._closed.set()
+        await self._close_stream()
+
+    async def _close_stream(self):
+        self._send_queue.put_nowait(None)
         if self._stream:
             try:
                 self._stream.cancel()

dapr/clients/grpc/client.py

@@ -13,6 +13,7 @@
 limitations under the License.
 """
 
+import logging
 import socket
 import threading
 import time
@@ -94,6 +95,10 @@
 from dapr.proto import api_service_v1, api_v1, common_v1
 from dapr.version import __version__
 
+logger = logging.getLogger(__name__)
+
+SUBSCRIPTION_CLOSE_TIMEOUT_SECONDS = 5
+
 
 class DaprGrpcClient:
     """The convenient layer implementation of Dapr gRPC APIs.
@@ -624,6 +629,10 @@ def subscribe_with_handler(
             handler_fn (Callable[..., TopicEventResponse]): The function to call when a message is received.
             metadata (Optional[MetadataTuple]): Additional metadata for the subscription.
             dead_letter_topic (Optional[str]): Name of the dead-letter topic.
+
+        Returns:
+            Callable: Closes the subscription and waits up to SUBSCRIPTION_CLOSE_TIMEOUT_SECONDS
+                for the handler thread to stop.
         """
         subscription = self.subscribe(pubsub_name, topic, metadata, dead_letter_topic)
 
@@ -640,9 +649,7 @@ def stream_messages(sub):
                             # No message received
                             continue
 
-                except StreamInactiveError:
-                    break
-                except StreamCancelledError:
+                except (StreamInactiveError, StreamCancelledError):
                     break
                 except Exception:
                     # Stream died — reconnect via the subscription's own
@@ -657,6 +664,13 @@ def stream_messages(sub):
 
         def close_subscription():
             subscription.close()
+            if threading.current_thread() is not streaming_thread:
+                streaming_thread.join(timeout=SUBSCRIPTION_CLOSE_TIMEOUT_SECONDS)
+                if streaming_thread.is_alive():
+                    logger.warning(
+                        'Subscription handler thread still running %ss after close',
+                        SUBSCRIPTION_CLOSE_TIMEOUT_SECONDS,
+                    )
 
         streaming_thread = threading.Thread(target=stream_messages, args=(subscription,))
         streaming_thread.start()