diff --git a/dapr/ext/workflow/AGENTS.md b/dapr/ext/workflow/AGENTS.md index 50e694005..e2465823b 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`. + +**One channel for the worker's lifetime.** `start()` opens the gRPC channel and `stop()` closes it; gRPC reconnects it on its own, so when the work-item stream ends the run loop only sends a new `Hello` and `GetWorkItems`. Every call goes through `_sidecar_stub()`, which keeps `stop()` from closing the channel while a call is starting on it (with interceptors, grpcio crashes with SIGSEGV on a call started on a closed channel). Result sends use `wait_for_ready=True` with the `_RESULT_SEND_TIMEOUT` deadline, so a result that is ready during a reconnect waits for the channel instead of failing fast. - **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 906422c14..4a92d2a68 100644 --- a/dapr/ext/workflow/_durabletask/worker.py +++ b/dapr/ext/workflow/_durabletask/worker.py @@ -229,6 +229,22 @@ 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. The Hello waits for +# the channel to be ready, so a sidecar that comes back is seen as soon as gRPC reconnects to it, +# rather than after the run loop's own retry delay on top of gRPC's reconnect backoff. +_CONNECT_HELLO_TIMEOUT = 30.0 +# Deadline for each call that sends a result or fetches history. These calls wait for the channel +# to be ready, so a result that is ready while the worker reconnects goes out once it is back. +# stop() closes the channel only once no call is in flight, so this also bounds stop(). +_RESULT_SEND_TIMEOUT = 30.0 +# gRPC's own reconnect backoff grows to 120s by default, which would leave a worker idle for up to +# two minutes after a long sidecar outage. Capped to match the run loop's own retry delay. +_MAX_RECONNECT_BACKOFF_MS = 15_000 +# How long stop() keeps the channel open for results that are still being sent or that become +# ready while it stops, before it closes the channel. +_STOP_DRAIN_TIMEOUT = 5.0 class _HistoryResolutionError(Exception): @@ -467,8 +483,14 @@ def __init__( self._channel_options = channel_options self._stop_timeout = stop_timeout self._keepalive_interval = keepalive_interval - self._current_channel: Optional[grpc.Channel] = None # Store channel reference for cleanup - self._channel_cleanup_threads: list[threading.Thread] = [] # Deferred channel close threads + # One channel for the worker's lifetime: gRPC reconnects it on its own, so a lost stream + # only needs a new GetWorkItems call. Opened by start() and closed by stop(). + self._channel: Optional[grpc.Channel] = None + self._stub: Optional[stubs.TaskHubSidecarServiceStub] = None + # Guards the channel, the stub and the calls in flight on the channel; notified + # when the last call in flight finishes. + self._connection_lock = threading.Condition() + self._calls_in_flight = 0 self._stream_ready = threading.Event() self._runLoop: Optional[Thread] = None # Use provided concurrency options or create default ones @@ -521,6 +543,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): @@ -534,6 +561,7 @@ def run_loop(): loop.run_until_complete(self._async_run_loop()) self._logger.info(f'Starting gRPC worker that connects to {self._host_address}') + self._open_channel() if not self._disable_stateful_history: self._history_janitor = Thread( target=self._sweep_history_cache_loop, name='WorkerHistoryJanitor', daemon=True @@ -550,7 +578,12 @@ def run_loop(): ) self._is_running = True - async def _keepalive_loop(self, stub): + def _send_keepalive(self) -> None: + with self._sidecar_stub() 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 +591,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}') @@ -575,61 +608,40 @@ async def _async_run_loop(self): """ This is the main async loop that runs the worker. It is responsible for: - - Creating a fresh connection to the sidecar + - Connecting to the sidecar - Reading work items from the sidecar - Executing the work items - Shutting down the worker """ 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 - # 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 + def connect(): + nonlocal connected, conn_retry_count + connected = False 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}') + with self._sidecar_stub() as stub: + if stub is None: + raise RuntimeError('the worker is stopping') + stub.Hello( + empty_pb2.Empty(), timeout=_CONNECT_HELLO_TIMEOUT, wait_for_ready=True + ) except Exception as e: - self._logger.warning(f'Failed to create connection: {e}') - current_channel = None - self._current_channel = None - current_stub = None + self._logger.warning(f'Failed to connect: {e}') raise + connected = True + conn_retry_count = 0 + self._logger.info(f'Connected 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 + def disconnect(): + nonlocal connected, current_reader_thread + 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() @@ -638,7 +650,7 @@ def invalidate_connection(): current_reader_thread.join(timeout=5) current_reader_thread = None - def should_invalidate_connection(rpc_error): + def is_connection_error(rpc_error): error_code = rpc_error.code() # type: ignore connection_level_errors = { grpc.StatusCode.UNAVAILABLE, @@ -653,9 +665,9 @@ 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() + connect() except Exception: conn_retry_count += 1 delay = min( @@ -670,8 +682,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 +690,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() 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 +805,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 +831,6 @@ def stream_reader(): self._async_worker_manager.submit_orchestration( self._execute_orchestrator, work_item.workflowRequest, - stub, work_item.completionToken, teardown_stream, ) @@ -836,7 +849,6 @@ def stream_reader(): activity_handler, activity_fn, work_item.activityRequest, - stub, work_item.completionToken, ) elif work_item.HasField('healthPing'): @@ -850,7 +862,7 @@ def stream_reader(): except Exception as e: if self._async_worker_manager._shutdown or loop.is_closed(): break - invalidate_connection() + disconnect() raise e current_reader_thread.join(timeout=1) await self._cancel_keepalive(keepalive_task) @@ -863,18 +875,18 @@ def stream_reader(): self._logger.info( f'Work item stream ended. Will attempt to reconnect to {self._host_address}...' ) - invalidate_connection() + disconnect() # Fall through to the top of the outer loop, which will - # create a fresh connection (with retry/backoff if needed) + # reconnect (with retry/backoff if needed) except grpc.RpcError as rpc_error: await self._cancel_keepalive(keepalive_task) # Check shutdown first - if shutting down, exit immediately if self._shutdown.is_set(): self._logger.debug('Shutdown detected during RPC error handling, exiting') break - should_invalidate = should_invalidate_connection(rpc_error) - if should_invalidate: - invalidate_connection() + connection_lost = is_connection_error(rpc_error) + if connection_lost: + disconnect() error_code = rpc_error.code() # type: ignore error_detail = ( rpc_error.details() if hasattr(rpc_error, 'details') else str(rpc_error) @@ -883,7 +895,7 @@ def stream_reader(): if error_code == grpc.StatusCode.CANCELLED: self._logger.info(f'Disconnected from {self._host_address}') break - elif should_invalidate: + elif connection_lost: self._logger.warning( f'Connection error ({error_code}): {error_detail} — resetting connection' ) @@ -932,9 +944,9 @@ def stream_reader(): # No event loop running, treat as shutdown self._logger.debug(f'No event loop running, exiting: {ex}') break - invalidate_connection() + disconnect() self._logger.warning(f'Unexpected error: {ex}') - invalidate_connection() + disconnect() self._logger.info('No longer listening for work items') # Cancel worker_task to ensure shutdown completes even if tasks are still running @@ -949,37 +961,80 @@ def stream_reader(): except Exception as e: self._logger.warning(f'Error while waiting for worker task shutdown: {e}') - def _schedule_deferred_channel_close( - self, old_channel: grpc.Channel, grace_timeout: float = 10.0 - ): - """Schedule a deferred close of an old gRPC channel. + def _open_channel(self) -> None: + options = [ + ('grpc.max_reconnect_backoff_ms', _MAX_RECONNECT_BACKOFF_MS), + *(self._channel_options or ()), + ] + channel = shared.get_grpc_channel( + self._host_address, self._secure_channel, self._interceptors, options=options + ) + stub = stubs.TaskHubSidecarServiceStub(channel) + with self._connection_lock: + self._channel = channel + self._stub = stub - Waits up to *grace_timeout* seconds for in-flight RPCs to complete - before closing the channel. This prevents orphaned TCP connections - while still allowing in-flight work (e.g. ``CompleteActivityTask`` - calls on another thread) to finish gracefully. + @contextlib.contextmanager + def _sidecar_stub(self) -> Iterator[Optional[stubs.TaskHubSidecarServiceStub]]: + """Yields the worker's stub and keeps the channel open until the block exits. - During ``stop()``, ``_shutdown`` is already set so the wait returns - immediately and the channel is closed at once. + Yields: + The stub, or None once stop() has started closing the channel. """ - # Prune already-finished cleanup threads to avoid unbounded growth - self._channel_cleanup_threads = [t for t in self._channel_cleanup_threads if t.is_alive()] - - def _deferred_close(): - try: - # Normal reconnect: wait grace period for RPCs to drain. - # Shutdown: _shutdown is already set, returns immediately. - self._shutdown.wait(timeout=grace_timeout) - finally: - try: - old_channel.close() - self._logger.debug('Deferred channel close completed') - except Exception as e: - self._logger.debug(f'Error during deferred channel close: {e}') - - thread = threading.Thread(target=_deferred_close, daemon=True, name='ChannelCleanup') - thread.start() - self._channel_cleanup_threads.append(thread) + with self._connection_lock: + stub = self._stub + if stub is not None: + self._calls_in_flight += 1 + try: + yield stub + finally: + if stub is not None: + self._release_call() + + def _release_call(self) -> None: + with self._connection_lock: + self._calls_in_flight -= 1 + if self._calls_in_flight == 0: + self._connection_lock.notify_all() + + def _close_channel(self) -> None: + """Closes the 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``, which stops handing out the stub here, so once the count reaches + zero no call can start on the channel 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. + """ + with self._connection_lock: + channel = self._channel + self._channel = None + self._stub = None + self._connection_lock.wait_for(lambda: self._calls_in_flight == 0) + if channel is not None: + channel.close() + + def _drain_calls(self, timeout: float) -> None: + """Waits up to ``timeout`` seconds for no call to be in flight on the channel.""" + with self._connection_lock: + self._connection_lock.wait_for(lambda: self._calls_in_flight == 0, timeout=timeout) + + def _send_result( + self, send: Callable[[stubs.TaskHubSidecarServiceStub], None], what: str + ) -> None: + """Calls ``send`` with the worker's stub, unless the worker has closed its channel.""" + with self._sidecar_stub() as stub: + if stub is None: + self._logger.warning( + f'Could not deliver {what}: the worker has stopped. ' + 'The sidecar will re-dispatch this work item.' + ) + return + send(stub) def stop(self): """Stops the worker and waits for any pending work items to complete.""" @@ -990,29 +1045,23 @@ def stop(self): self._logger.info('Stopping gRPC worker...') self._shutdown.set() - # Close the channel — propagates cancellation to all streams and cleans up resources - if self._current_channel is not None: - try: - self._current_channel.close() - 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) - if self._runLoop.is_alive(): - self._logger.warning( - f'Worker thread did not complete within {self._stop_timeout}s timeout. ' - 'Some resources may not be fully cleaned up.' - ) - else: - self._logger.debug('Worker thread completed successfully') + # Results that are being sent, or that become ready now, still go out on the channel. + self._drain_calls(_STOP_DRAIN_TIMEOUT) + # Closing the channel is what ends the work-item stream: through an interceptor, the + # stream is a plain generator with no cancel(). + try: + self._close_channel() + except Exception as e: + self._logger.exception(f'Error closing gRPC channel: {e}') - # Wait for any deferred channel-cleanup threads to finish - for t in self._channel_cleanup_threads: - t.join(timeout=5) - self._channel_cleanup_threads.clear() + self._runLoop.join(timeout=self._stop_timeout) + if self._runLoop.is_alive(): + self._logger.warning( + f'Worker thread did not complete within {self._stop_timeout}s timeout. ' + 'Some resources may not be fully cleaned up.' + ) + else: + self._logger.debug('Worker thread completed successfully') self._async_worker_manager.shutdown() self._logger.info('Worker shutdown completed') @@ -1068,9 +1117,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 +1136,12 @@ def _resolve_history( history_request = pb.GetInstanceHistoryRequest(instanceId=req.instanceId) try: - response = stub.GetInstanceHistory(history_request) + with self._sidecar_stub() as stub: + if stub is None: + raise RuntimeError('not connected to the sidecar') + response = stub.GetInstanceHistory( + history_request, timeout=_RESULT_SEND_TIMEOUT, wait_for_ready=True + ) except Exception as ex: raise _HistoryResolutionError( f"Failed to fetch the committed history for '{req.instanceId}': {ex}" @@ -1118,7 +1170,6 @@ def _update_history_cache( def _execute_orchestrator( self, req: pb.WorkflowRequest, - stub: stubs.TaskHubSidecarServiceStub, completionToken, teardown_stream: Callable[[], None], ): @@ -1129,7 +1180,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 +1225,22 @@ def _execute_orchestrator( completionToken=completionToken, ) + self._send_result( + 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, wait_for_ready=True) 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 +1267,9 @@ def _execute_orchestrator( completionToken=completionToken, ) try: - stub.CompleteOrchestratorTask(failure_res) + stub.CompleteOrchestratorTask( + failure_res, timeout=_RESULT_SEND_TIMEOUT, wait_for_ready=True + ) except Exception as ex: self._logger.exception( f"Failed to deliver orchestrator failure response for '{req.instanceId}': {ex}" @@ -1281,7 +1348,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, wait_for_ready=True) 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 +1364,9 @@ def _send_activity_response( req, instance_id, oversize_error, completion_token ) try: - stub.CompleteActivityTask(failure_res) + stub.CompleteActivityTask( + failure_res, timeout=_RESULT_SEND_TIMEOUT, wait_for_ready=True + ) except Exception as ex: self._logger.exception( f"Failed to deliver activity failure response for '{req.name}#{req.taskId}' " @@ -1317,11 +1386,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_result( + 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 +1423,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 +1457,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..078e981ad 100644 --- a/tests/ext/workflow/durabletask/test_async_dispatch_regression.py +++ b/tests/ext/workflow/durabletask/test_async_dispatch_regression.py @@ -17,9 +17,11 @@ """ import asyncio +from typing import cast 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 +42,12 @@ 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, + wait_for_ready: bool | None = None, + ) -> None: self.completions += 1 @@ -67,6 +74,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._stub = cast(stubs.TaskHubSidecarServiceStub, stub) async def activity(ctx, _inp) -> None: await asyncio.sleep(ACTIVITY_DURATION_SECONDS) @@ -78,7 +86,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..7588053d3 100644 --- a/tests/ext/workflow/durabletask/test_worker_history_cache.py +++ b/tests/ext/workflow/durabletask/test_worker_history_cache.py @@ -16,7 +16,7 @@ of full vs delta work items with a GetInstanceHistory fallback on a cache miss. """ -from typing import cast +from typing import Union, cast import pytest @@ -51,7 +51,12 @@ 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, + wait_for_ready: bool | None = None, + ): self.get_instance_history_calls += 1 return pb.GetInstanceHistoryResponse(events=self._events) @@ -62,10 +67,20 @@ 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, + wait_for_ready: bool | None = None, + ): raise RuntimeError('sidecar unavailable') - def CompleteOrchestratorTask(self, response: pb.WorkflowResponse) -> None: + def CompleteOrchestratorTask( + self, + response: pb.WorkflowResponse, + timeout: float | None = None, + wait_for_ready: bool | None = None, + ) -> None: self.responses.append(response) @@ -174,10 +189,15 @@ def _worker(**kwargs) -> TaskHubGrpcWorker: return TaskHubGrpcWorker(host_address='localhost:0', **kwargs) +def _connect(worker: TaskHubGrpcWorker, stub: Union[_FakeStub, _FailingHistoryStub]) -> None: + worker._stub = 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 +262,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 +279,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..049426229 --- /dev/null +++ b/tests/ext/workflow/durabletask/test_worker_reconnect.py @@ -0,0 +1,351 @@ +# -*- 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, + _MAX_RECONNECT_BACKOFF_MS, + _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: + """The worker's one stub: serves a work-item stream per connect, each open until ended. + + ``hello_errors`` and ``hello_gates`` apply to successive Hellos, so a test can fail or hold + a reconnect while earlier connects succeed. + """ + + def __init__( + self, + worker: TaskHubGrpcWorker, + hello_errors: Sequence[Optional[Exception]] = (), + hello_gates: Sequence[Optional[threading.Event]] = (), + work_items: Sequence[pb.WorkItem] = (), + ): + self._worker = worker + self._hello_errors = list(hello_errors) + self._hello_gates = list(hello_gates) + self._work_items = list(work_items) + self.hello_timeouts: list[Optional[float]] = [] + self.hello_waits_for_ready: list[bool] = [] + self.results: list[Union[pb.ActivityResponse, pb.WorkflowResponse]] = [] + self.result_kwargs: list[dict[str, object]] = [] + self.result_received = threading.Event() + self.streams_opened = 0 + self._stream_end = threading.Event() + + def Hello( # noqa: N802 + self, *_args: object, timeout: Optional[float] = None, wait_for_ready: bool = False + ) -> None: + self.hello_timeouts.append(timeout) + self.hello_waits_for_ready.append(wait_for_ready) + gate = self._hello_gates.pop(0) if self._hello_gates else None + if gate is not None: + assert gate.wait(5) + error = self._hello_errors.pop(0) if self._hello_errors else None + if error is not None: + raise error + + def GetWorkItems(self, _request: pb.GetWorkItemsRequest) -> Iterator[pb.WorkItem]: # noqa: N802 + self.streams_opened += 1 + self._stream_end.clear() + work_items, self._work_items = self._work_items, [] + yield from work_items + while not self._stream_end.wait(0.01): + if self._worker._shutdown.is_set(): + return + + def CompleteActivityTask( # noqa: N802 + self, response: pb.ActivityResponse, **kwargs: object + ) -> None: + self._record_result(response, kwargs) + + def CompleteOrchestratorTask( # noqa: N802 + self, response: pb.WorkflowResponse, **kwargs: object + ) -> None: + self._record_result(response, kwargs) + + def _record_result( + self, response: Union[pb.ActivityResponse, pb.WorkflowResponse], kwargs: dict[str, object] + ) -> None: + self.results.append(response) + self.result_kwargs.append(kwargs) + self.result_received.set() + + def end_stream(self) -> None: + self._stream_end.set() + + +@contextmanager +def _running(worker: TaskHubGrpcWorker, stub: _SidecarStub) -> Iterator[MagicMock]: + """Runs ``worker`` against ``stub``, yielding the patched channel factory.""" + with ( + patch('dapr.ext.workflow._durabletask.worker.shared.get_grpc_channel') as get_channel, + patch( + 'dapr.ext.workflow._durabletask.worker.stubs.TaskHubSidecarServiceStub', + return_value=stub, + ), + ): + worker.start() + try: + yield get_channel + finally: + worker.stop() + + +def _worker() -> TaskHubGrpcWorker: + return TaskHubGrpcWorker(keepalive_interval=0, stop_timeout=5) + + +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_connect_hello_has_a_deadline() -> None: + """A Hello that can block forever would stop the worker from ever reconnecting.""" + worker = _worker() + stub = _SidecarStub(worker) + + with _running(worker, stub): + assert _wait_until(lambda: stub.streams_opened == 1) + + assert stub.hello_timeouts == [_CONNECT_HELLO_TIMEOUT] + assert stub.hello_waits_for_ready == [True] + + +def test_reconnects_reuse_the_one_channel() -> None: + """gRPC reconnects a channel on its own, so only the work-item stream is re-opened.""" + worker = _worker() + timed_out = _RpcError(grpc.StatusCode.DEADLINE_EXCEEDED) + stub = _SidecarStub(worker, hello_errors=[None, timed_out]) + + with _running(worker, stub) as get_channel: + assert _wait_until(lambda: stub.streams_opened == 1) + stub.end_stream() + assert _wait_until(lambda: stub.streams_opened == 2, timeout=10) + channel = get_channel.return_value + channel.close.assert_not_called() + + get_channel.assert_called_once() + channel.close.assert_called_once() + + +def test_channel_caps_grpc_reconnect_backoff_and_keeps_caller_options() -> None: + worker = TaskHubGrpcWorker(channel_options=[('grpc.max_reconnect_backoff_ms', 1_000)]) + + with patch('dapr.ext.workflow._durabletask.worker.shared.get_grpc_channel') as get_channel: + worker._open_channel() + + options = get_channel.call_args.kwargs['options'] + assert options[0] == ('grpc.max_reconnect_backoff_ms', _MAX_RECONNECT_BACKOFF_MS) + assert options[-1] == ('grpc.max_reconnect_backoff_ms', 1_000), 'caller option must win' + + +def test_worker_is_not_ready_while_its_stream_is_down() -> None: + worker = _worker() + reconnect_gate = threading.Event() + stub = _SidecarStub(worker, hello_gates=[None, reconnect_gate]) + + with _running(worker, stub): + assert worker.is_worker_ready() + + stub.end_stream() + assert _wait_until(lambda: not worker.is_worker_ready()), 'still ready without a stream' + + reconnect_gate.set() + assert _wait_until(lambda: stub.streams_opened == 2) + 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_still_delivered(kind: str) -> None: + worker = _worker() + work = _SlowWork(worker, kind) + stub = _SidecarStub(worker, work_items=[work.work_item]) + + with _running(worker, stub): + assert work.started.wait(5) + stub.end_stream() + assert _wait_until(lambda: stub.streams_opened == 2) + + work.release.set() + assert stub.result_received.wait(5), 'result was not delivered' + + assert len(stub.results) == 1 + + +@pytest.mark.parametrize('kind', ['activity', 'workflow']) +def test_result_waits_for_the_channel_within_a_deadline(kind: str) -> None: + """A result ready during a reconnect waits for the channel instead of failing fast.""" + worker = _worker() + work = _SlowWork(worker, kind) + stub = _SidecarStub(worker, work_items=[work.work_item]) + + with _running(worker, stub): + assert work.started.wait(5) + work.release.set() + assert stub.result_received.wait(5) + + assert stub.result_kwargs == [{'timeout': _RESULT_SEND_TIMEOUT, 'wait_for_ready': True}] + + +def _worker_with_channel() -> tuple[TaskHubGrpcWorker, MagicMock, MagicMock]: + worker = _worker() + worker._runLoop = MagicMock() + worker._runLoop.is_alive.return_value = False + channel = MagicMock() + stub = MagicMock() + worker._channel = channel + worker._stub = stub + return worker, channel, stub + + +def test_result_is_dropped_once_the_channel_is_closed() -> None: + worker, _channel, stub = _worker_with_channel() + worker._close_channel() + + worker._deliver_activity_response(pb.ActivityRequest(), pb.ActivityResponse(), '', 'instance') + + stub.CompleteActivityTask.assert_not_called() + with worker._sidecar_stub() as handed_out: + assert handed_out is None + + +def test_channel_is_closed_only_after_the_calls_in_flight_on_it() -> None: + """Closing early could land between a holder's two calls, and the second would crash.""" + worker, channel, _stub = _worker_with_channel() + + with worker._sidecar_stub() as stub: + assert stub is not None + closer = threading.Thread(target=worker._close_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 test_stop_closes_the_channel_after_a_call_in_flight_finishes() -> None: + worker, channel, _stub = _worker_with_channel() + release = threading.Event() + call_started = threading.Event() + + def held_call() -> None: + with worker._sidecar_stub(): + call_started.set() + release.wait(5) + + caller = threading.Thread(target=held_call) + caller.start() + assert call_started.wait(5) + 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(), 'stop() closed the channel under a call in flight' + channel.close.assert_not_called() + + release.set() + caller.join(5) + stopper.join(5) + channel.close.assert_called_once() + + +def test_result_ready_while_the_worker_stops_is_still_delivered() -> None: + worker, _channel, stub = _worker_with_channel() + release = threading.Event() + call_started = threading.Event() + + def held_call() -> None: + with worker._sidecar_stub(): + call_started.set() + release.wait(5) + + caller = threading.Thread(target=held_call) + caller.start() + assert call_started.wait(5) + 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() diff --git a/tests/ext/workflow/durabletask/test_worker_stop.py b/tests/ext/workflow/durabletask/test_worker_stop.py index 009505c13..f7d3c73ba 100644 --- a/tests/ext/workflow/durabletask/test_worker_stop.py +++ b/tests/ext/workflow/durabletask/test_worker_stop.py @@ -19,10 +19,10 @@ def _make_running_worker(): def test_stop_with_grpc_future(): - """stop() closes the channel, which propagates cancellation to all streams.""" + """stop() cancels the work-item stream and closes the channel.""" worker = _make_running_worker() mock_channel = MagicMock() - worker._current_channel = mock_channel + worker._channel = mock_channel worker._response_stream = MagicMock(spec=grpc.Future) worker.stop() mock_channel.close.assert_called_once() @@ -32,7 +32,7 @@ def test_stop_with_generator_call(): """stop() closes the channel even when response stream has a .call attribute.""" worker = _make_running_worker() mock_channel = MagicMock() - worker._current_channel = mock_channel + worker._channel = mock_channel mock_stream = MagicMock() mock_stream.call = MagicMock() worker._response_stream = mock_stream @@ -44,7 +44,7 @@ def test_stop_with_unknown_stream_type(): """stop() closes the channel regardless of response stream type.""" worker = _make_running_worker() mock_channel = MagicMock() - worker._current_channel = mock_channel + worker._channel = mock_channel worker._response_stream = object() worker.stop() mock_channel.close.assert_called_once() @@ -53,7 +53,7 @@ def test_stop_with_unknown_stream_type(): def test_stop_with_none_stream(): worker = _make_running_worker() worker._response_stream = None - worker._current_channel = None + worker._channel = None # Should not raise worker.stop() @@ -72,85 +72,10 @@ def test_stop_channel_close_handles_exception(caplog): worker = _make_running_worker() mock_channel = MagicMock() mock_channel.close.side_effect = Exception('close failed') - worker._current_channel = mock_channel + worker._channel = mock_channel # Should not raise worker.stop() - assert worker._current_channel is None - - -def test_deferred_channel_close_waits_then_closes(): - """_schedule_deferred_channel_close waits grace period, then closes old channel.""" - worker = TaskHubGrpcWorker() - old_channel = MagicMock() - - worker._schedule_deferred_channel_close(old_channel, grace_timeout=0.1) - # Thread should be tracked - assert len(worker._channel_cleanup_threads) == 1 - - # Wait for the grace period to expire and the thread to finish - worker._channel_cleanup_threads[0].join(timeout=2) - old_channel.close.assert_called_once() - - -def test_deferred_channel_close_fires_immediately_on_shutdown(): - """Deferred close returns immediately when shutdown is already set.""" - worker = TaskHubGrpcWorker() - worker._shutdown.set() - old_channel = MagicMock() - - worker._schedule_deferred_channel_close(old_channel, grace_timeout=60) - # Even with a 60s grace, shutdown makes it return immediately - worker._channel_cleanup_threads[0].join(timeout=2) - old_channel.close.assert_called_once() - - -def test_deferred_channel_close_handles_close_exception(): - """Deferred close handles exceptions from channel.close() gracefully.""" - worker = TaskHubGrpcWorker() - worker._shutdown.set() - old_channel = MagicMock() - old_channel.close.side_effect = Exception('already closed') - - # Should not raise - worker._schedule_deferred_channel_close(old_channel, grace_timeout=0) - worker._channel_cleanup_threads[0].join(timeout=2) - old_channel.close.assert_called_once() - - -def test_stop_joins_deferred_cleanup_threads(): - """stop() joins all deferred channel cleanup threads.""" - worker = _make_running_worker() - mock_channel = MagicMock() - worker._current_channel = mock_channel - - # Pre-populate a cleanup thread (simulating a prior reconnection) - old_channel = MagicMock() - worker._schedule_deferred_channel_close(old_channel, grace_timeout=60) - assert len(worker._channel_cleanup_threads) == 1 - - worker.stop() - # stop() sets shutdown, which unblocks deferred close threads - # After stop(), cleanup threads list should be cleared - assert len(worker._channel_cleanup_threads) == 0 - old_channel.close.assert_called_once() - mock_channel.close.assert_called_once() - - -def test_deferred_close_prunes_finished_threads(): - """_schedule_deferred_channel_close prunes already-finished threads.""" - worker = TaskHubGrpcWorker() - worker._shutdown.set() # Make threads complete immediately - - ch1 = MagicMock() - ch2 = MagicMock() - worker._schedule_deferred_channel_close(ch1, grace_timeout=0) - worker._channel_cleanup_threads[0].join(timeout=2) - - # ch1's thread is finished; scheduling ch2 should prune it - worker._schedule_deferred_channel_close(ch2, grace_timeout=0) - worker._channel_cleanup_threads[-1].join(timeout=2) - # Only the still-alive (or just-finished ch2) thread remains; ch1's was pruned - assert len(worker._channel_cleanup_threads) <= 1 + assert worker._channel is None def test_stop_before_start_is_noop(): @@ -164,7 +89,7 @@ def test_stop_before_start_is_noop(): def test_stop_is_idempotent(): """A second stop() returns early because _runLoop was cleared by the first.""" worker = _make_running_worker() - worker._current_channel = MagicMock() + worker._channel = MagicMock() worker.stop() assert worker._runLoop is None with patch.object(worker._shutdown, 'set') as shutdown_set: