From 36d7b2c586cc1678b53eb1d409fe80e152e0ec5f Mon Sep 17 00:00:00 2001 From: Fabian Martinez <46371672+famarting@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:51:42 +0000 Subject: [PATCH] fix(workflow): keep the workflow worker working across sidecar reconnects (#1312) * fix(workflow): put a deadline on the worker's reconnect Hello The worker checks each new connection to the sidecar with a Hello call that had no deadline. After a host sleep or a network change the new connection can be left half-open, and that Hello then waits forever. It runs on the worker's event loop, so the worker never reconnects while the process looks healthy. Give the Hello a 30 second deadline, so a stuck attempt fails and the existing retry loop tries again. Also close the channel of a failed attempt instead of dropping it, so repeated failures do not leak channels. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Fabian Martinez <46371672+famarting@users.noreply.github.com> * fix(workflow): report a worker as not ready while its stream is down is_worker_ready() was set on the first connect and never cleared, so it kept returning True while the worker had lost its work-item stream and was trying to reconnect. Nothing outside the worker could tell a worker that receives work from one that has been cut off from the sidecar. Clear the ready flag whenever the stream is dropped; it is set again when the next stream opens. This matches what WorkflowRuntime documents for wait_for_worker_ready(): that the stream is ready to receive work items. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Fabian Martinez <46371672+famarting@users.noreply.github.com> * fix(workflow): send work-item results on the worker's current connection Each activity and workflow work item carried the stub of the connection it arrived on, and sent its result through that stub. When the sidecar ends the work-item stream the worker opens a new connection and closes the old channel ten seconds later. Any work item that runs longer than that sends its result on a closed channel: grpc raises ValueError, the worker logs it at debug level, and the result is lost. The sidecar then runs the activity again from the start. Send results (and the history fetch for a cache miss) through the worker's current connection instead. If the result is ready while the worker is between connections, wait up to 30 seconds for the next one; if the worker is stopping, give up at once and let the sidecar re-dispatch the work item. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Fabian Martinez <46371672+famarting@users.noreply.github.com> * fix(workflow): never close a channel while the worker may start a call on it grpcio does not refuse a call on a closed channel when the channel has client interceptors (the workflow runtime always adds one). The intercepted call registers its method on the underlying channel first, which locks a mutex inside the already freed native channel, and the whole process dies with SIGSEGV instead of getting a ValueError. The worker closed channels without regard to calls about to start on them: the old channel ten seconds after a reconnect, and the current one at once in stop(), racing result sends, the keepalive Hello, the connect Hello and GetWorkItems. Count the calls in flight on each channel. Every call the worker makes now takes the channel through _sidecar_stub, which only hands out the current connection, and a channel is closed only after it is retired and its count drops to zero. stop() waits up to five seconds for those calls and then closes the channel anyway; by then any such call has long started, so closing it only cancels it. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Fabian Martinez <46371672+famarting@users.noreply.github.com> * fix(workflow): verify a connection before results use it, and bound every call on it Addresses review feedback on the reconnect fix: - Hello runs on the new connection before it is published, so a result waiting for a reconnect is never sent on a connection that then fails. - Result sends and the history fetch get a 30s deadline. A retired channel is closed only once no call is in flight on it, so these deadlines are what bound that close and stop(); the close itself never runs early, since that could land between a result and its oversized-result fallback. - stop() keeps the connection open for its drain, so a result that becomes ready while the worker stops is still delivered. - Releasing a call wakes waiters only when the channel's count reaches zero. - The run loop tracks a connected flag instead of an unused stub. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Fabian Martinez <46371672+famarting@users.noreply.github.com> --------- Signed-off-by: Fabian Martinez <46371672+famarting@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (cherry picked from commit 9136b370298376dfd6023bae09703f20c359c044) Signed-off-by: dapr-bot --- dapr/ext/workflow/AGENTS.md | 4 +- dapr/ext/workflow/_durabletask/worker.py | 293 ++++++++++--- .../test_async_dispatch_regression.py | 10 +- .../durabletask/test_worker_history_cache.py | 27 +- .../durabletask/test_worker_reconnect.py | 404 ++++++++++++++++++ 5 files changed, 663 insertions(+), 75 deletions(-) create mode 100644 tests/ext/workflow/durabletask/test_worker_reconnect.py diff --git a/dapr/ext/workflow/AGENTS.md b/dapr/ext/workflow/AGENTS.md index e2606e8e9..d6b4e4246 100644 --- a/dapr/ext/workflow/AGENTS.md +++ b/dapr/ext/workflow/AGENTS.md @@ -120,7 +120,9 @@ Activities can be either `def my_activity(ctx, inp)` or `async def my_activity(c At dispatch time (the gRPC stream loop in `_durabletask/worker.py`), `is_async_callable(activity_fn)` on the wrapper selects between two handlers. -- **Async activities** go through `_execute_activity_async`, then `_ActivityExecutor.execute_async`, which awaits `fn(...)` directly on the event loop. The gRPC response is delivered via `loop.run_in_executor(self._async_worker_manager.thread_pool, stub.CompleteActivityTask, ...)` — the same pool sync activities use, sized by `maximum_thread_pool_workers`. +- **Async activities** go through `_execute_activity_async`, then `_ActivityExecutor.execute_async`, which awaits `fn(...)` directly on the event loop. The gRPC response is delivered via `loop.run_in_executor(self._async_worker_manager.thread_pool, self._deliver_activity_response, ...)` — the same pool sync activities use, sized by `maximum_thread_pool_workers`. + +**Results go on the current connection.** A work item can outlive the connection it arrived on: when the sidecar ends the work-item stream the worker reconnects and closes the old channel shortly after. Activity and workflow results are therefore sent through `_send_on_current_connection`, which uses the worker's connection at the moment the result is ready and waits up to `_RESULT_RECONNECT_WAIT` seconds for a reconnect. Never send on a stub captured when the work item arrived. - **Sync activities** go through `_execute_activity`, dispatched to the thread pool by `_AsyncWorkerManager._run_func`. The activity runs on a worker thread, and the response is delivered from the same thread. Workflow (orchestrator) functions must remain generators (`def` with `yield`). They cannot be `async def` because durabletask's deterministic replay depends on synchronous generator semantics. Only activities support async. diff --git a/dapr/ext/workflow/_durabletask/worker.py b/dapr/ext/workflow/_durabletask/worker.py index a9f5b3b0d..c1166bb8d 100644 --- a/dapr/ext/workflow/_durabletask/worker.py +++ b/dapr/ext/workflow/_durabletask/worker.py @@ -18,6 +18,7 @@ import threading import time import warnings +from collections import Counter from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta, timezone from threading import Event, Thread @@ -229,6 +230,19 @@ def _is_message_too_large(rpc_error: grpc.RpcError) -> bool: _DEFAULT_HISTORY_CACHE_TTL = 3600.0 _DEFAULT_HISTORY_CACHE_MAX_INSTANCES = 100_000 _HISTORY_CACHE_SWEEP_INTERVAL = 60.0 +# Deadline for the Hello that checks a new connection. Without one, a connection left half-open +# by a network change or a host sleep can block the Hello, and with it every reconnect, forever. +# Longer than the keepalive's, since it covers the TCP and TLS handshakes too. +_CONNECT_HELLO_TIMEOUT = 30.0 +# How long a finished work item waits for the worker to reconnect before giving up on sending its +# result. Covers the reconnect backoff (at most 16s) plus a connect Hello that needs a while. +_RESULT_RECONNECT_WAIT = 30.0 +# Deadline for each call that sends a result or fetches history. A retired channel is closed only +# once no call is in flight on it, so this deadline is also what bounds that close, and stop(). +_RESULT_SEND_TIMEOUT = 30.0 +# How long stop() keeps the current connection open for results that are still being sent or that +# become ready while it stops, before it retires the connection. +_STOP_DRAIN_TIMEOUT = 5.0 class _HistoryResolutionError(Exception): @@ -468,6 +482,11 @@ def __init__( self._stop_timeout = stop_timeout self._keepalive_interval = keepalive_interval self._current_channel: Optional[grpc.Channel] = None # Store channel reference for cleanup + self._current_stub: Optional[stubs.TaskHubSidecarServiceStub] = None + # Guards the current channel and stub and the calls in flight on each channel; notified + # whenever any of them change. + self._connection_lock = threading.Condition() + self._calls_in_flight: Counter[grpc.Channel] = Counter() self._channel_cleanup_threads: list[threading.Thread] = [] # Deferred channel close threads self._stream_ready = threading.Event() self._runLoop: Optional[Thread] = None @@ -521,6 +540,11 @@ def add_activity(self, fn: task.Activity) -> str: return self._registry.add_activity(fn) def is_worker_ready(self) -> bool: + """Whether the worker is running and its work-item stream to the sidecar is open. + + False while the worker reconnects after losing the stream, so a caller can tell a + worker that is receiving work items from one that has been cut off from the sidecar. + """ return self._stream_ready.is_set() and self._is_running def start(self): @@ -550,7 +574,12 @@ def run_loop(): ) self._is_running = True - async def _keepalive_loop(self, stub): + def _send_keepalive(self) -> None: + with self._sidecar_stub(wait=0) as stub: + if stub is not None: + stub.Hello(empty_pb2.Empty(), timeout=10) + + async def _keepalive_loop(self): """Background keepalive loop to keep the gRPC connection alive through L7 load balancers.""" loop = asyncio.get_running_loop() while not self._shutdown.is_set(): @@ -558,7 +587,7 @@ async def _keepalive_loop(self, stub): if self._shutdown.is_set(): return try: - await loop.run_in_executor(None, lambda: stub.Hello(empty_pb2.Empty(), timeout=10)) + await loop.run_in_executor(None, self._send_keepalive) except Exception as e: self._logger.debug(f'keepalive failed: {e}') @@ -582,54 +611,53 @@ async def _async_run_loop(self): """ worker_task = asyncio.create_task(self._async_worker_manager.run()) # Connection state management for retry fix - current_channel = None - current_stub = None + connected = False current_reader_thread = None conn_retry_count = 0 conn_max_retry_delay = 15 - def create_fresh_connection(): - nonlocal current_channel, current_stub, conn_retry_count + def retire_and_close_connection(): # Schedule deferred close of old channel to avoid orphaned TCP - # connections. In-flight RPCs on the old stub may still reference - # the channel from another thread, so we wait a grace period - # before closing instead of closing immediately. - if current_channel is not None: - self._schedule_deferred_channel_close(current_channel) - current_channel = None - current_stub = None + # connections. In-flight RPCs (e.g. CompleteActivityTask) may still + # be using the stub on another thread, so we defer the close by a + # grace period instead of closing immediately. A channel that + # stop() already retired is closed by stop(). + retired_channel = self._retire_connection() + if retired_channel is not None: + self._schedule_deferred_channel_close(retired_channel) + + def create_fresh_connection(): + nonlocal connected, conn_retry_count + connected = False + channel = shared.get_grpc_channel( + self._host_address, + self._secure_channel, + self._interceptors, + options=self._channel_options, + ) + stub = stubs.TaskHubSidecarServiceStub(channel) + # Verify before publishing, so a result waiting for a reconnect is never sent on a + # connection that fails its Hello. Nothing else holds the channel yet, so closing it + # directly is safe. try: - current_channel = shared.get_grpc_channel( - self._host_address, - self._secure_channel, - self._interceptors, - options=self._channel_options, - ) - # Store channel reference for cleanup in stop() - self._current_channel = current_channel - current_stub = stubs.TaskHubSidecarServiceStub(current_channel) - current_stub.Hello(empty_pb2.Empty()) - conn_retry_count = 0 - self._logger.info(f'Created fresh connection to {self._host_address}') + stub.Hello(empty_pb2.Empty(), timeout=_CONNECT_HELLO_TIMEOUT) except Exception as e: self._logger.warning(f'Failed to create connection: {e}') - current_channel = None - self._current_channel = None - current_stub = None + channel.close() raise + if not self._set_connection(channel, stub): + channel.close() + raise RuntimeError('the worker is stopping') + connected = True + conn_retry_count = 0 + self._logger.info(f'Created fresh connection to {self._host_address}') def invalidate_connection(): - nonlocal current_channel, current_stub, current_reader_thread - # Schedule deferred close of old channel to avoid orphaned TCP - # connections. In-flight RPCs (e.g. CompleteActivityTask) may still - # be using the stub on another thread, so we defer the close by a - # grace period instead of closing immediately. - if current_channel is not None: - self._schedule_deferred_channel_close(current_channel) - current_channel = None - self._current_channel = None - current_stub = None + nonlocal connected, current_reader_thread + retire_and_close_connection() + connected = False self._response_stream = None + self._stream_ready.clear() # The sidecar drops this stream's warm set on disconnect, so start the # next stream cold to stay in sync. self._history_cache.reset() @@ -653,7 +681,7 @@ def should_invalidate_connection(rpc_error): while True: if self._shutdown.is_set(): break - if current_stub is None: + if not connected: try: create_fresh_connection() except Exception: @@ -670,8 +698,6 @@ def should_invalidate_connection(rpc_error): continue keepalive_task = None try: - assert current_stub is not None - stub = current_stub get_work_items_request = pb.GetWorkItemsRequest( capabilities=[pb.WORKER_CAPABILITY_HEALTH_PING] ) @@ -680,7 +706,11 @@ def should_invalidate_connection(rpc_error): pb.WORKER_CAPABILITY_STATEFUL_HISTORY ) try: - self._response_stream = stub.GetWorkItems(get_work_items_request) + with self._sidecar_stub(wait=0) as stream_stub: + if stream_stub is None or self._shutdown.is_set(): + self._logger.debug('Worker is stopping, exiting') + break + self._response_stream = stream_stub.GetWorkItems(get_work_items_request) self._logger.info( f'Successfully connected to {self._host_address}. Waiting for work items...' ) @@ -791,7 +821,7 @@ def stream_reader(): loop = asyncio.get_running_loop() if self._keepalive_interval > 0: - keepalive_task = asyncio.ensure_future(self._keepalive_loop(stub)) + keepalive_task = asyncio.ensure_future(self._keepalive_loop()) # NOTE: This is a blocking call that will wait for a work item to become available or the shutdown sentinel while not self._shutdown.is_set(): @@ -817,7 +847,6 @@ def stream_reader(): self._async_worker_manager.submit_orchestration( self._execute_orchestrator, work_item.workflowRequest, - stub, work_item.completionToken, teardown_stream, ) @@ -836,7 +865,6 @@ def stream_reader(): activity_handler, activity_fn, work_item.activityRequest, - stub, work_item.completionToken, ) elif work_item.HasField('healthPing'): @@ -949,6 +977,115 @@ def stream_reader(): except Exception as e: self._logger.warning(f'Error while waiting for worker task shutdown: {e}') + def _set_connection(self, channel: grpc.Channel, stub: stubs.TaskHubSidecarServiceStub) -> bool: + """Makes ``channel`` the connection that work-item results are sent on. + + Returns: + False, without publishing, if the worker is stopping; the caller then owns the + channel and must close it. + """ + with self._connection_lock: + # stop() sets _shutdown before it retires the connection, so a channel published + # after this check is still retired and closed by stop(). + if self._shutdown.is_set(): + return False + self._current_channel = channel + self._current_stub = stub + self._connection_lock.notify_all() + return True + + def _retire_connection(self) -> Optional[grpc.Channel]: + """Stops handing out the current connection and returns its channel, if any.""" + with self._connection_lock: + channel = self._current_channel + self._current_channel = None + self._current_stub = None + self._connection_lock.notify_all() + return channel + + def _has_connection_or_is_stopping(self) -> bool: + return self._current_stub is not None or self._shutdown.is_set() + + @contextlib.contextmanager + def _sidecar_stub(self, wait: float) -> Iterator[Optional[stubs.TaskHubSidecarServiceStub]]: + """Yields the stub of the worker's current connection to the sidecar. + + A work item can outlive the connection it arrived on: when the sidecar ends the stream + the worker reconnects, and the old channel is closed soon after. Results are therefore + sent on whatever connection is current when they are ready, never on the work item's. + + Args: + wait: Seconds to wait for a connection while the worker is reconnecting. + + Yields: + The current stub, or None if the worker has no connection after ``wait`` seconds, or + has none while it is stopping. A stopping worker keeps its connection open for a + short drain, and results sent during it still go out. + """ + with self._connection_lock: + self._connection_lock.wait_for(self._has_connection_or_is_stopping, timeout=wait) + channel = self._current_channel + stub = self._current_stub + if stub is not None: + self._calls_in_flight[channel] += 1 + try: + yield stub + finally: + if stub is not None: + self._release_call(channel) + + def _release_call(self, channel: grpc.Channel) -> None: + with self._connection_lock: + self._calls_in_flight[channel] -= 1 + # Only the drain waits depend on the count, and only on it reaching zero. + if self._calls_in_flight[channel] <= 0: + del self._calls_in_flight[channel] + self._connection_lock.notify_all() + + def _drain_current_connection(self, timeout: float) -> None: + """Waits up to ``timeout`` seconds for no call to be in flight on the current channel.""" + with self._connection_lock: + channel = self._current_channel + if channel is None: + return + self._connection_lock.wait_for( + lambda: self._current_channel is not channel or self._calls_in_flight[channel] == 0, + timeout=timeout, + ) + + def _close_channel(self, channel: grpc.Channel) -> None: + """Closes a retired channel once no call that the worker started on it is in flight. + + grpcio does not guard against starting a call on a closed channel when the channel has + interceptors: registering the method touches the freed native channel and the process + dies with SIGSEGV. Every call the worker makes holds the channel through + ``_sidecar_stub``, and a retired channel is never handed out again, so once its count + reaches zero no call can start on it after the close. + + Gotcha: there is deliberately no timeout. Closing early could land between two calls a + holder makes, such as a result and its oversized-result fallback, and the second would + crash. Every call made under ``_sidecar_stub`` has a deadline, which bounds this wait. + + Args: + channel: A channel that ``_retire_connection`` already took out of use. + """ + with self._connection_lock: + self._connection_lock.wait_for(lambda: self._calls_in_flight[channel] == 0) + channel.close() + + def _send_on_current_connection( + self, send: Callable[[stubs.TaskHubSidecarServiceStub], None], what: str + ) -> None: + """Calls ``send`` with the current connection's stub, waiting out a reconnect.""" + with self._sidecar_stub(wait=_RESULT_RECONNECT_WAIT) as stub: + if stub is None: + self._logger.warning( + f'Could not deliver {what}: not connected to the sidecar. ' + 'The sidecar will re-dispatch this work item.' + ) + return + send(stub) + def _schedule_deferred_channel_close( self, old_channel: grpc.Channel, grace_timeout: float = 10.0 ): @@ -972,7 +1109,7 @@ def _deferred_close(): self._shutdown.wait(timeout=grace_timeout) finally: try: - old_channel.close() + self._close_channel(old_channel) self._logger.debug('Deferred channel close completed') except Exception as e: self._logger.debug(f'Error during deferred channel close: {e}') @@ -990,14 +1127,15 @@ def stop(self): self._logger.info('Stopping gRPC worker...') self._shutdown.set() + # Results that are being sent, or that become ready now, still go out on the connection. + self._drain_current_connection(_STOP_DRAIN_TIMEOUT) # Close the channel — propagates cancellation to all streams and cleans up resources - if self._current_channel is not None: + channel = self._retire_connection() + if channel is not None: try: - self._current_channel.close() + self._close_channel(channel) except Exception as e: self._logger.exception(f'Error closing gRPC channel: {e}') - finally: - self._current_channel = None if self._runLoop is not None: self._runLoop.join(timeout=self._stop_timeout) @@ -1068,9 +1206,7 @@ def _sweep_history_cache_loop(self): while not self._shutdown.wait(_HISTORY_CACHE_SWEEP_INTERVAL): self._history_cache.sweep_expired() - def _resolve_history( - self, req: pb.WorkflowRequest, stub: stubs.TaskHubSidecarServiceStub - ) -> list[pb.HistoryEvent]: + def _resolve_history(self, req: pb.WorkflowRequest) -> list[pb.HistoryEvent]: """Resolves the full committed history to replay for a workflow work item. For a full send it is the request's pastEvents. For a delta send (cachedHistory) it @@ -1089,7 +1225,10 @@ def _resolve_history( history_request = pb.GetInstanceHistoryRequest(instanceId=req.instanceId) try: - response = stub.GetInstanceHistory(history_request) + with self._sidecar_stub(wait=0) as stub: + if stub is None: + raise RuntimeError('not connected to the sidecar') + response = stub.GetInstanceHistory(history_request, timeout=_RESULT_SEND_TIMEOUT) except Exception as ex: raise _HistoryResolutionError( f"Failed to fetch the committed history for '{req.instanceId}': {ex}" @@ -1118,7 +1257,6 @@ def _update_history_cache( def _execute_orchestrator( self, req: pb.WorkflowRequest, - stub: stubs.TaskHubSidecarServiceStub, completionToken, teardown_stream: Callable[[], None], ): @@ -1129,7 +1267,7 @@ def _execute_orchestrator( if req.HasField('propagatedHistory') else None ) - old_events = self._resolve_history(req, stub) + old_events = self._resolve_history(req) result = executor.execute( req.instanceId, old_events, req.newEvents, propagated_history=propagated ) @@ -1174,8 +1312,22 @@ def _execute_orchestrator( completionToken=completionToken, ) + self._send_on_current_connection( + lambda stub: self._send_orchestrator_response(req, stub, res, completionToken), + f"orchestrator response for '{req.instanceId}'", + ) + + def _send_orchestrator_response( + self, + req: pb.WorkflowRequest, + stub: stubs.TaskHubSidecarServiceStub, + res: pb.WorkflowResponse, + completionToken, + ): + """Send an orchestrator response, falling back to a failure response when the + response is too large to deliver.""" try: - stub.CompleteOrchestratorTask(res) + stub.CompleteOrchestratorTask(res, timeout=_RESULT_SEND_TIMEOUT) except grpc.RpcError as rpc_error: # type: ignore if _is_message_too_large(rpc_error): # Response is too large to deliver - fail the orchestration immediately. @@ -1202,7 +1354,7 @@ def _execute_orchestrator( completionToken=completionToken, ) try: - stub.CompleteOrchestratorTask(failure_res) + stub.CompleteOrchestratorTask(failure_res, timeout=_RESULT_SEND_TIMEOUT) except Exception as ex: self._logger.exception( f"Failed to deliver orchestrator failure response for '{req.instanceId}': {ex}" @@ -1281,7 +1433,7 @@ def _send_activity_response( """Send an activity response, falling back to a failure response when the result is too large to deliver.""" try: - stub.CompleteActivityTask(res) + stub.CompleteActivityTask(res, timeout=_RESULT_SEND_TIMEOUT) except grpc.RpcError as rpc_error: # type: ignore if _is_message_too_large(rpc_error): # Result is too large to deliver - fail the activity immediately. @@ -1297,7 +1449,7 @@ def _send_activity_response( req, instance_id, oversize_error, completion_token ) try: - stub.CompleteActivityTask(failure_res) + stub.CompleteActivityTask(failure_res, timeout=_RESULT_SEND_TIMEOUT) except Exception as ex: self._logger.exception( f"Failed to deliver activity failure response for '{req.name}#{req.taskId}' " @@ -1317,11 +1469,24 @@ def _send_activity_response( f"Failed to deliver activity response for '{req.name}#{req.taskId}' of orchestration ID '{instance_id}' to sidecar: {ex}" ) + def _deliver_activity_response( + self, + req: pb.ActivityRequest, + res: pb.ActivityResponse, + completion_token, + instance_id: str, + ) -> None: + self._send_on_current_connection( + lambda stub: self._send_activity_response( + req, stub, res, completion_token, instance_id + ), + f"activity response for '{req.name}#{req.taskId}' of orchestration ID '{instance_id}'", + ) + def _execute_activity( self, fn: task.Activity | None, req: pb.ActivityRequest, - stub: stubs.TaskHubSidecarServiceStub, completionToken, ): instance_id = req.workflowInstance.instanceId @@ -1341,13 +1506,12 @@ def _execute_activity( ) except Exception as ex: res = self._build_activity_failure_response(req, instance_id, ex, completionToken) - self._send_activity_response(req, stub, res, completionToken, instance_id) + self._deliver_activity_response(req, res, completionToken, instance_id) async def _execute_activity_async( self, fn: task.Activity, req: pb.ActivityRequest, - stub: stubs.TaskHubSidecarServiceStub, completionToken, ): """Run an async activity on the event loop and send its result to the sidecar. @@ -1376,9 +1540,8 @@ async def _execute_activity_async( try: await loop.run_in_executor( self._async_worker_manager.thread_pool, - self._send_activity_response, + self._deliver_activity_response, req, - stub, res, completionToken, instance_id, diff --git a/tests/ext/workflow/durabletask/test_async_dispatch_regression.py b/tests/ext/workflow/durabletask/test_async_dispatch_regression.py index 8ea14172b..e7ed6d109 100644 --- a/tests/ext/workflow/durabletask/test_async_dispatch_regression.py +++ b/tests/ext/workflow/durabletask/test_async_dispatch_regression.py @@ -17,9 +17,12 @@ """ import asyncio +from typing import cast +from unittest.mock import MagicMock import pytest +import dapr.ext.workflow._durabletask.internal.orchestrator_service_pb2_grpc as stubs import dapr.ext.workflow._durabletask.internal.protos as pb from dapr.ext.workflow._durabletask.worker import ConcurrencyOptions, TaskHubGrpcWorker @@ -40,7 +43,9 @@ class _MockSidecarStub: def __init__(self) -> None: self.completions = 0 - def CompleteActivityTask(self, _response: pb.ActivityResponse) -> None: # noqa: N802 + def CompleteActivityTask( # noqa: N802 + self, _response: pb.ActivityResponse, timeout: float | None = None + ) -> None: self.completions += 1 @@ -67,6 +72,7 @@ async def _run_async_batch(n_items: int, timeout_s: float) -> int: worker = TaskHubGrpcWorker(host_address='in-process-mock', concurrency_options=options) manager = worker._async_worker_manager stub = _MockSidecarStub() + worker._set_connection(MagicMock(), cast(stubs.TaskHubSidecarServiceStub, stub)) async def activity(ctx, _inp) -> None: await asyncio.sleep(ACTIVITY_DURATION_SECONDS) @@ -78,7 +84,7 @@ async def activity(ctx, _inp) -> None: try: for task_id in range(n_items): req = _activity_request(task_id) - manager.submit_activity(worker._execute_activity_async, activity, req, stub, '') + manager.submit_activity(worker._execute_activity_async, activity, req, '') await asyncio.wait_for(manager.activity_queue.join(), timeout=timeout_s) finally: manager._shutdown = True diff --git a/tests/ext/workflow/durabletask/test_worker_history_cache.py b/tests/ext/workflow/durabletask/test_worker_history_cache.py index ec4b57e1b..0c4d729b9 100644 --- a/tests/ext/workflow/durabletask/test_worker_history_cache.py +++ b/tests/ext/workflow/durabletask/test_worker_history_cache.py @@ -16,7 +16,8 @@ of full vs delta work items with a GetInstanceHistory fallback on a cache miss. """ -from typing import cast +from typing import Union, cast +from unittest.mock import MagicMock import pytest @@ -51,7 +52,9 @@ def __init__(self, events: list[pb.HistoryEvent]) -> None: self._events = events self.get_instance_history_calls = 0 - def GetInstanceHistory(self, request: pb.GetInstanceHistoryRequest): + def GetInstanceHistory( + self, request: pb.GetInstanceHistoryRequest, timeout: float | None = None + ): self.get_instance_history_calls += 1 return pb.GetInstanceHistoryResponse(events=self._events) @@ -62,10 +65,14 @@ class _FailingHistoryStub: def __init__(self) -> None: self.responses: list[pb.WorkflowResponse] = [] - def GetInstanceHistory(self, request: pb.GetInstanceHistoryRequest): + def GetInstanceHistory( + self, request: pb.GetInstanceHistoryRequest, timeout: float | None = None + ): raise RuntimeError('sidecar unavailable') - def CompleteOrchestratorTask(self, response: pb.WorkflowResponse) -> None: + def CompleteOrchestratorTask( + self, response: pb.WorkflowResponse, timeout: float | None = None + ) -> None: self.responses.append(response) @@ -174,10 +181,15 @@ def _worker(**kwargs) -> TaskHubGrpcWorker: return TaskHubGrpcWorker(host_address='localhost:0', **kwargs) +def _connect(worker: TaskHubGrpcWorker, stub: Union[_FakeStub, _FailingHistoryStub]) -> None: + worker._set_connection(MagicMock(), cast(stubs.TaskHubSidecarServiceStub, stub)) + + def _resolve( worker: TaskHubGrpcWorker, req: pb.WorkflowRequest, stub: _FakeStub ) -> list[pb.HistoryEvent]: - return worker._resolve_history(req, cast(stubs.TaskHubSidecarServiceStub, stub)) + _connect(worker, stub) + return worker._resolve_history(req) def test_resolve_full_send_returns_past_events(): @@ -242,8 +254,9 @@ def test_resolve_fetch_failure_raises_history_resolution_error(): req = pb.WorkflowRequest(instanceId='a', pastEvents=_events(3)) req.cachedHistory.eventCount = 5 + _connect(worker, stub) with pytest.raises(_HistoryResolutionError): - worker._resolve_history(req, cast(stubs.TaskHubSidecarServiceStub, stub)) + worker._resolve_history(req) def test_fetch_failure_tears_down_stream_instead_of_failing_the_workflow(): @@ -258,9 +271,9 @@ def teardown_stream() -> None: nonlocal teardown_calls teardown_calls += 1 + _connect(worker, stub) worker._execute_orchestrator( req, - cast(stubs.TaskHubSidecarServiceStub, stub), 'completion-token', teardown_stream, ) diff --git a/tests/ext/workflow/durabletask/test_worker_reconnect.py b/tests/ext/workflow/durabletask/test_worker_reconnect.py new file mode 100644 index 000000000..70c926250 --- /dev/null +++ b/tests/ext/workflow/durabletask/test_worker_reconnect.py @@ -0,0 +1,404 @@ +# -*- coding: utf-8 -*- +# Copyright 2026 The Dapr Authors +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""How the worker behaves when its connection to the sidecar drops and comes back.""" + +import asyncio +import threading +import time +from collections.abc import Callable, Iterator, Sequence +from contextlib import contextmanager +from typing import Optional, Union +from unittest.mock import MagicMock, patch + +import grpc +import pytest + +import dapr.ext.workflow._durabletask.internal.helpers as ph +import dapr.ext.workflow._durabletask.internal.protos as pb +from dapr.ext.workflow._durabletask.worker import ( + _CONNECT_HELLO_TIMEOUT, + _RESULT_SEND_TIMEOUT, + TaskHubGrpcWorker, +) + + +class _RpcError(grpc.RpcError): + def __init__(self, code: grpc.StatusCode) -> None: + self._code = code + + def code(self) -> grpc.StatusCode: + return self._code + + def details(self) -> str: + return self._code.name + + +class _SidecarStub: + """One connection's stub: serves one work-item stream that stays open until ``end_stream``.""" + + def __init__( + self, + worker: TaskHubGrpcWorker, + hello_error: Optional[Exception] = None, + hello_gate: Optional[threading.Event] = None, + work_items: Sequence[pb.WorkItem] = (), + ): + self._worker = worker + self._hello_error = hello_error + self._hello_gate = hello_gate + self._work_items = list(work_items) + self.hello_timeouts: list[Optional[float]] = [] + self.results: list[Union[pb.ActivityResponse, pb.WorkflowResponse]] = [] + self.result_timeouts: list[Optional[float]] = [] + self.result_received = threading.Event() + self.stream_open = threading.Event() + self._stream_end = threading.Event() + + def Hello(self, *_args: object, timeout: Optional[float] = None) -> None: # noqa: N802 + self.hello_timeouts.append(timeout) + if self._hello_gate is not None: + assert self._hello_gate.wait(5) + if self._hello_error is not None: + raise self._hello_error + + def GetWorkItems(self, _request: pb.GetWorkItemsRequest) -> Iterator[pb.WorkItem]: # noqa: N802 + self.stream_open.set() + yield from self._work_items + while not self._stream_end.wait(0.01): + if self._worker._shutdown.is_set(): + return + yield from () + + def CompleteActivityTask( # noqa: N802 + self, response: pb.ActivityResponse, timeout: Optional[float] = None + ) -> None: + self._record_result(response, timeout) + + def CompleteOrchestratorTask( # noqa: N802 + self, response: pb.WorkflowResponse, timeout: Optional[float] = None + ) -> None: + self._record_result(response, timeout) + + def _record_result( + self, response: Union[pb.ActivityResponse, pb.WorkflowResponse], timeout: Optional[float] + ) -> None: + self.results.append(response) + self.result_timeouts.append(timeout) + self.result_received.set() + + def end_stream(self) -> None: + self._stream_end.set() + + +class _FakeSidecar: + """Hands the worker a new channel and stub for each connection it opens.""" + + def __init__(self, worker: TaskHubGrpcWorker, stubs: list[_SidecarStub]) -> None: + self._pending_stubs = list(stubs) + self.channels: list[MagicMock] = [] + self._stub_by_channel: dict[int, _SidecarStub] = {} + + def channel(self, *_args: object, **_kwargs: object) -> MagicMock: + new_channel = MagicMock(name=f'channel-{len(self.channels)}') + self._stub_by_channel[id(new_channel)] = self._pending_stubs.pop(0) + self.channels.append(new_channel) + return new_channel + + def stub(self, channel: MagicMock) -> _SidecarStub: + return self._stub_by_channel[id(channel)] + + +@contextmanager +def _running(worker: TaskHubGrpcWorker, sidecar: _FakeSidecar) -> Iterator[None]: + with ( + patch( + 'dapr.ext.workflow._durabletask.worker.shared.get_grpc_channel', + side_effect=sidecar.channel, + ), + patch( + 'dapr.ext.workflow._durabletask.worker.stubs.TaskHubSidecarServiceStub', + side_effect=sidecar.stub, + ), + ): + worker.start() + try: + yield + finally: + worker.stop() + + +def _worker() -> TaskHubGrpcWorker: + return TaskHubGrpcWorker(keepalive_interval=0, stop_timeout=5) + + +def test_connect_hello_has_a_deadline() -> None: + """A Hello that can block forever would stop the worker from ever reconnecting.""" + worker = _worker() + stub = _SidecarStub(worker) + sidecar = _FakeSidecar(worker, [stub]) + + with _running(worker, sidecar): + assert stub.stream_open.wait(5) + + assert stub.hello_timeouts == [_CONNECT_HELLO_TIMEOUT] + + +def test_failed_connect_closes_its_channel() -> None: + worker = _worker() + timed_out = _SidecarStub(worker, hello_error=_RpcError(grpc.StatusCode.DEADLINE_EXCEEDED)) + healthy = _SidecarStub(worker) + sidecar = _FakeSidecar(worker, [timed_out, healthy]) + + with _running(worker, sidecar): + assert healthy.stream_open.wait(10) + + sidecar.channels[0].close.assert_called_once() + + +def _wait_until(condition: Callable[[], bool], timeout: float = 5) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if condition(): + return True + time.sleep(0.01) + return condition() + + +def test_worker_is_not_ready_while_its_stream_is_down() -> None: + worker = _worker() + first = _SidecarStub(worker) + reconnect_gate = threading.Event() + second = _SidecarStub(worker, hello_gate=reconnect_gate) + sidecar = _FakeSidecar(worker, [first, second]) + + with _running(worker, sidecar): + assert worker.is_worker_ready() + + first.end_stream() + assert _wait_until(lambda: not worker.is_worker_ready()), 'still ready without a stream' + + reconnect_gate.set() + assert second.stream_open.wait(5) + assert _wait_until(worker.is_worker_ready), 'not ready again after reconnecting' + + assert not worker.is_worker_ready() + + +class _SlowWork: + """A registered activity or workflow that runs until the test releases it.""" + + def __init__(self, worker: TaskHubGrpcWorker, kind: str) -> None: + self.started = threading.Event() + self.release = threading.Event() + + def run(_ctx: object, _input: object) -> str: + self.started.set() + assert self.release.wait(5) + return 'done' + + async def run_async(_ctx: object, _input: object) -> str: + self.started.set() + while not self.release.is_set(): + await asyncio.sleep(0.01) + return 'done' + + if kind == 'workflow': + name = worker.add_orchestrator(run) + started = ph.new_execution_started_event(name, 'instance') + request = pb.WorkflowRequest(instanceId='instance', newEvents=[started]) + self.work_item = pb.WorkItem(workflowRequest=request) + return + name = worker.add_activity(run_async if kind == 'async activity' else run) + instance = pb.WorkflowInstance(instanceId='instance') + activity_request = pb.ActivityRequest(name=name, taskId=1, workflowInstance=instance) + self.work_item = pb.WorkItem(activityRequest=activity_request) + + +@pytest.mark.parametrize('kind', ['activity', 'async activity', 'workflow']) +def test_result_that_outlives_its_stream_is_sent_on_the_new_connection(kind: str) -> None: + """The work item's own channel is closed after a reconnect, so its result would be lost.""" + worker = _worker() + work = _SlowWork(worker, kind) + first = _SidecarStub(worker, work_items=[work.work_item]) + second = _SidecarStub(worker) + sidecar = _FakeSidecar(worker, [first, second]) + + with _running(worker, sidecar): + assert work.started.wait(5) + first.end_stream() + assert second.stream_open.wait(5) + + work.release.set() + assert second.result_received.wait(5), 'result never reached the new connection' + + assert first.results == [] + assert len(second.results) == 1 + assert second.result_timeouts == [_RESULT_SEND_TIMEOUT] + + +def test_result_ready_during_a_reconnect_waits_for_the_new_connection() -> None: + worker = _worker() + work = _SlowWork(worker, 'activity') + first = _SidecarStub(worker, work_items=[work.work_item]) + reconnect_gate = threading.Event() + second = _SidecarStub(worker, hello_gate=reconnect_gate) + sidecar = _FakeSidecar(worker, [first, second]) + + with _running(worker, sidecar): + assert work.started.wait(5) + first.end_stream() + assert _wait_until(lambda: not worker.is_worker_ready()) + work.release.set() + time.sleep(0.2) # the result is ready while there is no connection to send it on + assert second.results == [], 'result sent before the new connection passed its Hello' + + reconnect_gate.set() + assert second.result_received.wait(5), 'result was dropped instead of waiting' + + assert first.results == [] + + +def test_result_ready_during_a_failed_reconnect_waits_for_a_healthy_one() -> None: + worker = _worker() + work = _SlowWork(worker, 'activity') + first = _SidecarStub(worker, work_items=[work.work_item]) + reconnect_gate = threading.Event() + unavailable = _RpcError(grpc.StatusCode.UNAVAILABLE) + failing = _SidecarStub(worker, hello_error=unavailable, hello_gate=reconnect_gate) + healthy = _SidecarStub(worker) + sidecar = _FakeSidecar(worker, [first, failing, healthy]) + + with _running(worker, sidecar): + assert work.started.wait(5) + first.end_stream() + assert _wait_until(lambda: len(failing.hello_timeouts) == 1) + work.release.set() + time.sleep(0.2) # the result is ready while the failing connection is in its Hello + + reconnect_gate.set() + assert healthy.result_received.wait(10), 'result was lost on the failed connection' + + assert first.results == [] + assert failing.results == [] + + +def test_result_is_dropped_when_the_worker_stops_before_reconnecting() -> None: + worker = _worker() + stub = MagicMock() + worker._set_connection(MagicMock(), stub) + worker._retire_connection() + worker._shutdown.set() + + started = time.monotonic() + worker._deliver_activity_response(pb.ActivityRequest(), pb.ActivityResponse(), '', 'instance') + + assert time.monotonic() - started < 1, 'a stopping worker waited for a reconnect' + stub.CompleteActivityTask.assert_not_called() + + +def test_retired_connection_is_never_handed_out() -> None: + worker = _worker() + worker._set_connection(MagicMock(), MagicMock()) + worker._retire_connection() + + with worker._sidecar_stub(wait=0) as stub: + assert stub is None + + +def test_channel_is_closed_only_after_the_calls_in_flight_on_it() -> None: + """Starting a call on a closed channel with interceptors crashes the process in grpcio.""" + worker = _worker() + channel = MagicMock() + worker._set_connection(channel, MagicMock()) + + with worker._sidecar_stub(wait=0) as stub: + assert stub is not None + worker._retire_connection() + closer = threading.Thread(target=worker._close_channel, args=(channel,)) + closer.start() + closer.join(0.2) + assert closer.is_alive() + channel.close.assert_not_called() + + closer.join(5) + channel.close.assert_called_once() + + +def _stoppable_worker() -> TaskHubGrpcWorker: + worker = _worker() + worker._runLoop = MagicMock() + worker._runLoop.is_alive.return_value = False + return worker + + +def _hold_call(worker: TaskHubGrpcWorker, release: threading.Event) -> threading.Thread: + """Starts a thread that holds a call in flight on the current connection until released.""" + call_started = threading.Event() + + def held_call() -> None: + with worker._sidecar_stub(wait=0): + call_started.set() + release.wait(5) + + caller = threading.Thread(target=held_call) + caller.start() + assert call_started.wait(5) + return caller + + +def test_result_ready_while_the_worker_stops_is_still_delivered() -> None: + worker = _stoppable_worker() + stub = MagicMock() + worker._set_connection(MagicMock(), stub) + release = threading.Event() + caller = _hold_call(worker, release) + + stopper = threading.Thread(target=worker.stop) + stopper.start() + assert _wait_until(worker._shutdown.is_set) + worker._deliver_activity_response(pb.ActivityRequest(), pb.ActivityResponse(), '', 'instance') + release.set() + caller.join(5) + stopper.join(10) + + stub.CompleteActivityTask.assert_called_once() + + +def test_stop_never_closes_a_channel_with_a_call_in_flight() -> None: + """Closing early could land between a holder's two calls, and the second would crash.""" + worker = _stoppable_worker() + channel = MagicMock() + worker._set_connection(channel, MagicMock()) + release = threading.Event() + caller = _hold_call(worker, release) + + with patch('dapr.ext.workflow._durabletask.worker._STOP_DRAIN_TIMEOUT', 0.1): + stopper = threading.Thread(target=worker.stop) + stopper.start() + stopper.join(0.5) + assert stopper.is_alive() + channel.close.assert_not_called() + + release.set() + caller.join(5) + stopper.join(10) + channel.close.assert_called_once() + + +def test_connection_is_not_published_once_the_worker_is_stopping() -> None: + worker = _worker() + worker._shutdown.set() + + assert not worker._set_connection(MagicMock(), MagicMock()) + with worker._sidecar_stub(wait=0) as stub: + assert stub is None