Keep subscribe_with_handler running until the caller closes it

dapr/aio/clients/grpc/client.py

@@ -14,9 +14,22 @@
 """
 
 import asyncio
+import logging
 import socket
 import time
-from typing import Any, Awaitable, BinaryIO, Callable, Dict, List, Optional, Sequence, Text, Union
+from typing import (
+    Any,
+    Awaitable,
+    BinaryIO,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Text,
+    Union,
+)
 from urllib.parse import urlencode
 from warnings import warn
 from weakref import WeakSet
@@ -97,12 +110,21 @@
     UnlockResponseStatus,
 )
 from dapr.clients.grpc._state import StateItem, StateOptions
+from dapr.clients.grpc.client import (
+    SUBSCRIPTION_CLOSE_TIMEOUT_SECONDS,
+    SUBSCRIPTION_RECONNECT_BACKOFF_SECONDS,
+)
 from dapr.clients.health import DaprHealth
 from dapr.clients.retry import RetryPolicy
-from dapr.common.pubsub.subscription import StreamInactiveError
+from dapr.common.pubsub.subscription import StreamCancelledError, StreamInactiveError
 from dapr.proto import api_service_v1, api_v1, common_v1
 from dapr.proto.runtime.v1 import binarystore_pb2 as binary_v1
 
+logger = logging.getLogger(__name__)
+
+# The event loop only keeps weak references to tasks, so running handler tasks are kept here.
+_subscription_tasks: Set[asyncio.Task] = set()
+
 
 class DaprGrpcClientAsync:
     """The async convenient layer implementation of Dapr gRPC APIs.
