Fix channel handling

Author: berndverst

durabletask/worker.py

@@ -11,8 +11,8 @@
 import grpc
 from google.protobuf import empty_pb2
 
-import durabletask.internal.helpers as ph
 import durabletask.internal.helpers as pbh
+import durabletask.internal.helpers as ph
 import durabletask.internal.orchestrator_service_pb2 as pb
 import durabletask.internal.orchestrator_service_pb2_grpc as stubs
 import durabletask.internal.shared as shared
@@ -129,9 +129,6 @@ def add_activity(self, fn: task.Activity) -> str:
 
     def start(self):
         """Starts the worker on a background thread and begins listening for work items."""
-        channel = shared.get_grpc_channel(self._host_address, self._secure_channel, self._interceptors)
-        stub = stubs.TaskHubSidecarServiceStub(channel)
-
         if self._is_running:
             raise RuntimeError('The worker is already running.')
 
@@ -140,13 +137,30 @@ def run_loop():
             #       functions. We'd need to know ahead of time whether a function is async or not.
             # TODO: Max concurrency configuration settings
             with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
+                current_channel = None
+                current_stub = None
+
                 while not self._shutdown.is_set():
+                    # Only create a new channel and stub if we don't have a current one
+                    # or if we need to retry after a failure
+                    if current_channel is None or current_stub is None:
+                        # Clean up any existing channel before creating a new one
+                        if current_channel is not None:
+                            try:
+                                current_channel.close()
+                            except Exception:
+                                pass  # Ignore any errors during channel cleanup
+
+                        # Create a fresh channel and stub for this connection attempt
+                        current_channel = shared.get_grpc_channel(self._host_address, self._secure_channel, self._interceptors)
+                        current_stub = stubs.TaskHubSidecarServiceStub(current_channel)
+
                     try:
                         # send a "Hello" message to the sidecar to ensure that it's listening
-                        stub.Hello(empty_pb2.Empty())
+                        current_stub.Hello(empty_pb2.Empty())
 
                         # stream work items
-                        self._response_stream = stub.GetWorkItems(pb.GetWorkItemsRequest())
+                        self._response_stream = current_stub.GetWorkItems(pb.GetWorkItemsRequest())
                         self._logger.info(f'Successfully connected to {self._host_address}. Waiting for work items...')
 
                         # The stream blocks until either a work item is received or the stream is canceled
@@ -155,15 +169,19 @@ def run_loop():
                             request_type = work_item.WhichOneof('request')
                             self._logger.debug(f'Received "{request_type}" work item')
                             if work_item.HasField('orchestratorRequest'):
-                                executor.submit(self._execute_orchestrator, work_item.orchestratorRequest, stub, work_item.completionToken)
+                                executor.submit(self._execute_orchestrator, work_item.orchestratorRequest, current_stub, work_item.completionToken)
                             elif work_item.HasField('activityRequest'):
-                                executor.submit(self._execute_activity, work_item.activityRequest, stub, work_item.completionToken)
+                                executor.submit(self._execute_activity, work_item.activityRequest, current_stub, work_item.completionToken)
                             elif work_item.HasField('healthPing'):
                                 pass  # no-op
                             else:
                                 self._logger.warning(f'Unexpected work item type: {request_type}')
 
                     except grpc.RpcError as rpc_error:
+                        # Mark current channel/stub as invalid so they'll be recreated on next retry
+                        current_channel = None
+                        current_stub = None
+
                         if rpc_error.code() == grpc.StatusCode.CANCELLED:  # type: ignore
                             self._logger.info(f'Disconnected from {self._host_address}')
                         elif rpc_error.code() == grpc.StatusCode.UNAVAILABLE:  # type: ignore
@@ -172,10 +190,20 @@ def run_loop():
                         else:
                             self._logger.warning(f'Unexpected error: {rpc_error}')
                     except Exception as ex:
+                        # Mark current channel/stub as invalid so they'll be recreated on next retry
+                        current_channel = None
+                        current_stub = None
                         self._logger.warning(f'Unexpected error: {ex}')
 
                     # CONSIDER: exponential backoff
                     self._shutdown.wait(5)
+
+                # Clean up the final channel when the loop exits
+                if current_channel is not None:
+                    try:
+                        current_channel.close()
+                    except Exception:
+                        pass  # Ignore any errors during channel cleanup
                 self._logger.info("No longer listening for work items")
                 return