diff --git a/.github/lambda-layer-publish.toml b/.github/lambda-layer-publish.toml index f141c2a0..4fd02156 100644 --- a/.github/lambda-layer-publish.toml +++ b/.github/lambda-layer-publish.toml @@ -1,2 +1,2 @@ [layer] -sdk-version = "2.0.0" +sdk-version = "3.0.0" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md index 5ff978c8..9c23dd9b 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md @@ -35,9 +35,9 @@ tests/ # contract tests for the templates and handlers The 20 invocation and 20 execution requirements reuse the same scenario handlers; the view is selected per function through the `OTEL_PLUGIN_MODE` -environment variable, which `common.otel_plugin()` reads to pick -`InvocationOtelPlugin` or `ExecutionOtelPlugin`. `template.yaml` deploys only the -view named by its `OtelSuite` parameter. +environment variable, which `common.otel_plugin_factory()` reads to pick +`InvocationOtelPluginFactory` or `ExecutionOtelPluginFactory`. `template.yaml` +deploys only the view named by its `OtelSuite` parameter. ## Scenarios @@ -182,8 +182,9 @@ write access; the runner identity needs list, read, and cleanup access. 1. Find or add the requirement in the conformance repository under `test-requirements//.yaml`. New requirement IDs must be registered there first. -2. Add `src/otel__.py` exporting `handler`. Select the plugin with - `common.otel_plugin()` and guard the input with `common.require_scenario()`. +2. Add `src/otel__.py` exporting `handler`. Select the plugin factory + with `common.otel_plugin_factory()` and guard the input with + `common.require_scenario()`. Use the SDK's real API; never hand-roll behavior to force an expected result. 3. Register the function in `template.yaml` (or `template-long-running.yaml`) with `Handler: .handler` and `TestDescription: [""]`, and add a diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/pyproject.toml b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/pyproject.toml index 3c1f251f..3deffbdf 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/pyproject.toml @@ -8,8 +8,8 @@ version = "0.0.0" description = "OpenTelemetry conformance test handlers for the AWS Durable Execution SDK for Python, exercised by the aws-durable-execution-conformance-tests OTel suites." requires-python = ">=3.11" dependencies = [ - "aws-durable-execution-sdk-python==2.0.1", - "aws-durable-execution-sdk-python-otel==1.0.0", + "aws-durable-execution-sdk-python==3.0.0", + "aws-durable-execution-sdk-python-otel==2.0.0", ] [tool.hatch.build.targets.wheel] diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/common.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/common.py index 7c8b7c59..19ecbfd0 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/common.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/common.py @@ -9,20 +9,28 @@ from collections.abc import Mapping from typing import Any -from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPluginFactory, +) from aws_durable_execution_sdk_python_otel import ( - ExecutionOtelPlugin, - InvocationOtelPlugin, + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, OtelPluginConfig, ) -def otel_plugin() -> DurableInstrumentationPlugin: - """Select the telemetry view configured for this deployed function.""" +def otel_plugin_factory() -> DurableInstrumentationPluginFactory: + """Select the telemetry view configured for this deployed function. + + Returns a factory, which is what ``durable_execution(plugins=[...])`` takes: + the SDK calls its ``create_plugin`` once per invocation to build that + invocation's plugin. The view is still resolved once, when the handler module + is imported. + """ if os.environ.get("OTEL_PLUGIN_MODE") == "execution": - return ExecutionOtelPlugin(OtelPluginConfig()) - return InvocationOtelPlugin(OtelPluginConfig()) + return ExecutionOtelPluginFactory(OtelPluginConfig()) + return InvocationOtelPluginFactory(OtelPluginConfig()) def require_scenario(event: Mapping[str, Any], expected: str) -> None: diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_10_wait_for_callback.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_10_wait_for_callback.py index 3807d2ae..107362c9 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_10_wait_for_callback.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_10_wait_for_callback.py @@ -9,7 +9,7 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python.types import WaitForCallbackContext -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario def submit_callback( @@ -19,7 +19,7 @@ def submit_callback( return None -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "wait-for-callback") return context.wait_for_callback( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_11_chained_invoke.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_11_chained_invoke.py index 0a43b520..777d222c 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_11_chained_invoke.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_11_chained_invoke.py @@ -9,10 +9,10 @@ from typing import Any from aws_durable_execution_sdk_python import DurableContext, durable_execution -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler( event: dict[str, Any], context: DurableContext, @@ -25,7 +25,7 @@ def handler( ) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def target_handler( event: dict[str, Any], _context: DurableContext, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_12_child_context_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_12_child_context_failure.py index 2d06b9a3..365c2923 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_12_child_context_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_12_child_context_failure.py @@ -12,7 +12,7 @@ durable_execution, durable_with_child_context, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_with_child_context @@ -20,7 +20,7 @@ def fail_child_context(_context: DurableContext) -> None: raise RuntimeError("Intentional child-context failure") -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "child-context-failure") context.run_in_child_context( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_13_parallel_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_13_parallel_failure.py index c36fadf7..416b8bec 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_13_parallel_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_13_parallel_failure.py @@ -13,7 +13,7 @@ durable_parallel_branch, ) from aws_durable_execution_sdk_python.config import ParallelConfig -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_parallel_branch(name="otel-failed-parallel-branch") @@ -21,7 +21,7 @@ def fail_parallel_branch(_context: DurableContext) -> None: raise RuntimeError("Intentional parallel branch failure") -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "parallel-failure") result = context.parallel( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_14_map_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_14_map_failure.py index b77b1ddd..6e1bd35c 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_14_map_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_14_map_failure.py @@ -10,7 +10,7 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python.config import MapConfig -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario def fail_map_item( @@ -22,7 +22,7 @@ def fail_map_item( raise RuntimeError("Intentional map iteration failure") -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "map-failure") result = context.map( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_15_wait_interrupted.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_15_wait_interrupted.py index f8c225dc..584259c0 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_15_wait_interrupted.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_15_wait_interrupted.py @@ -9,10 +9,10 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python.config import Duration -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "wait-interrupted") context.wait( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_16_wait_for_condition_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_16_wait_for_condition_failure.py index db60d3c8..23aaa72a 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_16_wait_for_condition_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_16_wait_for_condition_failure.py @@ -14,7 +14,7 @@ WaitForConditionConfig, WaitForConditionDecision, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario def fail_condition_check( @@ -31,7 +31,7 @@ def continue_condition( return WaitForConditionDecision.continue_waiting(Duration.from_seconds(1)) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "wait-for-condition-failure") context.wait_for_condition( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_17_wait_for_callback_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_17_wait_for_callback_failure.py index 8aa55881..5b0a1b9f 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_17_wait_for_callback_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_17_wait_for_callback_failure.py @@ -9,7 +9,7 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python.types import WaitForCallbackContext -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario def submit_failed_callback( @@ -19,7 +19,7 @@ def submit_failed_callback( return None -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "wait-for-callback-failure") context.wait_for_callback( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_18_chained_invoke_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_18_chained_invoke_failure.py index c8acdd28..d7e4dd9b 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_18_chained_invoke_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_18_chained_invoke_failure.py @@ -9,10 +9,10 @@ from typing import Any from aws_durable_execution_sdk_python import DurableContext, durable_execution -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "chained-invoke-failure") context.invoke( @@ -22,7 +22,7 @@ def handler(event: dict[str, Any], context: DurableContext) -> None: ) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def target_handler( _event: dict[str, Any], _context: DurableContext, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_19_execution_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_19_execution_failure.py index 887a488c..cd14daf3 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_19_execution_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_19_execution_failure.py @@ -8,10 +8,10 @@ from typing import Any from aws_durable_execution_sdk_python import DurableContext, durable_execution -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], _context: DurableContext) -> None: require_scenario(event, "execution-failure") raise RuntimeError("Intentional execution failure") diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_1_success.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_1_success.py index 0409588f..1fb1ef2e 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_1_success.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_1_success.py @@ -13,7 +13,7 @@ durable_execution, durable_step, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -21,7 +21,7 @@ def complete_successfully(_step_context: StepContext) -> str: return "success" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "success") return context.step(complete_successfully(), name="otel-success") diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_20_virtual_context.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_20_virtual_context.py index 1f957bdb..2e867ae1 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_20_virtual_context.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_20_virtual_context.py @@ -13,7 +13,7 @@ durable_with_child_context, ) from aws_durable_execution_sdk_python.config import ChildConfig -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_with_child_context @@ -21,7 +21,7 @@ def run_virtual_context(_context: DurableContext) -> str: return "virtual-complete" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "virtual-context") return context.run_in_child_context( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_2_wait_resume.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_2_wait_resume.py index 21041a94..40e13596 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_2_wait_resume.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_2_wait_resume.py @@ -14,7 +14,7 @@ durable_step, ) from aws_durable_execution_sdk_python.config import Duration -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -22,7 +22,7 @@ def complete_after_resume(_step_context: StepContext) -> str: return "resumed" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "wait-resume") context.wait(Duration.from_seconds(1), name="otel-wait") diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_3_retry.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_3_retry.py index 761f76b4..87f51ad9 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_3_retry.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_3_retry.py @@ -18,7 +18,7 @@ RetryStrategyConfig, create_retry_strategy, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -28,7 +28,7 @@ def succeed_on_retry(step_context: StepContext) -> str: return "retried" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "retry") retry_strategy = create_retry_strategy( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_4_terminal_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_4_terminal_failure.py index 5850757e..cd4f81d0 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_4_terminal_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_4_terminal_failure.py @@ -18,7 +18,7 @@ RetryStrategyConfig, create_retry_strategy, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -26,7 +26,7 @@ def fail_terminally(_step_context: StepContext) -> None: raise RuntimeError("Intentional terminal failure") -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> None: require_scenario(event, "terminal-failure") retry_strategy = create_retry_strategy( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_5_child_context.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_5_child_context.py index b9d47534..61d049e0 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_5_child_context.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_5_child_context.py @@ -14,7 +14,7 @@ durable_step, durable_with_child_context, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -27,7 +27,7 @@ def run_child_workflow(context: DurableContext) -> str: return context.step(complete_child_step(), name="otel-child-step") -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "child-context") return context.run_in_child_context( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_6_parallel.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_6_parallel.py index e2ac6adb..e7b54e39 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_6_parallel.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_6_parallel.py @@ -15,7 +15,7 @@ durable_step, ) from aws_durable_execution_sdk_python.config import ParallelConfig -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -39,7 +39,7 @@ def run_parallel_branch_b(context: DurableContext) -> str: ) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> list[str]: require_scenario(event, "parallel-hierarchy") return context.parallel( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_7_map.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_7_map.py index b35ad6e7..e69c68ab 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_7_map.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_7_map.py @@ -15,7 +15,7 @@ durable_step, ) from aws_durable_execution_sdk_python.config import MapConfig -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -35,7 +35,7 @@ def process_map_item( ) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> list[int]: require_scenario(event, "map-hierarchy") return context.map( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_8_handled_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_8_handled_failure.py index 44fe595d..20407101 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_8_handled_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_8_handled_failure.py @@ -20,7 +20,7 @@ RetryStrategyConfig, create_retry_strategy, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario @durable_step @@ -33,7 +33,7 @@ def recover_after_failure(_step_context: StepContext) -> str: return "recovered" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "handled-failure") retry_strategy = create_retry_strategy( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_9_wait_for_condition.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_9_wait_for_condition.py index 5daf3b5e..578b9b10 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_9_wait_for_condition.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_9_wait_for_condition.py @@ -14,7 +14,7 @@ WaitForConditionConfig, WaitForConditionDecision, ) -from common import otel_plugin, require_scenario +from common import otel_plugin_factory, require_scenario def increment_condition( @@ -33,7 +33,7 @@ def stop_after_second_attempt( return WaitForConditionDecision.continue_waiting(Duration.from_seconds(1)) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> int: require_scenario(event, "wait-for-condition") return context.wait_for_condition( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_1_wait.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_1_wait.py index 3278f3c5..8b65a555 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_1_wait.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_1_wait.py @@ -14,7 +14,7 @@ durable_step, ) from aws_durable_execution_sdk_python.config import Duration -from common import long_delay_seconds, otel_plugin, require_scenario +from common import long_delay_seconds, otel_plugin_factory, require_scenario @durable_step @@ -22,7 +22,7 @@ def complete_after_long_wait(_step_context: StepContext) -> str: return "resumed" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "long-wait") context.wait( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_2_retry.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_2_retry.py index 58db3540..59d337b6 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_2_retry.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_2_retry.py @@ -18,7 +18,7 @@ RetryStrategyConfig, create_retry_strategy, ) -from common import long_delay_seconds, otel_plugin, require_scenario +from common import long_delay_seconds, otel_plugin_factory, require_scenario @durable_step @@ -28,7 +28,7 @@ def succeed_after_long_retry(step_context: StepContext) -> str: return "retried" -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "long-retry") retry_strategy = create_retry_strategy( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_3_callback.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_3_callback.py index a4078f10..d187b455 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_3_callback.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_3_callback.py @@ -9,7 +9,7 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python.types import WaitForCallbackContext -from common import long_delay_seconds, otel_plugin, require_scenario +from common import long_delay_seconds, otel_plugin_factory, require_scenario def submit_callback( @@ -19,7 +19,7 @@ def submit_callback( return None -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler(event: dict[str, Any], context: DurableContext) -> str: require_scenario(event, "long-callback") long_delay_seconds(event) diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_4_chained_invoke.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_4_chained_invoke.py index a490d44d..170fef6a 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_4_chained_invoke.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_long_running_4_chained_invoke.py @@ -10,10 +10,10 @@ from aws_durable_execution_sdk_python import DurableContext, durable_execution from aws_durable_execution_sdk_python.config import Duration -from common import long_delay_seconds, otel_plugin, require_scenario +from common import long_delay_seconds, otel_plugin_factory, require_scenario -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def handler( event: dict[str, Any], context: DurableContext, @@ -26,7 +26,7 @@ def handler( ) -@durable_execution(plugins=[otel_plugin()]) +@durable_execution(plugins=[otel_plugin_factory()]) def target_handler( event: dict[str, Any], context: DurableContext, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py index 5908dd7d..2d086d7c 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py @@ -519,5 +519,5 @@ def test_common_selects_the_plugin_from_the_deployed_view() -> None: common: str = (SRC_DIR / "common.py").read_text(encoding="utf-8") assert 'os.environ.get("OTEL_PLUGIN_MODE") == "execution"' in common - assert "ExecutionOtelPlugin(OtelPluginConfig())" in common - assert "InvocationOtelPlugin(OtelPluginConfig())" in common + assert "ExecutionOtelPluginFactory(OtelPluginConfig())" in common + assert "InvocationOtelPluginFactory(OtelPluginConfig())" in common diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_hooks_retry.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_hooks_retry.py index a8cdc04b..4b218725 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_hooks_retry.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_hooks_retry.py @@ -79,6 +79,21 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: ) +class AttemptPluginFactory: + """Builds one :class:`AttemptPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> AttemptPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return AttemptPlugin() + + @durable_step def unreliable_operation(step_context: StepContext) -> str: # Fail on the first attempt, succeed on the second, using the SDK's built-in @@ -89,7 +104,7 @@ def unreliable_operation(step_context: StepContext) -> str: return "Operation succeeded" -@durable_execution(plugins=[AttemptPlugin()]) +@durable_execution(plugins=[AttemptPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: retry_config = RetryStrategyConfig( max_attempts=3, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_info_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_info_shape.py index 1a770cc5..0bb85e96 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_info_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_attempt_info_shape.py @@ -79,6 +79,21 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: _emit(record, self._execution_arn) +class AttemptInfoShapePluginFactory: + """Builds one :class:`AttemptInfoShapePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> AttemptInfoShapePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return AttemptInfoShapePlugin() + + @durable_step def flaky(step_context: StepContext) -> str: if step_context.attempt < 2: @@ -86,7 +101,7 @@ def flaky(step_context: StepContext) -> str: return "ok" -@durable_execution(plugins=[AttemptInfoShapePlugin()]) +@durable_execution(plugins=[AttemptInfoShapePluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: retry_config = RetryStrategyConfig( max_attempts=3, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_context_info_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_context_info_shape.py index 0debf47b..9e40bb39 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_context_info_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_context_info_shape.py @@ -79,6 +79,21 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: _emit(record, self._execution_arn) +class ContextInfoShapePluginFactory: + """Builds one :class:`ContextInfoShapePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> ContextInfoShapePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return ContextInfoShapePlugin() + + @durable_step def inner(_step_context: StepContext) -> str: return "x" @@ -94,7 +109,7 @@ def branch_b(_context: DurableContext) -> str: return "b-done" -@durable_execution(plugins=[ContextInfoShapePlugin()]) +@durable_execution(plugins=[ContextInfoShapePluginFactory()]) def handler(_event: Any, context: DurableContext) -> list[str]: result: BatchResult[str] = context.parallel( [ diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_error_isolation.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_error_isolation.py index 0a116f88..b378607c 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_error_isolation.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_error_isolation.py @@ -82,12 +82,27 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: raise RuntimeError("faulty attempt-end") +class FaultyPluginFactory: + """Builds one :class:`FaultyPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> FaultyPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return FaultyPlugin() + + @durable_step def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[FaultyPlugin()]) +@durable_execution(plugins=[FaultyPluginFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_external_update_on_invoke.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_external_update_on_invoke.py index 7abbc401..50da44b4 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_external_update_on_invoke.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_external_update_on_invoke.py @@ -65,7 +65,22 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) -@durable_execution(plugins=[ExternalUpdatePlugin()]) +class ExternalUpdatePluginFactory: + """Builds one :class:`ExternalUpdatePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> ExternalUpdatePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return ExternalUpdatePlugin() + + +@durable_execution(plugins=[ExternalUpdatePluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return "Wait completed" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_faulty_and_healthy.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_faulty_and_healthy.py index ef11896a..1c178404 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_faulty_and_healthy.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_faulty_and_healthy.py @@ -1,9 +1,10 @@ """10-17: Faulty plugin does not affect a healthy plugin. -Two plugins are registered together, in order: a faulty plugin whose every -exercised hook logs a line and then raises, and a healthy plugin that logs -normally. The exercised hooks span the full lifecycle: invocation-start, -operation-start, attempt-start, attempt-end, operation-end, and invocation-end. +Two plugins are registered together, through their factories, in order: a faulty +plugin whose every exercised hook logs a line and then raises, and a healthy +plugin that logs normally. The exercised hooks span the full lifecycle: +invocation-start, operation-start, attempt-start, attempt-end, operation-end, and +invocation-end. In the Python SDK the per-attempt hooks are the real ``on_user_function_start`` / ``on_user_function_end`` callbacks (the latter carries the attempt ``outcome``), and the operation hooks are ``on_operation_start`` / ``on_operation_end``. @@ -103,6 +104,21 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: raise RuntimeError("faulty invocation-end") +class FaultyPluginFactory: + """Builds one :class:`FaultyPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> FaultyPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return FaultyPlugin() + + class HealthyPlugin(DurableInstrumentationPlugin): def __init__(self) -> None: # Operation/attempt hooks do not carry the execution ARN, so capture it @@ -182,12 +198,27 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) +class HealthyPluginFactory: + """Builds one :class:`HealthyPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> HealthyPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return HealthyPlugin() + + @durable_step def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[FaultyPlugin(), HealthyPlugin()]) +@durable_execution(plugins=[FaultyPluginFactory(), HealthyPluginFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_first_invocation_flag.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_first_invocation_flag.py index 8b5ecefd..4b6ea43a 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_first_invocation_flag.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_first_invocation_flag.py @@ -47,7 +47,22 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) -@durable_execution(plugins=[FirstInvocationPlugin()]) +class FirstInvocationPluginFactory: + """Builds one :class:`FirstInvocationPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> FirstInvocationPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return FirstInvocationPlugin() + + +@durable_execution(plugins=[FirstInvocationPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return "Wait completed" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_info_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_info_shape.py index 75b53f7b..19a20ddb 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_info_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_info_shape.py @@ -58,7 +58,22 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: _emit(record, info.execution_arn) -@durable_execution(plugins=[InvocationInfoShapePlugin()]) +class InvocationInfoShapePluginFactory: + """Builds one :class:`InvocationInfoShapePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> InvocationInfoShapePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return InvocationInfoShapePlugin() + + +@durable_execution(plugins=[InvocationInfoShapePluginFactory()]) def handler(event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return f"done-{event}" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_lifecycle.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_lifecycle.py index ddd22673..4deed16d 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_lifecycle.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_invocation_lifecycle.py @@ -1,9 +1,10 @@ """10-1: Plugin invocation lifecycle hooks (start and end on a single invocation). Registers an instrumentation plugin through the SDK's real ``plugins=[...]`` -parameter on ``durable_execution``. The plugin emits its lines from the SDK's -``on_invocation_start`` / ``on_invocation_end`` hooks; the step body logs its -running line via the SDK-provided step context logger (mirrors handler 1-7). +parameter on ``durable_execution``, which takes the plugin's factory. The plugin +emits its lines from the SDK's ``on_invocation_start`` / ``on_invocation_end`` +hooks; the step body logs its running line via the SDK-provided step context +logger (mirrors handler 1-7). """ import json @@ -50,13 +51,28 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) +class LifecyclePluginFactory: + """Builds one :class:`LifecyclePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> LifecyclePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return LifecyclePlugin() + + @durable_step def greet(step_context: StepContext, name: str) -> str: step_context.logger.info(f"Greeting step running for: {name}") return f"Hello, {name}!" -@durable_execution(plugins=[LifecyclePlugin()]) +@durable_execution(plugins=[LifecyclePluginFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_multiple_plugins.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_multiple_plugins.py index f9669a8e..2600ed32 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_multiple_plugins.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_multiple_plugins.py @@ -1,8 +1,8 @@ """10-5: Multiple registered plugins all receive lifecycle hooks. Two instrumentation plugins are registered together, in order A then B, through -the SDK's real ``plugins=[...]`` parameter. Each emits its own prefixed lines -from the invocation-start / invocation-end hooks. +the SDK's real ``plugins=[...]`` parameter, which takes their factories. Each +emits its own prefixed lines from the invocation-start / invocation-end hooks. """ import json @@ -45,6 +45,21 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) +class PluginAFactory: + """Builds one :class:`PluginA` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> PluginA: + """Return this invocation's plugin. ``info`` is unused.""" + return PluginA() + + class PluginB(DurableInstrumentationPlugin): def on_invocation_start(self, info: InvocationStartInfo) -> None: _emit( @@ -60,12 +75,27 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) +class PluginBFactory: + """Builds one :class:`PluginB` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> PluginB: + """Return this invocation's plugin. ``info`` is unused.""" + return PluginB() + + @durable_step def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[PluginA(), PluginB()]) +@durable_execution(plugins=[PluginAFactory(), PluginBFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_nested_parent_linkage.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_nested_parent_linkage.py index 82f6ad8b..ccb58191 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_nested_parent_linkage.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_nested_parent_linkage.py @@ -56,6 +56,21 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) +class ParentLinkagePluginFactory: + """Builds one :class:`ParentLinkagePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> ParentLinkagePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return ParentLinkagePlugin() + + @durable_step def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" @@ -66,7 +81,7 @@ def child_operation(ctx: DurableContext, name: str) -> str: return ctx.step(greet(name)) -@durable_execution(plugins=[ParentLinkagePlugin()]) +@durable_execution(plugins=[ParentLinkagePluginFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.run_in_child_context(child_operation(str(event))) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change.py index 14c36c2a..f3220bf7 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change.py @@ -58,12 +58,27 @@ def on_operation_change(self, info: OperationChangeInfo) -> None: ) +class OperationChangePluginFactory: + """Builds one :class:`OperationChangePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> OperationChangePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return OperationChangePlugin() + + @durable_step def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[OperationChangePlugin()]) +@durable_execution(plugins=[OperationChangePluginFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change_shape.py index bbb2988b..c1982a33 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_change_shape.py @@ -78,12 +78,27 @@ def on_operation_change(self, info: OperationChangeInfo) -> None: _emit(record, self._execution_arn) +class OperationChangeShapePluginFactory: + """Builds one :class:`OperationChangeShapePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> OperationChangeShapePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return OperationChangeShapePlugin() + + @durable_step def greet(_step_context: StepContext) -> str: return "task-a" -@durable_execution(plugins=[OperationChangeShapePlugin()]) +@durable_execution(plugins=[OperationChangeShapePluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: result: str = context.step(greet(), name="greet") return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_info_shape.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_info_shape.py index 07e4a55f..11de1570 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_info_shape.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_info_shape.py @@ -74,12 +74,27 @@ def on_operation_end(self, info: OperationEndInfo) -> None: _emit(_operation_record("operation-end", info), self._execution_arn) +class OperationInfoShapePluginFactory: + """Builds one :class:`OperationInfoShapePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> OperationInfoShapePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return OperationInfoShapePlugin() + + @durable_step def greet(_step_context: StepContext) -> str: return "task-a" -@durable_execution(plugins=[OperationInfoShapePlugin()]) +@durable_execution(plugins=[OperationInfoShapePluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: result: str = context.step(greet(), name="greet") return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_lifecycle.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_lifecycle.py index 48b25cf4..10b1a943 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_lifecycle.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_operation_lifecycle.py @@ -72,12 +72,27 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) +class OperationLifecyclePluginFactory: + """Builds one :class:`OperationLifecyclePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> OperationLifecyclePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return OperationLifecyclePlugin() + + @durable_step def greet(_step_context: StepContext, name: str) -> str: return f"Hello, {name}!" -@durable_execution(plugins=[OperationLifecyclePlugin()]) +@durable_execution(plugins=[OperationLifecyclePluginFactory()]) def handler(event: Any, context: DurableContext) -> str: result: str = context.step(greet(event)) return result diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_parallel_branch_hooks.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_parallel_branch_hooks.py index 2475d4be..f8ff4c01 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_parallel_branch_hooks.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_parallel_branch_hooks.py @@ -75,6 +75,21 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: ) +class ParallelBranchPluginFactory: + """Builds one :class:`ParallelBranchPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> ParallelBranchPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return ParallelBranchPlugin() + + def branch0(_ctx: DurableContext) -> str: return "task-1" @@ -83,7 +98,7 @@ def branch1(_ctx: DurableContext) -> str: return "task-2" -@durable_execution(plugins=[ParallelBranchPlugin()]) +@durable_execution(plugins=[ParallelBranchPluginFactory()]) def handler(_event: Any, context: DurableContext) -> list: result = context.parallel( [branch0, branch1], diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_replay_flags.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_replay_flags.py index 12a29056..e7e17372 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_replay_flags.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_replay_flags.py @@ -75,6 +75,21 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) +class ReplayFlagPluginFactory: + """Builds one :class:`ReplayFlagPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> ReplayFlagPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return ReplayFlagPlugin() + + @durable_step def step_a(_step_context: StepContext) -> str: return "a" @@ -90,7 +105,7 @@ def step_b(step_context: StepContext) -> str: return "Operation succeeded" -@durable_execution(plugins=[ReplayFlagPlugin()]) +@durable_execution(plugins=[ReplayFlagPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: context.step(step_a()) retry_config = RetryStrategyConfig( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_retry_exhaustion.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_retry_exhaustion.py index e97d4528..877b16af 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_retry_exhaustion.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_retry_exhaustion.py @@ -93,13 +93,28 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) +class RetryExhaustionPluginFactory: + """Builds one :class:`RetryExhaustionPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> RetryExhaustionPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return RetryExhaustionPlugin() + + @durable_step def always_fail(_step_context: StepContext) -> str: msg = "boom" raise RuntimeError(msg) -@durable_execution(plugins=[RetryExhaustionPlugin()]) +@durable_execution(plugins=[RetryExhaustionPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: retry_config = RetryStrategyConfig( max_attempts=2, diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_suspension_invocation_end.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_suspension_invocation_end.py index 327763ba..50f80739 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_suspension_invocation_end.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_suspension_invocation_end.py @@ -57,7 +57,22 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) -@durable_execution(plugins=[SuspensionPlugin()]) +class SuspensionPluginFactory: + """Builds one :class:`SuspensionPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> SuspensionPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return SuspensionPlugin() + + +@durable_execution(plugins=[SuspensionPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return "Wait completed" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_failure.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_failure.py index 2939d09e..3f3cec44 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_failure.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_failure.py @@ -52,13 +52,28 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: ) +class TerminalFailurePluginFactory: + """Builds one :class:`TerminalFailurePlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> TerminalFailurePlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return TerminalFailurePlugin() + + @durable_step def failing_step(_step_context: StepContext) -> str: msg = "Something went wrong" raise RuntimeError(msg) -@durable_execution(plugins=[TerminalFailurePlugin()]) +@durable_execution(plugins=[TerminalFailurePluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: result: str = context.step( failing_step(), diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_payloads.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_payloads.py index e54049c1..ff3af95d 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_payloads.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_terminal_payloads.py @@ -62,6 +62,21 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) +class TerminalPayloadPluginFactory: + """Builds one :class:`TerminalPayloadPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> TerminalPayloadPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return TerminalPayloadPlugin() + + @durable_step def step_a(_step_context: StepContext) -> str: return "task-a" @@ -73,7 +88,7 @@ def step_b(_step_context: StepContext) -> str: raise RuntimeError(msg) -@durable_execution(plugins=[TerminalPayloadPlugin()]) +@durable_execution(plugins=[TerminalPayloadPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: context.step(step_a()) result: str = context.step( diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_operation_hooks.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_operation_hooks.py index df7e79fa..c8e8e1b3 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_operation_hooks.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_operation_hooks.py @@ -68,7 +68,22 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) -@durable_execution(plugins=[WaitOperationPlugin()]) +class WaitOperationPluginFactory: + """Builds one :class:`WaitOperationPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> WaitOperationPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return WaitOperationPlugin() + + +@durable_execution(plugins=[WaitOperationPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(2)) return "Wait completed" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py index eef6d3ee..59deac3a 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py @@ -81,6 +81,21 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) +class WaitReplayFlagPluginFactory: + """Builds one :class:`WaitReplayFlagPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> WaitReplayFlagPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return WaitReplayFlagPlugin() + + def wait_short(ctx: DurableContext) -> str: ctx.wait(Duration.from_seconds(2), name="short") return "short-done" @@ -91,7 +106,7 @@ def wait_long(ctx: DurableContext) -> str: return "long-done" -@durable_execution(plugins=[WaitReplayFlagPlugin()]) +@durable_execution(plugins=[WaitReplayFlagPluginFactory()]) def handler(_event: Any, context: DurableContext) -> list: result = context.parallel( [wait_short, wait_long], diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/pyproject.toml b/packages/aws-durable-execution-sdk-python-conformance-tests/pyproject.toml index 07355556..a8901c90 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/pyproject.toml @@ -8,7 +8,7 @@ version = "0.0.0" description = "Cross-SDK conformance test handlers for the AWS Durable Execution SDK for Python, exercised by the aws-durable-execution-conformance-tests runner." requires-python = ">=3.11" dependencies = [ - "aws-durable-execution-sdk-python==2.0.1", + "aws-durable-execution-sdk-python==3.0.0", ] [tool.hatch.build.targets.wheel] diff --git a/packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py b/packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py index a727bf73..a94e1409 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py @@ -1,8 +1,9 @@ """Demonstrates OTel-enriched logging in a durable execution. -The InvocationOtelPlugin installs a logging filter on the root logger -(enrich_logger=True by default) when the plugin is constructed. The filter -stamps the active OpenTelemetry trace context (traceId, spanId, +InvocationOtelPluginFactory is the plugin factory the SDK registers; it builds +one InvocationOtelPlugin per invocation. Each plugin installs a logging filter +on the root logger (enrich_logger=True by default) when it is constructed. The +filter stamps the active OpenTelemetry trace context (traceId, spanId, otelTraceSampled) onto every log record that flows through the root handler. This includes logs emitted via context.logger / step_context.logger as well as direct logging.getLogger() calls and third-party library logs, so logs @@ -16,7 +17,7 @@ from typing import Any -from aws_durable_execution_sdk_python_otel import InvocationOtelPlugin +from aws_durable_execution_sdk_python_otel import InvocationOtelPluginFactory from aws_durable_execution_sdk_python import StepContext from aws_durable_execution_sdk_python.context import ( @@ -44,7 +45,7 @@ def greet_in_child(child_context: DurableContext, name: str) -> str: return result -@durable_execution(plugins=[InvocationOtelPlugin()]) +@durable_execution(plugins=[InvocationOtelPluginFactory()]) def handler(_event: Any, context: DurableContext) -> str: # Logged at the top level: enriched with the invocation span_id. context.logger.info("Workflow started") diff --git a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_otel.py b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_otel.py index 3d001d46..b61a6420 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_otel.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_otel.py @@ -2,7 +2,7 @@ from typing import Any -from aws_durable_execution_sdk_python_otel import InvocationOtelPlugin +from aws_durable_execution_sdk_python_otel import InvocationOtelPluginFactory from aws_durable_execution_sdk_python import StepContext from aws_durable_execution_sdk_python.config import Duration @@ -32,7 +32,7 @@ def add_numbers_in_child(child_context: DurableContext, a: int, b: int): return result -@durable_execution(plugins=[InvocationOtelPlugin()]) +@durable_execution(plugins=[InvocationOtelPluginFactory()]) def handler(_event: Any, context: DurableContext) -> int: result = 0 for i in range(3): diff --git a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py index d8858baa..903cef96 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py @@ -12,6 +12,7 @@ from aws_durable_execution_sdk_python.execution import durable_execution from aws_durable_execution_sdk_python.plugin import ( DurableInstrumentationPlugin, + InvocationStartInfo, ) @@ -37,6 +38,21 @@ def on_user_function_end(self, info) -> None: self.logger.info(f"User function ended: {info}") +class MyPluginFactory: + """Builds one :class:`MyPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> MyPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return MyPlugin() + + @durable_step def add_numbers(_step_context: StepContext, a: int, b: int) -> int: return a + b @@ -51,7 +67,7 @@ def add_numbers_in_child(child_context: DurableContext, a: int, b: int): return result -@durable_execution(plugins=[MyPlugin()]) +@durable_execution(plugins=[MyPluginFactory()]) def handler(_event: Any, context: DurableContext) -> int: result: int = context.run_in_child_context( add_numbers_in_child(6, 4), diff --git a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_wait_plugin.py b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_wait_plugin.py index fb13cc81..9851977c 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_wait_plugin.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_wait_plugin.py @@ -42,7 +42,22 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) -@durable_execution(plugins=[RecordingWaitPlugin()]) +class RecordingWaitPluginFactory: + """Builds one :class:`RecordingWaitPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, so registering one would + stop this handler from importing. This factory exists only to construct the + plugin. + """ + + def create_plugin(self, info: InvocationStartInfo) -> RecordingWaitPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return RecordingWaitPlugin() + + +@durable_execution(plugins=[RecordingWaitPluginFactory()]) def handler(_event: Any, context: DurableContext) -> dict[str, Any]: context.wait(Duration.from_seconds(1), name="plugin-wait") return { diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index 920e36a6..f57a8bbb 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -42,6 +42,13 @@ def handler(event, context): ... ``` +`workflow_insight()` returns a plugin *factory*, which is what the SDK's +`plugins` argument takes: the SDK calls its `create_plugin` once per invocation to +build that invocation's plugin instance. The factory holds the resolved +configuration and the exporters, so configuration is per handler while record +state is per invocation. Its type is `WorkflowInsightPluginFactory`, exported from +the package root for annotating a value you hold. + With no exporter configured, records are written to the function's own CloudWatch log group as single JSON lines (the `LambdaLogExporter` default), carrying the name-keyed `operationsByName` summary. The `S3Exporter` writes the @@ -215,11 +222,27 @@ Behavior is validated cross-SDK by the `insight` conformance suite (`aws-durable-execution-conformance-tests-insight`). > **Note (asynchronous export).** Export rendering, truncation, `export()`, and -> `flush()` run on one lazy background worker per plugin. Checkpoint hooks only -> replace the latest pending snapshot and wake the worker. Consecutive -> `on-change` snapshots may coalesce while an export is in flight. An invocation -> that emits a record drains the latest snapshot and flushes exporters before it -> returns; invocations that emit nothing do not start or flush the worker. +> `flush()` run on one lazy background worker per registered factory. Checkpoint +> hooks only replace the latest pending snapshot and wake the worker. Consecutive +> `on-change` snapshots may coalesce while an export is in flight. Every +> sampled-in invocation end drains the latest snapshot and flushes exporters +> before it returns, including an end that emitted no record: a buffering +> exporter therefore sees one flush per sampled-in invocation end, which is the +> cadence the JS and Java plugins have. Only a sampled-out execution neither +> exports nor flushes. + +> **Note (invocation-end latency under concurrency).** The drain an invocation +> end performs waits for every record any execution had pending when it was +> called, and one worker serializes all exports and all flushes, so every +> concurrently ending invocation is released together at the slowest one. The wait +> therefore grows with the number of executions the environment is running, not +> just with this execution's own work: measured with a 30 ms exporter, one +> execution ended in ~72 ms and 48 concurrent executions in ~1.8 s each. That is +> the deliberate trade against the alternative — releasing an end before its +> record reached the exporters, which is what silently lost terminal records +> before. It matters for an exporter that makes a network call per record: budget +> invocation-end time against the environment's concurrency, not against one +> execution. ## Requirements diff --git a/packages/aws-durable-execution-sdk-python-insight/pyproject.toml b/packages/aws-durable-execution-sdk-python-insight/pyproject.toml index 0691bc32..af270c89 100644 --- a/packages/aws-durable-execution-sdk-python-insight/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-insight/pyproject.toml @@ -21,9 +21,15 @@ classifiers = [ "Programming Language :: Python :: Implementation :: CPython", ] dependencies = [ - # >=2.0.0: first published release carrying the plugin invocation-hook fields + # >=3.0.0,<4: 3.0.0 is the first release whose `plugins` argument takes + # factories, and the ceiling is the same reasoning applied forwards -- the next + # core major that changes the plugin contract would install and then fail at + # handler initialization exactly as 2.x does below. + # `workflow_insight()` returns a factory, which core 2.x cannot call, so an + # install resolved against 2.x fails at handler initialization. 3.0.0 also + # carries the invocation-hook fields this plugin reads # (InvocationInfo.execution_input / InvocationEndInfo.execution_result). - "aws-durable-execution-sdk-python>=2.0.0", + "aws-durable-execution-sdk-python>=3.0.0,<4", ] [project.optional-dependencies] diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__about__.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__about__.py index c7c5adad..4faf00f1 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__about__.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__about__.py @@ -1,4 +1,4 @@ # SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. # # SPDX-License-Identifier: Apache-2.0 -__version__ = "0.0.1" +__version__ = "0.1.0" diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__init__.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__init__.py index 488c2ada..3cf30c9a 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__init__.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/__init__.py @@ -33,6 +33,7 @@ ) from aws_durable_execution_sdk_python_insight.plugin import ( WorkflowInsightPlugin, + WorkflowInsightPluginFactory, workflow_insight, ) from aws_durable_execution_sdk_python_insight.truncation import truncate_record @@ -77,6 +78,7 @@ "SQSExporter", "WorkflowInsightConfig", "WorkflowInsightPlugin", + "WorkflowInsightPluginFactory", "apply_operations_format", "build_operations_by_name", "truncate_record", diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py index ef22ea54..bea39fcb 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -1,7 +1,22 @@ # SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. # # SPDX-License-Identifier: Apache-2.0 -"""Latest-pending asynchronous export scheduling for Workflow Insight.""" +"""Per-execution latest-pending asynchronous export scheduling for Workflow Insight. + +One scheduler serves every execution its environment hosts -- it is owned by the +handler-lifetime plugin factory, because serializing export is a cross-execution +job -- and Lambda Managed Instances makes concurrent executions in one +environment routine. So the pending record and the bookkeeping that goes with it +live on the per-execution object the caller passes in, which is the caller's own +per-invocation plugin instance (:class:`_ExportState` is mixed into it): +coalescing happens only within a single execution and one execution's record can +never displace another's. The scheduler holds those objects; it has no notion of +an execution ARN and nothing to look up. + +Export itself stays strictly serialized -- one worker thread, one ``export()`` at +a time -- so exporters never see concurrent calls. Parallel export is a later +phase and a contract change. +""" from __future__ import annotations @@ -16,54 +31,279 @@ _logger = logging.getLogger("aws_durable_execution_sdk_python_insight") +# Consecutive export-worker deaths tolerated before asynchronous export is +# disabled for good. +# +# A replacement worker is started by whoever is waiting, so a worker that dies on +# every attempt is retried as fast as threads can be created, and each retry +# leaves the waiting drain -- an invocation thread -- exactly where it was. The +# bound converts that unbounded retry into a bounded one. +# +# The bound is not 1, because a single death can come from a transient condition +# that the next attempt would not hit, and disabling instrumentation for the rest +# of the environment's life on one transient is too coarse. A deterministic defect +# reproduces on every attempt, so a small constant separates the two cases. Every +# completed export attempt and every completed flush resets the count, so only +# deaths with no work completed in between accumulate. +_MAX_CONSECUTIVE_WORKER_FAULTS = 3 + + +class _ExportState: + """One execution's export bookkeeping, and its slot in the export queue. + + Mixed into the per-invocation plugin instance, so one object carries both an + execution's hook-facing state and its export bookkeeping. Those used to live + in three ARN-keyed structures -- the plugin's execution registry, the + scheduler's pending record and its per-execution lane -- three views of one + execution that had to agree about whether it still had work outstanding. In + Java the same shape produced a defect where two of the views disagreed. There + is one view now, and the scheduler holds the object itself. + + Every field here is guarded by ``_ExportScheduler._condition``. They belong to + the scheduler: nothing outside it reads or writes them, and it never touches + the hook-facing state on the same object. + + Every field is underscore-prefixed, because the mixin's fields land on + :class:`WorkflowInsightPlugin`, which the package exports. A public name there + would advertise scheduler bookkeeping as part of the plugin's API. The + scheduler reads these names directly, which is why they are not name-mangled: + :class:`_ExportScheduler` is declared in this module, so this class's private + fields are within its own module's reach. + """ + + def __init__(self) -> None: + # Newest sequence number the scheduler assigned to this execution. + self._scheduled_seq = 0 + # This execution's latest record, waiting for the export worker, or None + # when nothing of its own is outstanding. A repeat emission replaces it, + # which is what per-execution coalescing means; the scheduler's queue + # holds this object exactly while this field is set. + self._pending_record: dict[str, Any] | None = None + # Newest sequence number already handed to every exporter. + self._exported_seq = 0 + # Value of the scheduler's export counter when that export finished, so + # a waiter can tell whether a completed flush covered its own record. + self._exported_at = 0 + # drain() calls currently parked on this execution. Nothing depends on + # it: a waiter holds this object directly, so the bookkeeping it waits + # on can no longer be reclaimed from under it. It is kept because it is + # the only way to observe that a drain really parked rather than raced + # past. + self._waiters = 0 + + +# What a caller has to release once it is back outside the lock: the records the +# `_disabled` latch dropped, each with the execution object that was carrying it. +# Both can run customer finalizers. +_Dropped = list[tuple[_ExportState, dict[str, Any] | None]] + + class _ExportScheduler: - """Run all exporters on one lazy worker with one latest pending record.""" + """Run all exporters on one lazy worker, keeping the latest record per execution.""" def __init__(self, exporters: list[InsightExporter]) -> None: self._exporters = exporters self._condition = threading.Condition(threading.Lock()) - self._pending: dict[str, Any] | None = None + # Executions with a record waiting, oldest arrival first -- an ordered + # set, keyed by the execution object itself. A repeat schedule for an + # execution replaces the record the object carries and keeps the object's + # position, so coalescing never lets one execution jump the queue. An + # execution is in here exactly while its `_pending_record` is set. + self._pending: dict[_ExportState, None] = {} + self._seq = 0 + self._export_count = 0 + # Highest export counter value covered by a completed flush. + self._flushed_through = 0 + # Completed flushes, monotonic. Export coverage alone cannot express + # "a flush ran for this invocation end": a drain with nothing of its own + # to export -- an invocation end that emitted no record -- is trivially + # covered by an older flush, so it would return without flushing at all. + # JS and Java flush once per sampled-in invocation end whether or not a + # record was emitted, so a drain also requires a flush that COMPLETED + # AFTER it was called. Concurrent drains still share one flush: they all + # entered before it completed. + self._flushes_completed = 0 self._flush_requested = False - self._flush_event: threading.Event | None = None + # Export counter coverage of the flush the worker is running right now, or + # None when no flush is in flight. Presence and coverage are separate + # facts: a flush that covers zero exports is an ordinary flush -- it is + # what an invocation that emitted no record asks for -- and a single + # integer cannot say both "no flush is running" and "a flush covering + # nothing is running". Encoding the first as 0 made those two states + # identical, so a waiter needing zero coverage could not tell that its + # flush was already running and requested a second one that then ran + # after its invocation had returned. Published when the worker commits to + # a flush, so a waiter woken while that flush runs -- before its coverage + # reaches _flushed_through -- can tell it is already covered. + self._flush_in_flight: int | None = None + # Value of the global schedule counter (_seq) when a flush was requested. + # The worker defers the flush until no record scheduled at or before that + # point is still pending. That is deliberately wider than the requester's + # own record: a drain therefore also waits for records other executions + # had pending when it was called. It excludes records scheduled after the + # request, so a steady stream of other executions cannot starve a waiting + # drain. + self._flush_barrier = 0 self._worker: threading.Thread | None = None self._disabled = False + # Export workers that died without completing any work, counted since the + # last completed export attempt or flush. Only deaths accumulate here, so + # a worker that keeps making progress never approaches the bound. + self._worker_faults = 0 + + def schedule(self, execution: _ExportState, record: dict[str, Any]) -> list[Any]: + """Replace this execution's pending snapshot; never runs exporters inline. - def schedule(self, record: dict[str, Any]) -> None: - """Replace the pending snapshot and return without running exporters.""" + Returns the records this call displaced, for the caller to release once it + holds no lock. They are handed back rather than dropped here because + releasing one can run a customer ``__del__``, and this method is called + from inside the calling plugin's own lock. A finalizer that re-entered a + *different* execution's plugin would then block on that instance's lock + while holding this one, which deadlocks both invocations if the mirror + image happens on another thread at the same time. The plugin's hook frame + drops them when the outermost hook returns; see + ``WorkflowInsightPlugin._hook_frame``. + + A caller with no frame to defer to may drop the returned list + immediately: doing so is only unsafe while a plugin lock is held. + """ displaced: dict[str, Any] | None = None - failed_pending: dict[str, Any] | None = None + failed_pending: _Dropped | None = None start_error: Exception | None = None with self._condition: if self._disabled: - return - displaced = self._pending - self._pending = record + # The record is handed back rather than dropped here for the same + # reason a displaced one is: this runs inside the calling plugin's + # lock, and releasing the record can run a customer finalizer. + return [record] + self._seq += 1 + execution._scheduled_seq = self._seq + displaced = execution._pending_record + execution._pending_record = record + # Re-queuing an execution that is already queued is a no-op that + # keeps its arrival position. + self._pending[execution] = None failed_pending, start_error = self._ensure_worker_locked() - self._condition.notify() - # Releasing either record may run custom finalizers, so do it unlocked. - del displaced, failed_pending + self._condition.notify_all() if start_error is not None: _logger.warning( "workflow-insight: could not start export worker; disabling " "asynchronous export: %s", start_error, ) + return [item for item in (displaced, failed_pending) if item is not None] + + def drain(self, execution: _ExportState) -> None: + """Wait until this execution's latest record is exported and exporters flush. + + Returns once the calling execution's own record has reached every exporter + and a flush covering it has completed; a flush triggered by another + execution never releases a waiter whose record is still pending. + + Every call waits for a flush that completed after the call started, so an + invocation end that emitted no record still flushes -- the cadence JS and + Java have. Concurrent calls can share one flush, since they all started + before it completed. - def drain(self) -> None: - """Wait until the latest pending record is exported and exporters flush.""" - failed_pending: dict[str, Any] | None = None + The caller passes the execution object rather than an ARN, so there is + nothing to look up and nothing to create: an execution that never + scheduled a record simply carries zeroed bookkeeping, which is exactly + "nothing of my own is outstanding, flush and return". + + Two paths return without exporting or flushing anything, because the + permanent ``_disabled`` latch means no record will ever be exported: the + latch was already set when this call started, or it is set while this call + is parked. Two things set that latch, and a drain that meets either + returns without a flush: failing to start the export worker, and an export + worker that has died ``_MAX_CONSECUTIVE_WORKER_FAULTS`` times without + completing any work. + + A third path returns immediately: a call made on the export worker thread + itself. Only that worker exports records and completes flushes, so a wait + there would park the one thread able to release it. That happens when an + exporter re-enters a plugin hook and the hook reaches an invocation end. + The call is refused and reported rather than deadlocking the invocation; + the record stays queued and this same worker exports it once it resumes + its loop, and a flush covering it is requested on the way out. (Mirrors + the Java ``ExportScheduler.refuseWaitThatWouldBlockThePump``.) + """ + if self._is_export_worker(): + _logger.warning( + "workflow-insight: drain() was called on the export worker " + "thread, the only thread able to serve it, so the call was " + "refused rather than deadlocking the invocation; an exporter " + "re-entered a plugin hook" + ) + # The refused call still leaves a flush behind, but only when there is + # something for it to cover. The hook that re-entered may have queued a + # record, and this worker exits its loop once nothing is pending and no + # flush is requested -- so without the request the record would be + # handed to the exporters and the worker would stop, leaving a + # buffering exporter holding an execution's terminal telemetry when + # Lambda freezes the environment. + # + # Requesting one unconditionally would livelock instead: an exporter + # whose flush() re-enters a hook arrives here from inside a flush, and + # an unconditional request would ask for the next one, which re-enters + # again, for as long as the environment lives. A pending record is what + # distinguishes new work from that loop. + self._request_flush_for_pending_records() + return + failed_pending: _Dropped | None = None start_error: Exception | None = None with self._condition: if self._disabled: return - if not self._flush_requested: - self._flush_requested = True - self._flush_event = threading.Event() - flush_event = self._flush_event - assert flush_event is not None - failed_pending, start_error = self._ensure_worker_locked() - started = not self._disabled - self._condition.notify() + execution._waiters += 1 + try: + want_seq = execution._scheduled_seq + # A drain always flushes, so require a flush that covers every + # export completed before this call as well as our own. + want_flush = self._export_count + # ...and one that completed after this call, so an invocation end + # that emitted nothing still flushes exactly once instead of + # riding on a flush that finished before it started. + want_flushes = self._flushes_completed + while not self._disabled: + # Export counter value a flush has to cover to release us: + # our own record's export plus everything already exported + # when this call started. Recomputed every pass, because + # execution._exported_at only becomes ours once our record is + # out. + need = max(execution._exported_at, want_flush) + if ( + execution._exported_seq >= want_seq + and self._flushed_through >= need + and self._flushes_completed > want_flushes + ): + break + # A flush already in flight whose coverage reaches `need` was + # committed after our record was handed to the exporters, so + # its completion releases us. Requesting another one here -- + # which is what a waiter woken inside that flush would do, + # since the coverage is not published yet and the request it + # made has already been consumed -- runs an extra flush after + # this drain, and the invocation, returned. + # + # `_flush_in_flight` is None exactly while no flush is + # running, so a flush that covers zero exports is still a + # flush in flight. That case is the common one, not a corner: + # `need` is 0 for a drain whose invocation emitted no record, + # and the flush it asks for covers 0 exports when nothing has + # ever been exported. Two such drains at once both see the + # other's flush and neither asks for a second. + in_flight = self._flush_in_flight + covered = in_flight is not None and in_flight >= need + if not self._flush_requested and not covered: + self._flush_requested = True + self._flush_barrier = max(self._flush_barrier, self._seq) + failed_pending, start_error = self._ensure_worker_locked() + if start_error is not None: + break + self._condition.notify_all() + self._condition.wait() + finally: + execution._waiters -= 1 del failed_pending if start_error is not None: _logger.warning( @@ -71,12 +311,76 @@ def drain(self) -> None: "asynchronous export: %s", start_error, ) - if started: - flush_event.wait() - def _ensure_worker_locked( - self, - ) -> tuple[dict[str, Any] | None, Exception | None]: + # -- internals ------------------------------------------------------------ + + def _is_export_worker(self) -> bool: + """Report whether the calling thread is this scheduler's export worker. + + Read under the condition's lock, because ``_worker`` is replaced by the + waiter that starts a replacement and cleared by a worker that exits. + """ + with self._condition: + return self._worker is threading.current_thread() + + def _request_flush_for_pending_records(self) -> None: + """Ask the worker for a flush, unless a flush already in flight covers it. + + Returns without waiting, so it is safe to call from the worker itself. The + barrier is raised to the current schedule counter, which is what makes the + flush cover a record queued moments ago rather than running before it. + + The condition separates the two ways a refused drain is reached, because + they need opposite answers. Re-entered from an exporter's ``export()``, + the record has already left ``_pending`` -- the worker takes it before it + calls the exporter -- so a pending-only test would skip the request and + leave that snapshot buffered until the environment froze. Re-entered from + an exporter's ``flush()``, a flush is in flight and covers what was + exported before it, so requesting another would produce a flush that + re-enters, requests, and flushes again for as long as the environment + lived -- after the invocation returned. A flush in flight with nothing + newly pending is therefore the one case that asks for nothing. + """ + with self._condition: + if self._disabled: + return + if not self._pending and self._flush_in_flight is not None: + return + self._flush_requested = True + self._flush_barrier = max(self._flush_barrier, self._seq) + self._condition.notify_all() + + def _disable_locked(self) -> _Dropped: + """Latch asynchronous export off for good and surrender everything queued. + + The latch is permanent, so no record the scheduler still holds will ever + be exported: keeping any of them would pin customer objects for the + remaining life of the environment. Empty the queue, which is the + scheduler's only per-execution structure, and hand what came out back to + the CALLER to release once it is outside the lock -- a record can carry + customer objects whose finalizers run arbitrary code, and so can an + execution whose hook state the plugin has already discarded. + + Every parked waiter is woken, because the latch means the export and the + flush it is waiting for are never going to happen. + + Callers hold ``self._condition``. + """ + self._disabled = True + self._worker = None + dropped = [ + (execution, execution._pending_record) for execution in self._pending + ] + for execution, _ in dropped: + execution._pending_record = None + self._pending = {} + self._flush_requested = False + self._flush_barrier = 0 + self._flush_in_flight = None + self._condition.notify_all() + return dropped + + def _ensure_worker_locked(self) -> tuple[_Dropped | None, Exception | None]: if self._worker is not None and self._worker.is_alive(): return None, None worker = threading.Thread( @@ -88,45 +392,185 @@ def _ensure_worker_locked( try: worker.start() except Exception as exc: # noqa: BLE001 - instrumentation must not escape hooks - self._disabled = True - self._worker = None - failed_pending = self._pending - self._pending = None - failed_event = self._flush_event - self._flush_event = None - self._flush_requested = False - if failed_event is not None: - failed_event.set() + failed_pending = self._disable_locked() return failed_pending, exc return None, None + def _blocking_pending_locked(self) -> bool: + """True while a record scheduled at or before the flush barrier is pending.""" + barrier = self._flush_barrier + return any(execution._scheduled_seq <= barrier for execution in self._pending) + def _run(self) -> None: + # The worker slot must be empty whenever no worker is running, or + # _ensure_worker_locked() never starts a replacement and every later + # record sits pending forever. The loop's own exits clear it, but an + # exception that unwinds out of the loop passes them by, and a thread that + # is unwinding still reports is_alive(), so the slot would stay occupied + # by a dead thread. Vacate it here, on every exit path, and wake anyone + # parked so they can ask for the replacement. + # + # A replacement alone is not enough when the death repeats. The waiter + # that starts the replacement runs the same work again, so a fault the + # work reproduces every time is retried as fast as threads can be + # created, and the drain that keeps starting them never returns: the + # invocation hangs and the environment fills with dead threads. + # _export() and _flush() contain everything a customer exporter can + # raise, so a fault reaching here comes from the scheduler's own code or + # from a failure-reporting call that a customer object subverted, and + # neither is something a retry can be expected to clear. Count + # consecutive faults and give up on asynchronous export at the bound. + faulted = True + dropped: _Dropped | None = None + gave_up = False + try: + self._run_loop() + faulted = False + finally: + with self._condition: + if self._worker is threading.current_thread(): + self._worker = None + if faulted: + self._worker_faults += 1 + if self._worker_faults >= _MAX_CONSECUTIVE_WORKER_FAULTS: + # Releasing the waiters matters more than delivering the + # records. A waiter is an invocation thread inside + # on_invocation_end, so leaving it parked turns an + # instrumentation defect into a stalled customer + # execution; dropping records loses instrumentation data + # only. The drop is reported below, so the scheduler + # never claims delivery it did not make. + dropped = self._disable_locked() + gave_up = True + self._condition.notify_all() + # A dropped record can run customer finalizers, so release it outside + # the lock. The exception that brought us here keeps unwinding once + # this block finishes, into the thread's traceback, with nothing + # swallowed. + del dropped + if gave_up: + _logger.warning( + "workflow-insight: export worker died %d times without " + "completing any work; disabling asynchronous export and " + "dropping every record still queued", + self._worker_faults, + ) + + def _run_loop(self) -> None: while True: + execution: _ExportState | None = None + seq = 0 record: dict[str, Any] | None = None - flush_event: threading.Event | None = None + flush_covers = 0 with self._condition: - while self._pending is None and not self._flush_requested: + while True: + if self._flush_requested and not self._blocking_pending_locked(): + self._flush_requested = False + self._flush_barrier = 0 + flush_covers = self._export_count + # Publish what this flush will cover before releasing the + # lock, so a waiter that wakes while it runs can see that + # this flush releases it and skip asking for another. A + # coverage of 0 is published like any other: it means this + # flush covers every export so far, of which there are + # none, and it is still a flush in flight. + self._flush_in_flight = flush_covers + break + if self._pending: + execution = next(iter(self._pending)) + del self._pending[execution] + record = execution._pending_record + execution._pending_record = None + seq = execution._scheduled_seq + break self._condition.wait() - if self._pending is not None: - record = self._pending - self._pending = None - else: - flush_event = self._flush_event - self._flush_event = None - self._flush_requested = False if record is not None: - self._export(record) + assert execution is not None + # Taking the record consumed this execution's pending slot, so + # nothing will ever export that snapshot again. The bookkeeping + # must therefore advance whatever export() did: skip it and + # execution._exported_seq never reaches a waiter's want_seq, so a + # drain parked on this execution is never released. Count the + # attempt in a finally so that holds even if _export() raises. + # + # Advancing here means the record was OFFERED to every exporter, + # not that every exporter accepted it. _export() reports each + # exporter's own failure and moves to the next, so no exporter is + # skipped because another one failed, and coverage never stands + # for a delivery that was never attempted. + # + # The record and its bookkeeping are one object, so there is no + # second lookup left to come back empty: publishing cannot miss. + exported = False + try: + self._export(record) + exported = True + finally: + # Release the exported record before re-locking: a custom + # finalizer may re-enter schedule(). + del record + with self._condition: + self._export_count += 1 + if seq > execution._exported_seq: + execution._exported_seq = seq + execution._exported_at = self._export_count + if exported: + # This worker completed work, so any earlier worker + # death was not the start of a fault the work + # reproduces every time. + self._worker_faults = 0 + self._condition.notify_all() continue - self._flush() - if flush_event is not None: - flush_event.set() + flushed = False + try: + self._flush() + flushed = True + finally: + with self._condition: + # Retire the marker whatever happened: a stale one would park + # every later waiter that trusted this flush to cover it. Only + # a flush that ran to completion publishes its coverage. + self._flush_in_flight = None + if flushed: + self._flushes_completed += 1 + if flush_covers > self._flushed_through: + self._flushed_through = flush_covers + self._worker_faults = 0 + self._condition.notify_all() with self._condition: - if self._pending is None and not self._flush_requested: + if not self._pending and not self._flush_requested: self._worker = None return + # Isolation of one exporter's failure from the others is what both loops below + # exist for, and it holds for every exception type a customer exporter can + # raise, BaseException included. + # + # A BaseException is normally not caught, because KeyboardInterrupt, + # SystemExit and asyncio.CancelledError each mean that the operation the + # current thread is running must stop. None of those three can arrive here + # that way. This is the export worker thread: the interpreter raises + # KeyboardInterrupt only in the main thread, threading discards a SystemExit + # raised in a worker thread, and nothing cancels this thread because nothing + # outside the scheduler knows it exists. A BaseException seen at these call + # sites was therefore raised by the exporter itself, which makes it a report + # of a defective exporter rather than an instruction to this thread. An + # exporter that touches asyncio can raise CancelledError without writing + # `raise`, so the case is reachable without a customer intending it. + # + # Containing it here is what keeps the two guarantees the worker owes. Every + # remaining exporter still receives the record, so one exporter cannot make + # the others miss a snapshot that is then discarded. And the flush the waiting + # drain asked for still completes, so the drain is released by this worker + # instead of by a replacement that runs the same failing exporter and dies the + # same way. + # + # The containment is confined to these two call sites. Nowhere else does the + # scheduler catch a BaseException, and neither loop runs on an invocation + # thread. + def _export(self, record: dict[str, Any]) -> None: for exporter in self._exporters: try: @@ -134,7 +578,7 @@ def _export(self, record: dict[str, Any]) -> None: record, exporter.max_record_size_bytes, exporter.render ) exporter.export(shaped) - except Exception as exc: # noqa: BLE001 - one exporter must not break others + except BaseException as exc: # noqa: BLE001 - one exporter must not break others _logger.warning( "workflow-insight: exporter %s failed: %s", type(exporter).__name__, @@ -145,7 +589,7 @@ def _flush(self) -> None: for exporter in self._exporters: try: exporter.flush() - except Exception as exc: # noqa: BLE001 - one exporter must not break others + except BaseException as exc: # noqa: BLE001 - one exporter must not break others _logger.warning( "workflow-insight: exporter %s flush failed: %s", type(exporter).__name__, @@ -159,4 +603,4 @@ def _worker_alive(self) -> bool: def _pending_count(self) -> int: with self._condition: - return int(self._pending is not None) + return len(self._pending) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index f19796a0..d56bb5d4 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py @@ -9,8 +9,20 @@ record keeps the JS camelCase field names so records read identically across SDKs. +Two lifetimes: + ``workflow_insight(config)`` returns a FACTORY, which is what + ``@durable_execution(plugins=[...])`` takes. The factory lives as long as the + handler and owns everything that is not per-execution: the resolved immutable + config, the exporters, and the ``_ExportScheduler`` that serializes export + across executions. The SDK calls its ``create_plugin`` once per invocation and + drops the instance it returns when the invocation scope exits, so a + :class:`WorkflowInsightPlugin` instance serves exactly one invocation of one + execution. Everything this environment holds for that execution is therefore + ordinary instance state: no ARN-keyed registry, and no hook can reach an + instance other than its own. + Operation-map sourcing: - The Python SDK invocation hooks now carry the full operation map directly: + The Python SDK invocation hooks carry the full operation map directly: ``InvocationStartInfo.operations`` (a point-in-time snapshot at invocation start), ``InvocationEndInfo.operations`` (a fresh snapshot at invocation end), and ``OperationChangeInfo.operations`` (the full map at the change). Alongside @@ -19,8 +31,8 @@ snapshots as the authoritative operation state -- it does NOT reconstruct the map by accumulating per-operation ``on_operation_end`` events. Because every invocation start re-seeds the map from the snapshot, a cold resume in a fresh - Lambda environment (a brand-new plugin instance) still reports the prior - terminal operations. + Lambda environment (a brand-new instance, as every invocation now gets) still + reports the prior terminal operations. The Python SDK has no ``pluginsConfig.childOperationsDepth`` equivalent, so ``full-tree`` records rely on the child operations being present in the @@ -30,10 +42,12 @@ from __future__ import annotations +import contextlib import datetime import json import math import threading +from collections.abc import Iterator from typing import Any, Callable from aws_durable_execution_sdk_python.plugin import ( @@ -46,7 +60,10 @@ OperationType, ) -from aws_durable_execution_sdk_python_insight._export_scheduler import _ExportScheduler +from aws_durable_execution_sdk_python_insight._export_scheduler import ( + _ExportScheduler, + _ExportState, +) from aws_durable_execution_sdk_python_insight.exporters.lambda_log_exporter import ( LambdaLogExporter, ) @@ -159,191 +176,326 @@ def _apply_result_override( return None -class _ExecutionState: - __slots__ = ("start_time", "parsed_arn", "cached_input", "operations") +class _HookFrames(threading.local): + """Nested plugin hook frames on one thread, and the drains they owe. - def __init__(self, start_time: Any, parsed_arn: dict[str, str]) -> None: - self.start_time = start_time - self.parsed_arn = parsed_arn - self.cached_input: Any = None - # operation_id -> OperationInfo, adopted verbatim from the SDK's - # authoritative snapshot (invocation start/end and operation-change). - self.operations: dict[str, OperationInfo] = {} + Per thread, and shared by every plugin instance on that thread. A hook runs + customer code while holding a plugin's ``_lock``, and that code can call a + hook on *any* live instance, so the frame that must run a deferred drain is + the outermost one on the thread whatever instance it belongs to. + ``threading.local`` runs ``__init__`` once per thread, so each thread gets its + own counter and its own list. + """ -class WorkflowInsightPlugin(DurableInstrumentationPlugin): - def __init__(self, config: WorkflowInsightConfig) -> None: - self._sampling_rate = _resolve_sampling_rate(config.sampling_rate) - # config.emit_mode / operation_detail are already normalized to enum - # members (or None) by WorkflowInsightConfig.__post_init__; re-wrap to - # satisfy the static type of the union-typed config fields. - self._emit_mode: EmitMode = ( - EmitMode(config.emit_mode) - if config.emit_mode is not None - else EmitMode.ON_COMPLETE - ) - detail = ( - OperationDetail(config.operation_detail) - if config.operation_detail is not None - else OperationDetail.TOP_LEVEL + def __init__(self) -> None: + self.depth = 0 + self.pending: list[WorkflowInsightPlugin] = [] + # Records displaced from the scheduler's pending slots, held until every + # plugin lock on this thread is released. Releasing one can run a customer + # finalizer, and a finalizer that reaches another execution's plugin must + # not do so while this thread holds a plugin lock. + self.releases: list[Any] = [] + + +_hook_frames = _HookFrames() + + +class WorkflowInsightPlugin(DurableInstrumentationPlugin, _ExportState): + """Everything this environment holds for one invocation of one execution. + + Built by the factory ``workflow_insight()`` returns, once per invocation, + from that invocation's ``InvocationStartInfo`` -- the same object its + ``on_invocation_start`` then receives. Identity comes from that info and is + never revised afterwards: the execution ARN, the sampling decision, the + execution start time and the cached input. + + Per-execution state is plain instance state, and the object doubles as its + own export queue entry (via :class:`_ExportState`). Where three ARN-keyed + structures had to agree about one execution -- this plugin's registry, the + scheduler's pending record and its per-execution lane -- there is now one + object and no ARN to resolve. In Java the same three-view shape produced a + defect where two of the views disagreed. + + Two locks, disjoint field sets, and neither is ever taken to reach the + other's fields: + + * ``_lock`` guards ``_closed``, ``_build_revision``, the ``_operations`` + rebind and record emission. + * ``_ExportScheduler._condition``'s lock guards the export bookkeeping this + instance carries for the scheduler; nothing outside the scheduler reads or + writes those fields, and the scheduler never touches the fields above. + + Shared, handler-lifetime state -- resolved config, exporters, scheduler -- + lives on ``_shared`` and is read-only from here. + """ + + def __init__( + self, shared: WorkflowInsightPluginFactory, info: InvocationStartInfo + ) -> None: + _ExportState.__init__(self) + self._shared = shared + execution_arn = info.execution_arn or "" + self._execution_arn = execution_arn + # Deterministic per-ARN, so the decision could be recomputed on every + # hook; taken once here because the instance now has a place to keep it, + # and because an unsampled instance then never parses the ARN. + self._sampled_in = bool(execution_arn) and _should_sample( + execution_arn, shared._sampling_rate ) - self._top_level_only = detail != OperationDetail.FULL_TREE - content: ContentConfig | None = config.content - self._content = content - ops = content.operations if content and content.operations else None - self._include_errors = ( - True if ops is None or ops.include_errors is None else ops.include_errors + self._parsed_arn: dict[str, str] = ( + _parse_execution_arn(execution_arn) if self._sampled_in else {} ) - self._overrides_by_name: dict[str, OperationOverride] = {} - if ops is not None: - for override in ops.overrides: - self._overrides_by_name[override.operation_name] = override - # Default-exporter parity with the JS plugin: an omitted OR an explicitly - # empty exporter list falls back to the Lambda log exporter, so the - # plugin is never a silent no-op. A non-empty list is used verbatim. - self._exporters: list[InsightExporter] = ( - list(config.exporters) if config.exporters else [LambdaLogExporter()] + # Always the service-provided execution start time when present, + # including on a cold resume in a fresh environment -- never the resume + # time, which would corrupt duration and the date partition. `now` is the + # fallback for an info that carries no start time at all. + self._start_time: Any = ( + info.execution_start_time + if info.execution_start_time is not None + else datetime.datetime.now(datetime.UTC) ) - self._scheduler = _ExportScheduler(self._exporters) - self._state: dict[str, _ExecutionState] = {} - self._lock = threading.Lock() - - # -- sampling / state ----------------------------------------------------- - - def _sampled_in(self, execution_arn: str) -> bool: - # Deterministic per-ARN, so every hook for one execution agrees without - # needing to persist the decision in state. - return _should_sample(execution_arn, self._sampling_rate) - - def _ensure_state(self, execution_arn: str) -> _ExecutionState: - with self._lock: - state = self._state.get(execution_arn) - if state is None: - state = _ExecutionState( - start_time=datetime.datetime.now(datetime.UTC), - parsed_arn=_parse_execution_arn(execution_arn), - ) - self._state[execution_arn] = state - return state - - def _discard_state(self, execution_arn: str) -> None: - with self._lock: - self._state.pop(execution_arn, None) - - def _adopt_operations( - self, state: _ExecutionState, operations: dict[str, OperationInfo] - ) -> None: + self._cached_input: Any = info.execution_input + # operation_id -> OperationInfo, adopted verbatim from the SDK's + # authoritative snapshot (invocation start/end and operation-change). + self._operations: dict[str, OperationInfo] = {} + # Set once this invocation has ended. A hook that arrives afterwards (an + # operation-change for a checkpoint that completed just before the end) + # must emit nothing (mirrors the Java ExecutionState.closed flag). + self._closed = False + # Counts the record builds this instance has started. Never decremented. + # + # A record is a complete snapshot of one execution, so the scheduler's + # per-execution slot takes whichever record is handed over last and never + # compares ages. The build that hands its record over last is not the + # build that started last: `_emit` runs customer code (the input/output + # transforms and the operation result overrides) between the snapshot it + # takes and the hand-off, and that code can re-enter a hook on this thread + # and complete a newer build first. Without a comparison of build ages the + # outer frame's older snapshot then replaces the newer one in the slot, or + # is exported after it. Every non-terminal build takes the next value here + # before it starts, and `_emit` hands the record over only while that value + # is still the newest (mirrors the JS `buildRevision` and the Java + # `AtomicLong buildRevision`). + # + # A plain int, not an atomic: every read and every increment happens under + # `_lock`, so the read-modify-write cannot interleave, and the two builds + # this counter distinguishes are nested frames on one thread rather than + # two threads. Java needs an AtomicLong because its two builds really can + # run at once. + self._build_revision = 0 + # Guards `_closed`, `_build_revision`, the operations rebind and record + # emission, so a late hook can never slip a RUNNING record in after the + # terminal one. + # + # What it protects against is reentrancy, not two threads. The SDK + # dispatches every hook synchronously on the thread that produced the + # event, and it joins the checkpoint thread and the branch pools before + # the invocation-end hook is dispatched, so a checkpoint-path + # operation-change cannot overlap `on_invocation_end` -- an earlier + # version of this comment claimed it could. The lock still earns its place + # for the reason below, and it stays because a guard whose correctness + # rests on the SDK's join ordering is one refactor away from being wrong. + # + # Reentrant on purpose: `_emit` runs the scheduler's `schedule()` inside + # this hold, and `schedule()` releases the record it displaces, which can + # run a customer finalizer that re-enters a hook for this same execution + # on this same thread. A plain lock self-deadlocks the invocation thread + # there. (Java holds no such lock: its ExecutionState carries no + # operations map, and `cachedInput` is a bare volatile field.) + self._lock = threading.RLock() + + # -- state ---------------------------------------------------------------- + + def _adopt_operations_locked(self, operations: dict[str, OperationInfo]) -> None: # Adopt the authoritative point-in-time snapshot. Copy so plugin state # never aliases the SDK-owned map, and rebind the attribute so a # concurrent reader holding the prior reference iterates a stable dict. - with self._lock: - state.operations = dict(operations) + # Callers hold self._lock. + self._operations = dict(operations) # -- hooks ---------------------------------------------------------------- + @contextlib.contextmanager + def _hook_frame(self) -> Iterator[None]: + """Mark a hook frame on this thread and run the drains it owes on exit. + + A hook runs customer code -- the input/output transforms, a result + override, ``__del__`` on an object a displaced record carried -- while + holding ``_lock``, which is reentrant so that such code re-entering a hook + on this thread does not self-deadlock. A re-entrant + ``on_invocation_end`` therefore used to run its drain while the outer + frame still held ``_lock``. A drain waits for the export worker, and an + exporter that re-enters a hook blocks that worker on the very ``_lock`` + the waiting thread holds, so neither side can proceed and the invocation + hangs until Lambda times it out. + + A drain a nested frame asks for is therefore deferred to the outermost + frame, which runs it after every ``_lock`` hold on this thread has been + released. The unnested case is unchanged: the frame is the outermost one, + so its drain runs at the same point it always did. + + The frame state is per thread and shared by every instance, because + customer code inside one execution's build can call a hook on another + execution's instance, and a drain deferred to the outer frame of a + *different* instance is still a drain outside every lock. + """ + state = _hook_frames + state.depth += 1 + try: + yield + finally: + state.depth -= 1 + if state.depth == 0: + # Released before the drains, and with no lock held: a finalizer + # that schedules a record is then covered by the drain that + # follows it. + if state.releases: + state.releases.clear() + if state.pending: + owed, state.pending = state.pending, [] + for plugin in owed: + plugin._drain() + + def _request_drain(self) -> None: + """Ask for a drain once the outermost hook frame on this thread unwinds.""" + pending = _hook_frames.pending + if not any(plugin is self for plugin in pending): + pending.append(self) + + def _drain(self) -> None: + self._shared._scheduler.drain(self) + def on_invocation_start(self, info: InvocationStartInfo) -> None: - arn = info.execution_arn - if not arn or not self._sampled_in(arn): + if not self._sampled_in: return - state = self._ensure_state(arn) - # Always adopt the service-provided execution start time when present, - # including a cold resume in a fresh environment (never the resume time, - # which would corrupt duration and the date partition). - if info.execution_start_time is not None: - state.start_time = info.execution_start_time - state.cached_input = info.execution_input - # Seed the operation map from the full snapshot on every invocation. On a - # cold resume this rebuilds prior (terminal) operations that a fresh - # plugin instance never saw via per-operation hooks. - self._adopt_operations(state, info.operations) - if self._emit_mode == EmitMode.ON_CHANGE: - self._emit( - arn, - state, - status="RUNNING", - end_time=None, - output_raw=None, - error=None, - ) + with self._hook_frame(), self._lock: + if self._closed: + return + # Seed the operation map from the full snapshot. On a cold resume + # this rebuilds prior (terminal) operations that a fresh instance + # never saw via per-operation hooks. + self._adopt_operations_locked(info.operations) + if self._shared._emit_mode == EmitMode.ON_CHANGE: + self._emit(status="RUNNING", end_time=None, output_raw=None, error=None) def on_operation_change(self, info: OperationChangeInfo) -> None: - arn = info.execution_arn - if not arn or not self._sampled_in(arn): + # No ARN to resolve: this instance belongs to the invocation the change + # was raised in, so the hook's `execution_arn` is this execution's by + # construction. A change hook can no longer fabricate state for an + # execution whose invocation has ended -- there is no registry to + # fabricate it in, and the instance it reaches is its own. + if not self._sampled_in: return - state = self._ensure_state(arn) - # Replace state with the full operations snapshot carried by the hook. - self._adopt_operations(state, info.operations) - # on-change mode exports an updated RUNNING record on each change so - # mid-invocation progress is observable, not only at start/end. - if self._emit_mode == EmitMode.ON_CHANGE: - self._emit( - arn, - state, - status="RUNNING", - end_time=None, - output_raw=None, - error=None, - ) + with self._hook_frame(), self._lock: + if self._closed: + return + # Replace state with the full operations snapshot carried by the hook. + self._adopt_operations_locked(info.operations) + # on-change mode exports an updated RUNNING record on each change so + # mid-invocation progress is observable, not only at start/end. + if self._shared._emit_mode == EmitMode.ON_CHANGE: + self._emit(status="RUNNING", end_time=None, output_raw=None, error=None) def on_invocation_end(self, info: InvocationEndInfo) -> None: - arn = info.execution_arn - if not arn: - return - if not self._sampled_in(arn): - # Sampled-out executions process no operations and retain no state. - self._discard_state(arn) + if not self._sampled_in: + # Sampled-out executions process nothing: they neither export nor + # flush, so instrumenting a fraction of executions costs the rest + # nothing. return - state = self._ensure_state(arn) - # Refresh from the fresh end-of-invocation snapshot before emitting so - # the terminal record reflects the final operation map. - self._adopt_operations(state, info.operations) - status = _STATUS_MAP.get(info.status, "RUNNING") - is_terminal = status in ("SUCCEEDED", "FAILED") - is_failure = status == "FAILED" - - if self._emit_mode == EmitMode.ON_CHANGE: - should_emit = True - elif self._emit_mode == EmitMode.ON_FAILURE: - should_emit = is_failure - else: # on-complete - should_emit = is_terminal - - if should_emit: - # Only terminal (SUCCEEDED/FAILED) records carry an end time; a - # PENDING/RETRY invocation end maps to RUNNING (still in flight) and - # must omit endTime/durationMs. Passing end_time=None makes _emit - # drop both fields. Output and error likewise belong only to a - # terminal record. - self._emit( - arn, - state, - status=status, - end_time=datetime.datetime.now(datetime.UTC) if is_terminal else None, - output_raw=info.execution_result if is_terminal else None, - error=info.error if is_terminal else None, - ) - self._scheduler.drain() - - # Clear state after EVERY invocation end, including PENDING/RETRY. The - # next invocation rebuilds it from InvocationStartInfo.operations, so a - # suspended execution that never resumes in this environment (or that was - # sampled out) leaks nothing and state stays bounded. - self._discard_state(arn) + emit_mode = self._shared._emit_mode + with self._hook_frame(): + # The drain is registered before anything can fail, not after the + # record is built. `_emit` runs customer code -- the content and + # result transforms, and `__del__` on an object a displaced record + # carried -- and a failure there leaves this hook by way of the SDK's + # containment. Registering afterwards meant such a failure skipped + # the drain, so records this execution had already scheduled stayed + # in a buffering exporter when the environment froze. Registering + # here costs nothing when the hook succeeds: the frame runs the drain + # once, on the way out, either way. + self._request_drain() + with self._lock: + if not self._closed: + # Close the gate before emitting so a concurrent late hook for + # this execution cannot append a RUNNING record after the + # terminal one. + self._closed = True + # Refresh from the fresh end-of-invocation snapshot before + # emitting so the terminal record reflects the final operation + # map. + self._adopt_operations_locked(info.operations) + status = _STATUS_MAP.get(info.status, "RUNNING") + is_terminal = status in ("SUCCEEDED", "FAILED") + is_failure = status == "FAILED" + + if emit_mode == EmitMode.ON_CHANGE: + should_emit = True + elif emit_mode == EmitMode.ON_FAILURE: + should_emit = is_failure + else: # on-complete + should_emit = is_terminal + + if should_emit: + # Only terminal (SUCCEEDED/FAILED) records carry an end time; + # a PENDING/RETRY invocation end maps to RUNNING (still in + # flight) and must omit endTime/durationMs. Passing + # end_time=None makes _emit drop both fields. Output and + # error likewise belong only to a terminal record. + self._emit( + status=status, + end_time=datetime.datetime.now(datetime.UTC) + if is_terminal + else None, + output_raw=info.execution_result if is_terminal else None, + error=info.error if is_terminal else None, + # This is the emit that closed the gate, so it always runs + # with `_closed` already set and must never drop itself. + closing=True, + ) + + # Nothing has to be cleared after an invocation end, including a + # PENDING/RETRY one: this instance IS the state, and the SDK drops it + # when the invocation scope exits. A suspended execution that resumes + # here later gets a fresh instance, seeded from + # InvocationStartInfo.operations. + # + # Drain on EVERY sampled-in invocation end, emitted record or not: JS and + # Java flush once per sampled-in invocation end regardless, and a + # buffering exporter has to see the same rhythm in all three languages + # (an on-failure/on-complete mode that emits nothing for this invocation + # may still be holding records another execution handed it). A sampled-out + # execution returns above, so it neither exports nor flushes. + # + # The drain covers this execution only -- it names this instance, which + # carries its own export bookkeeping: it returns once this execution's + # own record, if any, reached the exporters and a flush that completed + # after this call is done, without waiting on records scheduled after the + # call by other executions. + # + # Asked for rather than performed, so it runs when the outermost hook + # frame on this thread unwinds and every `_lock` hold is released. See + # `_hook_frame`: an invocation end that customer code re-entered from + # inside another hook's build would otherwise wait for the export + # worker while holding the lock that worker may need. The request + # itself is made at the top of this hook, so a failure in the build + # below cannot skip it. # -- emission ------------------------------------------------------------- def _build_operations( self, operations: dict[str, OperationInfo] ) -> list[dict[str, Any]]: + shared = self._shared records: list[dict[str, Any]] = [] for op in operations.values(): if op.operation_type == OperationType.EXECUTION: continue if not op.name: continue - if self._top_level_only and op.parent_id: + if shared._top_level_only and op.parent_id: continue - override = self._overrides_by_name.get(op.name) + override = shared._overrides_by_name.get(op.name) if override is not None and override.exclude: continue @@ -365,7 +517,7 @@ def _build_operations( entry["durationMs"] = dur if op.attempt is not None: entry["attempt"] = op.attempt - if self._include_errors and op.error is not None: + if shared._include_errors and op.error is not None: entry["error"] = {"name": op.error.type, "message": op.error.message} if override is not None and override.result is not None: value = _apply_result_override(override.result, op.result) @@ -376,29 +528,41 @@ def _build_operations( def _emit( self, - execution_arn: str, - state: _ExecutionState, *, status: str, end_time: Any, output_raw: str | None, error: Any, + closing: bool = False, ) -> None: - arn = state.parsed_arn - start_time = state.start_time + arn = self._parsed_arn + start_time = self._start_time duration = _duration_ms(start_time, end_time) # Snapshot the operations reference once so a concurrent adopt() rebind # cannot change the map mid-build. - operations = state.operations - - content = self._content + operations = self._operations + + # The revision is taken here, before the build, never after. Customer code + # runs inside the build below and can re-enter a hook on this thread, + # which starts and finishes a newer build. A value read after the build + # would already be that newer build's, so this older record would pass the + # check and replace the newer one. Callers hold self._lock, so the + # increment cannot interleave with another build's. + # + # The closing emit takes no revision; see the hand-off below. + revision = 0 + if not closing: + self._build_revision += 1 + revision = self._build_revision + + content = self._shared._content record: dict[str, Any] = { "recordType": "WorkflowInsight", "schemaVersion": "1.0", "emittedAt": datetime.datetime.now(datetime.UTC) .isoformat() .replace("+00:00", "Z"), - "executionArn": execution_arn, + "executionArn": self._execution_arn, } if arn.get("executionName"): record["executionName"] = arn["executionName"] @@ -423,7 +587,7 @@ def _emit( except (json.JSONDecodeError, TypeError): parsed_output = output_raw input_value = _apply_data_content( - state.cached_input, content.input if content else None + self._cached_input, content.input if content else None ) output_value = _apply_data_content( parsed_output, content.output if content else None @@ -436,9 +600,137 @@ def _emit( record["error"] = {"name": error.type, "message": error.message} record["operations"] = self._build_operations(operations) - self._scheduler.schedule(record) + # INVARIANT: no record for an execution reaches the scheduler after that + # execution's closing record -- the exporters never see a RUNNING record + # follow the terminal one for the same execution. + # + # Re-check the gate here, because each hook's own `if self._closed` is a + # check-then-act and this is the act. One instance per invocation removed + # one of the two windows this used to cover -- state can no longer be + # discarded and recreated underneath a hook, because there is no registry + # to recreate it in -- but not the other: everything between the two runs + # customer code while holding self._lock (the input/output transforms + # above, a result override in _build_operations, and __del__ on any object + # the record carries), and self._lock is a reentrant RLock on purpose (so + # that customer code re-entering a hook on this thread does not + # self-deadlock). Such a re-entrant call can therefore run + # on_invocation_end all the way through -- set `_closed`, emit the + # terminal record and drain -- and then return here. Without this + # re-check the outer frame hands its already-built RUNNING record to the + # scheduler afterwards, and one execution exports + # ['SUCCEEDED', 'RUNNING'] from a single hook call, no concurrency + # required. + # + # `closing` marks the emit that set the gate (on_invocation_end's own), + # which by construction always runs with `_closed` set and must not drop + # itself. It is not the same test as "the record is terminal": in + # on-change mode a PENDING/RETRY invocation end legitimately emits a + # RUNNING record, and that record is the closing one. + # + # INVARIANT: the record handed to the scheduler below is the newest build + # this instance has started -- an exporter never stores an older snapshot + # over a newer one for the same execution. + # + # A build that customer code started from inside this one has already + # handed its own, newer record over by the time control returns here, so + # this record is superseded. Dropping it loses nothing: a record is a + # complete snapshot of one execution, so the newer record carries + # everything this one carries. That is the same property that makes the + # scheduler's per-execution coalescing sound. + # + # The check and the hand-off are one critical section. self._lock is held + # across both -- schedule() is called inside this hold -- and every build + # takes that same lock, so no build can start, finish and queue its record + # between this check and the schedule() below. A record that passes + # therefore cannot be queued after the record that supersedes it. (Java + # revalidates inside the scheduler's monitor instead, because there the + # scheduler's monitor is what guards `closed` and the record slot; here + # both facts already belong to self._lock.) + # + # The closing record is exempt from the revision check. Customer code + # inside its build can start a newer non-terminal build, which would make a + # revision taken here stale, and a checked hand-off would then drop the + # closing record and leave a RUNNING snapshot as this execution's last + # exported state. The exemption cannot let a stale record win, because + # `_closed` was set in the same self._lock hold that queues this record and + # every later non-terminal record is rejected above. + if not closing and (self._closed or revision != self._build_revision): + return + # The records this hand-off displaces are released by the hook frame, not + # here: this runs inside `_lock`, and a displaced record can carry a + # customer object whose `__del__` re-enters a hook. Re-entering *this* + # execution's hook is safe because `_lock` is reentrant; re-entering + # another execution's is not, and two threads doing it to each other at + # once would hang both invocations. See `_hook_frame`. + _hook_frames.releases.extend(self._shared._scheduler.schedule(self, record)) + + +class WorkflowInsightPluginFactory: + """The handler-lifetime half of the plugin: what is NOT per-execution. + + Built by :func:`workflow_insight`, which is the supported way to obtain one. + The class is public because it is the declared return type of that function, + and a ``py.typed`` consumer must be able to name the type it holds. + + Satisfies the SDK's ``DurableInstrumentationPluginFactory`` -- its + ``create_plugin`` is called with an ``InvocationStartInfo`` and returns the + plugin instance for that invocation. Everything it holds is either immutable + after construction (the resolved config) or deliberately shared across + executions: + + * the exporters, which are customer objects registered once, and + * the ``_ExportScheduler``, because export serialization is cross-execution: + one worker, one ``export()`` at a time, whatever the instance that + scheduled the record. + + A class rather than a closure so the resolved config stays inspectable + (``factory._emit_mode``, ``factory._exporters``) instead of being buried in + cell variables. + """ + + def __init__(self, config: WorkflowInsightConfig) -> None: + self._sampling_rate = _resolve_sampling_rate(config.sampling_rate) + # config.emit_mode / operation_detail are already normalized to enum + # members (or None) by WorkflowInsightConfig.__post_init__; re-wrap to + # satisfy the static type of the union-typed config fields. + self._emit_mode: EmitMode = ( + EmitMode(config.emit_mode) + if config.emit_mode is not None + else EmitMode.ON_COMPLETE + ) + detail = ( + OperationDetail(config.operation_detail) + if config.operation_detail is not None + else OperationDetail.TOP_LEVEL + ) + self._top_level_only = detail != OperationDetail.FULL_TREE + content: ContentConfig | None = config.content + self._content = content + ops = content.operations if content and content.operations else None + self._include_errors = ( + True if ops is None or ops.include_errors is None else ops.include_errors + ) + self._overrides_by_name: dict[str, OperationOverride] = {} + if ops is not None: + for override in ops.overrides: + self._overrides_by_name[override.operation_name] = override + # Default-exporter parity with the JS plugin: an omitted OR an explicitly + # empty exporter list falls back to the Lambda log exporter, so the + # plugin is never a silent no-op. A non-empty list is used verbatim. + self._exporters: list[InsightExporter] = ( + list(config.exporters) if config.exporters else [LambdaLogExporter()] + ) + self._scheduler = _ExportScheduler(self._exporters) + + def create_plugin(self, info: InvocationStartInfo) -> WorkflowInsightPlugin: + return WorkflowInsightPlugin(self, info) + +def workflow_insight(config: WorkflowInsightConfig) -> WorkflowInsightPluginFactory: + """Create a Workflow Insight plugin factory. Mirrors the JS ``workflowInsight()``. -def workflow_insight(config: WorkflowInsightConfig) -> WorkflowInsightPlugin: - """Create a Workflow Insight plugin. Mirrors the JS ``workflowInsight()`` factory.""" - return WorkflowInsightPlugin(config) + Pass the result straight to ``@durable_execution(plugins=[...])``: the SDK + calls its ``create_plugin`` once per invocation to build that invocation's + plugin instance. + """ + return WorkflowInsightPluginFactory(config) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/types.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/types.py index f813931c..5743d5ee 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/types.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/types.py @@ -63,7 +63,39 @@ def render(self, record: dict[str, Any]) -> Any: ... # pragma: no cover def export(self, record: dict[str, Any]) -> None: ... # pragma: no cover - def flush(self) -> None: ... # pragma: no cover + def flush(self) -> None: # pragma: no cover + """Push any records this exporter is buffering to their destination. + + Only an exporter that buffers needs a body here; one that writes + synchronously inside ``export()`` can leave it empty. The method itself is + *not* optional: it is part of this protocol, so an exporter without it + fails the static protocol check, and at run time the plugin's call raises + ``AttributeError``, which is caught and logged as an exporter failure on + every flush. + + When the plugin calls it: + + * Once per sampled-in invocation end, after that invocation's own record + -- if the emit mode produced one -- has been handed to every exporter. + An execution that is sampled out neither exports nor flushes. + Invocation ends that overlap in one environment may share a single + flush, so the call count is at most one per sampled-in invocation end. + * Never concurrently with ``export()`` on the same plugin instance: one + worker thread runs both, one call at a time. + + A flush may cover records belonging to other executions running in the + same environment, so it is not a per-execution barrier. + + It must return promptly. The invocation that triggered it cannot return + until it does, so a slow flush is billed to the customer's invocation. + + Failures are isolated, whatever is raised. An exception is logged, never + retried, never propagated into the execution, and never prevents another + exporter from flushing. That holds for a ``BaseException`` too + (``asyncio.CancelledError`` is one): the export worker is not the thread + such a signal is ever addressed to, so one raised here is a report of a + defective exporter and is contained exactly like any other failure. + """ @dataclass(frozen=True) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/e2e/wait_suspend_resume_int_test.py b/packages/aws-durable-execution-sdk-python-insight/tests/e2e/wait_suspend_resume_int_test.py index a52bf168..e5633c21 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/e2e/wait_suspend_resume_int_test.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/e2e/wait_suspend_resume_int_test.py @@ -66,10 +66,13 @@ def _insight_handler(event: Any, context: DurableContext) -> str: # noqa: ARG00 def test_terminal_record_includes_prior_step_and_completed_wait() -> None: capture = _CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[capture])) + factory = workflow_insight(WorkflowInsightConfig(exporters=[capture])) # Functional form (not the decorator-factory form) so the wrapped handler's - # static type stays a plain 2-arg callable for the runner. - handler = durable_execution(_insight_handler, plugins=[plugin]) + # static type stays a plain 2-arg callable for the runner. `factory` is the + # plugin factory the SDK calls once per invocation, so the two invocations + # this test drives run on two instances -- which is what makes the assertions + # below about the resuming invocation meaningful. + handler = durable_execution(_insight_handler, plugins=[factory]) with DurableFunctionTestRunner(handler=handler, execution_timeout=15) as runner: result: DurableFunctionTestResult = runner.run(input="{}") diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py index 34879059..7cd56b9d 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -4,15 +4,54 @@ from __future__ import annotations +import logging import threading import time from typing import Any from aws_durable_execution_sdk_python_insight._export_scheduler import ( + _MAX_CONSECUTIVE_WORKER_FAULTS, _ExportScheduler, + _ExportState, ) +ARN = "arn:aws:lambda:us-west-2:123456789012:function:my-fn:$LATEST/durable-execution/exec-{}/inv-1" +ARN_A = ARN.format("a") +ARN_B = ARN.format("b") + + +class _ArnScheduler(_ExportScheduler): + """Supplies the ARN -> execution map that the SDK's plugin lifecycle provides. + + ``schedule()`` and ``drain()`` take the per-execution object: the scheduler + holds the object the caller already has -- in production the caller's own + per-invocation plugin instance, which carries its export bookkeeping as an + ``_ExportState`` -- instead of resolving an ARN to bookkeeping of its own, so + it does not know what an ARN is. These tests drive the scheduler directly, so + they own an ARN map of their own and are otherwise unchanged. + """ + + def __init__(self, exporters: list[Any]) -> None: + super().__init__(exporters) + self.executions: dict[str, _ExportState] = {} + self._executions_lock = threading.Lock() + + def _execution(self, execution_arn: str) -> _ExportState: + with self._executions_lock: + execution = self.executions.get(execution_arn) + if execution is None: + execution = _ExportState() + self.executions[execution_arn] = execution + return execution + + def schedule(self, execution_arn: str, record: dict[str, Any]) -> None: # type: ignore[override] + super().schedule(self._execution(execution_arn), record) + + def drain(self, execution_arn: str) -> None: # type: ignore[override] + super().drain(self._execution(execution_arn)) + + def _record(value: str) -> dict[str, Any]: return {"status": "RUNNING", "value": value, "operations": []} @@ -62,20 +101,94 @@ def flush(self) -> None: raise RuntimeError("flush failed") -def test_latest_pending_coalesces_without_blocking_schedule() -> None: +class ExporterBaseException(BaseException): + """Stands in for the BaseExceptions a customer exporter can raise.""" + + +class BaseExceptionExporter(CaptureExporter): + """Raises a ``BaseException`` out of its first ``export()`` call.""" + + def __init__(self) -> None: + super().__init__() + self.exports = 0 + + def export(self, record: dict[str, Any]) -> None: + self.exports += 1 + if self.exports == 1: + raise ExporterBaseException("export exploded") + super().export(record) + + +class AlwaysBaseExceptionExportExporter(CaptureExporter): + """Raises a ``BaseException`` out of every ``export()`` call, and counts them.""" + + def __init__(self) -> None: + super().__init__() + self.lock = threading.Lock() + self.export_attempts = 0 + + def export(self, record: dict[str, Any]) -> None: + with self.lock: + self.export_attempts += 1 + raise ExporterBaseException("export exploded") + + def exports(self) -> int: + with self.lock: + return self.export_attempts + + +class AlwaysBaseExceptionFlushExporter(CaptureExporter): + """Raises a ``BaseException`` out of every ``flush()`` call, and counts them.""" + + def __init__(self) -> None: + super().__init__() + self.lock = threading.Lock() + self.flush_attempts = 0 + + def flush(self) -> None: + with self.lock: + self.flush_attempts += 1 + raise ExporterBaseException("flush exploded") + + def flushes(self) -> int: + with self.lock: + return self.flush_attempts + + +def _drain_off_thread( + scheduler: _ArnScheduler, arn: str +) -> tuple[threading.Thread, threading.Event]: + """Start drain() on its own thread and return it with its completion event. + + A regression that parks the drain would block whichever thread called it, so + no test may call drain() on the thread it asserts from. Waiting on the event + with a timeout turns such a regression into a failure instead of a hung run. + """ + returned = threading.Event() + + def drain() -> None: + scheduler.drain(arn) + returned.set() + + thread = threading.Thread(target=drain, daemon=True) + thread.start() + return thread, returned + + +def test_latest_pending_coalesces_within_one_execution() -> None: exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter]) - scheduler.schedule(_record("first")) + scheduler = _ArnScheduler([exporter]) + scheduler.schedule(ARN_A, _record("first")) assert exporter.started.wait(5.0) start = time.monotonic() - scheduler.schedule(_record("middle")) - scheduler.schedule(_record("latest")) + scheduler.schedule(ARN_A, _record("middle")) + scheduler.schedule(ARN_A, _record("latest")) assert time.monotonic() - start < 0.5 assert scheduler._pending_count() == 1 exporter.release.set() - scheduler.drain() + scheduler.drain(ARN_A) assert exporter.calls == [ ("export", "first"), ("export", "latest"), @@ -87,21 +200,175 @@ def test_latest_pending_coalesces_without_blocking_schedule() -> None: def test_exporter_failure_does_not_block_other_exporters() -> None: failing = FailingExporter() capture = CaptureExporter() - scheduler = _ExportScheduler([failing, capture]) + scheduler = _ArnScheduler([failing, capture]) - scheduler.schedule(_record("terminal")) + scheduler.schedule(ARN_A, _record("terminal")) - scheduler.drain() + scheduler.drain(ARN_A) assert capture.calls == [("export", "terminal"), ("flush", None)] +def test_base_exception_from_export_still_releases_drain() -> None: + # _export() contains every Exception, but a BaseException from a customer + # exporter -- asyncio.CancelledError is one -- unwinds out of the worker + # instead. The record has already been taken out of _pending by then and + # nothing will re-export that snapshot, so the export has to count and the + # worker slot has to be vacated anyway; otherwise the drain, and with it the + # invocation thread, parks forever. Every wait here is bounded so a + # regression fails instead of hanging the suite. + exporter = BaseExceptionExporter() + scheduler = _ArnScheduler([exporter]) + scheduler.schedule(ARN_A, _record("terminal")) + returned = threading.Event() + + def drain() -> None: + scheduler.drain(ARN_A) + returned.set() + + thread = threading.Thread(target=drain, daemon=True) + thread.start() + + assert returned.wait(10.0), "drain() never returned after export() raised" + thread.join(5.0) + assert not thread.is_alive() + # The snapshot was consumed, not retried, and the drain still got its flush. + assert exporter.exports == 1 + assert exporter.calls == [("flush", None)] + assert _wait_until(lambda: not scheduler._worker_alive()) + + +def test_always_failing_flush_releases_the_drain_without_a_respawn() -> None: + # A drain is only released by a flush that COMPLETED, and a worker that dies + # is replaced by whoever is waiting. An exporter whose flush() raises a + # BaseException on every call therefore used to make the drain start a + # replacement worker, which ran the same flush and died the same way, without + # bound: the flush was attempted thousands of times per second, thousands of + # threads were created, and the invocation parked on that drain never + # returned. The flush has to be attempted once, the failure reported, and the + # flush counted as completed so the drain returns. + # + # The wait is bounded, so the regression this pins fails the test rather than + # blocking the run. + exporter = AlwaysBaseExceptionFlushExporter() + scheduler = _ArnScheduler([exporter]) + scheduler.schedule(ARN_A, _record("terminal")) + thread, returned = _drain_off_thread(scheduler, ARN_A) + + assert returned.wait(10.0), "drain() never returned while flush() kept failing" + thread.join(5.0) + assert not thread.is_alive() + assert _wait_until(lambda: not scheduler._worker_alive()) + + # One attempt, not one per replacement worker. + assert exporter.flushes() == 1 + # The record still reached the exporter, and the failing flush was not retried + # after the drain returned either. + assert exporter.calls == [("export", "terminal")] + assert exporter.flushes() == 1 + + +def test_base_exception_from_one_exporter_never_skips_the_next() -> None: + # Consuming a record from the pending slot is what advances the completion + # bookkeeping, and nothing re-exports a consumed snapshot. A first exporter + # that raised a BaseException used to abort the fan-out loop, so every + # exporter after it missed that record permanently while the drain was + # released as though the record had been delivered. Each exporter's failure + # has to be contained at its own call, so the next exporter still receives the + # record that the bookkeeping counts as offered. + failing = AlwaysBaseExceptionExportExporter() + healthy = CaptureExporter() + scheduler = _ArnScheduler([failing, healthy]) + scheduler.schedule(ARN_A, _record("terminal")) + thread, returned = _drain_off_thread(scheduler, ARN_A) + + assert returned.wait(10.0), "drain() never returned after export() raised" + thread.join(5.0) + assert not thread.is_alive() + assert _wait_until(lambda: not scheduler._worker_alive()) + + assert healthy.calls == [("export", "terminal"), ("flush", None)] + # Offered once. The snapshot is gone from the pending slot, so a retry is not + # available and must not be implied. + assert failing.exports() == 1 + + +def test_base_exception_from_one_exporters_flush_never_skips_the_next() -> None: + # The same containment at the flush call. A first exporter whose flush() + # raises a BaseException must not stop a later exporter from flushing, and the + # flush must still count as completed so the waiting drain is released. + failing = AlwaysBaseExceptionFlushExporter() + healthy = CaptureExporter() + scheduler = _ArnScheduler([failing, healthy]) + scheduler.schedule(ARN_A, _record("terminal")) + thread, returned = _drain_off_thread(scheduler, ARN_A) + + assert returned.wait(10.0), "drain() never returned while flush() kept failing" + thread.join(5.0) + assert not thread.is_alive() + assert _wait_until(lambda: not scheduler._worker_alive()) + + assert healthy.calls == [("export", "terminal"), ("flush", None)] + assert failing.flushes() == 1 + + +class _FaultingScheduler(_ArnScheduler): + """Kills every export worker with a ``BaseException`` before it does any work. + + Stands in for a fault in the scheduler's own code rather than in an exporter: + exporter failures are contained at the exporter call, so they can no longer + reach the worker's exit path, and this is the only way left to drive it. + """ + + def __init__(self, exporters: list[Any]) -> None: + super().__init__(exporters) + self.runs = 0 + + def _run_loop(self) -> None: + with self._condition: + self.runs += 1 + raise ExporterBaseException("worker exploded") + + +def test_worker_deaths_stop_at_the_bound_instead_of_respawning_forever() -> None: + # A waiting drain starts a replacement worker for every worker that dies, so a + # fault the worker reproduces on every attempt is retried as fast as threads + # can be created and the drain never returns. Releasing the waiter matters + # more than delivering the records, because the waiter is an invocation + # thread: bound the replacements, then latch asynchronous export off, which + # wakes every waiter and drops what is queued. + exporter = CaptureExporter() + scheduler = _FaultingScheduler([exporter]) + scheduler.schedule(ARN_A, _record("terminal")) + thread, returned = _drain_off_thread(scheduler, ARN_A) + + assert returned.wait(10.0), "drain() never returned while the worker kept dying" + thread.join(5.0) + assert not thread.is_alive() + + with scheduler._condition: + assert scheduler.runs == _MAX_CONSECUTIVE_WORKER_FAULTS + assert scheduler._worker_faults == _MAX_CONSECUTIVE_WORKER_FAULTS + # The latch is what released the drain, and it retains nothing. + assert scheduler._disabled + assert scheduler._pending == {} + assert scheduler._flush_requested is False + assert scheduler._flush_in_flight is None + # No further worker is started once the latch is set, so the count cannot + # creep up after the drain returned. + scheduler.schedule(ARN_B, _record("after-the-latch")) + scheduler.drain(ARN_B) + with scheduler._condition: + assert scheduler.runs == _MAX_CONSECUTIVE_WORKER_FAULTS + assert exporter.calls == [] + + def test_drain_flushes_after_export() -> None: capture = CaptureExporter() - scheduler = _ExportScheduler([capture]) + scheduler = _ArnScheduler([capture]) - scheduler.schedule(_record("terminal")) + scheduler.schedule(ARN_A, _record("terminal")) - scheduler.drain() + scheduler.drain(ARN_A) assert capture.calls == [("export", "terminal"), ("flush", None)] @@ -110,45 +377,45 @@ def fail_start(self) -> None: # noqa: ARG001 raise RuntimeError("cannot start") monkeypatch.setattr(threading.Thread, "start", fail_start) - scheduler = _ExportScheduler([CaptureExporter()]) + scheduler = _ArnScheduler([CaptureExporter()]) - scheduler.schedule(_record("dropped")) - scheduler.drain() + scheduler.schedule(ARN_A, _record("dropped")) + scheduler.drain(ARN_A) assert scheduler._pending_count() == 0 def test_superseded_record_finalizes_after_lane_unlock() -> None: - scheduler = _ExportScheduler([BlockingExporter()]) + scheduler = _ArnScheduler([BlockingExporter()]) exporter = scheduler._exporters[0] assert isinstance(exporter, BlockingExporter) - scheduler.schedule(_record("inflight")) + scheduler.schedule(ARN_A, _record("inflight")) assert exporter.started.wait(5.0) finalized = threading.Event() class ReentrantValue: def __del__(self) -> None: - scheduler.schedule(_record("from-finalizer")) + scheduler.schedule(ARN_A, _record("from-finalizer")) finalized.set() pending = _record("superseded") pending["payload"] = ReentrantValue() - scheduler.schedule(pending) + scheduler.schedule(ARN_A, pending) del pending - scheduler.schedule(_record("replacement")) + scheduler.schedule(ARN_A, _record("replacement")) assert finalized.wait(5.0) exporter.release.set() - scheduler.drain() + scheduler.drain(ARN_A) def test_drain_waits_for_blocked_exporter() -> None: exporter = BlockingExporter() - scheduler = _ExportScheduler([exporter]) - scheduler.schedule(_record("terminal")) + scheduler = _ArnScheduler([exporter]) + scheduler.schedule(ARN_A, _record("terminal")) assert exporter.started.wait(5.0) - drain_thread = threading.Thread(target=scheduler.drain) + drain_thread = threading.Thread(target=scheduler.drain, args=(ARN_A,)) drain_thread.start() assert _wait_until(drain_thread.is_alive) @@ -157,3 +424,594 @@ def test_drain_waits_for_blocked_exporter() -> None: assert not drain_thread.is_alive() assert exporter.calls == [("export", "terminal"), ("flush", None)] + + +# -- concurrent executions in one environment (LMI) --------------------------- + + +class EventLogExporter: + """Records ``(kind, executionArn, status)`` events from every thread.""" + + max_record_size_bytes: int | None = None + + def __init__(self) -> None: + self.events: list[tuple[str, str | None, str | None]] = [] + self.lock = threading.Lock() + + def render(self, record: dict[str, Any]) -> dict[str, Any]: + return record + + def export(self, record: dict[str, Any]) -> None: + self.log("export", record["executionArn"], record["status"]) + + def flush(self) -> None: + self.log("flush", None, None) + + def log(self, kind: str, arn: str | None, status: str | None) -> None: + with self.lock: + self.events.append((kind, arn, status)) + + def snapshot(self) -> list[tuple[str, str | None, str | None]]: + with self.lock: + return list(self.events) + + +class GatedEventLogExporter(EventLogExporter): + """Blocks inside the first ``export()`` until released.""" + + def __init__(self) -> None: + super().__init__() + self.first_export_started = threading.Event() + self.release = threading.Event() + + def export(self, record: dict[str, Any]) -> None: + if not self.first_export_started.is_set(): + self.first_export_started.set() + self.release.wait(10.0) + super().export(record) + + +class FlushGateEventLogExporter(EventLogExporter): + """Logs flush begin/end and blocks inside the first ``flush()`` until released.""" + + def __init__(self) -> None: + super().__init__() + self.first_flush_started = threading.Event() + self.release_flush = threading.Event() + self.flush_count = 0 + + def flush(self) -> None: + with self.lock: + self.flush_count += 1 + first = self.flush_count == 1 + self.log("flush_begin", None, None) + if first: + self.first_flush_started.set() + self.release_flush.wait(10.0) + self.log("flush_end", None, None) + + def flushes(self) -> int: + with self.lock: + return self.flush_count + + +def _execution_record(arn: str, status: str) -> dict[str, Any]: + return {"executionArn": arn, "status": status, "operations": []} + + +def _scheduler_is_empty(scheduler: _ExportScheduler) -> bool: + with scheduler._condition: + return not scheduler._pending + + +def _drain_waiters(scheduler: _ArnScheduler, arn: str) -> int: + """How many drain() calls are currently parked on this execution.""" + with scheduler._condition: + return scheduler._execution(arn)._waiters + + +def test_drain_stays_parked_until_a_flush_covering_its_record_completes() -> None: + # The guarantee the per-execution scheduler exists for: drain() returns only + # after the calling execution's own record reached the exporters AND a flush + # that started after that export has itself completed. Gating the exporter + # inside flush() makes that deterministic -- while the gate is held the flush + # provably cannot have completed, so a drain that returns is a violation. + exporter = FlushGateEventLogExporter() + scheduler = _ArnScheduler([exporter]) + scheduler.schedule(ARN_A, _execution_record(ARN_A, "SUCCEEDED")) + returned = threading.Event() + + def drain() -> None: + scheduler.drain(ARN_A) + exporter.log("drain", ARN_A, None) + returned.set() + + thread = threading.Thread(target=drain, daemon=True) + thread.start() + assert exporter.first_flush_started.wait(5.0) + # Its own record went to the exporters before this flush was even started. + assert ("export", ARN_A, "SUCCEEDED") in exporter.snapshot() + assert not returned.wait(0.25), "drain returned while its flush was still running" + assert thread.is_alive() + + exporter.release_flush.set() + thread.join(5.0) + assert not thread.is_alive() + + events = exporter.snapshot() + exported = events.index(("export", ARN_A, "SUCCEEDED")) + flush_begin = events.index(("flush_begin", None, None)) + flush_end = events.index(("flush_end", None, None)) + drained = events.index(("drain", ARN_A, None)) + assert exported < flush_begin < flush_end < drained + + +def test_no_redundant_flush_runs_after_drain_returned() -> None: + # A drain woken while a flush is in flight must recognise that the in-flight + # flush already covers it. Otherwise it re-requests one (its own request has + # been consumed and the coverage is not published yet) and that second flush + # calls the exporters after drain(), and with it the invocation, returned. + exporter = FlushGateEventLogExporter() + scheduler = _ArnScheduler([exporter]) + scheduler.schedule(ARN_A, _execution_record(ARN_A, "SUCCEEDED")) + returned = threading.Event() + flushes_at_return: list[int] = [] + + def drain() -> None: + scheduler.drain(ARN_A) + flushes_at_return.append(exporter.flushes()) + returned.set() + + thread = threading.Thread(target=drain, daemon=True) + thread.start() + # The worker is inside the flush it committed to: it has consumed the drain's + # request and has not published the flush's coverage yet. + assert exporter.first_flush_started.wait(5.0) + # Wake the parked drain exactly inside that window, which is the interleaving + # a loaded environment produces by itself. + for _ in range(3): + with scheduler._condition: + scheduler._condition.notify_all() + time.sleep(0.01) + assert not returned.is_set() + + exporter.release_flush.set() + thread.join(5.0) + assert not thread.is_alive() + # The worker only retires once nothing is pending and no flush is requested, + # so this settles the question without sleeping for a late flush. + assert _wait_until(lambda: not scheduler._worker_alive()) + + assert flushes_at_return == [1] + assert exporter.flushes() == 1, "a second flush ran after drain() returned" + + +def test_drain_with_nothing_to_export_still_flushes_exactly_once() -> None: + # A drain for an execution that scheduled no record -- the invocation end + # emitted nothing -- has `need` == 0: no export has to be covered to release + # it. `_flush_in_flight` uses 0 for "no flush is running", so a naive + # `_flush_in_flight >= need` reads as "a flush already covers me" exactly + # when nothing is running at all; the drain would skip its request and park + # until some other execution happened to flush. It has to flush once and + # return. + capture = CaptureExporter() + scheduler = _ArnScheduler([capture]) + returned = threading.Event() + + def drain() -> None: + scheduler.drain(ARN_A) + returned.set() + + thread = threading.Thread(target=drain, daemon=True) + thread.start() + assert returned.wait(10.0), "drain() parked with nothing of its own to export" + thread.join(5.0) + assert not thread.is_alive() + assert capture.calls == [("flush", None)] + # Nothing arrives after it returned either. + assert _wait_until(lambda: not scheduler._worker_alive()) + assert capture.calls == [("flush", None)] + + +class GatedFlushExporter(CaptureExporter): + """Holds each ``flush()`` open until the test releases it, and counts them.""" + + def __init__(self) -> None: + super().__init__() + self._lock = threading.Lock() + self.flush_count = 0 + self.started: dict[int, threading.Event] = {} + self.release: dict[int, threading.Event] = {} + for index in (1, 2): + self.started[index] = threading.Event() + self.release[index] = threading.Event() + + def flush(self) -> None: + with self._lock: + self.flush_count += 1 + index = self.flush_count + if index in self.started: + self.started[index].set() + self.release[index].wait(10.0) + super().flush() + + +def test_concurrent_drains_with_nothing_to_export_share_one_flush() -> None: + # Two invocations that emitted no record drain at the same time. Neither has + # an export to be covered, so both need a flush that covers zero exports. + # While `_flush_in_flight` used 0 for "no flush is running", the second drain + # could not tell that the flush it needs was already running, and asked for + # another. The completion of the first flush then released both drains, and + # the second flush ran after both invocations had already returned -- customer + # exporter code running past the invocation boundary, which is what the flush + # contract forbids. One flush must serve both, and whichever drain a flush + # belongs to must stay parked until that flush completes. + exporter = GatedFlushExporter() + scheduler = _ArnScheduler([exporter]) + returned: list[str] = [] + returned_lock = threading.Lock() + + def drain(name: str, execution_arn: str) -> None: + scheduler.drain(execution_arn) + with returned_lock: + returned.append(name) + + threads = [ + threading.Thread(target=drain, args=("first", ARN_A), daemon=True), + threading.Thread(target=drain, args=("second", ARN_B), daemon=True), + ] + try: + threads[0].start() + assert exporter.started[1].wait(10.0), "the first drain never flushed" + threads[1].start() + + first = scheduler.executions[ARN_A] + + def both_parked() -> bool: + second = scheduler.executions.get(ARN_B) + if second is None: + return False + with scheduler._condition: + return first._waiters == 1 and second._waiters == 1 + + assert _wait_until(both_parked), "a drain raced past the flush it needs" + with returned_lock: + assert returned == [], "a drain returned before its flush completed" + + exporter.release[1].set() + for thread in threads: + thread.join(10.0) + assert not any(thread.is_alive() for thread in threads) + with returned_lock: + assert sorted(returned) == ["first", "second"] + + # The redundant request, if one was made, was recorded before either + # drain returned, so the worker starts that flush without further + # prompting. Nothing arriving here is what proves no second flush was + # requested. + assert not exporter.started[2].wait(0.75), ( + "a second flush ran after both invocations had returned" + ) + assert exporter.flush_count == 1 + finally: + exporter.release[1].set() + exporter.release[2].set() + _wait_until(lambda: not scheduler._worker_alive()) + + +def test_drain_never_rides_on_a_flush_that_finished_before_it_started() -> None: + # Export coverage alone would let the second drain return immediately: every + # export is already covered by the first drain's flush. A drain must wait for + # a flush that completed after it was called, so that one invocation end + # means one flush. + capture = CaptureExporter() + scheduler = _ArnScheduler([capture]) + scheduler.schedule(ARN_A, _record("terminal")) + + scheduler.drain(ARN_A) + scheduler.drain(ARN_A) + + assert capture.calls == [ + ("export", "terminal"), + ("flush", None), + ("flush", None), + ] + + +def test_disabled_latch_retains_no_lanes_or_pending_records(monkeypatch) -> None: + # The _disabled latch is permanent: nothing will ever be exported again, so + # the scheduler must not hold on to records or per-execution bookkeeping for + # the remaining life of the environment. + def fail_start(self) -> None: # noqa: ARG001 + raise RuntimeError("cannot start") + + monkeypatch.setattr(threading.Thread, "start", fail_start) + scheduler = _ArnScheduler([CaptureExporter()]) + + scheduler.schedule(ARN_A, _record("dropped")) + scheduler.drain(ARN_A) + scheduler.schedule(ARN_B, _record("also-dropped")) + scheduler.drain(ARN_B) + + with scheduler._condition: + assert scheduler._disabled + assert scheduler._pending == {} + # ...and no execution kept the record it was carrying. The per-execution + # bookkeeping that used to need clearing in a second map is on these + # objects now, so the queue and the records are one thing to release. + assert all( + execution._pending_record is None + for execution in scheduler.executions.values() + ) + assert scheduler._flush_requested is False + assert scheduler._flush_in_flight is None + + +def test_disabled_latch_clears_a_published_flush_in_flight_marker(monkeypatch) -> None: + # `_flush_in_flight` is the coverage of the flush the worker is running right + # now, published so a waiter woken during that flush can tell it is already + # covered and skip requesting another. None means "no flush is running", so + # any integer -- 0 included -- is a claim that a flush is in flight and will + # complete. + # + # The _disabled latch makes that claim permanently false: no worker exists and + # none will ever be started again, so the published flush can never complete. + # The latch therefore has to retire the marker along with the pending records + # and the lanes. test_disabled_latch_retains_no_lanes_or_pending_records + # asserts the same field, but reaches the latch with the marker already at 0, + # so it holds whether or not the latch clears it; this one arms the marker + # first. + scheduler = _ArnScheduler([CaptureExporter()]) + with scheduler._condition: + scheduler._flush_in_flight = 7 + + def fail_start(self) -> None: # noqa: ARG001 + raise RuntimeError("cannot start") + + monkeypatch.setattr(threading.Thread, "start", fail_start) + scheduler.schedule(ARN_A, _record("dropped")) + + with scheduler._condition: + assert scheduler._disabled + assert scheduler._flush_in_flight is None, ( + "the _disabled latch left a flush-in-flight marker behind for a flush " + "that can never run" + ) + + +def test_every_concurrent_execution_delivers_its_terminal_record_once() -> None: + # One plugin instance (one scheduler) serves every execution the environment + # hosts. Ten executions running at once must each land their terminal record + # exactly once: a pending record keyed per execution is never displaced by a + # different execution's record. + executions = 10 + exporter = EventLogExporter() + scheduler = _ArnScheduler([exporter]) + arns = [ARN.format(index) for index in range(executions)] + ready = threading.Barrier(executions) + + def run(arn: str) -> None: + ready.wait(10.0) + scheduler.schedule(arn, _execution_record(arn, "RUNNING")) + scheduler.schedule(arn, _execution_record(arn, "SUCCEEDED")) + scheduler.drain(arn) + + threads = [threading.Thread(target=run, args=(arn,)) for arn in arns] + for thread in threads: + thread.start() + for thread in threads: + thread.join(30.0) + assert not any(thread.is_alive() for thread in threads) + + terminal = [ + arn + for kind, arn, status in exporter.snapshot() + if kind == "export" and status == "SUCCEEDED" and arn is not None + ] + assert sorted(terminal) == sorted(arns) # each exactly once, none lost + # Nothing per-execution is retained once every execution has drained. + assert _wait_until(lambda: _scheduler_is_empty(scheduler)) + + +def test_blocked_export_never_loses_another_executions_terminal_record() -> None: + # The worker is busy exporting when two executions queue their terminal + # records. Neither may be dropped, and each drain must be released by its own + # record reaching the exporters -- not by another execution's flush. + exporter = GatedEventLogExporter() + scheduler = _ArnScheduler([exporter]) + scheduler.schedule(ARN_A, _execution_record(ARN_A, "RUNNING")) + assert exporter.first_export_started.wait(5.0) + + scheduler.schedule(ARN_A, _execution_record(ARN_A, "SUCCEEDED")) + scheduler.schedule(ARN_B, _execution_record(ARN_B, "SUCCEEDED")) + returned: dict[str, bool] = {} + + def drain(arn: str) -> None: + scheduler.drain(arn) + exporter.log("drain", arn, None) + returned[arn] = True + + drains = [threading.Thread(target=drain, args=(arn,)) for arn in (ARN_A, ARN_B)] + for thread in drains: + thread.start() + # Both drains really parked: each registered on its own lane (which drain() + # only does from inside its wait loop) and neither has returned while the + # exporter still holds the worker inside the very first export. + assert _wait_until( + lambda: ( + _drain_waiters(scheduler, ARN_A) == 1 + and _drain_waiters(scheduler, ARN_B) == 1 + ) + ) + assert returned == {} + assert all(thread.is_alive() for thread in drains) + + exporter.release.set() + for thread in drains: + thread.join(10.0) + assert not any(thread.is_alive() for thread in drains) + assert returned == {ARN_A: True, ARN_B: True} + + events = exporter.snapshot() + assert ("export", ARN_A, "SUCCEEDED") in events + assert ("export", ARN_B, "SUCCEEDED") in events + for arn in (ARN_A, ARN_B): + exported = events.index(("export", arn, "SUCCEEDED")) + returned_at = events.index(("drain", arn, None)) + # Each drain returned only after its own record was exported and a flush + # covering it completed. + assert exported < returned_at + assert ("flush", None, None) in events[exported:returned_at] + assert _wait_until(lambda: _scheduler_is_empty(scheduler)) + + +class ReentrantDrainExporter(CaptureExporter): + """Exporter that queues a record and drains from inside export(). + + This is what an exporter re-entering a plugin hook produces: the hook emits + its record and then, at an invocation end, drains. Both arrive on the export + worker, which is the one thread able to serve the wait. + """ + + def __init__(self) -> None: + super().__init__() + self.scheduler: _ArnScheduler | None = None + self.returned = threading.Event() + self._reentered = False + + def export(self, record: dict[str, Any]) -> None: + super().export(record) + assert self.scheduler is not None + if self._reentered: + return + self._reentered = True + self.scheduler.schedule(ARN_B, _record("r2")) + self.scheduler.drain(ARN_B) + self.returned.set() + + +def test_drain_from_the_export_worker_is_refused_rather_than_deadlocking( + caplog, +) -> None: + """A drain on the worker thread returns instead of parking it. + + Only the export worker exports records and completes flushes. A drain made on + that thread would wait for work only that thread can do, so the wait never + ends and the invocation hangs until Lambda times it out. The call is refused + and reported, and the worker goes back to its loop. + """ + exporter = ReentrantDrainExporter() + scheduler = _ArnScheduler([exporter]) + exporter.scheduler = scheduler + + with caplog.at_level(logging.WARNING): + scheduler.schedule(ARN_A, _record("r1")) + + assert exporter.returned.wait(timeout=10), "the refused drain must return" + assert _wait_until(lambda: "refused rather than deadlocking" in caplog.text) + + assert ("export", "r1") in exporter.calls + # The refused drain still leaves a flush behind, with no external drain to ask + # for one. The re-entering hook queued a record, and the worker exits once + # nothing is pending and no flush is requested, so without the request a + # buffering exporter would be holding that record when the environment froze. + assert _wait_until(lambda: ("export", "r2") in exporter.calls) + assert _wait_until( + lambda: exporter.calls.index(("flush", None)) + > exporter.calls.index(("export", "r2")) + ), "a refused drain must request a flush that covers the record it queued" + # The worker is still serving: a drain from any other thread completes. + scheduler.drain(ARN_B) + + +class ReentrantFlushExporter(CaptureExporter): + """Exporter whose flush() drains, as a hook re-entered from a flush would. + + An exporter that re-enters a plugin hook from ``flush()`` reaches + ``on_invocation_end``, which drains. The drain arrives on the export worker + from inside a flush, with nothing pending. + """ + + def __init__(self) -> None: + super().__init__() + self.scheduler: _ArnScheduler | None = None + self.flushes = 0 + + def flush(self) -> None: + super().flush() + self.flushes += 1 + assert self.scheduler is not None + self.scheduler.drain(ARN_A) + + +def test_a_drain_refused_from_inside_a_flush_does_not_re_arm_it() -> None: + """A refused drain with nothing pending asks for no further flush. + + Requesting one unconditionally would keep the worker flushing for as long as + the environment lived: the flush re-enters the hook, the hook drains, the + refused drain asks for the next flush. A pending record is what distinguishes + new work from that loop, so a drain refused from inside a flush leaves no + request behind. + """ + exporter = ReentrantFlushExporter() + scheduler = _ArnScheduler([exporter]) + exporter.scheduler = scheduler + + scheduler.schedule(ARN_A, _record("r1")) + scheduler.drain(ARN_A) + + flushes_after_drain = exporter.flushes + assert flushes_after_drain >= 1, "the drain must have flushed" + + # Give a re-armed flush time to appear. The worker retires when nothing is + # pending and no flush is requested, so a bounded count here is the whole + # assertion: an unconditional request never settles. + assert not _wait_until(lambda: exporter.flushes > flushes_after_drain, timeout=1.0) + assert _wait_until(lambda: not scheduler._worker_alive()) + + +class ReentrantDrainDuringExportExporter(CaptureExporter): + """Exporter whose export() drains without queueing anything new. + + The worker takes a record out of the pending map before it calls the + exporter, so a hook re-entered from inside ``export()`` and reaching an + invocation end that emits no record finds nothing pending -- while the record + it was just handed is still only in the exporter's buffer. + """ + + def __init__(self) -> None: + super().__init__() + self.scheduler: _ArnScheduler | None = None + self.returned = threading.Event() + self._reentered = False + + def export(self, record: dict[str, Any]) -> None: + super().export(record) + assert self.scheduler is not None + if self._reentered: + return + self._reentered = True + self.scheduler.drain(ARN_B) + self.returned.set() + + +def test_a_drain_refused_from_inside_an_export_still_flushes() -> None: + """A refused drain from inside an export asks for a flush. + + The record it must cover has already left the pending map, so a + pending-only condition would leave that snapshot in a buffering exporter + when the environment froze -- which is the loss the refusal path exists to + prevent, arriving by the other door. + """ + exporter = ReentrantDrainDuringExportExporter() + scheduler = _ArnScheduler([exporter]) + exporter.scheduler = scheduler + + scheduler.schedule(ARN_A, _record("r1")) + + assert exporter.returned.wait(timeout=10), "the refused drain must return" + assert _wait_until(lambda: ("flush", None) in exporter.calls), ( + "the exported record must be flushed even though nothing was pending" + ) + assert ("export", "r1") in exporter.calls + assert _wait_until(lambda: not scheduler._worker_alive()) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_package_metadata.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_package_metadata.py new file mode 100644 index 00000000..8134746f --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_package_metadata.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Packaging checks for the Workflow Insight plugin. + +The plugin and the core SDK are separate distributions, so pip resolves their +versions independently. A declared bound that admits a core release without the +contract this plugin uses is an install that succeeds and then fails at handler +initialization, which is why the bound is asserted here rather than left to +review. +""" + +from __future__ import annotations + +import re +import tomllib +from pathlib import Path + +from packaging.specifiers import SpecifierSet +from packaging.version import Version + + +PACKAGE_ROOT = Path(__file__).resolve().parents[1] +REPOSITORY_ROOT = PACKAGE_ROOT.parents[1] +CORE_DISTRIBUTION = "aws-durable-execution-sdk-python" + + +def _core_version() -> str: + """The core SDK version this repository builds, read from its source. + + Read from the file rather than imported, so the check keeps describing this + repository even when a published core is installed alongside these sources. + """ + about = ( + REPOSITORY_ROOT + / "packages" + / "aws-durable-execution-sdk-python" + / "src" + / "aws_durable_execution_sdk_python" + / "__about__.py" + ).read_text() + match = re.search(r'^__version__ = "([^"]+)"', about, re.MULTILINE) + assert match is not None, "core __about__.py has no __version__ assignment" + return match.group(1) + + +def _core_dependency_lower_bound() -> str: + with (PACKAGE_ROOT / "pyproject.toml").open("rb") as pyproject: + dependencies = tomllib.load(pyproject)["project"]["dependencies"] + + bounds = [ + dependency.removeprefix(CORE_DISTRIBUTION + ">=").split(",", 1)[0] + for dependency in dependencies + if dependency.startswith(CORE_DISTRIBUTION + ">=") + ] + assert len(bounds) == 1, f"expected one {CORE_DISTRIBUTION} bound, got {bounds}" + return bounds[0] + + +def _major(version: str) -> int: + return int(version.split(".", 1)[0]) + + +def test_core_dependency_bound_matches_the_core_major_in_this_repository() -> None: + """The declared bound must not admit a core major that predates the factory contract. + + ``workflow_insight()`` returns a plugin factory, and only the core major that + introduced factories calls it. A lower bound naming an earlier major is a + resolution pip accepts and that then fails at handler initialization, so the + bound tracks the core major this repository builds. The bound may lag within + that major -- a later core minor still satisfies the contract -- which is why + only the major is compared and the bound is required not to exceed the core + version. + """ + core_version = _core_version() + lower_bound = _core_dependency_lower_bound() + + assert _major(lower_bound) == _major(core_version) + assert Version(lower_bound) <= Version(core_version) + + +def test_core_dependency_excludes_the_next_core_major() -> None: + """A lower bound alone is the same defect one major later. + + The lower bound exists because this package's entry points resolve to plugin + factories, which the core major below cannot call: pip accepts the resolution + and the handler fails at initialization. Without a ceiling the next core major + that changes the plugin contract reproduces exactly that, so the specifier has + to reject it rather than only reject what came before. + """ + with (PACKAGE_ROOT / "pyproject.toml").open("rb") as pyproject: + dependencies = tomllib.load(pyproject)["project"]["dependencies"] + specifiers = [ + SpecifierSet(dependency.removeprefix(CORE_DISTRIBUTION)) + for dependency in dependencies + if dependency.startswith(CORE_DISTRIBUTION) + ] + assert len(specifiers) == 1 + + core_major = _major(_core_version()) + assert specifiers[0].contains(_core_version(), prereleases=True) + assert not specifiers[0].contains(f"{core_major + 1}.0.0", prereleases=True) + + +def test_core_dependency_admits_a_later_core_patch() -> None: + """The ceiling belongs on the major, not on the version built here. + + A ``<=`` ceiling looks equivalent and is not: it excludes the next core patch, + so the first core patch release puts this claim out of date for a change that + cannot have touched the plugin contract. + """ + with (PACKAGE_ROOT / "pyproject.toml").open("rb") as pyproject: + dependencies = tomllib.load(pyproject)["project"]["dependencies"] + specifier = next( + SpecifierSet(dependency.removeprefix(CORE_DISTRIBUTION)) + for dependency in dependencies + if dependency.startswith(CORE_DISTRIBUTION) + ) + + core = Version(_core_version()) + next_patch = f"{core.major}.{core.minor}.{core.micro + 1}" + assert specifier.contains(next_patch, prereleases=True) diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py index 9d7638eb..b55c59cc 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py @@ -8,13 +8,27 @@ exercised end to end, nothing about SDK behavior is mocked). Operations reach the plugin the way the real SDK delivers them: as the point-in-time ``operations`` map on ``InvocationStartInfo`` / ``InvocationEndInfo`` / ``OperationChangeInfo``. + +``workflow_insight()`` returns a factory, so these tests hold two things where +they used to hold one: the handler-lifetime factory (``factory``, carrying the +resolved config, the exporters and the export scheduler) and the per-invocation +instance the SDK builds from it (``plugin``). ``_invocation()`` does what the SDK +does -- build the instance from the invocation's start info, then dispatch that +same info to its first hook. """ from __future__ import annotations +import asyncio import datetime +import gc +import itertools +import threading +import time from typing import Any +import pytest + from aws_durable_execution_sdk_python.lambda_service import ( ErrorObject, OperationStatus, @@ -36,9 +50,14 @@ LambdaLogExporter, OperationOverride, WorkflowInsightConfig, + WorkflowInsightPluginFactory, workflow_insight, ) -from aws_durable_execution_sdk_python_insight.plugin import _resolve_sampling_rate +from aws_durable_execution_sdk_python_insight._export_scheduler import _ExportState +from aws_durable_execution_sdk_python_insight.plugin import ( + WorkflowInsightPlugin, + _resolve_sampling_rate, +) ARN = "arn:aws:lambda:us-west-2:123456789012:function:my-fn:$LATEST/durable-execution/exec-1/inv-1" ARN_B = "arn:aws:lambda:us-west-2:123456789012:function:my-fn:$LATEST/durable-execution/exec-2/inv-1" @@ -135,8 +154,20 @@ def _end( ) +def _invocation(factory, info: InvocationStartInfo) -> WorkflowInsightPlugin: + """Enter one invocation the way the SDK does. + + The SDK builds one plugin instance per invocation from that invocation's + start info and dispatches the very same object to its first hook. A test that + drives hooks directly does both. + """ + plugin = factory.create_plugin(info) + plugin.on_invocation_start(info) + return plugin + + def _run( - plugin, + factory, *, ops, status=InvocationStatus.SUCCEEDED, @@ -147,10 +178,53 @@ def _run( """Single-invocation drive: the full operation map is present in both the start and the end snapshot (the terminal record is built from the end one).""" operations = _ops(*ops) - plugin.on_invocation_start(_start(operations=operations, input_value=input_value)) + plugin = _invocation( + factory, _start(operations=operations, input_value=input_value) + ) plugin.on_invocation_end( _end(operations=operations, status=status, result=result, error=error) ) + return plugin + + +# -- public surface ---------------------------------------------------------- + + +def test_the_factory_type_is_public_and_re_exported(): + """A consumer must be able to name the type ``workflow_insight()`` returns. + + The package ships ``py.typed``, so a consumer annotating the value it holds + needs the name. A private name would force an import from a private module. + The sibling OTel package exports ``InvocationOtelPluginFactory`` and + ``ExecutionOtelPluginFactory`` for the same reason, so this keeps the two + plugin packages consistent. + """ + import aws_durable_execution_sdk_python_insight as pkg + + factory = workflow_insight(WorkflowInsightConfig(exporters=[CaptureExporter()])) + + assert type(factory) is WorkflowInsightPluginFactory + assert not WorkflowInsightPluginFactory.__name__.startswith("_") + assert "WorkflowInsightPluginFactory" in pkg.__all__ + assert pkg.WorkflowInsightPluginFactory is WorkflowInsightPluginFactory + + +def test_the_plugin_exposes_no_public_attributes(): + """Export bookkeeping must not become part of the plugin's public surface. + + ``WorkflowInsightPlugin`` is exported from the package root, and it mixes in + :class:`_ExportState`, so every field that mixin sets lands on the exported + class. Those fields belong to the export scheduler, which owns them under its + own lock. A public name among them would advertise scheduler bookkeeping as + something a consumer may read or set. So every field on the instance is + underscore-prefixed. + """ + exporter = CaptureExporter() + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + plugin = _invocation(factory, _start()) + + public = sorted(name for name in vars(plugin) if not name.startswith("_")) + assert public == [] # -- existing record-building coverage --------------------------------------- @@ -158,8 +232,8 @@ def _run( def test_basic_success_record(): exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) - _run(plugin, ops=[_step("greet")]) + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + _run(factory, ops=[_step("greet")]) assert len(exporter.records) == 1 rec = exporter.records[0] assert rec["recordType"] == "WorkflowInsight" @@ -182,21 +256,61 @@ def test_basic_success_record(): def test_on_failure_success_emits_nothing(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-failure") ) - _run(plugin, ops=[_step("greet")], status=InvocationStatus.SUCCEEDED) + _run(factory, ops=[_step("greet")], status=InvocationStatus.SUCCEEDED) + assert exporter.records == [] + # No record, but the invocation end still flushed once: a sampled-in + # invocation end flushes whether or not this emit mode produced a record + # (JS/Java cadence), because the exporter may be buffering another + # execution's records. + assert exporter.flush_count == 1 + assert _wait_until(lambda: not factory._scheduler._worker_alive()) + + +def test_invocation_end_that_emits_no_record_still_flushes_exactly_once(): + # on-complete mode with a PENDING end: the execution suspended, so nothing is + # emitted. JS and Java flush once per sampled-in invocation end regardless of + # whether a record was emitted, and a buffering exporter has to see the same + # rhythm in every SDK, so this end must still flush -- exactly once, not + # twice, and not zero times. + exporter = CaptureExporter() + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + plugin = _invocation(factory, _start(operations={})) + plugin.on_invocation_end( + _end(operations={}, status=InvocationStatus.PENDING, result=None) + ) + assert exporter.records == [] + assert exporter.flush_count == 1 + # The worker retires, so no later flush can arrive after the invocation + # returned. + assert _wait_until(lambda: not factory._scheduler._worker_alive()) + assert exporter.flush_count == 1 + + +def test_sampled_out_invocation_end_neither_exports_nor_flushes(): + # The sampled-out path is the one exception to the cadence above: a sampled + # out execution exports nothing and must not flush either, so instrumenting + # a fraction of executions costs the rest nothing. + exporter = CaptureExporter() + factory = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], sampling_rate=0) + ) + op = _step("s", op_id="1") + plugin = _invocation(factory, _start(operations={})) + plugin.on_invocation_end(_end(operations=_ops(op))) assert exporter.records == [] assert exporter.flush_count == 0 - assert not plugin._scheduler._worker_alive() + assert not factory._scheduler._worker_alive() def test_sampling_zero_emits_nothing(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], sampling_rate=0) ) - _run(plugin, ops=[_step("greet")]) + _run(factory, ops=[_step("greet")]) assert exporter.records == [] @@ -208,22 +322,22 @@ def test_resolve_sampling_rate_nan_fails_open_to_one(): def test_nan_sampling_rate_emits_instead_of_silently_disabling(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], sampling_rate=float("nan")) ) - _run(plugin, ops=[_step("greet")]) + _run(factory, ops=[_step("greet")]) # A NaN rate must not disable instrumentation: the record is still emitted. assert len(exporter.records) == 1 def test_content_omit_input_output_without_drop_flags(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig( exporters=[exporter], content=ContentConfig(input=False, output=False) ) ) - _run(plugin, ops=[_step("greet")]) + _run(factory, ops=[_step("greet")]) rec = exporter.records[0] assert "input" not in rec and "output" not in rec assert "droppedInput" not in rec and "droppedOutput" not in rec @@ -231,7 +345,7 @@ def test_content_omit_input_output_without_drop_flags(): def test_result_opt_in(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig( exporters=[exporter], content=ContentConfig( @@ -241,14 +355,14 @@ def test_result_opt_in(): ), ) ) - _run(plugin, ops=[_step("compute", result="42")], result="42") + _run(factory, ops=[_step("compute", result="42")], result="42") op = exporter.records[0]["operations"][0] assert op["result"] == 42 # checkpointed JSON string parsed def test_include_errors_false_drops_op_error_keeps_record_error(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig( exporters=[exporter], content=ContentConfig(operations=ContentOperations(include_errors=False)), @@ -259,7 +373,7 @@ def test_include_errors_false_drops_op_error_keeps_record_error(): message="boom", type="InsightTestError", data=None, stack_trace=None ) _run( - plugin, + factory, ops=[_step("failing-step", status=OperationStatus.FAILED, error=op_err)], status=InvocationStatus.FAILED, result=None, @@ -272,7 +386,7 @@ def test_include_errors_false_drops_op_error_keeps_record_error(): def test_top_level_only_drops_children(): exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) parent = _step( "parallel-work", op_id="p", @@ -280,14 +394,14 @@ def test_top_level_only_drops_children(): sub_type=OperationSubType.PARALLEL, ) child = _step("branch-a-step", parent_id="p", op_id="c") - _run(plugin, ops=[parent, child]) + _run(factory, ops=[parent, child]) names = [op["name"] for op in exporter.records[0]["operations"]] assert names == ["parallel-work"] def test_full_tree_includes_children_with_parent_id(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], operation_detail="full-tree") ) parent = _step( @@ -297,7 +411,7 @@ def test_full_tree_includes_children_with_parent_id(): sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, ) child = _step("child-step", parent_id="p", op_id="c") - _run(plugin, ops=[parent, child]) + _run(factory, ops=[parent, child]) ops = {op["name"]: op for op in exporter.records[0]["operations"]} assert set(ops) == {"parent-context", "child-step"} assert ops["child-step"]["parentId"] == "p" @@ -305,9 +419,9 @@ def test_full_tree_includes_children_with_parent_id(): def test_unnamed_operation_dropped(): exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) unnamed = _step(None, op_id="u") # type: ignore[arg-type] - _run(plugin, ops=[_step("named-step"), unnamed]) + _run(factory, ops=[_step("named-step"), unnamed]) names = [op["name"] for op in exporter.records[0]["operations"]] assert names == ["named-step"] @@ -316,9 +430,10 @@ def test_unnamed_operation_dropped(): def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): - # Invocation 1 (plugin A): a step completes, then a wait suspends -> PENDING. + # Invocation 1 (environment A): a step completes, then a wait suspends -> + # PENDING. exporter1 = CaptureExporter() - plugin1 = workflow_insight(WorkflowInsightConfig(exporters=[exporter1])) + factory1 = workflow_insight(WorkflowInsightConfig(exporters=[exporter1])) step = _step("greet", op_id="op-step") wait_pending = _step( "pause", @@ -328,7 +443,7 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): status=OperationStatus.PENDING, end_time=None, ) - plugin1.on_invocation_start(_start(operations={})) + plugin1 = _invocation(factory1, _start(operations={})) plugin1.on_invocation_end( _end( operations=_ops(step, wait_pending), @@ -337,12 +452,15 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): ) ) assert exporter1.records == [] # on-complete emits nothing for a suspend - assert plugin1._state == {} # and retains nothing + # And nothing is retained: the instance that served the suspending invocation + # is dropped by the SDK, and the scheduler holds no execution either. + assert _wait_until(lambda: _scheduler_is_empty(factory1)) - # Invocation 2 on a *fresh* plugin instance (new Lambda environment): the - # resume start snapshot carries the prior terminal step + resolved wait. + # Invocation 2 in a *fresh* Lambda environment -- a new factory, and so also a + # new instance: the resume start snapshot carries the prior terminal step + + # resolved wait. exporter2 = CaptureExporter() - plugin2 = workflow_insight(WorkflowInsightConfig(exporters=[exporter2])) + factory2 = workflow_insight(WorkflowInsightConfig(exporters=[exporter2])) step_done = _step("greet", op_id="op-step") wait_done = _step( "pause", @@ -352,8 +470,8 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): status=OperationStatus.SUCCEEDED, ) resume_ops = _ops(step_done, wait_done) - plugin2.on_invocation_start( - _start(operations=resume_ops, is_first=False, execution_start_time=T0) + plugin2 = _invocation( + factory2, _start(operations=resume_ops, is_first=False, execution_start_time=T0) ) plugin2.on_invocation_end( _end(operations=resume_ops, is_first=False, execution_start_time=T0) @@ -373,13 +491,13 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): def test_on_change_schedules_running_and_delivers_terminal(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") ) op1 = _step("s1", op_id="1") op2 = _step("s2", op_id="2") - plugin.on_invocation_start(_start(operations={})) + plugin = _invocation(factory, _start(operations={})) plugin.on_operation_change( OperationChangeInfo( execution_arn=ARN, updated_operations=_ops(op1), operations=_ops(op1) @@ -408,15 +526,17 @@ def test_on_change_schedules_running_and_delivers_terminal(): def test_concurrent_executions_do_not_cross_contaminate(): # A and B both suspend; B is the most-recently started (the old insertion- # order heuristic would have attributed A's resume to B). A then resumes to - # a terminal state. Its record must contain only A's data. + # a terminal state. Its record must contain only A's data. Each invocation + # gets its own instance from the one shared factory, which is what the SDK + # does for concurrent executions in one environment. exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) a_op = _step("a-step", op_id="a1") b_op = _step("b-step", op_id="b1") - plugin.on_invocation_start(_start(arn=ARN, operations={}, input_value="A")) - plugin.on_invocation_start(_start(arn=ARN_B, operations={}, input_value="B")) - plugin.on_invocation_end( + a_first = _invocation(factory, _start(arn=ARN, operations={}, input_value="A")) + b_first = _invocation(factory, _start(arn=ARN_B, operations={}, input_value="B")) + b_first.on_invocation_end( _end( arn=ARN_B, operations=_ops(b_op), @@ -424,7 +544,7 @@ def test_concurrent_executions_do_not_cross_contaminate(): result=None, ) ) - plugin.on_invocation_end( + a_first.on_invocation_end( _end( arn=ARN, operations=_ops(a_op), @@ -435,10 +555,11 @@ def test_concurrent_executions_do_not_cross_contaminate(): assert exporter.records == [] # both suspended, nothing terminal yet a_done = _step("a-step", op_id="a1") - plugin.on_invocation_start( - _start(arn=ARN, operations=_ops(a_done), is_first=False, input_value="A") + a_resume = _invocation( + factory, + _start(arn=ARN, operations=_ops(a_done), is_first=False, input_value="A"), ) - plugin.on_invocation_end( + a_resume.on_invocation_end( _end(arn=ARN, operations=_ops(a_done), is_first=False, result='"A-done"') ) @@ -449,22 +570,25 @@ def test_concurrent_executions_do_not_cross_contaminate(): assert [op["name"] for op in rec["operations"]] == ["a-step"] -# -- state lifecycle: clear after every invocation end (comment 4) ----------- +# -- nothing retained after an invocation end (comment 4) -------------------- -def test_state_cleared_after_pending_and_retry(): +def test_nothing_retained_after_pending_and_retry(): exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) op = _step("s", op_id="1") - plugin.on_invocation_start(_start(operations={})) - plugin.on_invocation_end( + suspend = _invocation(factory, _start(operations={})) + suspend.on_invocation_end( _end(operations=_ops(op), status=InvocationStatus.PENDING, result=None) ) - assert plugin._state == {} # no leak after suspend + # The instance that served the suspending invocation is dropped by the SDK, + # so the only thing that could retain anything for this execution is the + # scheduler, and it holds nothing either. + assert _wait_until(lambda: _scheduler_is_empty(factory)) - plugin.on_invocation_start(_start(operations=_ops(op), is_first=False)) - plugin.on_invocation_end( + retry = _invocation(factory, _start(operations=_ops(op), is_first=False)) + retry.on_invocation_end( _end( operations=_ops(op), status=InvocationStatus.RETRY, @@ -472,17 +596,17 @@ def test_state_cleared_after_pending_and_retry(): is_first=False, ) ) - assert plugin._state == {} # no leak after retry + assert _wait_until(lambda: _scheduler_is_empty(factory)) assert exporter.records == [] # on-complete emits nothing for non-terminal def test_sampled_out_processes_nothing_and_retains_no_state(): exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], sampling_rate=0) ) op = _step("s", op_id="1") - plugin.on_invocation_start(_start(operations={})) + plugin = _invocation(factory, _start(operations={})) plugin.on_operation_change( OperationChangeInfo( execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) @@ -490,33 +614,35 @@ def test_sampled_out_processes_nothing_and_retains_no_state(): ) plugin.on_invocation_end(_end(operations=_ops(op))) assert exporter.records == [] - assert plugin._state == {} + # A sampled-out invocation adopts no operations and schedules nothing. + assert plugin._operations == {} + assert _scheduler_is_empty(factory) # -- default exporter parity with JS (comment 6) ----------------------------- def test_default_exporter_when_config_omits_exporters(): - plugin = workflow_insight(WorkflowInsightConfig()) - assert len(plugin._exporters) == 1 - assert isinstance(plugin._exporters[0], LambdaLogExporter) + factory = workflow_insight(WorkflowInsightConfig()) + assert len(factory._exporters) == 1 + assert isinstance(factory._exporters[0], LambdaLogExporter) def test_default_exporter_when_exporters_explicitly_empty(): - plugin = workflow_insight(WorkflowInsightConfig(exporters=[])) - assert len(plugin._exporters) == 1 - assert isinstance(plugin._exporters[0], LambdaLogExporter) + factory = workflow_insight(WorkflowInsightConfig(exporters=[])) + assert len(factory._exporters) == 1 + assert isinstance(factory._exporters[0], LambdaLogExporter) def test_explicit_exporters_are_preserved(): exporter = CaptureExporter() - plugin = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) - assert plugin._exporters == [exporter] + factory = workflow_insight(WorkflowInsightConfig(exporters=[exporter])) + assert factory._exporters == [exporter] def test_default_exporter_actually_emits_to_stdout(capsys): - plugin = workflow_insight(WorkflowInsightConfig()) - _run(plugin, ops=[_step("greet")]) + factory = workflow_insight(WorkflowInsightConfig()) + _run(factory, ops=[_step("greet")]) out = capsys.readouterr().out assert '"recordType":"WorkflowInsight"' in out # compact JSON via LambdaLogExporter assert '"operationsByName"' in out @@ -533,11 +659,11 @@ def _last_record_for_end_status(status, *, emit_mode="on-change", result=None): emit nothing for a non-terminal end. """ exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode=emit_mode) ) op = _step("s", op_id="1") - plugin.on_invocation_start(_start(operations={})) + plugin = _invocation(factory, _start(operations={})) plugin.on_invocation_end(_end(operations=_ops(op), status=status, result=result)) return exporter.records[-1] @@ -569,11 +695,11 @@ def test_succeeded_end_is_terminal_with_end_time_and_duration(): def test_failed_end_is_terminal_with_end_time_and_duration(): err = ErrorObject(message="boom", type="StepError", data=None, stack_trace=None) exporter = CaptureExporter() - plugin = workflow_insight( + factory = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") ) op = _step("s", op_id="1", status=OperationStatus.FAILED) - plugin.on_invocation_start(_start(operations={})) + plugin = _invocation(factory, _start(operations={})) plugin.on_invocation_end( _end( operations=_ops(op), @@ -587,3 +713,938 @@ def test_failed_end_is_terminal_with_end_time_and_duration(): assert rec["endTime"] is not None assert rec["durationMs"] is not None assert rec["error"]["name"] == "StepError" + + +# -- concurrent executions in one environment (LMI) --------------------------- + + +class ConcurrentCaptureExporter: + """CaptureExporter for multi-threaded drives; appends under a lock.""" + + max_record_size_bytes = None + + def __init__(self) -> None: + self.records: list[dict[str, Any]] = [] + self.flushes = 0 + self._lock = threading.Lock() + + def render(self, record: dict[str, Any]) -> Any: + return record + + def export(self, record: dict[str, Any]) -> None: + with self._lock: + self.records.append(record) + + def flush(self) -> None: + with self._lock: + self.flushes += 1 + + def snapshot(self) -> list[dict[str, Any]]: + with self._lock: + return list(self.records) + + +def test_concurrent_executions_each_deliver_their_terminal_record(): + # One factory (one scheduler) serves every execution its environment hosts, + # and LMI runs several at once. Drive the real hooks concurrently, each + # execution on its own instance: every execution's terminal record must + # arrive exactly once. + executions = 5 + exporter = ConcurrentCaptureExporter() + factory = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") + ) + arns = [ + f"arn:aws:lambda:us-west-2:123456789012:function:my-fn:$LATEST/durable-execution/exec-{index}/inv-1" + for index in range(executions) + ] + ready = threading.Barrier(executions) + + def run(arn: str) -> None: + op = _step("s", op_id="1") + ready.wait(10.0) + plugin = _invocation(factory, _start(arn=arn, operations={})) + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=arn, updated_operations=_ops(op), operations=_ops(op) + ) + ) + plugin.on_invocation_end(_end(arn=arn, operations=_ops(op))) + + threads = [threading.Thread(target=run, args=(arn,)) for arn in arns] + for thread in threads: + thread.start() + for thread in threads: + thread.join(30.0) + assert not any(thread.is_alive() for thread in threads) + + terminal = [ + record["executionArn"] + for record in exporter.snapshot() + if record["status"] == "SUCCEEDED" + ] + assert sorted(terminal) == sorted(arns) # each exactly once, none lost + # Nothing per-execution is retained: each instance is dropped by the SDK with + # its invocation, and the scheduler holds no execution either. + assert _wait_until(lambda: _scheduler_is_empty(factory)) + + +def _wait_until(predicate, timeout: float = 5.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.005) + return predicate() + + +def _scheduler_is_empty(factory) -> bool: + scheduler = factory._scheduler + with scheduler._condition: + # One structure: the queue of executions with a record waiting. An + # execution's export bookkeeping lives on the per-invocation instance + # itself, which the SDK drops when the invocation scope exits, so an + # empty queue means the scheduler retains nothing. + return not scheduler._pending + + +def _force_drain(factory) -> None: + """Push everything this factory's instances scheduled out to the exporters. + + drain() takes the per-execution object, which in production is the plugin + instance itself. These tests present a bare ``_ExportState`` instead, exactly + as an execution that scheduled nothing of its own would: the flush it requests + is held back until every record pending when it was called has been exported. + """ + factory._scheduler.drain(_ExportState()) + + +class PinnedWorkerExporter: + """Blocks inside every ``export()`` until released, pinning the one worker.""" + + max_record_size_bytes = None + + def __init__(self) -> None: + self.entered = threading.Event() + self.release = threading.Event() + + def render(self, record: dict[str, Any]) -> Any: + return record + + def export(self, record: dict[str, Any]) -> None: + self.entered.set() + self.release.wait(30.0) + + def flush(self) -> None: + pass + + +def test_reentrant_finalizer_in_a_hook_does_not_deadlock(): + # _emit runs the scheduler's schedule() while holding the execution's lock, + # and schedule() releases the record it displaces inside that hold, on + # purpose: a record can carry customer objects whose finalizers run arbitrary + # code. A finalizer that re-enters a hook for the same execution therefore + # re-acquires that lock on the thread that already owns it, which a + # non-reentrant lock turns into a permanent hang of the invocation thread. + exporter = PinnedWorkerExporter() + reentered = threading.Event() + holder: dict[str, Any] = {} + + class ReentrantPayload: + """A customer object that reaches the record through a content transform.""" + + def __del__(self) -> None: + if reentered.is_set(): + return + reentered.set() + holder["plugin"].on_operation_change( + OperationChangeInfo( + execution_arn=ARN, + updated_operations=holder["ops"], + operations=holder["ops"], + ) + ) + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=lambda _value: ReentrantPayload()), + ) + ) + op = _step("s", op_id="1") + start = _start(operations={}) + plugin = factory.create_plugin(start) + holder["plugin"] = plugin + holder["ops"] = _ops(op) + change = OperationChangeInfo( + execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) + ) + try: + # The first emit pins the single worker inside export()... + plugin.on_invocation_start(start) + assert exporter.entered.wait(5.0) + # ...so this emit stays this execution's pending record... + plugin.on_operation_change(change) + + returned = threading.Event() + + def hook() -> None: + # ...and this one displaces it, releasing the displaced record (and + # running the payload's finalizer) on this thread, inside the lock. + plugin.on_operation_change(change) + returned.set() + + thread = threading.Thread(target=hook, daemon=True) + thread.start() + assert returned.wait(10.0), ( + "the hook never returned: a record finalizer that re-entered a hook " + "for the same execution deadlocked the invocation thread" + ) + assert reentered.is_set() # the finalizer really did re-enter a hook + finally: + exporter.release.set() + + +# -- late hooks after the invocation ended (closed gate) ---------------------- + + +def test_change_hook_reaching_the_lock_after_invocation_end_emits_nothing(): + # A change hook for a checkpoint that completed just before the invocation + # ended can reach the execution's lock while on_invocation_end still holds it. + # It must find the gate closed and emit nothing, so no RUNNING record can + # follow the terminal one. + exporter = ConcurrentCaptureExporter() + in_terminal_emit = threading.Event() + release_terminal = threading.Event() + + def blocking_output(value: Any) -> Any: + # Runs inside _emit, which runs inside the execution's lock, and only for + # a terminal record (a RUNNING record carries no output). + in_terminal_emit.set() + release_terminal.wait(10.0) + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(output=blocking_output), + ) + ) + op = _step("s", op_id="1") + plugin = _invocation(factory, _start(operations={})) + + end_returned = threading.Event() + + def end() -> None: + plugin.on_invocation_end(_end(operations=_ops(op))) + end_returned.set() + + end_thread = threading.Thread(target=end, daemon=True) + end_thread.start() + # on_invocation_end has closed the gate and is building the terminal record, + # still holding the execution's lock. + assert in_terminal_emit.wait(5.0) + + change_returned = threading.Event() + + def change() -> None: + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) + ) + ) + change_returned.set() + + change_thread = threading.Thread(target=change, daemon=True) + change_thread.start() + # It cannot get past the execution's lock while the end hook holds it. + assert not change_returned.wait(0.25) + + release_terminal.set() + end_thread.join(5.0) + change_thread.join(5.0) + assert end_returned.is_set() + assert change_returned.is_set() + _force_drain(factory) + + statuses = [record["status"] for record in exporter.snapshot()] + assert "SUCCEEDED" in statuses + # Nothing at all after the terminal record. + assert statuses[statuses.index("SUCCEEDED") + 1 :] == [] + + +def test_invocation_end_waits_for_an_in_flight_change_hook_emit(): + # The mirror interleaving: a change hook is already inside its emit, holding + # the execution's lock, when the invocation ends. Closing the gate and + # emitting the terminal record has to wait for it, otherwise the change hook + # finishes afterwards and appends a RUNNING record after the terminal one. + exporter = ConcurrentCaptureExporter() + in_change_emit = threading.Event() + release_change = threading.Event() + calls = itertools.count() + lock = threading.Lock() + + def blocking_input(value: Any) -> Any: + # Runs inside _emit for every record. Block only on the change hook's + # emit, which is the second one (invocation start emits the first). + with lock: + index = next(calls) + if index == 1: + in_change_emit.set() + release_change.wait(10.0) + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=blocking_input), + ) + ) + op = _step("s", op_id="1") + plugin = _invocation(factory, _start(operations={})) + + change_returned = threading.Event() + + def change() -> None: + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) + ) + ) + change_returned.set() + + change_thread = threading.Thread(target=change, daemon=True) + change_thread.start() + assert in_change_emit.wait(5.0) + + end_returned = threading.Event() + + def end() -> None: + plugin.on_invocation_end(_end(operations=_ops(op))) + end_returned.set() + + end_thread = threading.Thread(target=end, daemon=True) + end_thread.start() + # The terminal record cannot be emitted while the change hook holds the lock. + assert not end_returned.wait(0.25) + + release_change.set() + change_thread.join(5.0) + end_thread.join(10.0) + assert change_returned.is_set() + assert end_returned.is_set() + _force_drain(factory) + + statuses = [record["status"] for record in exporter.snapshot()] + assert "SUCCEEDED" in statuses + assert statuses[statuses.index("SUCCEEDED") + 1 :] == [] + + +def test_operation_change_after_invocation_end_emits_nothing(): + # A checkpoint that completed just before the invocation ended still delivers + # its operation-change hook. It reaches the instance for the invocation that + # just ended -- there is no registry it could recreate an entry in -- and must + # find the gate closed: no fabricated start time, and no RUNNING record after + # the terminal one. + exporter = CaptureExporter() + factory = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") + ) + op = _step("s", op_id="1") + plugin = _invocation(factory, _start(operations={})) + plugin.on_invocation_end(_end(operations=_ops(op))) + before = list(exporter.records) + assert before and before[-1]["status"] == "SUCCEEDED" + + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, updated_operations=_ops(op), operations=_ops(op) + ) + ) + _force_drain(factory) + + assert exporter.records == before # nothing emitted after the terminal record + assert [record["status"] for record in exporter.records][-1] == "SUCCEEDED" + # No fabricated start time: every record still reports the execution start. + assert {record["startTime"] for record in exporter.records} == { + "2026-01-01T00:00:00Z" + } + + +def test_invocation_start_after_the_gate_closed_changes_nothing(): + # Each invocation gets its own instance and exactly one invocation-start hook, + # so a start hook arriving on a closed instance is no longer reachable through + # a state registry -- there is none, and no instance can be handed to a second + # invocation. The `closed` gate is what holds that contract from the plugin's + # side: a start hook that arrives after the invocation ended must not re-seed + # the closed instance and must emit nothing. + exporter = CaptureExporter() + factory = workflow_insight( + WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") + ) + op = _step("s", op_id="1") + plugin = _invocation(factory, _start(operations={}, input_value="World")) + plugin.on_invocation_end(_end(operations=_ops(op))) + assert plugin._closed + + plugin.on_invocation_start( + _start( + operations=_ops(_step("late", op_id="2")), + input_value="late-input", + execution_start_time=T1, + ) + ) + _force_drain(factory) + + statuses = [record["status"] for record in exporter.records] + assert "SUCCEEDED" in statuses + # Nothing follows the terminal record... + assert statuses[statuses.index("SUCCEEDED") + 1 :] == [] + # ...and the closed instance was not re-seeded. A late start that got past the + # gate would adopt its own operation snapshot; the emission it would also have + # produced is stopped a second time by the re-check in _emit, so the adopted + # state is what pins the hook's own check. Input and start time are fixed by + # the constructor from the invocation's own start info, so no later hook can + # move them at all. + assert plugin._cached_input == "World" + assert plugin._start_time == T0 + assert [info.name for info in plugin._operations.values()] == ["s"] + + +def test_reentrant_invocation_end_stops_the_outer_running_record(): + # The gate at the top of each hook is a check-then-act, and _emit is the act. + # Between them _emit runs customer code while holding the execution's lock -- + # here a content transform -- and the lock is reentrant, so that customer code + # can run on_invocation_end to completion on this same thread: `_closed` set, + # terminal record scheduled and drained. The outer frame then resumes with a + # fully built RUNNING record, which must NOT reach the exporters after the + # terminal one. One hook call, one instance, no concurrency. + exporter = ConcurrentCaptureExporter() + holder: dict[str, Any] = {} + reentered = threading.Event() + + def reentering_input(value: Any) -> Any: + if not reentered.is_set(): + reentered.set() + holder["plugin"].on_invocation_end(_end(operations=_ops(_step("s")))) + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=reentering_input), + ) + ) + start = _start(operations={}) + plugin = factory.create_plugin(start) + holder["plugin"] = plugin + + # On a bounded thread, so a regression that makes the lock non-reentrant + # again fails here instead of hanging the suite. + returned = threading.Event() + + def hook() -> None: + plugin.on_invocation_start(start) + returned.set() + + thread = threading.Thread(target=hook, daemon=True) + thread.start() + assert returned.wait(10.0), ( + "the hook never returned: re-entering on_invocation_end from customer " + "code inside _emit deadlocked the invocation thread" + ) + thread.join(5.0) + assert not thread.is_alive() + assert reentered.is_set() # the re-entrant end hook really did run + # Force everything the plugin scheduled to reach the exporters, so a record + # that slipped past the gate is observed here rather than left pending. + _force_drain(factory) + + statuses = [record["status"] for record in exporter.snapshot()] + assert statuses == ["SUCCEEDED"], ( + "a non-terminal record reached the exporters after the terminal one for " + f"the same execution: {statuses}" + ) + assert _wait_until(lambda: not factory._scheduler._worker_alive()) + + +def test_a_reentrant_invocation_end_drains_after_the_lock_is_released(): + # The nested end hook used to drain from inside the outer frame's lock hold. A + # drain waits for the export worker, and an exporter that re-enters a hook + # blocks that worker on the very lock the waiting thread holds, so neither + # side can proceed and the invocation hangs until Lambda times it out. The + # drain a nested frame asks for is therefore deferred to the outermost hook + # frame, which runs it with every lock hold on this thread released. + exporter = ConcurrentCaptureExporter() + holder: dict[str, Any] = {} + reentered = threading.Event() + events: list[str] = [] + lock_free_at_drain: list[bool] = [] + + def reentering_input(value: Any) -> Any: + if not reentered.is_set(): + reentered.set() + holder["plugin"].on_invocation_end(_end(operations=_ops(_step("s")))) + events.append("nested-end-returned") + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=reentering_input), + ) + ) + start = _start(operations={}) + plugin = factory.create_plugin(start) + holder["plugin"] = plugin + + original_drain = factory._scheduler.drain + + def recording_drain(execution: Any) -> None: + events.append("drain") + lock_free_at_drain.append(_free_for_another_thread(plugin._lock)) + original_drain(execution) + + factory._scheduler.drain = recording_drain # type: ignore[method-assign] + + returned = threading.Event() + + def hook() -> None: + plugin.on_invocation_start(start) + returned.set() + + thread = threading.Thread(target=hook, daemon=True) + thread.start() + assert returned.wait(10.0), "the hook never returned" + thread.join(5.0) + assert not thread.is_alive() + assert reentered.is_set() + + assert events == ["nested-end-returned", "drain"], ( + "the nested end hook drained before its frame unwound, so the drain ran " + f"inside the outer lock hold: {events}" + ) + assert lock_free_at_drain == [True], ( + "the drain ran while this thread still held the execution's lock, which " + "an exporter re-entering a hook turns into a deadlock" + ) + + +def test_an_end_transform_that_fails_still_drains_what_was_scheduled(): + # `_emit` runs customer code between the snapshot and the hand-off: the content + # transforms, a result override, and __del__ on an object a displaced record + # carried. A failure there leaves the hook through the SDK's containment, and + # the drain used to be requested after the build, so that failure skipped it -- + # leaving records this execution had already scheduled in a buffering exporter + # when the environment froze. The drain is now requested on entry. + exporter = ConcurrentCaptureExporter() + calls: list[str] = [] + + def failing_on_the_second_call(value: Any) -> Any: + calls.append("input") + if len(calls) > 1: + # CancelledError rather than an ordinary exception on purpose: + # _apply_data_content contains Exception so that a failing redactor + # cannot leak the raw value, and a BaseException is what escapes the + # build and leaves the hook. + raise asyncio.CancelledError("transform cancelled") + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=failing_on_the_second_call), + ) + ) + start = _start(operations={}) + plugin = factory.create_plugin(start) + + # The first emit succeeds and schedules a RUNNING record. + plugin.on_invocation_start(start) + + # The terminal emit fails inside the transform, so this hook raises. The SDK + # contains that; here it is raised directly, which is the same code path. + with pytest.raises(asyncio.CancelledError): + plugin.on_invocation_end(_end(operations=_ops(_step("s")))) + + statuses = [record["status"] for record in exporter.snapshot()] + assert statuses == ["RUNNING"], ( + "the record scheduled before the failing transform must have been drained " + f"to the exporters, not left pending: {statuses}" + ) + assert exporter.flushes >= 1, "the drain must have flushed" + assert _wait_until(lambda: not factory._scheduler._worker_alive()) + + +def test_a_displaced_records_finalizer_runs_with_no_plugin_lock_held(): + # `_emit` hands the record to the scheduler while holding this instance's + # `_lock`, and that hand-off displaces the record already pending for this + # execution. Releasing the displaced one can run a customer `__del__` -- it is + # a snapshot of customer data -- and a finalizer that reaches ANOTHER + # execution's plugin blocks on that instance's lock. Two threads doing that to + # each other at once hang both invocations, which the reentrant lock does not + # help with: it only covers re-entry on the same instance. The hook frame + # therefore releases displaced records after every lock is dropped. + observed: list[bool] = [] + holder: dict[str, Any] = {} + + class _FinalizerProbe: + def __del__(self) -> None: + plugin = holder.get("plugin") + if plugin is None: + return + # `_is_owned()` answers "does the calling thread hold this lock", + # which is the question here. Acquiring it would not: an RLock lets + # its owner acquire it again. + observed.append(not plugin._lock._is_owned()) + + class _DiscardingExporter: + """Keeps no record, so the displaced one is the last reference to a probe. + + An exporter that retained records would decide this test by a race: if the + worker exported the first record before it was displaced, the probe stays + reachable from the exporter and no finalizer runs. + """ + + max_record_size_bytes = None + exports = 0 + + def render(self, record: dict[str, Any]) -> Any: + return record + + def export(self, record: dict[str, Any]) -> None: + type(self).exports += 1 + + def flush(self) -> None: + pass + + # A transform returning a fresh object per emit is what makes the record the + # only reference to it, so releasing the displaced record is what collects it. + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[_DiscardingExporter()], + emit_mode="on-change", + content=ContentConfig(input=lambda value: _FinalizerProbe()), + ) + ) + plugin = factory.create_plugin(_start(operations={})) + holder["plugin"] = plugin + + # The first emit queues a record holding a probe; the next displaces it, and + # the hook frame releases it on this thread. + plugin.on_invocation_start(_start(operations={})) + plugin.on_operation_change( + OperationChangeInfo( + execution_arn=ARN, + updated_operations=_ops(_step("s")), + operations=_ops(_step("s")), + ) + ) + plugin.on_invocation_end(_end(operations=_ops(_step("s")))) + gc.collect() + + assert observed, "the displaced record's finalizer must have run" + assert all(observed), ( + "a displaced record was released while this thread still held the " + "plugin's lock, which is what deadlocks two invocations whose finalizers " + "reach each other" + ) + assert _wait_until(lambda: not factory._scheduler._worker_alive()) + + +def _free_for_another_thread(lock: Any) -> bool: + """Report whether a lock is unheld, as seen from a thread that never took it. + + Asked from another thread on purpose: the lock is reentrant, so the thread + that owns it can always acquire it again and would learn nothing. + """ + acquired: list[bool] = [] + + def probe() -> None: + got = lock.acquire(blocking=False) + acquired.append(got) + if got: + lock.release() + + prober = threading.Thread(target=probe, daemon=True) + prober.start() + prober.join(5.0) + return acquired == [True] + + +# -- a build overtaken by one customer code started from inside it ------------- + + +class GatedCaptureExporter: + """Records every export; optionally blocks inside one execution's export. + + Blocking there pins the single export worker, so records scheduled while it + is held stay in their execution's pending slot and coalesce there. + """ + + max_record_size_bytes = None + + def __init__(self, hold_arn: str | None = None) -> None: + self._hold_arn = hold_arn + self.holding = threading.Event() + self.release = threading.Event() + self._lock = threading.Lock() + self._arrived = threading.Condition(self._lock) + self._records: list[dict[str, Any]] = [] + + def render(self, record: dict[str, Any]) -> Any: + return record + + def export(self, record: dict[str, Any]) -> None: + if self._hold_arn is not None and record["executionArn"] == self._hold_arn: + self.holding.set() + self.release.wait(30.0) + with self._arrived: + self._records.append(record) + self._arrived.notify_all() + + def flush(self) -> None: + pass + + def snapshot(self) -> list[dict[str, Any]]: + with self._lock: + return list(self._records) + + def wait_for(self, arn: str, count: int, timeout: float = 10.0) -> bool: + with self._arrived: + return self._arrived.wait_for( + lambda: ( + sum(1 for r in self._records if r["executionArn"] == arn) >= count + ), + timeout, + ) + + +def _records_for(exporter: GatedCaptureExporter, arn: str) -> list[dict[str, Any]]: + return [record for record in exporter.snapshot() if record["executionArn"] == arn] + + +def _force_drain_bounded(factory, timeout: float = 20.0) -> None: + """``_force_drain`` on a bounded thread, so a stalled drain fails the test. + + ``drain()`` parks on a condition with no deadline: that is correct in + production, where the only thing that can release it is the export worker, but + a regression that loses a record leaves it parked for good. Running it on + another thread turns that into an assertion failure instead of a hung suite. + """ + drained = threading.Event() + + def drain() -> None: + _force_drain(factory) + drained.set() + + thread = threading.Thread(target=drain, daemon=True) + thread.start() + assert drained.wait(timeout), "the drain never returned" + thread.join(5.0) + assert not thread.is_alive() + + +def _exported_operation_names( + exporter: GatedCaptureExporter, arn: str +) -> list[list[str]]: + return [ + [op["name"] for op in record["operations"]] + for record in _records_for(exporter, arn) + ] + + +def _newer_change() -> OperationChangeInfo: + """An operation-change carrying a strictly newer map than the drives below start from.""" + return OperationChangeInfo( + execution_arn=ARN, + updated_operations=_ops(_step("s1")), + operations=_ops(_step("s1")), + ) + + +def test_a_build_overtaken_by_a_nested_one_does_not_coalesce_the_newer_away(): + # _emit runs customer code -- here the content.input transform -- between the + # operation snapshot it takes and the hand-off to the scheduler, and the + # execution's lock is reentrant, so that code can run on_operation_change to + # completion on this same thread. The nested hook adopts a newer operation map + # and schedules its record first. The outer frame then hands over the older + # snapshot it built, and the pending slot takes whichever record arrives last, + # so without a comparison of build ages the newer record is coalesced away and + # the exporter only ever sees the older one. One hook call, one instance, no + # concurrency. + exporter = GatedCaptureExporter(hold_arn=ARN_B) + holder: dict[str, Any] = {} + reentered = threading.Event() + + def reentering_input(value: Any) -> Any: + # The primer execution below runs this same transform, so re-enter only + # once the instance under test has been published. + if holder.get("plugin") is None or reentered.is_set(): + return value + reentered.set() + holder["plugin"].on_operation_change(_newer_change()) + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=reentering_input), + ) + ) + # Pin the one export worker on another execution's record, so both records + # this execution builds are still in its pending slot when the outer frame + # hands its own over. + primer = factory.create_plugin(_start(arn=ARN_B)) + primer.on_invocation_start(_start(arn=ARN_B)) + assert exporter.holding.wait(10.0), "the export worker never reached the exporter" + + start = _start(operations={}) + plugin = factory.create_plugin(start) + holder["plugin"] = plugin + + # On a bounded thread, so a regression that deadlocks the hook fails here + # instead of hanging the suite. + returned = threading.Event() + + def hook() -> None: + plugin.on_invocation_start(start) + returned.set() + + thread = threading.Thread(target=hook, daemon=True) + thread.start() + assert returned.wait(10.0), "the hook never returned" + thread.join(5.0) + assert not thread.is_alive() + assert reentered.is_set() # the nested change hook really did run + exporter.release.set() + _force_drain_bounded(factory) + + assert _exported_operation_names(exporter, ARN) == [["s1"]], ( + "the newer snapshot the nested hook built was coalesced away by the " + "older one the outer frame built: " + f"{_exported_operation_names(exporter, ARN)}" + ) + + +def test_a_build_overtaken_by_a_nested_one_is_not_exported_after_it(): + # Same overtaking, with the newer record already handed to the exporter before + # the outer frame reaches the hand-off. Nothing coalesces, so without a + # comparison of build ages the exporter sees the newer snapshot and then the + # older one, and an exporter that upserts by execution ARN ends up storing the + # older state. + exporter = GatedCaptureExporter() + holder: dict[str, Any] = {} + reentered = threading.Event() + + def reentering_input(value: Any) -> Any: + if reentered.is_set(): + return value + reentered.set() + holder["plugin"].on_operation_change(_newer_change()) + # Wait for the newer record to reach the exporter, so this frame's older + # record cannot displace it in the pending slot and the two records are + # ordered at the exporter instead. + assert exporter.wait_for(ARN, 1), "the nested record never reached the exporter" + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=reentering_input), + ) + ) + start = _start(operations={}) + plugin = factory.create_plugin(start) + holder["plugin"] = plugin + + returned = threading.Event() + + def hook() -> None: + plugin.on_invocation_start(start) + returned.set() + + thread = threading.Thread(target=hook, daemon=True) + thread.start() + assert returned.wait(10.0), "the hook never returned" + thread.join(5.0) + assert not thread.is_alive() + assert reentered.is_set() + _force_drain_bounded(factory) + + assert _exported_operation_names(exporter, ARN) == [["s1"]], ( + "the exporter saw the older snapshot after the newer one for the same " + f"execution: {_exported_operation_names(exporter, ARN)}" + ) + + +def test_the_closing_record_is_exempt_from_the_build_age_check(): + # The closing record must reach the exporters whatever the build ages say. + # Customer code inside its build can start a newer non-terminal build, which + # would leave the closing record's age stale, and a checked hand-off would + # then drop it and leave a RUNNING snapshot as this execution's last exported + # state. Nothing else can rescue it: it is the last record this instance ever + # builds. + # + # The drive raises the build counter above the closing record's own first, by + # overtaking one build with a nested one exactly as the two tests above do, so + # a closing record that is age-checked against a counter it never incremented + # is dropped here. + exporter = GatedCaptureExporter() + holder: dict[str, Any] = {} + reentered = threading.Event() + + def reentering_input(value: Any) -> Any: + if reentered.is_set(): + return value + reentered.set() + holder["plugin"].on_operation_change(_newer_change()) + assert exporter.wait_for(ARN, 1), "the nested record never reached the exporter" + return value + + factory = workflow_insight( + WorkflowInsightConfig( + exporters=[exporter], + emit_mode="on-change", + content=ContentConfig(input=reentering_input), + ) + ) + start = _start(operations={}) + plugin = factory.create_plugin(start) + holder["plugin"] = plugin + + returned = threading.Event() + + def hook() -> None: + plugin.on_invocation_start(start) + # Push whatever the start hook handed over to the exporter before the end + # hook schedules the closing record, so a record the outer frame scheduled + # is observed here instead of being coalesced away by the closing one. + _force_drain(factory) + # on_invocation_end drains, so every record is at the exporter once this + # returns. + plugin.on_invocation_end(_end(operations=_ops(_step("s1"), _step("s2")))) + returned.set() + + thread = threading.Thread(target=hook, daemon=True) + thread.start() + assert returned.wait(10.0), "the hooks never returned" + thread.join(5.0) + assert not thread.is_alive() + assert reentered.is_set() + + records = _records_for(exporter, ARN) + assert [record["status"] for record in records] == ["RUNNING", "SUCCEEDED"], ( + "the closing record was dropped, or the superseded RUNNING record was " + f"exported: {[record['status'] for record in records]}" + ) + assert [op["name"] for op in records[0]["operations"]] == ["s1"] + assert sorted(op["name"] for op in records[1]["operations"]) == ["s1", "s2"] diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 5e78f2a9..6aa8d1ed 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -39,9 +39,15 @@ processors, and exporter. 1. Add the [ADOT Lambda Layer](#1-adot-lambda-layer) to your function and set `AWS_LAMBDA_EXEC_WRAPPER=/opt/otel-instrument` 2. Enable [X-Ray Active Tracing](#2-aws-x-ray-active-tracing) on the function -3. Pass `InvocationOtelPlugin` to your handler's `plugins` list +3. Pass `InvocationOtelPluginFactory()` to your handler's `plugins` list 4. Add X-Ray write permissions +The SDK's `plugins` list takes plugin *factories*, not plugin instances: it calls +each factory's `create_plugin` once per invocation and the plugin it returns +serves that one invocation. `InvocationOtelPluginFactory` and +`ExecutionOtelPluginFactory` are the factories for the two bundled plugins; each +takes the optional `OtelPluginConfig` that every plugin it builds will use. + Alternatively, install this package in the function artifact or a Lambda layer and select either OTel plugin by entry-point name: @@ -50,10 +56,11 @@ DURABLE_EXECUTION_PLUGINS=otel-invocation DURABLE_EXECUTION_PLUGINS=otel-execution ``` -`otel-invocation` creates `InvocationOtelPlugin`; `otel-execution` creates -`ExecutionOtelPlugin`. The SDK discovers the selected package entry point at -cold start, so the handler does not need to import or explicitly register the -plugin. +`otel-invocation` names a default-configured `InvocationOtelPluginFactory`; +`otel-execution` names a default-configured `ExecutionOtelPluginFactory`. The SDK +discovers the selected package entry point at cold start and calls its +`create_plugin` once per invocation, so the handler does not need to import or +explicitly register the plugin. ### 1. ADOT Lambda Layer @@ -161,10 +168,10 @@ lambda_.Function( ```python from aws_durable_execution_sdk_python import DurableContext from aws_durable_execution_sdk_python.execution import durable_execution -from aws_durable_execution_sdk_python_otel import InvocationOtelPlugin +from aws_durable_execution_sdk_python_otel import InvocationOtelPluginFactory -@durable_execution(plugins=[InvocationOtelPlugin()]) +@durable_execution(plugins=[InvocationOtelPluginFactory()]) def handler(event: dict, context: DurableContext) -> dict: result = context.step(lambda _: fetch_data(event["id"]), name="fetch-data") @@ -199,12 +206,12 @@ See the [ADOT sampling configuration](https://aws-otel.github.io/docs/getting-st ```python from aws_durable_execution_sdk_python_otel import ( - InvocationOtelPlugin, + InvocationOtelPluginFactory, OtelPluginConfig, xray_context_extractor, ) -plugin = InvocationOtelPlugin( +plugin_factory = InvocationOtelPluginFactory( OtelPluginConfig( # Use a custom context extractor (default: xray_context_extractor). context_extractor=xray_context_extractor, @@ -218,6 +225,9 @@ plugin = InvocationOtelPlugin( ) ``` +The config is resolved once and shared by every plugin the factory builds, so +configuration is per handler while plugin state is per invocation. + ### Context Extractors Context extractors return an `ExtractedContext` object, or `None` when no @@ -232,17 +242,19 @@ context: ```python from aws_durable_execution_sdk_python_otel import ( - InvocationOtelPlugin, + InvocationOtelPluginFactory, OtelPluginConfig, w3c_client_context_extractor, xray_context_extractor, ) # Default: X-Ray trace header (recommended for most Lambda deployments). -InvocationOtelPlugin(OtelPluginConfig(context_extractor=xray_context_extractor)) +InvocationOtelPluginFactory(OtelPluginConfig(context_extractor=xray_context_extractor)) # W3C Trace Context via clientContext (placeholder for backend propagation support). -InvocationOtelPlugin(OtelPluginConfig(context_extractor=w3c_client_context_extractor)) +InvocationOtelPluginFactory( + OtelPluginConfig(context_extractor=w3c_client_context_extractor) +) ``` Custom extractors should return `ExtractedContext`, not an OpenTelemetry @@ -300,15 +312,41 @@ each durable span in the same invocation. ### Log Correlation When `enrich_logger=True` (the default), the plugin installs a logging filter on -the root logger at invocation start. The filter stamps the active OTel trace -context onto every emitted log record using these attributes: +the root logger at invocation start. The filter stamps the trace context that is +active for the emitting invocation onto log records, using these attributes: - `traceId`: 32-char hex trace identifier - `spanId`: 16-char hex span identifier - `otelTraceSampled`: boolean indicating if the trace is sampled -These attributes are only set when a valid span context is active, so any log -formatter or schema must treat the fields as optional. +Which span a record carries follows the emitting thread: the operation attempt +span inside a step, the context span inside a child context, and the Invocation +span for top-level handler code. A span you start yourself on the execution +trace is used for records emitted inside it. A span the runtime already had +active when the invocation started — the Lambda invocation span the ADOT layer +creates under X-Ray active tracing — is the parent of the Invocation span rather +than a substitute for it, so a top-level record still names the Invocation span. +Correlation holds from the first statement of the handler, before any durable +operation, and holds when several invocations run concurrently in one +environment (Lambda Managed Instances): the plugin claims the invocation thread +at invocation start and the SDK runs the handler body in a copy of that thread's +context, so each record resolves to the invocation that emitted it rather than +to whichever invocation started last. A branch of a `map` or `parallel` runs on +a pool thread the invocation's context was not copied into, and the plugin +claims that thread from the hooks that run on it before your branch body does. + +Two cases are left unstamped, so any log formatter or schema must treat the +fields as optional: + +- No invocation is open — for example during environment initialization or + teardown. +- The record is emitted on a thread that carries no invocation claim. A thread + your code starts itself is such a thread, since Python does not copy context + into a new thread, as is the SDK's background checkpointing thread. The number + of invocations open in the environment does not change this: attributing an + unclaimed record to the single open invocation would be wrong whenever the + emitting thread belongs to a different invocation, and an uncorrelated record + is preferred over one carrying another execution's trace. ## Verification @@ -339,12 +377,15 @@ After deploying your function with the plugin configured: ## API Reference -### `InvocationOtelPlugin` +### `InvocationOtelPluginFactory` -Invocation-rooted view. Implements `DurableInstrumentationPlugin` from `aws_durable_execution_sdk_python`. +Factory for the invocation-rooted plugin, and what belongs in the SDK's `plugins` +list. Satisfies `DurableInstrumentationPluginFactory` from +`aws_durable_execution_sdk_python`: its `create_plugin(info)` takes an +`InvocationStartInfo` and returns the `InvocationOtelPlugin` for that invocation. ```python -InvocationOtelPlugin( +InvocationOtelPluginFactory( OtelPluginConfig( tracer_provider=None, context_extractor=None, @@ -356,13 +397,36 @@ InvocationOtelPlugin( ``` Pass `tracer_provider=...` when the application owns the OpenTelemetry SDK -provider. When omitted, the globally configured provider is used. +provider. When omitted, the globally configured provider is used, resolved per +invocation so a provider installed after the handler module is imported is still +picked up. + +`INVOCATION_OTEL_PLUGIN_FACTORY` is the default-configured instance the +`otel-invocation` entry point names. + +### `ExecutionOtelPluginFactory` + +Factory for the execution-rooted plugin, with the same construction and config as +`InvocationOtelPluginFactory`. `EXECUTION_OTEL_PLUGIN_FACTORY` is the +default-configured instance the `otel-execution` entry point names. + +### `InvocationOtelPlugin` + +Invocation-rooted view. Implements `DurableInstrumentationPlugin` from +`aws_durable_execution_sdk_python`. One instance serves exactly one invocation: +the SDK builds it from the factory before the first hook fires and drops it when +the invocation ends, so it holds its span registry and context tokens in ordinary +instance state. Construct it directly only when driving the hooks yourself; in a +handler, register the factory instead. ### `ExecutionOtelPlugin` Execution-rooted view. Uses the same execution ancestor and sampling behavior as `InvocationOtelPlugin`, but parents operation spans under Workflow and links -them to Invocation. +them to Invocation. Also one instance per invocation; the identities that must +agree across invocations (the trace ID, the Workflow span ID, and each +operation's span ID) are derived deterministically from the execution ARN rather +than carried in memory. ### `DeterministicIdGenerator` @@ -392,12 +456,15 @@ Structured trace context and sampling decision returned by context extractors. The logging filter (and its installer) used to stamp trace context onto log records. Installed automatically when `enrich_logger=True`; exported for manual -setups. +setups, where `install_log_filter(target_logger)` attaches it to a logger of your +choice. The filter carries no invocation identity of its own: it resolves the +invocation a record belongs to at emit time, so one filter serves every +invocation the environment runs, including concurrent ones. ## Requirements - Python >= 3.11 -- `aws-durable-execution-sdk-python` >= 2.0.0 +- `aws-durable-execution-sdk-python` >= 3.0.0 - An ADOT/community OpenTelemetry Lambda layer, or the `standalone` extra ## License diff --git a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml index cbc2f97a..574ff6fe 100644 --- a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml @@ -22,12 +22,19 @@ classifiers = [ "Programming Language :: Python :: Implementation :: PyPy", ] dependencies = [ - "aws-durable-execution-sdk-python>=2.0.0", + # >=3.0.0,<4: 3.0.0 is the first release whose `plugins` argument takes + # factories, and the ceiling is the same reasoning applied forwards -- the next + # core major that changes the plugin contract would install and then fail at + # handler initialization exactly as 2.x does below. + # DurableInstrumentationPluginProvider was removed in 3.0.0, and this package's + # entry points resolve to factories, so core 2.x accepts the install and then + # fails at handler initialization. + "aws-durable-execution-sdk-python>=3.0.0,<4", ] [project.entry-points."aws_durable_execution.plugins"] -otel-invocation = "aws_durable_execution_sdk_python_otel.plugin_provider:INVOCATION_OTEL_PLUGIN_PROVIDER" -otel-execution = "aws_durable_execution_sdk_python_otel.plugin_provider:EXECUTION_OTEL_PLUGIN_PROVIDER" +otel-invocation = "aws_durable_execution_sdk_python_otel.plugin_factory:INVOCATION_OTEL_PLUGIN_FACTORY" +otel-execution = "aws_durable_execution_sdk_python_otel.plugin_factory:EXECUTION_OTEL_PLUGIN_FACTORY" [project.optional-dependencies] # Lambda telemetry layers provide a version-aligned OpenTelemetry distribution. diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__about__.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__about__.py index fc21e413..82fa7f7e 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__about__.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__about__.py @@ -1,4 +1,4 @@ # SPDX-FileCopyrightText: 2025-present Amazon.com, Inc. or its affiliates. # # SPDX-License-Identifier: Apache-2.0 -__version__ = "1.0.0" +__version__ = "2.0.0" diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py index a8285dcb..774ae771 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py @@ -27,6 +27,12 @@ from aws_durable_execution_sdk_python_otel.invocation_plugin import ( InvocationOtelPlugin, ) +from aws_durable_execution_sdk_python_otel.plugin_factory import ( + EXECUTION_OTEL_PLUGIN_FACTORY, + INVOCATION_OTEL_PLUGIN_FACTORY, + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, +) from aws_durable_execution_sdk_python_otel.provider import ( ProviderResult, create_tracer_provider, @@ -35,12 +41,16 @@ __all__ = [ "__version__", + "EXECUTION_OTEL_PLUGIN_FACTORY", + "INVOCATION_OTEL_PLUGIN_FACTORY", "ContextExtractor", "DeterministicIdGenerator", "ExecutionOtelPlugin", + "ExecutionOtelPluginFactory", "ExtractedContext", "OtelPluginConfig", "InvocationOtelPlugin", + "InvocationOtelPluginFactory", "OtelContextLogFilter", "Sampling", "ProviderResult", diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py index c3d7afc1..1386dd2e 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py @@ -4,6 +4,7 @@ import contextvars import hashlib +import threading from collections.abc import Iterator from contextlib import contextmanager from dataclasses import dataclass @@ -23,6 +24,17 @@ class _IdOverride: span_id: int | None +# Serializes installation so two invocations binding to one tracer at the same +# time cannot both wrap its original generator. A TracerProvider caches tracers +# by instrumentation scope, so two plugins that ask for the same instrument name +# get the same tracer object. Without this lock both can read the original +# generator, both wrap it, and the second assignment replaces the first: the +# plugin that assigned first then holds a wrapper the tracer no longer uses, its +# deterministic overrides are never consulted, and its workflow and operation +# span IDs are random, which breaks cross-invocation stitching. +_install_lock = threading.Lock() + + def _to_otel_trace_id(execution_arn: str, start_timestamp: datetime) -> int: """Build a deterministic OTel-compatible execution trace ID (128 bits). @@ -123,16 +135,25 @@ def install_on_tracer(cls, tracer: SdkTracer) -> DeterministicIdGenerator: """Return the tracer's deterministic generator, installing one if needed. Installing on the plugin's tracer keeps unrelated instrumentation scopes - on the provider's original generator. Reusing an installed generator also - supports SDK versions that cache tracers by instrumentation scope. + on the provider's original generator. + + The returned generator is always the one the tracer holds. The check and + the install are one critical section, so a caller that arrives while + another is installing waits and then finds the installed generator + instead of wrapping the original a second time. A caller that acted on a + generator the tracer does not hold would set its deterministic overrides + somewhere the tracer never reads. Returning the installed generator also + supports SDK versions that cache tracers by instrumentation scope, which + is what makes two invocations share one tracer in the first place. """ - current_generator = tracer.id_generator - if isinstance(current_generator, cls): - return current_generator - - generator = cls(fallback_id_generator=current_generator) - tracer.id_generator = generator - return generator + with _install_lock: + current_generator = tracer.id_generator + if isinstance(current_generator, cls): + return current_generator + + generator = cls(fallback_id_generator=current_generator) + tracer.id_generator = generator + return generator @contextmanager def use_ids(self, *, trace_id: int | None, span_id: int | None) -> Iterator[None]: diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py index 9d2fa610..9689b341 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py @@ -4,6 +4,7 @@ import functools import inspect +import threading from dataclasses import dataclass from typing import Any @@ -24,6 +25,16 @@ ) +# Serializes installation so two invocations binding to one tracer at the same +# time cannot both wrap its original sampler. A TracerProvider caches tracers by +# instrumentation scope, so two plugins that ask for the same instrument name get +# the same tracer object. Without this lock both can read the original sampler, +# both wrap it, and the second assignment replaces the first: the plugin that +# assigned first then holds a wrapper the tracer no longer uses, and the delegate +# it took from that wrapper is not the delegate the tracer consults. +_install_lock = threading.Lock() + + @dataclass(frozen=True) class DurableSamplingIntent: """Sampling result to apply to each durable span in one invocation.""" @@ -39,12 +50,22 @@ def __init__(self, delegate: Sampler) -> None: @classmethod def install_on_tracer(cls, tracer: SdkTracer) -> "DurableSampler": - current_sampler = tracer.sampler - if isinstance(current_sampler, cls): - return current_sampler - sampler = cls(current_sampler) - tracer.sampler = sampler - return sampler + """Return the tracer's durable sampler, installing one if needed. + + The returned sampler is always the one the tracer holds. The check and the + install are one critical section, so a caller that arrives while another + is installing waits and then finds the installed sampler instead of + wrapping the original a second time. A caller that kept a sampler the + tracer does not hold would also keep a delegate the tracer never + consults. + """ + with _install_lock: + current_sampler = tracer.sampler + if isinstance(current_sampler, cls): + return current_sampler + sampler = cls(current_sampler) + tracer.sampler = sampler + return sampler def should_sample( self, diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py index 1a46cafc..4654a826 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -90,7 +90,11 @@ canonical_trace_id, ) from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig -from aws_durable_execution_sdk_python_otel.log_filter import install_log_filter +from aws_durable_execution_sdk_python_otel.log_filter import ( + bind_invocation, + install_log_filter, + unbind_invocation, +) from aws_durable_execution_sdk_python_otel.provider import create_tracer_provider @@ -117,6 +121,15 @@ def _to_otel_timestamp(dt: datetime.datetime | None) -> int | None: class ExecutionOtelPlugin(DurableInstrumentationPlugin): """OTel plugin that renders a durable execution as one Workflow-rooted trace. + Lifetime: one instance per invocation. The SDK builds it from + :class:`~aws_durable_execution_sdk_python_otel.plugin_factory.ExecutionOtelPluginFactory` + before the first hook fires and drops it when the invocation scope exits, so + every field below is per-invocation state that no other invocation can + observe. The execution-scoped identities the plugin needs across invocations + -- the canonical trace ID, the Workflow span ID and each operation's span ID + -- are derived deterministically from the execution ARN, so a fresh instance + rejoins the same trace without carrying anything over. + Args: config: Shared plugin configuration. When omitted, defaults are used (globally configured provider, X-Ray extractor, "Workflow" root @@ -139,7 +152,10 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._sampling_delegate: Sampler | None = None self._bind_sdk_tracer() - # Per-invocation state. + # Per-invocation state. The SDK builds one plugin instance per + # invocation through ExecutionOtelPluginFactory and drops it when the + # invocation scope exits, so these are ordinary instance fields that + # never have to be cleared for reuse. self._execution_arn = "" self._execution_trace_id: int | None = None self._execution_start_time: datetime.datetime | None = None @@ -148,6 +164,10 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._sampling_intent: DurableSamplingIntent | None = None self._workflow_span: Span | None = None self._invocation_span: Span | None = None + # The span that was already current when this invocation's body began, + # recorded by _record_enclosing_span. Used only to resolve log + # correlation; see get_current_span_context. + self._enclosing_span_context: SpanContext | None = None self._operation_spans: dict[str, Span] = {} # CONTEXT operations that emitted a durable START hook this invocation. # A context absent from this set is checkpointless (for example, a FLAT @@ -165,7 +185,12 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._tracing_enabled = False if self._config.enrich_logger: - install_log_filter(self) + # Install the root-logger filter so every log record is stamped with + # the active span context. On a warm environment the handler already + # carries the filter a previous invocation installed and it is reused + # as is: the filter holds no invocation identity, and this plugin + # claims the invocation in on_invocation_start instead. + install_log_filter() def _bind_sdk_tracer(self) -> bool: """Bind to an SDK tracer, retrying a deferred global provider.""" @@ -267,11 +292,10 @@ def _detach_context(self, key: str) -> None: otel_context.detach(token) # type: ignore[arg-type] def _detach_remaining_contexts(self) -> None: - """Release scopes still open, newest first, so nothing outlives the plugin. + """Release scopes still open, newest first, so nothing outlives the invocation. Reached when a lifecycle end hook never fires -- for example a user - function that suspends, or a warm invocation that starts before the - previous one was cleaned up. + function that suspends. """ with self._lock: keys = list(reversed(self._context_tokens)) @@ -279,9 +303,21 @@ def _detach_remaining_contexts(self) -> None: self._detach_context(key) def get_current_span_context(self) -> SpanContext | None: - """Return the active span context for log correlation (see log_filter).""" + """Return the active span context for log correlation (see log_filter). + + A span that became current *inside* this invocation wins: the attempt + span inside a step, the context span inside a child context, or a span + the handler body starts itself. The span that enclosed this invocation is + excluded, so top-level handler records resolve to the Invocation span + rather than to the Workflow span this plugin makes current at invocation + start, which the SDK carries into the thread running the handler body. + """ span_context = trace.get_current_span().get_span_context() - if span_context and span_context.is_valid: + if ( + span_context + and span_context.is_valid + and not self._is_enclosing_span(span_context) + ): return span_context for candidate in (self._invocation_span, self._workflow_span): if candidate is not None: @@ -290,6 +326,26 @@ def get_current_span_context(self) -> SpanContext | None: return ctx return None + def _record_enclosing_span(self) -> None: + """Record the span that is current now, as this invocation's body begins. + + Called at the end of ``on_invocation_start``, on the invocation thread, + which is the context the SDK copies into the thread that runs the handler + body. That is the Workflow span this plugin just attached, so it is less + specific than the Invocation span for a top-level record. + """ + span_context = trace.get_current_span().get_span_context() + self._enclosing_span_context = span_context if span_context.is_valid else None + + def _is_enclosing_span(self, span_context: SpanContext) -> bool: + """Whether ``span_context`` is the span that enclosed this invocation.""" + enclosing = self._enclosing_span_context + return ( + enclosing is not None + and enclosing.trace_id == span_context.trace_id + and enclosing.span_id == span_context.span_id + ) + # ------------------------------------------------------------------ # Links # ------------------------------------------------------------------ @@ -413,7 +469,12 @@ def _with_sampling(self, parent_context: Context) -> Context: # ------------------------------------------------------------------ def on_invocation_start(self, info: InvocationStartInfo) -> None: logger.debug("Durable invocation started: %s", info) - self._reset_state() + # Claim log correlation for this invocation before anything can fail + # below: the claim is what keeps a concurrent invocation's records off + # this invocation's trace, and it is registered even when tracing turns + # out to be disabled so that the filter can still tell how many + # invocations are open. + bind_invocation(self) if info.execution_start_time is None: logger.warning( "ExecutionOtelPlugin requires InvocationStartInfo.execution_start_time " @@ -477,8 +538,8 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: # Make the Workflow span the active span so auto-instrumented spans # created during the invocation become its children. The token is - # released in _reset_state at invocation end, restoring the context that - # was active before the invocation started. + # released in _release_invocation_scope at invocation end, restoring the + # context that was active before the invocation started. if self._workflow_span is not None: self._attach_context( _INVOCATION_CONTEXT_KEY, @@ -487,6 +548,9 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: ), ) + # Last, so that everything this hook makes current is accounted for. + self._record_enclosing_span() + def _start_workflow_span(self, info: InvocationStartInfo) -> None: """Install a non-recording placeholder for the execution-scoped Workflow span. @@ -590,11 +654,28 @@ def _start_invocation_span(self, info: InvocationStartInfo) -> None: self._set_span(_INVOCATION_KEY, self._invocation_span) def on_invocation_end(self, info: InvocationEndInfo) -> None: + """End this invocation's spans, then release its scope and flush. + + Ending a span and exporting the Workflow span call the configured tracer + and span processors, which are customer-supplied and can raise. The + invocation must still be released: until it is, the log filter counts it + as open and the OTel scope this plugin attached at invocation start stays + current on a warm environment's thread. The release and the flush + therefore run in a ``finally``. + """ logger.debug("Durable invocation ended: %s", info) if not self._tracing_enabled: - self._reset_state() + self._release_invocation_scope() return + try: + self._end_invocation_spans(info) + finally: + self._release_invocation_scope() + self._force_flush() + + def _end_invocation_spans(self, info: InvocationEndInfo) -> None: + """End this invocation's open spans and export the Workflow span.""" # End the invocation span regardless of terminal status. Record the # invocation status and map it to a span status: # SUCCEEDED/PENDING -> OK (this invocation did its work, whether it @@ -628,28 +709,46 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if info.status in _TERMINAL_INVOCATION_STATUSES: self._export_workflow_span(info) - self._reset_state() + def _force_flush(self) -> None: + """Flush pending spans, containing any error the flush raises. - if hasattr(self._provider, "force_flush"): - try: - self._provider.force_flush() - except Exception: # noqa: BLE001 - logger.exception("force_flush failed at invocation end") - - def _reset_state(self) -> None: + A flush calls the configured span processors and exporters, which are + customer-supplied and can raise. This runs in a ``finally`` block, where + an escaping exception would replace the exception that ended the + invocation and hide its cause, so the error is logged and dropped. + """ + if not hasattr(self._provider, "force_flush"): + return + try: + self._provider.force_flush() + except Exception: # noqa: BLE001 + logger.exception("force_flush failed at invocation end") + + def _release_invocation_scope(self) -> None: + """Release what this invocation attached, and stop instrumenting. + + Not a state reset. The instance serves exactly one invocation and is + dropped afterwards, so its fields never have to be cleared for a warm + environment's next invocation. Two things still have to happen at the + invocation boundary: + + * The OpenTelemetry context stack belongs to the thread, not to the + plugin, so any scope this plugin attached and did not release must be + detached here or it would stay current on a warm environment's thread + after the invocation returns. + * The log filter's record of open invocations is process-global, so this + invocation must be removed from it. A thread this invocation claimed + keeps that claim, because a context variable can only be reset by the + thread that set it, so until the invocation is removed a record emitted + on such a thread is still correlated to this finished invocation's + spans. + * ``_tracing_enabled`` is cleared so a hook that arrives after the + invocation end -- one dispatched off the checkpointing path, for + instance -- cannot start a span after the invocation span was ended and + the provider flushed. + """ self._detach_remaining_contexts() - self._execution_arn = "" - self._execution_trace_id = None - self._extracted_context = None - self._execution_trace_context = None - self._sampling_intent = None - self._execution_start_time = None - self._workflow_span = None - self._invocation_span = None - with self._lock: - self._operation_spans = {} - self._checkpointed_context_ids = set() - self._ended_operation_ids = set() + unbind_invocation(self) self._tracing_enabled = False # ------------------------------------------------------------------ @@ -659,6 +758,13 @@ def on_operation_start(self, info: OperationStartInfo) -> None: logger.debug("Durable operation started: %s", info) if not self._tracing_enabled: return + # Runs on the thread that drives the durable operation. The thread + # running the handler body already carries this invocation's claim, + # propagated from the invocation thread, but a branch of a map or + # parallel runs on a pool thread that does not, so claim it here. + # Claimed after the tracing-enabled gate, so a hook arriving after the + # invocation ended cannot re-register a finished invocation. + bind_invocation(self) if info.operation_type is OperationType.CONTEXT: with self._lock: self._checkpointed_context_ids.add(info.operation_id) @@ -791,6 +897,9 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: logger.debug("Durable user function started: %s", info) if not self._tracing_enabled: return + # Runs on the thread executing user code -- a parallel branch runs on its + # own thread -- so claim that thread for this invocation as well. + bind_invocation(self) if info.operation_type not in (OperationType.CONTEXT, OperationType.STEP): raise RuntimeError( "on_user_function_start only supports CONTEXT and STEP operations" diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index 012f17f2..923733f8 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -58,7 +58,11 @@ ExecutionTraceContext, canonical_trace_id, ) -from aws_durable_execution_sdk_python_otel.log_filter import install_log_filter +from aws_durable_execution_sdk_python_otel.log_filter import ( + bind_invocation, + install_log_filter, + unbind_invocation, +) from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig from aws_durable_execution_sdk_python_otel.provider import create_tracer_provider @@ -96,6 +100,12 @@ class InvocationOtelPlugin(DurableInstrumentationPlugin): use newly generated span IDs. Operation attributes and links to the Workflow span provide execution-scoped correlation across invocations. + Lifetime: one instance per invocation. The SDK builds it from + :class:`~aws_durable_execution_sdk_python_otel.plugin_factory.InvocationOtelPluginFactory` + before the first hook fires and drops it when the invocation scope exits, so + every field below is per-invocation state that no other invocation can + observe. + Args: config: Shared plugin configuration (the same OtelPluginConfig accepted by ExecutionOtelPlugin). When omitted, defaults are used (X-Ray @@ -120,7 +130,10 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: When ``enrich_logger`` is enabled (default), the plugin installs a logging filter that stamps the active OTel trace context onto every - emitted log record. + emitted log record. The filter is installed on the root logger's + handlers, which outlive this instance, so installing rebinds an + already-present filter to this invocation's plugin rather than stacking a + second one. """ self._config = config or OtelPluginConfig() self._context_extractor: ContextExtractor = ( @@ -137,7 +150,10 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._sampling_delegate: Sampler | None = None self._bind_sdk_tracer() - # per invocation status: + # Per-invocation state. The SDK builds one plugin instance per + # invocation through InvocationOtelPluginFactory and drops it when the + # invocation scope exits, so these are ordinary instance fields that + # never have to be cleared for reuse. self._execution_arn = "" self._execution_trace_id: int | None = None self._execution_start_time: datetime.datetime | None = None @@ -146,6 +162,10 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._sampling_intent: DurableSamplingIntent | None = None self._workflow_span: Span | None = None self._span_time_floor_ns: int | None = None + # The span that was already current when this invocation's body began, + # recorded by _record_enclosing_span. Used only to resolve log + # correlation; see get_current_span_context. + self._enclosing_span_context: SpanContext | None = None # Maps operation ID (None for root) to the active span. self._operation_spans: dict[str | None, Span] = {} # Replay state supplied by CONTEXT operation START hooks. Missing @@ -166,8 +186,12 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # Install the root-logger filter so every log record is stamped with # the active span context. The Lambda runtime attaches its root # handler before the handler module is imported (and thus before the - # plugin is constructed), so the handlers are available here. - install_log_filter(self) + # plugin is constructed), so the handlers are available here. On a + # warm environment the handler already carries the filter a previous + # invocation installed and is reused as is: the filter holds no + # invocation identity, and this plugin claims the invocation in + # on_invocation_start instead. + install_log_filter() def _bind_sdk_tracer(self) -> bool: """Bind to an SDK tracer, retrying a deferred global provider.""" @@ -266,11 +290,10 @@ def _detach_context(self, key: str) -> None: context.detach(token) # type: ignore[arg-type] def _detach_remaining_contexts(self) -> None: - """Release scopes still open, newest first, so nothing outlives the plugin. + """Release scopes still open, newest first, so nothing outlives the invocation. Reached when a lifecycle end hook never fires -- for example a user - function that suspends, or a warm invocation that starts before the - previous one was cleaned up. + function that suspends. """ with self._operation_spans_lock: keys = list(reversed(self._context_tokens)) @@ -281,11 +304,13 @@ def get_current_span_context(self) -> SpanContext | None: """Return the span context to use for log correlation. Resolution order: - 1. The same-trace span attached to the OTel thread-local context. + 1. A same-trace span that became current *inside* this invocation. Inside a step this is the active attempt span, and inside a child context this is the active context span (attached in - on_user_function_start). Unrelated ambient spans are ignored so logs - stay correlated to the durable execution trace. + on_user_function_start); a span the handler body starts itself also + lands here. Such a span is more specific than the Invocation span, so + it wins. Unrelated ambient spans are ignored so logs stay correlated + to the durable execution trace. 2. The invocation span from the plugin registry. This is the path used for top-level handler code: the invocation span is never attached to the worker thread's context, so the registry is the only way to @@ -293,6 +318,16 @@ def get_current_span_context(self) -> SpanContext | None: detaching the operation scope restores a context with no durable span. + The span that enclosed this invocation is deliberately excluded from + step 1. Under X-Ray active tracing with the ADOT layer, the layer's + Lambda invocation span is current before this invocation starts and is + on the execution trace, and the SDK carries it into the thread running + the handler body. It is also the parent of this plugin's Invocation + span, so preferring it would point top-level records one level up the + tree from the invocation they were emitted by. Anything that becomes + current after the enclosing span was recorded is inside the invocation + and still takes precedence. + Returns: A valid SpanContext, or None if no span is active. """ @@ -301,6 +336,7 @@ def get_current_span_context(self) -> SpanContext | None: span_context and span_context.is_valid and span_context.trace_id == self._execution_trace_id + and not self._is_enclosing_span(span_context) ): return span_context @@ -312,6 +348,28 @@ def get_current_span_context(self) -> SpanContext | None: return None + def _record_enclosing_span(self) -> None: + """Record the span that is current now, as this invocation's body begins. + + Called at the end of ``on_invocation_start``, on the invocation thread, + which is the context the SDK copies into the thread that runs the + handler body. Whatever span is current at that moment existed before the + invocation did -- the ADOT layer's Lambda invocation span, in the X-Ray + active tracing shape -- and is therefore less specific than this + plugin's own Invocation span. + """ + span_context = trace.get_current_span().get_span_context() + self._enclosing_span_context = span_context if span_context.is_valid else None + + def _is_enclosing_span(self, span_context: SpanContext) -> bool: + """Whether ``span_context`` is the span that enclosed this invocation.""" + enclosing = self._enclosing_span_context + return ( + enclosing is not None + and enclosing.trace_id == span_context.trace_id + and enclosing.span_id == span_context.span_id + ) + # ------------------------------------------------------------------ # Context resolution # ------------------------------------------------------------------ @@ -513,7 +571,12 @@ def _end_span( def on_invocation_start(self, info: InvocationStartInfo) -> None: """Called at the start of each invocation. Creates the invocation span.""" logger.debug("Durable invocation started: %s", info) - self._reset_state() + # Claim log correlation for this invocation before anything can fail + # below: the claim is what keeps a concurrent invocation's records off + # this invocation's trace, and it is registered even when tracing turns + # out to be disabled so that the filter can still tell how many + # invocations are open. + bind_invocation(self) if info.execution_start_time is None: logger.warning( "InvocationOtelPlugin requires InvocationStartInfo.execution_start_time " @@ -579,6 +642,9 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: attributes=self._extract_attributes(info), ) + # Last, so that everything this hook makes current is accounted for. + self._record_enclosing_span() + def _start_workflow_span(self, info: InvocationStartInfo) -> None: """Install a non-recording placeholder for the execution-scoped Workflow span. @@ -649,12 +715,29 @@ def _export_workflow_span(self, info: InvocationEndInfo) -> None: workflow_span.end() def on_invocation_end(self, info: InvocationEndInfo) -> None: - """Called at the end of each invocation. Ends the invocation span and flushes.""" + """Called at the end of each invocation. Ends the invocation span and flushes. + + Ending a span and exporting the Workflow span call the configured tracer + and span processors, which are customer-supplied and can raise. The + invocation must still be released: until it is, the log filter counts it + as open and any OTel scope this plugin attached stays current on a warm + environment's thread. The release and the flush therefore run in a + ``finally``. + """ logger.debug("Durable invocation ended: %s", info) if not self._tracing_enabled: - self._reset_state() + self._release_invocation_scope() return + try: + self._end_invocation_spans(info) + finally: + self._release_invocation_scope() + # Flush before Lambda freeze. + self._force_flush() + + def _end_invocation_spans(self, info: InvocationEndInfo) -> None: + """End this invocation's open spans and export the Workflow span.""" # Spans are registered parent-first, so close pending spans in reverse # order to keep every child contained within its parent. with self._operation_spans_lock: @@ -700,27 +783,46 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if info.status in _TERMINAL_INVOCATION_STATUSES: self._export_workflow_span(info) - self._reset_state() + def _force_flush(self) -> None: + """Flush pending spans, containing any error the flush raises. - # Flush before Lambda freeze - if hasattr(self._provider, "force_flush"): + A flush calls the configured span processors and exporters, which are + customer-supplied and can raise. This runs in a ``finally`` block, where + an escaping exception would replace the exception that ended the + invocation and hide its cause, so the error is logged and dropped. + """ + if not hasattr(self._provider, "force_flush"): + return + try: self._provider.force_flush() - - def _reset_state(self) -> None: - """Clear per-invocation state for warm Lambda environment reuse.""" + except Exception: # noqa: BLE001 + logger.exception("force_flush failed at invocation end") + + def _release_invocation_scope(self) -> None: + """Release what this invocation attached, and stop instrumenting. + + Not a state reset. The instance serves exactly one invocation and is + dropped afterwards, so its fields never have to be cleared for a warm + environment's next invocation. Two things still have to happen at the + invocation boundary: + + * The OpenTelemetry context stack belongs to the thread, not to the + plugin, so any scope this plugin attached and did not release must be + detached here or it would stay current on a warm environment's thread + after the invocation returns. + * The log filter's record of open invocations is process-global, so this + invocation must be removed from it. A thread this invocation claimed + keeps that claim, because a context variable can only be reset by the + thread that set it, so until the invocation is removed a record emitted + on such a thread is still correlated to this finished invocation's + spans. + * ``_tracing_enabled`` is cleared so a hook that arrives after the + invocation end -- one dispatched off the checkpointing path, for + instance -- cannot start a span after the invocation span was ended and + the provider flushed. + """ self._detach_remaining_contexts() - self._execution_arn = "" - self._execution_trace_id = None - self._extracted_context = None - self._execution_trace_context = None - self._sampling_intent = None - self._execution_start_time = None - self._workflow_span = None - self._span_time_floor_ns = None - with self._operation_spans_lock: - self._operation_spans = {} - self._context_operation_replays = {} - self._incomplete_attempt_span_keys = set() + unbind_invocation(self) self._tracing_enabled = False def on_operation_start(self, info: OperationStartInfo) -> None: @@ -728,6 +830,13 @@ def on_operation_start(self, info: OperationStartInfo) -> None: logger.debug("Durable operation started: %s", info) if not self._tracing_enabled: return + # Runs on the thread that drives the durable operation. The thread + # running the handler body already carries this invocation's claim, + # propagated from the invocation thread, but a branch of a map or + # parallel runs on a pool thread that does not, so claim it here. + # Claimed after the tracing-enabled gate, so a hook arriving after the + # invocation ended cannot re-register a finished invocation. + bind_invocation(self) if info.operation_type is OperationType.CONTEXT: # The user-function hook owns the span, but this durable START hook # distinguishes checkpoint-backed contexts from virtual branches. @@ -800,6 +909,9 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: logger.debug("Durable user function started: %s", info) if not self._tracing_enabled: return + # Runs on the thread executing user code -- a parallel branch runs on its + # own thread -- so claim that thread for this invocation as well. + bind_invocation(self) # Context and Step operations are tracked using on_user_function_start if info.operation_type not in [OperationType.CONTEXT, OperationType.STEP]: raise RuntimeError( diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/log_filter.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/log_filter.py index 7e11ebc9..afc07060 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/log_filter.py @@ -13,11 +13,67 @@ These attributes are only set when a valid span context is active. Records emitted outside an active invocation (e.g. during Lambda teardown) pass through unmodified, so any log formatter or schema must treat the fields as optional. + +Resolving *which* invocation a record belongs to +----------------------------------------------- + +A logging handler is process-global and outlives every invocation, while a +plugin instance serves exactly one invocation. Concurrent executions in one +environment (Lambda Managed Instances) therefore have several plugin instances +alive at once, all reachable from one installed filter. The filter must not hold +a single mutable reference to "the" plugin: whichever invocation started last +would win, and every other invocation's records would be stamped with its trace. + +So the binding is per invocation, not per filter: + + - ``bind_invocation`` marks an invocation open and claims the calling + thread/task for it, through a :class:`contextvars.ContextVar`. A record + emitted on a claimed thread resolves to the invocation that claimed it, + which is per-thread and per-task and so cannot be overwritten by a + concurrent invocation. The claim is a weak reference, so a pool thread + that is never used again cannot pin a finished invocation's plugin. + - ``unbind_invocation`` marks the invocation closed. + - The claim reaches the thread running the handler body because the SDK + submits that work with a copy of the invocation thread's context, taken + after the invocation-start hook has run. Records from top-level handler + code therefore resolve to their own invocation, including the first + statement of the handler, before any durable operation has claimed the + thread directly. + - A branch of a map or parallel runs on a pool thread the invocation's + context was not copied into, so the plugins claim that thread from the + hooks that run on it before the branch body does. + - A record emitted on a thread that carries no live claim is left unstamped, + whatever the number of open invocations. Two threads carry no claim: the + SDK's background checkpointing thread, and a thread customer code starts + itself, since Python does not copy context into a new thread. Correlation + is lost for those records. + +Resolving an unclaimed record against the single open invocation, when exactly +one is open, would be wrong in two orderings. A thread claimed by invocation A +keeps that claim after A ends, because a :class:`contextvars.ContextVar` can only +be reset by the thread that set it, so a record A's thread emits while B is the +only open invocation would be stamped with B's trace. A record emitted on +invocation B's own thread before B reaches its invocation-start hook carries no +claim at all, so while A is the only open invocation it would be stamped with A's +trace. Both stamp one customer's execution onto another's, which is a larger +defect than an unattributed record, so a live claim is required and the count of +open invocations is never consulted. + +Reading the active span straight from the OTel context (as the Java plugin's +static ``MdcSpanEnricher`` does) is not sufficient here: the invocation span is +never attached to the OTel context, so a record emitted between durable +operations would find no durable span current and silently lose correlation. The +plugin's ``get_current_span_context()`` still reads the OTel context first, so +records emitted inside a step or child context resolve to the active operation +span exactly as before. """ from __future__ import annotations +import contextvars import logging +import threading +import weakref from typing import TYPE_CHECKING, Protocol from opentelemetry.trace import TraceFlags @@ -38,30 +94,132 @@ class _SpanContextProvider(Protocol): def get_current_span_context(self) -> SpanContext | None: ... +# Guards the open-invocation registry. Held for the length of a set membership +# test or a single mutation, never while a plugin is called. +_registry_lock = threading.Lock() + +# Invocations that have started and not yet ended. Weak so that a plugin whose +# end hook never ran (a process torn down mid-invocation) cannot keep itself, +# and the spans it holds, alive for the life of the environment. +_open_invocations: weakref.WeakSet[_SpanContextProvider] = weakref.WeakSet() + +# The invocation owning the current thread/task. Set by bind_invocation on every +# thread the owning plugin is given control on. +# +# A weak reference, because a claim outlives the invocation that made it on every +# thread except the one that ends it: a ContextVar can only be reset by the +# thread that set it, so a pool thread that is never used again keeps whatever +# the claim holds. A strong claim would therefore pin a finished invocation's +# plugin, its spans and its context tokens for the life of the execution +# environment, and would also defeat _open_invocations being weak, since a +# plugin whose end hook never ran would stay reachable through the claim. While +# an invocation is open the SDK holds its plugin, which is what keeps the +# referent alive for every record the filter resolves. +_current_invocation: contextvars.ContextVar[ + weakref.ref[_SpanContextProvider] | None +] = contextvars.ContextVar("durable_execution_otel_invocation", default=None) + +# Serializes installation so two invocations starting at once cannot both find +# a handler filterless and both add a filter to it. +_install_lock = threading.Lock() + + +def bind_invocation(provider: _SpanContextProvider) -> None: + """Mark ``provider``'s invocation open and claim this thread/task for it. + + Called by a plugin when it takes control on a thread: at invocation start on + the Lambda handler thread, and again from the hooks that run on the threads + executing user code. Idempotent, so a plugin can call it from every such + hook without tracking which threads it has already claimed. + + A thread that already carries this provider's claim returns before taking the + registry lock. Every operation-start and user-function-start hook calls this, + and after the first call on a thread there is nothing to add: membership in + the open-invocation set is idempotent, and the claim is already in place. + + Args: + provider: The plugin serving the invocation that owns this thread. + """ + claim = _current_invocation.get() + if claim is not None and claim() is provider: + return + with _registry_lock: + _open_invocations.add(provider) + _current_invocation.set(weakref.ref(provider)) + + +def unbind_invocation(provider: _SpanContextProvider) -> None: + """Mark ``provider``'s invocation closed and release its claim on this thread. + + Claims made on other threads are not released here -- a context can only be + reset from the thread that set it -- so :func:`_resolve_provider` also checks + that a claim names a still-open invocation. That check is what keeps a + pooled thread outliving its invocation from correlating a later record to a + finished one. + + Args: + provider: The plugin whose invocation has ended. + """ + with _registry_lock: + _open_invocations.discard(provider) + claim = _current_invocation.get() + if claim is not None and claim() is provider: + _current_invocation.set(None) + + +def _resolve_provider() -> _SpanContextProvider | None: + """Return the invocation to correlate a record emitted right here against. + + A claim is only honoured while the invocation naming it is still open. A + :class:`contextvars.ContextVar` can only be reset by the thread that set it, + so ``unbind_invocation`` leaves the claim in place on every other thread the + invocation claimed; without the liveness check a pooled thread would keep + correlating records to a finished invocation. + + A thread with no live claim resolves to nothing, even when exactly one + invocation is open. Deciding by count would attribute the record to that + invocation, and two orderings make that the wrong one: a thread still + carrying a finished invocation's claim, and an invocation's own thread that + has not reached its invocation-start hook yet. + + A claim whose referent has been collected resolves to nothing as well. The + claim is weak, so a plugin the SDK has released can be gone while the claim + that named it remains on a pool thread. A collected referent means the + invocation is over, which is the same answer the liveness check gives. + """ + claim = _current_invocation.get() + if claim is None: + return None + claimed = claim() + if claimed is None: + return None + with _registry_lock: + if claimed in _open_invocations: + return claimed + return None + + class OtelContextLogFilter(logging.Filter): """Logging filter that injects the active OTel span context onto records. - The filter is a pure reader of the plugin's current span context. It - resolves the span at emit time, on the thread that emits the record, via - ``plugin.get_current_span_context()``. That method returns the active - operation span inside steps and child contexts (attached to the worker - thread's OTel context) and falls back to the invocation span for top-level - handler code. + The filter holds no state: it resolves the invocation and the span at emit + time, on the thread that emits the record, so one installed filter serves + any number of concurrent invocations. Resolution is described in the module + docstring; the span itself comes from that invocation's + ``get_current_span_context()``, which returns the active operation span + inside steps and child contexts and falls back to the invocation span for + top-level handler code. The filter never caches identifiers and always returns ``True`` so it never drops a record. - - Args: - plugin: The OTel plugin instance that resolves the current span context. """ - def __init__(self, plugin: _SpanContextProvider) -> None: - super().__init__() - self._plugin = plugin - def filter(self, record: logging.LogRecord) -> bool: """Stamp the active span context onto the record, then allow it through.""" - span_context = self._plugin.get_current_span_context() + provider = _resolve_provider() + if provider is None: + return True + span_context = provider.get_current_span_context() if span_context and span_context.is_valid: record.traceId = format(span_context.trace_id, "032x") record.spanId = format(span_context.span_id, "016x") @@ -72,7 +230,6 @@ def filter(self, record: logging.LogRecord) -> bool: def install_log_filter( - plugin: _SpanContextProvider, target_logger: logging.Logger | None = None, ) -> OtelContextLogFilter | None: """Attach an OtelContextLogFilter to a logger's handlers, idempotently. @@ -82,13 +239,14 @@ def install_log_filter( records propagated from child loggers are also enriched, since handler filters run for every record reaching the handler. - This is safe to call on every invocation: if a handler already has an - OtelContextLogFilter, it is left as-is, so warm Lambda reuse will not stack - duplicate filters. A single shared filter instance is reused across all - handlers. + This is safe to call on every invocation, and from several at once: the + check for an already-installed filter and the install that follows it happen + under one lock, so concurrent first-time callers cannot stack duplicate + filters on a handler. Installation carries no invocation identity -- see + :func:`bind_invocation` for that -- so a warm environment reuses the + filter installed by the first invocation as is. Args: - plugin: The OTel plugin that resolves the current span context. target_logger: Logger whose handlers receive the filter. Defaults to the root logger, which in AWS Lambda is where runtime log handlers live. @@ -98,18 +256,20 @@ def install_log_filter( """ logger = target_logger if target_logger is not None else logging.getLogger() - context_filter: OtelContextLogFilter | None = None - for handler in logger.handlers: - existing = next( - (f for f in handler.filters if isinstance(f, OtelContextLogFilter)), - None, - ) - if existing is not None: - # Reuse the already-installed filter so a single instance is shared. - context_filter = existing - continue - if context_filter is None: - context_filter = OtelContextLogFilter(plugin) - handler.addFilter(context_filter) - - return context_filter + with _install_lock: + context_filter: OtelContextLogFilter | None = None + for handler in logger.handlers: + existing = next( + (f for f in handler.filters if isinstance(f, OtelContextLogFilter)), + None, + ) + if existing is not None: + # Reuse the already-installed filter so a single instance is + # shared by every handler. + context_filter = existing + continue + if context_filter is None: + context_filter = OtelContextLogFilter() + handler.addFilter(context_filter) + + return context_filter diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py new file mode 100644 index 00000000..028b9998 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_factory.py @@ -0,0 +1,95 @@ +"""Plugin factories for the bundled durable-execution OTel plugins. + +The SDK's plugin contract is a factory object whose ``create_plugin`` is called +once per invocation: ``DurableInstrumentationPluginFactory`` declares +``create_plugin(info: InvocationStartInfo) -> DurableInstrumentationPlugin``. The +instance a factory returns serves exactly that one invocation and is dropped when +the invocation scope exits, so a plugin keeps its per-invocation state in ordinary +instance attributes. + +Both factories are classes rather than closures so the configuration they were +built with stays inspectable (``factory.config``) and so the entry points below +name an object with a readable type. + +Everything else the plugins need is resolved per invocation inside the plugin +itself: the tracer provider (which for the global-provider case may only be +installed after the handler module is imported), the tracer, and the +deterministic id generator and sampler installed on it. Those installs are +scoped to the plugin's own tracer, so building a plugin per invocation does not +disturb other instrumentation scopes. They are also atomic and return whatever +the tracer holds, which matters because a provider caches tracers by +instrumentation scope: two invocations starting at once get one tracer, and a +plugin that lost the install race must use the wrapper on that tracer rather than +one of its own that the tracer would never consult. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from aws_durable_execution_sdk_python_otel.execution_plugin import ( + ExecutionOtelPlugin, +) +from aws_durable_execution_sdk_python_otel.invocation_plugin import ( + InvocationOtelPlugin, +) + + +if TYPE_CHECKING: + from aws_durable_execution_sdk_python.plugin import InvocationStartInfo + + from aws_durable_execution_sdk_python_otel.otel_plugin_config import ( + OtelPluginConfig, + ) + + +class InvocationOtelPluginFactory: + """Builds one :class:`InvocationOtelPlugin` per invocation. + + Register the factory itself, not a plugin, with + ``@durable_execution(plugins=[...])``:: + + @durable_execution(plugins=[InvocationOtelPluginFactory()]) + def handler(event, context): ... + + Args: + config: Shared plugin configuration handed to every plugin this factory + builds. When omitted, each plugin uses defaults (globally configured + tracer provider, X-Ray extractor, "Workflow" span name, log + enrichment on). + """ + + def __init__(self, config: OtelPluginConfig | None = None) -> None: + self.config = config + + def create_plugin(self, info: InvocationStartInfo) -> InvocationOtelPlugin: + """Return this invocation's plugin. + + ``info`` is accepted because the SDK passes it, and is unused: the + plugin reads the same object again in ``on_invocation_start``, which is + where all of its invocation identity is derived. + """ + return InvocationOtelPlugin(self.config) + + +class ExecutionOtelPluginFactory: + """Builds one :class:`ExecutionOtelPlugin` per invocation. + + Args: + config: Shared plugin configuration handed to every plugin this factory + builds. When omitted, each plugin uses defaults. + """ + + def __init__(self, config: OtelPluginConfig | None = None) -> None: + self.config = config + + def create_plugin(self, info: InvocationStartInfo) -> ExecutionOtelPlugin: + """Return this invocation's plugin. ``info`` is unused; see above.""" + return ExecutionOtelPlugin(self.config) + + +INVOCATION_OTEL_PLUGIN_FACTORY = InvocationOtelPluginFactory() +"""Default-configured factory named by the ``otel-invocation`` entry point.""" + +EXECUTION_OTEL_PLUGIN_FACTORY = ExecutionOtelPluginFactory() +"""Default-configured factory named by the ``otel-execution`` entry point.""" diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_provider.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_provider.py deleted file mode 100644 index c4be734c..00000000 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin_provider.py +++ /dev/null @@ -1,23 +0,0 @@ -from aws_durable_execution_sdk_python.plugin import ( - DurableInstrumentationPluginProvider, -) - -from aws_durable_execution_sdk_python_otel.execution_plugin import ( - ExecutionOtelPlugin, -) -from aws_durable_execution_sdk_python_otel.invocation_plugin import ( - InvocationOtelPlugin, -) - - -INVOCATION_OTEL_PLUGIN_PROVIDER = DurableInstrumentationPluginProvider( - plugin_type=InvocationOtelPlugin, - factory=InvocationOtelPlugin, - plugin_api_version=1, -) - -EXECUTION_OTEL_PLUGIN_PROVIDER = DurableInstrumentationPluginProvider( - plugin_type=ExecutionOtelPlugin, - factory=ExecutionOtelPlugin, - plugin_api_version=1, -) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_concurrent_log_correlation_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_concurrent_log_correlation_int.py new file mode 100644 index 00000000..198b648d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_concurrent_log_correlation_int.py @@ -0,0 +1,290 @@ +"""End-to-end log correlation coverage for concurrent invocations. + +Drives the decorated handler, not the filter in isolation. The SDK runs the +handler body on a worker thread it submits to a pool, so only a test that goes +through ``durable_execution`` exercises the path from the invocation-start hook +to a record emitted by top-level handler code. + +Both invocations are held at a barrier until each has started, so a record is +only emitted while two invocations are open. That is the case the log filter +cannot resolve from the number of open invocations alone. +""" + +from __future__ import annotations + +import json +import logging +import threading +from dataclasses import replace +from datetime import UTC, datetime +from typing import Any +from unittest.mock import Mock, patch + +from aws_durable_execution_sdk_python.context import DurableContext, durable_step +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + ExecutionDetails, + Operation, + OperationAction, + OperationStatus, + OperationType, + StepDetails, +) +from aws_durable_execution_sdk_python_otel.log_filter import OtelContextLogFilter +from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from aws_durable_execution_sdk_python_otel.plugin_factory import ( + InvocationOtelPluginFactory, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + +EXECUTION_START = datetime(2026, 8, 27, 5, 11, 47, tzinfo=UTC) +OWNERS = ("first", "second") +TOP_LEVEL_MESSAGE = "top-level-handler-log" + + +def _execution_arn(owner: str) -> str: + return f"test-arn/{_execution_operation_id(owner)}" + + +def _execution_operation_id(owner: str) -> str: + """The execution operation is keyed by the last segment of the ARN.""" + return f"concurrent-log-correlation-{owner}" + + +def _lambda_context(owner: str) -> Mock: + context = Mock() + context.aws_request_id = f"request-{owner}" + context.client_context = None + context.identity = None + context._epoch_deadline_time_in_ms = 0 # noqa: SLF001 + context.invoked_function_arn = "test-arn" + context.tenant_id = None + return context + + +def _execution_operation(owner: str) -> Operation: + return Operation( + operation_id=_execution_operation_id(owner), + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + start_timestamp=EXECUTION_START, + # The handler receives the parsed input payload, so the owner travels + # through the execution input rather than the invocation event. + execution_details=ExecutionDetails(input_payload=json.dumps({"owner": owner})), + ) + + +def _event(owner: str) -> dict[str, Any]: + return { + "DurableExecutionArn": _execution_arn(owner), + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [_execution_operation(owner).to_json_dict()], + "NextMarker": "", + }, + "LocalRunner": True, + } + + +def _shared_checkpoint_store(): + """Return a checkpoint callable serving several concurrent executions. + + One mocked Lambda client is shared by both invocations, so operations are + kept per execution ARN and mutated under a lock. + """ + operations: dict[str, dict[str, Operation]] = { + _execution_arn(owner): { + _execution_operation_id(owner): _execution_operation(owner) + } + for owner in OWNERS + } + lock = threading.Lock() + + def checkpoint( + durable_execution_arn, + checkpoint_token, # noqa: ARG001 + updates, + client_token="token", # noqa: S107, ARG001 + ) -> CheckpointOutput: + with lock: + execution_operations = operations.setdefault(durable_execution_arn, {}) + for update in updates: + now = datetime.now(UTC) + previous = execution_operations.get(update.operation_id) + if update.action is OperationAction.START: + execution_operations[update.operation_id] = Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + name=update.name, + sub_type=update.sub_type, + start_timestamp=now, + ) + elif update.action is OperationAction.SUCCEED: + base = previous or Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + name=update.name, + sub_type=update.sub_type, + start_timestamp=now, + ) + execution_operations[update.operation_id] = replace( + base, + status=OperationStatus.SUCCEEDED, + end_timestamp=now, + step_details=( + StepDetails(result=update.payload, attempt=1) + if update.operation_type is OperationType.STEP + else base.step_details + ), + ) + snapshot = list(execution_operations.values()) + + return CheckpointOutput( + checkpoint_token="new-token", + new_execution_state=CheckpointUpdatedExecutionState(operations=snapshot), + ) + + return checkpoint + + +class _RecordCollector(logging.Handler): + """Collects the top-level handler records the filter has stamped.""" + + def __init__(self) -> None: + super().__init__() + self.records: list[logging.LogRecord] = [] + self._lock = threading.Lock() + + def emit(self, record: logging.LogRecord) -> None: + if record.getMessage().startswith(TOP_LEVEL_MESSAGE): + with self._lock: + self.records.append(record) + + +def _remove_otel_filters(handler: logging.Handler) -> None: + for installed in [ + f for f in handler.filters if isinstance(f, OtelContextLogFilter) + ]: + handler.removeFilter(installed) + + +def test_overlapping_invocations_stamp_top_level_logs_with_their_own_span() -> None: + """A handler's first log carries its own invocation's trace and span ids. + + Both invocations are open when either record is emitted, and neither record + is emitted from a thread the plugin hooks have run on, so the record can + only be correlated if the SDK carried the invocation's context into the + thread it runs the handler body on. + """ + collector = _RecordCollector() + root = logging.getLogger() + root.addHandler(collector) + # The record must reach the root handlers, so the emitting logger opts in to + # INFO explicitly rather than inheriting the root level. + probe_logger = logging.getLogger("probe.handler") + previous_level = probe_logger.level + probe_logger.setLevel(logging.INFO) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + # enrich_logger is the documented default: the plugin installs the filter on + # the root logger's handlers, including the collector added above. + factory = InvocationOtelPluginFactory( + OtelPluginConfig(tracer_provider=provider, enrich_logger=True) + ) + + both_started = threading.Barrier(len(OWNERS), timeout=30) + + @durable_step + def after_log(_step_context) -> str: + return "done" + + def handler_impl(event: Any, context: DurableContext) -> str: + owner = event["owner"] + # Hold here until every invocation has started, so the log below is + # emitted with more than one invocation open. Nothing durable has run. + both_started.wait() + logging.getLogger("probe.handler").info("%s %s", TOP_LEVEL_MESSAGE, owner) + return context.step(after_log(), name=f"after-log-{owner}") + + handler = durable_execution(handler_impl, plugins=[factory]) + results: dict[str, Any] = {} + failures: list[BaseException] = [] + results_lock = threading.Lock() + + def invoke(owner: str) -> None: + try: + result = handler(_event(owner), _lambda_context(owner)) + with results_lock: + results[owner] = result + except BaseException as error: # noqa: BLE001 + with results_lock: + failures.append(error) + both_started.abort() + + try: + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = _shared_checkpoint_store() + mock_client_class.initialize_client.return_value = mock_client + + threads = [ + threading.Thread(target=invoke, args=(owner,), name=f"invoke-{owner}") + for owner in OWNERS + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=60) + + assert not failures, failures + assert not any(thread.is_alive() for thread in threads) + for owner in OWNERS: + assert results[owner]["Status"] == InvocationStatus.SUCCEEDED.value + + # Each invocation's Invocation span identifies its execution by ARN, so + # the expected identifiers are read back from the exported spans rather + # than recomputed. + expected: dict[str, tuple[str, str]] = {} + for span in exporter.get_finished_spans(): + if span.name != "Invocation": + continue + assert span.attributes is not None + arn = span.attributes["durable.execution.arn"] + assert span.context is not None + expected[str(arn)] = ( + format(span.context.trace_id, "032x"), + format(span.context.span_id, "016x"), + ) + assert set(expected) == {_execution_arn(owner) for owner in OWNERS} + + stamped = { + record.getMessage().rsplit(" ", 1)[1]: ( + getattr(record, "traceId", None), + getattr(record, "spanId", None), + ) + for record in collector.records + } + assert set(stamped) == set(OWNERS), collector.records + assert stamped == { + owner: expected[_execution_arn(owner)] for owner in OWNERS + }, f"records carried {stamped}" + finally: + for installed_on in list(root.handlers): + _remove_otel_filters(installed_on) + root.removeHandler(collector) + probe_logger.setLevel(previous_level) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 63563493..ebabc49b 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -28,9 +28,11 @@ from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( derive_workflow_span_id, ) -from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin -from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from aws_durable_execution_sdk_python_otel.plugin_factory import ( + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, +) from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -138,18 +140,21 @@ def checkpoint( @pytest.mark.parametrize( - "plugin_type", - [InvocationOtelPlugin, ExecutionOtelPlugin], + "factory_type", + [InvocationOtelPluginFactory, ExecutionOtelPluginFactory], ) def test_otel_wait_resume_spans_share_default_xray_execution_trace( monkeypatch: pytest.MonkeyPatch, - plugin_type: type[InvocationOtelPlugin] | type[ExecutionOtelPlugin], + factory_type: type[InvocationOtelPluginFactory] | type[ExecutionOtelPluginFactory], ) -> None: monkeypatch.setenv("_X_AMZN_TRACE_ID", XRAY_TRACE_HEADER) exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) - plugin = plugin_type( + # The SDK takes a factory and calls its create_plugin once per invocation, so + # the two invocations below are served by two plugin instances sharing this + # provider. + factory = factory_type( OtelPluginConfig( tracer_provider=provider, enrich_logger=False, @@ -164,7 +169,7 @@ def handler_impl(_event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(1), name="otel-wait") return context.step(complete_after_resume(), name="otel-after-resume") - handler = durable_execution(handler_impl, plugins=[plugin]) + handler = durable_execution(handler_impl, plugins=[factory]) initial_operations = [_execution_operation()] first_checkpoint, first_operations = _checkpoint_store(initial_operations) @@ -226,7 +231,7 @@ def handler_impl(_event: Any, context: DurableContext) -> str: after_resume = next(span for span in spans if span.name == "otel-after-resume") assert len(invocations) >= 2 - if plugin_type is InvocationOtelPlugin: + if factory_type is InvocationOtelPluginFactory: assert len(waits) >= 2 # one segment per invocation else: assert len(waits) == 1 # one span per operation @@ -238,7 +243,7 @@ def handler_impl(_event: Any, context: DurableContext) -> str: } assert after_resume.parent is not None - if plugin_type is InvocationOtelPlugin: + if factory_type is InvocationOtelPluginFactory: assert after_resume.parent.span_id in { span.context.span_id for span in invocations } diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py index 5c3cc3ec..3fd0d9e9 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging import threading from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime, timedelta @@ -32,6 +33,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import ( NonRecordingSpan, + Span, SpanContext, SpanKind, TraceFlags, @@ -48,6 +50,7 @@ from aws_durable_execution_sdk_python_otel.durable_parent_span import ( DurableParentSpan, ) +from aws_durable_execution_sdk_python_otel.log_filter import OtelContextLogFilter from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig @@ -73,9 +76,16 @@ def _assert_otel_context_balanced(): def _create_plugin( context_extractor=lambda _: None, + exporter: InMemorySpanExporter | None = None, ) -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: - """Create an ExecutionOtelPlugin wired to an in-memory exporter.""" - exporter = InMemorySpanExporter() + """Create an ExecutionOtelPlugin wired to an in-memory exporter. + + One plugin instance serves exactly one invocation, so a test that spans + invocations creates a plugin per invocation and passes the same ``exporter`` + to each -- the way the SDK's factory hands successive invocations distinct + instances that publish to one provider. + """ + exporter = exporter if exporter is not None else InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) plugin = ExecutionOtelPlugin( @@ -721,7 +731,13 @@ def test_suspended_operation_held_as_non_recording_placeholder(): def test_suspend_then_resume_operation_exports_one_deterministic_span(): - """An operation spanning invocations exports one deterministic span.""" + """An operation spanning invocations exports one deterministic span. + + Each invocation gets its own plugin instance, so nothing about the operation + is carried in memory from one invocation to the next: the single exported + span and its ID come from the deterministic derivation off the execution ARN, + and the replay guard is the hook's own ``is_replayed`` flag. + """ plugin, exporter = _create_plugin() operation_id = "wait-across-invocations" @@ -745,6 +761,7 @@ def test_suspend_then_resume_operation_exports_one_deterministic_span(): assert not [s for s in exporter.get_finished_spans() if s.name == "long-wait"] # Invocation N+1: the still-open operation is replayed, then completes. + plugin, _ = _create_plugin(exporter=exporter) plugin.on_invocation_start(_invocation_start_info()) plugin.on_operation_start( OperationStartInfo( @@ -785,6 +802,7 @@ def test_suspend_then_resume_operation_exports_one_deterministic_span(): # ReplayChildren/virtual child completion callbacks are replay-only and # must not re-export the terminal deterministic span in a later invocation. + plugin, _ = _create_plugin(exporter=exporter) plugin.on_invocation_start(_invocation_start_info()) plugin.on_operation_end( OperationEndInfo( @@ -836,7 +854,8 @@ def test_suspended_child_context_exports_one_span_on_replay(): # Nothing exported for the suspended context. assert not [s for s in exporter.get_finished_spans() if s.name == context_id] - # Invocation 2: the context replays and completes. + # Invocation 2: the context replays and completes, in a fresh plugin instance. + plugin, _ = _create_plugin(exporter=exporter) plugin.on_invocation_start(_invocation_start_info()) plugin.on_operation_start( OperationStartInfo( @@ -903,10 +922,12 @@ def test_checkpointless_context_end_uses_a_non_negative_duration(): def test_virtual_context_replay_uses_unique_linked_segments(): - plugin, exporter = _create_plugin() + exporter = InMemorySpanExporter() context_id = "flat-branch" for _ in range(2): + # Each invocation is served by its own plugin instance. + plugin, _ = _create_plugin(exporter=exporter) plugin.on_invocation_start(_invocation_start_info()) # Virtual contexts have no durable START hook. plugin.on_user_function_start(_context_start_info(context_id)) @@ -1401,15 +1422,21 @@ def test_invocation_end_releases_scope_of_suspended_user_function(): assert plugin._context_tokens == {} -def test_warm_invocation_reuse_restores_ambient_span_each_time(): - """Verify repeated invocations leave the ambient Lambda span current.""" - plugin, _ = _create_plugin() +def test_successive_invocations_restore_ambient_span_each_time(): + """Verify each invocation's own plugin leaves the ambient Lambda span current. + + The SDK builds a plugin per invocation, so this drives three invocations + through three instances against one warm environment. What is asserted is + that no instance leaves a context attached behind it -- state carried in the + instance is irrelevant, because the instance is gone. + """ ambient_provider = TracerProvider() ambient = ambient_provider.get_tracer("ambient").start_span("AmbientLambda") token = otel_context.attach(trace.set_span_in_context(ambient)) try: warm_context = otel_context.get_current() for index in range(3): + plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) operation_id = f"step-{index}" plugin.on_user_function_start(_step_start_info(operation_id)) @@ -1620,3 +1647,73 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): plugin.on_invocation_end(_invocation_end_info()) assert plugin._context_tokens == {} + + +# --------------------------------------------------------------------------- +# Log correlation +# --------------------------------------------------------------------------- +def _stamped_span_id() -> str: + """Return the span ID the log filter stamps on a record emitted right here.""" + record = logging.LogRecord( + name="test", + level=logging.INFO, + pathname=__file__, + lineno=1, + msg="message", + args=(), + exc_info=None, + ) + OtelContextLogFilter().filter(record) + return str(getattr(record, "spanId", None)) + + +def _span_id_hex(span: Span) -> str: + """Return a span's ID in the hex form the log filter stamps.""" + return format(span.get_span_context().span_id, "016x") + + +def test_top_level_log_names_invocation_span_not_the_attached_workflow_span(): + """A top-level handler record names the Invocation span, not the Workflow span. + + This plugin makes the Workflow span current at invocation start so + auto-instrumented spans join the execution trace, and the SDK carries the + invocation thread's context into the thread that runs the handler body. The + Workflow span spans the whole execution, so it is less specific than the + Invocation span for a record emitted by one invocation's top-level code. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + try: + assert plugin._invocation_span is not None + assert plugin._workflow_span is not None + # Confirm the shape under test: the Workflow span is the current span. + assert trace.get_current_span() is plugin._workflow_span + + stamped = _stamped_span_id() + + assert stamped == _span_id_hex(plugin._invocation_span) + assert stamped != _span_id_hex(plugin._workflow_span) + finally: + plugin.on_invocation_end(_invocation_end_info()) + + +def test_step_attempt_log_names_the_attempt_span(): + """A record inside a step names the attempt span, not the Invocation span. + + Excluding the Workflow span from log correlation must not also exclude spans + that become current inside the invocation. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + try: + plugin.on_user_function_start(_step_start_info("step-1")) + attempt_span = plugin._get_span("step-1:attempt:1") + assert attempt_span is not None + assert plugin._invocation_span is not None + + stamped = _stamped_span_id() + + assert stamped == _span_id_hex(attempt_span) + assert stamped != _span_id_hex(plugin._invocation_span) + finally: + plugin.on_invocation_end(_invocation_end_info()) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py index df3e4b65..d559625f 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py @@ -50,6 +50,9 @@ ) from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from aws_durable_execution_sdk_python_otel.plugin_factory import ( + ExecutionOtelPluginFactory, +) START_TIME = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) @@ -303,7 +306,9 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( monkeypatch.setattr(trace, "_TRACER_PROVIDER", None) current_provider: list[ApiTracerProvider] = [ProxyTracerProvider()] monkeypatch.setattr(trace, "get_tracer_provider", lambda: current_provider[0]) - plugin = ExecutionOtelPlugin( + # A factory, because the two invocations below need two plugin instances and + # the second one must resolve the provider installed after the first ran. + factory = ExecutionOtelPluginFactory( OtelPluginConfig( context_extractor=lambda _: None, enrich_logger=False, @@ -311,6 +316,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( ) provider, exporter = _provider() + plugin = factory.create_plugin(_invocation_start()) plugin.on_invocation_start(_invocation_start()) assert "telemetry is disabled for this invocation" in caplog.text @@ -319,6 +325,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( plugin.on_invocation_end(_invocation_end()) assert exporter.get_finished_spans() == () + plugin = factory.create_plugin(_invocation_start()) plugin.on_invocation_start(_invocation_start()) _run_step_lifecycle(plugin) plugin.on_invocation_end(_invocation_end()) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index f108ef72..efc765d3 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -35,6 +35,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import ( NonRecordingSpan, + Span, SpanContext, SpanKind, StatusCode, @@ -77,17 +78,26 @@ def _assert_otel_context_balanced(): ) -def _create_plugin() -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: +def _create_plugin( + exporter: InMemorySpanExporter | None = None, +) -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: """Create a plugin wired to an in-memory span exporter.""" - return _create_plugin_with_sampler() + return _create_plugin_with_sampler(exporter=exporter) def _create_plugin_with_sampler( sampler: Sampler | None = None, context_extractor=lambda _: None, + exporter: InMemorySpanExporter | None = None, ) -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: - """Create a plugin wired to an in-memory span exporter.""" - exporter = InMemorySpanExporter() + """Create a plugin wired to an in-memory span exporter. + + One plugin instance serves exactly one invocation, so a test that spans + invocations creates a plugin per invocation and passes the same ``exporter`` + to each -- the way the SDK's factory hands successive invocations distinct + instances that publish to one provider. + """ + exporter = exporter if exporter is not None else InMemorySpanExporter() trace_provider = TracerProvider(sampler=sampler) trace_provider.add_span_processor(SimpleSpanProcessor(exporter)) plugin = InvocationOtelPlugin( @@ -125,6 +135,45 @@ def _invocation_end_info( ) +def _same_trace_ambient_context() -> SpanContext: + """Return the span context of an ADOT-style enclosing Lambda span. + + Under X-Ray active tracing the execution trace is the ambient trace, so the + span the ADOT layer has current before the invocation starts is on the + execution trace and becomes the parent of the Invocation span. + """ + return SpanContext( + trace_id=_to_otel_trace_id(EXECUTION_ARN, START_TIME), + span_id=int("1234567890abcdef", 16), + is_remote=True, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_state=TraceState(), + ) + + +def _stamped_span_id(plugin: InvocationOtelPlugin) -> str: + """Return the span ID the log filter stamps on a record emitted right here.""" + record = logging.LogRecord( + name="test", + level=logging.INFO, + pathname=__file__, + lineno=1, + msg="message", + args=(), + exc_info=None, + ) + OtelContextLogFilter().filter(record) + assert getattr(record, "traceId", None) == format( + plugin._execution_trace_id or 0, "032x" + ) + return str(getattr(record, "spanId", None)) + + +def _span_id_hex(span: Span) -> str: + """Return a span's ID in the hex form the log filter stamps.""" + return format(span.get_span_context().span_id, "016x") + + def _user_function_start_info( operation_id: str, attempt: int = 1, @@ -372,7 +421,7 @@ def test_log_filter_uses_invocation_trace_when_ambient_trace_is_rejected(): exc_info=None, ) - OtelContextLogFilter(plugin).filter(record) + OtelContextLogFilter().filter(record) invocation_span = plugin._get_span(None) assert invocation_span is not None @@ -417,6 +466,119 @@ def test_invocation_span_parents_to_same_trace_ambient_span(): assert workflow.parent.span_id == derive_execution_root_span_id(EXECUTION_ARN) +def test_top_level_log_names_invocation_span_not_the_enclosing_ambient_span(): + """A top-level handler record names the Invocation span, not its parent. + + In the X-Ray active tracing plus ADOT shape the layer's Lambda span is + current before the invocation starts, is on the execution trace, and is the + parent of the Invocation span. The SDK carries the invocation thread's + context into the thread that runs the handler body, so that span is current + where top-level handler code runs. It is one level up the tree from the + invocation that emitted the record, so the Invocation span is the more + specific answer and the one a top-level record must carry. + """ + plugin, _ = _create_plugin() + ambient_context = _same_trace_ambient_context() + ambient = NonRecordingSpan(ambient_context) + token = otel_context.attach(trace.set_span_in_context(ambient, Context())) + try: + plugin.on_invocation_start(_invocation_start_info()) + invocation_span = plugin._get_span(None) + assert invocation_span is not None + # Confirm the shape under test: same trace, and the ambient span really + # is the parent of the Invocation span. + assert plugin._execution_trace_id == ambient_context.trace_id + assert invocation_span.parent is not None + assert invocation_span.parent.span_id == ambient_context.span_id + + stamped = _stamped_span_id(plugin) + + assert stamped == _span_id_hex(invocation_span) + assert stamped != format(ambient_context.span_id, "016x") + + # The same holds between top-level operations: once a step's scope is + # released the enclosing ambient span is current again. + plugin.on_user_function_start(_user_function_start_info("step-1")) + plugin.on_user_function_end(_user_function_end_info("step-1")) + assert trace.get_current_span().get_span_context().span_id == ( + ambient_context.span_id + ) + assert _stamped_span_id(plugin) == _span_id_hex(invocation_span) + finally: + plugin.on_invocation_end(_invocation_end_info()) + otel_context.detach(token) + + +def test_durable_operation_spans_win_over_the_enclosing_ambient_span(): + """A record inside a durable operation names that operation's span. + + The enclosing ambient span is excluded from log correlation, but a span that + becomes current inside the invocation is more specific than the Invocation + span and still wins. Both operation shapes that attach a scope are covered: + a STEP attempt and a child CONTEXT. + """ + plugin, _ = _create_plugin() + ambient = NonRecordingSpan(_same_trace_ambient_context()) + token = otel_context.attach(trace.set_span_in_context(ambient, Context())) + try: + plugin.on_invocation_start(_invocation_start_info()) + invocation_span = plugin._get_span(None) + assert invocation_span is not None + + plugin.on_user_function_start(_user_function_start_info("step-1")) + attempt_span = plugin._get_span("step-1:attempt:1") + assert attempt_span is not None + assert _stamped_span_id(plugin) == _span_id_hex(attempt_span) + assert _stamped_span_id(plugin) != _span_id_hex(invocation_span) + plugin.on_user_function_end(_user_function_end_info("step-1")) + + plugin.on_user_function_start( + _user_function_start_info("ctx-1", operation_type=OperationType.CONTEXT) + ) + context_span = plugin._get_span("ctx-1") + assert context_span is not None + assert _stamped_span_id(plugin) == _span_id_hex(context_span) + assert _stamped_span_id(plugin) != _span_id_hex(invocation_span) + finally: + # The child context never ends, so invocation cleanup releases its scope. + plugin.on_invocation_end(_invocation_end_info()) + otel_context.detach(token) + + +def test_customer_span_started_in_the_handler_wins_over_the_invocation_span(): + """A record inside a span the handler body started names that span. + + A span the customer creates while the invocation is running is on the + execution trace, because the enclosing ambient span is current when it + starts, and it is not the enclosing span itself, so it is the most specific + span for a record emitted inside it. + """ + plugin, _ = _create_plugin() + ambient = NonRecordingSpan(_same_trace_ambient_context()) + token = otel_context.attach(trace.set_span_in_context(ambient, Context())) + try: + plugin.on_invocation_start(_invocation_start_info()) + invocation_span = plugin._get_span(None) + assert invocation_span is not None + + customer_span = plugin._provider.get_tracer("customer").start_span( + "customer-work" + ) + customer_token = otel_context.attach( + trace.set_span_in_context(customer_span, otel_context.get_current()) + ) + try: + stamped = _stamped_span_id(plugin) + assert stamped == _span_id_hex(customer_span) + assert stamped != _span_id_hex(invocation_span) + finally: + otel_context.detach(customer_token) + customer_span.end() + finally: + plugin.on_invocation_end(_invocation_end_info()) + otel_context.detach(token) + + def test_pre_terminal_placeholder_preserves_same_trace_tracestate(): """The Workflow placeholder and operation links carry ambient tracestate.""" plugin, exporter = _create_plugin() @@ -1759,12 +1921,14 @@ def test_checkpointed_context_first_span_uses_deterministic_id(): def test_virtual_context_replay_uses_unique_linked_segments(): - plugin, exporter = _create_plugin() + exporter = InMemorySpanExporter() operation_id = "flat-branch" span_name = f"step-{operation_id}" workflow_span_id = derive_workflow_span_id(EXECUTION_ARN) for _ in range(2): + # Each invocation is served by its own plugin instance. + plugin, _ = _create_plugin(exporter=exporter) plugin.on_invocation_start(_invocation_start_info()) # Virtual contexts have no durable START hook. plugin.on_user_function_start( @@ -1937,15 +2101,21 @@ def test_ambient_span_is_current_again_after_full_lifecycle(): ambient.end() -def test_warm_invocation_reuse_does_not_accumulate_scopes(): - """Verify repeated invocations on one plugin instance stay balanced.""" - plugin, _ = _create_plugin() +def test_successive_invocations_do_not_accumulate_scopes(): + """Verify each invocation's own plugin leaves the warm environment balanced. + + The SDK builds a plugin per invocation, so this drives three invocations + through three instances. What is asserted is that no instance leaves a + context attached behind it; state carried in the instance is irrelevant, + because the instance is gone. + """ ambient_provider = TracerProvider() ambient = ambient_provider.get_tracer("ambient").start_span("AmbientLambda") token = otel_context.attach(trace.set_span_in_context(ambient)) try: warm_context = otel_context.get_current() for index in range(3): + plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) operation_id = f"step-{index}" plugin.on_user_function_start(_user_function_start_info(operation_id)) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py index a50a06ac..7f71ee72 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py @@ -53,6 +53,9 @@ ) from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from aws_durable_execution_sdk_python_otel.plugin_factory import ( + InvocationOtelPluginFactory, +) START_TIME = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) @@ -253,7 +256,9 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( monkeypatch.setattr(trace, "_TRACER_PROVIDER", None) current_provider: list[ApiTracerProvider] = [ProxyTracerProvider()] monkeypatch.setattr(trace, "get_tracer_provider", lambda: current_provider[0]) - plugin = InvocationOtelPlugin( + # A factory, because the two invocations below need two plugin instances and + # the second one must resolve the provider installed after the first ran. + factory = InvocationOtelPluginFactory( OtelPluginConfig( context_extractor=lambda _: None, enrich_logger=False, @@ -261,6 +266,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( ) provider, exporter = _provider() + plugin = factory.create_plugin(_invocation_start()) plugin.on_invocation_start(_invocation_start()) assert "telemetry is disabled for this invocation" in caplog.text @@ -269,6 +275,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( plugin.on_invocation_end(_invocation_end()) assert exporter.get_finished_spans() == () + plugin = factory.create_plugin(_invocation_start()) plugin.on_invocation_start(_invocation_start()) _run_step_lifecycle(plugin) plugin.on_invocation_end(_invocation_end()) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py index 0cce516e..770a0f53 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py @@ -2,21 +2,42 @@ from __future__ import annotations +import contextvars +import gc import logging +import threading +import weakref from datetime import UTC, datetime +import opentelemetry.context as otel_context +import pytest from aws_durable_execution_sdk_python.lambda_service import ( OperationStatus, ) from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, InvocationStartInfo, + InvocationStatus, OperationType, UserFunctionStartInfo, ) +from opentelemetry import trace +from opentelemetry.context import Context from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import ( + NonRecordingSpan, + SpanContext, + TraceFlags, + TraceState, +) +from aws_durable_execution_sdk_python_otel import log_filter as log_filter_module +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + _to_otel_trace_id, +) +from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin from aws_durable_execution_sdk_python_otel.log_filter import ( OtelContextLogFilter, install_log_filter, @@ -29,6 +50,25 @@ EXECUTION_ARN = "arn:aws:lambda:us-west-2:123456789012:function:workflow:$LATEST" +@pytest.fixture(autouse=True) +def _isolated_invocation_registry(): + """Run each test against an empty open-invocation registry. + + The registry is process-global by design -- one installed filter serves every + invocation in the environment -- so a test that starts an invocation without + ending it would otherwise change what later tests resolve. + """ + saved = list(log_filter_module._open_invocations) + log_filter_module._open_invocations.clear() + token = log_filter_module._current_invocation.set(None) + try: + yield + finally: + log_filter_module._current_invocation.reset(token) + log_filter_module._open_invocations.clear() + log_filter_module._open_invocations.update(saved) + + def _create_plugin( enrich_logger: bool = True, ) -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: @@ -46,16 +86,26 @@ def _create_plugin( return plugin, exporter -def _invocation_start_info() -> InvocationStartInfo: +def _invocation_start_info(suffix: str = "") -> InvocationStartInfo: """Create standard invocation start info for tests.""" return InvocationStartInfo( - request_id="request-1", - execution_arn=EXECUTION_ARN, + request_id=f"request-1{suffix}", + execution_arn=f"{EXECUTION_ARN}{suffix}", execution_start_time=START_TIME, is_first_invocation=True, ) +def _invocation_end_info(suffix: str = "") -> InvocationEndInfo: + """Create standard invocation end info for tests.""" + return InvocationEndInfo( + request_id=f"request-1{suffix}", + execution_arn=f"{EXECUTION_ARN}{suffix}", + is_first_invocation=True, + status=InvocationStatus.SUCCEEDED, + ) + + def _user_function_start_info(operation_id: str) -> UserFunctionStartInfo: """Create standard user function start info for tests.""" return UserFunctionStartInfo( @@ -85,6 +135,21 @@ def _make_record() -> logging.LogRecord: ) +def _stamped(record: logging.LogRecord) -> tuple[str | None, str | None]: + """Return the trace and span identifiers a filter stamped on a record.""" + return getattr(record, "traceId", None), getattr(record, "spanId", None) + + +def _own_identifiers(plugin: InvocationOtelPlugin) -> tuple[str, str]: + """Return the trace and span identifiers of the plugin's current span.""" + span_context = plugin.get_current_span_context() + assert span_context is not None + return ( + format(span_context.trace_id, "032x"), + format(span_context.span_id, "016x"), + ) + + def _remove_otel_filters(handler: logging.Handler) -> None: """Remove any OtelContextLogFilter from a handler (test cleanup).""" for log_filter in [ @@ -95,16 +160,15 @@ def _remove_otel_filters(handler: logging.Handler) -> None: def test_filter_always_returns_true(): """The filter never drops a record, even with no active span.""" - plugin, _ = _create_plugin() - log_filter = OtelContextLogFilter(plugin) + log_filter = OtelContextLogFilter() assert log_filter.filter(_make_record()) is True def test_filter_does_not_set_fields_without_active_span(): - """With no invocation active, the filter leaves the record unmodified.""" - plugin, _ = _create_plugin() - log_filter = OtelContextLogFilter(plugin) + """With no invocation open, the filter leaves the record unmodified.""" + _create_plugin() + log_filter = OtelContextLogFilter() record = _make_record() log_filter.filter(record) @@ -118,7 +182,7 @@ def test_filter_injects_trace_context_from_invocation_span(): """The filter stamps the invocation span context for top-level code.""" plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) - log_filter = OtelContextLogFilter(plugin) + log_filter = OtelContextLogFilter() record = _make_record() log_filter.filter(record) @@ -134,24 +198,485 @@ def test_filter_uses_attempt_span_inside_user_function(): plugin.on_invocation_start(_invocation_start_info()) operation_id = "step-1" plugin.on_user_function_start(_user_function_start_info(operation_id)) + try: + record = _make_record() + OtelContextLogFilter().filter(record) + + attempt_span = plugin._get_span("step-1:attempt:1") + assert attempt_span is not None + expected_span_id = format(attempt_span.get_span_context().span_id, "016x") + assert record.spanId == expected_span_id + finally: + # Ends the invocation, which detaches the attempt scope this test + # attached to the running thread's OTel context. Left attached, it would + # stay current for every later test on this thread. + plugin.on_invocation_end(_invocation_end_info()) + + +def test_concurrent_invocations_each_stamp_their_own_span_context(): + """Two invocations open at once each correlate to their own trace. + + Logging handlers are process-global, so both invocations are served by one + filter instance. Each record must carry the trace and span identifiers of the + invocation that emitted it, not of whichever invocation started most + recently. + """ + shared_filter = OtelContextLogFilter() + both_started = threading.Barrier(2, timeout=10) + stamped: dict[str, tuple[str | None, str | None]] = {} + own: dict[str, tuple[str, str]] = {} + failures: list[BaseException] = [] + lock = threading.Lock() + + def invocation(owner: str) -> None: + try: + plugin, _ = _create_plugin(enrich_logger=False) + plugin.on_invocation_start(_invocation_start_info(suffix=owner)) + try: + with lock: + own[owner] = _own_identifiers(plugin) + # Emit only once both invocations are open, so a filter holding + # one mutable plugin reference is guaranteed to have been + # overwritten by the other invocation. + both_started.wait() + record = _make_record() + shared_filter.filter(record) + with lock: + stamped[owner] = _stamped(record) + finally: + plugin.on_invocation_end(_invocation_end_info()) + except BaseException as error: # noqa: BLE001 + with lock: + failures.append(error) + both_started.abort() + + threads = [ + threading.Thread(target=invocation, args=(owner,), name=f"invocation-{owner}") + for owner in ("a", "b") + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + assert not failures + assert own["a"] != own["b"] + assert stamped["a"] == own["a"] + assert stamped["b"] == own["b"] + + +def test_record_on_an_unclaimed_thread_is_not_stamped_with_the_open_invocation(): + """A thread carrying no claim is not correlated, even to a lone invocation. + + A thread that carries no claim -- one customer code started itself, since + ``threading.Thread`` does not copy the starting thread's context -- has no + invocation of its own. Resolving it to the single open invocation would be + right only when that invocation is the one that emitted the record, and two + orderings make it the wrong one: a thread still carrying a finished + invocation's claim, and an invocation's own thread before its start hook has + run. Correlation for such records is given up so neither ordering can stamp + one execution's record with another's trace. + """ + plugin, _ = _create_plugin(enrich_logger=False) + plugin.on_invocation_start(_invocation_start_info()) + try: + open_identifiers = _own_identifiers(plugin) + stamped: list[tuple[str | None, str | None]] = [] + + def emit() -> None: + record = _make_record() + OtelContextLogFilter().filter(record) + stamped.append(_stamped(record)) + + worker = threading.Thread(target=emit, name="unclaimed") + worker.start() + worker.join(timeout=10) + + assert stamped == [(None, None)] + assert stamped[0] != open_identifiers + finally: + plugin.on_invocation_end(_invocation_end_info()) + + +def test_a_stale_claim_is_not_resolved_to_the_only_open_invocation(): + """A thread still claimed by a finished invocation is not given a live one. + + ``unbind_invocation`` cannot reset the claim on any thread but its own, so a + thread the first invocation claimed still names that invocation after it + ends. A second invocation is then the only open one. Resolving by the number + of open invocations would stamp the first invocation's record with the + second's trace, which is one customer execution's identifiers on another's + record. + """ + first, _ = _create_plugin(enrich_logger=False) + first.on_invocation_start(_invocation_start_info(suffix="first")) + first_identifiers = _own_identifiers(first) + + # A worker running in a copy of the claiming thread's context, which is what + # the SDK submits the handler body as, so the claim reaches it. + claimed_context = contextvars.copy_context() + second_is_open = threading.Event() + stamped: list[tuple[str | None, str | None]] = [] + + def emit() -> None: + assert second_is_open.wait(timeout=10) + record = _make_record() + OtelContextLogFilter().filter(record) + stamped.append(_stamped(record)) + + worker = threading.Thread( + target=claimed_context.run, args=(emit,), name="claimed-by-first" + ) + worker.start() + try: + first.on_invocation_end(_invocation_end_info(suffix="first")) + second, _ = _create_plugin(enrich_logger=False) + second.on_invocation_start(_invocation_start_info(suffix="second")) + try: + second_identifiers = _own_identifiers(second) + second_is_open.set() + worker.join(timeout=10) + + assert len(stamped) == 1 + assert stamped[0] != second_identifiers + assert stamped[0] != first_identifiers + assert stamped[0] == (None, None) + finally: + second.on_invocation_end(_invocation_end_info(suffix="second")) + finally: + second_is_open.set() + worker.join(timeout=10) + + +def test_a_claimed_worker_does_not_pin_the_finished_invocation(): + """A claim left on a pool thread does not keep the plugin alive. + + ``unbind_invocation`` resets the claim only on the thread that ends the + invocation, so a pool thread that is never used again keeps the claim it was + given. The claim is a weak reference, so what it keeps is nothing: once the + invocation ends and the SDK releases the plugin, the plugin is collectable + even while the claimed thread is still running. + """ + plugin, _ = _create_plugin(enrich_logger=False) + plugin.on_invocation_start(_invocation_start_info(suffix="only")) + + # A worker running in a copy of the claiming thread's context, held alive for + # the length of the test, as a pooled worker would be. + claimed_context = contextvars.copy_context() + release_worker = threading.Event() + resolved_while_open: list[bool] = [] + + def hold_the_claim() -> None: + resolved_while_open.append(log_filter_module._resolve_provider() is not None) + assert release_worker.wait(timeout=10) + + worker = threading.Thread( + target=claimed_context.run, args=(hold_the_claim,), name="claimed-worker" + ) + worker.start() + try: + collected = threading.Event() + weakref.finalize(plugin, collected.set) + plugin.on_invocation_end(_invocation_end_info(suffix="only")) + + del plugin + gc.collect() + + assert resolved_while_open == [True], "the claim must resolve while open" + assert collected.is_set(), "the claimed worker must not pin the plugin" + finally: + release_worker.set() + worker.join(timeout=10) + + +def test_a_record_emitted_before_its_invocation_binds_is_not_given_the_open_one(): + """A record emitted before its own invocation binds is not correlated. + + An invocation emits records on its own thread before its invocation-start + hook runs -- while the SDK fetches initial state, for instance -- and that + thread carries no claim yet. Another invocation can be open at that moment. + Resolving by the number of open invocations would stamp the starting + invocation's record with the open invocation's trace. + """ + first, _ = _create_plugin(enrich_logger=False) + first.on_invocation_start(_invocation_start_info(suffix="first")) + second, _ = _create_plugin(enrich_logger=False) + stamped: list[tuple[str | None, str | None]] = [] + try: + first_identifiers = _own_identifiers(first) + + def start_second_invocation() -> None: + # This thread belongs to the second invocation, which has not bound + # itself yet, so nothing here carries a claim. + record = _make_record() + OtelContextLogFilter().filter(record) + stamped.append(_stamped(record)) + second.on_invocation_start(_invocation_start_info(suffix="second")) + + worker = threading.Thread(target=start_second_invocation, name="second-wrapper") + worker.start() + worker.join(timeout=10) + + assert len(stamped) == 1 + assert stamped[0] != first_identifiers + assert stamped[0] != _own_identifiers(second) + assert stamped[0] == (None, None) + finally: + first.on_invocation_end(_invocation_end_info(suffix="first")) + second.on_invocation_end(_invocation_end_info(suffix="second")) + + +def test_context_propagated_into_a_worker_resolves_the_claiming_invocation(): + """A worker started from a copy of the claiming thread's context resolves it. + + This is what the SDK does for the thread it runs the handler body on: the + invocation-start hook claims the invocation thread, and the handler body runs + in a copy of that thread's context. Both invocations are open when either + record is emitted, so the number of open invocations cannot resolve them and + only the propagated claim can. + """ + both_emitted = threading.Barrier(2, timeout=10) + stamped: dict[str, tuple[str | None, str | None]] = {} + own: dict[str, tuple[str, str]] = {} + failures: list[BaseException] = [] + lock = threading.Lock() + + def invocation(owner: str) -> None: + """Run one invocation the way the SDK does, on its own thread.""" + try: + plugin, _ = _create_plugin(enrich_logger=False) + plugin.on_invocation_start(_invocation_start_info(suffix=owner)) + try: + with lock: + own[owner] = _own_identifiers(plugin) + + def emit() -> None: + record = _make_record() + # Emit only once both invocations are open, so the + # single-open-invocation fallback cannot resolve the record. + both_emitted.wait() + OtelContextLogFilter().filter(record) + with lock: + stamped[owner] = _stamped(record) + + # A fresh copy per submission: one Context cannot be entered + # twice concurrently, which is why the SDK copies at each submit. + worker = threading.Thread( + target=contextvars.copy_context().run, + args=(emit,), + name=f"worker-{owner}", + ) + worker.start() + worker.join(timeout=10) + finally: + plugin.on_invocation_end(_invocation_end_info()) + except BaseException as error: # noqa: BLE001 + with lock: + failures.append(error) + both_emitted.abort() + + threads = [ + threading.Thread(target=invocation, args=(owner,), name=f"invocation-{owner}") + for owner in ("a", "b") + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + assert not failures + assert own["a"] != own["b"] + assert stamped == own + + +def test_concurrent_invocations_with_enclosing_ambient_spans_stay_separate(): + """Excluding the enclosing ambient span does not blur two open invocations. + + Each invocation runs with its own ADOT-style enclosing span current, on its + own execution trace, and both are open when either record is emitted. Each + record must carry its own invocation's Invocation span: not the other + invocation's span, and not its own enclosing span. + """ + both_started = threading.Barrier(2, timeout=10) + stamped: dict[str, tuple[str | None, str | None]] = {} + invocation_spans: dict[str, tuple[str, str]] = {} + ambient_span_ids: dict[str, str] = {} + failures: list[BaseException] = [] + lock = threading.Lock() + ambient_span_id_by_owner = { + "a": int("1111aaaa1111aaaa", 16), + "b": int("2222bbbb2222bbbb", 16), + } + + def invocation(owner: str) -> None: + try: + plugin, _ = _create_plugin(enrich_logger=False) + # The enclosing span sits on this execution's own trace, which is + # what the ADOT layer produces under X-Ray active tracing. + ambient_context = SpanContext( + trace_id=_to_otel_trace_id(f"{EXECUTION_ARN}{owner}", START_TIME), + span_id=ambient_span_id_by_owner[owner], + is_remote=True, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_state=TraceState(), + ) + token = otel_context.attach( + trace.set_span_in_context(NonRecordingSpan(ambient_context), Context()) + ) + try: + plugin.on_invocation_start(_invocation_start_info(suffix=owner)) + try: + invocation_span = plugin._get_span(None) + assert invocation_span is not None + span_context = invocation_span.get_span_context() + with lock: + invocation_spans[owner] = ( + format(span_context.trace_id, "032x"), + format(span_context.span_id, "016x"), + ) + ambient_span_ids[owner] = format( + ambient_context.span_id, "016x" + ) + both_started.wait() + record = _make_record() + OtelContextLogFilter().filter(record) + with lock: + stamped[owner] = _stamped(record) + finally: + plugin.on_invocation_end(_invocation_end_info(suffix=owner)) + finally: + otel_context.detach(token) + except BaseException as error: # noqa: BLE001 + with lock: + failures.append(error) + both_started.abort() + + threads = [ + threading.Thread(target=invocation, args=(owner,), name=f"invocation-{owner}") + for owner in ("a", "b") + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + assert not failures, failures + assert invocation_spans["a"] != invocation_spans["b"] + assert stamped == invocation_spans + for owner in ("a", "b"): + assert stamped[owner][1] != ambient_span_ids[owner] + + +def test_record_on_a_customer_thread_is_not_attributed_to_another_invocation(): + """A thread customer code starts itself is never given the wrong invocation. + + ``threading.Thread`` does not copy the starting thread's context, so a thread + a handler creates directly carries no claim. This case is not the SDK's + handler worker, which is submitted with a copy of the invocation's context. + With two invocations open there is nothing to resolve such a record against, + and it is left uncorrelated rather than attributed to either invocation. + """ + first, _ = _create_plugin(enrich_logger=False) + second, _ = _create_plugin(enrich_logger=False) + first.on_invocation_start(_invocation_start_info(suffix="first")) + second.on_invocation_start(_invocation_start_info(suffix="second")) + try: + wrong_identifiers = {_own_identifiers(first), _own_identifiers(second)} + stamped: list[tuple[str | None, str | None]] = [] + + def emit() -> None: + record = _make_record() + OtelContextLogFilter().filter(record) + stamped.append(_stamped(record)) + + worker = threading.Thread(target=emit, name="customer-created") + worker.start() + worker.join(timeout=10) + + assert len(stamped) == 1 + # The rule that matters: never another execution's trace. + assert stamped[0] not in wrong_identifiers + # And with no claim and no single open invocation, nothing is stamped. + assert stamped[0] == (None, None) + finally: + first.on_invocation_end(_invocation_end_info()) + second.on_invocation_end(_invocation_end_info()) + + +def test_finished_invocation_does_not_correlate_later_records(): + """A thread claimed by an invocation stops correlating once it ends. + + A pooled thread can outlive the invocation that claimed it, so liveness is + checked at emit time rather than assumed from the claim. + """ + plugin, _ = _create_plugin(enrich_logger=False) + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info()) record = _make_record() - OtelContextLogFilter(plugin).filter(record) + OtelContextLogFilter().filter(record) + + assert not hasattr(record, "traceId") + assert not hasattr(record, "spanId") - attempt_span = plugin._get_span("step-1:attempt:1") - assert attempt_span is not None - expected_span_id = format(attempt_span.get_span_context().span_id, "016x") - assert record.spanId == expected_span_id + +class _SpanProcessorFailingOnEnd(SpanProcessor): + """Raises when a span ends, standing in for customer span-processor code.""" + + def on_start(self, span, parent_context=None) -> None: + pass + + def on_end(self, span) -> None: + raise RuntimeError("span processor on_end failed") + + +@pytest.mark.parametrize( + "plugin_class", + [InvocationOtelPlugin, ExecutionOtelPlugin], + ids=lambda c: c.__name__, +) +def test_invocation_is_released_when_span_shutdown_fails(plugin_class): + """A failing span shutdown still releases the invocation. + + Ending a span calls the configured span processors, which are customer code + and can raise. The plugin executor contains that exception, so the invocation + survives it. An invocation left registered would stay open for the life of + the environment, and a thread still carrying its claim would keep correlating + records to its finished spans. + """ + provider = TracerProvider() + provider.add_span_processor(_SpanProcessorFailingOnEnd()) + plugin = plugin_class( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + plugin.on_invocation_start(_invocation_start_info()) + + with pytest.raises(RuntimeError, match="span processor on_end failed"): + plugin.on_invocation_end(_invocation_end_info()) + + assert plugin not in log_filter_module._open_invocations + # This thread carried the claim, so it is the thread a later record would be + # mis-attributed on. + assert log_filter_module._resolve_provider() is not plugin + record = _make_record() + OtelContextLogFilter().filter(record) + assert _stamped(record) == (None, None) + # The OTel context stack belongs to the thread, so a scope the plugin left + # attached would stay current after the invocation returned. + assert plugin._context_tokens == {} def test_install_log_filter_attaches_to_handlers(): """install_log_filter adds the filter to each handler on the target logger.""" - plugin, _ = _create_plugin() target = logging.getLogger("test.install") handler = logging.NullHandler() target.addHandler(handler) try: - installed = install_log_filter(plugin, target_logger=target) + installed = install_log_filter(target_logger=target) assert isinstance(installed, OtelContextLogFilter) assert any(isinstance(f, OtelContextLogFilter) for f in handler.filters) @@ -161,13 +686,59 @@ def test_install_log_filter_attaches_to_handlers(): def test_install_log_filter_is_idempotent(): """Repeated installs do not stack duplicate filters on a handler.""" - plugin, _ = _create_plugin() target = logging.getLogger("test.idempotent") handler = logging.NullHandler() target.addHandler(handler) try: - install_log_filter(plugin, target_logger=target) - install_log_filter(plugin, target_logger=target) + install_log_filter(target_logger=target) + install_log_filter(target_logger=target) + + otel_filters = [ + f for f in handler.filters if isinstance(f, OtelContextLogFilter) + ] + assert len(otel_filters) == 1 + finally: + target.removeHandler(handler) + + +def test_concurrent_first_time_installs_attach_one_filter(): + """Two invocations installing at once cannot both add a filter. + + The handler holds the first install open at the moment it attaches, so a + second caller that was not serialized behind it still sees a filterless + handler and attaches a second filter. + """ + + class HandlerRacingOnAttach(logging.NullHandler): + """Holds the first attach open until a second caller reaches it.""" + + def __init__(self) -> None: + super().__init__() + self._barrier = threading.Barrier(2, timeout=0.2) + self._released = False + + def addFilter(self, filter) -> None: # noqa: A002 - stdlib signature + if not self._released: + try: + self._barrier.wait() + except threading.BrokenBarrierError: + # Installation was serialized, so no second caller arrived. + pass + self._released = True + super().addFilter(filter) + + target = logging.getLogger("test.install.race") + handler = HandlerRacingOnAttach() + target.addHandler(handler) + try: + threads = [ + threading.Thread(target=install_log_filter, args=(target,)) + for _ in range(2) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(10) otel_filters = [ f for f in handler.filters if isinstance(f, OtelContextLogFilter) @@ -177,16 +748,45 @@ def test_install_log_filter_is_idempotent(): target.removeHandler(handler) +def test_a_later_invocation_takes_over_log_correlation(): + """The invocation that is open now owns correlation, not the first one. + + A logging handler lives as long as the Lambda environment while a plugin + instance lives for one invocation, so a filter tied to the first + invocation's plugin would stop correlating logs after that invocation ended. + """ + first, _ = _create_plugin() + second, _ = _create_plugin() + target = logging.getLogger("test.rebind") + handler = logging.NullHandler() + target.addHandler(handler) + try: + installed = install_log_filter(target_logger=target) + assert installed is not None + assert install_log_filter(target_logger=target) is installed + + first.on_invocation_start(_invocation_start_info(suffix="first")) + first.on_invocation_end(_invocation_end_info()) + second.on_invocation_start(_invocation_start_info(suffix="second")) + + record = _make_record() + installed.filter(record) + + assert _stamped(record) == _own_identifiers(second) + finally: + second.on_invocation_end(_invocation_end_info()) + target.removeHandler(handler) + + def test_install_log_filter_reuses_single_instance_across_handlers(): """A single filter instance is shared across all handlers.""" - plugin, _ = _create_plugin() target = logging.getLogger("test.shared") handler_a = logging.NullHandler() handler_b = logging.NullHandler() target.addHandler(handler_a) target.addHandler(handler_b) try: - installed = install_log_filter(plugin, target_logger=target) + installed = install_log_filter(target_logger=target) filter_a = next( f for f in handler_a.filters if isinstance(f, OtelContextLogFilter) @@ -202,10 +802,9 @@ def test_install_log_filter_reuses_single_instance_across_handlers(): def test_install_log_filter_returns_none_without_handlers(): """With no handlers, install_log_filter has nothing to attach to.""" - plugin, _ = _create_plugin() target = logging.getLogger("test.nohandlers") - assert install_log_filter(plugin, target_logger=target) is None + assert install_log_filter(target_logger=target) is None def test_plugin_installs_filter_on_root_logger_at_construction(): diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py index 22b41638..5a59d9bf 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py @@ -1,10 +1,15 @@ +import re import tomllib from pathlib import Path +from packaging.specifiers import SpecifierSet +from packaging.version import Version + PACKAGE_ROOT = Path(__file__).resolve().parents[1] REPOSITORY_ROOT = PACKAGE_ROOT.parents[1] -CORE_DEPENDENCY = "aws-durable-execution-sdk-python>=2.0.0" +CORE_DISTRIBUTION = "aws-durable-execution-sdk-python" +CORE_DEPENDENCY = "aws-durable-execution-sdk-python>=3.0.0,<4" TEST_OTEL_DEPENDENCIES = { "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", @@ -80,3 +85,121 @@ def test_pypi_compatibility_environment_uses_compatible_core_sdk() -> None: ]["test-pypi-otel"]["dependencies"] assert CORE_DEPENDENCY in dependencies + + +def _core_version() -> str: + """The core SDK version this repository builds, read from its source. + + Read from the file rather than imported. The ``test-pypi-otel`` environment + installs a *published* core alongside this package's source, so an import + would report that release's version and the checks below would stop saying + anything about this repository's version story. + """ + about = ( + REPOSITORY_ROOT + / "packages" + / "aws-durable-execution-sdk-python" + / "src" + / "aws_durable_execution_sdk_python" + / "__about__.py" + ).read_text() + match = re.search(r'^__version__ = "([^"]+)"', about, re.MULTILINE) + assert match is not None, "core __about__.py has no __version__ assignment" + return match.group(1) + + +def _core_dependency_lower_bound() -> str: + dependencies = _load_pyproject(PACKAGE_ROOT / "pyproject.toml")["project"][ + "dependencies" + ] + bounds = [ + dependency.removeprefix(CORE_DISTRIBUTION + ">=").split(",", 1)[0] + for dependency in dependencies + if dependency.startswith(CORE_DISTRIBUTION + ">=") + ] + assert len(bounds) == 1, f"expected one {CORE_DISTRIBUTION} bound, got {bounds}" + return bounds[0] + + +def _major(version: str) -> int: + return int(version.split(".", 1)[0]) + + +def test_core_dependency_bound_matches_the_core_major_in_this_repository() -> None: + """The declared bound must not admit a core major that predates this plugin's contract. + + This package's entry points resolve to plugin factories, and the ``plugins`` + argument only accepts factories from the core major that introduced them. A + lower bound naming an earlier major is a resolution pip accepts and that then + fails at handler initialization, so the bound has to track the core major this + repository builds. The bound may lag within that major -- a later core minor + still satisfies the contract -- which is why only the major is compared and + the bound is required not to exceed the core version. + """ + core_version = _core_version() + lower_bound = _core_dependency_lower_bound() + + assert _major(lower_bound) == _major(core_version) + assert Version(lower_bound) <= Version(core_version) + + +def test_layer_sdk_pin_matches_the_core_version_in_this_repository() -> None: + """The OTel layer pin selects the core wheel bundled into the published layer. + + A combined SDK and OTel release fails outright when the pin disagrees with the + released SDK version, and an OTel-only release downloads exactly the pinned + version from PyPI. A stale pin therefore either blocks the release or ships a + layer whose core cannot run this plugin, so the pin tracks the core version + this repository builds. + """ + metadata_path = REPOSITORY_ROOT / ".github" / "lambda-layer-publish.toml" + + with metadata_path.open("rb") as metadata_file: + pinned_version = tomllib.load(metadata_file)["layer"]["sdk-version"] + + assert pinned_version == _core_version() + + +def test_core_dependency_excludes_the_next_core_major() -> None: + """A lower bound alone is the same defect one major later. + + The lower bound exists because this package's entry points resolve to plugin + factories, which the core major below cannot call: pip accepts the resolution + and the handler fails at initialization. Without a ceiling the next core major + that changes the plugin contract reproduces exactly that, so the specifier has + to reject it rather than only reject what came before. + """ + dependencies = _load_pyproject(PACKAGE_ROOT / "pyproject.toml")["project"][ + "dependencies" + ] + specifiers = [ + SpecifierSet(dependency.removeprefix(CORE_DISTRIBUTION)) + for dependency in dependencies + if dependency.startswith(CORE_DISTRIBUTION) + ] + assert len(specifiers) == 1 + + core_major = _major(_core_version()) + assert specifiers[0].contains(_core_version(), prereleases=True) + assert not specifiers[0].contains(f"{core_major + 1}.0.0", prereleases=True) + + +def test_core_dependency_admits_a_later_core_patch() -> None: + """The ceiling belongs on the major, not on the version built here. + + A ``<=`` ceiling looks equivalent and is not: it excludes the next core patch, + so the first core patch release puts this claim out of date for a change that + cannot have touched the plugin contract. + """ + dependencies = _load_pyproject(PACKAGE_ROOT / "pyproject.toml")["project"][ + "dependencies" + ] + specifier = next( + SpecifierSet(dependency.removeprefix(CORE_DISTRIBUTION)) + for dependency in dependencies + if dependency.startswith(CORE_DISTRIBUTION) + ) + + core = Version(_core_version()) + next_patch = f"{core.major}.{core.minor}.{core.micro + 1}" + assert specifier.contains(next_patch, prereleases=True) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py new file mode 100644 index 00000000..74ca4dea --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_factory.py @@ -0,0 +1,276 @@ +"""Tests for the bundled OTel plugin factories and the entry points naming them. + +The SDK plugin contract is a factory object whose ``create_plugin`` is called once +per invocation, so what these tests have to establish is that the objects the +package exposes -- and the ones its entry points name -- carry ``create_plugin``, +that it builds a plugin, and that each call builds a NEW plugin. The old +provider-shaped assertions (a declared ``plugin_type`` and an API version) have no +counterpart: the contract carries neither. +""" + +from __future__ import annotations + +import importlib +import logging +import sys +import threading +import tomllib +from datetime import UTC, datetime +from pathlib import Path + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStartInfo, +) +from opentelemetry.sdk.trace import Tracer as SdkTracer +from opentelemetry.sdk.trace import TracerProvider + +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + DeterministicIdGenerator, +) +from aws_durable_execution_sdk_python_otel.durable_sampling import DurableSampler +from aws_durable_execution_sdk_python_otel.execution_plugin import ( + ExecutionOtelPlugin, +) +from aws_durable_execution_sdk_python_otel.invocation_plugin import ( + InvocationOtelPlugin, +) +from aws_durable_execution_sdk_python_otel.log_filter import OtelContextLogFilter +from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from aws_durable_execution_sdk_python_otel.plugin_factory import ( + EXECUTION_OTEL_PLUGIN_FACTORY, + INVOCATION_OTEL_PLUGIN_FACTORY, + ExecutionOtelPluginFactory, + InvocationOtelPluginFactory, +) + + +PACKAGE_ROOT = Path(__file__).resolve().parents[1] +PLUGIN_ENTRY_POINT_GROUP = "aws_durable_execution.plugins" +EXECUTION_ARN = "arn:aws:lambda:us-west-2:123456789012:function:workflow:$LATEST" + + +@pytest.fixture(autouse=True) +def _remove_installed_log_filters(): + """Detach any log filter a default-configured plugin installed. + + ``OtelPluginConfig.enrich_logger`` defaults to True, so building a plugin + from a default-configured factory attaches a filter to the root logger's + handlers, which outlive the test. + """ + yield + for handler in logging.getLogger().handlers: + for installed in [ + f for f in handler.filters if isinstance(f, OtelContextLogFilter) + ]: + handler.removeFilter(installed) + + +def _invocation_start_info() -> InvocationStartInfo: + return InvocationStartInfo( + request_id="request-1", + execution_arn=EXECUTION_ARN, + execution_start_time=datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC), + is_first_invocation=True, + ) + + +def _declared_entry_points() -> dict[str, str]: + with (PACKAGE_ROOT / "pyproject.toml").open("rb") as pyproject: + project = tomllib.load(pyproject)["project"] + return project["entry-points"][PLUGIN_ENTRY_POINT_GROUP] + + +def _resolve(spec: str) -> object: + """Resolve a ``module:attribute`` entry-point value the way the SDK does. + + Resolved from the declared value rather than from installed distribution + metadata so the assertion holds wherever the tests run, installed or not. + """ + module_name, _, attribute = spec.partition(":") + return getattr(importlib.import_module(module_name), attribute) + + +@pytest.mark.parametrize( + ("factory_type", "plugin_type"), + [ + (InvocationOtelPluginFactory, InvocationOtelPlugin), + (ExecutionOtelPluginFactory, ExecutionOtelPlugin), + ], +) +def test_factory_builds_its_plugin_type( + factory_type: type, plugin_type: type[DurableInstrumentationPlugin] +) -> None: + factory = factory_type(OtelPluginConfig(enrich_logger=False)) + + assert isinstance(factory.create_plugin(_invocation_start_info()), plugin_type) + + +@pytest.mark.parametrize( + "factory_type", + [InvocationOtelPluginFactory, ExecutionOtelPluginFactory], +) +def test_factory_builds_a_fresh_plugin_per_invocation(factory_type: type) -> None: + """Two calls must never hand back the same instance. + + This is the whole point of the factory contract: the returned plugin holds + one invocation's state in ordinary instance fields, so a shared instance + would leak span registries and context tokens from one invocation into the + next. + """ + factory = factory_type(OtelPluginConfig(enrich_logger=False)) + + first = factory.create_plugin(_invocation_start_info()) + second = factory.create_plugin(_invocation_start_info()) + + assert first is not second + + +@pytest.mark.parametrize( + "factory_type", + [InvocationOtelPluginFactory, ExecutionOtelPluginFactory], +) +def test_factory_passes_its_config_to_every_plugin(factory_type: type) -> None: + config = OtelPluginConfig(workflow_span_name="Custom", enrich_logger=False) + factory = factory_type(config) + + assert factory.config is config + assert factory.create_plugin(_invocation_start_info())._config is config + assert factory.create_plugin(_invocation_start_info())._config is config + + +@pytest.mark.parametrize( + "factory_type", + [InvocationOtelPluginFactory, ExecutionOtelPluginFactory], +) +def test_factory_without_config_builds_a_default_configured_plugin( + factory_type: type, +) -> None: + factory = factory_type() + + assert factory.config is None + assert factory.create_plugin(_invocation_start_info())._config == OtelPluginConfig() + + +def test_module_level_factories_are_default_configured() -> None: + assert INVOCATION_OTEL_PLUGIN_FACTORY.config is None + assert EXECUTION_OTEL_PLUGIN_FACTORY.config is None + + +def test_declared_entry_points_name_the_bundled_factories() -> None: + entry_points = _declared_entry_points() + + assert set(entry_points) == {"otel-invocation", "otel-execution"} + assert _resolve(entry_points["otel-invocation"]) is INVOCATION_OTEL_PLUGIN_FACTORY + assert _resolve(entry_points["otel-execution"]) is EXECUTION_OTEL_PLUGIN_FACTORY + + +def test_declared_entry_points_resolve_to_factories_that_build_plugins() -> None: + """The entry points must satisfy the SDK's factory contract. + + ``plugin_discovery._load_factory`` requires an object with a callable + ``create_plugin``, so a target resolving to a plugin class, to a plugin + instance, or to a plain function now fails at handler initialization. The + check here mirrors the SDK's own, then calls the method to confirm what it + builds. + """ + expected = { + "otel-invocation": InvocationOtelPlugin, + "otel-execution": ExecutionOtelPlugin, + } + + for name, spec in _declared_entry_points().items(): + factory = _resolve(spec) + create_plugin = getattr(factory, "create_plugin", None) + assert callable(create_plugin) + assert not isinstance(factory, type) + assert isinstance(create_plugin(_invocation_start_info()), expected[name]) + + +# -- concurrent first invocations sharing one cached tracer -------------------- + +# Trials per factory. A single trial reproduces the loss roughly one time in ten +# (measured at 16 threads with the switch interval below), so a handful of trials +# would pass with an unsynchronized install still in place. +_SHARED_TRACER_TRIALS = 200 + +# Concurrent create_plugin calls per trial, standing in for concurrent first +# invocations in one Lambda Managed Instances environment. +_SHARED_TRACER_PLUGINS = 16 + + +def _build_plugins_concurrently( + factory_type: type, provider: TracerProvider, count: int +) -> list[InvocationOtelPlugin | ExecutionOtelPlugin]: + config = OtelPluginConfig(tracer_provider=provider, enrich_logger=False) + factory = factory_type(config) + ready = threading.Barrier(count) + plugins: list[InvocationOtelPlugin | ExecutionOtelPlugin | None] = [None] * count + + def build(index: int) -> None: + ready.wait(10.0) + plugins[index] = factory.create_plugin(_invocation_start_info()) + + workers = [threading.Thread(target=build, args=(index,)) for index in range(count)] + for worker in workers: + worker.start() + for worker in workers: + worker.join(10.0) + assert not any(worker.is_alive() for worker in workers) + built = [plugin for plugin in plugins if plugin is not None] + assert len(built) == count + return built + + +@pytest.mark.parametrize( + "factory_type", + [InvocationOtelPluginFactory, ExecutionOtelPluginFactory], +) +def test_concurrent_plugins_share_the_wrappers_their_tracer_holds( + factory_type: type, +) -> None: + """Every plugin must hold the generator and sampler the tracer actually uses. + + A TracerProvider caches tracers by instrumentation scope, so plugins built for + concurrent first invocations ask for one instrument name and get one tracer + object. Installing the deterministic generator and the durable sampler on that + tracer is a check-then-set: two plugins can each read the original generator, + each wrap it, and the second assignment replaces the first. The plugin that + assigned first then holds a wrapper the tracer no longer consults, so its + deterministic ID overrides are ignored and its workflow and operation span IDs + come out random, which breaks cross-invocation stitching. + + The switch interval is lowered so the interpreter preempts inside that + check-then-set often enough for the loss to appear within the trial count; it + does not create the window, it only makes an existing one likely to be hit. + """ + previous_interval = sys.getswitchinterval() + sys.setswitchinterval(1e-9) + try: + for trial in range(_SHARED_TRACER_TRIALS): + provider = TracerProvider() + plugins = _build_plugins_concurrently( + factory_type, provider, _SHARED_TRACER_PLUGINS + ) + tracer = provider.get_tracer(OtelPluginConfig().instrument_name) + assert isinstance(tracer, SdkTracer) + # One tracer for every plugin: the premise the rest of the assertions + # rest on, and the reason a lost install is not simply harmless. + assert all(plugin._tracer is tracer for plugin in plugins) + assert isinstance(tracer.id_generator, DeterministicIdGenerator) + assert isinstance(tracer.sampler, DurableSampler) + # Exactly one wrapper of each kind exists, and it is the tracer's. + assert {id(plugin._id_generator) for plugin in plugins} == { + id(tracer.id_generator) + }, f"trial {trial}: a plugin holds an id generator the tracer discarded" + assert {id(plugin._sampling_delegate) for plugin in plugins} == { + id(tracer.sampler.delegate) + }, f"trial {trial}: a plugin holds a sampling delegate the tracer discarded" + # The wrapper wraps the provider's original, not another wrapper. + assert not isinstance(tracer.sampler.delegate, DurableSampler) + assert not isinstance( + tracer.id_generator._fallback_id_generator, DeterministicIdGenerator + ) + finally: + sys.setswitchinterval(previous_interval) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_provider.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_provider.py deleted file mode 100644 index c47098b2..00000000 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin_provider.py +++ /dev/null @@ -1,56 +0,0 @@ -from aws_durable_execution_sdk_python.plugin import ( - DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION, -) -from aws_durable_execution_sdk_python.plugin_discovery import ( - PLUGIN_ENVIRONMENT_VARIABLE, - load_configured_plugins, -) - -from aws_durable_execution_sdk_python_otel.execution_plugin import ( - ExecutionOtelPlugin, -) -from aws_durable_execution_sdk_python_otel.invocation_plugin import ( - InvocationOtelPlugin, -) -from aws_durable_execution_sdk_python_otel.plugin_provider import ( - EXECUTION_OTEL_PLUGIN_PROVIDER, - INVOCATION_OTEL_PLUGIN_PROVIDER, -) - - -def test_invocation_otel_plugin_provider_uses_current_plugin_api() -> None: - assert INVOCATION_OTEL_PLUGIN_PROVIDER.plugin_type is InvocationOtelPlugin - assert ( - INVOCATION_OTEL_PLUGIN_PROVIDER.plugin_api_version - == DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION - ) - - -def test_invocation_otel_plugin_provider_creates_invocation_plugin() -> None: - assert isinstance(INVOCATION_OTEL_PLUGIN_PROVIDER.factory(), InvocationOtelPlugin) - - -def test_execution_otel_plugin_provider_uses_current_plugin_api() -> None: - assert EXECUTION_OTEL_PLUGIN_PROVIDER.plugin_type is ExecutionOtelPlugin - assert ( - EXECUTION_OTEL_PLUGIN_PROVIDER.plugin_api_version - == DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION - ) - - -def test_execution_otel_plugin_provider_creates_execution_plugin() -> None: - assert isinstance(EXECUTION_OTEL_PLUGIN_PROVIDER.factory(), ExecutionOtelPlugin) - - -def test_installed_otel_entry_points_load_both_plugin_types() -> None: - plugins = load_configured_plugins( - None, - environment={ - PLUGIN_ENVIRONMENT_VARIABLE: "otel-invocation,otel-execution", - }, - ) - - assert [type(plugin) for plugin in plugins] == [ - InvocationOtelPlugin, - ExecutionOtelPlugin, - ] diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py index 41789171..b7a4f137 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py @@ -37,7 +37,12 @@ class RecordingWaitPlugin(DurableInstrumentationPlugin): - """Records end notifications for the wait and counts invocations.""" + """Records end notifications for the wait and counts invocations. + + State is class-level on purpose: the SDK builds a fresh instance per + invocation, and this test spans two invocations, so what it asserts on has + to outlive any single instance. + """ invocation_count: ClassVar[int] = 0 wait_end_infos: ClassVar[list[OperationEndInfo]] = [] @@ -55,13 +60,29 @@ def on_operation_end(self, info: OperationEndInfo) -> None: self.wait_end_infos.append(info) +class RecordingWaitPluginFactory: + """Builds one :class:`RecordingWaitPlugin` for each invocation. + + ``durable_execution(plugins=[...])`` takes factory objects whose + ``create_plugin`` the SDK calls once per invocation. A bare callable is + rejected while the handler is being initialized, and this handler is built at + module import, so registering one would abort collection of this module. This + factory exists only to construct the plugin; the state the test asserts on is + class-level on the plugin, so the factory holds nothing. + """ + + def create_plugin(self, info: InvocationStartInfo) -> RecordingWaitPlugin: + """Return this invocation's plugin. ``info`` is unused.""" + return RecordingWaitPlugin() + + def _wait_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 """Suspend on a top-level wait, then finish.""" context.wait(Duration.from_seconds(1), name=_WAIT_NAME) return "done" -wait_handler = durable_execution(_wait_handler, plugins=[RecordingWaitPlugin()]) +wait_handler = durable_execution(_wait_handler, plugins=[RecordingWaitPluginFactory()]) def test_wait_completed_during_suspend_is_delivered_as_new() -> None: diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index bf7776f4..d2df0d1f 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -40,15 +40,60 @@ DURABLE_EXECUTION_PLUGINS=otel-invocation,example_audit The SDK resolves those names from the `aws_durable_execution.plugins` Python entry-point group when the decorated handler is initialized. An unset or blank variable preserves the existing behavior. The decorator's `plugins` argument -remains supported; explicit plugins run first and take precedence when a -dynamic provider creates the same concrete plugin type. +remains supported; explicit factories run first, and a factory passed to the +decorator is not registered a second time through the environment. + +A plugin is registered as a *factory*, not as an instance. A factory is an object +with a `create_plugin(info)` method taking the invocation's `InvocationStartInfo` +and returning a `DurableInstrumentationPlugin`; the SDK calls that method once per +invocation, so the instance it returns serves that one invocation only and can +hold **per-invocation** state in ordinary attributes. + +Per-invocation is narrower than per-execution, and the difference matters. A +durable execution spans as many invocations as it waits, retries or resumes, and +the instance is dropped when each of those returns — so anything a plugin keeps in +its attributes is gone by the next invocation of the same execution. State that +has to survive that has two honest homes: rebuild it from the operation map the +invocation hooks carry (`InvocationStartInfo.operations` is a full snapshot, +which is how the bundled Insight plugin reports operations that completed in an +earlier invocation), or put it on the factory, which outlives every invocation — +keyed by execution ARN, and pruned by the owner, because the SDK will not tell the +factory when an execution ends for good. + +A factory is an object with a method rather than a plain callable so the +registration type can grow a second, optional member later -- a process-level +flush on execution-environment shutdown, for example -- without a second breaking +change to this surface. + +Write a small factory class and construct the plugin in its `create_plugin`. The +factory holds what outlives an invocation, such as an exporter or a resolved +configuration, and setup work that can fail belongs in the factory's own +constructor rather than the plugin's: -Provider packages expose a versioned factory: +```python +class AuditPluginFactory: + def __init__(self, sink): + self._sink = sink + + def create_plugin(self, info): + return AuditPlugin(self._sink) + + +plugins=[AuditPluginFactory(sink)] +``` + +A plugin class is not a factory, and neither is a bare callable. `plugins=[MyPlugin]` +and `plugins=[lambda info: MyPlugin(sink)]` raise `PluginLoadError` during handler +initialization, because neither carries `create_plugin`. A class that declares +`create_plugin` as a `@classmethod` is accepted, since the requirement is the +member and not the kind of object. + +Provider packages expose such a factory: ```python from aws_durable_execution_sdk_python.plugin import ( DurableInstrumentationPlugin, - DurableInstrumentationPluginProvider, + InvocationStartInfo, ) @@ -56,26 +101,24 @@ class AuditPlugin(DurableInstrumentationPlugin): pass -AUDIT_PLUGIN_PROVIDER = DurableInstrumentationPluginProvider( - plugin_type=AuditPlugin, - factory=AuditPlugin, - plugin_api_version=1, -) +class AuditPluginFactory: + def create_plugin(self, info: InvocationStartInfo) -> AuditPlugin: + return AuditPlugin() + + +AUDIT_PLUGIN_FACTORY = AuditPluginFactory() ``` -Register the provider in the package's `pyproject.toml`: +Register the factory instance in the package's `pyproject.toml`: ```toml [project.entry-points."aws_durable_execution.plugins"] -example_audit = "example_audit:AUDIT_PLUGIN_PROVIDER" +example_audit = "example_audit:AUDIT_PLUGIN_FACTORY" ``` -Set `plugin_api_version` to the literal API version the provider implements. -Update it only after verifying the provider against that API version. - Provider names must be unique across installed distributions. Missing, -ambiguous, incompatible, or invalid providers raise `PluginLoadError` during -handler initialization with the provider and distribution details. +ambiguous, or wrongly shaped providers raise `PluginLoadError` during handler +initialization with the provider and distribution details. ## 🚀 Quick Start diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__about__.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__about__.py index 0d5aa99f..92b4bb99 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__about__.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__about__.py @@ -1,4 +1,4 @@ # SPDX-FileCopyrightText: 2025-present Amazon.com, Inc. or its affiliates. # # SPDX-License-Identifier: Apache-2.0 -__version__ = "2.0.1" +__version__ = "3.0.0" diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py index 8ca5ef90..517eca30 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py @@ -1,6 +1,7 @@ from __future__ import annotations import contextlib +import contextvars import functools import json import logging @@ -27,8 +28,9 @@ OperationUpdate, ) from aws_durable_execution_sdk_python.plugin import ( - DurableInstrumentationPlugin, + DurableInstrumentationPluginFactory, PluginExecutor, + PluginHost, ) from aws_durable_execution_sdk_python.plugin_discovery import ( load_configured_plugins, @@ -169,7 +171,7 @@ def durable_execution( func: Callable[[Any, DurableContext], Any] | None = None, *, boto3_client: Boto3LambdaClient | None = None, - plugins: list[DurableInstrumentationPlugin] | None = None, + plugins: list[DurableInstrumentationPluginFactory] | None = None, ) -> Callable[[Any, LambdaContext], Any]: """ Decorator to create a durable execution handler. @@ -177,7 +179,10 @@ def durable_execution( Args: func: The user function to decorate boto3_client: Optional boto3 Lambda client to use - plugins: Optional list of instrumentation plugins to use + plugins: Optional list of instrumentation plugin factories. Each + factory's ``create_plugin`` is called once per invocation with that + invocation's ``InvocationStartInfo``, and the instance it returns + serves only that invocation. """ # Decorator called with parameters if func is None: @@ -188,10 +193,26 @@ def durable_execution( logger.debug("Starting durable execution handler...") - plugin_executor = PluginExecutor(load_configured_plugins(plugins)) - - @plugin_executor.handle_durable_output - def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: + # Only the resolved factory list is handler-lifetime. The plugin instances, + # and the invocation metadata the hooks read, are built per invocation by + # PluginHost.invocation() and live in that invocation's frame -- see the + # plugin_executor parameter below. + plugin_host = PluginHost(load_configured_plugins(plugins)) + + @plugin_host.handle_durable_output + # The metadata of whatever function the host wrapper wraps becomes the + # decorated handler's own, because the host wrapper applies + # functools.wraps() to it. This function takes a third argument the returned + # handler does not accept, so without this line inspect.signature() on the + # handler advertises a required `plugin_executor` parameter, and a + # signature-aware runtime or test harness rejects or misinvokes a handler + # that in fact takes (event, context). Copying the user function's metadata + # here puts it at the head of the chain the host wrapper then extends, so + # the handler reports the user function's signature, name and docstring. + @functools.wraps(func) + def wrapper( + event: Any, context: LambdaContext, plugin_executor: PluginExecutor + ) -> MutableMapping[str, Any]: invocation_input: DurableExecutionInvocationInput service_client: DurableServiceClient @@ -306,14 +327,41 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: operations_provider=lambda: execution_state.operations, updated_operation_ids=invocation_input.updated_operation_ids, ) - # Thread 1: Run background checkpoint processing + # Thread 1: Run background checkpoint processing. + # + # Submitted without the invocation's context, deliberately. This + # thread runs SDK checkpointing, never user code, and nothing it + # does needs a contextvar the invocation thread set. Copying the + # context here would also make the invocation thread's ambient OTel + # context current on a background thread that starts and ends + # plugin spans, which widens the change with no caller-visible + # benefit. executor.submit(execution_state.checkpoint_batches_forever) # Thread 2: Execute user function logger.debug( "%s entering user-space...", invocation_input.durable_execution_arn ) - user_future = executor.submit(func, input_event, durable_context) + # Carry this thread's context into the worker that runs the handler + # body. A callable submitted to a ThreadPoolExecutor runs on a worker + # thread, whose context is not the submitting thread's, so without + # this the handler body starts from a context in which no contextvar + # set before the handler is visible -- neither those set by the + # plugins that ran in on_invocation_start just above, nor those set + # by customer code around the decorator. Log correlation depends on + # it, since a plugin that claims the invocation for the calling + # thread at invocation start has no other way to reach the thread the + # handler body runs on. + # + # The copy is taken here, per submission: a Context cannot be entered + # twice concurrently, so the checkpoint thread above could not share + # one with this. Copying does not couple the two threads -- the + # worker mutates its own copy, and the invocation thread's context is + # unchanged either way, exactly as it was when the worker started + # from an unrelated context. + user_future = executor.submit( + contextvars.copy_context().run, func, input_event, durable_context + ) logger.debug( "%s waiting for user code completion...", diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py index eb5ee78b..797ae255 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py @@ -235,7 +235,7 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> ErrorObject: ) @classmethod - def from_exception(cls, exception: Exception) -> ErrorObject: + def from_exception(cls, exception: BaseException) -> ErrorObject: # SerDesError and subclasses pin to the base discriminator so replay # always reconstructs them as SerDesError. if isinstance(exception, SerDesError): diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py index e9549d30..1a1f4692 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py @@ -5,11 +5,10 @@ import datetime import functools import logging -from collections.abc import Mapping, Sequence -from concurrent.futures import ThreadPoolExecutor +from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass, field from enum import Enum -from typing import Any, Callable, MutableMapping, cast +from typing import Any, Callable, MutableMapping, Protocol, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -28,8 +27,6 @@ logger = logging.getLogger(__name__) -DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION = 1 - class InvocationStatus(Enum): """Invocation outcomes exposed to instrumentation plugins.""" @@ -451,41 +448,299 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: pass -@dataclass(frozen=True) -class DurableInstrumentationPluginProvider: - """Versioned factory exposed through the plugin entry-point group.""" +class DurableInstrumentationPluginFactory(Protocol): + """Builds one plugin instance for one invocation. + + An object with a method rather than a bare callable, because the SDK will + grow process-level plugin hooks -- a flush when the execution environment + shuts down, for example. A callable type has no member to add such a hook to, + so growing one would have to change the registration type from callable to + object, which is a second breaking change on the same public surface. One + method on an object leaves room for an optional second member, which is + additive. + + Not ``@runtime_checkable``. Three facts decide it. An ``isinstance`` check + against a runtime-checkable protocol tests only that the member name is + present, not that it is callable, so it would accept an object whose + ``create_plugin`` is a string; the SDK needs the stronger test. The SDK also + has to name what an invalid entry actually was, which a boolean + ``isinstance`` result cannot supply. And ``isinstance`` against a + runtime-checkable protocol requires *every* declared member, so publishing + one would make the optional second member above non-additive for any caller + who wrote such a check. :func:`plugin_discovery._is_plugin_factory` performs + the check instead. + + Register the factory, not a plugin:: + + class MyPluginFactory: + def __init__(self, exporter: Exporter) -> None: + self._exporter = exporter + + def create_plugin(self, info: InvocationStartInfo) -> MyPlugin: + return MyPlugin(self._exporter) + + + plugins = [MyPluginFactory(exporter)] + + A plugin class is not a factory. ``plugins=[MyPlugin]`` used to work because + calling a class constructs an instance, and it now fails at handler + initialization because a class carries no ``create_plugin``. A class that + declares ``create_plugin`` itself -- as a ``@classmethod`` -- does satisfy the + shape, because the requirement is the member and not the kind of object. + + The factory holds what outlives an invocation: an exporter, a resolved + configuration, a shared worker. The plugin instance holds what does not. + Setup work that can fail or that reads the environment belongs in the + factory's own constructor rather than in the plugin's, because a plugin + constructor should only assign fields (see ``CONTRIBUTING.md``, + "Initialization and conversion"). + """ + + def create_plugin( + self, info: InvocationStartInfo, / + ) -> DurableInstrumentationPlugin: + """Return the plugin instance that serves the described invocation. + + Called once per invocation, with that invocation's + :class:`InvocationStartInfo` -- the same object the returned instance's + ``on_invocation_start`` then receives -- and before any hook fires. The + instance serves only that invocation and is dropped when it returns, so a + plugin can hold that invocation's state in ordinary instance attributes + without keying it by execution ARN. + + Per invocation is narrower than per execution. A durable execution spans + as many invocations as it waits, retries or resumes, so state a plugin + leaves in its attributes is gone by the next invocation of the same + execution. Anything that has to survive that is rebuilt from the operation + map the invocation hooks carry -- ``InvocationStartInfo.operations`` is a + full snapshot, including operations that completed in an earlier + invocation -- or kept on the factory, which outlives every invocation and + is therefore the caller's to key and to prune. + + ``info`` is positional-only, so an implementation may name the parameter + whatever reads best; a named protocol parameter would pin that name for + every implementation. + + A call that raises, or that returns ``None``, is logged and skipped for + that invocation and never disrupts the execution. + + Args: + info: The invocation the returned plugin instance will observe. + + Returns: + The plugin instance for this invocation. + """ + ... + + +def _type_name(value: object) -> str: + """Best available name for a value's type, for log messages. + + Reads the type rather than the value, so no instance ``__getattr__`` or + ``__getattribute__`` runs, and wraps the lookup, because these names are built + while a plugin failure is being contained: a name that raises would turn a + contained failure into a failed execution. + """ + try: + return type(value).__qualname__ + except BaseException: # noqa: BLE001 - a name is never worth failing a hook for + return "" - plugin_type: type[DurableInstrumentationPlugin] - factory: Callable[[], DurableInstrumentationPlugin] - plugin_api_version: int + +def _factory_name(factory: object) -> str: + """Best available name for a factory, for log messages. + + Read from the factory's *type*, not from the factory. This runs while a + factory failure is being contained, and a ``getattr`` on the instance would + call a custom ``__getattr__`` or ``__getattribute__`` -- so a factory whose + attribute hook raises would make the containment itself raise, turning a + contained plugin failure into a failed execution. A class registered directly + as a factory is read through the class object, which carries its own + ``__qualname__``. + + Wrapped as well, because diagnostics must not be the thing that fails: a name + that cannot be produced is reported as unavailable rather than raised. + """ + try: + if isinstance(factory, type): + return factory.__qualname__ + except BaseException: # noqa: BLE001 - a name is never worth failing a hook for + return "" + return _type_name(factory) + + +# Raised out of plugin code, these three are not reports of a plugin defect but +# instructions to the thread that is running: stop. Containing one would drop the +# instruction and return a thread that was told to unwind to the work after the +# plugin. They are re-raised; every other BaseException is contained. +# +# asyncio.CancelledError is deliberately NOT here. It derives from BaseException +# and it does mean "stop" for the task that was cancelled, but the task here is +# the SDK's, not the plugin's: nothing cancels the invocation thread or the +# single-worker plugin pool. A CancelledError arriving from plugin code therefore +# came from the plugin's own asyncio use -- an awaited task it let be cancelled -- +# which is a plugin defect and belongs on the contained side, or the plugin's +# failure would fail an execution it was only observing. +_PLUGIN_THREAD_CONTROL_EXCEPTIONS = (KeyboardInterrupt, SystemExit, GeneratorExit) + + +def _contain_plugin_failure( + error: BaseException, message: str, *message_args: object +) -> BaseException | None: + """Log what plugin code raised, and return the part that must not be contained. + + Returns ``None`` when the whole failure was contained, or the part that + instructs the calling thread to stop, which the caller re-raises. + + A :class:`BaseExceptionGroup` is split rather than tested, because it is + neither of the two cases a plain ``isinstance`` chain covers: a group carrying + a :class:`KeyboardInterrupt` is not an instance of one, so a tuple handler + naming the three does not match it and a broad handler would contain the + interrupt inside it. Plugin code produces such a group without asking for it + -- an ``asyncio.TaskGroup`` whose task is interrupted raises one -- so the + group is partitioned: the control leaves are returned to be re-raised, and + what remains is logged like any other contained plugin failure. + """ + control: BaseException | None + contained: BaseException | None + if isinstance(error, BaseExceptionGroup): + control, contained = error.split(_PLUGIN_THREAD_CONTROL_EXCEPTIONS) + elif isinstance(error, _PLUGIN_THREAD_CONTROL_EXCEPTIONS): + control, contained = error, None + else: + control, contained = None, error + if contained is not None: + logger.error(message, *message_args, exc_info=contained) + return control class PluginExecutor: - def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): - self._plugins = plugins or [] - self._executor: ThreadPoolExecutor | None = None + """One invocation's plugin instances, metadata and dispatch. + + Scoped to a single invocation, not to the handler. Everything mutable here -- + the instances built for this invocation, the start info the end hook derives + from, the operations provider, the dispatch pool -- describes one invocation, + so a single instance shared by two of them would let each overwrite the + other's state. Concurrent executions in one environment (Lambda Managed + Instances) are exactly that case: they run in separate threads of one + process, against one decorated handler. :class:`PluginHost` therefore builds + a fresh executor per invocation and holds it only in that invocation's frame. + + Single-use by construction: :meth:`run` refuses a second entry, so the + lifetime is an invariant of the class rather than a convention its callers + have to keep. Only the factory list is handler-lifetime, and it is copied in + rather than shared mutably. + """ + + def __init__(self, plugins: list[DurableInstrumentationPluginFactory] | None): + # Factories outlive this executor -- the list is copied, never aliased. + # The instances they build do not: _plugins is populated in + # on_invocation_start and emptied when the invocation scope exits. + self._plugin_factories = list(plugins or []) + self._plugins: list[DurableInstrumentationPlugin] = [] + # The subset of _plugins whose invocation-start hook has been dispatched. + # Every later hook is dispatched to this list, so a plugin that never + # received its start hook never receives its end hook. + self._started: list[DurableInstrumentationPlugin] = [] self._invocation_status: InvocationStartInfo | None = None self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None + self._run_entered = False @contextlib.contextmanager def run(self): - if self._plugins: - self._executor = ThreadPoolExecutor( - max_workers=1, - thread_name_prefix="plugin-executor", + """Open this executor's one invocation scope. + + Raises: + RuntimeError: if entered more than once. A second entry would mean an + executor is serving two invocations, which is the shape this + class exists to prevent; failing loudly here keeps the bug from + reappearing as silent crosstalk. + """ + if self._run_entered: + msg = ( + "PluginExecutor.run() is single-use: this executor has already " + "served an invocation. Build one executor per invocation." ) + raise RuntimeError(msg) + self._run_entered = True try: yield finally: self._invocation_status = None self._operations_provider = None - # Shut down the thread pool, waiting for pending tasks to complete. - if self._executor: - self._executor.shutdown(wait=True) + # Drop this invocation's plugin instances: nothing outlives the + # invocation. Every dispatch is synchronous, so there is no queued + # work still holding one. + self._plugins = [] + self._started = [] + + def _create_plugins(self, info: InvocationStartInfo) -> None: + """Build this invocation's plugin instances from its start info. + + Called once per invocation, before the first hook is dispatched. A + factory whose ``create_plugin`` raises or returns ``None`` is contained + exactly as a failing hook is -- logged and skipped -- so a broken plugin + cannot disrupt the execution. The remaining factories still produce their + instances. + + An entry without a usable ``create_plugin`` raises ``AttributeError`` + here, which this containment then swallows once per invocation. + :func:`plugin_discovery.load_configured_plugins` rejects such an entry + while the handler is being initialized, so the silent case is not + reachable through ``durable_execution()``. + + Containment covers every ``BaseException`` except the parts that instruct + the calling thread to stop; see + :data:`_PLUGIN_THREAD_CONTROL_EXCEPTIONS` and + :func:`_contain_plugin_failure`. Narrowing it to ``Exception`` left the + contract conditional on a factory never raising outside that hierarchy, + and a factory that awaits a cancelled task raises + ``asyncio.CancelledError``, which is outside it. + """ + plugins: list[DurableInstrumentationPlugin] = [] + for factory in self._plugin_factories: + try: + plugin = factory.create_plugin(info) + except BaseException as error: # noqa: BLE001 - a factory must not fail the execution + control = _contain_plugin_failure( + error, "Plugin factory %s exception ignored", _factory_name(factory) + ) + if control is not None: + raise control from None + continue + if plugin is None: + logger.error( + "Plugin factory %s returned None; plugin ignored", + _factory_name(factory), + ) + continue + # The load-time shape check can only establish that the factory has + # a callable create_plugin; what that call returns is knowable only + # here. A value that is not a plugin fails every hook inside + # _dispatch_plugin, so registering it would produce one logged error + # per hook per invocation for the life of the function while + # providing no telemetry. Reject it once instead. + if not isinstance(plugin, DurableInstrumentationPlugin): + logger.error( + "Plugin factory %s returned %s, which is not a " + "DurableInstrumentationPlugin; plugin ignored", + _factory_name(factory), + _type_name(plugin), + ) + continue + plugins.append(plugin) + self._plugins = plugins @staticmethod def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: - """Invoke the appropriate plugin callback. Runs inside the thread pool.""" + """Invoke the appropriate plugin callback. Runs inside the thread pool. + + Contains every ``BaseException`` except the parts that instruct the calling + thread to stop, the same rule the factory boundary uses. The thread here + is the executor's own single worker, which nothing outside this class + cancels or interrupts, so an exception outside the ``Exception`` hierarchy + arriving here was raised by the plugin. + """ try: match info: case InvocationStartInfo(): @@ -504,20 +759,73 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: plugin.on_user_function_end(info) case _: raise RuntimeError(f"Unknown info type: {type(info)}") - except Exception: - # log and ignore the exception - logger.exception("Plugin %s exception ignored", plugin.__class__.__name__) - - def execute_plugins(self, info, sync): - if not self._executor: + except BaseException as error: # noqa: BLE001 - a hook must not fail the execution + control = _contain_plugin_failure( + error, "Plugin %s exception ignored", plugin.__class__.__name__ + ) + if control is not None: + raise control from None + + def execute_plugins(self, info): + """Dispatch one hook to this invocation's plugins. + + A plugin receives a hook only once it has received the invocation-start + hook, which makes the pairing an invariant rather than a coincidence. + Without it one dispatch order breaks the pairing: a start hook that raises + one of the three exceptions :data:`_PLUGIN_THREAD_CONTROL_EXCEPTIONS` + names propagates out of this loop, so plugins later in the list never + receive their start hook -- and the invocation-end hook that the + propagating exception then triggers used to reach them anyway, leaving a + plugin to tear down state it had never been told to build. + + A plugin is counted as started before its start hook is dispatched rather + than after, because a hook that begins and then fails may already have + allocated what its end hook releases. + + The invocation-end hook is the one hook that finishes dispatching even + when a plugin raises one of those. Every plugin it reaches has already + started, so cutting the loop short costs a plugin its only chance to + finish: Insight would not drain, and OTel would leave spans unended. The + first such exception is held and re-raised once every plugin has been + called, so the thread still stops and nothing is swallowed. No other hook + defers: stopping a start-hook loop early leaves later plugins with nothing + to clean up, because the pairing rule above then withholds their end hook + too. + + Anything :meth:`_dispatch_plugin` raises is already a thread-control + failure -- it contains everything else -- so the end path catches + ``BaseException`` rather than naming the three again. Naming them would + miss a :class:`BaseExceptionGroup` carrying one, which is what + :func:`_contain_plugin_failure` hands back. + + Every hook is dispatched on the calling thread. That is what lets a + plugin set a ``ThreadLocal`` or an MDC key the SDK's own logging then + reads, and it is what makes the pairing and re-raise rules above + enforceable: a hook dispatched to a pool would land in a + :class:`~concurrent.futures.Future` nobody reads, so a control exception + raised there would be swallowed and the end-hook fan-out could not hold + it. An earlier ``sync`` parameter offered the pool path; no caller ever + passed it, and it is removed rather than left as a way to opt out of + those rules. + """ + if not self._plugin_factories: return - for plugin in self._plugins: - if sync: - # this is called synchronously, so plugins will be able to manipulate thread local objects + starting = isinstance(info, InvocationStartInfo) + ending = isinstance(info, InvocationEndInfo) + deferred_control: BaseException | None = None + for plugin in self._plugins if starting else self._started: + if starting: + self._started.append(plugin) + if not ending: + self._dispatch_plugin(plugin, info) + continue + try: self._dispatch_plugin(plugin, info) - else: - # this is called asynchronously, so plugins cannot manipulate thread local objects - self._executor.submit(self._dispatch_plugin, plugin, info) + except BaseException as control: # noqa: BLE001 - held and re-raised below + if deferred_control is None: + deferred_control = control + if deferred_control is not None: + raise deferred_control def _snapshot_operation_infos( self, @@ -535,11 +843,13 @@ def _snapshot_operation_infos( plugin that stashes the info and reads it later still sees the state as of its own hook. - Skipped entirely when no plugins are registered -- ``durable_execution()`` + Skipped entirely when no plugins are configured -- ``durable_execution()`` passes a provider unconditionally, so without this gate a plugin-free - execution would pay for a view nothing can read. + execution would pay for a view nothing can read. The gate reads the + factory list, not the instances: this runs while the start info is being + built, before any instance exists. """ - if not self._plugins or operations_provider is None: + if not self._plugin_factories or operations_provider is None: return {} try: return _to_operation_info_map(operations_provider()) @@ -572,7 +882,9 @@ def on_invocation_start( ``UpdatedOperationIds`` -- those updated while suspended. """ aws_request_id = lambda_context.aws_request_id if lambda_context else None - self._operations_provider = operations_provider if self._plugins else None + self._operations_provider = ( + operations_provider if self._plugin_factories else None + ) operations = self._snapshot_operation_infos(operations_provider) self._invocation_status = InvocationStartInfo( execution_arn=execution_arn, @@ -587,7 +899,10 @@ def on_invocation_start( if operation_id in operations }, ) - self.execute_plugins(self._invocation_status, sync=True) + # Build this invocation's plugin instances from the very info their first + # hook receives, and before that hook is dispatched. + self._create_plugins(self._invocation_status) + self.execute_plugins(self._invocation_status) def _snapshot_execution_input(self, execution_input: Any) -> Any: """Deep-copy the execution input so the plugin view is isolated. @@ -602,12 +917,12 @@ def _snapshot_execution_input(self, execution_input: Any) -> Any: The copy is eager rather than deferred: the handler starts running immediately after this hook, so a lazily-taken snapshot could already have observed the handler's mutations. It is skipped when no plugins are - registered, so non-plugin executions pay nothing. + configured, so non-plugin executions pay nothing. The snapshot is shared by all plugins for this invocation; plugins should still treat it as read-only with respect to each other. """ - if not self._plugins or execution_input is None: + if not self._plugin_factories or execution_input is None: return execution_input try: return copy.deepcopy(execution_input) @@ -635,7 +950,7 @@ def on_invocation_end( operations=self._snapshot_operation_infos(self._operations_provider), ) ) - self.execute_plugins(invocation_end_info, sync=True) + self.execute_plugins(invocation_end_info) def on_user_function_start( self, @@ -656,7 +971,7 @@ def on_user_function_start( is_replay_children=is_replay_children, attempt=attempt, ) - self.execute_plugins(start_info, sync=True) + self.execute_plugins(start_info) return start_info def on_user_function_end( @@ -669,7 +984,6 @@ def on_user_function_end( """Execute plugins when a user function returns, fails, or is incomplete.""" self.execute_plugins( UserFunctionEndInfo.from_start_info(start_info, error, outcome=outcome), - sync=True, ) def on_operation_action( @@ -699,7 +1013,6 @@ def on_operation_action( is_replayed=previous_operation is not None, status=OperationStatus.STARTED, ), - sync=True, ) def on_operation_replay(self, operation: Operation) -> None: @@ -717,7 +1030,7 @@ def on_operation_replay(self, operation: Operation) -> None: is_replayed=True, status=operation.status, ) - self.execute_plugins(start_info, sync=True) + self.execute_plugins(start_info) def on_child_context_end( self, @@ -742,7 +1055,6 @@ def on_child_context_end( error=error, is_replayed=is_replayed, ), - sync=True, ) def on_operation_update( @@ -792,7 +1104,6 @@ def on_operation_update( ), is_replayed=False, ), - sync=True, ) if ( @@ -820,7 +1131,6 @@ def on_operation_update( }, operations=_to_operation_info_map(operations), ), - sync=True, ) @staticmethod @@ -837,26 +1147,117 @@ def _is_terminal_status(status): OperationStatus.STOPPED, ] + +class PluginHost: + """Handler-lifetime owner of the configured plugin factories. + + The factory list is the only plugin state that may span invocations: a + factory is resolved once when the handler is initialized and is, by + definition, environment-lifetime. Everything a factory produces is + invocation-lifetime, so this class never holds an instance, a start info or a + dispatch pool -- :meth:`invocation` hands out a fresh + :class:`PluginExecutor` and the caller keeps it in the invocation's own + frame. + + The handler holds factories and the invocation holds instances. The other + SDKs are moving to the same split, in aws/aws-durable-execution-sdk-js#924 + (``createInvocationPluginRunner``) and aws/aws-durable-execution-sdk-java#721 + (``PluginRunner`` constructed from factories). Both are open pull requests, so + neither shape is on those repositories' default branches: ``createPluginRunner`` + on JS and ``PluginRunner`` on Java both still hold plugin instances directly. + """ + + def __init__(self, plugins: list[DurableInstrumentationPluginFactory] | None): + self._plugin_factories = list(plugins or []) + + @contextlib.contextmanager + def invocation(self) -> Iterator[PluginExecutor]: + """Open one invocation's plugin scope and yield its executor. + + The executor is created here rather than at handler-initialization time + so that two invocations sharing this process -- concurrent executions on + a Lambda Managed Instance, or successive executions on a warm + environment -- never write to the same slot. Teardown on scope exit + touches only the executor yielded here. + """ + executor = PluginExecutor(self._plugin_factories) + with executor.run(): + yield executor + @property def handle_durable_output(self): - def decorator(func: Callable[[Any, LambdaContext], MutableMapping[str, Any]]): + """Wrap an invocation body so plugins see its outcome. + + The wrapped function receives this invocation's :class:`PluginExecutor` + as a third argument. Passing it in, rather than closing over one, is what + keeps the instances out of handler-lifetime state: the executor is + reachable only from the frames of the invocation it belongs to. + """ + + def decorator( + func: Callable[ + [Any, LambdaContext, PluginExecutor], MutableMapping[str, Any] + ], + ): @functools.wraps(func) def wrapper(event: Any, context: LambdaContext): - with self.run(): + with self.invocation() as plugin_executor: + # The end hook is dispatched exactly once per invocation, so + # the success dispatch sits outside the try. Inside it, an + # end hook that raised would be caught as though the handler + # had failed, and the hook would run a second time with a + # RETRY outcome -- telling later plugins the wrong thing about + # an invocation that succeeded, and letting an exporter export + # twice. A hook can raise: _dispatch_plugin re-raises the three + # exceptions that instruct the calling thread to stop, and + # from_dict below can reject an output the handler built. try: - output = func(event, context) - - self.on_invocation_end( - output=DurableExecutionInvocationOutput.from_dict(output), - ) - return output - except Exception as e: - self.on_invocation_end( - output=DurableExecutionInvocationOutput.create_retry( - ErrorObject.from_exception(e) - ), - ) + output = func(event, context, plugin_executor) + completed = DurableExecutionInvocationOutput.from_dict(output) + except BaseException as e: + # Every exit fires the end hook, not only the ones that + # derive from Exception. A handler that surfaces an + # asyncio.CancelledError -- user code that awaited a + # cancelled task, most simply -- leaves the invocation by + # a BaseException, and an invocation that ends without + # its end hook costs the plugins the only point at which + # they can finish: Insight never drains, so the records it + # holds for this execution are dropped, and OTel never + # ends the spans it opened, so they are never exported. + # The teardown below still runs either way, which is why + # the gap was silent rather than a leak. + # + # KeyboardInterrupt and SystemExit reach here too, and + # they also fire the hook. The hook is what a plugin needs + # to flush, and a process being torn down is when flushing + # matters; the cost is the same bounded work any + # invocation end does. + # + # The handler's exception is what the caller sees, + # whatever the hook does. The end-hook dispatch can raise + # -- it re-raises the control exceptions it holds through + # the fan-out -- and letting that replace the handler's + # failure would report an instrumentation problem as the + # execution's outcome and leave the real failure reachable + # only as __context__. Instrumentation does not decide + # what an execution failed with, so the hook's exception + # is contained here and the original is re-raised + # unchanged. + try: + plugin_executor.on_invocation_end( + output=DurableExecutionInvocationOutput.create_retry( + ErrorObject.from_exception(e) + ), + ) + except BaseException: # noqa: BLE001 - the handler's failure wins + logger.exception( + "Plugin invocation-end hook failed while the " + "invocation was already failing; the original " + "failure is raised" + ) raise + plugin_executor.on_invocation_end(output=completed) + return output return wrapper diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py index ffb07260..281e7275 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin_discovery.py @@ -1,16 +1,15 @@ from __future__ import annotations +import inspect import logging import os from collections.abc import Mapping, Sequence from importlib import metadata +from typing import cast -from aws_durable_execution_sdk_python.__about__ import __version__ from aws_durable_execution_sdk_python.exceptions import PluginLoadError from aws_durable_execution_sdk_python.plugin import ( - DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION, - DurableInstrumentationPlugin, - DurableInstrumentationPluginProvider, + DurableInstrumentationPluginFactory, ) @@ -19,6 +18,10 @@ PLUGIN_ENTRY_POINT_GROUP = "aws_durable_execution.plugins" PLUGIN_ENVIRONMENT_VARIABLE = "DURABLE_EXECUTION_PLUGINS" +# Stands in for the InvocationStartInfo when a factory's signature is checked at +# registration. Only the bind is performed, so nothing reads it. +_ARGUMENT_PROBE = object() + def _parse_configured_plugin_names(environment: Mapping[str, str]) -> list[str]: configured_plugins = environment.get(PLUGIN_ENVIRONMENT_VARIABLE) @@ -51,103 +54,246 @@ def _distribution_name(entry_point: metadata.EntryPoint) -> str: return distribution.metadata.get("Name", "unknown distribution") -def _qualified_type_name(value: object) -> str: - value_type = type(value) - return f"{value_type.__module__}.{value_type.__qualname__}" +def _module_qualified_name(named: object) -> str: + """Join a class's or function's module and qualified name. + Both attributes are read with a default, because a rejected entry can be any + object at all. A missing name must not replace the configuration error with an + ``AttributeError``. + """ + module = getattr(named, "__module__", None) or "unknown module" + qualified = getattr(named, "__qualname__", None) or getattr(named, "__name__", None) + return f"{module}.{qualified or ''}" -def _qualified_class_name(value_type: type[object]) -> str: - return f"{value_type.__module__}.{value_type.__qualname__}" +def _describe_value(value: object) -> str: + """Describe a rejected registration entry so the caller can identify it. -def _load_provider( - plugin_name: str, entry_point: metadata.EntryPoint -) -> DurableInstrumentationPluginProvider: - try: - provider = entry_point.load() - except Exception as error: - raise PluginLoadError( - f"Failed to load durable instrumentation plugin provider " - f"'{plugin_name}' from '{entry_point.value}' " - f"({_distribution_name(entry_point)}): {error}" - ) from error + Two kinds of value are named by themselves rather than by their type. A + class's type is its metaclass, which is ``builtins.type`` for an ordinary + class. A function's type is ``builtins.function``. Neither of those two names + says which value was passed. - if not isinstance(provider, DurableInstrumentationPluginProvider): - raise PluginLoadError( - f"Durable instrumentation plugin entry point '{plugin_name}' must " - "resolve to DurableInstrumentationPluginProvider, but resolved to " - f"{_qualified_type_name(provider)}." - ) + Both of the likeliest migration mistakes are classes: ``plugins=[MyPlugin]`` + is the shape the previous major accepted, and ``plugins=[MyPluginFactory]`` is + this major's shape with the parentheses left off. Naming the type would report + ``builtins.type`` for both. So a class and a function are named directly, and + every other value is named by its type. - if provider.plugin_api_version != DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION: - raise PluginLoadError( - f"Durable instrumentation plugin provider '{plugin_name}' declares " - f"plugin API version {provider.plugin_api_version}, but " - f"aws-durable-execution-sdk-python {__version__} supports plugin API " - f"version {DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION}. Install " - "compatible SDK and plugin package versions." - ) + The kind is named alongside the name -- "the class", "the function", "an + instance of" -- because a factory class and an instance of that factory class + share one qualified name, and passing the class where an instance is required + is itself one of the rejected shapes. + """ + if isinstance(value, type): + return f"the class {_module_qualified_name(value)}" + if inspect.isroutine(value): + return f"the function {_module_qualified_name(value)}" + return f"an instance of {_module_qualified_name(type(value))}" + + +def _is_plugin_factory(value: object) -> bool: + """Report whether a value has the shape of a plugin factory. + + :class:`DurableInstrumentationPluginFactory` declares one method, so the + shape is one member: a ``create_plugin`` that can be called with one + positional argument. The attribute is fetched and tested rather than merely + checked for presence, because an object carrying a non-callable + ``create_plugin`` would otherwise pass here and fail at invocation time. + + Callability alone is not enough. ``plugins=[MyFactory]`` -- the factory + *class* rather than an instance of it -- resolves ``create_plugin`` to a + plain function whose first parameter is ``self``, which is callable. The + per-invocation call supplies only the info, Python binds it to ``self``, and + the resulting :exc:`TypeError` is contained like any other factory failure: + telemetry is silently absent for the lifetime of the function. Two checks + reject that at registration instead. For a class, the *kind* of the member + decides, because a signature bind cannot tell ``create_plugin(self, info)`` + from ``create_plugin(info)``: see :func:`_is_unbound_instance_method`. For + everything else, binding one positional argument to the signature rejects a + member that cannot receive the info. Neither check runs factory code. + + A callable with no introspectable signature -- a C-implemented callable, for + example -- is accepted on the member alone. ``inspect.signature`` raises for + it, and refusing a factory because its signature could not be read would + reject a usable factory over a missing description of it. + + Structural rather than nominal, so a factory need not import the SDK + protocol to satisfy it. The protocol is deliberately not + ``@runtime_checkable``; see its docstring. + + The check stops there. Whether ``create_plugin`` returns a plugin is only + knowable by calling it, and calling it at load time is what the + per-invocation factory design avoids: there is no invocation yet. That case + is checked per invocation by :meth:`PluginExecutor._create_plugins`. + """ + create_plugin = getattr(value, "create_plugin", None) + if not callable(create_plugin): + return False + if isinstance(value, type) and _is_unbound_instance_method(value, create_plugin): + return False + return _accepts_one_positional_argument(create_plugin) + + +def _is_unbound_instance_method(cls: type, create_plugin: object) -> bool: + """Report whether a class's ``create_plugin`` is an instance method. + + Read off the class, an instance method is a plain function whose first + parameter is ``self``, so the per-invocation call binds the info to ``self`` + and the factory never sees it. The signature bind below cannot catch every + such shape: ``create_plugin(self, info=None)`` and + ``create_plugin(self, *args)`` both bind one argument to ``self`` and leave + the rest satisfied. The kind of the member decides it instead. + + Three shapes read off a class are usable and none of them is a plain + function. A ``@classmethod`` is already bound to the class, so it carries + ``__self__``. A ``@staticmethod`` is a plain function, but its descriptor says + it takes no implicit first argument. And an attribute holding a callable + object -- ``create_plugin = SomeCallable()`` -- is not a function at all and + takes no implicit first argument either. Anything else read off a class takes + ``self`` and cannot serve. + + :func:`inspect.getattr_static` is what distinguishes the ``@staticmethod``, + because it returns the descriptor rather than what reading the attribute + produces. It walks the MRO without running any descriptor, so no factory code + runs here. + """ + if getattr(create_plugin, "__self__", None) is not None: + return False + if not inspect.isfunction(create_plugin): + return False + try: + declared = inspect.getattr_static(cls, "create_plugin") + except AttributeError: + return False + return not isinstance(declared, staticmethod) - declared_plugin_type: object = provider.plugin_type - if not isinstance(declared_plugin_type, type) or not issubclass( - declared_plugin_type, DurableInstrumentationPlugin - ): - declared_type_name = ( - _qualified_class_name(declared_plugin_type) - if isinstance(declared_plugin_type, type) - else _qualified_type_name(declared_plugin_type) - ) - raise PluginLoadError( - f"Durable instrumentation plugin provider '{plugin_name}' declares " - f"invalid plugin type {declared_type_name}; " - "expected a DurableInstrumentationPlugin subclass." - ) - return provider +def _accepts_one_positional_argument(create_plugin: object) -> bool: + """Report whether one positional argument can be bound to a callable. + + A bound method, a ``@classmethod`` or ``@staticmethod`` read off a class, and + a ``__call__`` on an instance all present the signature the SDK calls, so all + three bind. An instance method read off the class does not: its first + parameter is ``self``, so one argument leaves the info unbound. + """ + try: + signature = inspect.signature(create_plugin) # type: ignore[arg-type] + except (TypeError, ValueError): + return True + try: + signature.bind(_ARGUMENT_PROBE) + except TypeError: + return False + return True -def _create_plugin( - plugin_name: str, - entry_point: metadata.EntryPoint, - provider: DurableInstrumentationPluginProvider, -) -> DurableInstrumentationPlugin: +def _load_factory( + plugin_name: str, entry_point: metadata.EntryPoint +) -> DurableInstrumentationPluginFactory: + """Resolve an entry point to a plugin factory. + + Only two things can still be checked here. The entry point has to import, + and what it resolves to has to have the factory shape. Nothing more is + knowable without calling the factory, and calling it at load time is + precisely what this design avoids: the instance belongs to an invocation, + and there is no invocation yet. A factory that then misbehaves at invocation + time is contained by :meth:`PluginExecutor._create_plugins`. + """ try: - plugin = provider.factory() + factory = entry_point.load() except Exception as error: raise PluginLoadError( - f"Failed to create durable instrumentation plugin '{plugin_name}' " - f"from '{entry_point.value}' ({_distribution_name(entry_point)}): " - f"{error}" + f"Failed to load durable instrumentation plugin factory " + f"'{plugin_name}' from '{entry_point.value}' " + f"({_distribution_name(entry_point)}): {error}" ) from error - if type(plugin) is not provider.plugin_type: + if not _is_plugin_factory(factory): raise PluginLoadError( - f"Durable instrumentation plugin provider '{plugin_name}' returned " - f"{_qualified_type_name(plugin)}; expected " - f"{_qualified_class_name(provider.plugin_type)}." + f"Durable instrumentation plugin entry point '{plugin_name}' must " + "resolve to a plugin factory -- an object with a " + "create_plugin(info) method returning a " + "DurableInstrumentationPlugin -- but resolved to " + f"{_describe_value(factory)}. Name the factory instance, not a " + "plugin, not a plugin class, and not the factory class." ) - return plugin + return cast(DurableInstrumentationPluginFactory, factory) + + +def _validate_explicit_factories( + explicit_plugins: Sequence[DurableInstrumentationPluginFactory] | None, +) -> list[DurableInstrumentationPluginFactory]: + """Check that every explicitly passed plugin entry has the factory shape. + + Each entry's ``create_plugin`` is called once per invocation to build that + invocation's plugin instance. An entry without one can never be called, so + :meth:`PluginExecutor._create_plugins` raises ``AttributeError`` on every + invocation, logs it and continues without that plugin -- telemetry is lost + for the lifetime of the function, and nothing fails. Raising here converts + that into one configuration failure while the handler is being initialized. + The position is named because a caller passing several entries cannot + otherwise tell which one is wrong. The entry itself is named too, by + :func:`_describe_value`, which names a class and a function directly rather + than by type: the type of a class is ``builtins.type``, and that would + identify no particular class. + + A plugin *class* is rejected, and so is any bare callable. Both were + accepted while the registration type was ``Callable``: a lambda satisfied it + directly, and a class satisfied it because calling a class constructs an + instance. Neither carries a ``create_plugin`` the SDK can call, so + ``plugins=[MyPlugin]`` and ``plugins=[lambda info: MyPlugin()]`` now fail + here. A *factory* class passed instead of an instance of it fails here too: + ``MyFactory.create_plugin`` is callable, but its first parameter is ``self``, + so the per-invocation call binds the info to ``self``. The replacement is a + small factory class, instantiated, which is also where setup work that can + fail belongs. A class that declares ``create_plugin`` as a ``@classmethod`` + or a ``@staticmethod`` is accepted, because that member presents the + signature the SDK calls. + """ + factories = list(explicit_plugins or []) + for index, factory in enumerate(factories): + if not _is_plugin_factory(factory): + raise PluginLoadError( + f"Durable instrumentation plugin at plugins[{index}] must be a " + "plugin factory -- an object with a create_plugin(info) method " + "returning a DurableInstrumentationPlugin -- but is " + f"{_describe_value(factory)}. Pass a factory rather than a " + "plugin, a plugin class, or a plain callable, and pass a factory " + "instance rather than the factory class, for example " + "plugins=[MyPluginFactory(exporter)]." + ) + return factories def load_configured_plugins( - explicit_plugins: Sequence[DurableInstrumentationPlugin] | None, + explicit_plugins: Sequence[DurableInstrumentationPluginFactory] | None, *, environment: Mapping[str, str] | None = None, -) -> list[DurableInstrumentationPlugin]: - """Combine explicit plugins with providers selected through the environment. - - Explicit plugins retain their order. Dynamically selected plugins follow in - configured order. When discovery creates a plugin whose concrete type is - already registered, the first registration wins, so explicit registration - takes precedence. +) -> list[DurableInstrumentationPluginFactory]: + """Combine explicit plugin factories with those selected through the environment. + + Explicit factories retain their order. Dynamically selected factories follow + in configured order. Every returned factory has its ``create_plugin`` called + once per invocation. + + A factory already registered explicitly is not registered a second time + through the environment. The check is by factory identity, which is what is + knowable here: the old shape declared a ``plugin_type`` and could dedup on + it, but what a factory builds is unknown until ``create_plugin`` is called, + and calling it at load time is what this design avoids. Identity still covers + the case the plugin packages document -- the same factory object both passed + to the decorator and named in ``DURABLE_EXECUTION_PLUGINS``. Two *different* + factories that happen to build the same plugin type will now both be + registered. """ - resolved_plugins = list(explicit_plugins or []) + resolved_factories = _validate_explicit_factories(explicit_plugins) resolved_environment = os.environ if environment is None else environment plugin_names = _parse_configured_plugin_names(resolved_environment) if not plugin_names: - return resolved_plugins + return resolved_factories try: discovered_entry_points = list( @@ -163,10 +309,6 @@ def load_configured_plugins( for entry_point in discovered_entry_points: entry_points_by_name.setdefault(entry_point.name, []).append(entry_point) - registered_types: dict[type[DurableInstrumentationPlugin], str] = { - type(plugin): "the decorator's plugins argument" for plugin in resolved_plugins - } - for plugin_name in plugin_names: matching_entry_points = entry_points_by_name.get(plugin_name, []) if not matching_entry_points: @@ -190,20 +332,15 @@ def load_configured_plugins( "duplicate provider package." ) - entry_point = matching_entry_points[0] - provider = _load_provider(plugin_name, entry_point) - if existing_registration := registered_types.get(provider.plugin_type): + factory = _load_factory(plugin_name, matching_entry_points[0]) + if any(factory is registered for registered in resolved_factories): logger.warning( - "Skipping dynamically configured plugin '%s' because %s is " - "already registered by %s.", + "Skipping dynamically configured plugin '%s' because the same " + "plugin factory is already registered.", plugin_name, - _qualified_class_name(provider.plugin_type), - existing_registration, ) continue - plugin = _create_plugin(plugin_name, entry_point, provider) - resolved_plugins.append(plugin) - registered_types[provider.plugin_type] = f"dynamic provider '{plugin_name}'" + resolved_factories.append(factory) - return resolved_plugins + return resolved_factories diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py index f5ce7214..42200017 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py @@ -374,6 +374,9 @@ def __init__( self._current_checkpoint_token: str = initial_checkpoint_token self._operations: dict[str, Operation] = dict(operations) self._service_client: DurableServiceClient = service_client + # Invocation-scoped, like this state object itself: PluginHost builds one + # executor per invocation, so holding it here cannot make one + # invocation's plugin instances visible to another. self._plugin_executor: PluginExecutor = plugin_executor self._operations_lock: Lock = Lock() diff --git a/packages/aws-durable-execution-sdk-python/tests/context_test.py b/packages/aws-durable-execution-sdk-python/tests/context_test.py index 1b75325b..110ef919 100644 --- a/packages/aws-durable-execution-sdk-python/tests/context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/context_test.py @@ -63,7 +63,11 @@ WaitForConditionDecision, ) from tests.serdes_test import CustomDictSerDes -from tests.test_helpers import operation_id_sequence +from tests.test_helpers import ( + operation_id_sequence, + plugin_factory, + plugin_invocation, +) def create_test_context( @@ -3014,9 +3018,9 @@ def on_operation_start(self, info): def on_operation_end(self, info): captured.append(f"end:{info.operation_id}") - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) step_body_calls: list[bool] = [] - with plugin_executor.run(): + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="arn", initial_checkpoint_token="token", # noqa: S106 @@ -3062,8 +3066,8 @@ def on_operation_start(self, info): def on_operation_end(self, info): captured.append(("end", info.operation_id, info.is_replayed, info.status)) - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="arn", initial_checkpoint_token="token", # noqa: S106 @@ -3100,8 +3104,8 @@ def on_operation_start(self, info): def on_operation_end(self, info): captured.append(("end", info.operation_id, info.is_replayed, info.status)) - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="arn", initial_checkpoint_token="token", # noqa: S106 diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/handler_context_propagation_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/handler_context_propagation_int_test.py new file mode 100644 index 00000000..37e2614b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/handler_context_propagation_int_test.py @@ -0,0 +1,263 @@ +"""Integration tests for context propagation into the user-function thread. + +The SDK runs the handler body on a worker thread it submits to a pool. A +submitted callable is given a context of its own, so anything a contextvar +carries -- set by a plugin in ``on_invocation_start``, or by customer code +around the decorator -- is only visible to top-level handler code if the +invocation thread's context is carried into that worker. + +Log-correlating plugins depend on this: a plugin that claims the invocation for +the calling thread at invocation start has no hook that runs on the worker +before the handler body, so the claim reaches top-level handler code by context +propagation or not at all. +""" + +from __future__ import annotations + +import contextvars +import threading +from dataclasses import replace +from datetime import UTC, datetime +from typing import Any +from unittest.mock import Mock, patch + +from aws_durable_execution_sdk_python.context import DurableContext, durable_step +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + Operation, + OperationAction, + OperationStatus, + OperationType, + StepDetails, +) +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from tests.test_helpers import plugin_factory + + +UNSET = "unset" +_probe: contextvars.ContextVar[str] = contextvars.ContextVar( + "handler_context_propagation_probe", default=UNSET +) + + +def _lambda_context() -> Mock: + ctx = Mock() + ctx.aws_request_id = "test-request-id" + ctx.client_context = None + ctx.identity = None + ctx._epoch_deadline_time_in_ms = 0 # noqa: SLF001 + ctx.invoked_function_arn = "test-arn" + ctx.tenant_id = None + return ctx + + +def _event() -> dict: + return { + "DurableExecutionArn": "test-arn/execution-1", + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [ + { + "Id": "execution-1", + "Type": "EXECUTION", + "Status": "STARTED", + "ExecutionDetails": {"InputPayload": "{}"}, + } + ], + "NextMarker": "", + }, + "LocalRunner": True, + } + + +def _tracking_checkpoint(): + """Checkpoint mock that accumulates operations, as the service would. + + SUCCEED actions are recorded as SUCCEEDED so the SDK dispatches + operation-end hooks, which is how a hook reaches the background + checkpointing thread. + """ + operations: dict[str, Operation] = {} + + def mock_checkpoint( + durable_execution_arn, # noqa: ARG001 + checkpoint_token, # noqa: ARG001 + updates, + client_token="token", # noqa: S107, ARG001 + ) -> CheckpointOutput: + for update in updates: + previous = operations.get(update.operation_id) + base = previous or Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + name=update.name, + sub_type=update.sub_type, + start_timestamp=datetime.now(UTC), + ) + if update.action is OperationAction.SUCCEED: + operations[update.operation_id] = replace( + base, + status=OperationStatus.SUCCEEDED, + end_timestamp=datetime.now(UTC), + step_details=( + StepDetails(result=update.payload, attempt=1) + if update.operation_type is OperationType.STEP + else base.step_details + ), + ) + else: + operations[update.operation_id] = base + return CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState( + operations=list(operations.values()) + ), + ) + + return mock_checkpoint + + +def _run(handler, event: dict | None = None) -> dict: + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = _tracking_checkpoint() + mock_client_class.initialize_client.return_value = mock_client + + return handler(event if event is not None else _event(), _lambda_context()) + + +class _ContextvarSettingPlugin(DurableInstrumentationPlugin): + """Sets a contextvar on the invocation thread, as a log-correlating plugin does.""" + + def __init__(self, value: str) -> None: + self._value = value + self.operation_end_observations: list[tuple[str, str]] = [] + self._lock = threading.Lock() + + def on_invocation_start(self, info) -> None: # noqa: ARG002 + _probe.set(self._value) + + def on_operation_end(self, info) -> None: # noqa: ARG002 + # Dispatched from the background checkpointing thread, which reads back + # the terminal status of a checkpointed operation. + with self._lock: + self.operation_end_observations.append( + (threading.current_thread().name, _probe.get()) + ) + + +def test_handler_body_sees_a_contextvar_set_by_a_plugin_at_invocation_start(): + """The invocation-start hook's contextvar reaches top-level handler code. + + The hook runs on the invocation thread and the handler body runs on a worker + thread, and no hook runs on that worker before the handler's first + statement, so this only holds if the invocation's context is propagated. + """ + plugin = _ContextvarSettingPlugin("claimed-by-plugin") + seen: list[str] = [] + + @durable_execution(plugins=[plugin_factory(plugin)]) + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + seen.append(_probe.get()) + return "ok" + + result = _run(my_handler) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + assert seen == ["claimed-by-plugin"] + + +def test_handler_body_sees_a_contextvar_set_by_the_caller(): + """A contextvar set before the handler is visible inside the handler body. + + This matches ``asyncio.to_thread``, which also runs the callable in a copy + of the caller's context. + """ + seen: list[str] = [] + + @durable_execution + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + seen.append(_probe.get()) + return "ok" + + token = _probe.set("set-by-caller") + try: + result = _run(my_handler) + finally: + _probe.reset(token) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + assert seen == ["set-by-caller"] + + +def test_handler_body_contextvar_writes_do_not_leak_to_the_caller(): + """The worker mutates its own copy, so the caller's context is untouched. + + A worker thread has a context of its own whether it starts empty or from a + copy, so propagation adds no path from the handler back to the invocation + thread. + """ + observed_in_step: list[str] = [] + + @durable_step + def read_probe(_step_context) -> str: + observed_in_step.append(_probe.get()) + return _probe.get() + + @durable_execution + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + _probe.set("set-inside-handler") + return context.step(read_probe(), name="read-probe") + + token = _probe.set("set-by-caller") + try: + result = _run(my_handler) + after = _probe.get() + finally: + _probe.reset(token) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + # The write is visible to code the handler drives, and nowhere else. + assert observed_in_step == ["set-inside-handler"] + assert after == "set-by-caller" + + +def test_checkpointing_thread_does_not_carry_the_invocation_context(): + """The background checkpointing thread is submitted without the context. + + It runs SDK checkpointing rather than user code, so it is left starting from + an empty context. This test pins that choice: a plugin hook dispatched from + that thread sees the contextvar's default, not the value the invocation + thread set. + """ + plugin = _ContextvarSettingPlugin("claimed-by-plugin") + handler_thread_names: list[str] = [] + + @durable_step + def a_step(_step_context) -> str: + return "stepped" + + @durable_execution(plugins=[plugin_factory(plugin)]) + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + handler_thread_names.append(threading.current_thread().name) + return context.step(a_step(), name="a-step") + + result = _run(my_handler) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + off_handler_thread = [ + value + for thread_name, value in plugin.operation_end_observations + if thread_name not in handler_thread_names + ] + assert off_handler_thread, plugin.operation_end_observations + assert set(off_handler_thread) == {UNSET} diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py index 3cad3412..311f8648 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py @@ -26,8 +26,11 @@ OperationStatus, OperationType, ) -from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin -from tests.test_helpers import operation_id_sequence +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStartInfo, +) +from tests.test_helpers import operation_id_sequence, plugin_factory class _MapRecordingPlugin(DurableInstrumentationPlugin): @@ -114,7 +117,7 @@ def test_operation_maps_on_a_completing_invocation(): """The start map holds the prior state; the end map sees the step added.""" plugin = _MapRecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 return context.step(lambda _ctx: "stepped", name="greet") @@ -150,17 +153,30 @@ def test_operation_maps_across_suspend_and_replay(): Invocation 1 suspends on a wait. Invocation 2 replays with the wait already SUCCEEDED and its id in ``UpdatedOperationIds``, which is exactly what ``updated_operations`` is derived from. + + One handler serves both invocations, as in production. Previously this test + had to declare a second handler with its own plugin instance, because the + single shared instance would have interleaved both invocations' records into + one list. Now the factory builds an instance per invocation, so each + instance's ``starts[0]``/``ends[0]`` unambiguously describes its own + invocation -- and the test asserts that separation directly. """ wait_id = next(operation_id_sequence()) - # --- Invocation 1: the wait starts and the execution suspends. - first = _MapRecordingPlugin() + built: list[_MapRecordingPlugin] = [] - @durable_execution(plugins=[first]) - def suspending_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + class _BuildingFactory: + def create_plugin(self, info: InvocationStartInfo) -> _MapRecordingPlugin: + plugin = _MapRecordingPlugin() + built.append(plugin) + return plugin + + @durable_execution(plugins=[_BuildingFactory()]) + def wait_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 context.wait(Duration.from_seconds(60)) return "done" + # --- Invocation 1: the wait starts and the execution suspends. with patch( "aws_durable_execution_sdk_python.execution.LambdaClient" ) as mock_client_class: @@ -168,9 +184,11 @@ def suspending_handler(event: Any, context: DurableContext) -> str: # noqa: ARG mock_client.checkpoint = _tracking_checkpoint() mock_client_class.initialize_client.return_value = mock_client - first_result = suspending_handler(_event(), _lambda_context()) + first_result = wait_handler(_event(), _lambda_context()) assert first_result["Status"] == InvocationStatus.PENDING.value + assert len(built) == 1 + first = built[0] start_operations, start_updated = first.starts[0] assert start_operations == ["execution-1"] assert start_updated == [] @@ -179,13 +197,6 @@ def suspending_handler(event: Any, context: DurableContext) -> str: # noqa: ARG assert wait_id in end_operations # --- Invocation 2: replay with the wait completed externally. - replay = _MapRecordingPlugin() - - @durable_execution(plugins=[replay]) - def replayed_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 - context.wait(Duration.from_seconds(60)) - return "done" - completed_wait = { "Id": wait_id, "Type": OperationType.WAIT.value, @@ -200,13 +211,21 @@ def replayed_handler(event: Any, context: DurableContext) -> str: # noqa: ARG00 mock_client.checkpoint = _tracking_checkpoint() mock_client_class.initialize_client.return_value = mock_client - replay_result = replayed_handler( + replay_result = wait_handler( _event(extra_operations=[completed_wait], updated_operation_ids=[wait_id]), _lambda_context(), ) assert replay_result["Status"] == InvocationStatus.SUCCEEDED.value + # The second invocation got its own instance, and the first one recorded + # nothing further after it returned. + assert len(built) == 2 + replay = built[1] + assert replay is not first + assert len(first.starts) == 1 + assert len(first.ends) == 1 + start_operations, start_updated = replay.starts[0] # The replay start map carries the prior state, including the wait. assert sorted(["execution-1", wait_id]) == start_operations @@ -222,7 +241,7 @@ def test_updated_operations_ignores_ids_absent_from_the_map(): """An id the execution state does not carry must not appear in the subset.""" plugin = _MapRecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 return "ok" diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py index 6b934423..e0c86b17 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py @@ -27,7 +27,7 @@ OperationType, ) from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin -from tests.test_helpers import operation_id_sequence +from tests.test_helpers import operation_id_sequence, plugin_factory class _PayloadRecordingPlugin(DurableInstrumentationPlugin): @@ -116,7 +116,7 @@ def test_plugin_sees_execution_input_and_result_end_to_end(): """A completing invocation surfaces the input on both hooks and the result.""" plugin = _PayloadRecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def my_handler(event: Any, context: DurableContext) -> dict: # noqa: ARG001 return {"greeting": f"Hello, {event['name']}!"} @@ -146,7 +146,7 @@ def test_plugin_payload_surfaces_on_suspending_invocation(): """A suspending invocation carries the input but no execution result.""" plugin = _PayloadRecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def my_handler(event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(60)) return f"done-{event['name']}" @@ -175,7 +175,7 @@ def test_plugin_payload_surfaces_on_replay_invocation(): """A replay past a completed wait carries the input and the terminal result.""" plugin = _PayloadRecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def my_handler(event: Any, context: DurableContext) -> str: context.wait(Duration.from_seconds(60)) return f"done-{event['name']}" @@ -232,7 +232,7 @@ def on_invocation_end(self, info) -> None: plugin = _MutatingPlugin() handler_saw: dict[str, Any] = {} - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 handler_saw.update( {"top": dict(event), "nested_items": list(event["nested"]["items"])} diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py index 7b30828d..71eaeca2 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py @@ -28,6 +28,7 @@ UserFunctionOutcome, UserFunctionStartInfo, ) +from tests.test_helpers import plugin_factory @dataclass(frozen=True) @@ -156,7 +157,7 @@ def child_function(context: DurableContext) -> str: def user_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 return context.run_in_child_context(child_function, name="charge") - handler = durable_execution(user_handler, plugins=[plugin]) + handler = durable_execution(user_handler, plugins=[plugin_factory(plugin)]) first_checkpoint, first_operations = _tracking_checkpoint() with patch( diff --git a/packages/aws-durable-execution-sdk-python/tests/execution_test.py b/packages/aws-durable-execution-sdk-python/tests/execution_test.py index 85ea72a9..7f30b255 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -1,7 +1,9 @@ """Tests for execution.""" import datetime +import inspect import json +import threading import time import warnings from collections.abc import Sequence @@ -60,6 +62,7 @@ WaitDetails, ) from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from tests.test_helpers import plugin_factory LARGE_RESULT = "large_success" * 1024 * 1024 @@ -2728,10 +2731,10 @@ def _make_invocation_input(mock_client, next_marker="", input_payload="{}"): ) -def _make_lambda_context(): +def _make_lambda_context(request_id: str = "test-request"): """Helper to create a standard mock Lambda context.""" ctx = Mock() - ctx.aws_request_id = "test-request" + ctx.aws_request_id = request_id ctx.client_context = None ctx.identity = None ctx._epoch_deadline_time_in_ms = 1000000 # noqa: SLF001 @@ -3482,24 +3485,24 @@ def on_operation_attempt_end(self, info): def test_durable_execution_loads_plugins_when_handler_is_initialized(): - """Configured plugins are resolved once while the decorator initializes.""" - explicit_plugin = _RecordingPlugin() - resolved_plugin = _RecordingPlugin() + """Configured factories are resolved once while the decorator initializes.""" + explicit_factory = plugin_factory(_RecordingPlugin()) + resolved_factory = plugin_factory(_RecordingPlugin()) with ( warnings.catch_warnings(), patch( "aws_durable_execution_sdk_python.execution.load_configured_plugins", - return_value=[explicit_plugin, resolved_plugin], + return_value=[explicit_factory, resolved_factory], ) as load_plugins, ): warnings.simplefilter("error", FutureWarning) - @durable_execution(plugins=[explicit_plugin]) + @durable_execution(plugins=[explicit_factory]) def test_handler(event: Any, context: DurableContext) -> dict: return {"result": "success"} - load_plugins.assert_called_once_with([explicit_plugin]) + load_plugins.assert_called_once_with([explicit_factory]) assert callable(test_handler) @@ -3514,7 +3517,7 @@ def test_durable_execution_with_plugins_success(): plugin = _RecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: return {"result": "success"} @@ -3545,7 +3548,7 @@ def test_durable_execution_forwards_execution_input_to_plugins(): plugin = _RecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: return {"echoed": event["name"]} @@ -3573,7 +3576,7 @@ def test_durable_execution_surfaces_empty_input_as_empty_mapping(): plugin = _RecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> str: return "ok" @@ -3608,7 +3611,7 @@ def on_invocation_start(self, info): def on_invocation_end(self, info): observed["end_input"] = dict(info.execution_input) - @durable_execution(plugins=[_MutatingPlugin()]) + @durable_execution(plugins=[plugin_factory(_MutatingPlugin())]) def test_handler(event: Any, context: DurableContext) -> dict: observed["handler_saw"] = dict(event) # Direction B: handler mutates its event after the start hook fired. @@ -3642,7 +3645,7 @@ class _NestedMutatingPlugin(DurableInstrumentationPlugin): def on_invocation_start(self, info): info.execution_input["outer"]["inner"].append("from_plugin") - @durable_execution(plugins=[_NestedMutatingPlugin()]) + @durable_execution(plugins=[plugin_factory(_NestedMutatingPlugin())]) def test_handler(event: Any, context: DurableContext) -> dict: observed["handler_saw"] = deepcopy(event) return {"ok": True} @@ -3668,7 +3671,7 @@ def test_durable_execution_with_plugins_failure(): plugin = _RecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: msg = "user error" raise ValueError(msg) @@ -3694,7 +3697,7 @@ def test_durable_execution_with_plugins_pending(): plugin = _RecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: raise SuspendExecution("test") @@ -3717,7 +3720,7 @@ def test_durable_execution_with_plugins_retryable_error(): plugin = _RecordingPlugin() - @durable_execution(plugins=[plugin]) + @durable_execution(plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: msg = "Retriable error" raise InvocationError(msg) @@ -3744,7 +3747,7 @@ def test_durable_execution_with_multiple_plugins(): plugin1 = _RecordingPlugin() plugin2 = _RecordingPlugin() - @durable_execution(plugins=[plugin1, plugin2]) + @durable_execution(plugins=[plugin_factory(plugin1), plugin_factory(plugin2)]) def test_handler(event: Any, context: DurableContext) -> dict: return {"result": "success"} @@ -3772,7 +3775,9 @@ def test_durable_execution_with_failing_plugin_does_not_break_execution(): failing_plugin = _FailingPlugin() recording_plugin = _RecordingPlugin() - @durable_execution(plugins=[failing_plugin, recording_plugin]) + @durable_execution( + plugins=[plugin_factory(failing_plugin), plugin_factory(recording_plugin)] + ) def test_handler(event: Any, context: DurableContext) -> dict: return {"result": "success"} @@ -3788,6 +3793,199 @@ def test_handler(event: Any, context: DurableContext) -> dict: assert "invocation_end:SUCCEEDED" in recording_plugin.calls +def test_durable_execution_builds_a_plugin_per_invocation(): + """One handler, two invocations, two instances -- and no crosstalk. + + This is the LMI-driven property: a plugin may hold per-execution state in + ordinary instance attributes because the instance never outlives the + invocation it was built for. + """ + mock_client = Mock(spec=DurableServiceClient) + mock_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + built: list[_RecordingPlugin] = [] + factory_arns: list[str | None] = [] + + class _BuildingFactory: + def create_plugin(self, info) -> _RecordingPlugin: + factory_arns.append(info.execution_arn) + plugin = _RecordingPlugin() + built.append(plugin) + return plugin + + @durable_execution(plugins=[_BuildingFactory()]) + def test_handler(event: Any, context: DurableContext) -> dict: + return {"result": "success"} + + for _ in range(2): + result = test_handler( + _make_invocation_input(mock_client), + _make_lambda_context(), + ) + assert result["Status"] == InvocationStatus.SUCCEEDED.value + + assert len(built) == 2 + assert built[0] is not built[1] + # The factory saw the execution it was being built for. + assert factory_arns == [ + "arn:test:execution/exec1", + "arn:test:execution/exec1", + ] + # Each instance recorded exactly one invocation's worth of hooks. + for plugin in built: + assert plugin.calls.count("invocation_start") == 1 + assert plugin.calls.count("invocation_end:SUCCEEDED") == 1 + + +class _TaggedRecordingPlugin(DurableInstrumentationPlugin): + """Records the hooks it receives, tagging each with its own identity.""" + + def __init__(self) -> None: + self.invocation_starts: list[str | None] = [] + self.invocation_ends: list[str] = [] + self.operation_names: list[str | None] = [] + + def on_invocation_start(self, info): + self.invocation_starts.append(info.request_id) + + def on_invocation_end(self, info): + self.invocation_ends.append(f"{info.request_id}:{info.status.value}") + + def on_operation_start(self, info): + self.operation_names.append(info.name) + + def on_operation_end(self, info): + self.operation_names.append(info.name) + + def on_user_function_start(self, info): + self.operation_names.append(info.name) + + def on_user_function_end(self, info): + self.operation_names.append(info.name) + + +def test_durable_execution_keeps_overlapping_invocations_isolated(): + """Two invocations in flight at once never see each other's hooks. + + This is the Lambda Managed Instances case: one decorated handler, one + process, two concurrent executions in separate threads. A barrier holds both + invocations inside their user function at the same time, so both + ``on_invocation_start`` hooks have already fired before either operation + runs, and neither invocation returns until the other has run its operation. + Running the two invocations sequentially would not pin this -- the defect it + guards against is a per-invocation slot being overwritten while both + invocations are live. + """ + mock_client = Mock(spec=DurableServiceClient) + mock_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + built: dict[str, _TaggedRecordingPlugin] = {} + built_lock = threading.Lock() + + class _BuildingFactory: + def create_plugin(self, info) -> _TaggedRecordingPlugin: + plugin = _TaggedRecordingPlugin() + with built_lock: + built[str(info.request_id)] = plugin + return plugin + + timeout = 30 + # Released only once both invocations are inside their user function, so both + # invocation-start hooks have fired and both invocations are live. + both_in_user_code = threading.Barrier(2, timeout=timeout) + # b runs its operation first; a runs its operation afterwards, while b is + # still inside its user function waiting on a. + b_ran_operation = threading.Event() + a_ran_operation = threading.Event() + + @durable_execution(plugins=[_BuildingFactory()]) + def test_handler(event: Any, context: DurableContext) -> dict: + tag = event["tag"] + both_in_user_code.wait() + if tag == "b": + context.step(lambda _: "ok", name="step-b") + b_ran_operation.set() + assert a_ran_operation.wait(timeout) + else: + assert b_ran_operation.wait(timeout) + context.step(lambda _: "ok", name="step-a") + a_ran_operation.set() + return {"result": tag} + + results: dict[str, Any] = {} + + def invoke(tag: str) -> None: + results[tag] = test_handler( + _make_invocation_input(mock_client, input_payload=f'{{"tag": "{tag}"}}'), + _make_lambda_context(request_id=f"request-{tag}"), + ) + + threads = [ + threading.Thread(target=invoke, args=(tag,), name=f"invocation-{tag}") + for tag in ("a", "b") + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=timeout) + for thread in threads: + assert not thread.is_alive(), f"{thread.name} did not finish" + + assert results["a"]["Status"] == InvocationStatus.SUCCEEDED.value + assert results["b"]["Status"] == InvocationStatus.SUCCEEDED.value + + # One instance per invocation, and each one keyed to its own request. + assert sorted(built) == ["request-a", "request-b"] + + for tag in ("a", "b"): + plugin = built[f"request-{tag}"] + # Exactly its own invocation hooks: not the other invocation's request + # id, not two copies of its own, not zero because the other invocation's + # teardown got there first. + assert plugin.invocation_starts == [f"request-{tag}"] + assert plugin.invocation_ends == [f"request-{tag}:SUCCEEDED"] + # Only its own operation. The other invocation's step ran while this one + # was live, so a shared plugin slot would show up here. + assert f"step-{tag}" in plugin.operation_names + other = "b" if tag == "a" else "a" + assert f"step-{other}" not in plugin.operation_names + + +def test_durable_execution_with_failing_plugin_factory_does_not_break_execution(): + """A factory that raises is contained exactly as a failing hook is.""" + mock_client = Mock(spec=DurableServiceClient) + mock_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + recording_plugin = _RecordingPlugin() + + class _ExplodingFactory: + def create_plugin(self, info) -> DurableInstrumentationPlugin: + raise RuntimeError("factory boom") + + @durable_execution(plugins=[_ExplodingFactory(), plugin_factory(recording_plugin)]) + def test_handler(event: Any, context: DurableContext) -> dict: + return {"result": "success"} + + result = test_handler( + _make_invocation_input(mock_client), + _make_lambda_context(), + ) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + # The other plugin is unaffected. + assert "invocation_start" in recording_plugin.calls + assert "invocation_end:SUCCEEDED" in recording_plugin.calls + + def test_durable_execution_with_no_plugins(): """Test that passing no plugins (None) works correctly.""" mock_client = Mock(spec=DurableServiceClient) @@ -3843,7 +4041,7 @@ def test_durable_execution_decorator_with_plugins_and_boto3_client(): # When using DurableExecutionInvocationInputWithClient, boto3_client is ignored # but we verify the decorator accepts both parameters - @durable_execution(boto3_client=None, plugins=[plugin]) + @durable_execution(boto3_client=None, plugins=[plugin_factory(plugin)]) def test_handler(event: Any, context: DurableContext) -> dict: return {"result": "success"} @@ -3857,3 +4055,82 @@ def test_handler(event: Any, context: DurableContext) -> dict: # endregion Plugin Integration Tests + + +# region Handler Metadata Tests + + +def _bare_decorated_handler(): + """The ``@durable_execution`` form, desugared so the test keeps both objects.""" + + def test_handler(event: Any, context: DurableContext) -> dict: + """Handler docstring.""" + return {"result": "success"} + + return durable_execution(test_handler), test_handler + + +def _plugin_decorated_handler(): + """The ``@durable_execution(plugins=[...])`` form, desugared the same way. + + This form takes a different path through the decorator: the first call + returns a ``functools.partial``, which the second call applies to the user + function. + """ + + def test_handler(event: Any, context: DurableContext) -> dict: + """Handler docstring.""" + return {"result": "success"} + + decorator = durable_execution(plugins=[plugin_factory(_RecordingPlugin())]) + return decorator(test_handler), test_handler + + +@pytest.mark.parametrize( + "build_handler", + [_bare_decorated_handler, _plugin_decorated_handler], + ids=["bare", "with_plugins"], +) +def test_decorated_handler_signature_is_event_and_context(build_handler) -> None: + """The decorated handler accepts exactly the two Lambda handler arguments. + + A Lambda runtime, or a test harness that inspects a handler before calling + it, reads ``inspect.signature()``. Any parameter reported there that the + callable does not accept makes the handler look uninvokable, or invites a + caller to pass an argument that raises. The SDK's own invocation body takes a + third argument -- this invocation's ``PluginExecutor`` -- so this pins the + public signature to the two arguments the handler really takes, for both + decorator forms. + """ + handler, _user_function = build_handler() + + signature = inspect.signature(handler) + + assert list(signature.parameters) == ["event", "context"] + assert "plugin_executor" not in signature.parameters + # Binding is the operation a signature-aware caller actually performs. + signature.bind({}, _make_lambda_context()) + + +@pytest.mark.parametrize( + "build_handler", + [_bare_decorated_handler, _plugin_decorated_handler], + ids=["bare", "with_plugins"], +) +def test_decorated_handler_reports_user_function_metadata(build_handler) -> None: + """The handler identifies itself as the user's function, not as SDK internals. + + Logging, ``help()`` and error messages read ``__name__`` and ``__doc__``. The + SDK wraps the user function twice, so without the user function's metadata at + the head of the chain the handler names an internal wrapper instead. + ``inspect.unwrap()`` reaching the user function is what makes + ``inspect.signature()`` report that function's parameters. + """ + handler, user_function = build_handler() + + assert handler.__name__ == user_function.__name__ + assert handler.__doc__ == user_function.__doc__ + assert inspect.unwrap(handler) is user_function + + +# endregion Handler Metadata Tests diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py index 51200ba3..83aac8ec 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_discovery_test.py @@ -1,18 +1,17 @@ from __future__ import annotations +import inspect import logging import os -from collections.abc import Callable -from typing import cast +import time from unittest.mock import Mock, patch import pytest from aws_durable_execution_sdk_python.exceptions import PluginLoadError from aws_durable_execution_sdk_python.plugin import ( - DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION, DurableInstrumentationPlugin, - DurableInstrumentationPluginProvider, + InvocationStartInfo, ) from aws_durable_execution_sdk_python.plugin_discovery import ( PLUGIN_ENTRY_POINT_GROUP, @@ -21,6 +20,13 @@ ) +INVOCATION_START_INFO = InvocationStartInfo( + request_id="req-1", + execution_arn="arn:exec", + is_first_invocation=True, +) + + class _PluginA(DurableInstrumentationPlugin): pass @@ -29,6 +35,34 @@ class _PluginB(DurableInstrumentationPlugin): pass +class _PluginAFactory: + def create_plugin(self, info: InvocationStartInfo) -> _PluginA: + return _PluginA() + + +class _PluginBFactory: + def create_plugin(self, info: InvocationStartInfo) -> _PluginB: + return _PluginB() + + +class _DefaultedArgumentFactory: + """Instance method whose info parameter has a default, so one argument binds.""" + + def create_plugin(self, info: InvocationStartInfo | None = None) -> _PluginA: + return _PluginA() + + +class _VariadicFactory: + """Instance method taking ``*args``, so any argument count binds.""" + + def create_plugin(self, *args: object) -> _PluginA: + return _PluginA() + + +_plugin_a_factory = _PluginAFactory() +_plugin_b_factory = _PluginBFactory() + + class _FakeDistribution: def __init__(self, name: str) -> None: self.metadata = {"Name": name} @@ -59,32 +93,10 @@ def load(self) -> object: return self._loaded_value -def _provider( - factory: Callable[[], object], - *, - plugin_type: type[DurableInstrumentationPlugin] = _PluginA, - plugin_api_version: int = DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION, -) -> DurableInstrumentationPluginProvider: - return DurableInstrumentationPluginProvider( - plugin_type=plugin_type, - factory=cast(Callable[[], DurableInstrumentationPlugin], factory), - plugin_api_version=plugin_api_version, - ) - - -def test_plugin_provider_requires_authored_api_version() -> None: - with pytest.raises(TypeError, match="plugin_api_version"): - DurableInstrumentationPluginProvider( - plugin_type=_PluginA, - factory=_PluginA, - ) # type: ignore[call-arg] - - @pytest.mark.parametrize("configured_value", [None, "", " "]) -def test_unconfigured_discovery_preserves_explicit_plugins( +def test_unconfigured_discovery_preserves_explicit_factories( configured_value: str | None, ) -> None: - explicit_plugin = _PluginA() environment = ( {} if configured_value is None @@ -95,16 +107,16 @@ def test_unconfigured_discovery_preserves_explicit_plugins( "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points" ) as entry_points: result = load_configured_plugins( - [explicit_plugin], + [_plugin_a_factory], environment=environment, ) - assert result == [explicit_plugin] + assert result == [_plugin_a_factory] entry_points.assert_not_called() def test_discovery_uses_process_environment_by_default() -> None: - entry_point = _FakeEntryPoint("a", _provider(_PluginA)) + entry_point = _FakeEntryPoint("a", _plugin_a_factory) with ( patch.dict( @@ -119,25 +131,37 @@ def test_discovery_uses_process_environment_by_default() -> None: ): result = load_configured_plugins(None) - assert len(result) == 1 - assert isinstance(result[0], _PluginA) + assert result == [_plugin_a_factory] entry_points.assert_called_once_with(group=PLUGIN_ENTRY_POINT_GROUP) -def test_discovery_preserves_configured_order() -> None: - factory_calls: list[str] = [] +def test_discovery_returns_factories_without_calling_them() -> None: + """Discovery resolves factories only; instances belong to an invocation. - def create_a() -> _PluginA: - factory_calls.append("a") - return _PluginA() + Nothing is constructed at load time, so no plugin instance exists outside + the invocation that will use it. + """ + factory = Mock() + factory.create_plugin = Mock(return_value=_PluginA()) + entry_point = _FakeEntryPoint("a", factory) - def create_b() -> _PluginB: - factory_calls.append("b") - return _PluginB() + with patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], + ): + result = load_configured_plugins( + None, + environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, + ) + assert result == [factory] + factory.create_plugin.assert_not_called() + + +def test_discovery_preserves_configured_order() -> None: entry_points = [ - _FakeEntryPoint("b", _provider(create_b, plugin_type=_PluginB)), - _FakeEntryPoint("a", _provider(create_a)), + _FakeEntryPoint("b", _plugin_b_factory), + _FakeEntryPoint("a", _plugin_a_factory), ] with patch( @@ -149,8 +173,83 @@ def create_b() -> _PluginB: environment={PLUGIN_ENVIRONMENT_VARIABLE: " a, b "}, ) - assert [type(plugin) for plugin in result] == [_PluginA, _PluginB] - assert factory_calls == ["a", "b"] + assert result == [_plugin_a_factory, _plugin_b_factory] + assert [ + type(factory.create_plugin(INVOCATION_START_INFO)) for factory in result + ] == [ + _PluginA, + _PluginB, + ] + + +def test_explicit_factories_precede_discovered_factories() -> None: + entry_point = _FakeEntryPoint("b", _plugin_b_factory) + + with patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], + ): + result = load_configured_plugins( + [_plugin_a_factory], + environment={PLUGIN_ENVIRONMENT_VARIABLE: "b"}, + ) + + assert result == [_plugin_a_factory, _plugin_b_factory] + + +def test_explicit_registration_wins_over_the_same_discovered_factory( + caplog: pytest.LogCaptureFixture, +) -> None: + """The same factory passed explicitly and named in the env registers once. + + This is the narrowed form of the old type-based precedence rule. Dedup by + declared plugin type is gone with the provider object; identity still covers + the documented double-registration case. + """ + entry_point = _FakeEntryPoint("a", _plugin_a_factory) + + with ( + patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], + ), + caplog.at_level( + logging.WARNING, + logger="aws_durable_execution_sdk_python.plugin_discovery", + ), + ): + result = load_configured_plugins( + [_plugin_a_factory], + environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, + ) + + assert result == [_plugin_a_factory] + assert "already registered" in caplog.text + + +def test_distinct_factories_for_one_plugin_type_are_both_registered() -> None: + """Type-level dedup is gone: two distinct factories both register. + + Recorded deliberately. The provider object declared a ``plugin_type`` that + discovery could compare without constructing anything; what a factory builds + is unknown until ``create_plugin`` is called, and calling it at load time + would build an instance outside any invocation. Callers that both pass a + factory and name a different one in the environment now get both plugins. + """ + another_plugin_a_factory = _PluginAFactory() + + entry_point = _FakeEntryPoint("a", another_plugin_a_factory) + + with patch( + "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", + return_value=[entry_point], + ): + result = load_configured_plugins( + [_plugin_a_factory], + environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, + ) + + assert result == [_plugin_a_factory, another_plugin_a_factory] @pytest.mark.parametrize("configured_value", ["a,,b", ",a", "a,"]) @@ -177,7 +276,7 @@ def test_discovery_rejects_duplicate_configured_names() -> None: def test_discovery_reports_missing_provider_and_available_names() -> None: - entry_point = _FakeEntryPoint("available", _provider(_PluginA)) + entry_point = _FakeEntryPoint("available", _plugin_a_factory) with ( patch( @@ -216,12 +315,12 @@ def test_discovery_rejects_ambiguous_provider_name() -> None: entry_points = [ _FakeEntryPoint( "duplicate", - _provider(_PluginA), + _plugin_a_factory, distribution_name="package-a", ), _FakeEntryPoint( "duplicate", - _provider(_PluginB), + _plugin_b_factory, distribution_name="package-b", ), ] @@ -256,10 +355,10 @@ def test_discovery_wraps_entry_point_enumeration_failure() -> None: ) -def test_discovery_wraps_provider_load_failure() -> None: +def test_discovery_wraps_factory_load_failure() -> None: entry_point = _FakeEntryPoint( "a", - _provider(_PluginA), + _plugin_a_factory, load_error=ImportError("missing dependency"), ) @@ -275,25 +374,27 @@ def test_discovery_wraps_provider_load_failure() -> None: environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, ) - assert "Failed to load durable instrumentation plugin provider 'a'" in str( + assert "Failed to load durable instrumentation plugin factory 'a'" in str( error.value ) assert "test-plugin-package" in str(error.value) assert isinstance(error.value.__cause__, ImportError) -def test_discovery_rejects_invalid_provider_type() -> None: - entry_point = _FakeEntryPoint("a", _PluginA) +def test_discovery_names_unknown_distribution_in_load_failure() -> None: + entry_point = _FakeEntryPoint( + "a", + _plugin_a_factory, + distribution_name=None, + load_error=ImportError("missing dependency"), + ) with ( patch( "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", return_value=[entry_point], ), - pytest.raises( - PluginLoadError, - match="must resolve to DurableInstrumentationPluginProvider", - ), + pytest.raises(PluginLoadError, match="unknown distribution"), ): load_configured_plugins( None, @@ -301,11 +402,30 @@ def test_discovery_rejects_invalid_provider_type() -> None: ) -def test_discovery_rejects_incompatible_plugin_api_version() -> None: - entry_point = _FakeEntryPoint( - "a", - _provider(_PluginA, plugin_api_version=99), - ) +@pytest.mark.parametrize( + ("resolved_value", "expected_description"), + [ + (_PluginA(), "_PluginA"), + (object(), "builtins.object"), + ("not-a-factory", "builtins.str"), + (None, "builtins.NoneType"), + # A function is named by its own qualified name, not by its type. Its type + # is ``builtins.function`` for every function ever written, so naming the + # type would identify no particular one. + (lambda info: _PluginA(), ""), + ], +) +def test_discovery_rejects_entry_point_without_create_plugin( + resolved_value: object, + expected_description: str, +) -> None: + """A plugin *instance* at the entry point is now the common mistake. + + A bare callable is the other one, and it is rejected too: the registration + type is an object with ``create_plugin``, so a function that builds a plugin + no longer satisfies it. The message names what the target actually was. + """ + entry_point = _FakeEntryPoint("a", resolved_value) with ( patch( @@ -319,70 +439,209 @@ def test_discovery_rejects_incompatible_plugin_api_version() -> None: environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, ) - assert "declares plugin API version 99" in str(error.value) - assert ( - f"supports plugin API version {DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION}" - in str(error.value) - ) + assert "must resolve to a plugin factory" in str(error.value) + assert "create_plugin(info) method" in str(error.value) + assert expected_description in str(error.value) -def test_discovery_rejects_invalid_declared_plugin_type() -> None: - provider = DurableInstrumentationPluginProvider( - plugin_type=cast(type[DurableInstrumentationPlugin], object), - factory=_PluginA, - plugin_api_version=DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION, - ) - entry_point = _FakeEntryPoint("a", provider) +def test_discovery_rejects_a_plugin_class_at_the_entry_point() -> None: + """A plugin class carries no ``create_plugin``, so it is not a factory.""" + + class _InfoAwarePlugin(DurableInstrumentationPlugin): + def __init__(self, info: InvocationStartInfo) -> None: + self.info = info + + entry_point = _FakeEntryPoint("a", _InfoAwarePlugin) with ( patch( "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", return_value=[entry_point], ), - pytest.raises( - PluginLoadError, - match="declares invalid plugin type builtins.object", - ), + pytest.raises(PluginLoadError) as error, ): load_configured_plugins( None, environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, ) + assert "must resolve to a plugin factory" in str(error.value) + assert "not a plugin class" in str(error.value) -def test_discovery_rejects_non_class_declared_plugin_type() -> None: - provider = DurableInstrumentationPluginProvider( - plugin_type=cast(type[DurableInstrumentationPlugin], _PluginA()), - factory=_PluginA, - plugin_api_version=DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION, - ) - entry_point = _FakeEntryPoint("a", provider) - with ( - patch( - "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", - return_value=[entry_point], - ), - pytest.raises( - PluginLoadError, - match="declares invalid plugin type .*_PluginA", - ), - ): +def test_explicit_plugin_instance_is_rejected_with_its_position() -> None: + """A plugin instance in ``plugins`` fails configuration, not every invocation. + + An instance has no ``create_plugin``, so the per-invocation factory call + raises ``AttributeError``, which the executor logs and swallows -- the plugin + silently never runs. The position is asserted because a caller passing several + entries has no other way to tell which one is wrong. + """ + with pytest.raises(PluginLoadError) as error: load_configured_plugins( - None, - environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, + [_plugin_a_factory, _PluginB(), _plugin_b_factory], # type: ignore[list-item] + environment={}, ) + assert "plugins[1]" in str(error.value) + assert "must be a plugin factory" in str(error.value) + assert "_PluginB" in str(error.value) + # An instance and the class it was built from share one qualified name, so the + # message states which of the two was passed. + assert "an instance of" in str(error.value) + + +@pytest.mark.parametrize( + ("invalid_entry", "expected_type_name"), + [ + (_PluginA(), "_PluginA"), + (object(), "builtins.object"), + ("not-a-factory", "builtins.str"), + (None, "builtins.NoneType"), + ], +) +def test_explicit_entries_without_create_plugin_are_rejected( + invalid_entry: object, + expected_type_name: str, +) -> None: + """One callable member is all that is checkable without building an instance.""" + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([invalid_entry], environment={}) # type: ignore[list-item] -def test_discovery_wraps_plugin_factory_failure() -> None: - def fail_factory() -> _PluginA: - raise RuntimeError("factory failed") + assert "plugins[0]" in str(error.value) + assert expected_type_name in str(error.value) - entry_point = _FakeEntryPoint( - "a", - _provider(fail_factory), - distribution_name=None, - ) + +def test_explicit_non_callable_create_plugin_is_rejected() -> None: + """The attribute is tested for callability, not merely for presence. + + An object whose ``create_plugin`` is data would otherwise pass here and raise + ``TypeError`` on every invocation, where it is logged and swallowed. + """ + + class _NotAFactory: + create_plugin = "not callable" + + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([_NotAFactory()], environment={}) # type: ignore[list-item] + + assert "plugins[0]" in str(error.value) + assert "_NotAFactory" in str(error.value) + + +@pytest.mark.parametrize( + "bare_callable", + [ + lambda info: _PluginA(), + _PluginAFactory.create_plugin, + ], +) +def test_explicit_bare_callable_is_rejected(bare_callable: object) -> None: + """A callable is no longer a factory, which reverses the previous rule. + + The registration type was ``Callable[[InvocationStartInfo], + DurableInstrumentationPlugin]``, so a lambda or a plain function was a valid + factory. It is now an object with ``create_plugin``, and no compatibility + path accepts both: a bare callable fails at handler initialization with + guidance naming the replacement. + """ + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([bare_callable], environment={}) # type: ignore[list-item] + + assert "plugins[0]" in str(error.value) + assert "must be a plugin factory" in str(error.value) + assert "plugins=[MyPluginFactory(exporter)]" in str(error.value) + + +def test_explicit_plugin_class_is_rejected_as_a_factory() -> None: + """A plugin class is not a factory, reversing what this branch documented. + + Calling a class constructs an instance, so a class satisfied the previous + ``Callable`` registration type and ``plugins=[MyPlugin]`` was accepted. A + class carries no ``create_plugin`` attribute, so it is now rejected at + handler initialization. The replacement is a factory class whose + ``create_plugin`` constructs the plugin, which also keeps setup work out of + the plugin's ``__init__``. + """ + + class _InfoAwarePlugin(DurableInstrumentationPlugin): + def __init__(self, info: InvocationStartInfo) -> None: + self.info = info + + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([_InfoAwarePlugin], environment={}) # type: ignore[list-item] + + assert "plugins[0]" in str(error.value) + assert "a plugin class" in str(error.value) + + +def test_explicit_class_declaring_create_plugin_is_accepted() -> None: + """The requirement is the member, not the kind of object. + + A class that declares ``create_plugin`` as a ``@classmethod`` carries the + attribute, so the class object itself is a valid factory. Nothing in the + contract requires a factory to be an instance. + """ + + class _ClassFactoryPlugin(DurableInstrumentationPlugin): + def __init__(self, info: InvocationStartInfo) -> None: + self.info = info + + @classmethod + def create_plugin(cls, info: InvocationStartInfo) -> _ClassFactoryPlugin: + return cls(info) + + result = load_configured_plugins([_ClassFactoryPlugin], environment={}) + + assert result == [_ClassFactoryPlugin] + plugin = result[0].create_plugin(INVOCATION_START_INFO) + assert isinstance(plugin, _ClassFactoryPlugin) + assert plugin.info is INVOCATION_START_INFO + + +@pytest.mark.parametrize( + "factory_class", + [ + _PluginAFactory, + _DefaultedArgumentFactory, + _VariadicFactory, + ], + ids=["plain", "defaulted", "variadic"], +) +def test_explicit_factory_class_with_an_instance_method_is_rejected( + factory_class: type, +) -> None: + """The factory class is not the factory, and no signature shape rescues it. + + ``MyFactory.create_plugin`` read off the class is a plain function whose first + parameter is ``self``, so the per-invocation call binds the info to ``self`` + and the factory never sees it. A signature bind alone does not catch every + such shape: ``create_plugin(self, info=None)`` and + ``create_plugin(self, *args)`` both bind one argument to ``self`` and leave + the rest satisfied, so they used to pass and then fail on every invocation + where the error is swallowed. The kind of the member decides instead. + """ + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([factory_class], environment={}) # type: ignore[list-item] + + assert "plugins[0]" in str(error.value) + assert "a factory instance rather than the factory class" in str(error.value) + + +@pytest.mark.parametrize( + "factory_class", + [ + _PluginAFactory, + _DefaultedArgumentFactory, + _VariadicFactory, + ], + ids=["plain", "defaulted", "variadic"], +) +def test_discovery_rejects_a_factory_class_at_the_entry_point( + factory_class: type, +) -> None: + """The entry-point path applies the same rule.""" + entry_point = _FakeEntryPoint("a", factory_class) with ( patch( @@ -396,84 +655,190 @@ def fail_factory() -> _PluginA: environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, ) - assert "Failed to create durable instrumentation plugin 'a'" in str(error.value) - assert "unknown distribution" in str(error.value) - assert isinstance(error.value.__cause__, RuntimeError) + assert "must resolve to a plugin factory" in str(error.value) + assert "not the factory class" in str(error.value) -def test_discovery_rejects_invalid_plugin_type() -> None: - entry_point = _FakeEntryPoint("a", _provider(lambda: object())) +@pytest.mark.parametrize( + "rejected_class", + [_PluginA, _PluginAFactory], + ids=["plugin-class", "factory-class"], +) +def test_a_rejected_class_is_named_by_itself_not_by_its_metaclass( + rejected_class: type, +) -> None: + """The message has to name the class that was passed. + + The type of a class is its metaclass, which is ``builtins.type`` for both + classes here. Both are rejected shapes a caller reaches by accident: + ``plugins=[MyPlugin]`` is what the previous major accepted, and + ``plugins=[MyPluginFactory]`` is this major's shape with the parentheses left + off. Naming the type would report ``builtins.type`` for either one and + distinguish neither. So the class is named directly, and the message says it + was a class rather than an instance. + """ + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([rejected_class], environment={}) # type: ignore[list-item] + + message = str(error.value) + qualified = f"{rejected_class.__module__}.{rejected_class.__qualname__}" + assert f"the class {qualified}" in message + assert "builtins.type" not in message + + +@pytest.mark.parametrize( + "rejected_class", + [_PluginA, _PluginAFactory], + ids=["plugin-class", "factory-class"], +) +def test_a_rejected_class_at_the_entry_point_is_named_by_itself( + rejected_class: type, +) -> None: + """The entry-point path applies the same naming rule.""" + entry_point = _FakeEntryPoint("a", rejected_class) with ( patch( "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", return_value=[entry_point], ), - pytest.raises( - PluginLoadError, - match="expected .*_PluginA", - ), + pytest.raises(PluginLoadError) as error, ): load_configured_plugins( None, environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, ) + message = str(error.value) + qualified = f"{rejected_class.__module__}.{rejected_class.__qualname__}" + assert f"the class {qualified}" in message + assert "builtins.type" not in message -def test_explicit_plugin_registration_takes_precedence( - caplog: pytest.LogCaptureFixture, -) -> None: - explicit_plugin = _PluginA() - factory = Mock(return_value=_PluginA()) - entry_point = _FakeEntryPoint("a", _provider(factory)) - with ( - patch( - "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", - return_value=[entry_point], - ), - caplog.at_level( - logging.WARNING, - logger="aws_durable_execution_sdk_python.plugin_discovery", - ), - ): - result = load_configured_plugins( - [explicit_plugin], - environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, - ) +def test_explicit_class_holding_a_callable_create_plugin_is_accepted() -> None: + """A class attribute holding a callable takes no implicit first argument. - assert result == [explicit_plugin] - factory.assert_not_called() - assert "already registered by the decorator's plugins argument" in caplog.text + Reading it off the class produces the callable itself, so the info reaches it. + """ + class _CallableMember: + def __call__(self, info: InvocationStartInfo) -> _PluginA: + return _PluginA() -def test_first_dynamic_registration_wins_for_duplicate_plugin_type( - caplog: pytest.LogCaptureFixture, -) -> None: - first_factory = Mock(return_value=_PluginA()) - second_factory = Mock(return_value=_PluginA()) - entry_points = [ - _FakeEntryPoint("first", _provider(first_factory)), - _FakeEntryPoint("second", _provider(second_factory)), - ] + class _MemberFactory: + create_plugin = _CallableMember() + + result = load_configured_plugins([_MemberFactory], environment={}) # type: ignore[list-item] + + assert result == [_MemberFactory] + assert isinstance( + _MemberFactory.create_plugin(INVOCATION_START_INFO), + _PluginA, + ) + + +def test_explicit_class_declaring_a_static_create_plugin_is_accepted() -> None: + """A ``@staticmethod`` presents the signature the SDK calls, so it binds.""" + + class _StaticFactoryPlugin(DurableInstrumentationPlugin): + def __init__(self, info: InvocationStartInfo) -> None: + self.info = info + + @staticmethod + def create_plugin(info: InvocationStartInfo) -> _StaticFactoryPlugin: + return _StaticFactoryPlugin(info) + + result = load_configured_plugins([_StaticFactoryPlugin], environment={}) + + assert result == [_StaticFactoryPlugin] + assert isinstance( + result[0].create_plugin(INVOCATION_START_INFO), _StaticFactoryPlugin + ) + + +def test_explicit_factory_without_an_introspectable_signature_is_accepted() -> None: + """A signature that cannot be read is not evidence of a broken factory. + + ``inspect.signature`` raises ``ValueError`` for some C-implemented callables, + ``time.strftime`` among them. Rejecting such a factory would refuse a usable + one over a missing description of it, so the member alone decides. + """ + + class _UnreadableSignatureFactory: + create_plugin = staticmethod(time.strftime) + + factory = _UnreadableSignatureFactory() + + with pytest.raises(ValueError, match="no signature"): + inspect.signature(time.strftime) + + assert load_configured_plugins([factory], environment={}) == [factory] # type: ignore[list-item, comparison-overlap] + + +def test_explicit_factory_taking_no_argument_is_rejected() -> None: + """A ``create_plugin`` that takes nothing cannot receive the info.""" + + class _NoArgumentFactory: + def create_plugin(self) -> _PluginA: + return _PluginA() + + with pytest.raises(PluginLoadError) as error: + load_configured_plugins([_NoArgumentFactory()], environment={}) # type: ignore[list-item] + + assert "plugins[0]" in str(error.value) + assert "create_plugin(info) method" in str(error.value) + + +def test_explicit_factory_object_is_accepted() -> None: + result = load_configured_plugins([_plugin_a_factory], environment={}) + + assert result == [_plugin_a_factory] + assert isinstance(result[0].create_plugin(INVOCATION_START_INFO), _PluginA) + + +def test_explicit_factory_holds_handler_lifetime_state() -> None: + """A factory instance is where state spanning invocations belongs.""" + + class _StatefulFactory: + def __init__(self) -> None: + self.calls: list[InvocationStartInfo] = [] + + def create_plugin(self, info: InvocationStartInfo) -> _PluginA: + self.calls.append(info) + return _PluginA() + + factory = _StatefulFactory() + + result = load_configured_plugins([factory], environment={}) + + assert result == [factory] + assert isinstance(result[0].create_plugin(INVOCATION_START_INFO), _PluginA) + assert factory.calls == [INVOCATION_START_INFO] + + +def test_explicit_entries_are_validated_before_entry_points_are_imported() -> None: + """Explicit entries are checked first, so a valid provider is not imported. + + Importing a provider runs third-party module code. A configuration that is + already invalid should fail before that happens, and the failure should name + the invalid entry rather than whatever the import did. + """ + entry_point = _FakeEntryPoint("a", _plugin_a_factory) with ( patch( "aws_durable_execution_sdk_python.plugin_discovery.metadata.entry_points", - return_value=entry_points, - ), - caplog.at_level( - logging.WARNING, - logger="aws_durable_execution_sdk_python.plugin_discovery", - ), + return_value=[entry_point], + ) as entry_points, + patch.object( + _FakeEntryPoint, "load", side_effect=AssertionError("must not import") + ) as load, + pytest.raises(PluginLoadError, match=r"plugins\[0\]"), ): - result = load_configured_plugins( - None, - environment={PLUGIN_ENVIRONMENT_VARIABLE: "first,second"}, + load_configured_plugins( + [_PluginA()], # type: ignore[list-item] + environment={PLUGIN_ENVIRONMENT_VARIABLE: "a"}, ) - assert len(result) == 1 - assert isinstance(result[0], _PluginA) - first_factory.assert_called_once_with() - second_factory.assert_not_called() - assert "already registered by dynamic provider 'first'" in caplog.text + entry_points.assert_not_called() + load.assert_not_called() diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 69cdfc50..517eded2 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -1,7 +1,11 @@ +import asyncio +import contextlib import datetime import logging import pickle +import threading import unittest +from collections.abc import Iterator from copy import deepcopy from dataclasses import asdict, fields from unittest.mock import MagicMock, patch @@ -20,6 +24,7 @@ ) from aws_durable_execution_sdk_python.plugin import ( DurableInstrumentationPlugin, + DurableInstrumentationPluginFactory, InvocationEndInfo, InvocationInfo, InvocationStatus, @@ -30,10 +35,13 @@ OperationStartInfo, OperationType, PluginExecutor, + PluginHost, UserFunctionEndInfo, UserFunctionOutcome, UserFunctionStartInfo, ) +from aws_durable_execution_sdk_python.plugin_discovery import _is_plugin_factory +from tests.test_helpers import plugin_factory # region Dataclass Tests @@ -44,6 +52,35 @@ LAMBDA_CTX = MagicMock() LAMBDA_CTX.aws_request_id = "req-1" + +@contextlib.contextmanager +def _invocation( + executor: PluginExecutor, + *recorders: "_TrackingPlugin", +) -> Iterator[None]: + """Open the executor's per-invocation scope around a single-hook test. + + Plugin instances are built by ``on_invocation_start`` and dropped when the + scope exits, so a hook dispatched outside an invocation reaches no plugin at + all. A test that exercises one hook therefore has to establish the invocation + that hook belongs to. + + The invocation-start hook this fires is scaffolding for the hook under test, + so it is cleared from the given recorders before the body runs. The + assertions that follow are about the hook under test only. + """ + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=False, + ) + for recorder in recorders: + recorder.calls.clear() + yield + + OPERATION_START_INFO = OperationStartInfo( operation_id="op-2", operation_type=OperationType.CALLBACK, @@ -522,39 +559,686 @@ def test_subclass_override(self): # endregion DurableInstrumentationPlugin Tests +# region DurableInstrumentationPluginFactory Tests +class TestDurableInstrumentationPluginFactory(unittest.TestCase): + def test_protocol_is_not_runtime_checkable(self): + """``isinstance`` against the protocol must stay unavailable. + + A runtime-checkable protocol requires every declared member, so an + optional second member added later would break any caller's isinstance + check -- the additive extensibility this protocol exists for. The SDK + checks the shape with ``plugin_discovery._is_plugin_factory`` instead, + which also tests that ``create_plugin`` is callable rather than merely + present. + """ + with self.assertRaises(TypeError): + isinstance( # type: ignore[misc] # noqa: B018 + plugin_factory(_NoOpPlugin()), DurableInstrumentationPluginFactory + ) + + def test_factory_shape_check_requires_a_callable_create_plugin(self): + """The shape is one callable member, structurally, without importing it.""" + + class _Data: + create_plugin = "not callable" + + self.assertTrue(_is_plugin_factory(plugin_factory(_NoOpPlugin()))) + self.assertFalse(_is_plugin_factory(_Data())) + self.assertFalse(_is_plugin_factory(lambda info: _NoOpPlugin())) + self.assertFalse(_is_plugin_factory(_NoOpPlugin())) + + +# endregion DurableInstrumentationPluginFactory Tests + + # region PluginExecutor Tests class TestPluginExecutorInit(unittest.TestCase): def test_init_with_none(self): executor = PluginExecutor(plugins=None) + self.assertEqual(executor._plugin_factories, []) self.assertEqual(executor._plugins, []) def test_init_with_empty_list(self): executor = PluginExecutor(plugins=[]) + self.assertEqual(executor._plugin_factories, []) self.assertEqual(executor._plugins, []) - def test_init_with_plugins(self): + def test_init_records_factories_without_building_plugins(self): + """Construction stores factories only; instances belong to an invocation.""" p1 = _NoOpPlugin() p2 = _TrackingPlugin() - executor = PluginExecutor(plugins=[p1, p2]) - self.assertEqual(len(executor._plugins), 2) + executor = PluginExecutor(plugins=[plugin_factory(p1), plugin_factory(p2)]) + self.assertEqual(len(executor._plugin_factories), 2) + self.assertEqual(executor._plugins, []) + + +class TestPluginLifetime(unittest.TestCase): + """The per-invocation plugin lifetime.""" + + def test_each_invocation_gets_its_own_instance(self): + """Two invocations of one handler never share a plugin instance.""" + built: list[_TrackingPlugin] = [] + + class _BuildingFactory: + def create_plugin(self, info: InvocationStartInfo) -> _TrackingPlugin: + plugin = _TrackingPlugin() + built.append(plugin) + return plugin + + # One host for the handler, one executor per invocation -- the shape + # durable_execution() uses. + host = PluginHost(plugins=[_BuildingFactory()]) + + for request_id in ("req-1", "req-2"): + lambda_context = MagicMock() + lambda_context.aws_request_id = request_id + with host.invocation() as executor: + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=lambda_context, + execution_start_time=START_TS, + is_first_invocation=False, + ) + + self.assertEqual(len(built), 2) + self.assertIsNot(built[0], built[1]) + # Each instance saw only its own invocation. + self.assertEqual(built[0].calls, ["invocation_start:req-1"]) + self.assertEqual(built[1].calls, ["invocation_start:req-2"]) + + def test_a_hostile_factory_attribute_hook_does_not_escape_containment(self): + """Naming a failing factory must not be what fails the invocation. + + ``_factory_name`` runs while a factory failure is being contained, so it + reads the factory's type rather than the factory: an instance + ``__getattr__`` belongs to customer code and can raise. + """ + surviving = _TrackingPlugin() + executor = PluginExecutor( + plugins=[_HostileAttributeFactory(), plugin_factory(surviving)], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + self.assertIn("_HostileAttributeFactory", "\n".join(logs.output)) + self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + + def test_a_factory_raising_a_group_with_a_control_exception_propagates(self): + """A group carrying a control exception is not contained. + + ``BaseExceptionGroup`` is neither of the two cases an ``isinstance`` chain + covers: a group holding a ``KeyboardInterrupt`` is not an instance of one, + so naming the three in a handler does not match it and a broad handler + would swallow the interrupt inside it. Plugin code produces such a group + without asking for one -- an ``asyncio.TaskGroup`` whose task is + interrupted raises it. + """ + group = BaseExceptionGroup( + "plugin group", [ValueError("contained"), KeyboardInterrupt()] + ) + executor = PluginExecutor(plugins=[_GroupRaisingFactory(group)]) + + with ( + self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs, + executor.run(), + self.assertRaises(BaseExceptionGroup) as raised, + ): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + # Only the control leaf propagates; the rest was logged as a contained + # plugin failure. + self.assertEqual( + [type(leaf) for leaf in raised.exception.exceptions], [KeyboardInterrupt] + ) + self.assertIn("contained", "\n".join(logs.output)) + + def test_a_factory_raising_a_group_without_a_control_exception_is_contained(self): + """A group of ordinary failures is contained like any other.""" + surviving = _TrackingPlugin() + group = BaseExceptionGroup("plugin group", [ValueError("boom")]) + executor = PluginExecutor( + plugins=[_GroupRaisingFactory(group), plugin_factory(surviving)], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + self.assertIn("boom", "\n".join(logs.output)) + self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + + def test_a_failing_end_hook_does_not_replace_the_handlers_failure(self): + """Instrumentation does not decide what an execution failed with. + + The end-hook dispatch can raise: it re-raises the control exceptions it + holds through the fan-out. On the failure exit that raise used to leave + the ``except`` block before the handler's own exception was re-raised, so + the caller saw the plugin's exception and the real failure survived only + as ``__context__``. + """ + plugin = _ControlOnEndPlugin() + host = PluginHost(plugins=[plugin_factory(plugin)]) + + @host.handle_durable_output + def handler(event, context, plugin_executor): + plugin_executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + raise ValueError("the real failure") + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ): + with self.assertRaises(ValueError) as raised: + handler({}, LAMBDA_CTX) + + self.assertEqual(str(raised.exception), "the real failure") + # The hook still ran, and still saw the failing outcome. + self.assertEqual(plugin.end_statuses, [InvocationStatus.RETRY]) + + def test_an_end_hook_that_stops_the_thread_still_finishes_the_dispatch(self): + """Every started plugin receives the end hook, then the thread stops. + + The end hook is a plugin's only chance to finish -- Insight drains there + and OTel ends its spans -- so a plugin raising one of the three control + exceptions must not cost the plugins after it in the list their own end + hook. The exception is held and re-raised once every plugin has been + called. + """ + later = _TrackingPlugin() + host = PluginHost( + plugins=[plugin_factory(_ControlOnEndPlugin()), plugin_factory(later)] + ) + + @host.handle_durable_output + def handler(event, context, plugin_executor): + plugin_executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + return { + "Status": ServiceInvocationStatus.SUCCEEDED.value, + "Result": None, + } + + with self.assertRaises(KeyboardInterrupt): + handler({}, LAMBDA_CTX) + + self.assertEqual( + later.calls, ["invocation_start:req-1", "invocation_end:req-1"] + ) + + def test_a_start_hook_that_stops_the_thread_leaves_later_plugins_unpaired(self): + """A plugin that never received the start hook never receives the end hook. + + A start hook raising one of the three control exceptions propagates out of + the dispatch loop, so plugins after it in the list never receive their + start hook. The invocation-end hook that the propagating exception then + triggers used to reach them anyway, leaving a plugin to tear down state it + had never been told to build. + """ + later = _TrackingPlugin() + host = PluginHost( + plugins=[plugin_factory(_ControlOnStartPlugin()), plugin_factory(later)] + ) + + @host.handle_durable_output + def handler(event, context, plugin_executor): + plugin_executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + raise AssertionError("unreachable: the start hook stops the thread") + + with self.assertRaises(KeyboardInterrupt): + handler({}, LAMBDA_CTX) + + self.assertEqual(later.calls, []) + + def test_an_end_hook_that_stops_the_thread_is_not_reported_twice(self): + """Exactly one end notification per invocation, whatever the hook does. + + ``_dispatch_plugin`` re-raises the three exceptions that instruct the + calling thread to stop, so a hook can raise. With the success dispatch + inside the try, that raise was caught as a handler failure and the hook + ran a second time with a RETRY outcome -- the wrong outcome for an + invocation that succeeded, and a second export for an exporter. + """ + plugin = _ControlOnEndPlugin() + host = PluginHost(plugins=[plugin_factory(plugin)]) + + @host.handle_durable_output + def handler(event, context, plugin_executor): + plugin_executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + return { + "Status": ServiceInvocationStatus.SUCCEEDED.value, + "Result": None, + } + + with self.assertRaises(KeyboardInterrupt): + handler({}, LAMBDA_CTX) + + self.assertEqual(plugin.end_statuses, [InvocationStatus.SUCCEEDED]) + + def test_a_body_raising_cancellation_still_fires_the_end_hook(self): + """Every exit fires the end hook, not only the ones deriving from Exception. + + A handler that surfaces an ``asyncio.CancelledError`` left the invocation + without its end hook, so Insight never drained the records it held for the + execution and OTel never ended the spans it had opened. Nothing failed + visibly, which is why the gap was silent. + """ + for raised in ( + asyncio.CancelledError("cancelled"), + KeyboardInterrupt(), + SystemExit(), + ): + with self.subTest(raised=type(raised).__name__): + plugin = _TrackingPlugin() + host = PluginHost(plugins=[plugin_factory(plugin)]) + + @host.handle_durable_output + def handler(event, context, plugin_executor): + plugin_executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + raise raised + + with self.assertRaises(type(raised)): + handler({}, LAMBDA_CTX) + + self.assertEqual( + plugin.calls, + ["invocation_start:req-1", "invocation_end:req-1"], + ) + + def test_host_hands_out_a_new_executor_per_invocation(self): + """The host itself holds no per-invocation state to overwrite.""" + host = PluginHost(plugins=[plugin_factory(_TrackingPlugin())]) + + with host.invocation() as first, host.invocation() as second: + self.assertIsNot(first, second) + + def test_executor_refuses_a_second_invocation(self): + """An executor that has served an invocation cannot serve another. + + The lifetime is enforced by the class, not just by how + ``durable_execution()`` happens to call it, so reintroducing a shared + executor fails loudly instead of silently mixing two invocations. + """ + executor = PluginExecutor(plugins=[plugin_factory(_TrackingPlugin())]) + + with executor.run(): + pass + + with self.assertRaisesRegex(RuntimeError, "single-use"): + with executor.run(): + pass + + def test_instances_are_dropped_when_the_invocation_returns(self): + """Nothing on the handler-lifetime executor still references the instance.""" + plugin = _TrackingPlugin() + executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=False, + ) + self.assertEqual(executor._plugins, [plugin]) + + self.assertEqual(executor._plugins, []) + + def test_factory_receives_the_info_the_first_hook_receives(self): + """The factory argument is the identical object, not a copy.""" + factory_infos: list[InvocationStartInfo] = [] + hook_infos: list[InvocationStartInfo] = [] + + class _RecordingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info: InvocationStartInfo) -> None: + hook_infos.append(info) + + class _RecordingFactory: + def create_plugin(self, info: InvocationStartInfo) -> _RecordingPlugin: + factory_infos.append(info) + return _RecordingPlugin() + + executor = PluginExecutor(plugins=[_RecordingFactory()]) + + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + execution_input={"name": "World"}, + ) + + self.assertEqual(len(factory_infos), 1) + self.assertEqual(len(hook_infos), 1) + self.assertIs(factory_infos[0], hook_infos[0]) + self.assertEqual(factory_infos[0].execution_arn, "arn:exec") + self.assertEqual(factory_infos[0].execution_input, {"name": "World"}) + + def test_factories_run_before_the_first_hook_is_dispatched(self): + """Every instance exists before any of them receives a hook.""" + events: list[str] = [] + + class _OrderedPlugin(DurableInstrumentationPlugin): + def __init__(self, label: str) -> None: + self.label = label + + def on_invocation_start(self, info: InvocationStartInfo) -> None: + events.append(f"hook:{self.label}") + + class _OrderedFactory: + def __init__(self, label: str) -> None: + self._label = label + + def create_plugin(self, info: InvocationStartInfo) -> _OrderedPlugin: + events.append(f"build:{self._label}") + return _OrderedPlugin(self._label) + + executor = PluginExecutor(plugins=[_OrderedFactory("a"), _OrderedFactory("b")]) + + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + self.assertEqual(events, ["build:a", "build:b", "hook:a", "hook:b"]) + + def test_failing_factory_is_contained(self): + """A factory whose create_plugin raises is logged and skipped.""" + surviving = _TrackingPlugin() + + executor = PluginExecutor( + plugins=[_ExplodingFactory(), plugin_factory(surviving)], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + self.assertIn("factory boom", "\n".join(logs.output)) + # The other factory's plugin still receives its hooks. + self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + + def test_bare_callable_builds_no_plugin(self): + """The executor calls ``create_plugin``, so a bare callable yields nothing. + + A callable used to be a factory. It is not one now, and an entry that + reaches the executor anyway raises ``AttributeError`` there, which the + executor logs and skips. ``load_configured_plugins`` rejects such an entry + at handler initialization, so this path is only reachable by constructing + an executor directly. + """ + surviving = _TrackingPlugin() + + executor = PluginExecutor( + plugins=[ + (lambda info: _TrackingPlugin()), # type: ignore[list-item] + plugin_factory(surviving), + ], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + self.assertEqual(executor._plugins, [surviving]) + + self.assertIn("create_plugin", "\n".join(logs.output)) + + def test_factory_returning_none_is_contained(self): + """A factory that returns nothing is logged and skipped.""" + surviving = _TrackingPlugin() + + class _NoneFactory: + def create_plugin(self, info: InvocationStartInfo): + return None + + executor = PluginExecutor( + plugins=[_NoneFactory(), plugin_factory(surviving)], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + self.assertEqual(executor._plugins, [surviving]) + + self.assertIn("returned None", "\n".join(logs.output)) + self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + + def test_factory_returning_a_non_plugin_is_contained(self): + """A factory whose return is not a plugin is logged once and skipped. + + The load-time shape check establishes only that the factory has a + callable ``create_plugin``; what that call returns is knowable only + here. A value that is not a plugin fails every hook, so registering it + would log one error per hook per invocation and still provide no + telemetry. + """ + surviving = _TrackingPlugin() + + class _WrongTypeFactory: + def create_plugin(self, info: InvocationStartInfo): + return object() + + executor = PluginExecutor( + plugins=[_WrongTypeFactory(), plugin_factory(surviving)], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + self.assertEqual(executor._plugins, [surviving]) + + output = "\n".join(logs.output) + self.assertIn("not a DurableInstrumentationPlugin", output) + self.assertIn("object", output) + # One error, not one per hook: the value never reaches the dispatch. + self.assertEqual(len(logs.output), 1) + self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + + def test_every_failing_factory_leaves_the_executor_usable(self): + """All factories failing is not distinguishable from having no plugins.""" + executor = PluginExecutor(plugins=[_ExplodingFactory()]) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ): + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + self.assertEqual(executor._plugins, []) + # Later hooks are no-ops rather than errors. + executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=ServiceInvocationStatus.SUCCEEDED, + result=None, + error=None, + ), + ) + + def test_a_factory_raising_cancellation_is_contained(self): + """Containment is not limited to ``Exception``. + + ``asyncio.CancelledError`` derives from ``BaseException``, so a factory + that awaits a cancelled task used to abort the invocation it was only + instrumenting and stop the remaining factories from running. + """ + surviving = _TrackingPlugin() + + executor = PluginExecutor( + plugins=[_CancellingFactory(), plugin_factory(surviving)], + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + self.assertIn("factory cancelled", "\n".join(logs.output)) + self.assertEqual(surviving.calls, ["invocation_start:req-1"]) + + def test_a_factory_raising_thread_control_still_propagates(self): + """The three that tell the thread to stop are not contained. + + Containing one would drop the instruction and hand the thread back to the + work that follows the plugin. + """ + for control in (KeyboardInterrupt, SystemExit, GeneratorExit): + with self.subTest(control=control.__name__): + executor = PluginExecutor( + plugins=[_ThreadControlFactory(control)], + ) + + with executor.run(), self.assertRaises(control): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + def test_a_hook_raising_cancellation_is_contained(self): + """The hook boundary uses the same rule as the factory boundary.""" + tracking = _TrackingPlugin() + executor = PluginExecutor( + plugins=[plugin_factory(_CancellingPlugin()), plugin_factory(tracking)] + ) + + with self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level=logging.ERROR + ) as logs: + with _invocation(executor, tracking): + executor.execute_plugins(OPERATION_START_INFO) + + self.assertIn("hook cancelled", "\n".join(logs.output)) + self.assertIn("operation_start:op-2", tracking.calls) class TestPluginExecutor(unittest.TestCase): - def test_no_thread_pool_when_plugins_is_none(self): - """Tests that PluginExecutor does not create a thread pool when plugins is empty.""" - executor = PluginExecutor(plugins=None) - self.assertIsNone(executor._executor) + def test_dispatch_is_a_no_op_when_no_factory_is_registered(self): + """An executor with no factories dispatches nothing and needs no thread. + + Hooks are dispatched on the calling thread, which is what lets a plugin + set thread-affine state the SDK's own logging then reads, and what makes + the end-hook pairing and re-raise rules enforceable. There is therefore no + pool to create, and an executor with nothing registered returns before it + touches anything. + """ + for plugins in (None, []): + with self.subTest(plugins=plugins): + executor = PluginExecutor(plugins=plugins) + self.assertEqual(executor._plugin_factories, []) + with executor.run(): + # Nothing registered, so no hook reaches a plugin and no + # dispatch raises on the way through. + executor.execute_plugins(INVOCATION_START_INFO) + executor.execute_plugins(OPERATION_START_INFO) - def test_no_thread_pool_when_plugins_is_empty_list(self): - executor = PluginExecutor(plugins=[]) - self.assertIsNone(executor._executor) + def test_hooks_are_dispatched_on_the_calling_thread(self): + """Thread affinity is the contract, so it is asserted rather than assumed.""" + seen: list[str] = [] + + class _ThreadRecordingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info: InvocationStartInfo) -> None: + seen.append(threading.current_thread().name) - def test_thread_pool_created_when_plugins_provided(self): - executor = PluginExecutor(plugins=[_NoOpPlugin()]) + executor = PluginExecutor(plugins=[plugin_factory(_ThreadRecordingPlugin())]) with executor.run(): - self.assertIsNotNone(executor._executor) + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + ) + + self.assertEqual(seen, [threading.current_thread().name]) def test_start_is_noop_when_empty(self): executor = PluginExecutor(plugins=[]) @@ -629,41 +1313,41 @@ class TestPluginExecutorExecutePlugins(unittest.TestCase): def setUp(self): self.plugin = _TrackingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) def test_dispatch_invocation_start_info(self): - with self.executor.run(): - self.executor.execute_plugins(INVOCATION_START_INFO, sync=True) + with _invocation(self.executor, self.plugin): + self.executor.execute_plugins(INVOCATION_START_INFO) self.assertIn("invocation_start:req-1", self.plugin.calls) def test_dispatch_invocation_end_info(self): - with self.executor.run(): - self.executor.execute_plugins(INVOCATION_END_INFO, sync=True) + with _invocation(self.executor, self.plugin): + self.executor.execute_plugins(INVOCATION_END_INFO) self.assertIn("invocation_end:req-1", self.plugin.calls) def test_dispatch_operation_end_info(self): - with self.executor.run(): - self.executor.execute_plugins(OPERATION_END_INFO, sync=False) + with _invocation(self.executor, self.plugin): + self.executor.execute_plugins(OPERATION_END_INFO) self.assertIn("operation_end:op-1", self.plugin.calls) def test_dispatch_operation_start_info(self): - with self.executor.run(): - self.executor.execute_plugins(OPERATION_START_INFO, sync=False) + with _invocation(self.executor, self.plugin): + self.executor.execute_plugins(OPERATION_START_INFO) self.assertIn("operation_start:op-2", self.plugin.calls) def test_dispatch_operation_change_info(self): - with self.executor.run(): - self.executor.execute_plugins(OPERATION_CHANGE_INFO, sync=False) + with _invocation(self.executor, self.plugin): + self.executor.execute_plugins(OPERATION_CHANGE_INFO) self.assertIn("operation_change:op-1", self.plugin.calls) def test_dispatch_user_function_start_info(self): - with self.executor.run(): - self.executor.execute_plugins(USER_FUNCTION_START_INFO, sync=True) + with _invocation(self.executor, self.plugin): + self.executor.execute_plugins(USER_FUNCTION_START_INFO) self.assertIn("user_function_start:op-1", self.plugin.calls) def test_dispatch_user_function_end_info(self): - with self.executor.run(): - self.executor.execute_plugins(USER_FUNCTION_END_INFO, sync=True) + with _invocation(self.executor, self.plugin): + self.executor.execute_plugins(USER_FUNCTION_END_INFO) self.assertIn("user_function_end:op-1", self.plugin.calls) def test_dispatch_unknown_type_logs_exception(self): @@ -671,20 +1355,22 @@ def test_dispatch_unknown_type_logs_exception(self): with self.assertLogs( "aws_durable_execution_sdk_python.plugin", level=logging.ERROR ): - with self.executor.run(): - self.executor.execute_plugins("not a valid info type", sync=True) + with _invocation(self.executor, self.plugin): + self.executor.execute_plugins("not a valid info type") def test_plugin_exception_is_swallowed(self): """If a plugin raises, the exception is logged and execution continues.""" failing_plugin = _FailingPlugin() tracking_plugin = _TrackingPlugin() - executor = PluginExecutor(plugins=[failing_plugin, tracking_plugin]) + executor = PluginExecutor( + plugins=[plugin_factory(failing_plugin), plugin_factory(tracking_plugin)] + ) with self.assertLogs( "aws_durable_execution_sdk_python.plugin", level=logging.ERROR ): - with executor.run(): - executor.execute_plugins(OPERATION_START_INFO, sync=True) + with _invocation(executor, tracking_plugin): + executor.execute_plugins(OPERATION_START_INFO) # The second plugin should still have been called self.assertIn("operation_start:op-2", tracking_plugin.calls) @@ -692,10 +1378,10 @@ def test_plugin_exception_is_swallowed(self): def test_multiple_plugins_all_called(self): p1 = _TrackingPlugin() p2 = _TrackingPlugin() - executor = PluginExecutor(plugins=[p1, p2]) + executor = PluginExecutor(plugins=[plugin_factory(p1), plugin_factory(p2)]) - with executor.run(): - executor.execute_plugins(OPERATION_START_INFO, sync=True) + with _invocation(executor, p1, p2): + executor.execute_plugins(OPERATION_START_INFO) self.assertIn("operation_start:op-2", p1.calls) self.assertIn("operation_start:op-2", p2.calls) @@ -706,7 +1392,7 @@ class TestPluginExecutorOnInvocationStart(unittest.TestCase): def setUp(self): self.plugin = _TrackingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) self.ts = datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC) def _make_operation(self, start_time=None): @@ -780,7 +1466,7 @@ class TestPluginExecutorOnInvocationEnd(unittest.TestCase): def setUp(self): self.plugin = _TrackingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) self.ts = datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC) def _make_operation(self, start_ts=None, end_ts=None): @@ -857,7 +1543,7 @@ def on_invocation_start(_self, info): # noqa: N805 def on_invocation_end(_self, info): # noqa: N805 self.captured.append(info) - self.executor = PluginExecutor(plugins=[_CapturingPlugin()]) + self.executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) @staticmethod def _operation(operation_id, status=OperationStatus.SUCCEEDED): @@ -1030,7 +1716,7 @@ class _CapturingPlugin(DurableInstrumentationPlugin): def on_invocation_start(_self, info): # noqa: N805 seen.append(info) - executor = PluginExecutor(plugins=[_CapturingPlugin()]) + executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) with executor.run(): executor.on_invocation_start( execution_arn="arn:exec", @@ -1209,7 +1895,7 @@ def on_invocation_start(_self, info): # noqa: N805 def on_invocation_end(_self, info): # noqa: N805 self.captured.append(info) - self.executor = PluginExecutor(plugins=[_CapturingPlugin()]) + self.executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) def _fire_hooks(self): with self.executor.run(): @@ -1281,7 +1967,7 @@ class TestPluginExecutorOnOperationAction(unittest.TestCase): def setUp(self): self.plugin = _TrackingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) def test_start_action_fires_operation_start(self): captured: list[OperationStartInfo] = [] @@ -1292,7 +1978,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: captured.append(info) self.plugin = _CapturingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) update = MagicMock() update.action = OperationAction.START update.operation_id = "op-1" @@ -1301,7 +1987,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: update.name = "my-step" update.parent_id = "parent-1" - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_action(update) self.assertIn("operation_start:op-1", self.plugin.calls) @@ -1318,7 +2004,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: captured.append(info) self.plugin = _CapturingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) update = MagicMock() update.action = OperationAction.START update.operation_id = "op-1" @@ -1334,7 +2020,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: start_timestamp=START_TS, ) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_action(update, operation) self.assertEqual(captured[0].start_time, START_TS) @@ -1348,7 +2034,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: captured.append(info) self.plugin = _CapturingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) update = MagicMock() update.action = OperationAction.START update.operation_id = "op-1" @@ -1368,7 +2054,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: status=OperationStatus.READY, ) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_action( update, operation=current_operation, @@ -1411,28 +2097,28 @@ def test_terminal_operation_does_not_fire_callbacks(self): for status in terminal_statuses: with self.subTest(status=status): plugin = _TrackingPlugin() - executor = PluginExecutor(plugins=[plugin]) + executor = PluginExecutor(plugins=[plugin_factory(plugin)]) operation = Operation( operation_id="op-1", operation_type=ServiceOperationType.STEP, status=status, ) - with executor.run(): + with _invocation(executor, plugin): executor.on_operation_replay(operation) self.assertEqual(plugin.calls, []) def test_non_terminal_operation_fires_operation_start(self): plugin = _TrackingPlugin() - executor = PluginExecutor(plugins=[plugin]) + executor = PluginExecutor(plugins=[plugin_factory(plugin)]) operation = Operation( operation_id="op-1", operation_type=ServiceOperationType.WAIT, status=OperationStatus.STARTED, ) - with executor.run(): + with _invocation(executor, plugin): executor.on_operation_replay(operation) self.assertEqual(plugin.calls, ["operation_start:op-1"]) @@ -1450,7 +2136,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: captured.append(info) plugin = _CapturingPlugin() - executor = PluginExecutor(plugins=[plugin]) + executor = PluginExecutor(plugins=[plugin_factory(plugin)]) identifier = OperationIdentifier( operation_id="context-1", sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, @@ -1459,7 +2145,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: ) before = datetime.datetime.now(datetime.UTC) - with executor.run(): + with _invocation(executor, plugin): executor.on_child_context_end( identifier, OperationStatus.FAILED, @@ -1487,13 +2173,13 @@ def on_operation_end(self, info: OperationEndInfo) -> None: class TestPluginExecutorOnUserFunction(unittest.TestCase): def test_user_function_info_uses_plugin_operation_type(self): - executor = PluginExecutor(plugins=[_TrackingPlugin()]) + executor = PluginExecutor(plugins=[plugin_factory(_TrackingPlugin())]) identifier = OperationIdentifier( operation_id="step-1", sub_type=OperationSubType.STEP, ) - with executor.run(): + with _invocation(executor): info = executor.on_user_function_start(identifier) self.assertIs(info.operation_type, OperationType.STEP) @@ -1504,7 +2190,7 @@ class TestPluginExecutorOnOperationUpdate(unittest.TestCase): def setUp(self): self.plugin = _TrackingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) def _make_operation( self, @@ -1532,7 +2218,7 @@ def _make_operation( def test_terminal_status_without_step_details_fires_operation_only(self): op = self._make_operation(status=OperationStatus.FAILED, step_details=None) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_update(op) self.assertIn("operation_end:op-1", self.plugin.calls) @@ -1540,7 +2226,7 @@ def test_terminal_status_without_step_details_fires_operation_only(self): def test_non_terminal_status_without_step_details_fires_nothing(self): op = self._make_operation(status=OperationStatus.STARTED, step_details=None) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_update(op) self.assertEqual(self.plugin.calls, []) @@ -1548,7 +2234,7 @@ def test_non_terminal_status_without_step_details_fires_nothing(self): def test_ready_status_fires_nothing(self): op = self._make_operation(status=OperationStatus.READY, step_details=None) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_update(op) self.assertEqual(self.plugin.calls, []) @@ -1556,7 +2242,7 @@ def test_ready_status_fires_nothing(self): def test_timed_out_is_terminal(self): op = self._make_operation(status=OperationStatus.TIMED_OUT, step_details=None) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_update(op) self.assertIn("operation_end:op-1", self.plugin.calls) @@ -1564,7 +2250,7 @@ def test_timed_out_is_terminal(self): def test_cancelled_is_terminal(self): op = self._make_operation(status=OperationStatus.CANCELLED, step_details=None) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_update(op) self.assertIn("operation_end:op-1", self.plugin.calls) @@ -1572,7 +2258,7 @@ def test_cancelled_is_terminal(self): def test_stopped_is_terminal(self): op = self._make_operation(status=OperationStatus.STOPPED, step_details=None) - with self.executor.run(): + with _invocation(self.executor, self.plugin): self.executor.on_operation_update(op) self.assertIn("operation_end:op-1", self.plugin.calls) @@ -1583,7 +2269,7 @@ class TestPluginExecutorOnOperationChange(unittest.TestCase): def setUp(self): self.plugin = _TrackingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) def test_operation_change_uses_invocation_and_operation_maps(self): updated_operation = Operation( @@ -1612,7 +2298,7 @@ def on_operation_change(self, info: OperationChangeInfo) -> None: captured.append(info) self.plugin = _CapturingPlugin() - self.executor = PluginExecutor(plugins=[self.plugin]) + self.executor = PluginExecutor(plugins=[plugin_factory(self.plugin)]) with self.executor.run(): self.executor.on_invocation_start( @@ -1644,11 +2330,19 @@ def on_operation_change(self, info: OperationChangeInfo) -> None: self.assertEqual(updated_info.end_time, END_TS) self.assertFalse(updated_info.is_replayed) - def test_operation_change_without_invocation_start_is_noop(self): + def test_hooks_outside_an_invocation_reach_no_plugin(self): + """Nothing is dispatched before an invocation establishes the instances. + + This used to assert the ``_invocation_status is None`` guard inside + ``on_operation_update``. That guard is now redundant with the plugin + lifetime itself: instances are built by ``on_invocation_start`` and + dropped when the invocation scope exits, so outside an invocation there + is no instance to dispatch to in the first place. + """ operation = Operation( operation_id="op-1", operation_type=ServiceOperationType.STEP, - status=OperationStatus.STARTED, + status=OperationStatus.SUCCEEDED, ) with self.executor.run(): @@ -1658,6 +2352,7 @@ def test_operation_change_without_invocation_start_is_noop(self): previous_operations={}, ) + self.assertEqual(self.executor._plugins, []) self.assertEqual(self.plugin.calls, []) @@ -1802,6 +2497,81 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: self.calls.append(f"user_function_end:{info.operation_id}") +class _ExplodingFactory: + """Factory whose ``create_plugin`` raises, for containment tests.""" + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + raise RuntimeError("factory boom") + + +class _CancellingFactory: + """Factory whose ``create_plugin`` raises outside the ``Exception`` hierarchy.""" + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + raise asyncio.CancelledError("factory cancelled") + + +class _ThreadControlFactory: + """Factory that raises one of the three exceptions that must propagate.""" + + def __init__(self, control: type[BaseException]) -> None: + self._control = control + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + raise self._control + + +class _GroupRaisingFactory: + """Factory that raises a ``BaseExceptionGroup``, as an ``asyncio.TaskGroup`` does.""" + + def __init__(self, group: BaseExceptionGroup) -> None: + self._group = group + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + raise self._group + + +class _HostileAttributeFactory: + """Factory that fails, and whose attribute hook fails too. + + The second failure is the point: naming a failing factory for the log must not + reach customer code that raises, or containment raises instead of containing. + """ + + def __getattr__(self, name: str) -> object: + msg = f"attribute hook refuses {name}" + raise RuntimeError(msg) + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + msg = "factory boom" + raise RuntimeError(msg) + + +class _CancellingPlugin(DurableInstrumentationPlugin): + """Plugin whose hook raises outside the ``Exception`` hierarchy.""" + + def on_operation_start(self, info): + raise asyncio.CancelledError("hook cancelled") + + +class _ControlOnEndPlugin(DurableInstrumentationPlugin): + """Plugin whose end hook records the outcome and then stops the thread.""" + + def __init__(self) -> None: + self.end_statuses: list[InvocationStatus] = [] + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + self.end_statuses.append(info.status) + raise KeyboardInterrupt + + +class _ControlOnStartPlugin(DurableInstrumentationPlugin): + """Plugin whose start hook stops the thread, so later plugins never start.""" + + def on_invocation_start(self, info: InvocationStartInfo) -> None: + raise KeyboardInterrupt + + class _FailingPlugin(DurableInstrumentationPlugin): """Plugin that raises on every hook call.""" diff --git a/packages/aws-durable-execution-sdk-python/tests/state_test.py b/packages/aws-durable-execution-sdk-python/tests/state_test.py index 713f63fb..b978689a 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -63,6 +63,7 @@ QueuedOperation, ) from aws_durable_execution_sdk_python.threading import CompletionEvent +from tests.test_helpers import plugin_factory, plugin_invocation def test_checkpointed_result_create_from_operation_step(): @@ -4584,7 +4585,7 @@ def test_execution_state_accepts_plugin_executor_parameter(): """Test that ExecutionState can be created with a plugin_executor parameter.""" mock_client = Mock(spec=LambdaClient) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) state = ExecutionState( durable_execution_arn="test_arn", @@ -4617,8 +4618,8 @@ def test_plugin_executor_on_operation_action_called_on_checkpoint(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -4651,8 +4652,8 @@ def test_plugin_executor_on_operation_action_called_on_checkpoint(): def test_async_operation_start_precedes_user_function_start(): """Async START notifies plugins before the user function begins.""" plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -4672,7 +4673,12 @@ def test_async_operation_start_precedes_user_function_start(): ) plugin_executor.on_user_function_start(operation_identifier, attempt=1) - assert plugin.calls[:2] == [ + # Drop the per-invocation instance's own invocation-start hook, which now + # always precedes the operation hooks under test. + operation_calls = [ + call for call in plugin.calls if not call.startswith("invocation_") + ] + assert operation_calls[:2] == [ f"operation_start:{operation_id}", f"user_function_start:{operation_id}", ] @@ -4702,8 +4708,8 @@ def test_existing_operation_start_is_reported_as_replayed(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -4750,8 +4756,8 @@ def test_plugin_executor_on_operation_update_called_for_terminal_operations(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -4800,8 +4806,8 @@ def test_plugin_executor_not_called_for_non_terminal_operations(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -4858,8 +4864,8 @@ def test_plugin_executor_called_for_multiple_updates_in_batch(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): config = CheckpointBatcherConfig( max_batch_time_seconds=0.2, max_batch_operations=10, @@ -4930,7 +4936,7 @@ def test_plugin_executor_on_operation_change_called_for_status_changes(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) with plugin_executor.run(): plugin_executor.on_invocation_start( execution_arn="test_arn", @@ -4979,8 +4985,8 @@ def test_operation_start_plugin_hook_fires_before_checkpoint_failure(): mock_client.checkpoint.side_effect = RuntimeError("API error") plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5037,8 +5043,8 @@ def on_operation_end(self, info): raise RuntimeError("plugin exploded") exploding_plugin = _ExplodingPlugin() - plugin_executor = PluginExecutor(plugins=[exploding_plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(exploding_plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5080,8 +5086,8 @@ class _CapturingPlugin(DurableInstrumentationPlugin): def on_user_function_end(self, info: UserFunctionEndInfo) -> None: captured.append(info) - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5135,8 +5141,8 @@ def test_plugin_executor_not_called_for_pending_operations(): ) plugin = _RecordingPlugin() - plugin_executor = PluginExecutor(plugins=[plugin]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(plugin)]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5202,8 +5208,8 @@ def on_operation_start(self, info): def on_operation_end(self, info): captured.append(("end", info.operation_id, info.is_replayed, info.status)) - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5237,8 +5243,8 @@ def on_operation_start(self, info): def on_operation_end(self, info): captured.append(("end", info.operation_id, info.is_replayed, info.status)) - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5260,8 +5266,8 @@ class _CapturingPlugin(DurableInstrumentationPlugin): def on_operation_start(self, info): captured.append(info.operation_id) - plugin_executor = PluginExecutor(plugins=[_CapturingPlugin()]) - with plugin_executor.run(): + plugin_executor = PluginExecutor(plugins=[plugin_factory(_CapturingPlugin())]) + with plugin_invocation(plugin_executor): state = ExecutionState( durable_execution_arn="test_arn", initial_checkpoint_token="token123", # noqa: S106 @@ -5867,7 +5873,7 @@ def _wrapping_state(plugin: _RecordingPlugin) -> ExecutionState: initial_checkpoint_token="token123", # noqa: S106 operations={}, service_client=Mock(spec=LambdaClient), - plugin_executor=PluginExecutor(plugins=[plugin]), + plugin_executor=PluginExecutor(plugins=[plugin_factory(plugin)]), ) @@ -5901,10 +5907,13 @@ def test_wrap_user_function_reports_incomplete_when_no_outcome(raised): def user_function(): raise raised - with state._plugin_executor.run(), pytest.raises(type(raised)): + with plugin_invocation(state._plugin_executor), pytest.raises(type(raised)): _wrapped(state, user_function)() assert plugin.calls == [ + # The plugin instance is built for this invocation, so its + # invocation-start hook always precedes the operation hooks. + "invocation_start", "user_function_start:step-1", "user_function_end:step-1", ] @@ -5919,10 +5928,13 @@ def test_wrap_user_function_does_not_report_incomplete_on_success(): plugin = _RecordingPlugin() state = _wrapping_state(plugin) - with state._plugin_executor.run(): + with plugin_invocation(state._plugin_executor): assert _wrapped(state, lambda: "done")() == "done" assert plugin.calls == [ + # The plugin instance is built for this invocation, so its + # invocation-start hook always precedes the operation hooks. + "invocation_start", "user_function_start:step-1", "user_function_end:step-1", ] @@ -5939,10 +5951,16 @@ def test_wrap_user_function_does_not_report_incomplete_on_failure(): def user_function(): raise ValueError("boom") - with state._plugin_executor.run(), pytest.raises(ValueError, match="boom"): + with ( + plugin_invocation(state._plugin_executor), + pytest.raises(ValueError, match="boom"), + ): _wrapped(state, user_function)() assert plugin.calls == [ + # The plugin instance is built for this invocation, so its + # invocation-start hook always precedes the operation hooks. + "invocation_start", "user_function_start:step-1", "user_function_end:step-1", ] @@ -5966,7 +5984,7 @@ def on_user_function_end(self, info) -> None: initial_checkpoint_token="token123", # noqa: S106 operations={}, service_client=Mock(spec=LambdaClient), - plugin_executor=PluginExecutor(plugins=[plugin]), + plugin_executor=PluginExecutor(plugins=[plugin_factory(plugin)]), ) def user_function(): @@ -5979,7 +5997,10 @@ def run_on_worker() -> None: with contextlib.suppress(SuspendExecution): _wrapped(state, user_function)() - with state._plugin_executor.run(), ThreadPoolExecutor(max_workers=1) as worker: + with ( + plugin_invocation(state._plugin_executor), + ThreadPoolExecutor(max_workers=1) as worker, + ): worker.submit(run_on_worker).result() assert hook_threads == worker_threads diff --git a/packages/aws-durable-execution-sdk-python/tests/test_helpers.py b/packages/aws-durable-execution-sdk-python/tests/test_helpers.py index 77611a34..f77a11c2 100644 --- a/packages/aws-durable-execution-sdk-python/tests/test_helpers.py +++ b/packages/aws-durable-execution-sdk-python/tests/test_helpers.py @@ -1,9 +1,17 @@ """Test helpers for generating expected step IDs.""" +import contextlib +from collections.abc import Iterator from unittest.mock import Mock from aws_durable_execution_sdk_python.context import DurableContext, ExecutionContext from aws_durable_execution_sdk_python.execution import ExecutionState +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + DurableInstrumentationPluginFactory, + InvocationStartInfo, + PluginExecutor, +) def operation_id_sequence(parent_id: str | None = None): @@ -18,3 +26,54 @@ def operation_id_sequence(parent_id: str | None = None): while True: yield context._create_step_id() # noqa: SLF001 + + +class _FixedPluginFactory: + """Returns one plugin instance the test already holds, for every invocation. + + Production factories build a fresh instance per invocation. A test that has + to read what the plugin recorded needs the instance it passed in, so it + supplies a factory that returns that one. Only valid for a single + invocation, which is all these tests run. + """ + + def __init__(self, plugin: DurableInstrumentationPlugin) -> None: + self._plugin = plugin + + def create_plugin(self, info: InvocationStartInfo) -> DurableInstrumentationPlugin: + return self._plugin + + +def plugin_factory( + plugin: DurableInstrumentationPlugin, +) -> DurableInstrumentationPluginFactory: + """Wrap a plugin instance a test already holds a reference to as a factory.""" + return _FixedPluginFactory(plugin) + + +@contextlib.contextmanager +def plugin_invocation( + plugin_executor: PluginExecutor, + *, + execution_arn: str = "test_arn", + is_first_invocation: bool = True, +) -> Iterator[None]: + """Open a plugin executor's per-invocation scope. + + Plugin instances are built by ``on_invocation_start`` and dropped when the + ``run()`` scope exits, so a hook dispatched outside an invocation reaches no + plugin at all. Tests that exercise a single hook still have to establish the + invocation the hook belongs to; this does that and nothing else. + + The invocation-start hook this fires is scaffolding. Tests asserting on an + exact list of recorded calls should record only the hooks they care about, or + clear their recorder after entering the scope. + """ + with plugin_executor.run(): + plugin_executor.on_invocation_start( + execution_arn=execution_arn, + is_first_invocation=is_first_invocation, + execution_start_time=None, + lambda_context=None, + ) + yield diff --git a/pyproject.toml b/pyproject.toml index 428d680e..ab0f0b85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,7 +129,7 @@ test = "pytest packages/aws-durable-execution-sdk-python-examples/test {args}" [tool.hatch.envs.test-pypi-otel] dependencies = [ - "aws-durable-execution-sdk-python>=2.0.0", + "aws-durable-execution-sdk-python>=3.0.0,<4", "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", "pytest",