From 3fed0cce54f08f7cee01249f8f326e10e96b3d8f Mon Sep 17 00:00:00 2001 From: Fabian Martinez <46371672+famarting@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:14:59 +0200 Subject: [PATCH 1/6] 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> --- dapr/ext/workflow/_durabletask/worker.py | 8 +- .../durabletask/test_worker_reconnect.py | 128 ++++++++++++++++++ 2 files changed, 135 insertions(+), 1 deletion(-) create mode 100644 tests/ext/workflow/durabletask/test_worker_reconnect.py diff --git a/dapr/ext/workflow/_durabletask/worker.py b/dapr/ext/workflow/_durabletask/worker.py index 906422c14..905ebd104 100644 --- a/dapr/ext/workflow/_durabletask/worker.py +++ b/dapr/ext/workflow/_durabletask/worker.py @@ -229,6 +229,10 @@ 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 class _HistoryResolutionError(Exception): @@ -608,11 +612,13 @@ def create_fresh_connection(): # Store channel reference for cleanup in stop() self._current_channel = current_channel current_stub = stubs.TaskHubSidecarServiceStub(current_channel) - current_stub.Hello(empty_pb2.Empty()) + current_stub.Hello(empty_pb2.Empty(), timeout=_CONNECT_HELLO_TIMEOUT) conn_retry_count = 0 self._logger.info(f'Created fresh connection to {self._host_address}') except Exception as e: self._logger.warning(f'Failed to create connection: {e}') + if current_channel is not None: + self._schedule_deferred_channel_close(current_channel) current_channel = None self._current_channel = None current_stub = None 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..504f35cc5 --- /dev/null +++ b/tests/ext/workflow/durabletask/test_worker_reconnect.py @@ -0,0 +1,128 @@ +# -*- 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 threading +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Optional +from unittest.mock import MagicMock, patch + +import grpc + +import dapr.ext.workflow._durabletask.internal.protos as pb +from dapr.ext.workflow._durabletask.worker import _CONNECT_HELLO_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): + self._worker = worker + self._hello_error = hello_error + self.hello_timeouts: list[Optional[float]] = [] + 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_error is not None: + raise self._hello_error + + def GetWorkItems(self, _request: pb.GetWorkItemsRequest) -> Iterator[pb.WorkItem]: # noqa: N802 + self.stream_open.set() + while not self._stream_end.wait(0.01): + if self._worker._shutdown.is_set(): + return + yield from () + + 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: + worker = TaskHubGrpcWorker(keepalive_interval=0, stop_timeout=5) + worker._logger = MagicMock() + return worker + + +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() From d7bcbc5ee90b0be387b52a90f62a899acc8d120a Mon Sep 17 00:00:00 2001 From: Fabian Martinez <46371672+famarting@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:15:39 +0200 Subject: [PATCH 2/6] 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> --- dapr/ext/workflow/_durabletask/worker.py | 6 +++ .../durabletask/test_worker_reconnect.py | 42 ++++++++++++++++++- 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/dapr/ext/workflow/_durabletask/worker.py b/dapr/ext/workflow/_durabletask/worker.py index 905ebd104..21fd16165 100644 --- a/dapr/ext/workflow/_durabletask/worker.py +++ b/dapr/ext/workflow/_durabletask/worker.py @@ -525,6 +525,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): @@ -636,6 +641,7 @@ def invalidate_connection(): self._current_channel = None current_stub = None 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() diff --git a/tests/ext/workflow/durabletask/test_worker_reconnect.py b/tests/ext/workflow/durabletask/test_worker_reconnect.py index 504f35cc5..64a860d39 100644 --- a/tests/ext/workflow/durabletask/test_worker_reconnect.py +++ b/tests/ext/workflow/durabletask/test_worker_reconnect.py @@ -13,7 +13,8 @@ """How the worker behaves when its connection to the sidecar drops and comes back.""" import threading -from collections.abc import Iterator +import time +from collections.abc import Callable, Iterator from contextlib import contextmanager from typing import Optional from unittest.mock import MagicMock, patch @@ -38,15 +39,23 @@ def details(self) -> str: 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): + def __init__( + self, + worker: TaskHubGrpcWorker, + hello_error: Optional[Exception] = None, + hello_gate: Optional[threading.Event] = None, + ): self._worker = worker self._hello_error = hello_error + self._hello_gate = hello_gate self.hello_timeouts: list[Optional[float]] = [] 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 @@ -126,3 +135,32 @@ def test_failed_connect_closes_its_channel() -> None: 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() From b96c73cb84b10f4d3c6c516ea95d8aef5c4e0bc6 Mon Sep 17 00:00:00 2001 From: Fabian Martinez <46371672+famarting@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:20:03 +0200 Subject: [PATCH 3/6] 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> --- dapr/ext/workflow/AGENTS.md | 4 +- dapr/ext/workflow/_durabletask/worker.py | 122 +++++++++++++++--- .../test_async_dispatch_regression.py | 6 +- .../durabletask/test_worker_history_cache.py | 15 ++- .../durabletask/test_worker_reconnect.py | 110 +++++++++++++++- 5 files changed, 225 insertions(+), 32 deletions(-) diff --git a/dapr/ext/workflow/AGENTS.md b/dapr/ext/workflow/AGENTS.md index 50e694005..3ef1b13b9 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 21fd16165..9486582a8 100644 --- a/dapr/ext/workflow/_durabletask/worker.py +++ b/dapr/ext/workflow/_durabletask/worker.py @@ -233,6 +233,9 @@ def _is_message_too_large(rpc_error: grpc.RpcError) -> bool: # 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 class _HistoryResolutionError(Exception): @@ -472,6 +475,9 @@ 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; notified whenever they change. + self._connection_lock = threading.Condition() self._channel_cleanup_threads: list[threading.Thread] = [] # Deferred channel close threads self._stream_ready = threading.Event() self._runLoop: Optional[Thread] = None @@ -614,9 +620,9 @@ def create_fresh_connection(): self._interceptors, options=self._channel_options, ) - # Store channel reference for cleanup in stop() - self._current_channel = current_channel current_stub = stubs.TaskHubSidecarServiceStub(current_channel) + # Publish the connection for result sends, and for cleanup in stop() + self._set_connection(current_channel, current_stub) current_stub.Hello(empty_pb2.Empty(), timeout=_CONNECT_HELLO_TIMEOUT) conn_retry_count = 0 self._logger.info(f'Created fresh connection to {self._host_address}') @@ -625,7 +631,7 @@ def create_fresh_connection(): if current_channel is not None: self._schedule_deferred_channel_close(current_channel) current_channel = None - self._current_channel = None + self._retire_connection() current_stub = None raise @@ -638,7 +644,7 @@ def invalidate_connection(): if current_channel is not None: self._schedule_deferred_channel_close(current_channel) current_channel = None - self._current_channel = None + self._retire_connection() current_stub = None self._response_stream = None self._stream_ready.clear() @@ -829,7 +835,6 @@ def stream_reader(): self._async_worker_manager.submit_orchestration( self._execute_orchestrator, work_item.workflowRequest, - stub, work_item.completionToken, teardown_stream, ) @@ -848,7 +853,6 @@ def stream_reader(): activity_handler, activity_fn, work_item.activityRequest, - stub, work_item.completionToken, ) elif work_item.HasField('healthPing'): @@ -961,6 +965,58 @@ 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) -> None: + """Makes ``channel`` the connection that work-item results are sent on.""" + with self._connection_lock: + self._current_channel = channel + self._current_stub = stub + self._connection_lock.notify_all() + + 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 + is stopping. + """ + with self._connection_lock: + self._connection_lock.wait_for(self._has_connection_or_is_stopping, timeout=wait) + stub = None if self._shutdown.is_set() else self._current_stub + yield stub + + 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 ): @@ -1003,13 +1059,12 @@ 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: + channel = self._retire_connection() + if channel is not None: try: - self._current_channel.close() + 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) @@ -1080,9 +1135,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 @@ -1101,7 +1154,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) except Exception as ex: raise _HistoryResolutionError( f"Failed to fetch the committed history for '{req.instanceId}': {ex}" @@ -1130,7 +1186,6 @@ def _update_history_cache( def _execute_orchestrator( self, req: pb.WorkflowRequest, - stub: stubs.TaskHubSidecarServiceStub, completionToken, teardown_stream: Callable[[], None], ): @@ -1141,7 +1196,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 ) @@ -1186,6 +1241,20 @@ 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) except grpc.RpcError as rpc_error: # type: ignore @@ -1329,11 +1398,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 @@ -1353,13 +1435,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. @@ -1388,9 +1469,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..0294cb00d 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 @@ -67,6 +70,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 +82,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..f60c3694b 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 @@ -174,10 +175,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 +248,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 +265,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 index 64a860d39..1df46c95a 100644 --- a/tests/ext/workflow/durabletask/test_worker_reconnect.py +++ b/tests/ext/workflow/durabletask/test_worker_reconnect.py @@ -12,15 +12,18 @@ """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 +from collections.abc import Callable, Iterator, Sequence from contextlib import contextmanager -from typing import Optional +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, TaskHubGrpcWorker @@ -44,11 +47,15 @@ def __init__( 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_received = threading.Event() self.stream_open = threading.Event() self._stream_end = threading.Event() @@ -61,11 +68,20 @@ def Hello(self, *_args: object, timeout: Optional[float] = None) -> None: # noq 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(self, response: pb.ActivityResponse) -> None: # noqa: N802 + self.results.append(response) + self.result_received.set() + + def CompleteOrchestratorTask(self, response: pb.WorkflowResponse) -> None: # noqa: N802 + self.results.append(response) + self.result_received.set() + def end_stream(self) -> None: self._stream_end.set() @@ -108,9 +124,7 @@ def _running(worker: TaskHubGrpcWorker, sidecar: _FakeSidecar) -> Iterator[None] def _worker() -> TaskHubGrpcWorker: - worker = TaskHubGrpcWorker(keepalive_interval=0, stop_timeout=5) - worker._logger = MagicMock() - return worker + return TaskHubGrpcWorker(keepalive_interval=0, stop_timeout=5) def test_connect_hello_has_a_deadline() -> None: @@ -164,3 +178,89 @@ def test_worker_is_not_ready_while_its_stream_is_down() -> None: 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 + + +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 + + reconnect_gate.set() + assert second.result_received.wait(5), 'result was dropped instead of waiting' + + assert first.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() From 72b85840b2c75db74339c86de35557435ccec9bb Mon Sep 17 00:00:00 2001 From: Fabian Martinez <46371672+famarting@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:21:49 +0200 Subject: [PATCH 4/6] 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> --- dapr/ext/workflow/_durabletask/worker.py | 116 ++++++++++++------ .../durabletask/test_worker_reconnect.py | 53 ++++++++ 2 files changed, 132 insertions(+), 37 deletions(-) diff --git a/dapr/ext/workflow/_durabletask/worker.py b/dapr/ext/workflow/_durabletask/worker.py index 9486582a8..4960e2a3b 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 @@ -236,6 +237,8 @@ def _is_message_too_large(rpc_error: grpc.RpcError) -> bool: # 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 +# How long stop() lets calls in flight on the current channel finish before it closes the channel. +_STOP_DRAIN_TIMEOUT = 5.0 class _HistoryResolutionError(Exception): @@ -476,8 +479,10 @@ def __init__( 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; notified whenever they change. + # 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 @@ -565,7 +570,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(): @@ -573,7 +583,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}') @@ -597,54 +607,48 @@ 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 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_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 + # 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 current_stub, conn_retry_count current_stub = None try: - current_channel = shared.get_grpc_channel( + channel = shared.get_grpc_channel( self._host_address, self._secure_channel, self._interceptors, options=self._channel_options, ) - current_stub = stubs.TaskHubSidecarServiceStub(current_channel) # Publish the connection for result sends, and for cleanup in stop() - self._set_connection(current_channel, current_stub) - current_stub.Hello(empty_pb2.Empty(), timeout=_CONNECT_HELLO_TIMEOUT) + self._set_connection(channel, stubs.TaskHubSidecarServiceStub(channel)) + with self._sidecar_stub(wait=0) as hello_stub: + if hello_stub is None: + raise RuntimeError('the worker is stopping') + hello_stub.Hello(empty_pb2.Empty(), timeout=_CONNECT_HELLO_TIMEOUT) + current_stub = hello_stub conn_retry_count = 0 self._logger.info(f'Created fresh connection to {self._host_address}') except Exception as e: self._logger.warning(f'Failed to create connection: {e}') - if current_channel is not None: - self._schedule_deferred_channel_close(current_channel) - current_channel = None - self._retire_connection() - current_stub = None + retire_connection() raise 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._retire_connection() + nonlocal current_stub, current_reader_thread + retire_connection() current_stub = None self._response_stream = None self._stream_ready.clear() @@ -688,8 +692,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] ) @@ -698,7 +700,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: + self._logger.debug('Connection retired by stop(), 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...' ) @@ -809,7 +815,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(): @@ -1001,8 +1007,44 @@ def _sidecar_stub(self, wait: float) -> Iterator[Optional[stubs.TaskHubSidecarSe """ with self._connection_lock: self._connection_lock.wait_for(self._has_connection_or_is_stopping, timeout=wait) + channel = self._current_channel stub = None if self._shutdown.is_set() else self._current_stub - yield 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 + if self._calls_in_flight[channel] <= 0: + del self._calls_in_flight[channel] + self._connection_lock.notify_all() + + def _close_channel(self, channel: grpc.Channel, drain_timeout: Optional[float]) -> 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. + + Args: + channel: A channel that ``_retire_connection`` already took out of use. + drain_timeout: Seconds to wait for calls in flight; None waits until they finish. + After a timeout the channel is closed anyway, which cancels those calls. + """ + with self._connection_lock: + drained = self._connection_lock.wait_for( + lambda: self._calls_in_flight[channel] == 0, timeout=drain_timeout + ) + if not drained: + self._logger.debug('Closing a gRPC channel with calls still in flight') + channel.close() def _send_on_current_connection( self, send: Callable[[stubs.TaskHubSidecarServiceStub], None], what: str @@ -1040,7 +1082,7 @@ def _deferred_close(): self._shutdown.wait(timeout=grace_timeout) finally: try: - old_channel.close() + self._close_channel(old_channel, drain_timeout=None) self._logger.debug('Deferred channel close completed') except Exception as e: self._logger.debug(f'Error during deferred channel close: {e}') @@ -1062,7 +1104,7 @@ def stop(self): channel = self._retire_connection() if channel is not None: try: - channel.close() + self._close_channel(channel, drain_timeout=_STOP_DRAIN_TIMEOUT) except Exception as e: self._logger.exception(f'Error closing gRPC channel: {e}') diff --git a/tests/ext/workflow/durabletask/test_worker_reconnect.py b/tests/ext/workflow/durabletask/test_worker_reconnect.py index 1df46c95a..921d99125 100644 --- a/tests/ext/workflow/durabletask/test_worker_reconnect.py +++ b/tests/ext/workflow/durabletask/test_worker_reconnect.py @@ -264,3 +264,56 @@ def test_result_is_dropped_when_the_worker_stops_before_reconnecting() -> None: 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, None)) + 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_once_a_stuck_call_runs_past_the_drain_timeout() -> None: + worker = _worker() + worker._runLoop = MagicMock() + worker._runLoop.is_alive.return_value = False + channel = MagicMock() + worker._set_connection(channel, MagicMock()) + call_started = threading.Event() + call_finish = threading.Event() + + def stuck_call() -> None: + with worker._sidecar_stub(wait=0): + call_started.set() + call_finish.wait(5) + + caller = threading.Thread(target=stuck_call) + caller.start() + assert call_started.wait(5) + with patch('dapr.ext.workflow._durabletask.worker._STOP_DRAIN_TIMEOUT', 0.1): + worker.stop() + call_finish.set() + caller.join(5) + + channel.close.assert_called_once() From e68ae575696e2939182eba5caee96ef614e5cbc8 Mon Sep 17 00:00:00 2001 From: Fabian Martinez <46371672+famarting@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:42:50 +0200 Subject: [PATCH 5/6] 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> --- dapr/ext/workflow/_durabletask/worker.py | 125 +++++++++++------- .../test_async_dispatch_regression.py | 4 +- .../durabletask/test_worker_history_cache.py | 12 +- .../durabletask/test_worker_reconnect.py | 117 +++++++++++++--- 4 files changed, 190 insertions(+), 68 deletions(-) diff --git a/dapr/ext/workflow/_durabletask/worker.py b/dapr/ext/workflow/_durabletask/worker.py index 4960e2a3b..6d14accde 100644 --- a/dapr/ext/workflow/_durabletask/worker.py +++ b/dapr/ext/workflow/_durabletask/worker.py @@ -237,7 +237,11 @@ def _is_message_too_large(rpc_error: grpc.RpcError) -> bool: # 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 -# How long stop() lets calls in flight on the current channel finish before it closes the channel. +# 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 @@ -607,12 +611,12 @@ async def _async_run_loop(self): """ worker_task = asyncio.create_task(self._async_worker_manager.run()) # Connection state management for retry fix - current_stub = None + connected = False current_reader_thread = None conn_retry_count = 0 conn_max_retry_delay = 15 - def retire_connection(): + def retire_and_close_connection(): # 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 @@ -623,33 +627,35 @@ def retire_connection(): self._schedule_deferred_channel_close(retired_channel) def create_fresh_connection(): - nonlocal current_stub, conn_retry_count - current_stub = None + 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: - channel = shared.get_grpc_channel( - self._host_address, - self._secure_channel, - self._interceptors, - options=self._channel_options, - ) - # Publish the connection for result sends, and for cleanup in stop() - self._set_connection(channel, stubs.TaskHubSidecarServiceStub(channel)) - with self._sidecar_stub(wait=0) as hello_stub: - if hello_stub is None: - raise RuntimeError('the worker is stopping') - hello_stub.Hello(empty_pb2.Empty(), timeout=_CONNECT_HELLO_TIMEOUT) - current_stub = hello_stub - 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}') - retire_connection() + 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_stub, current_reader_thread - retire_connection() - 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 @@ -675,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: @@ -701,8 +707,8 @@ def should_invalidate_connection(rpc_error): ) try: with self._sidecar_stub(wait=0) as stream_stub: - if stream_stub is None: - self._logger.debug('Connection retired by stop(), exiting') + 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( @@ -971,12 +977,22 @@ 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) -> None: - """Makes ``channel`` the connection that work-item results are sent on.""" + 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.""" @@ -1002,13 +1018,14 @@ def _sidecar_stub(self, wait: float) -> Iterator[Optional[stubs.TaskHubSidecarSe 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 - is stopping. + 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 = None if self._shutdown.is_set() else self._current_stub + stub = self._current_stub if stub is not None: self._calls_in_flight[channel] += 1 try: @@ -1020,11 +1037,23 @@ def _sidecar_stub(self, wait: float) -> Iterator[Optional[stubs.TaskHubSidecarSe 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() + self._connection_lock.notify_all() - def _close_channel(self, channel: grpc.Channel, drain_timeout: Optional[float]) -> None: + 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 @@ -1033,17 +1062,15 @@ def _close_channel(self, channel: grpc.Channel, drain_timeout: Optional[float]) ``_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. - drain_timeout: Seconds to wait for calls in flight; None waits until they finish. - After a timeout the channel is closed anyway, which cancels those calls. """ with self._connection_lock: - drained = self._connection_lock.wait_for( - lambda: self._calls_in_flight[channel] == 0, timeout=drain_timeout - ) - if not drained: - self._logger.debug('Closing a gRPC channel with calls still in flight') + self._connection_lock.wait_for(lambda: self._calls_in_flight[channel] == 0) channel.close() def _send_on_current_connection( @@ -1082,7 +1109,7 @@ def _deferred_close(): self._shutdown.wait(timeout=grace_timeout) finally: try: - self._close_channel(old_channel, drain_timeout=None) + 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}') @@ -1100,11 +1127,13 @@ 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 channel = self._retire_connection() if channel is not None: try: - self._close_channel(channel, drain_timeout=_STOP_DRAIN_TIMEOUT) + self._close_channel(channel) except Exception as e: self._logger.exception(f'Error closing gRPC channel: {e}') @@ -1199,7 +1228,7 @@ def _resolve_history(self, req: pb.WorkflowRequest) -> list[pb.HistoryEvent]: with self._sidecar_stub(wait=0) as stub: if stub is None: raise RuntimeError('not connected to the sidecar') - response = stub.GetInstanceHistory(history_request) + 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}" @@ -1298,7 +1327,7 @@ def _send_orchestrator_response( """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. @@ -1325,7 +1354,7 @@ def _send_orchestrator_response( 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}" @@ -1404,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. @@ -1420,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}' " diff --git a/tests/ext/workflow/durabletask/test_async_dispatch_regression.py b/tests/ext/workflow/durabletask/test_async_dispatch_regression.py index 0294cb00d..e7ed6d109 100644 --- a/tests/ext/workflow/durabletask/test_async_dispatch_regression.py +++ b/tests/ext/workflow/durabletask/test_async_dispatch_regression.py @@ -43,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 diff --git a/tests/ext/workflow/durabletask/test_worker_history_cache.py b/tests/ext/workflow/durabletask/test_worker_history_cache.py index f60c3694b..0c4d729b9 100644 --- a/tests/ext/workflow/durabletask/test_worker_history_cache.py +++ b/tests/ext/workflow/durabletask/test_worker_history_cache.py @@ -52,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) @@ -63,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) diff --git a/tests/ext/workflow/durabletask/test_worker_reconnect.py b/tests/ext/workflow/durabletask/test_worker_reconnect.py index 921d99125..70c926250 100644 --- a/tests/ext/workflow/durabletask/test_worker_reconnect.py +++ b/tests/ext/workflow/durabletask/test_worker_reconnect.py @@ -25,7 +25,11 @@ 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, TaskHubGrpcWorker +from dapr.ext.workflow._durabletask.worker import ( + _CONNECT_HELLO_TIMEOUT, + _RESULT_SEND_TIMEOUT, + TaskHubGrpcWorker, +) class _RpcError(grpc.RpcError): @@ -55,6 +59,7 @@ def __init__( 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() @@ -74,12 +79,21 @@ def GetWorkItems(self, _request: pb.GetWorkItemsRequest) -> Iterator[pb.WorkItem return yield from () - def CompleteActivityTask(self, response: pb.ActivityResponse) -> None: # noqa: N802 - self.results.append(response) - self.result_received.set() + 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 CompleteOrchestratorTask(self, response: pb.WorkflowResponse) -> None: # noqa: N802 + 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: @@ -229,6 +243,7 @@ def test_result_that_outlives_its_stream_is_sent_on_the_new_connection(kind: str 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: @@ -245,6 +260,7 @@ def test_result_ready_during_a_reconnect_waits_for_the_new_connection() -> None: 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' @@ -252,6 +268,30 @@ def test_result_ready_during_a_reconnect_waits_for_the_new_connection() -> None: 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() @@ -284,7 +324,7 @@ def test_channel_is_closed_only_after_the_calls_in_flight_on_it() -> None: 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, None)) + closer = threading.Thread(target=worker._close_channel, args=(channel,)) closer.start() closer.join(0.2) assert closer.is_alive() @@ -294,26 +334,71 @@ def test_channel_is_closed_only_after_the_calls_in_flight_on_it() -> None: channel.close.assert_called_once() -def test_stop_closes_the_channel_once_a_stuck_call_runs_past_the_drain_timeout() -> None: +def _stoppable_worker() -> TaskHubGrpcWorker: worker = _worker() worker._runLoop = MagicMock() worker._runLoop.is_alive.return_value = False - channel = MagicMock() - worker._set_connection(channel, MagicMock()) + 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() - call_finish = threading.Event() - def stuck_call() -> None: + def held_call() -> None: with worker._sidecar_stub(wait=0): call_started.set() - call_finish.wait(5) + release.wait(5) - caller = threading.Thread(target=stuck_call) + 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): - worker.stop() - call_finish.set() + 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 From 3669201b8d097d39ad24bd3b6285d2615001f41a Mon Sep 17 00:00:00 2001 From: Fabian Martinez <46371672+famarting@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:56:44 +0200 Subject: [PATCH 6/6] refactor(workflow): keep one gRPC channel for the worker's lifetime The worker replaced its gRPC channel every time the work-item stream ended. That forced the retire/publish handshake, per-channel call counts and the deferred-close threads, because results and other calls could still be using the old channel when it was closed. A grpc.Channel reconnects on its own, so only the GetWorkItems call needs to be re-issued. start() now opens one channel and stop() closes it. The run loop only sends Hello and opens a new stream on it. Result sends and the history fetch use wait_for_ready=True within their 30s deadline, so a result that is ready during a reconnect waits for the channel instead of a separate reconnect wait. The connect Hello waits for the channel too, so a returning sidecar is seen as soon as gRPC reconnects rather than after the run loop's retry delay on top of gRPC's own backoff, which is capped at 15s (default 120s). stop() still drains calls in flight for up to 5s and closes the channel only once no call is in flight, since grpcio crashes on a call started on a closed channel with interceptors. Closing the channel is also what ends the work-item stream: through an interceptor it is a generator with no cancel(). Co-Authored-By: Claude Opus 5.5 Signed-off-by: Fabian Martinez <46371672+famarting@users.noreply.github.com> --- dapr/ext/workflow/AGENTS.md | 2 +- dapr/ext/workflow/_durabletask/worker.py | 307 ++++++----------- .../test_async_dispatch_regression.py | 8 +- .../durabletask/test_worker_history_cache.py | 18 +- .../durabletask/test_worker_reconnect.py | 319 ++++++++---------- .../workflow/durabletask/test_worker_stop.py | 91 +---- 6 files changed, 272 insertions(+), 473 deletions(-) diff --git a/dapr/ext/workflow/AGENTS.md b/dapr/ext/workflow/AGENTS.md index 3ef1b13b9..e2465823b 100644 --- a/dapr/ext/workflow/AGENTS.md +++ b/dapr/ext/workflow/AGENTS.md @@ -122,7 +122,7 @@ At dispatch time (the gRPC stream loop in `_durabletask/worker.py`), `is_async_c - **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. +**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 6d14accde..4a92d2a68 100644 --- a/dapr/ext/workflow/_durabletask/worker.py +++ b/dapr/ext/workflow/_durabletask/worker.py @@ -18,7 +18,6 @@ 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 @@ -232,16 +231,19 @@ def _is_message_too_large(rpc_error: grpc.RpcError) -> bool: _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. +# 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 -# 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(). +# 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 -# 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. +# 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 @@ -481,13 +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._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. + # 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: Counter[grpc.Channel] = Counter() - self._channel_cleanup_threads: list[threading.Thread] = [] # Deferred channel close threads + self._calls_in_flight = 0 self._stream_ready = threading.Event() self._runLoop: Optional[Thread] = None # Use provided concurrency options or create default ones @@ -558,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 @@ -575,7 +579,7 @@ def run_loop(): self._is_running = True def _send_keepalive(self) -> None: - with self._sidecar_stub(wait=0) as stub: + with self._sidecar_stub() as stub: if stub is not None: stub.Hello(empty_pb2.Empty(), timeout=10) @@ -604,7 +608,7 @@ 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 @@ -616,45 +620,25 @@ async def _async_run_loop(self): conn_retry_count = 0 conn_max_retry_delay = 15 - def retire_and_close_connection(): - # 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. 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(): + def connect(): 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: - stub.Hello(empty_pb2.Empty(), timeout=_CONNECT_HELLO_TIMEOUT) + 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}') - channel.close() + self._logger.warning(f'Failed to connect: {e}') 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}') + self._logger.info(f'Connected to {self._host_address}') - def invalidate_connection(): + def disconnect(): nonlocal connected, current_reader_thread - retire_and_close_connection() connected = False self._response_stream = None self._stream_ready.clear() @@ -666,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, @@ -683,7 +667,7 @@ def should_invalidate_connection(rpc_error): break if not connected: try: - create_fresh_connection() + connect() except Exception: conn_retry_count += 1 delay = min( @@ -706,7 +690,7 @@ def should_invalidate_connection(rpc_error): pb.WORKER_CAPABILITY_STATEFUL_HISTORY ) try: - with self._sidecar_stub(wait=0) as stream_stub: + 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 @@ -878,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) @@ -891,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) @@ -911,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' ) @@ -960,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 @@ -977,147 +961,81 @@ 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.""" + 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: - 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() + self._channel = channel + self._stub = stub @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. + def _sidecar_stub(self) -> Iterator[Optional[stubs.TaskHubSidecarServiceStub]]: + """Yields the worker's stub and keeps the channel open until the block exits. 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. + The stub, or None once stop() has started closing the channel. """ 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 + stub = self._stub if stub is not None: - self._calls_in_flight[channel] += 1 + self._calls_in_flight += 1 try: yield stub finally: if stub is not None: - self._release_call(channel) + self._release_call() - def _release_call(self, channel: grpc.Channel) -> None: + def _release_call(self) -> 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._calls_in_flight -= 1 + if self._calls_in_flight == 0: 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. + 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``, and a retired channel is never handed out again, so once its count - reaches zero no call can start on it after the close. + ``_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. - - 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() + 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_on_current_connection( + def _send_result( 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: + """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}: not connected to the sidecar. ' + f'Could not deliver {what}: the worker has stopped. ' '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 - ): - """Schedule a deferred close of an old gRPC channel. - - 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. - - During ``stop()``, ``_shutdown`` is already set so the wait returns - immediately and the channel is closed at once. - """ - # 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: - 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}') - - thread = threading.Thread(target=_deferred_close, daemon=True, name='ChannelCleanup') - thread.start() - self._channel_cleanup_threads.append(thread) - def stop(self): """Stops the worker and waits for any pending work items to complete.""" # Guards on _runLoop rather than _is_running so stop() can unblock a start() @@ -1127,30 +1045,23 @@ 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 - channel = self._retire_connection() - if channel is not None: - try: - self._close_channel(channel) - except Exception as e: - self._logger.exception(f'Error closing gRPC channel: {e}') - - 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') @@ -1225,10 +1136,12 @@ def _resolve_history(self, req: pb.WorkflowRequest) -> list[pb.HistoryEvent]: history_request = pb.GetInstanceHistoryRequest(instanceId=req.instanceId) try: - with self._sidecar_stub(wait=0) as stub: + 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) + 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}" @@ -1312,7 +1225,7 @@ def _execute_orchestrator( completionToken=completionToken, ) - self._send_on_current_connection( + self._send_result( lambda stub: self._send_orchestrator_response(req, stub, res, completionToken), f"orchestrator response for '{req.instanceId}'", ) @@ -1327,7 +1240,7 @@ def _send_orchestrator_response( """Send an orchestrator response, falling back to a failure response when the response is too large to deliver.""" try: - stub.CompleteOrchestratorTask(res, timeout=_RESULT_SEND_TIMEOUT) + 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. @@ -1354,7 +1267,9 @@ def _send_orchestrator_response( completionToken=completionToken, ) try: - stub.CompleteOrchestratorTask(failure_res, timeout=_RESULT_SEND_TIMEOUT) + 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}" @@ -1433,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, timeout=_RESULT_SEND_TIMEOUT) + 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. @@ -1449,7 +1364,9 @@ def _send_activity_response( req, instance_id, oversize_error, completion_token ) try: - stub.CompleteActivityTask(failure_res, timeout=_RESULT_SEND_TIMEOUT) + 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}' " @@ -1476,7 +1393,7 @@ def _deliver_activity_response( completion_token, instance_id: str, ) -> None: - self._send_on_current_connection( + self._send_result( lambda stub: self._send_activity_response( req, stub, res, completion_token, 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 e7ed6d109..078e981ad 100644 --- a/tests/ext/workflow/durabletask/test_async_dispatch_regression.py +++ b/tests/ext/workflow/durabletask/test_async_dispatch_regression.py @@ -18,7 +18,6 @@ import asyncio from typing import cast -from unittest.mock import MagicMock import pytest @@ -44,7 +43,10 @@ def __init__(self) -> None: self.completions = 0 def CompleteActivityTask( # noqa: N802 - self, _response: pb.ActivityResponse, timeout: float | None = None + self, + _response: pb.ActivityResponse, + timeout: float | None = None, + wait_for_ready: bool | None = None, ) -> None: self.completions += 1 @@ -72,7 +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._set_connection(MagicMock(), cast(stubs.TaskHubSidecarServiceStub, stub)) + worker._stub = cast(stubs.TaskHubSidecarServiceStub, stub) async def activity(ctx, _inp) -> None: await asyncio.sleep(ACTIVITY_DURATION_SECONDS) diff --git a/tests/ext/workflow/durabletask/test_worker_history_cache.py b/tests/ext/workflow/durabletask/test_worker_history_cache.py index 0c4d729b9..7588053d3 100644 --- a/tests/ext/workflow/durabletask/test_worker_history_cache.py +++ b/tests/ext/workflow/durabletask/test_worker_history_cache.py @@ -17,7 +17,6 @@ """ from typing import Union, cast -from unittest.mock import MagicMock import pytest @@ -53,7 +52,10 @@ def __init__(self, events: list[pb.HistoryEvent]) -> None: self.get_instance_history_calls = 0 def GetInstanceHistory( - self, request: pb.GetInstanceHistoryRequest, timeout: float | None = None + 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) @@ -66,12 +68,18 @@ def __init__(self) -> None: self.responses: list[pb.WorkflowResponse] = [] def GetInstanceHistory( - self, request: pb.GetInstanceHistoryRequest, timeout: float | None = None + self, + request: pb.GetInstanceHistoryRequest, + timeout: float | None = None, + wait_for_ready: bool | None = None, ): raise RuntimeError('sidecar unavailable') def CompleteOrchestratorTask( - self, response: pb.WorkflowResponse, timeout: float | None = None + self, + response: pb.WorkflowResponse, + timeout: float | None = None, + wait_for_ready: bool | None = None, ) -> None: self.responses.append(response) @@ -182,7 +190,7 @@ def _worker(**kwargs) -> TaskHubGrpcWorker: def _connect(worker: TaskHubGrpcWorker, stub: Union[_FakeStub, _FailingHistoryStub]) -> None: - worker._set_connection(MagicMock(), cast(stubs.TaskHubSidecarServiceStub, stub)) + worker._stub = cast(stubs.TaskHubSidecarServiceStub, stub) def _resolve( diff --git a/tests/ext/workflow/durabletask/test_worker_reconnect.py b/tests/ext/workflow/durabletask/test_worker_reconnect.py index 70c926250..049426229 100644 --- a/tests/ext/workflow/durabletask/test_worker_reconnect.py +++ b/tests/ext/workflow/durabletask/test_worker_reconnect.py @@ -27,6 +27,7 @@ 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, ) @@ -44,95 +45,86 @@ def details(self) -> str: class _SidecarStub: - """One connection's stub: serves one work-item stream that stays open until ``end_stream``.""" + """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_error: Optional[Exception] = None, - hello_gate: Optional[threading.Event] = None, + hello_errors: Sequence[Optional[Exception]] = (), + hello_gates: Sequence[Optional[threading.Event]] = (), work_items: Sequence[pb.WorkItem] = (), ): self._worker = worker - self._hello_error = hello_error - self._hello_gate = hello_gate + 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_timeouts: list[Optional[float]] = [] + self.result_kwargs: list[dict[str, object]] = [] self.result_received = threading.Event() - self.stream_open = threading.Event() + self.streams_opened = 0 self._stream_end = threading.Event() - def Hello(self, *_args: object, timeout: Optional[float] = None) -> None: # noqa: N802 + def Hello( # noqa: N802 + self, *_args: object, timeout: Optional[float] = None, wait_for_ready: bool = False + ) -> None: 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 + 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.stream_open.set() - yield from self._work_items + 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 - yield from () def CompleteActivityTask( # noqa: N802 - self, response: pb.ActivityResponse, timeout: Optional[float] = None + self, response: pb.ActivityResponse, **kwargs: object ) -> None: - self._record_result(response, timeout) + self._record_result(response, kwargs) def CompleteOrchestratorTask( # noqa: N802 - self, response: pb.WorkflowResponse, timeout: Optional[float] = None + self, response: pb.WorkflowResponse, **kwargs: object ) -> None: - self._record_result(response, timeout) + self._record_result(response, kwargs) def _record_result( - self, response: Union[pb.ActivityResponse, pb.WorkflowResponse], timeout: Optional[float] + self, response: Union[pb.ActivityResponse, pb.WorkflowResponse], kwargs: dict[str, object] ) -> None: self.results.append(response) - self.result_timeouts.append(timeout) + self.result_kwargs.append(kwargs) 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]: +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', - side_effect=sidecar.channel, - ), + patch('dapr.ext.workflow._durabletask.worker.shared.get_grpc_channel') as get_channel, patch( 'dapr.ext.workflow._durabletask.worker.stubs.TaskHubSidecarServiceStub', - side_effect=sidecar.stub, + return_value=stub, ), ): worker.start() try: - yield + yield get_channel finally: worker.stop() @@ -141,54 +133,68 @@ 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) - sidecar = _FakeSidecar(worker, [stub]) - with _running(worker, sidecar): - assert stub.stream_open.wait(5) + 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_failed_connect_closes_its_channel() -> None: +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 = _SidecarStub(worker, hello_error=_RpcError(grpc.StatusCode.DEADLINE_EXCEEDED)) - healthy = _SidecarStub(worker) - sidecar = _FakeSidecar(worker, [timed_out, healthy]) + 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() - with _running(worker, sidecar): - assert healthy.stream_open.wait(10) + get_channel.assert_called_once() + channel.close.assert_called_once() - sidecar.channels[0].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)]) -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() + 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() - first = _SidecarStub(worker) reconnect_gate = threading.Event() - second = _SidecarStub(worker, hello_gate=reconnect_gate) - sidecar = _FakeSidecar(worker, [first, second]) + stub = _SidecarStub(worker, hello_gates=[None, reconnect_gate]) - with _running(worker, sidecar): + with _running(worker, stub): assert worker.is_worker_ready() - first.end_stream() + stub.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(lambda: stub.streams_opened == 2) assert _wait_until(worker.is_worker_ready), 'not ready again after reconnecting' assert not worker.is_worker_ready() @@ -225,106 +231,66 @@ async def run_async(_ctx: object, _input: object) -> str: @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.""" +def test_result_that_outlives_its_stream_is_still_delivered(kind: str) -> None: worker = _worker() work = _SlowWork(worker, kind) - first = _SidecarStub(worker, work_items=[work.work_item]) - second = _SidecarStub(worker) - sidecar = _FakeSidecar(worker, [first, second]) + stub = _SidecarStub(worker, work_items=[work.work_item]) - with _running(worker, sidecar): + with _running(worker, stub): assert work.started.wait(5) - first.end_stream() - assert second.stream_open.wait(5) + stub.end_stream() + assert _wait_until(lambda: stub.streams_opened == 2) work.release.set() - assert second.result_received.wait(5), 'result never reached the new connection' + assert stub.result_received.wait(5), 'result was not delivered' - assert first.results == [] - assert len(second.results) == 1 - assert second.result_timeouts == [_RESULT_SEND_TIMEOUT] + assert len(stub.results) == 1 -def test_result_ready_during_a_reconnect_waits_for_the_new_connection() -> None: +@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, '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]) + work = _SlowWork(worker, kind) + stub = _SidecarStub(worker, work_items=[work.work_item]) - with _running(worker, sidecar): + with _running(worker, stub): 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 stub.result_received.wait(5) - assert first.results == [] + assert stub.result_kwargs == [{'timeout': _RESULT_SEND_TIMEOUT, 'wait_for_ready': True}] -def test_result_ready_during_a_failed_reconnect_waits_for_a_healthy_one() -> None: +def _worker_with_channel() -> tuple[TaskHubGrpcWorker, MagicMock, MagicMock]: 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 == [] + 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_when_the_worker_stops_before_reconnecting() -> None: - worker = _worker() - stub = MagicMock() - worker._set_connection(MagicMock(), stub) - worker._retire_connection() - worker._shutdown.set() +def test_result_is_dropped_once_the_channel_is_closed() -> None: + worker, _channel, stub = _worker_with_channel() + worker._close_channel() - 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 + 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: - """Starting a call on a closed channel with interceptors crashes the process in grpcio.""" - worker = _worker() - channel = MagicMock() - worker._set_connection(channel, MagicMock()) + """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(wait=0) as stub: + with worker._sidecar_stub() as stub: assert stub is not None - worker._retire_connection() - closer = threading.Thread(target=worker._close_channel, args=(channel,)) + closer = threading.Thread(target=worker._close_channel) closer.start() closer.join(0.2) assert closer.is_alive() @@ -334,71 +300,52 @@ def test_channel_is_closed_only_after_the_calls_in_flight_on_it() -> None: 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.""" +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(wait=0): + with worker._sidecar_stub(): call_started.set() release.wait(5) caller = threading.Thread(target=held_call) caller.start() assert call_started.wait(5) - return caller + 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 = _stoppable_worker() - stub = MagicMock() - worker._set_connection(MagicMock(), stub) + worker, _channel, stub = _worker_with_channel() release = threading.Event() - caller = _hold_call(worker, release) + 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() - - -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 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: