diff --git a/CHANGELOG.md b/CHANGELOG.md index be3444177..a3ca40991 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,10 @@ to include examples, links to docs, or any other relevant information. ### Changed +- System Nexus Signal-with-Start Workflow operations now use the typed + `WorkflowOutboundInterceptor.start_signal_with_start_workflow` interception point instead of + the generic `WorkflowOutboundInterceptor.start_nexus_operation` method. + ### Deprecated ### :boom: Breaking Changes diff --git a/scripts/gen_nexus_system_api.py b/scripts/gen_nexus_system_api.py index 82f40a8f0..0c410f009 100644 --- a/scripts/gen_nexus_system_api.py +++ b/scripts/gen_nexus_system_api.py @@ -131,6 +131,7 @@ def generate_nexus_system_api() -> None: str(wit_path), str(wit_deps_dir), "--native-api", + "--system-nexus", "--support-file", str(python_support_path), "--descriptors", diff --git a/temporalio/bridge/sdk-core b/temporalio/bridge/sdk-core index b860e3f14..85b71d7ec 160000 --- a/temporalio/bridge/sdk-core +++ b/temporalio/bridge/sdk-core @@ -1 +1 @@ -Subproject commit b860e3f14f68af64bd4d0e372b3c78b2ce9ecff1 +Subproject commit 85b71d7ecd4f2bf677fa1cee17f3fbc1ab10f1b9 diff --git a/temporalio/contrib/opentelemetry/_interceptor.py b/temporalio/contrib/opentelemetry/_interceptor.py index eb22f8be6..6dca4596e 100644 --- a/temporalio/contrib/opentelemetry/_interceptor.py +++ b/temporalio/contrib/opentelemetry/_interceptor.py @@ -33,6 +33,7 @@ import temporalio.client import temporalio.converter import temporalio.exceptions +import temporalio.nexus.system.workflow_service.models import temporalio.worker import temporalio.workflow from temporalio.exceptions import ApplicationError, ApplicationErrorCategory @@ -433,6 +434,10 @@ class _InputWithStringHeaders(Protocol): headers: Mapping[str, str] | None +class _InputWithModelHeaders(Protocol): + headers: Mapping[str, Any] | None + + class _InputWithOperationContext(Generic[_ContextT], Protocol): ctx: _ContextT @@ -684,6 +689,18 @@ def _context_carrier_to_headers( } return headers + def _context_carrier_to_model_headers( + self, + carrier: _CarrierDict, + headers: Mapping[str, Any] | None, + ) -> Mapping[str, Any]: + if carrier: + return { + **(headers or {}), + self.header_key: carrier, + } + return headers or {} + def _completed_span( self, span_name: str, @@ -691,6 +708,7 @@ def _completed_span( link_context_carrier: _CarrierDict | None = None, add_to_outbound: _InputWithHeaders | None = None, add_to_outbound_str: _InputWithStringHeaders | None = None, + add_to_outbound_model: _InputWithModelHeaders | None = None, new_span_even_on_replay: bool = False, additional_attributes: opentelemetry.util.types.Attributes = None, exception: Exception | None = None, @@ -742,6 +760,11 @@ def _completed_span( updated_context_carrier, add_to_outbound_str.headers ) + if add_to_outbound_model: + add_to_outbound_model.headers = self._context_carrier_to_model_headers( + updated_context_carrier, add_to_outbound_model.headers + ) + def _set_on_context( self, context: opentelemetry.context.Context ) -> opentelemetry.context.Context: @@ -830,6 +853,19 @@ async def start_nexus_operation( return await super().start_nexus_operation(input) + async def start_signal_with_start_workflow( + self, + request: temporalio.nexus.system.workflow_service.models.SignalWithStartWorkflowRequest, + ) -> temporalio.workflow.NexusOperationHandle[ + temporalio.nexus.system.workflow_service.models.SignalWithStartWorkflowResponse + ]: + self.root._completed_span( + "SignalWithStartWorkflow", + kind=opentelemetry.trace.SpanKind.CLIENT, + add_to_outbound_model=request, + ) + return await super().start_signal_with_start_workflow(request) + def _carrier_to_nexus_headers( carrier: _CarrierDict, initial: Mapping[str, str] | None = None diff --git a/temporalio/contrib/opentelemetry/_otel_interceptor.py b/temporalio/contrib/opentelemetry/_otel_interceptor.py index c120fcd03..ff07f0f1d 100644 --- a/temporalio/contrib/opentelemetry/_otel_interceptor.py +++ b/temporalio/contrib/opentelemetry/_otel_interceptor.py @@ -32,6 +32,7 @@ import temporalio.api.common.v1 import temporalio.client import temporalio.converter +import temporalio.nexus.system.workflow_service.models import temporalio.worker import temporalio.workflow from temporalio.contrib.opentelemetry._tracer_provider import ( @@ -600,3 +601,15 @@ async def start_nexus_operation( ): input.headers = _context_to_nexus_headers(input.headers or {}) return await super().start_nexus_operation(input) + + async def start_signal_with_start_workflow( + self, + request: temporalio.nexus.system.workflow_service.models.SignalWithStartWorkflowRequest, + ) -> temporalio.workflow.NexusOperationHandle[ + temporalio.nexus.system.workflow_service.models.SignalWithStartWorkflowResponse + ]: + with self._workflow_maybe_span( + "SignalWithStartWorkflow", kind=opentelemetry.trace.SpanKind.CLIENT + ): + request.headers = _context_to_headers(request.headers or {}) + return await super().start_signal_with_start_workflow(request) diff --git a/temporalio/nexus/system/workflow_service/_system_nexus_interceptor.py b/temporalio/nexus/system/workflow_service/_system_nexus_interceptor.py new file mode 100644 index 000000000..8f21b9bdb --- /dev/null +++ b/temporalio/nexus/system/workflow_service/_system_nexus_interceptor.py @@ -0,0 +1,96 @@ +# Generated by nexgen v0.2.2. DO NOT EDIT! + +from __future__ import annotations + +import abc +import typing + +from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT + +from . import models + +if typing.TYPE_CHECKING: + import temporalio.workflow + from temporalio.worker._interceptor import StartNexusOperationInput + + +__all__ = [ + "_start_system_nexus_operation", + "_SystemNexusWorkflowOutboundInterceptorBase", + "_SystemNexusWorkflowOutboundInterceptorTerminal", +] + + +_InputT = typing.TypeVar("_InputT") +_OutputT = typing.TypeVar("_OutputT") + + +async def _start_system_nexus_operation( + interceptor: _SystemNexusWorkflowOutboundInterceptorBase, + input: StartNexusOperationInput[_InputT, _OutputT], +) -> temporalio.workflow.NexusOperationHandle[_OutputT]: + if ( + input.service == "temporal.api.workflowservice.v1.WorkflowService" + and input.operation_name == "SignalWithStartWorkflowExecution" + ): + typed_input = typing.cast( + "StartNexusOperationInput[models.SignalWithStartWorkflowRequest, models.SignalWithStartWorkflowResponse]", + input, + ) + # The dispatch check above establishes that this operation's response type is _OutputT. + return typing.cast( + "temporalio.workflow.NexusOperationHandle[_OutputT]", + await interceptor.start_signal_with_start_workflow(typed_input.input), + ) + raise ValueError( + f"unsupported System Nexus operation: {input.service}/{input.operation_name}" + ) + + +class _SystemNexusWorkflowOutboundInterceptorBase(abc.ABC): + @abc.abstractmethod + def _next_system_nexus_interceptor( + self, + ) -> _SystemNexusWorkflowOutboundInterceptorBase: ... + + async def start_signal_with_start_workflow( + self, request: models.SignalWithStartWorkflowRequest + ) -> temporalio.workflow.NexusOperationHandle[ + models.SignalWithStartWorkflowResponse + ]: + """Intercept the SignalWithStartWorkflow operation.""" + return await self._next_system_nexus_interceptor().start_signal_with_start_workflow( + request + ) + + +class _SystemNexusWorkflowOutboundInterceptorTerminal(abc.ABC): + @abc.abstractmethod + async def _outbound_start_nexus_operation( + self, + input: StartNexusOperationInput[_InputT, _OutputT], + ) -> temporalio.workflow.NexusOperationHandle[_OutputT]: ... + + async def start_signal_with_start_workflow( + self, request: models.SignalWithStartWorkflowRequest + ) -> temporalio.workflow.NexusOperationHandle[ + models.SignalWithStartWorkflowResponse + ]: + from temporalio.worker._interceptor import StartNexusOperationInput + from temporalio.workflow import NexusOperationCancellationType + + return await self._outbound_start_nexus_operation( + StartNexusOperationInput( + endpoint=TEMPORAL_SYSTEM_ENDPOINT, + service="temporal.api.workflowservice.v1.WorkflowService", + operation="SignalWithStartWorkflowExecution", + input=request, + output_type=models.SignalWithStartWorkflowResponse, + schedule_to_close_timeout=None, + schedule_to_start_timeout=None, + start_to_close_timeout=None, + cancellation_type=NexusOperationCancellationType.WAIT_COMPLETED, + headers=None, + summary=None, + ) + ) diff --git a/temporalio/worker/_interceptor.py b/temporalio/worker/_interceptor.py index 4acf3c5d1..1bac84673 100644 --- a/temporalio/worker/_interceptor.py +++ b/temporalio/worker/_interceptor.py @@ -21,6 +21,9 @@ import temporalio.nexus import temporalio.nexus._util import temporalio.workflow +from temporalio.nexus.system.workflow_service._system_nexus_interceptor import ( + _SystemNexusWorkflowOutboundInterceptorBase, +) from temporalio.workflow import ContinueAsNewVersioningBehavior, VersioningIntent @@ -414,12 +417,14 @@ async def handle_update_handler(self, input: HandleUpdateInput) -> Any: return await self.next.handle_update_handler(input) -class WorkflowOutboundInterceptor: +class WorkflowOutboundInterceptor(_SystemNexusWorkflowOutboundInterceptorBase): """Outbound interceptor to wrap calls made from within workflows. This should be extended by any workflow outbound interceptors. """ + next: WorkflowOutboundInterceptor + def __init__(self, next: WorkflowOutboundInterceptor) -> None: """Create the outbound interceptor. @@ -429,6 +434,11 @@ def __init__(self, next: WorkflowOutboundInterceptor) -> None: """ self.next = next + def _next_system_nexus_interceptor( + self, + ) -> _SystemNexusWorkflowOutboundInterceptorBase: + return self.next + def continue_as_new(self, input: ContinueAsNewInput) -> NoReturn: """Called for every :py:func:`temporalio.workflow.continue_as_new` call.""" self.next.continue_as_new(input) diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 69b416f26..fca73c266 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -60,6 +60,10 @@ import temporalio.nexus.system import temporalio.workflow from temporalio.converter import StorageDriverStoreContext, StorageDriverWorkflowInfo +from temporalio.nexus.system.workflow_service._system_nexus_interceptor import ( + _start_system_nexus_operation, + _SystemNexusWorkflowOutboundInterceptorTerminal, +) from temporalio.service import __version__ from ..api.failure.v1.message_pb2 import Failure @@ -1732,7 +1736,23 @@ async def workflow_start_nexus_operation( headers: Mapping[str, str] | None, summary: str | None, ) -> temporalio.workflow.NexusOperationHandle[OutputT]: - # start_nexus_operation + if temporalio.nexus.system.is_system_endpoint(endpoint): + return await _start_system_nexus_operation( + self._outbound, + StartNexusOperationInput( + endpoint=temporalio.nexus.system.TEMPORAL_SYSTEM_ENDPOINT, + service=service, + operation=operation, + input=input, + output_type=output_type, + schedule_to_close_timeout=schedule_to_close_timeout, + schedule_to_start_timeout=schedule_to_start_timeout, + start_to_close_timeout=start_to_close_timeout, + cancellation_type=cancellation_type, + headers=None, + summary=summary, + ), + ) return await self._outbound.start_nexus_operation( StartNexusOperationInput( endpoint=endpoint, @@ -3131,11 +3151,18 @@ async def handle_update_handler(self, input: HandleUpdateInput) -> Any: return handler(*input.args) -class _WorkflowOutboundImpl(WorkflowOutboundInterceptor): +class _WorkflowOutboundImpl( + _SystemNexusWorkflowOutboundInterceptorTerminal, WorkflowOutboundInterceptor +): def __init__(self, instance: _WorkflowInstanceImpl) -> None: # type: ignore # We are intentionally not calling the base class's __init__ here self._instance = instance + async def _outbound_start_nexus_operation( + self, input: StartNexusOperationInput[InputT, OutputT] + ) -> temporalio.workflow.NexusOperationHandle[OutputT]: + return await self._instance._outbound_start_nexus_operation(input) + def continue_as_new(self, input: ContinueAsNewInput) -> NoReturn: self._instance._outbound_continue_as_new(input) diff --git a/tests/contrib/opentelemetry/test_opentelemetry.py b/tests/contrib/opentelemetry/test_opentelemetry.py index 1bab931ac..0ea9530e6 100644 --- a/tests/contrib/opentelemetry/test_opentelemetry.py +++ b/tests/contrib/opentelemetry/test_opentelemetry.py @@ -88,6 +88,34 @@ class TracingWorkflowActionActivity: fail_on_non_replay_before_complete: bool = False +@workflow.defn +class SignalWithStartHeaderWorkflow: + def __init__(self) -> None: + self._signaled = False + + @workflow.run + async def run(self) -> bool: + await workflow.wait_condition(lambda: self._signaled) + return "_tracer-data" in workflow.info().headers + + @workflow.signal + def notify(self) -> None: + self._signaled = True + + +@workflow.defn +class SignalWithStartCallerWorkflow: + @workflow.run + async def run(self, target_id: str, task_queue: str) -> str: + handle = await workflow.signal_with_start_workflow( + SignalWithStartHeaderWorkflow.run, + id=target_id, + task_queue=task_queue, + signal=SignalWithStartHeaderWorkflow.notify, + ) + return handle.id + + @dataclass class TracingWorkflowActionContinueAsNew: param: TracingWorkflowParam @@ -229,6 +257,41 @@ def update_validator(self) -> None: pass +@pytest.mark.requires_local_server +async def test_workflow_signal_with_start_propagates_trace_headers( + client: Client, env: WorkflowEnvironment +): + if env.supports_time_skipping: + pytest.skip("Nexus tests don't work with the Java test server") + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = get_tracer(__name__, tracer_provider=provider) + config = client.config() + config["interceptors"] = [TracingInterceptor(tracer)] + client = Client(**config) + + async with Worker( + client, + task_queue=f"signal-with-start-{uuid.uuid4()}", + workflows=[SignalWithStartCallerWorkflow, SignalWithStartHeaderWorkflow], + workflow_runner=UnsandboxedWorkflowRunner(), + ) as worker: + target_id = f"signal-with-start-target-{uuid.uuid4()}" + with tracer.start_as_current_span("signal-with-start"): + caller = await client.start_workflow( + SignalWithStartCallerWorkflow.run, + args=[target_id, worker.task_queue], + id=f"signal-with-start-caller-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await caller.result() == target_id + assert await client.get_workflow_handle(target_id).result() is True + assert any( + span.name == "SignalWithStartWorkflow" for span in exporter.get_finished_spans() + ) + + async def test_opentelemetry_tracing(client: Client, env: WorkflowEnvironment): # TODO(cretz): Fix if env.supports_time_skipping: diff --git a/tests/contrib/opentelemetry/test_opentelemetry_plugin.py b/tests/contrib/opentelemetry/test_opentelemetry_plugin.py index 12d0c5972..022645fee 100644 --- a/tests/contrib/opentelemetry/test_opentelemetry_plugin.py +++ b/tests/contrib/opentelemetry/test_opentelemetry_plugin.py @@ -122,6 +122,34 @@ async def run(self): return +@workflow.defn +class SignalWithStartHeaderWorkflow: + def __init__(self) -> None: + self._signaled = False + + @workflow.run + async def run(self) -> bool: + await workflow.wait_condition(lambda: self._signaled) + return "_tracer-data" in workflow.info().headers + + @workflow.signal + def notify(self) -> None: + self._signaled = True + + +@workflow.defn +class SignalWithStartCallerWorkflow: + @workflow.run + async def run(self, target_id: str, task_queue: str) -> str: + handle = await workflow.signal_with_start_workflow( + SignalWithStartHeaderWorkflow.run, + id=target_id, + task_queue=task_queue, + signal=SignalWithStartHeaderWorkflow.notify, + ) + return handle.id + + async def test_otel_tracing_basic(client: Client, reset_otel_tracer_provider: Any): # type: ignore[reportUnusedParameter] exporter = InMemorySpanExporter() provider = create_tracer_provider() @@ -169,6 +197,41 @@ async def test_otel_tracing_basic(client: Client, reset_otel_tracer_provider: An ) +@pytest.mark.requires_local_server +async def test_workflow_signal_with_start_propagates_trace_headers( + client: Client, + env: WorkflowEnvironment, + reset_otel_tracer_provider: Any, # type: ignore[reportUnusedParameter] +): + if env.supports_time_skipping: + pytest.skip("Nexus tests don't work with the Java test server") + exporter = InMemorySpanExporter() + provider = create_tracer_provider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + opentelemetry.trace.set_tracer_provider(provider) + config = client.config() + config["plugins"] = [OpenTelemetryPlugin(add_temporal_spans=True)] + client = Client(**config) + + async with new_worker( + client, SignalWithStartCallerWorkflow, SignalWithStartHeaderWorkflow + ) as worker: + target_id = f"signal-with-start-target-{uuid.uuid4()}" + with get_tracer(__name__).start_as_current_span("signal-with-start"): + caller = await client.start_workflow( + SignalWithStartCallerWorkflow.run, + args=[target_id, worker.task_queue], + id=f"signal-with-start-caller-{uuid.uuid4()}", + task_queue=worker.task_queue, + execution_timeout=timedelta(seconds=3), + ) + assert await caller.result() == target_id + assert await client.get_workflow_handle(target_id).result() is True + assert any( + span.name == "SignalWithStartWorkflow" for span in exporter.get_finished_spans() + ) + + @workflow.defn class ComprehensiveWorkflow: def __init__(self) -> None: diff --git a/tests/nexus/test_temporal_system_nexus.py b/tests/nexus/test_temporal_system_nexus.py index 6a9bc9959..69bed1a9c 100644 --- a/tests/nexus/test_temporal_system_nexus.py +++ b/tests/nexus/test_temporal_system_nexus.py @@ -36,7 +36,6 @@ from temporalio.testing import WorkflowEnvironment from temporalio.worker import ( Interceptor, - StartNexusOperationInput, Worker, WorkflowInboundInterceptor, WorkflowInterceptorClassInput, @@ -251,11 +250,16 @@ def init(self, outbound: WorkflowOutboundInterceptor) -> None: class _TracingWorkflowOutboundInterceptor(WorkflowOutboundInterceptor): - async def start_nexus_operation( - self, input: StartNexusOperationInput[Any, Any] - ) -> workflow.NexusOperationHandle[Any]: - interceptor_traces.append(("workflow.start_nexus_operation", input)) - return await super().start_nexus_operation(input) + async def start_signal_with_start_workflow( + self, request: workflow_service_models.SignalWithStartWorkflowRequest + ) -> workflow.NexusOperationHandle[ + workflow_service_models.SignalWithStartWorkflowResponse + ]: + request.headers = {**(request.headers or {}), "interceptor-header": "value"} + interceptor_traces.append( + ("workflow.start_signal_with_start_workflow", request) + ) + return await super().start_signal_with_start_workflow(request) def _assert_stored_payloads_include( @@ -270,15 +274,15 @@ def _assert_stored_payloads_include( assert expected_payload_data.issubset(stored_payload_data) -def _assert_start_nexus_operation_interceptor_trace() -> None: +def _assert_signal_with_start_workflow_interceptor_trace() -> None: assert len(interceptor_traces) == 1 trace_name, trace_value = interceptor_traces.pop() - assert trace_name == "workflow.start_nexus_operation" - trace_input = cast(StartNexusOperationInput[Any, Any], trace_value) - request = trace_input.input + assert trace_name == "workflow.start_signal_with_start_workflow" + request = cast(workflow_service_models.SignalWithStartWorkflowRequest, trace_value) assert request.id == "system-nexus-workflow-id" assert request.signal == "test-signal" assert request.workflow == "test-workflow" + assert request.headers == {"interceptor-header": "value"} class _MarkingPayloadVisitor(VisitorFunctions): @@ -711,7 +715,7 @@ async def test_external_workflow_handle_signal_with_start_workflow_uses_system_n b'"details-value"', }, ) - _assert_start_nexus_operation_interceptor_trace() + _assert_signal_with_start_workflow_interceptor_trace() # Cloud namespaces created by CI do not have the System Nexus dynamic config.