Introduce activity hooks; update on-demand worker

Add lifecycle hooks around activity execution and adapt the on-demand sandbox worker to use them. TaskHubGrpcWorker now defines _durabletask_on_activity_execution_started/completed and invokes them around activity execution; _execute_activity was refactored to adjust payload (de/externalization) and error handling. OnDemandSandboxWorker was updated to use new hook semantics, track active activity counts, expose add_activity wrapper for name resolution, use the internal shared logger, and consolidate on-demand-specific host/secure-channel attributes. Tests updated accordingly and a new test file verifies the activity hook behavior and active-activity counting.

Author: YunchuWang

durabletask-azuremanaged/durabletask/azuremanaged/preview/on_demand_sandbox/worker.py

@@ -18,6 +18,7 @@
     resolve_activity_names,
 )
 from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker
+import durabletask.internal.shared as shared
 from durabletask.worker import (
     ActivityWorkItemFilter,
     ConcurrencyOptions,
@@ -41,7 +42,10 @@ def __init__(self):
         concurrency_options = ConcurrencyOptions(
             maximum_concurrent_activity_work_items=resolved_max_concurrent_activities)
 
+        self._on_demand_sandbox_host_address = resolved_host_address
+        self._on_demand_sandbox_secure_channel = resolved_secure_channel
         self._on_demand_sandbox_token_credential = resolved_token_credential
+        self._on_demand_sandbox_logger = shared.get_logger("worker")
 
         super().__init__(
             host_address=resolved_host_address,
@@ -62,10 +66,13 @@ def __init__(self):
         self._on_demand_sandbox_active_activities = 0
         self._on_demand_sandbox_active_activities_lock = threading.Lock()
 
-    def start(self) -> None:
-        if self._is_running:
-            raise RuntimeError("The worker is already running.")
+    def add_activity(self, fn):
+        activity_name = super().add_activity(fn)
+        self._on_demand_sandbox_activity_names = resolve_activity_names(
+            [*self._on_demand_sandbox_activity_names, activity_name])
+        return activity_name
 
+    def start(self) -> None:
         self._configure_on_demand_sandbox_activity_filters()
         super().start()
         self._start_on_demand_sandbox_registration()
@@ -74,17 +81,16 @@ def stop(self) -> None:
         self._stop_on_demand_sandbox_registration()
         super().stop()
 
-    def _execute_activity(self, req, stub, completionToken):
+    def _durabletask_on_activity_execution_started(self, req) -> None:
         with self._on_demand_sandbox_active_activities_lock:
             self._on_demand_sandbox_active_activities += 1
-        try:
-            return super()._execute_activity(req, stub, completionToken)
-        finally:
-            with self._on_demand_sandbox_active_activities_lock:
-                self._on_demand_sandbox_active_activities = max(0, self._on_demand_sandbox_active_activities - 1)
+
+    def _durabletask_on_activity_execution_completed(self, req) -> None:
+        with self._on_demand_sandbox_active_activities_lock:
+            self._on_demand_sandbox_active_activities = max(0, self._on_demand_sandbox_active_activities - 1)
 
     def _configure_on_demand_sandbox_activity_filters(self) -> None:
-        activity_names = resolve_activity_names(self._registry.activities.keys())
+        activity_names = resolve_activity_names(self._on_demand_sandbox_activity_names)
         if not activity_names:
             raise RuntimeError(
                 "On-demand sandbox worker requires at least one registered activity before it can register.")
@@ -114,12 +120,10 @@ def _run_on_demand_sandbox_registration_loop(self) -> None:
         while not self._on_demand_sandbox_registration_stop.is_set():
             try:
                 client = OnDemandSandboxActivitiesClient(
-                    host_address=self._host_address,
+                    host_address=self._on_demand_sandbox_host_address,
                     taskhub=self._on_demand_sandbox_taskhub,
                     token_credential=self._on_demand_sandbox_token_credential,
-                    channel=self._channel,
-                    secure_channel=self._secure_channel,
-                    channel_options=self._channel_options)
+                    secure_channel=self._on_demand_sandbox_secure_channel)
                 try:
                     client.connect_on_demand_sandbox_activity_worker(self._registration_messages())
                     retry_delay = 1.0
@@ -128,7 +132,7 @@ def _run_on_demand_sandbox_registration_loop(self) -> None:
             except Exception as ex:
                 if self._on_demand_sandbox_registration_stop.is_set():
                     break
-                self._logger.warning("On-demand sandbox activity worker registration failed: %s", ex)
+                self._on_demand_sandbox_logger.warning("On-demand sandbox activity worker registration failed: %s", ex)
                 delay = random.uniform(0, retry_delay)
                 self._on_demand_sandbox_registration_stop.wait(delay)
                 retry_delay = min(retry_delay * 2, 30.0)

durabletask/worker.py

@@ -973,62 +973,71 @@ def _cancel_orchestrator(
         )
         self._logger.info(f"Cancelled orchestration task for invocation ID: {req.instanceId}")
 
+    def _durabletask_on_activity_execution_started(self, req: pb.ActivityRequest) -> None:
+        pass
+
+    def _durabletask_on_activity_execution_completed(self, req: pb.ActivityRequest) -> None:
+        pass
+
     def _execute_activity(
             self,
             req: pb.ActivityRequest,
             stub: Union[stubs.TaskHubSidecarServiceStub, ProtoTaskHubSidecarServiceStub],
             completionToken,
     ):
         instance_id = req.orchestrationInstance.instanceId
-
-        # De-externalize any large-payload tokens in the incoming request
-        if self._payload_store is not None:
-            payload_helpers.deexternalize_payloads(req, self._payload_store)
-        try:
-            executor = _ActivityExecutor(self._registry, self._logger)
-            with tracing.start_span(
-                tracing.create_span_name("activity", req.name),
-                trace_context=req.parentTraceContext,
-                kind=tracing.SpanKind.SERVER,
-                attributes={
-                    tracing.ATTR_TASK_TYPE: "activity",
-                    tracing.ATTR_TASK_INSTANCE_ID: instance_id,
-                    tracing.ATTR_TASK_NAME: req.name,
-                    tracing.ATTR_TASK_TASK_ID: str(req.taskId),
-                },
-            ) as span:
-                try:
-                    result = executor.execute(
-                        instance_id, req.name, req.taskId, req.input.value
-                    )
-                except Exception as ex:
-                    tracing.set_span_error(span, ex)
-                    raise
-            res = pb.ActivityResponse(
-                instanceId=instance_id,
-                taskId=req.taskId,
-                result=ph.get_string_value(result),
-                completionToken=completionToken,
-            )
-        except Exception as ex:
-            res = pb.ActivityResponse(
-                instanceId=instance_id,
-                taskId=req.taskId,
-                failureDetails=ph.new_failure_details(ex),
-                completionToken=completionToken,
-            )
-
+        self._durabletask_on_activity_execution_started(req)
         try:
-            # Externalize any large payloads in the response
+            # De-externalize any large-payload tokens in the incoming request
             if self._payload_store is not None:
-                payload_helpers.externalize_payloads(
-                    res, self._payload_store, instance_id=instance_id,
+                payload_helpers.deexternalize_payloads(req, self._payload_store)
+            try:
+                executor = _ActivityExecutor(self._registry, self._logger)
+                with tracing.start_span(
+                    tracing.create_span_name("activity", req.name),
+                    trace_context=req.parentTraceContext,
+                    kind=tracing.SpanKind.SERVER,
+                    attributes={
+                        tracing.ATTR_TASK_TYPE: "activity",
+                        tracing.ATTR_TASK_INSTANCE_ID: instance_id,
+                        tracing.ATTR_TASK_NAME: req.name,
+                        tracing.ATTR_TASK_TASK_ID: str(req.taskId),
+                    },
+                ) as span:
+                    try:
+                        result = executor.execute(
+                            instance_id, req.name, req.taskId, req.input.value
+                        )
+                    except Exception as ex:
+                        tracing.set_span_error(span, ex)
+                        raise
+                res = pb.ActivityResponse(
+                    instanceId=instance_id,
+                    taskId=req.taskId,
+                    result=ph.get_string_value(result),
+                    completionToken=completionToken,
                 )
-            stub.CompleteActivityTask(res)
-        except Exception as ex:
-            self._logger.exception(
-                f"Failed to deliver activity response for '{req.name}#{req.taskId}' of orchestration ID '{instance_id}' to sidecar: {ex}"
-            )
+            except Exception as ex:
+                res = pb.ActivityResponse(
+                    instanceId=instance_id,
+                    taskId=req.taskId,
+                    failureDetails=ph.new_failure_details(ex),
+                    completionToken=completionToken,
+                )
+
+            try:
+                # Externalize any large payloads in the response
+                if self._payload_store is not None:
+                    payload_helpers.externalize_payloads(
+                        res, self._payload_store, instance_id=instance_id,
+                    )
+                stub.CompleteActivityTask(res)
+            except Exception as ex:
+                self._logger.exception(
+                    f"Failed to deliver activity response for '{req.name}#{req.taskId}' of orchestration ID '{instance_id}' to sidecar: {ex}"
+                )
+        finally:
+            self._durabletask_on_activity_execution_completed(req)
 
     def _cancel_activity(
             self,

tests/durabletask-azuremanaged/test_on_demand_sandbox_extension.py

@@ -290,6 +290,7 @@ def test_generated_stub_uses_on_demand_sandbox_rpc_paths() -> None:
 
 def test_on_demand_sandbox_worker_constructor_does_not_expose_runtime_contract() -> None:
     assert list(inspect.signature(OnDemandSandboxWorker).parameters) == []
+    assert "_execute_activity" not in OnDemandSandboxWorker.__dict__
 
 
 def test_on_demand_sandbox_worker_does_not_own_legacy_wakeup_server(monkeypatch) -> None:
@@ -311,17 +312,23 @@ def test_on_demand_sandbox_worker_reads_sandbox_environment_and_registered_activ
     monkeypatch.setenv("DTS_SUBSTRATE", "AcaSessionPool")
     monkeypatch.setenv("DTS_SANDBOX_ID", "env-sandbox")
 
+    def EnvActivity(_ctx, value):
+        return value
+
+    def OtherActivity(_ctx, value):
+        return value
+
     worker = OnDemandSandboxWorker()
-    worker._registry.add_named_activity("EnvActivity", lambda _ctx, value: value)
-    worker._registry.add_named_activity("OtherActivity", lambda _ctx, value: value)
+    worker.add_activity(EnvActivity)
+    worker.add_activity(OtherActivity)
     worker._configure_on_demand_sandbox_activity_filters()
     start = next(worker._registration_messages())
 
-    assert worker._host_address == "http://localhost:8080"
+    assert worker._on_demand_sandbox_host_address == "http://localhost:8080"
     assert worker._on_demand_sandbox_token_credential is None
     assert worker._on_demand_sandbox_taskhub == "env-hub"
     assert worker._on_demand_sandbox_worker_profile_id == "env-profile"
-    assert worker._concurrency_options.maximum_concurrent_activity_work_items == 7
+    assert worker.concurrency_options.maximum_concurrent_activity_work_items == 7
     assert worker._work_item_filters is not None
     assert [activity.name for activity in worker._work_item_filters.activities] == [
         "EnvActivity",
@@ -353,7 +360,23 @@ def test_on_demand_sandbox_worker_ignores_legacy_max_activities(monkeypatch) ->
 
     worker = OnDemandSandboxWorker()
 
-    assert worker._concurrency_options.maximum_concurrent_activity_work_items == 100
+    assert worker.concurrency_options.maximum_concurrent_activity_work_items == 100
+
+
+def test_on_demand_sandbox_worker_tracks_active_activity_count_with_hooks(monkeypatch) -> None:
+    monkeypatch.setenv("DTS_ENDPOINT", "https://example.scheduler")
+    monkeypatch.setenv("DTS_TASK_HUB", "env-hub")
+
+    worker = OnDemandSandboxWorker()
+
+    worker._durabletask_on_activity_execution_started(object())
+    assert worker._on_demand_sandbox_active_activities == 1
+
+    worker._durabletask_on_activity_execution_completed(object())
+    assert worker._on_demand_sandbox_active_activities == 0
+
+    worker._durabletask_on_activity_execution_completed(object())
+    assert worker._on_demand_sandbox_active_activities == 0
 
 
 def test_on_demand_sandbox_worker_uses_managed_identity_credential_when_injected(monkeypatch) -> None:

tests/durabletask/test_worker_activity_hooks.py

@@ -0,0 +1,68 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from google.protobuf import wrappers_pb2
+
+import durabletask.internal.orchestrator_service_pb2 as pb
+from durabletask.worker import TaskHubGrpcWorker
+
+
+class RecordingStub:
+    def __init__(self):
+        self.completed = []
+
+    def CompleteActivityTask(self, response):
+        self.completed.append(response)
+
+
+class HookedWorker(TaskHubGrpcWorker):
+    def __init__(self):
+        super().__init__()
+        self.events = []
+
+    def _durabletask_on_activity_execution_started(self, req):
+        self.events.append(("started", req.name))
+
+    def _durabletask_on_activity_execution_completed(self, req):
+        self.events.append(("completed", req.name))
+
+
+def test_activity_execution_hooks_run_around_successful_activity() -> None:
+    def echo(_ctx, value):
+        return value
+
+    worker = HookedWorker()
+    worker.add_activity(echo)
+    stub = RecordingStub()
+    request = _activity_request("echo", "hello")
+
+    worker._execute_activity(request, stub, b"token")
+
+    assert worker.events == [("started", "echo"), ("completed", "echo")]
+    assert len(stub.completed) == 1
+    assert stub.completed[0].result.value == "\"hello\""
+
+
+def test_activity_execution_completed_hook_runs_after_failed_activity() -> None:
+    def fails(_ctx, _value):
+        raise ValueError("boom")
+
+    worker = HookedWorker()
+    worker.add_activity(fails)
+    stub = RecordingStub()
+    request = _activity_request("fails", "hello")
+
+    worker._execute_activity(request, stub, b"token")
+
+    assert worker.events == [("started", "fails"), ("completed", "fails")]
+    assert len(stub.completed) == 1
+    assert stub.completed[0].failureDetails.errorType == "ValueError"
+
+
+def _activity_request(name: str, value: str) -> pb.ActivityRequest:
+    return pb.ActivityRequest(
+        orchestrationInstance=pb.OrchestrationInstance(instanceId="instance-1"),
+        name=name,
+        taskId=1,
+        input=wrappers_pb2.StringValue(value=f"\"{value}\""),
+    )