@@ -592,28 +614,72 @@ async def subscribe_with_handler(
             dead_letter_topic (Optional[str]): Name of the dead-letter topic.
 
         Returns:
-            Callable[[], Awaitable[None]]: An async function to close the subscription.
+            Callable[[], Awaitable[None]]: An async function that closes the subscription and
+                waits up to SUBSCRIPTION_CLOSE_TIMEOUT_SECONDS for the handler task to stop.
         """
         subscription = await self.subscribe(pubsub_name, topic, metadata, dead_letter_topic)
+        closed = asyncio.Event()
 
         async def stream_messages(sub: Subscription):
             while True:
                 try:
                     async for message in subscription:
                         if message:
-                            response = await handler_fn(message)
+                            try:
+                                response = await handler_fn(message)
+                            except Exception:
+                                logger.exception('Subscription handler failed, retrying message')
+                                await subscription.respond_retry(message)
+                                continue
                             if response:
                                 await subscription.respond(message, response.status)
                         else:
                             continue
 
-                except StreamInactiveError:
+                except (StreamInactiveError, StreamCancelledError):
+                    pass
+                except asyncio.CancelledError:
+                    # grpc.aio raises this from a read on a stream that close() cancelled
+                    if not closed.is_set():
+                        raise
+                except Exception:
+                    logger.warning(
+                        'Subscription stream failed, reconnecting in %s seconds',
+                        SUBSCRIPTION_RECONNECT_BACKOFF_SECONDS,
+                        exc_info=True,
+                    )
+                    try:
+                        await asyncio.wait_for(
+                            closed.wait(), timeout=SUBSCRIPTION_RECONNECT_BACKOFF_SECONDS
+                        )
+                    except asyncio.TimeoutError:
+                        pass
+                if closed.is_set():
                     break
+                try:
+                    await sub.reconnect_stream()
+                except Exception:
+                    logger.warning(
+                        'Subscription reconnect failed, retrying in %s seconds',
+                        SUBSCRIPTION_RECONNECT_BACKOFF_SECONDS,
+                        exc_info=True,
+                    )
+                    try:
+                        await asyncio.wait_for(
+                            closed.wait(), timeout=SUBSCRIPTION_RECONNECT_BACKOFF_SECONDS
+                        )
+                    except asyncio.TimeoutError:
+                        pass
 
         async def close_subscription():
+            closed.set()
             await subscription.close()
+            if asyncio.current_task() is not task:
+                await asyncio.wait({task}, timeout=SUBSCRIPTION_CLOSE_TIMEOUT_SECONDS)
 
-        asyncio.create_task(stream_messages(subscription))
+        task = asyncio.create_task(stream_messages(subscription))
+        _subscription_tasks.add(task)
+        task.add_done_callback(_subscription_tasks.discard)
 
         return close_subscription
 

dapr/clients/grpc/client.py

@@ -106,6 +106,7 @@
 logger = logging.getLogger(__name__)
 
 SUBSCRIPTION_CLOSE_TIMEOUT_SECONDS = 5
+SUBSCRIPTION_RECONNECT_BACKOFF_SECONDS = 5
 
 
 class DaprGrpcClient:
@@ -643,34 +644,52 @@ def subscribe_with_handler(
                 for the handler thread to stop.
         """
         subscription = self.subscribe(pubsub_name, topic, metadata, dead_letter_topic)
+        closed = threading.Event()
 
         def stream_messages(sub):
             while True:
                 try:
                     for message in sub:
                         if message:
                             # Process the message
-                            response = handler_fn(message)
+                            try:
+                                response = handler_fn(message)
+                            except Exception:
+                                logger.exception('Subscription handler failed, retrying message')
+                                subscription.respond_retry(message)
+                                continue
                             if response:
                                 subscription.respond(message, response.status)
                         else:
                             # No message received
                             continue
 
                 except (StreamInactiveError, StreamCancelledError):
+                    pass
+                except Exception:
+                    logger.warning(
+                        'Subscription stream failed, reconnecting in %s seconds',
+                        SUBSCRIPTION_RECONNECT_BACKOFF_SECONDS,
+                        exc_info=True,
+                    )
+                    closed.wait(SUBSCRIPTION_RECONNECT_BACKOFF_SECONDS)
+                if closed.is_set():
                     break
+                # Reconnect via the subscription's own reconnect logic (which waits for the
+                # sidecar to be healthy).
+                try:
+                    sub.reconnect_stream()
                 except Exception:
-                    # Stream died — reconnect via the subscription's own
-                    # reconnect logic (which waits for the sidecar to be healthy).
-                    try:
-                        sub.reconnect_stream()
-                    except Exception:
-                        # Sidecar still unavailable — back off before retrying
-                        # TODO: Make this configurable
-                        time.sleep(5)
-                    continue
+                    # Sidecar still unavailable — back off before retrying
+                    logger.warning(
+                        'Subscription reconnect failed, retrying in %s seconds',
+                        SUBSCRIPTION_RECONNECT_BACKOFF_SECONDS,
+                        exc_info=True,
+                    )
+                    closed.wait(SUBSCRIPTION_RECONNECT_BACKOFF_SECONDS)
 
         def close_subscription():
+            closed.set()
             subscription.close()
             if threading.current_thread() is not streaming_thread:
                 streaming_thread.join(timeout=SUBSCRIPTION_CLOSE_TIMEOUT_SECONDS)

tests/clients/fake_dapr_server.py

@@ -75,9 +75,10 @@ def __init__(self, grpc_port: int = 0, http_port: int = 0):
         self.actor_stream_replies: List[api_v1.SubscribeActorEventsRequestAlpha1] = []
         # SubscribeTopicEventsAlpha1: one status code is popped per stream connection, which then
         # fails right after the initial response (OK ends it cleanly). IDs of acked events are
-        # recorded in order.
+        # recorded in order, and with the response status in topic_stream_responses.
         self.topic_stream_failures: List[grpc.StatusCode] = []
         self.topic_stream_acks: List[str] = []
+        self.topic_stream_responses: List[Tuple[str, int]] = []
 
     def set_bulk_publish_unimplemented_on_stable_next(self) -> None:
         """Make the next BulkPublishEvent (stable) call return UNIMPLEMENTED.
@@ -313,6 +314,9 @@ def SubscribeTopicEventsAlpha1(self, request_iterator, context):
         for request in request_iterator:
             if request.HasField('event_processed'):
                 self.topic_stream_acks.append(request.event_processed.id)
+                self.topic_stream_responses.append(
+                    (request.event_processed.id, request.event_processed.status.status)
+                )
                 break
 
         msg2 = appcallback_v1.TopicEventRequest(
@@ -331,6 +335,9 @@ def SubscribeTopicEventsAlpha1(self, request_iterator, context):
         for request in request_iterator:
             if request.HasField('event_processed'):
                 self.topic_stream_acks.append(request.event_processed.id)
+                self.topic_stream_responses.append(
+                    (request.event_processed.id, request.event_processed.status.status)
+                )
                 break
 
         # On the third message simulate a disconnection