Add patch activation callback (#1639)

Author: Sushisource

CHANGELOG.md

@@ -20,6 +20,10 @@ to include examples, links to docs, or any other relevant information.
 
 ### Added
 
+- Added the experimental `Worker` `patch_activation_callback` option, allowing workers
+  to decide whether a first non-replay `workflow.patched` call should activate a patch
+  during rolling deployments.
+
 ### Changed
 
 ### Deprecated

temporalio/worker/__init__.py

@@ -58,6 +58,7 @@
     WorkerDeploymentConfig,
 )
 from ._workflow_instance import (
+    PatchActivationInput,
     UnsandboxedWorkflowRunner,
     WorkflowInstance,
     WorkflowInstanceDetails,
@@ -77,6 +78,7 @@
     "PollerBehavior",
     "PollerBehaviorSimpleMaximum",
     "PollerBehaviorAutoscaling",
+    "PatchActivationInput",
     # Interceptor base classes
     "Interceptor",
     "ActivityInboundInterceptor",

temporalio/worker/_replayer.py

@@ -255,6 +255,7 @@ def on_eviction_hook(
                 workflow_failure_exception_types=self._config.get(
                     "workflow_failure_exception_types", []
                 ),
+                patch_activation_callback=None,
                 debug_mode=self._config.get("debug_mode", False),
                 metric_meter=runtime.metric_meter,
                 on_eviction_hook=on_eviction_hook,

temporalio/worker/_worker.py

@@ -40,7 +40,11 @@
     _DEFAULT_WORKFLOW_TASK_EXTERNAL_STORAGE_CONCURRENCY,
     _WorkflowWorker,
 )
-from ._workflow_instance import UnsandboxedWorkflowRunner, WorkflowRunner
+from ._workflow_instance import (
+    PatchActivationInput,
+    UnsandboxedWorkflowRunner,
+    WorkflowRunner,
+)
 from .workflow_sandbox import SandboxedWorkflowRunner
 
 logger = logging.getLogger(__name__)
@@ -135,6 +139,7 @@ def __init__(
         use_worker_versioning: bool = False,
         disable_safe_workflow_eviction: bool = False,
         deployment_config: WorkerDeploymentConfig | None = None,
+        patch_activation_callback: Callable[[PatchActivationInput], bool] | None = None,
         workflow_task_poller_behavior: PollerBehavior = PollerBehaviorSimpleMaximum(
             maximum=5
         ),
@@ -307,6 +312,12 @@ def __init__(
             deployment_config: Deployment config for the worker. Exclusive with ``build_id`` and
                 ``use_worker_versioning``.
                 WARNING: This is an experimental feature and may change in the future.
+            patch_activation_callback: Callback to decide whether the first non-replay
+                call to :py:func:`workflow.patched<temporalio.workflow.patched>` for a
+                patch ID should activate that patch. The callback receives a
+                :py:class:`PatchActivationInput` and must return ``True`` to activate the
+                patch or ``False`` to leave it inactive.
+                WARNING: This is an experimental feature and may change in the future.
             workflow_task_poller_behavior: Specify the behavior of workflow task polling.
                 Defaults to a 5-poller maximum.
             activity_task_poller_behavior: Specify the behavior of activity task polling.
@@ -368,6 +379,7 @@ def __init__(
             use_worker_versioning=use_worker_versioning,
             disable_safe_workflow_eviction=disable_safe_workflow_eviction,
             deployment_config=deployment_config,
+            patch_activation_callback=patch_activation_callback,
             workflow_task_poller_behavior=workflow_task_poller_behavior,
             activity_task_poller_behavior=activity_task_poller_behavior,
             nexus_task_poller_behavior=nexus_task_poller_behavior,
@@ -528,6 +540,7 @@ def check_activity(activity: str):
                 workflow_failure_exception_types=config[
                     "workflow_failure_exception_types"
                 ],  # type: ignore[reportTypedDictNotRequiredAccess]
+                patch_activation_callback=config.get("patch_activation_callback"),
                 debug_mode=config["debug_mode"],  # type: ignore[reportTypedDictNotRequiredAccess]
                 disable_eager_activity_execution=config[
                     "disable_eager_activity_execution"
@@ -982,6 +995,7 @@ class WorkerConfig(TypedDict, total=False):
     use_worker_versioning: bool
     disable_safe_workflow_eviction: bool
     deployment_config: WorkerDeploymentConfig | None
+    patch_activation_callback: Callable[[PatchActivationInput], bool] | None
     workflow_task_poller_behavior: PollerBehavior
     activity_task_poller_behavior: PollerBehavior
     nexus_task_poller_behavior: PollerBehavior

temporalio/worker/_workflow.py

@@ -42,6 +42,7 @@
     WorkflowInterceptorClassInput,
 )
 from ._workflow_instance import (
+    PatchActivationInput,
     WorkflowInstance,
     WorkflowInstanceDetails,
     WorkflowRunner,
@@ -78,6 +79,7 @@ def __init__(
         data_converter: temporalio.converter.DataConverter,
         interceptors: Sequence[Interceptor],
         workflow_failure_exception_types: Sequence[type[BaseException]],
+        patch_activation_callback: Callable[[PatchActivationInput], bool] | None,
         debug_mode: bool,
         disable_eager_activity_execution: bool,
         metric_meter: temporalio.common.MetricMeter,
@@ -145,6 +147,7 @@ def __init__(
         )
 
         self._workflow_failure_exception_types = workflow_failure_exception_types
+        self._patch_activation_callback = patch_activation_callback
         self._running_workflows: dict[str, _RunningWorkflow] = {}
         self._disable_eager_activity_execution = disable_eager_activity_execution
         self._on_eviction_hook = on_eviction_hook
@@ -798,6 +801,7 @@ def _create_workflow_instance(
             extern_functions=self._extern_functions,
             disable_eager_activity_execution=self._disable_eager_activity_execution,
             worker_level_failure_exception_types=self._workflow_failure_exception_types,
+            patch_activation_callback=self._patch_activation_callback,
             last_completion_result=init.last_completion_result,
             last_failure=last_failure,
         )

temporalio/worker/_workflow_instance.py

@@ -131,6 +131,17 @@ def set_worker_level_failure_exception_types(
         pass
 
 
+@dataclass(frozen=True)
+class PatchActivationInput:
+    """Input for the worker patch activation callback."""
+
+    workflow_info: temporalio.workflow.Info
+    """Information about the workflow execution calling ``patched``."""
+
+    patch_id: str
+    """Patch ID passed to ``patched``."""
+
+
 @dataclass(frozen=True)
 class WorkflowInstanceDetails:
     """Immutable details for creating a workflow instance."""
@@ -144,6 +155,7 @@ class WorkflowInstanceDetails:
     extern_functions: Mapping[str, Callable]
     disable_eager_activity_execution: bool
     worker_level_failure_exception_types: Sequence[type[BaseException]]
+    patch_activation_callback: Callable[[PatchActivationInput], bool] | None
     last_completion_result: temporalio.api.common.v1.Payloads
     last_failure: Failure | None
 
@@ -264,6 +276,7 @@ def __init__(self, det: WorkflowInstanceDetails) -> None:
         self._worker_level_failure_exception_types = (
             det.worker_level_failure_exception_types
         )
+        self._patch_activation_callback = det.patch_activation_callback
         self._primary_task: asyncio.Task[None] | None = None
         self._time_ns = 0
         self._cancel_reason: str | None = None
@@ -1363,7 +1376,20 @@ def workflow_patch(self, id: str, *, deprecated: bool) -> bool:
         if use_patch is not None:
             return use_patch
 
-        use_patch = not self._is_replaying or id in self._patches_notified
+        # Replay and history markers already determine the branch, and deprecation must
+        # keep existing patch semantics, so only a genuinely new patch consults the
+        # callback.
+        if deprecated or self._is_replaying or id in self._patches_notified:
+            use_patch = not self._is_replaying or id in self._patches_notified
+        elif self._patch_activation_callback is not None:
+            with self._as_read_only(in_query_or_validator=False):
+                use_patch = self._patch_activation_callback(
+                    PatchActivationInput(workflow_info=self._info, patch_id=id)
+                )
+            if type(use_patch) is not bool:
+                raise TypeError("Patch activation callback must return true or false")
+        else:
+            use_patch = True
         self._patches_memoized[id] = use_patch
         if use_patch:
             command = self._add_command()
@@ -1873,6 +1899,7 @@ def workflow_random_seed(self) -> int:
     def workflow_register_random_seed_callback(
         self, callback: Callable[[int], None]
     ) -> None:
+        self._assert_not_read_only("register random seed callback")
         self._seed_callbacks.append(callback)
 
     #### Calls from outbound impl ####

temporalio/worker/workflow_sandbox/_runner.py

@@ -89,6 +89,7 @@ def prepare_workflow(self, defn: temporalio.workflow._Definition) -> None:
                 extern_functions={},
                 disable_eager_activity_execution=False,
                 worker_level_failure_exception_types=self._worker_level_failure_exception_types,
+                patch_activation_callback=None,
                 last_completion_result=Payloads(),
                 last_failure=Failure(),
             ),