Migrating shared access token logic to new grpc class

Signed-off-by: Ryan Lettieri <ryanLettieri@microsoft.com>

Author: RyanLettieri

durabletask/internal/shared.py

@@ -10,6 +10,7 @@
 import grpc
 
 from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
+from externalpackages.durabletaskscheduler.durabletask_grpc_interceptor import DTSDefaultClientInterceptorImpl
 
 # Field name used to indicate that an object was automatically serialized
 # and should be deserialized as a SimpleNamespace
@@ -50,7 +51,12 @@ def get_grpc_channel(
         channel = grpc.insecure_channel(host_address)
 
     if metadata is not None and len(metadata) > 0:
-        interceptors = [DefaultClientInterceptorImpl(metadata)]
+        for key, _ in metadata:
+            # Check if we are using DTS as the backend and if so, construct the DTS specific interceptors
+            if key == "dts":
+                interceptors = [DTSDefaultClientInterceptorImpl(metadata)]
+            else:
+                interceptors = [DefaultClientInterceptorImpl(metadata)]
         channel = grpc.intercept_channel(channel, *interceptors)
     return channel
 

examples/dts/dts_activity_sequence.py

@@ -4,7 +4,6 @@
 from durabletask import client, task
 from externalpackages.durabletaskscheduler.durabletask_scheduler_worker import DurableTaskSchedulerWorker
 from externalpackages.durabletaskscheduler.durabletask_scheduler_client import DurableTaskSchedulerClient
-from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
 
 def hello(ctx: task.ActivityContext, name: str) -> str:
     """Activity function that returns a greeting"""
@@ -48,15 +47,15 @@ def sequence(ctx: task.OrchestrationContext, _):
 
 
 # configure and start the worker
-with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True, client_id="", taskhub=taskhub_name) as w:
+with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True, use_managed_identity=False, client_id="", taskhub=taskhub_name) as w:
     w.add_orchestrator(sequence)
     w.add_activity(hello)
     w.start()
 
     # Construct the client and run the orchestrations
     c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True, taskhub=taskhub_name)
     instance_id = c.schedule_new_orchestration(sequence)
-    state = c.wait_for_orchestration_completion(instance_id, timeout=45)
+    state = c.wait_for_orchestration_completion(instance_id, timeout=60)
     if state and state.runtime_status == client.OrchestrationStatus.COMPLETED:
         print(f'Orchestration completed! Result: {state.serialized_output}')
     elif state:

examples/dts/dts_fanout_fanin.py

@@ -8,7 +8,6 @@
 from durabletask import client, task
 from externalpackages.durabletaskscheduler.durabletask_scheduler_worker import DurableTaskSchedulerWorker
 from externalpackages.durabletaskscheduler.durabletask_scheduler_client import DurableTaskSchedulerClient
-from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
 
 
 def get_work_items(ctx: task.ActivityContext, _) -> list[str]:
@@ -74,7 +73,7 @@ def orchestrator(ctx: task.OrchestrationContext, _):
     exit()
 
 # configure and start the worker
-with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True, client_id="", taskhub=taskhub_name) as w:
+with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True, taskhub=taskhub_name) as w:
     w.add_orchestrator(orchestrator)
     w.add_activity(process_work_item)
     w.add_activity(get_work_items)
@@ -88,3 +87,4 @@ def orchestrator(ctx: task.OrchestrationContext, _):
         print(f'Orchestration completed! Result: {state.serialized_output}')
     elif state:
         print(f'Orchestration failed: {state.failure_details}')
+    exit()

externalpackages/durabletaskscheduler/access_token_manager.py

@@ -1,22 +1,33 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from azure.identity import DefaultAzureCredential, ManagedIdentityCredential
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Optional
 
+# By default, when there's 10minutes left before the token expires, refresh the token
 class AccessTokenManager:
-    def __init__(self, refresh_buffer: int = 60, use_managed_identity: bool = False, client_id: Optional[str] = None):
+    def __init__(self, refresh_buffer: int = 600, metadata: Optional[list[tuple[str, str]]] = None):
         self.scope = "https://durabletask.io/.default"
         self.refresh_buffer = refresh_buffer
-        
+        self._use_managed_identity = False
+        self._metadata = metadata
+        self._client_id = None
+
+        if metadata:  # Ensure metadata is not None
+            for key, value in metadata:
+                if key == "use_managed_identity":
+                    self._use_managed_identity = value.lower() == "true"  # Properly convert string to bool
+                elif key == "client_id":
+                    self._client_id = value  # Directly assign string
+
         # Choose the appropriate credential based on use_managed_identity
-        if use_managed_identity:
-            if not client_id:
+        if self._use_managed_identity:
+            if not self._client_id:
                 print("Using System Assigned Managed Identity for authentication.")
                 self.credential = ManagedIdentityCredential()
             else:
                 print("Using User Assigned Managed Identity for authentication.")
