feat: gRPC interceptor improvements (channel options, retry policy, custom channel, workerid, thread safety) (#132)

* feat: gRPC interceptor improvements (channel options, retry policy, custom channel, workerid, thread safety)

- Add GrpcChannelOptions and GrpcRetryPolicyOptions for transport-level
  configuration including message size limits, keepalive settings, and
  gRPC service-config-based retry policies.
- Add optional channel parameter to TaskHubGrpcClient,
  AsyncTaskHubGrpcClient, and TaskHubGrpcWorker for pre-configured
  channel passthrough.
- Add optional channel_options parameter to the same classes for
  low-level channel customization.
- Expose interceptors, channel, and channel_options on Azure
  Managed DurableTaskSchedulerClient, AsyncDurableTaskSchedulerClient,
  and DurableTaskSchedulerWorker so custom interceptors can be combined
  with the built-in DTS auth interceptor.
- Add workerid (hostname:pid:uuid) gRPC metadata header to Azure
  Managed worker calls for improved service-side observability.
- Fix sync AccessTokenManager to use double-checked locking under a
  threading.Lock to prevent duplicate token refresh under concurrent load.
- Export GrpcChannelOptions and GrpcRetryPolicyOptions from the
  top-level durabletask package.
- Add unit tests covering all new behaviors and backward compatibility.

* fix: address PR review comments

- Track channel ownership in sync/async clients; close() is a no-op
  when the channel was provided externally via the channel= parameter.
- Raise ValueError in GrpcRetryPolicyOptions.__post_init__ when either
  backoff duration rounds to zero at 9-decimal-place precision, rather
  than silently emitting a zero-duration string that gRPC rejects.

* Fix merge issue

Author: andystaples

CHANGELOG.md

@@ -9,6 +9,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ADDED
 
+- Added `GrpcChannelOptions` and `GrpcRetryPolicyOptions` for configuring
+  gRPC transport behavior, including message-size limits, keepalive settings,
+  and channel-level retry policy service configuration.
+- Added optional `channel` and `channel_options` parameters to
+  `TaskHubGrpcClient`, `AsyncTaskHubGrpcClient`, and `TaskHubGrpcWorker` to
+  support pre-configured channel passthrough and low-level gRPC channel
+  customization.
 - Added `get_orchestration_history()` and `list_instance_ids()` to the sync and async gRPC clients.
 - Added in-memory backend support for `StreamInstanceHistory` and `ListInstanceIds` so local orchestration tests can retrieve history and page terminal instance IDs by completion window.
 

durabletask-azuremanaged/CHANGELOG.md

@@ -7,6 +7,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- Added optional `interceptors`, `channel`, and `channel_options` parameters to
+  `DurableTaskSchedulerClient`, `AsyncDurableTaskSchedulerClient`, and
+  `DurableTaskSchedulerWorker` to allow combining custom gRPC interceptors with
+  DTS defaults and to support pre-configured/customized gRPC channels.
+- Added `workerid` gRPC metadata on Durable Task Scheduler worker calls for
+  improved worker identity and observability.
+- Improved sync access token refresh concurrency handling to avoid duplicate
+  refresh operations under concurrent access.
+
 ## v1.4.0
 
 - Updates base dependency to durabletask v1.4.0

durabletask-azuremanaged/durabletask/azuremanaged/client.py

@@ -3,8 +3,10 @@
 
 import logging
 
-from typing import Optional
+from typing import Optional, Sequence
 
+import grpc
+import grpc.aio
 from azure.core.credentials import TokenCredential
 from azure.core.credentials_async import AsyncTokenCredential
 
@@ -13,6 +15,8 @@
     DTSDefaultClientInterceptorImpl,
 )
 from durabletask.client import AsyncTaskHubGrpcClient, TaskHubGrpcClient
+from durabletask.grpc_options import GrpcChannelOptions
+import durabletask.internal.shared as shared
 from durabletask.payload.store import PayloadStore
 
 
