diff --git a/temporalio/bridge/extension.py b/temporalio/bridge/extension.py new file mode 100644 index 0000000000..5c45160776 --- /dev/null +++ b/temporalio/bridge/extension.py @@ -0,0 +1,28 @@ +"""Worker extensions that this bridge build links. (unstable) + +Nothing in this module should be considered stable. The API may change. +""" + +from __future__ import annotations + +from typing import Any + +import temporalio.bridge.runtime +import temporalio.bridge.temporal_sdk_bridge + + +class WorkerExtension: + """An extension that the bridge builds by its name.""" + + def __init__( + self, runtime: temporalio.bridge.runtime.Runtime, name: str, config: bytes + ) -> None: + """Build the extension with ``name`` from its encoded config.""" + self._ref: Any = temporalio.bridge.temporal_sdk_bridge.new_worker_extension( + runtime._ref, name, config + ) + + @property + def name(self) -> str: + """Name the extension gives itself.""" + return self._ref.name diff --git a/temporalio/bridge/proto/activity_task/__init__.py b/temporalio/bridge/proto/activity_task/__init__.py index 779ebf3535..9e69a0f894 100644 --- a/temporalio/bridge/proto/activity_task/__init__.py +++ b/temporalio/bridge/proto/activity_task/__init__.py @@ -4,6 +4,7 @@ ActivityTask, Cancel, Start, + WorkflowTaskAttempt, ) __all__ = [ @@ -12,4 +13,5 @@ "ActivityTask", "Cancel", "Start", + "WorkflowTaskAttempt", ] diff --git a/temporalio/bridge/proto/activity_task/activity_task_pb2.py b/temporalio/bridge/proto/activity_task/activity_task_pb2.py index 0e09839dc9..f269eb0cac 100644 --- a/temporalio/bridge/proto/activity_task/activity_task_pb2.py +++ b/temporalio/bridge/proto/activity_task/activity_task_pb2.py @@ -26,7 +26,7 @@ ) DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( - b'\n3temporal/sdk/core/activity_task/activity_task.proto\x12\x15\x63oresdk.activity_task\x1a\x1egoogle/protobuf/duration.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a$temporal/api/common/v1/message.proto\x1a%temporal/sdk/core/common/common.proto"\x8d\x01\n\x0c\x41\x63tivityTask\x12\x12\n\ntask_token\x18\x01 \x01(\x0c\x12-\n\x05start\x18\x03 \x01(\x0b\x32\x1c.coresdk.activity_task.StartH\x00\x12/\n\x06\x63\x61ncel\x18\x04 \x01(\x0b\x32\x1d.coresdk.activity_task.CancelH\x00\x42\t\n\x07variant"\xb1\x07\n\x05Start\x12\x1a\n\x12workflow_namespace\x18\x01 \x01(\t\x12\x15\n\rworkflow_type\x18\x02 \x01(\t\x12\x45\n\x12workflow_execution\x18\x03 \x01(\x0b\x32).temporal.api.common.v1.WorkflowExecution\x12\x13\n\x0b\x61\x63tivity_id\x18\x04 \x01(\t\x12\x15\n\ractivity_type\x18\x05 \x01(\t\x12\x45\n\rheader_fields\x18\x06 \x03(\x0b\x32..coresdk.activity_task.Start.HeaderFieldsEntry\x12.\n\x05input\x18\x07 \x03(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12:\n\x11heartbeat_details\x18\x08 \x03(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12\x32\n\x0escheduled_time\x18\t \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x42\n\x1e\x63urrent_attempt_scheduled_time\x18\n \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x30\n\x0cstarted_time\x18\x0b \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07\x61ttempt\x18\x0c \x01(\r\x12<\n\x19schedule_to_close_timeout\x18\r \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x39\n\x16start_to_close_timeout\x18\x0e \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x34\n\x11heartbeat_timeout\x18\x0f \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x39\n\x0cretry_policy\x18\x10 \x01(\x0b\x32#.temporal.api.common.v1.RetryPolicy\x12\x32\n\x08priority\x18\x12 \x01(\x0b\x32 .temporal.api.common.v1.Priority\x12\x10\n\x08is_local\x18\x11 \x01(\x08\x12\x0e\n\x06run_id\x18\x13 \x01(\t\x1aT\n\x11HeaderFieldsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12.\n\x05value\x18\x02 \x01(\x0b\x32\x1f.temporal.api.common.v1.Payload:\x02\x38\x01"\x8a\x01\n\x06\x43\x61ncel\x12;\n\x06reason\x18\x01 \x01(\x0e\x32+.coresdk.activity_task.ActivityCancelReason\x12\x43\n\x07\x64\x65tails\x18\x02 \x01(\x0b\x32\x32.coresdk.activity_task.ActivityCancellationDetails"\xa0\x01\n\x1b\x41\x63tivityCancellationDetails\x12\x14\n\x0cis_not_found\x18\x01 \x01(\x08\x12\x14\n\x0cis_cancelled\x18\x02 \x01(\x08\x12\x11\n\tis_paused\x18\x03 \x01(\x08\x12\x14\n\x0cis_timed_out\x18\x04 \x01(\x08\x12\x1a\n\x12is_worker_shutdown\x18\x05 \x01(\x08\x12\x10\n\x08is_reset\x18\x06 \x01(\x08*o\n\x14\x41\x63tivityCancelReason\x12\r\n\tNOT_FOUND\x10\x00\x12\r\n\tCANCELLED\x10\x01\x12\r\n\tTIMED_OUT\x10\x02\x12\x13\n\x0fWORKER_SHUTDOWN\x10\x03\x12\n\n\x06PAUSED\x10\x04\x12\t\n\x05RESET\x10\x05\x42\x32\xea\x02/Temporalio::Internal::Bridge::Api::ActivityTaskb\x06proto3' + b'\n3temporal/sdk/core/activity_task/activity_task.proto\x12\x15\x63oresdk.activity_task\x1a\x1egoogle/protobuf/duration.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a$temporal/api/common/v1/message.proto\x1a%temporal/sdk/core/common/common.proto"\x8d\x01\n\x0c\x41\x63tivityTask\x12\x12\n\ntask_token\x18\x01 \x01(\x0c\x12-\n\x05start\x18\x03 \x01(\x0b\x32\x1c.coresdk.activity_task.StartH\x00\x12/\n\x06\x63\x61ncel\x18\x04 \x01(\x0b\x32\x1d.coresdk.activity_task.CancelH\x00\x42\t\n\x07variant"\xf4\x07\n\x05Start\x12\x1a\n\x12workflow_namespace\x18\x01 \x01(\t\x12\x15\n\rworkflow_type\x18\x02 \x01(\t\x12\x45\n\x12workflow_execution\x18\x03 \x01(\x0b\x32).temporal.api.common.v1.WorkflowExecution\x12\x13\n\x0b\x61\x63tivity_id\x18\x04 \x01(\t\x12\x15\n\ractivity_type\x18\x05 \x01(\t\x12\x45\n\rheader_fields\x18\x06 \x03(\x0b\x32..coresdk.activity_task.Start.HeaderFieldsEntry\x12.\n\x05input\x18\x07 \x03(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12:\n\x11heartbeat_details\x18\x08 \x03(\x0b\x32\x1f.temporal.api.common.v1.Payload\x12\x32\n\x0escheduled_time\x18\t \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x42\n\x1e\x63urrent_attempt_scheduled_time\x18\n \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x30\n\x0cstarted_time\x18\x0b \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07\x61ttempt\x18\x0c \x01(\r\x12<\n\x19schedule_to_close_timeout\x18\r \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x39\n\x16start_to_close_timeout\x18\x0e \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x34\n\x11heartbeat_timeout\x18\x0f \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x39\n\x0cretry_policy\x18\x10 \x01(\x0b\x32#.temporal.api.common.v1.RetryPolicy\x12\x32\n\x08priority\x18\x12 \x01(\x0b\x32 .temporal.api.common.v1.Priority\x12\x10\n\x08is_local\x18\x11 \x01(\x08\x12\x0e\n\x06run_id\x18\x13 \x01(\t\x12\x41\n\rworkflow_task\x18\x14 \x01(\x0b\x32*.coresdk.activity_task.WorkflowTaskAttempt\x1aT\n\x11HeaderFieldsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12.\n\x05value\x18\x02 \x01(\x0b\x32\x1f.temporal.api.common.v1.Payload:\x02\x38\x01"r\n\x13WorkflowTaskAttempt\x12\x18\n\x10started_event_id\x18\x01 \x01(\x03\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\r\x12\x30\n\x0cstarted_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp"\x8a\x01\n\x06\x43\x61ncel\x12;\n\x06reason\x18\x01 \x01(\x0e\x32+.coresdk.activity_task.ActivityCancelReason\x12\x43\n\x07\x64\x65tails\x18\x02 \x01(\x0b\x32\x32.coresdk.activity_task.ActivityCancellationDetails"\xa0\x01\n\x1b\x41\x63tivityCancellationDetails\x12\x14\n\x0cis_not_found\x18\x01 \x01(\x08\x12\x14\n\x0cis_cancelled\x18\x02 \x01(\x08\x12\x11\n\tis_paused\x18\x03 \x01(\x08\x12\x14\n\x0cis_timed_out\x18\x04 \x01(\x08\x12\x1a\n\x12is_worker_shutdown\x18\x05 \x01(\x08\x12\x10\n\x08is_reset\x18\x06 \x01(\x08*o\n\x14\x41\x63tivityCancelReason\x12\r\n\tNOT_FOUND\x10\x00\x12\r\n\tCANCELLED\x10\x01\x12\r\n\tTIMED_OUT\x10\x02\x12\x13\n\x0fWORKER_SHUTDOWN\x10\x03\x12\n\n\x06PAUSED\x10\x04\x12\t\n\x05RESET\x10\x05\x42\x32\xea\x02/Temporalio::Internal::Bridge::Api::ActivityTaskb\x06proto3' ) _ACTIVITYCANCELREASON = DESCRIPTOR.enum_types_by_name["ActivityCancelReason"] @@ -42,6 +42,7 @@ _ACTIVITYTASK = DESCRIPTOR.message_types_by_name["ActivityTask"] _START = DESCRIPTOR.message_types_by_name["Start"] _START_HEADERFIELDSENTRY = _START.nested_types_by_name["HeaderFieldsEntry"] +_WORKFLOWTASKATTEMPT = DESCRIPTOR.message_types_by_name["WorkflowTaskAttempt"] _CANCEL = DESCRIPTOR.message_types_by_name["Cancel"] _ACTIVITYCANCELLATIONDETAILS = DESCRIPTOR.message_types_by_name[ "ActivityCancellationDetails" @@ -78,6 +79,17 @@ _sym_db.RegisterMessage(Start) _sym_db.RegisterMessage(Start.HeaderFieldsEntry) +WorkflowTaskAttempt = _reflection.GeneratedProtocolMessageType( + "WorkflowTaskAttempt", + (_message.Message,), + { + "DESCRIPTOR": _WORKFLOWTASKATTEMPT, + "__module__": "temporal.sdk.core.activity_task.activity_task_pb2", + # @@protoc_insertion_point(class_scope:coresdk.activity_task.WorkflowTaskAttempt) + }, +) +_sym_db.RegisterMessage(WorkflowTaskAttempt) + Cancel = _reflection.GeneratedProtocolMessageType( "Cancel", (_message.Message,), @@ -107,16 +119,18 @@ ) _START_HEADERFIELDSENTRY._options = None _START_HEADERFIELDSENTRY._serialized_options = b"8\001" - _ACTIVITYCANCELREASON._serialized_start = 1616 - _ACTIVITYCANCELREASON._serialized_end = 1727 + _ACTIVITYCANCELREASON._serialized_start = 1799 + _ACTIVITYCANCELREASON._serialized_end = 1910 _ACTIVITYTASK._serialized_start = 221 _ACTIVITYTASK._serialized_end = 362 _START._serialized_start = 365 - _START._serialized_end = 1310 - _START_HEADERFIELDSENTRY._serialized_start = 1226 - _START_HEADERFIELDSENTRY._serialized_end = 1310 - _CANCEL._serialized_start = 1313 - _CANCEL._serialized_end = 1451 - _ACTIVITYCANCELLATIONDETAILS._serialized_start = 1454 - _ACTIVITYCANCELLATIONDETAILS._serialized_end = 1614 + _START._serialized_end = 1377 + _START_HEADERFIELDSENTRY._serialized_start = 1293 + _START_HEADERFIELDSENTRY._serialized_end = 1377 + _WORKFLOWTASKATTEMPT._serialized_start = 1379 + _WORKFLOWTASKATTEMPT._serialized_end = 1493 + _CANCEL._serialized_start = 1496 + _CANCEL._serialized_end = 1634 + _ACTIVITYCANCELLATIONDETAILS._serialized_start = 1637 + _ACTIVITYCANCELLATIONDETAILS._serialized_end = 1797 # @@protoc_insertion_point(module_scope) diff --git a/temporalio/bridge/proto/activity_task/activity_task_pb2.pyi b/temporalio/bridge/proto/activity_task/activity_task_pb2.pyi index 1a2f8c7d2f..dead4d805e 100644 --- a/temporalio/bridge/proto/activity_task/activity_task_pb2.pyi +++ b/temporalio/bridge/proto/activity_task/activity_task_pb2.pyi @@ -160,6 +160,7 @@ class Start(google.protobuf.message.Message): PRIORITY_FIELD_NUMBER: builtins.int IS_LOCAL_FIELD_NUMBER: builtins.int RUN_ID_FIELD_NUMBER: builtins.int + WORKFLOW_TASK_FIELD_NUMBER: builtins.int workflow_namespace: builtins.str """The namespace the workflow lives in""" workflow_type: builtins.str @@ -227,6 +228,11 @@ class Start(google.protobuf.message.Message): """ run_id: builtins.str """Run ID of this activity execution. Only set for standalone activities.""" + @property + def workflow_task(self) -> global___WorkflowTaskAttempt: + """For a local activity, the workflow task in which Core dispatched this execution. Retries + that Core runs without a new task keep it. Unset for other activities. + """ def __init__( self, *, @@ -258,6 +264,7 @@ class Start(google.protobuf.message.Message): priority: temporalio.api.common.v1.message_pb2.Priority | None = ..., is_local: builtins.bool = ..., run_id: builtins.str = ..., + workflow_task: global___WorkflowTaskAttempt | None = ..., ) -> None: ... def HasField( self, @@ -280,6 +287,8 @@ class Start(google.protobuf.message.Message): b"started_time", "workflow_execution", b"workflow_execution", + "workflow_task", + b"workflow_task", ], ) -> builtins.bool: ... def ClearField( @@ -321,6 +330,8 @@ class Start(google.protobuf.message.Message): b"workflow_execution", "workflow_namespace", b"workflow_namespace", + "workflow_task", + b"workflow_task", "workflow_type", b"workflow_type", ], @@ -328,6 +339,44 @@ class Start(google.protobuf.message.Message): global___Start = Start +class WorkflowTaskAttempt(google.protobuf.message.Message): + """One attempt of a workflow task, with the values from the poll response that started it.""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + STARTED_EVENT_ID_FIELD_NUMBER: builtins.int + ATTEMPT_FIELD_NUMBER: builtins.int + STARTED_TIME_FIELD_NUMBER: builtins.int + started_event_id: builtins.int + """The id of the task's WorkflowTaskStarted event.""" + attempt: builtins.int + @property + def started_time(self) -> google.protobuf.timestamp_pb2.Timestamp: + """When the server started the task.""" + def __init__( + self, + *, + started_event_id: builtins.int = ..., + attempt: builtins.int = ..., + started_time: google.protobuf.timestamp_pb2.Timestamp | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["started_time", b"started_time"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "attempt", + b"attempt", + "started_event_id", + b"started_event_id", + "started_time", + b"started_time", + ], + ) -> None: ... + +global___WorkflowTaskAttempt = WorkflowTaskAttempt + class Cancel(google.protobuf.message.Message): """Attempt to cancel a running activity""" diff --git a/temporalio/bridge/sdk-core b/temporalio/bridge/sdk-core index 0a5d646279..3a9203ac7c 160000 --- a/temporalio/bridge/sdk-core +++ b/temporalio/bridge/sdk-core @@ -1 +1 @@ -Subproject commit 0a5d646279824adcf07eeaa52d9179dcd934ab46 +Subproject commit 3a9203ac7cb15171097d7f58e80b648f517ea658 diff --git a/temporalio/bridge/src/extension.rs b/temporalio/bridge/src/extension.rs new file mode 100644 index 0000000000..cce72dce16 --- /dev/null +++ b/temporalio/bridge/src/extension.rs @@ -0,0 +1,158 @@ +//! Worker extensions that this bridge build links, so Python can pass one to Core by its name. + +use std::sync::Arc; + +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use temporalio_sdk_core::WorkerExtension; + +use crate::runtime; + +/// An extension that a Worker config can carry. Python keeps it to reach the same instance that +/// Core calls. +#[pyclass(frozen, from_py_object)] +#[derive(Clone)] +pub struct WorkerExtensionRef { + pub(crate) extension: Arc, +} + +#[pymethods] +impl WorkerExtensionRef { + #[getter] + fn name(&self) -> String { + self.extension.name().to_string() + } +} + +/// Builds the extension with `name` from its encoded config. Only the extensions that this build +/// links are known. +// A release build links no extension yet, so it never reads the config. +#[cfg_attr(not(debug_assertions), allow(unused_variables))] +pub fn new_worker_extension( + runtime_ref: &runtime::RuntimeRef, + name: &str, + config: Vec, +) -> PyResult { + let runtime = runtime_ref.runtime.clone(); + runtime.assert_same_process("create a worker extension")?; + // An extension may start tasks on the runtime that it lives on. + let _guard = runtime.core.tokio_handle().enter(); + let extension: Arc = match name { + #[cfg(debug_assertions)] + test_extension::NAME => Arc::new(test_extension::TestExtension::new(&config)?), + _ => { + return Err(PyValueError::new_err(format!( + "This build has no worker extension named {name}" + ))) + } + }; + Ok(WorkerExtensionRef { extension }) +} + +/// An extension for the SDK's own tests. Release builds leave it out. +#[cfg(debug_assertions)] +mod test_extension { + use std::sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }; + + use pyo3::exceptions::PyValueError; + use pyo3::PyResult; + use temporalio_sdk_core::{ + ExtensionError, WorkerExtension, WorkflowTaskCompletionInfo, WorkflowTaskInterceptError, + WorkflowTaskInterceptor, + }; + + pub(super) const NAME: &str = "temporalio-test"; + + pub(super) struct TestExtension { + interceptor: Arc, + } + + impl TestExtension { + /// The config is empty, or `fail_before_complete` to fail every workflow task. + pub(super) fn new(config: &[u8]) -> PyResult { + let fail_before_complete = match config { + b"" => false, + b"fail_before_complete" => true, + _ => return Err(PyValueError::new_err("Unknown test extension config")), + }; + Ok(Self { + interceptor: Arc::new(CountingInterceptor { + fail_before_complete, + before: AtomicU64::new(0), + after: AtomicU64::new(0), + accepted: AtomicU64::new(0), + }), + }) + } + } + + #[async_trait::async_trait] + impl WorkerExtension for TestExtension { + fn name(&self) -> &str { + NAME + } + + fn workflow_task_interceptor(&self) -> Option> { + Some(self.interceptor.clone()) + } + + fn call(&self, method: &str, input: &[u8]) -> Result, ExtensionError> { + match method { + "Echo" => Ok(input.to_vec()), + "Fail" => Err(ExtensionError::new(tonic::Code::NotFound, "test failure") + .with_details(input.to_vec())), + // Comma separated, so a test reads it with no proto. + "Counts" => Ok(format!( + "{},{},{}", + self.interceptor.before.load(Ordering::SeqCst), + self.interceptor.after.load(Ordering::SeqCst), + self.interceptor.accepted.load(Ordering::SeqCst) + ) + .into_bytes()), + _ => Err(ExtensionError::unknown_method(method)), + } + } + + async fn call_async( + &self, + method: String, + input: Vec, + ) -> Result, ExtensionError> { + tokio::task::yield_now().await; + self.call(&method, &input) + } + } + + struct CountingInterceptor { + fail_before_complete: bool, + before: AtomicU64, + after: AtomicU64, + accepted: AtomicU64, + } + + #[async_trait::async_trait] + impl WorkflowTaskInterceptor for CountingInterceptor { + async fn before_complete( + &self, + _task: &WorkflowTaskCompletionInfo, + ) -> Result<(), WorkflowTaskInterceptError> { + self.before.fetch_add(1, Ordering::SeqCst); + if self.fail_before_complete { + return Err(WorkflowTaskInterceptError::new( + "test extension refused the task", + )); + } + Ok(()) + } + + async fn after_complete(&self, _task: &WorkflowTaskCompletionInfo, accepted: bool) { + self.after.fetch_add(1, Ordering::SeqCst); + if accepted { + self.accepted.fetch_add(1, Ordering::SeqCst); + } + } + } +} diff --git a/temporalio/bridge/src/lib.rs b/temporalio/bridge/src/lib.rs index a4fda26db6..ae2eed5860 100644 --- a/temporalio/bridge/src/lib.rs +++ b/temporalio/bridge/src/lib.rs @@ -4,6 +4,7 @@ use pyo3::types::PyTuple; mod client; mod client_rpc_generated; mod envconfig; +mod extension; mod metric; mod runtime; mod testing; @@ -69,6 +70,8 @@ fn temporal_sdk_bridge(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_function(wrap_pyfunction!(new_worker, m)?)?; m.add_function(wrap_pyfunction!(new_replay_worker, m)?)?; + m.add_class::()?; + m.add_function(wrap_pyfunction!(new_worker_extension, m)?)?; // envconfig let envconfig_module = PyModule::new(py, "envconfig")?; @@ -131,6 +134,15 @@ fn new_worker( worker::new_worker(runtime_ref, client, config) } +#[pyfunction] +fn new_worker_extension( + runtime_ref: &runtime::RuntimeRef, + name: &str, + config: Vec, +) -> PyResult { + extension::new_worker_extension(runtime_ref, name, config) +} + #[pyfunction] fn new_replay_worker<'a>( py: Python<'a>, diff --git a/temporalio/bridge/src/worker.rs b/temporalio/bridge/src/worker.rs index e74698e944..f838a0384e 100644 --- a/temporalio/bridge/src/worker.rs +++ b/temporalio/bridge/src/worker.rs @@ -26,6 +26,7 @@ use tokio_stream::wrappers::ReceiverStream; use tracing::error; use crate::client; +use crate::extension::WorkerExtensionRef; use crate::runtime; pyo3::create_exception!(temporal_sdk_bridge, PollShutdownError, PyException); @@ -65,6 +66,16 @@ pub struct WorkerConfig { plugins: Vec, storage_drivers: HashSet, disable_payload_error_limit: bool, + #[pyo3(from_py_with = extract_extension)] + extension: Option, +} + +/// Python wraps the extension, so the config has the wrapper. +fn extract_extension(obj: &Bound<'_, PyAny>) -> PyResult> { + if obj.is_none() { + return Ok(None); + } + obj.getattr("_ref")?.extract().map(Some).map_err(Into::into) } #[derive(FromPyObject)] @@ -780,6 +791,7 @@ fn convert_worker_config( .collect::>(), ) .disable_payload_error_limit(conf.disable_payload_error_limit) + .maybe_extension(conf.extension.map(|ext| ext.extension)) .build() .map_err(|err| PyValueError::new_err(format!("Invalid worker config: {err}"))) } diff --git a/temporalio/bridge/worker.py b/temporalio/bridge/worker.py index 4b7f55d092..0243e3627a 100644 --- a/temporalio/bridge/worker.py +++ b/temporalio/bridge/worker.py @@ -12,6 +12,7 @@ ) import temporalio.bridge.client +import temporalio.bridge.extension import temporalio.bridge.proto import temporalio.bridge.proto.activity_task import temporalio.bridge.proto.nexus @@ -60,6 +61,7 @@ class WorkerConfig: plugins: Sequence[str] storage_drivers: set[str] disable_payload_error_limit: bool + extension: temporalio.bridge.extension.WorkerExtension | None @dataclass diff --git a/temporalio/worker/__init__.py b/temporalio/worker/__init__.py index 4f6efe68c3..061175df8f 100644 --- a/temporalio/worker/__init__.py +++ b/temporalio/worker/__init__.py @@ -2,6 +2,7 @@ from ..common import WorkerDeploymentVersion from ._activity import SharedHeartbeatSender, SharedStateManager +from ._extension import WorkerExtension from ._interceptor import ( ActivityInboundInterceptor, ActivityOutboundInterceptor, @@ -79,6 +80,7 @@ "PollerBehaviorSimpleMaximum", "PollerBehaviorAutoscaling", "PatchActivationInput", + "WorkerExtension", # Interceptor base classes "Interceptor", "ActivityInboundInterceptor", diff --git a/temporalio/worker/_extension.py b/temporalio/worker/_extension.py new file mode 100644 index 0000000000..cbf218491c --- /dev/null +++ b/temporalio/worker/_extension.py @@ -0,0 +1,52 @@ +"""Native extensions that a Worker passes to Core.""" + +from __future__ import annotations + +import temporalio.bridge.extension +import temporalio.runtime + + +class WorkerExtension: + """A native extension that the SDK's bridge build links, which a Worker passes to Core. + + Core calls the extension at fixed points of each Workflow Task. A plugin builds the extension + and passes it in the Worker's ``extension`` option. The Worker must use the extension's + runtime. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + name: str, + config: bytes = b"", + *, + runtime: temporalio.runtime.Runtime | None = None, + ) -> None: + """Build the extension with a name that the bridge build knows. + + Args: + name: The name of the extension in the bridge build. + config: The extension's encoded config. Each extension has its own format. + runtime: The runtime the extension runs on. Defaults to + :py:meth:`temporalio.runtime.Runtime.default`. + + Raises: + ValueError: The bridge build has no extension with this name, or the config is not + valid. + """ + self._runtime = runtime or temporalio.runtime.Runtime.default() + self._bridge_extension = temporalio.bridge.extension.WorkerExtension( + self._runtime._core_runtime, name, config + ) + + @property + def name(self) -> str: + """Name the extension gives itself.""" + return self._bridge_extension.name + + @property + def runtime(self) -> temporalio.runtime.Runtime: + """Runtime the extension runs on.""" + return self._runtime diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index b3eb1a4d12..036cc42053 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -326,6 +326,7 @@ def on_eviction_hook( nonsticky_to_sticky_poll_ratio=1, no_remote_activities=True, disable_payload_error_limit=True, + extension=None, task_types=temporalio.bridge.worker.WorkerTaskTypes( enable_workflows=True, enable_local_activities=False, diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index 3c6effd80e..adb638feea 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -31,6 +31,7 @@ ) from ._activity import SharedStateManager, _ActivityWorker +from ._extension import WorkerExtension from ._interceptor import Interceptor from ._nexus import _NexusWorker from ._plugin import Plugin @@ -152,6 +153,7 @@ def __init__( ), disable_payload_error_limit: bool = False, max_workflow_task_external_storage_concurrency: int = _DEFAULT_WORKFLOW_TASK_EXTERNAL_STORAGE_CONCURRENCY, + extension: WorkerExtension | None = None, ) -> None: """Create a worker to process workflows and/or activities. @@ -345,6 +347,10 @@ def __init__( Please report any issues you encounter with this setting or if you feel the default should be changed. WARNING: This setting is experimental. + extension: A native extension that Core calls at fixed points of each + Workflow Task. It must use the same runtime as the client. Plugins + usually set it. + WARNING: This setting is experimental. """ config = WorkerConfig( @@ -393,6 +399,7 @@ def __init__( nexus_task_poller_behavior=nexus_task_poller_behavior, disable_payload_error_limit=disable_payload_error_limit, max_workflow_task_external_storage_concurrency=max_workflow_task_external_storage_concurrency, + extension=extension, ) plugins_from_client = cast( @@ -495,6 +502,9 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf or bridge_client.config.runtime or temporalio.runtime.Runtime.default() ) + extension = config.get("extension") + if extension is not None and extension.runtime is not self._runtime: + raise ValueError("Worker extension must use the same runtime as the client") activities = config.get("activities") if activities: _warn_if_activity_executor_max_workers_is_inconsistent(config) @@ -707,6 +717,7 @@ def check_activity(activity: str): disable_payload_error_limit=config.get( "disable_payload_error_limit", False ), + extension=extension._bridge_extension if extension else None, ), ) @@ -1030,6 +1041,7 @@ class WorkerConfig(TypedDict, total=False): nexus_task_poller_behavior: PollerBehavior disable_payload_error_limit: bool max_workflow_task_external_storage_concurrency: int + extension: WorkerExtension | None def _warn_if_activity_executor_max_workers_is_inconsistent( diff --git a/tests/worker/test_extension.py b/tests/worker/test_extension.py new file mode 100644 index 0000000000..1b9ca9ac7a --- /dev/null +++ b/tests/worker/test_extension.py @@ -0,0 +1,68 @@ +import uuid + +import pytest + +from temporalio import workflow +from temporalio.client import Client +from temporalio.runtime import Runtime, TelemetryConfig +from temporalio.worker import WorkerExtension +from tests.helpers import assert_task_fail_eventually, new_worker + +# Debug bridge builds link this extension for the SDK's tests. +TEST_EXTENSION = "temporalio-test" + + +def load_test_extension(config: bytes = b"") -> WorkerExtension: + try: + return WorkerExtension(TEST_EXTENSION, config) + except ValueError: + pytest.skip("Bridge build has no test extension") + + +@workflow.defn +class SimpleWorkflow: + @workflow.run + async def run(self) -> str: + return "done" + + +def test_unknown_extension_name(): + with pytest.raises(ValueError, match="no worker extension named not-linked"): + WorkerExtension("not-linked") + + +def test_extension_name(): + assert load_test_extension().name == TEST_EXTENSION + + +async def test_extension_lets_workflow_task_complete(client: Client): + async with new_worker( + client, SimpleWorkflow, extension=load_test_extension() + ) as worker: + result = await client.execute_workflow( + SimpleWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert result == "done" + + +async def test_extension_failure_fails_workflow_task(client: Client): + extension = load_test_extension(b"fail_before_complete") + async with new_worker(client, SimpleWorkflow, extension=extension) as worker: + handle = await client.start_workflow( + SimpleWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await assert_task_fail_eventually( + handle, message_contains="test extension refused the task" + ) + + +async def test_extension_must_use_client_runtime(client: Client): + load_test_extension() + other = Runtime(telemetry=TelemetryConfig()) + extension = WorkerExtension(TEST_EXTENSION, runtime=other) + with pytest.raises(ValueError, match="same runtime as the client"): + new_worker(client, SimpleWorkflow, extension=extension)