-                self.credential = ManagedIdentityCredential(client_id)
+                self.credential = ManagedIdentityCredential(client_id=self._client_id)
         else:
             self.credential = DefaultAzureCredential()
             print("Using Default Azure Credentials for authentication.")
@@ -29,13 +40,18 @@ def get_access_token(self) -> str:
             self.refresh_token()
         return self.token
 
+    # Checks if the token is expired, or if it will expire in the next "refresh_buffer" seconds.
+    # For example, if the token is created to have a lifespan of 2 hours, and the refresh buffer is set to 30 minutes,
+    # We will grab a new token when there're 30minutes left on the lifespan of the token
     def is_token_expired(self) -> bool:
         if self.expiry_time is None:
             return True
-        return datetime.utcnow() >= (self.expiry_time - timedelta(seconds=self.refresh_buffer))
+        return datetime.now(timezone.utc) >= (self.expiry_time - timedelta(seconds=self.refresh_buffer))
 
     def refresh_token(self):
         new_token = self.credential.get_token(self.scope)
         self.token = f"Bearer {new_token.token}"
-        self.expiry_time = datetime.utcnow() + timedelta(seconds=new_token.expires_on - int(datetime.utcnow().timestamp()))
+        
+        # Convert UNIX timestamp to timezone-aware datetime
+        self.expiry_time = datetime.fromtimestamp(new_token.expires_on, tz=timezone.utc)
         print(f"Token refreshed. Expires at: {self.expiry_time}")

externalpackages/durabletaskscheduler/durabletask_grpc_interceptor.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from durabletask.internal.grpc_interceptor import _ClientCallDetails, DefaultClientInterceptorImpl
+from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
+
+import grpc
+
+class DTSDefaultClientInterceptorImpl (DefaultClientInterceptorImpl):
+    """The class implements a UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor,
+    StreamUnaryClientInterceptor and StreamStreamClientInterceptor from grpc to add an 
+    interceptor to add additional headers to all calls as needed."""
+
+    def __init__(self, metadata: list[tuple[str, str]]):
+        super().__init__(metadata)
+        self._token_manager = AccessTokenManager(metadata=self._metadata)
+
+    def _intercept_call(
+                self, client_call_details: _ClientCallDetails) -> grpc.ClientCallDetails:
+        """Internal intercept_call implementation which adds metadata to grpc metadata in the RPC
+            call details."""
+        # Refresh the auth token if it is present and needed
+        if self._metadata is not None:
+            for i, (key, _) in enumerate(self._metadata):
+                if key.lower() == "authorization":  # Ensure case-insensitive comparison
+                    new_token = self._token_manager.get_access_token()  # Get the new token
+                    self._metadata[i] = ("authorization", new_token)  # Update the token
+
+        return super()._intercept_call(client_call_details)

externalpackages/durabletaskscheduler/durabletask_scheduler_client.py

@@ -1,21 +1,31 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
 from typing import Optional
 from durabletask.client import TaskHubGrpcClient
 from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
 
 class DurableTaskSchedulerClient(TaskHubGrpcClient):