@@ -22,7 +26,10 @@ def __init__(self, *,
                  host_address: str,
                  taskhub: str,
                  token_credential: Optional[TokenCredential],
+                 channel: Optional[grpc.Channel] = None,
                  secure_channel: bool = True,
+                 interceptors: Optional[Sequence[shared.ClientInterceptor]] = None,
+                 channel_options: Optional[GrpcChannelOptions] = None,
                  default_version: Optional[str] = None,
                  payload_store: Optional[PayloadStore] = None,
                  log_handler: Optional[logging.Handler] = None,
@@ -31,17 +38,22 @@ def __init__(self, *,
         if not taskhub:
             raise ValueError("Taskhub value cannot be empty. Please provide a value for your taskhub")
 
-        interceptors = [DTSDefaultClientInterceptorImpl(token_credential, taskhub)]
+        resolved_interceptors: list[shared.ClientInterceptor] = (
+            list(interceptors) if interceptors is not None else []
+        )
+        resolved_interceptors.append(DTSDefaultClientInterceptorImpl(token_credential, taskhub))
 
         # We pass in None for the metadata so we don't construct an additional interceptor in the parent class
         # Since the parent class doesn't use anything metadata for anything else, we can set it as None
         super().__init__(
             host_address=host_address,
+            channel=channel,
             secure_channel=secure_channel,
             metadata=None,
             log_handler=log_handler,
             log_formatter=log_formatter,
-            interceptors=interceptors,
+            interceptors=resolved_interceptors,
+            channel_options=channel_options,
             default_version=default_version,
             payload_store=payload_store)
 
@@ -88,7 +100,10 @@ def __init__(self, *,
                  host_address: str,
                  taskhub: str,
                  token_credential: Optional[AsyncTokenCredential],
+                 channel: Optional[grpc.aio.Channel] = None,
                  secure_channel: bool = True,
+                 interceptors: Optional[Sequence[shared.AsyncClientInterceptor]] = None,
+                 channel_options: Optional[GrpcChannelOptions] = None,
                  default_version: Optional[str] = None,
                  payload_store: Optional[PayloadStore] = None,
                  log_handler: Optional[logging.Handler] = None,
@@ -97,16 +112,21 @@ def __init__(self, *,
         if not taskhub:
             raise ValueError("Taskhub value cannot be empty. Please provide a value for your taskhub")
 
-        interceptors = [DTSAsyncDefaultClientInterceptorImpl(token_credential, taskhub)]
+        resolved_interceptors: list[shared.AsyncClientInterceptor] = (
+            list(interceptors) if interceptors is not None else []
+        )
+        resolved_interceptors.append(DTSAsyncDefaultClientInterceptorImpl(token_credential, taskhub))
 
         # We pass in None for the metadata so we don't construct an additional interceptor in the parent class
         # Since the parent class doesn't use anything metadata for anything else, we can set it as None
         super().__init__(
             host_address=host_address,
+            channel=channel,
             secure_channel=secure_channel,
             metadata=None,
             log_handler=log_handler,
             log_formatter=log_formatter,
-            interceptors=interceptors,
+            interceptors=resolved_interceptors,
+            channel_options=channel_options,
             default_version=default_version,
             payload_store=payload_store)

durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from datetime import datetime, timedelta, timezone
+from threading import Lock
 from typing import Optional
 
 from azure.core.credentials import AccessToken, TokenCredential
@@ -20,6 +21,7 @@ def __init__(self, token_credential: Optional[TokenCredential], refresh_interval
         self._logger = shared.get_logger("token_manager")
 
         self._credential = token_credential
+        self._refresh_lock = Lock()
 
         if self._credential is not None:
             self._token = self._credential.get_token(self._scope)
@@ -30,7 +32,9 @@ def __init__(self, token_credential: Optional[TokenCredential], refresh_interval
 
     def get_access_token(self) -> Optional[AccessToken]:
         if self._token is None or self.is_token_expired():
-            self.refresh_token()
+            with self._refresh_lock:
+                if self._token is None or self.is_token_expired():
+                    self.refresh_token()
         return self._token
 
     # Checks if the token is expired, or if it will expire in the next "refresh_interval_seconds" seconds.

durabletask-azuremanaged/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py

@@ -25,7 +25,11 @@ class DTSDefaultClientInterceptorImpl (DefaultClientInterceptorImpl):
     StreamUnaryClientInterceptor and StreamStreamClientInterceptor from grpc to add an
     interceptor to add additional headers to all calls as needed."""
 
-    def __init__(self, token_credential: Optional[TokenCredential], taskhub_name: str):
+    def __init__(
+            self,
+            token_credential: Optional[TokenCredential],
+            taskhub_name: str,
+            worker_id: Optional[str] = None):
         try:
             # Get the version of the azuremanaged package
             sdk_version = version('durabletask-azuremanaged')
@@ -35,7 +39,9 @@ def __init__(self, token_credential: Optional[TokenCredential], taskhub_name: st
         user_agent = f"durabletask-python/{sdk_version}"
         self._metadata = [
             ("taskhub", taskhub_name),
-            ("x-user-agent", user_agent)]  # 'user-agent' is a reserved header in grpc, so we use 'x-user-agent' instead
+            ("x-user-agent", user_agent)]  # 'user-agent' is a reserved header; use 'x-user-agent'
+        if worker_id is not None:
+            self._metadata.append(("workerid", worker_id))
         super().__init__(self._metadata)
 
         self._token_manager = None
@@ -44,7 +50,17 @@ def __init__(self, token_credential: Optional[TokenCredential], taskhub_name: st
             self._token_manager = AccessTokenManager(token_credential=self._token_credential)
             access_token = self._token_manager.get_access_token()
             if access_token is not None:
-                self._metadata.append(("authorization", f"Bearer {access_token.token}"))
+                self._upsert_authorization_header(access_token.token)
+
+    def _upsert_authorization_header(self, token: str) -> None:
+        found = False
+        for i, (key, _) in enumerate(self._metadata):
+            if key.lower() == "authorization":
+                self._metadata[i] = ("authorization", f"Bearer {token}")
+                found = True
+                break
+        if not found:
+            self._metadata.append(("authorization", f"Bearer {token}"))
 
     def _intercept_call(
             self, client_call_details: _ClientCallDetails) -> grpc.ClientCallDetails:
@@ -56,15 +72,7 @@ def _intercept_call(
         if self._token_manager is not None:
             access_token = self._token_manager.get_access_token()
             if access_token is not None:
-                # Update the existing authorization header
-                found = False
-                for i, (key, _) in enumerate(self._metadata):
-                    if key.lower() == "authorization":
-                        self._metadata[i] = ("authorization", f"Bearer {access_token.token}")
-                        found = True
-                        break
-                if not found:
-                    self._metadata.append(("authorization", f"Bearer {access_token.token}"))
+                self._upsert_authorization_header(access_token.token)
 
         return super()._intercept_call(client_call_details)
 
@@ -96,6 +104,16 @@ def __init__(self, token_credential: Optional[AsyncTokenCredential], taskhub_nam
             self._token_credential = token_credential
             self._token_manager = AsyncAccessTokenManager(token_credential=self._token_credential)
 
+    def _upsert_authorization_header(self, token: str) -> None:
+        found = False
+        for i, (key, _) in enumerate(self._metadata):
+            if key.lower() == "authorization":
+                self._metadata[i] = ("authorization", f"Bearer {token}")
+                found = True
+                break
+        if not found:
+            self._metadata.append(("authorization", f"Bearer {token}"))
+
     async def _intercept_call(
             self, client_call_details: _AsyncClientCallDetails) -> grpc.aio.ClientCallDetails:
         """Internal intercept_call implementation which adds metadata to grpc metadata in the RPC
@@ -106,16 +124,6 @@ async def _intercept_call(
         if self._token_manager is not None:
             access_token = await self._token_manager.get_access_token()
             if access_token is not None:
-                # Update the existing authorization header, or append one if this
-                # is the first successful token acquisition (token is lazily
-                # fetched on the first call since async constructors aren't possible).
-                found = False
-                for i, (key, _) in enumerate(self._metadata):
-                    if key.lower() == "authorization":
-                        self._metadata[i] = ("authorization", f"Bearer {access_token.token}")
-                        found = True
-                        break
-                if not found:
-                    self._metadata.append(("authorization", f"Bearer {access_token.token}"))
+                self._upsert_authorization_header(access_token.token)
 
         return await super()._intercept_call(client_call_details)

durabletask-azuremanaged/durabletask/azuremanaged/worker.py

@@ -2,13 +2,19 @@
 # Licensed under the MIT License.
 
 import logging
+import os
+import socket
+import uuid
 
-from typing import Optional
+from typing import Optional, Sequence
 
+import grpc
 from azure.core.credentials import TokenCredential
 
 from durabletask.azuremanaged.internal.durabletask_grpc_interceptor import \
     DTSDefaultClientInterceptorImpl
+from durabletask.grpc_options import GrpcChannelOptions
+import durabletask.internal.shared as shared
 from durabletask.payload.store import PayloadStore
 from durabletask.worker import ConcurrencyOptions, TaskHubGrpcWorker
 
@@ -64,7 +70,10 @@ def __init__(self, *,
                  host_address: str,
                  taskhub: str,
                  token_credential: Optional[TokenCredential],
+                 channel: Optional[grpc.Channel] = None,
                  secure_channel: bool = True,
+                 interceptors: Optional[Sequence[shared.ClientInterceptor]] = None,
+                 channel_options: Optional[GrpcChannelOptions] = None,
                  concurrency_options: Optional[ConcurrencyOptions] = None,
                  payload_store: Optional[PayloadStore] = None,
                  log_handler: Optional[logging.Handler] = None,
@@ -73,17 +82,25 @@ def __init__(self, *,
         if not taskhub:
             raise ValueError("The taskhub value cannot be empty.")
 
-        interceptors = [DTSDefaultClientInterceptorImpl(token_credential, taskhub)]
+        worker_id = f"{socket.gethostname()}:{os.getpid()}:{uuid.uuid4()}"
+        resolved_interceptors: list[shared.ClientInterceptor] = (
+            list(interceptors) if interceptors is not None else []
+        )
+        resolved_interceptors.append(
+            DTSDefaultClientInterceptorImpl(token_credential, taskhub, worker_id=worker_id)
+        )
 
         # We pass in None for the metadata so we don't construct an additional interceptor in the parent class
         # Since the parent class doesn't use anything metadata for anything else, we can set it as None
         super().__init__(
             host_address=host_address,
+            channel=channel,
             secure_channel=secure_channel,
             metadata=None,
             log_handler=log_handler,
             log_formatter=log_formatter,
-            interceptors=interceptors,
+            interceptors=resolved_interceptors,
+            channel_options=channel_options,
             concurrency_options=concurrency_options,
             # DTS natively supports long timers so chunking is unnecessary
             maximum_timer_interval=None,

durabletask/__init__.py

@@ -3,6 +3,7 @@
 
 """Durable Task SDK for Python"""
 
+from durabletask.grpc_options import GrpcChannelOptions, GrpcRetryPolicyOptions
 from durabletask.payload.store import LargePayloadStorageOptions, PayloadStore
 from durabletask.worker import (
     ActivityWorkItemFilter,
@@ -17,6 +18,8 @@
     "ActivityWorkItemFilter",
     "ConcurrencyOptions",
     "EntityWorkItemFilter",
+    "GrpcChannelOptions",
+    "GrpcRetryPolicyOptions",
     "LargePayloadStorageOptions",
     "OrchestrationWorkItemFilter",
     "PayloadStore",