diff --git a/.github/scripts/install_otel_test_wheels.py b/.github/scripts/install_otel_test_wheels.py new file mode 100644 index 000000000..938996aca --- /dev/null +++ b/.github/scripts/install_otel_test_wheels.py @@ -0,0 +1,101 @@ +"""Install and verify built artifacts without silently using editable sources.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import importlib.metadata +import json +import subprocess +import sys +import zipfile +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] +CORE = "aws-durable-execution-sdk-python" +OTEL = CORE + "-otel" + + +def built_wheel(package: str) -> Path: + directory = ROOT / "packages" / package + module = package.replace("-", "_") + about = ast.parse((directory / "src" / module / "__about__.py").read_text()) + version = next( + ast.literal_eval(node.value) + for node in about.body + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "__version__" + for target in node.targets + ) + ) + wheels = list((directory / "dist").glob(f"{module}-{version}-*.whl")) + if len(wheels) != 1: + raise ValueError(f"Build exactly one {package} {version} wheel first: {wheels}") + return wheels[0] + + +def verify(wheel: Path, package: str) -> None: + installed = importlib.metadata.distribution(package) + direct = json.loads(installed.read_text("direct_url.json") or "{}") + assert not direct.get("dir_info", {}).get("editable"), direct + digest = hashlib.sha256(wheel.read_bytes()).hexdigest() + assert direct["archive_info"]["hashes"]["sha256"] == digest, direct + module = package.replace("-", "_") + with zipfile.ZipFile(wheel) as archive: + sources = [ + name + for name in archive.namelist() + if name.startswith(module + "/") and name.endswith(".py") + ] + assert sources + for name in sources: + path = Path(installed.locate_file(name)).resolve() + assert "site-packages" in path.parts, path + assert path.read_bytes() == archive.read(name), path + print( + json.dumps( + { + "package": package, + "version": installed.version, + "wheel": str(wheel), + "sha256": digest, + "verified_sources": len(sources), + } + ) + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--legacy-plugin", action="store_true") + args = parser.parse_args() + packages = [CORE] if args.legacy_plugin else [CORE, OTEL] + wheels = [built_wheel(package) for package in packages] + subprocess.run( + [ + sys.executable, + "-m", + "pip", + "install", + "--no-index", + "--no-deps", + "--force-reinstall", + *map(str, wheels), + ], + check=True, + ) + for wheel, package in zip(wheels, packages, strict=True): + verify(wheel, package) + if args.legacy_plugin: + assert importlib.metadata.version(OTEL) == "1.0.0" + import aws_durable_execution_sdk_python_otel as otel + + assert "site-packages" in Path(otel.__file__).resolve().parts + subprocess.run([sys.executable, "-m", "pip", "check"], check=True) + + +if __name__ == "__main__": + main() diff --git a/.github/tests/otel_lifecycle_compatibility_test.py b/.github/tests/otel_lifecycle_compatibility_test.py new file mode 100644 index 000000000..8e9ce857b --- /dev/null +++ b/.github/tests/otel_lifecycle_compatibility_test.py @@ -0,0 +1,162 @@ +"""Exercise real installed version pairs through the public durable handler.""" + +from __future__ import annotations + +import contextvars +from datetime import UTC, datetime +from importlib.metadata import version +import os +from pathlib import Path +import threading +from types import SimpleNamespace + +import pytest +from aws_durable_execution_sdk_python import durable_execution +from aws_durable_execution_sdk_python import execution as core_execution +from aws_durable_execution_sdk_python.lambda_service import ( + ExecutionDetails, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +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 opentelemetry import baggage, context, trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from packaging.version import Version + + +class NoNetworkClient: + def __getattr__(self, name: str): + raise AssertionError(f"Unexpected service access: {name}") + + +@pytest.mark.parametrize("view", [InvocationOtelPlugin, ExecutionOtelPlugin]) +@pytest.mark.parametrize("order", ["alone", "baggage-first", "baggage-last"]) +@pytest.mark.parametrize("ambient_present", [False, True]) +def test_installed_pair_preserves_context_and_documents_legacy_fallback( + view, order, ambient_present +): + legacy = os.environ.get("OTEL_COMPAT_LEGACY") == "1" + assert Version(version("aws-durable-execution-sdk-python")) >= Version("2.1.0") + assert version("aws-durable-execution-sdk-python-otel") == ( + "1.0.0" if legacy else "1.1.0" + ) + assert "site-packages" in Path(core_execution.__file__).resolve().parts + import aws_durable_execution_sdk_python_otel as installed_otel + + assert "site-packages" in Path(installed_otel.__file__).resolve().parts + + def run() -> None: + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer("installed-lifecycle") + phases = [] + worker_threads = [] + caller_thread = threading.get_ident() + + class BaggagePlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + self.token = context.attach(baggage.set_baggage("customer", "present")) + phases.append("baggage-start") + worker_threads.append(threading.get_ident()) + + def on_invocation_end(self, info): + context.detach(self.token) + phases.append("baggage-end") + worker_threads.append(threading.get_ident()) + + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + plugins = { + "alone": [plugin], + "baggage-first": [BaggagePlugin(), plugin], + "baggage-last": [plugin, BaggagePlugin()], + }[order] + body_parents = [] + + @durable_execution(plugins=plugins, boto3_client=NoNetworkClient()) + def handler(event, durable_context): + phases.append("body") + worker_threads.append(threading.get_ident()) + assert baggage.get_baggage("incoming") == "keep" + assert baggage.get_baggage("customer") == ( + None if order == "alone" else "present" + ) + parent = trace.get_current_span().get_span_context() + body_parents.append(parent) + with tracer.start_as_current_span("customer-span"): + pass + return "ok" + + operation = Operation( + operation_id="installed", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + start_timestamp=datetime(2026, 10, 8, tzinfo=UTC), + execution_details=ExecutionDetails(input_payload="{}"), + ) + event = { + "DurableExecutionArn": "test-arn/installed", + "CheckpointToken": "token", + "InitialExecutionState": { + "Operations": [operation.to_json_dict()], + "NextMarker": "", + }, + } + lambda_context = SimpleNamespace( + aws_request_id="installed", + client_context=None, + identity=None, + _epoch_deadline_time_in_ms=0, + invoked_function_arn="test-arn", + tenant_id=None, + ) + ambient = tracer.start_span("host") if ambient_present else trace.INVALID_SPAN + host = baggage.set_baggage( + "incoming", "keep", trace.set_span_in_context(ambient) + ) + token = context.attach(host) + try: + for _ in range(2): + assert handler(event, lambda_context)["Status"] == "SUCCEEDED" + assert context.get_current() is host + assert caller_thread not in worker_threads + assert phases == ( + ["body"] * 2 + if order == "alone" + else ["baggage-start", "body", "baggage-end"] * 2 + ) + spans = exporter.get_finished_spans() + if legacy and view is InvocationOtelPlugin: + # The released plugin does not attach an Invocation fallback. + assert body_parents == [ambient.get_span_context()] * 2 + else: + name = "Invocation" if view is InvocationOtelPlugin else "Workflow" + contexts = [span.context for span in spans if span.name == name] + assert contexts + assert all( + parent.is_valid and parent in contexts for parent in body_parents + ) + users = [span for span in spans if span.name == "customer-span"] + assert len(users) == 2 + assert [span.parent for span in users] == [ + parent if parent.is_valid else None for parent in body_parents + ] + assert plugin._context_tokens == {} + finally: + context.detach(token) + ambient.end() + provider.shutdown() + + contextvars.Context().run(run) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f3e682d08..6cbca4dd8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -84,8 +84,6 @@ jobs: run: hatch run types:check - name: Run tests + coverage run: hatch run test:cov - - name: Verify supported legacy core compatibility - run: hatch run test-pypi-otel-legacy:test - name: Build distribution run: | for pkg in packages/*/; do @@ -96,6 +94,10 @@ jobs: cd "$GITHUB_WORKSPACE" fi done + - name: Test installed core and OTel wheels + run: hatch run test-wheel-otel:test + - name: Test released OTel with the new core wheel + run: hatch run test-wheel-otel-legacy:test - name: Verify OTel wheel dependency contract run: | OTEL_WHEEL=$(find packages/aws-durable-execution-sdk-python-otel/dist \ diff --git a/.github/workflows/cloud-tests.yml b/.github/workflows/cloud-tests.yml index 8e54be4fd..da80a43e2 100644 --- a/.github/workflows/cloud-tests.yml +++ b/.github/workflows/cloud-tests.yml @@ -133,7 +133,9 @@ jobs: echo "Could not resolve the latest ADOT Python layer for $AWS_REGION" exit 1 fi - aws lambda get-layer-version-by-arn \ + # Parallel jobs can throttle this read; keep retries local to the lookup. + AWS_RETRY_MODE=standard AWS_MAX_ATTEMPTS=8 \ + aws lambda get-layer-version-by-arn \ --arn "$ADOT_LAYER_ARN" \ --region "$AWS_REGION" \ --query LayerVersionArn \ diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 216490be3..4b4f2215a 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -66,7 +66,7 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@a66037abbbfa55fde97f714e30f0bc262edefd63 + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@f18bd0b5f28c5c90e288d0fb8bca08a849b51863 with: language: python runs_on: codebuild-github-actions-runner-${{ github.run_id }}-${{ github.run_attempt }} diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 94176be2c..a3b6bc2f7 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -70,15 +70,25 @@ hatch run dev-otel:typecheck # type check otel only hatch run dev-examples:test # run examples tests only ``` -### PyPI release testing +### Installed package compatibility testing -To verify packages work against the published PyPI version of the core SDK (rather than the local workspace): +Build the core and OTel distributions with `hatch build` in each package, then run +these commands from the repository root: ```bash -hatch run test-pypi-otel:test # test new OTel capabilities against capable installed core -hatch run test-pypi-otel-legacy:test # valid registrations/lifecycles on supported core 2.0.x -hatch run test-pypi-examples:test # test examples against PyPI core SDK -``` +hatch run test-wheel-otel:test # full OTel suite on the two built wheels +hatch run test-wheel-otel-legacy:test # released OTel 1.0.0 with the built core +hatch run test-pypi-examples:test # examples against the published core +``` + +The wheel lanes have no editable workspace members. They verify installed source +bytes and artifact hashes before exercising public handlers. OTel 1.1 requires +the redesigned core 2.1.0 lifecycle; it no longer claims compatibility with core +2.0.x. Publish core first. Building both wheels lets CI verify the intended pair +before that minimum is available on PyPI. The legacy-plugin lane documents the +actual core-only upgrade: host isolation is provided by the new core, while old +Invocation OTel does not gain the new fallback. Workspace tests continue to cover +the complete current implementation with `hatch run dev-otel:test`. ### Package-level commands 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 eef6d3eed..2998740d8 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 @@ -18,6 +18,7 @@ """ import json +from threading import Lock from typing import Any from aws_durable_execution_sdk_python.config import Duration, ParallelConfig @@ -31,13 +32,19 @@ ) +_log_lock = Lock() + + def _emit(record: dict[str, Any], execution_arn: str | None) -> None: # Prefix every plugin record with the execution ARN as a top-level field so # the conformance runner's CloudWatch JSON filter can scope logs to a single # execution. Omit the field when the ARN is unset (never invent a value). if execution_arn: record = {"durableExecutionArn": execution_arn, **record} - print(json.dumps(record), flush=True) + # Start and end hooks can run on different threads. Keep print's separate + # body/newline writes together so the runner receives one JSON per line. + with _log_lock: + print(json.dumps(record), flush=True) class WaitReplayFlagPlugin(DurableInstrumentationPlugin): diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/tests/test_plugin_wait_replay_flag.py b/packages/aws-durable-execution-sdk-python-conformance-tests/tests/test_plugin_wait_replay_flag.py new file mode 100644 index 000000000..0c6006043 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/tests/test_plugin_wait_replay_flag.py @@ -0,0 +1,72 @@ +"""Concurrent plugin callbacks must emit separate parseable JSON records.""" + +from __future__ import annotations + +import importlib.util +import json +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Event +from typing import Any + +import pytest + +import aws_durable_execution_sdk_python.execution as execution + + +def test_concurrent_wait_hooks_keep_complete_stdout_records( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Import the real fixture without constructing a Lambda client: this test + # exercises its stdout producer, not a deployed durable invocation. + monkeypatch.setattr(execution, "durable_execution", lambda **_: lambda fn: fn) + path = ( + Path(__file__).resolve().parents[1] + / "handlers/plugin/plugin_wait_replay_flag.py" + ) + spec = importlib.util.spec_from_file_location("wait_replay_log_fixture", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + first_body = Event() + second_ready = Event() + second_done = Event() + chunks: list[str] = [] + + class FragmentingStdout: + def write(self, text: str) -> int: + chunks.append(text) + if '"operation-start"' in text: + first_body.set() + # print writes its body and newline separately. Permit the + # other real hook to run between them unless _emit serializes it. + second_done.wait(0.2) + return len(text) + + def flush(self) -> None: + pass + + records: list[dict[str, Any]] = [ + {"plugin": "CONFPLUGIN", "hook": "operation-start", "name": "long"}, + {"plugin": "CONFPLUGIN", "hook": "operation-end", "name": "short"}, + ] + + def emit_end() -> None: + second_ready.set() + assert first_body.wait(2) + module._emit(records[1], "execution-arn") + second_done.set() + + with monkeypatch.context() as capture: + capture.setattr("sys.stdout", FragmentingStdout()) + with ThreadPoolExecutor(max_workers=2) as executor: + end = executor.submit(emit_end) + assert second_ready.wait(2) + start = executor.submit(module._emit, records[0], "execution-arn") + start.result(timeout=2) + end.result(timeout=2) + actual = [json.loads(line) for line in "".join(chunks).splitlines() if line] + assert actual == [ + {"durableExecutionArn": "execution-arn", **record} for record in records + ] diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 2c7e1bb1b..2a095bf9c 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -55,9 +55,11 @@ DURABLE_EXECUTION_PLUGINS=otel-execution cold start, so the handler does not need to import or explicitly register the plugin. -Automatic mutual-exclusion validation requires core SDK 2.1.0 or later together -with OTel 1.1.0 or later. OTel 1.1 remains compatible with core 2.0.x for existing -valid registrations; those older cores do not enforce the new group metadata. +OTel 1.1.0 requires core SDK 2.1.0 or later for the invocation-worker lifecycle +and automatic mutual-exclusion validation. Publish the redesigned core first, +then the plugin. Core 2.0.x runs invocation hooks on the caller and cannot provide +the new plugin's handler propagation and host-context isolation; that version +pair is not supported by OTel 1.1.0. Configure only one OTel view on every core version. `InvocationOtelPlugin` shows work within each Lambda invocation; `ExecutionOtelPlugin` shows logical operations across the whole execution. They emit overlapping telemetry and @@ -180,6 +182,38 @@ lambda_.Function( ) ``` +### Handler context propagation + +The core runs the existing Start hooks, handler, output preparation, resource +cleanup and End hooks on one invocation worker. Start finishes before checkpoint +processing begins. End follows registered-branch joins and checkpoint shutdown, +and output serialization or checkpoint errors retain their normal classification. +There is no separate handler-context plugin API. + +During Start, Invocation view preserves a valid active span on the canonical +execution trace; when it is absent or unrelated, it attaches the Invocation span. +Execution view attaches the Workflow span, including when an incoming span is on +the same trace. The Invocation span's own ambient parenting is separate from the +active context supplied to handler instrumentation. Both views preserve baggage. + +Start hooks run in registration order. A later successful plugin that deliberately +sets or clears the active span wins; OTel does not apply a second correction pass. +Place OTel after a span-replacing plugin when OTel's view-specific context is +desired. Baggage-only plugins that extend the current context can appear on either +side. The SDK also copies the coordinator's current bindings for each `map` or +`parallel` branch admission/resume when plugins are registered. Baggage and +other successful bindings reach these SDK-managed callbacks, while a branch's +changes cannot leak into siblings, the coordinator or a reused worker. +End hooks also retain registration order and reset tokens in their owning +Context; they do not promise a reverse-stack observation of other plugins' spans. +Existing registration, factory lifetime and checkpoint formats are unchanged. + +Upgrade both core to 2.1+ and OTel to 1.1+ for these guarantees. The released +OTel 1.0 plugin can run on the new core with worker/host isolation, but its +Invocation view does not attach the new fallback; a core-only upgrade does not +supply that plugin behavior. CI tests the new pair as installed wheels and tests +the actual released plugin separately, including before core 2.1 is on PyPI. + ### 3. In your Lambda handler (index.py) ```python @@ -330,6 +364,17 @@ OTel `OK` only for `SUCCEEDED`, `ERROR` when error details are delivered, and `CANCELLED`, `TIMED_OUT`, and `STOPPED`. The original durable operation status remains in `durable.operation.status`. +### Invocation context isolation + +With core 2.1+, invocation hooks and the handler run on one worker in an +invocation-local Context initialized from the host's bindings. Successful Start +bindings are visible to later hooks, the handler and resource cleanup. The host's +bindings remain unchanged after return, even if a plugin fails during Start or +End. A failed Start's bindings are discarded for subsequent work, while End runs +in that Start's original Context to preserve token ownership. This isolates +context-variable bindings; it does not undo mutations to shared objects or +external side effects. The isolation applies when plugins are registered. + ### Log Correlation When `enrich_logger=True` (the default), the plugin installs a logging filter on @@ -430,7 +475,7 @@ setups. ## Requirements - Python >= 3.11 -- `aws-durable-execution-sdk-python` >= 2.0.0 (core >= 2.1.0 with OTel >= 1.1.0 for automatic view-exclusivity validation) +- `aws-durable-execution-sdk-python` >= 2.1.0 (release the redesigned core before OTel 1.1.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 cbc2f97a2..461718714 100644 --- a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml @@ -22,7 +22,7 @@ classifiers = [ "Programming Language :: Python :: Implementation :: PyPy", ] dependencies = [ - "aws-durable-execution-sdk-python>=2.0.0", + "aws-durable-execution-sdk-python>=2.1.0", ] [project.entry-points."aws_durable_execution.plugins"] 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 efaf3b965..0fb974fe1 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 @@ -73,6 +73,7 @@ {InvocationStatus.SUCCEEDED, InvocationStatus.FAILED} ) _TIMESTAMP_STEP_NANOS = 1_000 +_INVOCATION_CONTEXT_KEY = "__invocation_context__" _SpanAttributes = dict[str, str | bool | int] @@ -299,12 +300,8 @@ def get_current_span_context(self) -> SpanContext | None: 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. - 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 - resolve it. It also covers code between top-level operations, where - detaching the operation scope restores a context with no durable - span. + 2. The invocation span from the plugin registry, including lifecycle + phases where another plugin changed the active context. Returns: A valid SpanContext, or None if no span is active. @@ -593,6 +590,17 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: attributes=self._extract_attributes(info), ) + # Start and End share the invocation worker and token-owning Context. + # Retain a valid same-trace ambient parent; otherwise bind this invocation. + ambient = trace.get_current_span().get_span_context() + invocation_span = self._get_span(None) + if invocation_span is not None and ( + not ambient.is_valid or ambient.trace_id != self._execution_trace_id + ): + self._attach_context( + _INVOCATION_CONTEXT_KEY, trace.set_span_in_context(invocation_span) + ) + # Cover handlers installed after construction as well. if self._enrich_logger: install_log_filter(self) @@ -673,6 +681,10 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._reset_state() return + # User work and checkpoint cleanup have finished. Release Start bindings + # in the Context that created their tokens before closing the spans. + self._detach_remaining_contexts() + # 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: diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_branch_context_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_branch_context_int.py new file mode 100644 index 000000000..c2ae779af --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_branch_context_int.py @@ -0,0 +1,179 @@ +"""Public SDK branch workers inherit instrumented bindings without sharing them.""" + +from __future__ import annotations + +import contextvars +import json +import threading +from typing import Any + +import pytest +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.config import Duration, MapConfig, ParallelConfig +from aws_durable_execution_sdk_python.concurrency.models import BatchResult +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStatus, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import baggage, context, trace +from opentelemetry.sdk.trace import TracerProvider + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + ("view", "failed_start"), + [ + (None, False), + (InvocationOtelPlugin, False), + (ExecutionOtelPlugin, False), + (InvocationOtelPlugin, True), + (ExecutionOtelPlugin, True), + ], +) +@pytest.mark.parametrize("kind", ["map", "parallel"]) +@pytest.mark.parametrize("concurrency", [1, 2]) +@pytest.mark.parametrize("resume", [False, True]) +def test_public_branch_context_propagation_and_isolation( + monkeypatch: pytest.MonkeyPatch, + view: type[InvocationOtelPlugin] | type[ExecutionOtelPlugin] | None, + failed_start: bool, + caplog: pytest.LogCaptureFixture, + kind: str, + concurrency: int, + resume: bool, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("branch-invocation-marker", default="unset") + observations: list[tuple[int, str, Any, bool]] = [] + bodies: list[int] = [] + visits = {0: 0, 1: 0} + statuses: list[InvocationStatus] = [] + starts: list[int] = [] + completed_ends: list[int] = [] + lock = threading.Lock() + barrier = threading.Barrier(2) if concurrency == 2 else None + provider = TracerProvider() + host = context.get_current() + + class Bind(DurableInstrumentationPlugin): + def on_invocation_start(self, info: Any) -> None: + starts.append(threading.get_ident()) + self.token = marker.set("plugin") + self.bag = context.attach(baggage.set_baggage("tenant", "present")) + + def on_invocation_end(self, info: Any) -> None: + context.detach(self.bag) + marker.reset(self.token) + # Record completion only after both tokens were reset in their owner. + completed_ends.append(threading.get_ident()) + statuses.append(info.status) + + class Broken(DurableInstrumentationPlugin): + def on_invocation_start(self, _info: Any) -> None: + marker.set("partial") + context.attach(baggage.set_baggage("tenant", "partial")) + raise ValueError("expected failed Start") + + plugins: list[DurableInstrumentationPlugin] = [] + otel_plugin: InvocationOtelPlugin | ExecutionOtelPlugin | None = None + if view is not None: + plugins.append(Bind()) + if failed_start: + plugins.append(Broken()) + otel_plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + plugins.append(otel_plugin) + + def branch(child: DurableContext, index: int) -> str: + with lock: + visits[index] += 1 + first_visit = visits[index] == 1 + observations.append( + ( + index, + marker.get(), + baggage.get_baggage("tenant"), + trace.get_current_span().get_span_context().is_valid, + ) + ) + if view is not None: + # Deliberately leave a binding behind: another logical branch or + # resume on this pool thread must still start with its parent's copy. + marker.set(f"branch-{index}") + if barrier is not None and first_visit: + barrier.wait(timeout=10) + + def step(_step: Any) -> str: + with lock: + bodies.append(index) + return f"saved-{index}" + + saved = child.step(step, name="save") + return saved + + def handler(_event: Any, durable: DurableContext) -> list[str]: + expected = "unset" if view is None else "plugin" + assert marker.get() == expected + token = marker.set("handler" if view is None else "plugin") + try: + result: BatchResult[str] = ( + durable.map( + [0, 1], + lambda child, item, index, items: branch(child, index), + name="mapped", + config=MapConfig(max_concurrency=concurrency), + ) + if kind == "map" + else durable.parallel( + [lambda child: branch(child, 0), lambda child: branch(child, 1)], + name="parallel", + config=ParallelConfig(max_concurrency=concurrency), + ) + ) + assert marker.get() == ("handler" if view is None else "plugin") + if resume: + durable.wait(Duration.from_seconds(1), name="resume") + return result.get_results() + finally: + marker.reset(token) + + wrapped = durable_execution(handler, plugins=plugins) + try: + with DurableFunctionTestRunner(handler=wrapped) as runner: + result = runner.run(input="{}", timeout=30) + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ["saved-0", "saved-1"] + assert sorted(bodies) == [0, 1] + assert {item[0] for item in observations} == {0, 1} + expected = "unset" if view is None else "plugin" + assert all( + item[1:] + == (expected, None if view is None else "present", view is not None) + for item in observations + ) + # The final wait replays the completed batch without rerunning branches. + assert visits == {0: 1, 1: 1} + if view is not None and resume: + assert InvocationStatus.PENDING in statuses + assert starts == completed_ends + if view is not None: + assert starts and statuses[-1] is InvocationStatus.SUCCEEDED + assert otel_plugin is not None and otel_plugin._context_tokens == {} + assert not any( + r.exc_info and r.name == "opentelemetry.context" for r in caplog.records + ) + assert marker.get() == "unset" + assert context.get_current() == host + finally: + provider.shutdown() 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 635634938..fa7036920 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 @@ -2,6 +2,7 @@ from __future__ import annotations +from contextlib import nullcontext from dataclasses import replace from datetime import UTC, datetime from typing import Any @@ -25,12 +26,16 @@ OperationType, StepDetails, ) +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin 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 opentelemetry import context as otel_context +from opentelemetry import trace +from opentelemetry.propagators.aws.aws_xray_propagator import AwsXRayPropagator from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -253,3 +258,242 @@ def handler_impl(_event: Any, context: DurableContext) -> str: assert completed_wait_span.end_time is not None assert after_resume.start_time is not None assert completed_wait_span.end_time <= after_resume.start_time + + +class InheritedInvocationPlugin(InvocationOtelPlugin): + pass + + +class InheritedExecutionPlugin(ExecutionOtelPlugin): + pass + + +@pytest.mark.parametrize( + ("plugin_type", "extra_context_plugin"), + [ + (view, extra) + for view in (InvocationOtelPlugin, ExecutionOtelPlugin) + for extra in (False, True, "same", "unrelated", "absent") + ] + + [(InheritedInvocationPlugin, False), (InheritedExecutionPlugin, "absent")], +) +@pytest.mark.parametrize("reverse_plugins", [False, True]) +@pytest.mark.parametrize("fail_after_resume", [False, True]) +@pytest.mark.parametrize("ambient_kind", ["same", "unrelated", "absent"]) +def test_handler_user_spans_inherit_context_across_resume_and_failure( + monkeypatch: pytest.MonkeyPatch, + plugin_type: type[InvocationOtelPlugin] | type[ExecutionOtelPlugin], + fail_after_resume: bool, + ambient_kind: str, + extra_context_plugin: bool | str, + reverse_plugins: bool, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + monkeypatch.setenv("_X_AMZN_TRACE_ID", XRAY_TRACE_HEADER) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig(tracer_provider=provider, enrich_logger=False) + ) + tracer = provider.get_tracer("customer") + before_context = otel_context.get_current() + calls: list[str] = [] + + def user_span(name: str) -> None: + # Ordinary instrumentation: the SDK/caller supplies the active parent. + span = tracer.start_span(name) + span.end() + + def step_body(_step_context: Any) -> str: + calls.append("step") + user_span("step-user") + return "saved" + + def handler_body(_event: Any, context: DurableContext) -> str: + if extra_context_plugin: + assert baggage.get_baggage("customer") == "present" + user_span("handler-entry") + saved = context.step(step_body, name="before-wait") + user_span("handler-after-step") + context.wait(Duration.from_seconds(1), name="context-wait") + user_span("handler-after-resume") + if fail_after_resume: + raise ValueError("handler failed after resume") + return saved + + # Start hooks share the invocation worker, in registration order. A later + # successful plugin may deliberately replace or clear the active span. + from opentelemetry import baggage + + class BaggagePlugin(DurableInstrumentationPlugin): + token: Any = None + + def on_invocation_start(self, _info: Any) -> None: + current = baggage.set_baggage("customer", "present") + if isinstance(extra_context_plugin, str): + # A real third-party invocation hook can bind or clear a span. + # User functions below still use ordinary implicit parenting. + parent = trace.SpanContext( + trace_id=(XRAY_TRACE_ID if extra_context_plugin == "same" else 1), + span_id=0xCAFE, + is_remote=False, + trace_flags=trace.TraceFlags(1), + ) + current = trace.set_span_in_context( + trace.INVALID_SPAN + if extra_context_plugin == "absent" + else trace.NonRecordingSpan(parent), + current, + ) + self.token = otel_context.attach(current) + + def on_invocation_end(self, _info: Any) -> None: + otel_context.detach(self.token) + self.token = None + + plugins: list[DurableInstrumentationPlugin] = [plugin] + if extra_context_plugin: + plugins.append(BaggagePlugin()) + if reverse_plugins: + plugins.reverse() + handler = durable_execution(handler_body, plugins=plugins) + remote = AwsXRayPropagator().extract({"X-Amzn-Trace-Id": XRAY_TRACE_HEADER}) + assert trace.get_current_span(remote).get_span_context().trace_id == XRAY_TRACE_ID + initial_operations = [_execution_operation()] + checkpoint, operations = _checkpoint_store(initial_operations) + ambient_ids: list[int] = [] + host_context = remote if ambient_kind == "same" else otel_context.Context() + try: + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as client_class: + client = Mock() + client.checkpoint = checkpoint + client_class.initialize_client.return_value = client + # Standard host instrumentation supplies a same-trace Lambda span. + host_scope = ( + tracer.start_as_current_span("lambda-first", context=host_context) + if ambient_kind != "absent" + else nullcontext() + ) + with host_scope: + host = trace.get_current_span() + ambient_ids.append(host.get_span_context().span_id) + first = handler(_event(initial_operations), _lambda_context()) + assert ( + trace.get_current_span().get_span_context() + == host.get_span_context() + ) + assert first["Status"] == InvocationStatus.PENDING.value + assert otel_context.get_current() == before_context + resumed_operations = [ + replace( + operation, + status=OperationStatus.SUCCEEDED, + end_timestamp=datetime.now(UTC), + ) + if operation.name == "context-wait" + else operation + for operation in operations.values() + ] + wait_id = next( + operation.operation_id + for operation in resumed_operations + if operation.name == "context-wait" + ) + checkpoint, _ = _checkpoint_store(resumed_operations) + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as client_class: + client = Mock() + client.checkpoint = checkpoint + client_class.initialize_client.return_value = client + host_scope = ( + tracer.start_as_current_span("lambda-resume", context=host_context) + if ambient_kind != "absent" + else nullcontext() + ) + with host_scope: + host = trace.get_current_span() + ambient_ids.append(host.get_span_context().span_id) + resumed = handler( + _event(resumed_operations, updated_operation_ids=[wait_id]), + _lambda_context(), + ) + assert ( + trace.get_current_span().get_span_context() + == host.get_span_context() + ) + assert resumed["Status"] == ( + InvocationStatus.FAILED.value + if fail_after_resume + else InvocationStatus.SUCCEEDED.value + ) + assert calls == ["step"] + assert otel_context.get_current() == before_context + spans = exporter.get_finished_spans() + expected_parents: list[int | None] + if isinstance(extra_context_plugin, str) and not reverse_plugins: + # No second OTel correction pass: the later successful Start wins. + expected_parents = [ + None if extra_context_plugin == "absent" else 0xCAFE + ] * 2 + expected_trace_id = {"same": XRAY_TRACE_ID, "unrelated": 1, "absent": None}[ + extra_context_plugin + ] + elif issubclass(plugin_type, ExecutionOtelPlugin): + expected_parents = [derive_workflow_span_id(EXECUTION_ARN)] * 2 + expected_trace_id = XRAY_TRACE_ID + elif reverse_plugins and isinstance(extra_context_plugin, str): + expected_parents = ( + [0xCAFE] * 2 + if extra_context_plugin == "same" + else [ + span.context.span_id + for span in spans + if span.name == "Invocation" and span.context is not None + ] + ) + expected_trace_id = XRAY_TRACE_ID + else: + expected_parents = ( + [*ambient_ids] + if ambient_kind == "same" + else [ + span.context.span_id + for span in spans + if span.name == "Invocation" and span.context is not None + ] + ) + expected_trace_id = XRAY_TRACE_ID + for name in ("handler-entry", "handler-after-step"): + users = [span for span in spans if span.name == name] + assert len(users) == 2 + assert [span.parent.span_id if span.parent else None for span in users] == ( + expected_parents + ) + assert all(span.context is not None for span in users) + if expected_trace_id is None: + assert all( + span.parent is None and span.context.trace_id != XRAY_TRACE_ID + for span in users + ) + else: + assert all(span.context.trace_id == expected_trace_id for span in users) + after_resume = next( + span for span in spans if span.name == "handler-after-resume" + ) + assert ( + after_resume.parent.span_id if after_resume.parent else None + ) == expected_parents[1] + step_user = next(span for span in spans if span.name == "step-user") + assert step_user.parent is not None + assert any( + span.name == "before-wait attempt 1" + and span.context is not None + and span.context.span_id == step_user.parent.span_id + for span in spans + ) + finally: + provider.shutdown() 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 12edcc0ae..214a3a89d 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 @@ -28,6 +28,7 @@ from opentelemetry import baggage, trace from opentelemetry.context import Context from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.sampling import ALWAYS_ON, ALWAYS_OFF, Sampler from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import ( @@ -1622,3 +1623,66 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): plugin.on_invocation_end(_invocation_end_info()) assert plugin._context_tokens == {} + + +@pytest.mark.parametrize("ambient_kind", ["same", "unrelated", "absent"]) +@pytest.mark.parametrize("raises", [False, True]) +@pytest.mark.parametrize("sampler", [ALWAYS_ON, ALWAYS_OFF]) +def test_invocation_hooks_bind_parent_and_restore_baggage( + ambient_kind: str, raises: bool, sampler: Sampler +) -> None: + provider = TracerProvider(sampler=sampler) + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + caller = baggage.set_baggage("tenant", "hook-test", Context()) + ambient = None + if ambient_kind != "absent": + ambient = SpanContext( + trace_id=_to_otel_trace_id(EXECUTION_ARN, START_TIME) + if ambient_kind == "same" + else 1, + span_id=0x42, + is_remote=False, + trace_flags=TraceFlags(1), + ) + caller = trace.set_span_in_context(NonRecordingSpan(ambient), caller) + token = otel_context.attach(caller) + error = ValueError("handler error") + try: + plugin.on_invocation_start(_invocation_start_info()) + expected = trace.get_current_span().get_span_context() + assert expected.trace_id == _to_otel_trace_id(EXECUTION_ARN, START_TIME) + assert expected.span_id == derive_workflow_span_id(EXECUTION_ARN) + + def body() -> None: + active = trace.get_current_span().get_span_context() + assert active == expected + assert active.is_valid + assert baggage.get_baggage("tenant") == "hook-test" + if raises: + raise error + + try: + if raises: + with pytest.raises(ValueError) as caught: + body() + assert caught.value is error + else: + body() + finally: + plugin.on_invocation_end( + _invocation_end_info( + InvocationStatus.FAILED if raises else InvocationStatus.SUCCEEDED + ) + ) + assert plugin._context_tokens == {} + assert otel_context.get_current() is caller + assert baggage.get_baggage("tenant") == "hook-test" + finally: + otel_context.detach(token) + provider.shutdown() 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 1894d2b48..ba373ad8d 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 @@ -724,6 +724,7 @@ def test_operation_end_without_start_links_previous_logical_operation(): assert ( span.attributes["durable.operation.status"] == OperationStatus.SUCCEEDED.value ) + plugin.on_invocation_end(_invocation_end_info()) def test_continuation_span_uses_current_start_and_end_times(): @@ -753,6 +754,7 @@ def test_continuation_span_uses_current_start_and_end_times(): span = exporter.get_finished_spans()[0] assert invocation_span.start_time <= span.start_time assert before_callback <= span.start_time <= span.end_time <= after_callback + plugin.on_invocation_end(_invocation_end_info()) def test_resume_operation_timestamps_do_not_precede_current_invocation(): @@ -812,6 +814,7 @@ def test_resume_operation_timestamps_do_not_precede_current_invocation(): assert invocation_span.start_time <= after_resume_span.start_time assert after_resume_span.parent is not None assert after_resume_span.parent.span_id == invocation_span.context.span_id + plugin.on_invocation_end(_invocation_end_info()) def test_ordered_timestamps_are_thread_safe(): @@ -878,6 +881,7 @@ def test_retried_operation_uses_fresh_id_and_links_previous_logical_operation(): derive_workflow_span_id(EXECUTION_ARN), operation_id_to_span_id(EXECUTION_ARN, operation_id), } + plugin.on_invocation_end(_invocation_end_info()) def test_step_operation_span_parents_attempt_span(): @@ -1028,6 +1032,7 @@ def test_user_function_callbacks_emit_attempt_span_attributes(): == UserFunctionOutcome.SUCCEEDED.value ) assert "durable.operation.status" not in span.attributes + plugin.on_invocation_end(_invocation_end_info()) def test_step_attempt_span_name_includes_attempt_number(): @@ -1070,6 +1075,7 @@ def test_step_attempt_span_name_includes_attempt_number(): span = exporter.get_finished_spans()[0] assert span.name == "fetch-user attempt 2" + plugin.on_invocation_end(_invocation_end_info()) def test_step_attempt_span_name_defaults_to_first_attempt(): @@ -1112,6 +1118,7 @@ def test_step_attempt_span_name_defaults_to_first_attempt(): span = exporter.get_finished_spans()[0] assert span.name == "fetch-user attempt 1" + plugin.on_invocation_end(_invocation_end_info()) @pytest.mark.parametrize( @@ -1213,6 +1220,7 @@ def test_context_span_waits_for_terminal_status_and_omits_attempt_attributes( assert "durable.attempt.number" not in span.attributes assert "durable.attempt.outcome" not in span.attributes assert span.status.status_code is expected_span_status + plugin.on_invocation_end(_invocation_end_info()) def test_span_registry_helpers_can_be_called_from_multiple_threads(): @@ -1257,8 +1265,9 @@ def test_user_function_end_restores_enclosing_context(): # After the step, the enclosing context is restored and no scope is left # behind. Log correlation resolves the invocation span from the registry. assert otel_context.get_current() == enclosing_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} assert plugin.get_current_span_context().span_id == invocation_span_id + plugin.on_invocation_end(_invocation_end_info()) def test_user_function_start_preserves_baggage_in_current_context(): @@ -1292,7 +1301,8 @@ def test_user_function_end_restores_enclosing_context_on_failure(): ) assert otel_context.get_current() == enclosing_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} + plugin.on_invocation_end(_invocation_end_info()) def test_user_function_end_restores_enclosing_context_across_multiple_steps(): @@ -1309,8 +1319,9 @@ def test_user_function_end_restores_enclosing_context_across_multiple_steps(): # Between each step the context is back to where it started, and log # correlation still resolves the invocation span. assert otel_context.get_current() == enclosing_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} assert plugin.get_current_span_context().span_id == invocation_span_id + plugin.on_invocation_end(_invocation_end_info()) # ---------------------------------------------------------------------- @@ -1332,6 +1343,7 @@ def test_get_current_span_context_returns_invocation_span_at_top_level(): invocation_span = plugin._get_span(None) assert span_context is not None assert span_context.span_id == invocation_span.get_span_context().span_id + plugin.on_invocation_end(_invocation_end_info()) def test_get_current_span_context_returns_operation_span_inside_step(): @@ -1363,6 +1375,7 @@ def test_get_current_span_context_returns_invocation_span_between_steps(): invocation_span = plugin._get_span(None) assert span_context is not None assert span_context.span_id == invocation_span.get_span_context().span_id + plugin.on_invocation_end(_invocation_end_info()) # ---------------------------------------------------------------------- @@ -1441,11 +1454,12 @@ def test_top_level_step_end_falls_back_to_invocation_for_correlation(): plugin.on_user_function_start(_user_function_start_info(operation_id)) plugin.on_user_function_end(_user_function_end_info(operation_id)) - # No durable span is attached at the top level, so the registry fallback - # supplies the invocation span for log correlation. + # The invocation Start binding is restored after the step, including + # between-step user instrumentation and log correlation. assert otel_context.get_current() == enclosing_context - assert not trace.get_current_span().get_span_context().is_valid + assert trace.get_current_span().get_span_context().span_id == invocation_span_id assert plugin.get_current_span_context().span_id == invocation_span_id + plugin.on_invocation_end(_invocation_end_info()) def test_get_current_span_context_returns_context_span_between_nested_steps(): @@ -1709,6 +1723,7 @@ def test_replayed_context_span_links_previous_logical_operation(): derive_workflow_span_id(EXECUTION_ARN), operation_id_to_span_id(EXECUTION_ARN, operation_id), } + plugin.on_invocation_end(_invocation_end_info()) def test_checkpointed_context_first_span_uses_deterministic_id(): @@ -1867,7 +1882,8 @@ def test_child_context_end_restores_context_active_before_it(): ) assert otel_context.get_current() == enclosing_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} + plugin.on_invocation_end(_invocation_end_info()) def test_nested_scopes_are_released_without_accumulating(): @@ -1891,13 +1907,14 @@ def test_nested_scopes_are_released_without_accumulating(): ) # The inner step restored the child-context scope, not a copy of it. assert otel_context.get_current() == inside_context - assert set(plugin._context_tokens) == {context_id} + assert set(plugin._context_tokens) == {context_id, "__invocation_context__"} plugin.on_user_function_end( _user_function_end_info(context_id, operation_type=OperationType.CONTEXT) ) assert otel_context.get_current() == before_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} + plugin.on_invocation_end(_invocation_end_info()) def test_invocation_end_releases_scope_of_suspended_user_function(): @@ -1981,7 +1998,7 @@ def test_detach_ignores_token_attached_on_another_thread(): plugin._detach_context("step-1:attempt:1") - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} assert otel_context.get_current() == before_context plugin.on_invocation_end(_invocation_end_info()) @@ -2072,7 +2089,7 @@ def test_reentered_step_attempt_releases_the_previous_scope(): plugin.on_user_function_end(_user_function_end_info(operation_id)) assert otel_context.get_current() == before_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} plugin.on_invocation_end(_invocation_end_info()) @@ -2166,7 +2183,7 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): ) ) assert otel_context.get_current() == before_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} # Neither span is ended: both operations are still in flight. assert not exporter.get_finished_spans() @@ -2199,6 +2216,69 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): _user_function_end_info("ctx-outer", operation_type=OperationType.CONTEXT) ) assert otel_context.get_current() == before_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} plugin.on_invocation_end(_invocation_end_info()) + + +@pytest.mark.parametrize("ambient_kind", ["same", "unrelated", "absent"]) +@pytest.mark.parametrize("raises", [False, True]) +@pytest.mark.parametrize("sampler", [ALWAYS_ON, ALWAYS_OFF]) +def test_invocation_hooks_bind_parent_and_restore_baggage( + ambient_kind: str, raises: bool, sampler: Sampler +) -> None: + provider = TracerProvider(sampler=sampler) + plugin = InvocationOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + caller = baggage.set_baggage("tenant", "hook-test", Context()) + ambient = None + if ambient_kind != "absent": + ambient = SpanContext( + trace_id=_to_otel_trace_id(EXECUTION_ARN, START_TIME) + if ambient_kind == "same" + else 1, + span_id=0x42, + is_remote=False, + trace_flags=TraceFlags(1), + ) + caller = trace.set_span_in_context(NonRecordingSpan(ambient), caller) + token = otel_context.attach(caller) + error = ValueError("handler error") + try: + plugin.on_invocation_start(_invocation_start_info()) + expected = ( + ambient if ambient_kind == "same" else plugin.get_current_span_context() + ) + + def body() -> None: + active = trace.get_current_span().get_span_context() + assert active == expected + assert active.is_valid + assert baggage.get_baggage("tenant") == "hook-test" + if raises: + raise error + + try: + if raises: + with pytest.raises(ValueError) as caught: + body() + assert caught.value is error + else: + body() + finally: + plugin.on_invocation_end( + _invocation_end_info( + InvocationStatus.FAILED if raises else InvocationStatus.SUCCEEDED + ) + ) + assert plugin._context_tokens == {} + assert otel_context.get_current() is caller + assert baggage.get_baggage("tenant") == "hook-test" + finally: + otel_context.detach(token) + provider.shutdown() 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 ac9af3939..b7a34ebda 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 @@ -6,8 +6,7 @@ PACKAGE_ROOT = Path(__file__).resolve().parents[1] REPOSITORY_ROOT = PACKAGE_ROOT.parents[1] -CORE_DEPENDENCY = "aws-durable-execution-sdk-python>=2.0.0" -EXCLUSIVITY_TEST_CORE = "aws-durable-execution-sdk-python>=2.1.0" +CORE_DEPENDENCY = "aws-durable-execution-sdk-python>=2.1.0" TEST_OTEL_DEPENDENCIES = { "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", @@ -68,7 +67,7 @@ def test_test_environments_install_layer_provided_dependencies() -> None: "test", "dev-otel", "dev-examples", - "test-pypi-otel", + "test-wheel-otel", "test-pypi-examples", ): assert TEST_OTEL_DEPENDENCIES <= set( @@ -77,22 +76,7 @@ def test_test_environments_install_layer_provided_dependencies() -> None: assert TEST_OTEL_DEPENDENCIES <= set(environments["types"]["extra-dependencies"]) -def test_pypi_compatibility_environment_uses_compatible_core_sdk() -> None: - dependencies = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ - "envs" - ]["test-pypi-otel"]["dependencies"] - - assert EXCLUSIVITY_TEST_CORE in dependencies - - -def test_pypi_otel_environment_installs_lifecycle_test_runner() -> None: - dependencies = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ - "envs" - ]["test-pypi-otel"]["dependencies"] - assert "aws-durable-execution-sdk-python-testing>=1.2.1" in dependencies - - -def test_core_dependency_preserves_previously_supported_releases() -> None: +def test_core_dependency_requires_worker_lifecycle_release() -> None: dependencies = _load_pyproject(PACKAGE_ROOT / "pyproject.toml")["project"][ "dependencies" ] @@ -101,32 +85,23 @@ def test_core_dependency_preserves_previously_supported_releases() -> None: for value in dependencies if Requirement(value).name == "aws-durable-execution-sdk-python" ) - assert requirement.specifier.contains("2.0.0") - assert requirement.specifier.contains("2.0.1") + assert not requirement.specifier.contains("2.0.0") + assert not requirement.specifier.contains("2.0.1") assert requirement.specifier.contains("2.1.0") -def test_pypi_otel_environment_does_not_shadow_installed_core() -> None: - environment = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ +def test_wheel_lanes_do_not_shadow_installed_artifacts() -> None: + environments = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ "envs" - ]["test-pypi-otel"] - assert environment["workspace"]["members"] == [ - "packages/aws-durable-execution-sdk-python-otel" ] - - -def test_legacy_lane_retains_supported_core_20_range() -> None: - environment = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ - "envs" - ]["test-pypi-otel-legacy"] - requirement = next( - Requirement(value) - for value in environment["dependencies"] - if Requirement(value).name == "aws-durable-execution-sdk-python" + for name in ("test-wheel-otel", "test-wheel-otel-legacy"): + assert environments[name]["workspace"]["members"] == [] + assert environments[name]["detached"] is True + assert ( + "aws-durable-execution-sdk-python-otel==1.0.0" + in environments["test-wheel-otel-legacy"]["dependencies"] + ) + assert ( + "aws-durable-execution-sdk-python-testing>=1.2.1" + in environments["test-wheel-otel"]["dependencies"] ) - assert requirement.specifier.contains("2.0.0") - assert requirement.specifier.contains("2.0.1") - assert not requirement.specifier.contains("2.1.0") - assert environment["workspace"]["members"] == [ - "packages/aws-durable-execution-sdk-python-otel" - ] diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index 18bc0d5e7..4b2463e52 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -85,9 +85,8 @@ constraint applies to the combined `plugins=[...]` argument and Core 2.1+ with OTel 1.1+ rejects both at cold start with `PluginLoadError` naming the conflicting views; keep only one. Choose Invocation for work within each Lambda invocation or Execution for logical operations across the durable execution. -Unrelated instrumentation plugins can run alongside either view. Existing valid -registrations remain supported with OTel 1.1 on older core 2.0.x; those cores do -not implement the new exclusivity validation. +Unrelated instrumentation plugins can run alongside either view. OTel 1.1 requires core 2.1 or later for the invocation-worker lifecycle as well +as view-exclusivity validation. Release core 2.1 before releasing OTel 1.1. Plugin authors explicitly opt in by declaring `__durable_registration_api__ = 1` on a plugin class. That class and its subclasses @@ -112,6 +111,40 @@ in their hierarchy retain their existing attributes/helpers. The generic plugin neither attribute nor hook, and provider API version 1 and existing plugin lifecycle order are unchanged. +### Invocation worker lifecycle + +The existing `on_invocation_start` and `on_invocation_end` hooks run on the +handler's invocation worker, both in registration order. Start hooks finish +before background checkpoint processing begins. The same worker runs the handler +(including its `finally` blocks), prepares the existing serialized output or +error, checkpoints large results when necessary, joins registered branches while +checkpointing is available, stops and waits for checkpoint processing, then calls +End before returning the outcome. The caller shuts down the handler executor. +There is no separate handler-context hook or context-manager plugin API. + +When plugins are registered, the worker begins with a copy of the caller's +context-variable bindings. Successful Start bindings flow into later Start hooks, +the handler, serialization and resource cleanup. A later successful Start may +replace an earlier binding. If a Start hook raises, its new bindings are discarded +for subsequent work; its End still runs in the original Context so its tokens can +be reset. End hooks retain forward registration order, not reverse stack order, +so they must not rely on observing a stack-like unwind of other plugins' contexts. +Each SDK-managed `map` or `parallel` branch admission, including an in-process +resume, receives a fresh copy of the coordinator's bindings when plugins are +registered. Branch changes remain local even when pool threads are reused. +User-created threads retain normal Python context-variable behavior; without +registered plugins, SDK branch submission keeps its existing behavior. +The worker's invocation context is discarded on return, including after plugin +cleanup failures, leaving the host's bindings unchanged. This isolates bindings, +not mutations to shared objects or external side effects. Without plugins, the +handler retains its existing fresh-worker context behavior. + +These lifecycle guarantees require core 2.1.0 or later. The OTel 1.1 plugin +requires that core version and uses the existing Start/End hooks. Upgrading the +core alone with OTel 1.0 isolates worker bindings, but does not add the newer +Invocation-view fallback to that older plugin; upgrade both packages for it. +Provider API version 1 and the independent registration API are unchanged. + ## 🚀 Quick Start Install the execution SDK: diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py index 2a39c4eff..bafa43b19 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py @@ -2,6 +2,7 @@ from __future__ import annotations +import contextvars import heapq import logging import queue @@ -291,13 +292,25 @@ def execute( def submit(branch: Branch[CallableType, ResultType]) -> None: branch.start() - pool.submit( - self._branch_worker, - execution_state, - executor_context, - events, - branch.executable, - ) + if execution_state._plugin_executor._plugins: # noqa: SLF001 + # Every admission/resume gets its own Context. Branch bindings + # cannot leak into siblings, the coordinator, or reused workers. + pool.submit( + contextvars.copy_context().run, + self._branch_worker, + execution_state, + executor_context, + events, + branch.executable, + ) + else: + pool.submit( + self._branch_worker, + execution_state, + executor_context, + events, + branch.executable, + ) try: # Only rebuild the items snapshot after a terminal event changes 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 8ca5ef903..e59820007 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,10 +1,11 @@ from __future__ import annotations import contextlib +import contextvars import functools import json import logging -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import Future, ThreadPoolExecutor, wait from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any @@ -188,9 +189,9 @@ def durable_execution( logger.debug("Starting durable execution handler...") - plugin_executor = PluginExecutor(load_configured_plugins(plugins)) + configured_plugins = load_configured_plugins(plugins) + plugin_executor = PluginExecutor(configured_plugins) - @plugin_executor.handle_durable_output def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: invocation_input: DurableExecutionInvocationInput service_client: DurableServiceClient @@ -279,50 +280,13 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: ), ) - # Use ThreadPoolExecutor for concurrent execution of user code and background checkpoint processing - with ( - ThreadPoolExecutor( - max_workers=2, thread_name_prefix="dex-handler" - ) as executor, - contextlib.closing(execution_state) as execution_state, - ): - execution_operation = execution_state.get_execution_operation() - - # execute the plugins - plugin_executor.on_invocation_start( - execution_arn=invocation_input.durable_execution_arn, - lambda_context=context, - execution_start_time=( - execution_operation.start_timestamp - if execution_operation is not None - else None - ), - is_first_invocation=not has_prior_operations, - execution_input=input_event, - # Read the map through a callable rather than snapshotting it - # here: the invocation-end hook needs the state as of the end of - # the invocation, and neither hook pays for the conversion until - # a plugin actually reads it. - operations_provider=lambda: execution_state.operations, - updated_operation_ids=invocation_input.updated_operation_ids, - ) - # Thread 1: Run background checkpoint processing - executor.submit(execution_state.checkpoint_batches_forever) - - # Thread 2: Execute user function + def invoke_handler() -> MutableMapping[str, Any]: logger.debug( "%s entering user-space...", invocation_input.durable_execution_arn ) - user_future = executor.submit(func, input_event, durable_context) - - logger.debug( - "%s waiting for user code completion...", - invocation_input.durable_execution_arn, - ) - try: # Background checkpointing errors will propagate through CompletionEvent.wait() as BackgroundThreadError - result = user_future.result() + result = func(input_event, durable_context) # done with userland logger.debug( @@ -458,6 +422,78 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: return result + def invoke_worker() -> MutableMapping[str, Any]: + # This task owns Start, user cleanup, output preparation, resource + # cleanup and End. The caller owns the executor that runs this task. + with plugin_executor.run(): + checkpoint_future: Future[None] | None = None + + def wait_for_checkpoint() -> None: + if checkpoint_future is not None: + # Match the former executor join: observed checkpoint + # errors retain their existing classification paths. + wait((checkpoint_future,)) + + try: + with contextlib.ExitStack() as resources: + # LIFO cleanup: branches join with checkpointing alive; + # then close stops checkpointing and the wait completes. + resources.callback(wait_for_checkpoint) + resources.callback( + plugin_executor._run_in_invocation_context, + execution_state.close, + ) + execution_operation = execution_state.get_execution_operation() + + # execute the plugins + plugin_executor.on_invocation_start( + execution_arn=invocation_input.durable_execution_arn, + lambda_context=context, + execution_start_time=( + execution_operation.start_timestamp + if execution_operation is not None + else None + ), + is_first_invocation=not has_prior_operations, + execution_input=input_event, + # Read the map through a callable rather than snapshotting it + # here: the invocation-end hook needs the state as of the end of + # the invocation, and neither hook pays for the conversion until + # a plugin actually reads it. + operations_provider=lambda: execution_state.operations, + updated_operation_ids=invocation_input.updated_operation_ids, + ) + # No checkpoint work starts until all Start hooks finish. + checkpoint_future = executor.submit( + execution_state.checkpoint_batches_forever + ) + output = plugin_executor._run_in_invocation_context( + invoke_handler + ) + plugin_executor.on_invocation_end( + DurableExecutionInvocationOutput.from_dict(output) + ) + return output + except Exception as error: + plugin_executor.on_invocation_end( + DurableExecutionInvocationOutput.create_retry( + ErrorObject.from_exception(error) + ) + ) + raise + + with ThreadPoolExecutor( + max_workers=2, thread_name_prefix="dex-handler" + ) as executor: + if configured_plugins: + invocation_future = executor.submit( + contextvars.copy_context().run, invoke_worker + ) + else: + # Preserve the original fresh-worker context without plugins. + invocation_future = executor.submit(invoke_worker) + return invocation_future.result() + return wrapper 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 e9549d308..856fc68c6 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 @@ -1,15 +1,15 @@ from __future__ import annotations import contextlib +import contextvars import copy import datetime -import functools import logging from collections.abc import Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum -from typing import Any, Callable, MutableMapping, cast +from typing import Any, Callable, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -466,6 +466,10 @@ def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None + # Non-None only after a Start hook fails: later setup and invocation + # work use its pre-hook snapshot. End still uses each hook's token owner. + self._startup_context: contextvars.Context | None = None + self._invocation_contexts: list[contextvars.Context | None] = [] @contextlib.contextmanager def run(self): @@ -479,12 +483,14 @@ def run(self): finally: self._invocation_status = None self._operations_provider = None + self._startup_context = None + self._invocation_contexts.clear() # Shut down the thread pool, waiting for pending tasks to complete. if self._executor: self._executor.shutdown(wait=True) @staticmethod - def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: + def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> bool: """Invoke the appropriate plugin callback. Runs inside the thread pool.""" try: match info: @@ -507,18 +513,54 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: except Exception: # log and ignore the exception logger.exception("Plugin %s exception ignored", plugin.__class__.__name__) + return False + return True def execute_plugins(self, info, sync): if not self._executor: return - for plugin in self._plugins: - if sync: - # this is called synchronously, so plugins will be able to manipulate thread local objects - self._dispatch_plugin(plugin, info) + if sync and isinstance(info, InvocationStartInfo): + self._startup_context = None + self._invocation_contexts.clear() + for index, plugin in enumerate(self._plugins): + if sync and isinstance(info, InvocationStartInfo): + owner = self._startup_context + before = ( + owner.copy() if owner is not None else contextvars.copy_context() + ) + self._invocation_contexts.append(owner) + succeeded = ( + owner.run(self._dispatch_plugin, plugin, info) + if owner is not None + else self._dispatch_plugin(plugin, info) + ) + if not succeeded: + # A failing hook may have left new bindings with no reset token. + # Continue setup and the handler in the pre-hook snapshot. + self._startup_context = before + elif sync: + # End hooks must reset tokens in the Context that created them, + # even when a failed start hook moved later setup to a snapshot. + owner = ( + self._invocation_contexts[index] + if isinstance(info, InvocationEndInfo) + and index < len(self._invocation_contexts) + else None + ) + if owner is not None: + owner.run(self._dispatch_plugin, plugin, info) + else: + 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) + def _run_in_invocation_context(self, invoke: Callable[[], Any]) -> Any: + """Continue in the pre-hook Context only after a failed Start hook.""" + if self._startup_context is not None: + return self._startup_context.run(invoke) + return invoke() + def _snapshot_operation_infos( self, operations_provider: Callable[[], Mapping[str, Operation]] | None, @@ -836,28 +878,3 @@ def _is_terminal_status(status): OperationStatus.CANCELLED, OperationStatus.STOPPED, ] - - @property - def handle_durable_output(self): - def decorator(func: Callable[[Any, LambdaContext], MutableMapping[str, Any]]): - @functools.wraps(func) - def wrapper(event: Any, context: LambdaContext): - with self.run(): - 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) - ), - ) - raise - - return wrapper - - return decorator diff --git a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py index a93135663..c30980c73 100644 --- a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py @@ -1,5 +1,6 @@ """Tests for the concurrency module.""" +import contextvars import hashlib import json import queue @@ -77,7 +78,10 @@ from aws_durable_execution_sdk_python.operation.parallel import ( ParallelExecutor, ) -from aws_durable_execution_sdk_python.plugin import PluginExecutor +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + PluginExecutor, +) from aws_durable_execution_sdk_python.state import ( CheckpointedResult, ExecutionState, @@ -5275,3 +5279,44 @@ def predicate(s: CompletionStatus) -> CompletionDecision: # endregion Custom completion predicate (should_complete) integration tests + + +def test_instrumented_executor_isolates_bindings_on_a_reused_worker() -> None: + """Exercise actual admission/worker execution independently of durable I/O.""" + binding = contextvars.ContextVar("executor-unit-binding", default="empty") + seen: list[tuple[str, int]] = [] + + class RecordingExecutor(ConcurrentExecutor[Callable[[], str], str]): + def _execute_item_in_child_context( + self, + executor_context: DurableContext, + executable: Executable[Callable[[], str]], + ) -> str: + seen.append((binding.get(), threading.get_ident())) + binding.set(f"branch-{executable.index}") + return executable.func() + + executor = RecordingExecutor( + executables=[Executable(index, lambda: "ok") for index in range(2)], + max_concurrency=1, + completion_config=CompletionConfig(min_successful=2), + sub_type_top=OperationSubType.PARALLEL, + sub_type_iteration=OperationSubType.PARALLEL_BRANCH, + name_prefix="branch-", + serdes=None, + operation_id_namespace=_StubNamespace(), + ) + state = Mock(spec=ExecutionState) + state._plugin_executor = PluginExecutor([DurableInstrumentationPlugin()]) + token = binding.set("coordinator") + try: + result = executor.execute(state, Mock(spec=DurableContext)) + assert result.get_results() == ["ok", "ok"] + assert binding.get() == "coordinator" + assert [value for value, _thread in seen] == ["coordinator", "coordinator"] + assert seen[0][1] == seen[1][1] != threading.get_ident() + finally: + for call in state.register_branch_pool.call_args_list: + call.args[0].shutdown(wait=True) + binding.reset(token) + assert binding.get() == "empty" diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/branch_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/branch_worker_context_test.py new file mode 100644 index 000000000..a9104a745 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/branch_worker_context_test.py @@ -0,0 +1,129 @@ +"""Registered Start bindings reach independently resumed public branch workers.""" + +from __future__ import annotations + +import contextvars +import json +import threading +from typing import Any + +import pytest +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.config import Duration, MapConfig, ParallelConfig +from aws_durable_execution_sdk_python.concurrency.models import BatchResult +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner + + +@pytest.mark.parametrize("kind", ["map", "parallel"]) +@pytest.mark.parametrize("concurrency", [1, 2]) +@pytest.mark.parametrize("failed_start", [False, True]) +def test_partial_branch_resume_uses_fresh_successful_start_bindings( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + kind: str, + concurrency: int, + failed_start: bool, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("resumed-branch", default="host") + partial = contextvars.ContextVar[str]("partial-start") + observations: list[tuple[int, str]] = [] + visits = {0: 0, 1: 0} + bodies: list[int] = [] + statuses: list[InvocationStatus] = [] + lock = threading.Lock() + barrier = threading.Barrier(2) if concurrency == 2 else None + + class Bind(DurableInstrumentationPlugin): + token: contextvars.Token[str] | None = None + + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + self.token = marker.set("plugin") + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + statuses.append(info.status) + assert self.token is not None + marker.reset(self.token) + self.token = None + + class Broken(DurableInstrumentationPlugin): + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + marker.set("failed") + partial.set("failed") + raise ValueError("expected failed Start") + + plugin = Bind() + plugins: list[DurableInstrumentationPlugin] = [plugin] + if failed_start: + plugins.append(Broken()) + + def branch(child: DurableContext, index: int) -> str: + with lock: + visits[index] += 1 + first_visit = visits[index] == 1 + observations.append((index, marker.get())) + with pytest.raises(LookupError): + partial.get() + marker.set(f"branch-{index}") + if barrier is not None and first_visit: + barrier.wait(timeout=10) + + def step(_step: Any) -> str: + with lock: + bodies.append(index) + return f"saved-{index}" + + saved = child.step(step, name="save") + child.wait(Duration.from_seconds(index + 1), name="branch-wait") + return saved + + def handler(_event: Any, durable: DurableContext) -> list[str]: + assert marker.get() == "plugin" + result: BatchResult[str] + if kind == "map": + result = durable.map( + [0, 1], + lambda child, item, index, items: branch(child, index), + name="mapped", + config=MapConfig(max_concurrency=concurrency), + ) + else: + result = durable.parallel( + [lambda child: branch(child, 0), lambda child: branch(child, 1)], + name="parallel", + config=ParallelConfig(max_concurrency=concurrency), + ) + assert marker.get() == "plugin" + durable.wait(Duration.from_seconds(1), name="after-batch") + return result.get_results() + + wrapped = durable_execution(handler, plugins=plugins) + with DurableFunctionTestRunner(handler=wrapped) as runner: + result = runner.run(input="{}", timeout=30) + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ["saved-0", "saved-1"] + assert sorted(bodies) == [0, 1] + assert all(count >= 2 for count in visits.values()) + assert all(binding == "plugin" for _, binding in observations) + assert InvocationStatus.PENDING in statuses + assert statuses[-1] is InvocationStatus.SUCCEEDED + assert marker.get() == "host" and plugin.token is None + with pytest.raises(LookupError): + partial.get() + errors = [ + record + for record in caplog.records + if record.exc_info and record.name == "aws_durable_execution_sdk_python.plugin" + ] + assert len(errors) == (len(statuses) if failed_start else 0) + assert all( + record.exc_info is not None + and str(record.exc_info[1]) == "expected failed Start" + for record in errors + ) diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/handler_invocation_context_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/handler_invocation_context_int_test.py new file mode 100644 index 000000000..59f5c9b52 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/handler_invocation_context_int_test.py @@ -0,0 +1,79 @@ +"""Invocation context isolation across real suspension and replay.""" + +from __future__ import annotations + +import contextvars +import json +from typing import Any + +import pytest +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner + +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStartInfo, +) + + +@pytest.mark.parametrize("fail_after_resume", [False, True]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_invocation_plugins_restore_host_context_across_resume( + monkeypatch: pytest.MonkeyPatch, fail_after_resume: bool, reverse: bool +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("invocation-host", default="host") + boundaries: list[tuple[str, str]] = [] + body_calls: list[str] = [] + + class ScopePlugin(DurableInstrumentationPlugin): + def __init__(self, name: str) -> None: + self.name = name + self.token: contextvars.Token[str] | None = None + + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + self.token = marker.set(self.name) + + def on_invocation_end(self, _info: InvocationEndInfo) -> None: + assert self.token is not None + marker.reset(self.token) + self.token = None + + names = ["first", "second"] + if reverse: + names.reverse() + + def body(_event: Any, context: DurableContext) -> str: + assert marker.get() == names[-1] + + def step(_step_context: Any) -> str: + body_calls.append("step") + return "saved" + + saved = context.step(step, name="before-wait") + context.wait(Duration.from_seconds(1), name="resume") + assert marker.get() == names[-1] + if fail_after_resume: + raise ValueError("failure after resume") + return saved + + durable_handler = durable_execution( + body, plugins=[ScopePlugin(name) for name in names] + ) + + def host(event: Any, context: Any) -> Any: + before = marker.get() + try: + return durable_handler(event, context) + finally: + boundaries.append((before, marker.get())) + + with DurableFunctionTestRunner(handler=host) as runner: + result = runner.run(input="{}", timeout=15) + assert result.status.value == ("FAILED" if fail_after_resume else "SUCCEEDED") + if not fail_after_resume: + assert json.loads(result.result) == "saved" + assert body_calls == ["step"] + assert boundaries == [("host", "host"), ("host", "host")] diff --git a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py new file mode 100644 index 000000000..8cb081029 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py @@ -0,0 +1,589 @@ +"""Real public invocation-worker lifecycle and ContextVar ownership controls.""" + +from __future__ import annotations + +import contextvars +import threading +from concurrent.futures import ThreadPoolExecutor +from typing import Any +from unittest.mock import Mock + +import pytest + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.exceptions import ( + CheckpointError, + CheckpointErrorCategory, + InvocationError, + SuspendExecution, +) +from aws_durable_execution_sdk_python.execution import ( + DurableExecutionInvocationInputWithClient, + InitialExecutionState, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + DurableServiceClient, + ExecutionDetails, + Operation, + OperationStatus, + OperationType, + OperationUpdate, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, +) +from aws_durable_execution_sdk_python.state import ExecutionState + + +def invocation(client: Any = None) -> tuple[Any, Any]: + client = client or Mock(spec=DurableServiceClient) + client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="next", new_execution_state=CheckpointUpdatedExecutionState() + ) + event = DurableExecutionInvocationInputWithClient( + durable_execution_arn="test-arn/worker-lifecycle", + checkpoint_token="initial", + initial_execution_state=InitialExecutionState( + operations=[ + Operation( + operation_id="execution", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + execution_details=ExecutionDetails(input_payload="{}"), + ) + ], + next_marker="", + ), + service_client=client, + ) + context = Mock() + context.aws_request_id = "worker-request" + context.client_context = context.identity = context.invoked_function_arn = None + context._epoch_deadline_time_in_ms = 0 + context.tenant_id = None + return event, context + + +@pytest.mark.parametrize("outcome", ["success", "failure", "retry", "pending"]) +@pytest.mark.parametrize("mode", ["none", "healthy", "failed-start", "failed-end"]) +def test_public_worker_lifecycle_owns_context_and_outcome( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + outcome: str, + mode: str, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("worker-owner", default="fresh-worker") + phases: list[tuple[str, int, str]] = [] + boundary: list[tuple[str, int, str, str]] = [] + retry = InvocationError("original retry") + fail = ValueError("original body failure") + statuses: list[InvocationStatus] = [] + caller_thread = threading.get_ident() + + class Observer(ThreadPoolExecutor): + def submit(self, fn: Any, /, *args: Any, **kwargs: Any) -> Any: + role = ( + "checkpoint" + if getattr(fn, "__name__", "") == "checkpoint_batches_forever" + else "invocation" + ) + + def observed() -> Any: + before = marker.get() + try: + return fn(*args, **kwargs) + finally: + boundary.append((role, threading.get_ident(), before, marker.get())) + + return super().submit(observed) + + monkeypatch.setattr( + "aws_durable_execution_sdk_python.execution.ThreadPoolExecutor", Observer + ) + + class Plugin(DurableInstrumentationPlugin): + token: Any = None + + def on_invocation_start(self, info: InvocationStartInfo) -> None: + phases.append(("start", threading.get_ident(), marker.get())) + self.token = marker.set("start-binding") + if mode == "failed-start": + raise ValueError("startup failure") + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + phases.append(("end", threading.get_ident(), marker.get())) + statuses.append(info.status) + assert self.token is not None + marker.reset(self.token) + self.token = None + if mode == "failed-end": + raise ValueError("end failure") + + def body(_event: Any, _context: DurableContext) -> str: + phases.append(("body", threading.get_ident(), marker.get())) + marker.set("body-binding") + try: + if outcome == "failure": + raise fail + if outcome == "retry": + raise retry + if outcome == "pending": + raise SuspendExecution("test suspension") + return "ok" + finally: + phases.append(("finally", threading.get_ident(), marker.get())) + + handler = durable_execution(body, plugins=[] if mode == "none" else [Plugin()]) + event, context = invocation() + token = marker.set("host") + try: + for _ in range(2): + if outcome == "retry": + with pytest.raises(InvocationError) as caught: + handler(event, context) + assert caught.value is retry + else: + result = handler(event, context) + assert ( + result["Status"] + == { + "success": "SUCCEEDED", + "failure": "FAILED", + "pending": "PENDING", + }[outcome] + ) + if outcome == "failure": + assert result["Error"]["ErrorMessage"] == "original body failure" + assert marker.get() == "host" + finally: + marker.reset(token) + per_call = 2 if mode == "none" else 4 + for offset in range(0, len(phases), per_call): + batch = phases[offset : offset + per_call] + assert [p[0] for p in batch] == ( + ["body", "finally"] + if mode == "none" + else ["start", "body", "finally", "end"] + ) + assert len({p[1] for p in batch}) == 1 + assert batch[0][1] != caller_thread + bodies = [p for p in phases if p[0] == "body"] + expected = ( + "fresh-worker" + if mode == "none" + else "host" + if mode == "failed-start" + else "start-binding" + ) + assert [p[2] for p in bodies] == [expected, expected] + for role, tid, before, after in boundary: + if role == "invocation": + assert after == ("body-binding" if mode == "none" else before) + if mode != "none": + assert ( + statuses + == [ + { + "success": InvocationStatus.SUCCEEDED, + "failure": InvocationStatus.FAILED, + "retry": InvocationStatus.RETRY, + "pending": InvocationStatus.PENDING, + }[outcome] + ] + * 2 + ) + assert "different Context" not in caplog.text + + +@pytest.mark.parametrize("bad_first", [False, True]) +@pytest.mark.parametrize("outcome", ["success", "failure", "retry", "pending"]) +def test_failed_start_preserves_clean_bindings_and_original_token_owners( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + bad_first: bool, + outcome: str, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("failed-start-owner", default="host") + added = contextvars.ContextVar[str]("failed-start-added") + events: list[tuple[str, str, int]] = [] + cleanup_contexts: list[str] = [] + original_close = ExecutionState.close + + def close(state: ExecutionState) -> None: + cleanup_contexts.append(marker.get()) + with pytest.raises(LookupError): + added.get() + original_close(state) + + monkeypatch.setattr(ExecutionState, "close", close) + + class Plugin(DurableInstrumentationPlugin): + def __init__(self, name: str) -> None: + self.name = name + self.token: contextvars.Token[str] | None = None + self.added: contextvars.Token[str] | None = None + + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + events.append(("start", self.name, threading.get_ident())) + self.token = marker.set(self.name) + if self.name == "bad": + self.added = added.set("partial") + raise ValueError("partial startup") + + def on_invocation_end(self, _info: InvocationEndInfo) -> None: + events.append(("end", self.name, threading.get_ident())) + assert self.token is not None + marker.reset(self.token) + if self.added is not None: + added.reset(self.added) + + failure = InvocationError("retry") + + def body(_event: Any, _context: DurableContext) -> str: + events.append(("body", "handler", threading.get_ident())) + assert marker.get() == "healthy" + with pytest.raises(LookupError): + added.get() + if outcome == "retry": + raise failure + if outcome == "failure": + raise ValueError("body") + if outcome == "pending": + raise SuspendExecution("test suspension") + return "ok" + + names = ["bad", "healthy"] if bad_first else ["healthy", "bad"] + handler = durable_execution(body, plugins=[Plugin(name) for name in names]) + event, context = invocation() + for _ in range(2): + if outcome == "retry": + with pytest.raises(InvocationError) as caught: + handler(event, context) + assert caught.value is failure + else: + assert ( + handler(event, context)["Status"] + == {"success": "SUCCEEDED", "failure": "FAILED", "pending": "PENDING"}[ + outcome + ] + ) + assert marker.get() == "host" + with pytest.raises(LookupError): + added.get() + assert cleanup_contexts == ["healthy"] * 2 + for offset in (0, 5): + batch = events[offset : offset + 5] + assert [(x[0], x[1]) for x in batch] == [("start", n) for n in names] + [ + ("body", "handler") + ] + [("end", n) for n in names] + assert len({x[2] for x in batch}) == 1 + errors = [ + r.exc_info + for r in caplog.records + if r.exc_info and r.name == "aws_durable_execution_sdk_python.plugin" + ] + assert len(errors) == 2 and all(str(e[1]) == "partial startup" for e in errors) + + +@pytest.mark.parametrize("result_kind", ["normal", "bad-json", "large", "large-error"]) +def test_worker_prepares_output_joins_branches_waits_checkpoint_then_ends( + monkeypatch: pytest.MonkeyPatch, + result_kind: str, +) -> None: + import aws_durable_execution_sdk_python.execution as execution + + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + events: list[tuple[str, int]] = [] + started = threading.Event() + checkpoint_started = threading.Event() + closing = threading.Event() + branch_done = threading.Event() + real_checkpoint = ExecutionState.checkpoint_batches_forever + real_close = ExecutionState.close + real_stop = ExecutionState.stop_checkpointing + real_dumps = execution.json.dumps + result: Any = { + "normal": {"ok": True}, + "bad-json": object(), + "large": "x" * 256, + "large-error": None, + }[result_kind] + + def record(name: str) -> None: + events.append((name, threading.get_ident())) + + def checkpoint(state: ExecutionState) -> None: + assert started.is_set() + record("checkpoint-start") + checkpoint_started.set() + try: + real_checkpoint(state) + finally: + record("checkpoint-end") + + def stop(state: ExecutionState) -> None: + assert branch_done.is_set() + record("checkpoint-stop") + real_stop(state) + + def close(state: ExecutionState) -> None: + record("close") + closing.set() + real_close(state) + + def dumps(value: Any, *args: Any, **kwargs: Any) -> str: + if value is result or ( + isinstance(value, dict) and value.get("Status") == "FAILED" + ): + record("serialize") + return real_dumps(value, *args, **kwargs) + + monkeypatch.setattr(ExecutionState, "checkpoint_batches_forever", checkpoint) + monkeypatch.setattr(ExecutionState, "stop_checkpointing", stop) + monkeypatch.setattr(ExecutionState, "close", close) + monkeypatch.setattr(execution.json, "dumps", dumps) + if result_kind in ("large", "large-error"): + monkeypatch.setattr(execution, "LAMBDA_RESPONSE_SIZE_LIMIT", 64) + + class Plugin(DurableInstrumentationPlugin): + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + record("start") + assert not checkpoint_started.is_set() + started.set() + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + record("end") + assert branch_done.is_set() + assert any(x[0] == "checkpoint-end" for x in events) + assert info.status is ( + InvocationStatus.FAILED + if result_kind in ("bad-json", "large-error") + else InvocationStatus.SUCCEEDED + ) + + def body(_event: Any, ctx: DurableContext) -> Any: + assert checkpoint_started.wait(5) + record("body") + pool = ThreadPoolExecutor(max_workers=1) + ctx.state.register_branch_pool(pool) + + def late_branch() -> None: + assert closing.wait(5) + assert not ctx.state._checkpointing_stopped.is_set() + ctx.state.create_checkpoint( + OperationUpdate.create_execution_succeed(payload='"branch"'), + is_sync=True, + ) + record("branch-done") + branch_done.set() + + pool.submit(late_branch) + try: + if result_kind == "large-error": + raise ValueError("x" * 256) + return result + finally: + record("finally") + + handler = durable_execution(body, plugins=[Plugin()]) + event, context = invocation() + output = handler(event, context) + record("caller") + names = [x[0] for x in events] + assert ( + names.index("start") + < names.index("checkpoint-start") + < names.index("body") + < names.index("finally") + < names.index("serialize") + < names.index("close") + ) + assert ( + names.index("branch-done") + < names.index("checkpoint-stop") + < names.index("checkpoint-end") + < names.index("end") + < names.index("caller") + ) + worker_events = { + tid + for name, tid in events + if name + in {"start", "body", "finally", "serialize", "close", "checkpoint-stop", "end"} + } + assert len(worker_events) == 1 and threading.get_ident() not in worker_events + assert output["Status"] == ( + "FAILED" if result_kind in ("bad-json", "large-error") else "SUCCEEDED" + ) + + +@pytest.mark.parametrize("shape", ["method", "property", "dynamic"]) +def test_removed_handler_scope_api_is_not_inspected( + monkeypatch: pytest.MonkeyPatch, shape: str +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + calls = [] + + def legacy(*_args: Any) -> Any: + calls.append("unexpected") + raise AssertionError("Removed API called") + + choices: dict[str, dict[str, Any]] = { + "method": {"handler_context": legacy}, + "property": {"handler_context": property(legacy)}, + "dynamic": { + "__getattr__": lambda self, name: legacy() + if name == "handler_context" + else (_ for _ in ()).throw(AttributeError(name)) + }, + } + plugin = type("LegacyPlugin", (DurableInstrumentationPlugin,), choices[shape])() + handler = durable_execution(lambda _e, _c: "ok", plugins=[plugin]) + event, context = invocation() + assert handler(event, context)["Status"] == "SUCCEEDED" + assert calls == [] + + +@pytest.mark.parametrize( + "checkpoint_path", ["step-start", "large-result", "large-error"] +) +@pytest.mark.parametrize("retryable", [False, True]) +def test_checkpoint_failure_reports_prepared_outcome_after_worker_cleanup( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + checkpoint_path: str, + retryable: bool, +) -> None: + import aws_durable_execution_sdk_python.execution as execution + + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + if checkpoint_path != "step-start": + monkeypatch.setattr(execution, "LAMBDA_RESPONSE_SIZE_LIMIT", 128) + marker = contextvars.ContextVar("checkpoint-worker", default="host") + failure = CheckpointError( + "actual checkpoint failure", + error_category=( + CheckpointErrorCategory.INVOCATION + if retryable + else CheckpointErrorCategory.EXECUTION + ), + ) + events: list[tuple[str, int]] = [] + ends: list[tuple[InvocationEndInfo, str]] = [] + real_checkpoint = ExecutionState.checkpoint_batches_forever + real_close = ExecutionState.close + + def record(name: str) -> None: + events.append((name, threading.get_ident())) + + def checkpoint(state: ExecutionState) -> None: + record("checkpoint-start") + try: + real_checkpoint(state) + finally: + record("checkpoint-end") + + def close(state: ExecutionState) -> None: + record("close") + assert marker.get() == "plugin" + real_close(state) + record("closed") + + def service_checkpoint(*_args: Any, **_kwargs: Any) -> Any: + record("service-failure") + raise failure + + monkeypatch.setattr(ExecutionState, "checkpoint_batches_forever", checkpoint) + monkeypatch.setattr(ExecutionState, "close", close) + + class Plugin(DurableInstrumentationPlugin): + token: contextvars.Token[str] | None = None + + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + record("start") + self.token = marker.set("plugin") + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + record("end") + ends.append((info, marker.get())) + assert self.token is not None + marker.reset(self.token) + self.token = None + + def body(_event: Any, ctx: DurableContext) -> str: + record("body") + try: + if checkpoint_path == "step-start": + ctx.step(lambda _step: "ok", name="failed-checkpoint") + if checkpoint_path == "large-error": + raise ValueError("x" * 256) + return "x" * 256 + finally: + record("finally") + + plugin = Plugin() + handler = durable_execution(body, plugins=[plugin]) + client = Mock(spec=DurableServiceClient) + client.checkpoint.side_effect = service_checkpoint + event, context = invocation(client) + for _ in range(2): + events.clear() + ends.clear() + if retryable: + with pytest.raises(CheckpointError) as caught: + handler(event, context) + assert caught.value is failure + else: + output = handler(event, context) + assert output["Status"] == "FAILED" + assert output["Error"]["ErrorMessage"] == str(failure) + assert output["Error"]["ErrorType"].endswith(".CheckpointError") + record("caller") + assert marker.get() == "host" and plugin.token is None + assert len(ends) == 1 + info, active = ends[0] + assert active == "plugin" + assert info.status is ( + InvocationStatus.RETRY if retryable else InvocationStatus.FAILED + ) + assert info.error is not None and info.error.message == str(failure) + assert info.error.type is not None and info.error.type.endswith( + ".CheckpointError" + ) + names = [name for name, _tid in events] + assert ( + names.index("start") + < names.index("checkpoint-start") + < names.index("service-failure") + ) + assert ( + names.index("body") + < names.index("finally") + < names.index("close") + < names.index("closed") + < names.index("end") + < names.index("caller") + ) + assert ( + names.index("service-failure") + < names.index("checkpoint-end") + < names.index("end") + ) + workers = { + tid + for name, tid in events + if name in {"start", "body", "finally", "close", "closed", "end"} + } + assert len(workers) == 1 and threading.get_ident() not in workers + assert not any( + record.exc_info and record.name == "aws_durable_execution_sdk_python.plugin" + for record in caplog.records + ) diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py index 2670e6747..926b1dd0b 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py @@ -7,6 +7,7 @@ import pytest # Mock the executor.execute method +from aws_durable_execution_sdk_python.plugin import PluginExecutor from aws_durable_execution_sdk_python.concurrency.models import ( BatchItem, BatchItemStatus, @@ -184,6 +185,8 @@ def mock_run_in_child_context(func, name, config): # Create a minimal ExecutionState mock class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -222,6 +225,8 @@ def mock_run_in_child_context(func, name, config): return func("mock_context") class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -354,6 +359,8 @@ def callable_func(ctx, item, idx, items): ) as mock_execute: class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -407,6 +414,8 @@ def callable_func(ctx, item, idx, items): executor_context.create_child_context = lambda *args, **kwargs: Mock() class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -461,6 +470,8 @@ def callable_func(ctx, item, idx, items): executor_context.create_child_context = lambda *args, **kwargs: child_context class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -510,6 +521,8 @@ def mock_summary_generator(result): executor_context.create_child_context = Mock(return_value=_child_ctx) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -550,6 +563,8 @@ def callable_func(ctx, item, idx, items): executor_context.create_child_context = Mock(return_value=Mock()) # SLF001 class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -589,6 +604,8 @@ def func(ctx, item, index, array): config = MapConfig(summary_generator=None) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -634,6 +651,8 @@ def callable_func(ctx, item, idx, items): # Mock execution state that indicates operation already succeeded class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -703,6 +722,8 @@ def callable_func(ctx, item, idx, items): # Mock execution state that indicates operation succeeded but children need replay class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -802,6 +823,8 @@ def test_func(ctx, item, idx, items): execution_count = 0 class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1032,6 +1055,8 @@ def func(ctx, item, idx, items): return {"item": item.upper(), "index": idx} class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1481,6 +1506,8 @@ def test_map_handler_defaults_summary_generator_for_user_config(): """ class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1556,6 +1583,8 @@ def predicate(s: CompletionStatus) -> CompletionDecision: ) as mock_execute: class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py index 12231cfab..87d4d5a89 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py @@ -8,6 +8,7 @@ import pytest +from aws_durable_execution_sdk_python.plugin import PluginExecutor from aws_durable_execution_sdk_python.concurrency.executor import ConcurrentExecutor from aws_durable_execution_sdk_python.identifier import OperationIdNamespace @@ -211,6 +212,8 @@ def func2(ctx): config = ParallelConfig(max_concurrency=2) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -255,6 +258,8 @@ def func1(ctx): callables = [func1] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -299,6 +304,8 @@ def func1(ctx): config = ParallelConfig(max_concurrency=5) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -351,6 +358,8 @@ def func1(ctx): callables = [func1] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -466,6 +475,8 @@ def func1(ctx): callables = [func1] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -513,6 +524,8 @@ def mock_summary_generator(result): config = ParallelConfig(summary_generator=mock_summary_generator) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -556,6 +569,8 @@ def func2(ctx): callables = [func1, func2] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -606,6 +621,8 @@ def func3(ctx): config = ParallelConfig(summary_generator=None) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -654,6 +671,8 @@ def func2(ctx): # Mock execution state that indicates operation already succeeded class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -723,6 +742,8 @@ def func1(ctx): # Mock execution state that indicates operation succeeded but children need replay class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -824,6 +845,8 @@ def task2(ctx): execution_count = 0 class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1052,6 +1075,8 @@ def func3(ctx): callables = [func1, func2, func3] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1474,6 +1499,8 @@ def test_parallel_handler_defaults_summary_generator_for_user_config(): """A user config without a summary generator gets the default (JS parity).""" class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1551,6 +1578,8 @@ def predicate(s: CompletionStatus) -> CompletionDecision: ) as mock_execute: class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass diff --git a/pyproject.toml b/pyproject.toml index 5291b0c5e..80eb21e0c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -131,44 +131,48 @@ dependencies = [ [tool.hatch.envs.dev-examples.scripts] test = "pytest packages/aws-durable-execution-sdk-python-examples/test {args}" -[tool.hatch.envs.test-pypi-otel] -# Test new exclusivity capability against an installed capable core. -# Legacy valid-registration combinations are verified separately against 2.0.x. -# Override inherited workspace membership: core must come from an installed distribution. -workspace.members = ["packages/aws-durable-execution-sdk-python-otel"] +# Install immutable artifacts built from this checkout, including before the +# first compatible core is published. No editable workspace packages in this lane. +[tool.hatch.envs.test-wheel-otel] +detached = true +workspace.members = [] dependencies = [ "aws-durable-execution-sdk-python-testing>=1.2.1", - "aws-durable-execution-sdk-python>=2.1.0", "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", "pytest", - "pytest-cov", - "coverage[toml]", -] -pre-install-commands = [ - "pip install -e packages/aws-durable-execution-sdk-python-otel", + "packaging", ] -[tool.hatch.envs.test-pypi-otel.scripts] -test = "pytest packages/aws-durable-execution-sdk-python-otel/tests {args}" +[tool.hatch.envs.test-wheel-otel.scripts] +test = [ + "python .github/scripts/install_otel_test_wheels.py", + "pytest packages/aws-durable-execution-sdk-python-otel/tests {args}", + "pytest .github/tests/otel_lifecycle_compatibility_test.py {args}", +] -[tool.hatch.envs.test-pypi-otel-legacy] -template = "test-pypi-otel" -workspace.members = ["packages/aws-durable-execution-sdk-python-otel"] +# A core-only upgrade retains the released plugin's more limited Invocation +# fallback. Validate the real published plugin separately from the new pair. +[tool.hatch.envs.test-wheel-otel-legacy] +template = "test-wheel-otel" +detached = true +workspace.members = [] dependencies = [ - "aws-durable-execution-sdk-python>=2.0.0,<2.1.0", + "aws-durable-execution-sdk-python-otel==1.0.0", "aws-durable-execution-sdk-python-testing>=1.2.1", "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", "pytest", - "pytest-cov", - "coverage[toml]", + "packaging", ] -[tool.hatch.envs.test-pypi-otel-legacy.scripts] +[tool.hatch.envs.test-wheel-otel-legacy.env-vars] +OTEL_COMPAT_LEGACY = "1" + +[tool.hatch.envs.test-wheel-otel-legacy.scripts] test = [ - "python -c 'from pathlib import Path; from importlib.metadata import version; from packaging.version import Version; import aws_durable_execution_sdk_python.execution as core; assert Version(version(\"aws-durable-execution-sdk-python\")).release[:2] == (2, 0); assert \"site-packages\" in Path(core.__file__).resolve().parts; print(version(\"aws-durable-execution-sdk-python\"), core.__file__)' ", - "pytest packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_view_registration.py -k 'one_view_and_unrelated_plugin_suspend_resume or no_otel_plugin_remains_valid or execution_constructor_retains_ambient_log_correlation' {args}" + "python .github/scripts/install_otel_test_wheels.py --legacy-plugin", + "pytest .github/tests/otel_lifecycle_compatibility_test.py {args}", ] [tool.hatch.envs.test-pypi-examples]