-    def __init__(self, *args, 
+    def __init__(self,
+                 host_address: str,
+                 secure_channel: bool,
                  metadata: Optional[list[tuple[str, str]]] = None,
+                 use_managed_identity: Optional[bool] = False,
                  client_id: Optional[str] = None,
-                 taskhub: str,
+                 taskhub: str = None,
                  **kwargs):
         if metadata is None:
             metadata = []  # Ensure metadata is initialized
         self._metadata = metadata
+        self._use_managed_identity = use_managed_identity
         self._client_id = client_id
         self._metadata.append(("taskhub", taskhub))
-        self._access_token_manager = AccessTokenManager(client_id=self._client_id)
+        self._metadata.append(("dts", "True"))
+        self._metadata.append(("use_managed_identity", str(use_managed_identity)))
+        self._metadata.append(("client_id", str(client_id)))
+        self._access_token_manager = AccessTokenManager(metadata=self._metadata)
         self.__update_metadata_with_token()
-        super().__init__(*args, metadata=self._metadata, **kwargs)
+        super().__init__(host_address=host_address, secure_channel=secure_channel, metadata=self._metadata, **kwargs)
 
     def __update_metadata_with_token(self):
         """
@@ -37,40 +47,4 @@ def __update_metadata_with_token(self):
 
         # If not updated, add a new entry
         if not updated:
-            self._metadata.append(("authorization", token))
-
-    def schedule_new_orchestration(self, *args, **kwargs) -> str:
-        self.__update_metadata_with_token()
-        return super().schedule_new_orchestration(*args, **kwargs)
-
-    def get_orchestration_state(self, *args, **kwargs):
-        self.__update_metadata_with_token()
-        super().get_orchestration_state(*args, **kwargs)
-
-    def wait_for_orchestration_start(self, *args, **kwargs):
-        self.__update_metadata_with_token()
-        super().wait_for_orchestration_start(*args, **kwargs)
-
-    def wait_for_orchestration_completion(self, *args, **kwargs):
-        self.__update_metadata_with_token()
-        super().wait_for_orchestration_completion(*args, **kwargs)
-
-    def raise_orchestration_event(self, *args, **kwargs):
-        self.__update_metadata_with_token()
-        super().raise_orchestration_event(*args, **kwargs)
-
-    def terminate_orchestration(self, *args, **kwargs):
-        self.__update_metadata_with_token()
-        super().terminate_orchestration(*args, **kwargs)
-
-    def suspend_orchestration(self, *args, **kwargs):
-        self.__update_metadata_with_token()
-        super().suspend_orchestration(*args, **kwargs)
-
-    def resume_orchestration(self, *args, **kwargs):
-        self.__update_metadata_with_token()
-        super().resume_orchestration(*args, **kwargs)
-
-    def purge_orchestration(self, *args, **kwargs):
-        self.__update_metadata_with_token()
-        super().purge_orchestration(*args, **kwargs)
+            self._metadata.append(("authorization", token))

externalpackages/durabletaskscheduler/durabletask_scheduler_worker.py

@@ -1,30 +1,32 @@
-import concurrent.futures
-from threading import Thread
-from google.protobuf import empty_pb2
-import grpc
-import durabletask.internal.orchestrator_service_pb2 as pb
-import durabletask.internal.orchestrator_service_pb2_grpc as stubs
-import durabletask.internal.shared as shared
-from typing import Optional
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
 
+from typing import Optional
 from durabletask.worker import TaskHubGrpcWorker
 from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
 
+# Worker class used for Durable Task Scheduler (DTS)
 class DurableTaskSchedulerWorker(TaskHubGrpcWorker):
-    def __init__(self, *args, 
+    def __init__(self,
+                 host_address: str,
+                 secure_channel: bool,
                  metadata: Optional[list[tuple[str, str]]] = None,
+                 use_managed_identity: Optional[bool] = False,
                  client_id: Optional[str] = None,
-                 taskhub: str,
+                 taskhub: str = None,
                  **kwargs):
         if metadata is None:
             metadata = []  # Ensure metadata is initialized
         self._metadata = metadata
+        self._use_managed_identity = use_managed_identity
         self._client_id = client_id
         self._metadata.append(("taskhub", taskhub))
-        self._access_token_manager = AccessTokenManager(client_id=self._client_id)
+        self._metadata.append(("dts", "True"))
+        self._metadata.append(("use_managed_identity", str(use_managed_identity)))
+        self._metadata.append(("client_id", str(client_id)))
+        self._access_token_manager = AccessTokenManager(metadata=self._metadata)
         self.__update_metadata_with_token()
-        super().__init__(*args, metadata=self._metadata, **kwargs)
-
+        super().__init__(host_address=host_address, secure_channel=secure_channel, metadata=self._metadata, **kwargs)
 
     def __update_metadata_with_token(self):
         """
@@ -47,61 +49,3 @@ def __update_metadata_with_token(self):
         # If not updated, add a new entry
         if not updated:
             self._metadata.append(("authorization", token))
-
-    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._metadata, self._secure_channel)
-        stub = stubs.TaskHubSidecarServiceStub(channel)
-
-        if self._is_running:
-            raise RuntimeError('The worker is already running.')
-
-        def run_loop():
-            # TODO: Investigate whether asyncio could be used to enable greater concurrency for async activity
-            #       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:
-                while not self._shutdown.is_set():
-                    try:
-                        self.__update_metadata_with_token()
-                        # send a "Hello" message to the sidecar to ensure that it's listening
-                        stub.Hello(empty_pb2.Empty())
-
-                        # stream work items
-                        self._response_stream = 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
-                        # by another thread (see the stop() method).
-                        for work_item in self._response_stream:  # type: ignore
-                            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)
-                            elif work_item.HasField('activityRequest'):
-                                executor.submit(self._execute_activity, work_item.activityRequest, 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:
-                        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
-                            self._logger.warning(
-                                f'The sidecar at address {self._host_address} is unavailable - will continue retrying')
-                        else:
-                            self._logger.warning(f'Unexpected error: {rpc_error}')
-                    except Exception as ex:
-                        self._logger.warning(f'Unexpected error: {ex}')
-
-                    # CONSIDER: exponential backoff
-                    self._shutdown.wait(5)
-                self._logger.info("No longer listening for work items")
-                return
-
-        self._logger.info(f"Starting gRPC worker that connects to {self._host_address}")
-        self._runLoop = Thread(target=run_loop)
-        self._runLoop.start()
-        self._is_running = True