diff --git a/cadence/_internal/context_propagation.py b/cadence/_internal/context_propagation.py new file mode 100644 index 0000000..05ac0a3 --- /dev/null +++ b/cadence/_internal/context_propagation.py @@ -0,0 +1,117 @@ +"""Shared Cadence header codecs and context-propagation scopes.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping, Sequence +from contextlib import ExitStack, contextmanager + +from cadence.api.v1.common_pb2 import Header +from cadence.context import ContextPropagator +from cadence.error import ContextPropagationError + + +def normalize_context_propagators( + propagators: Sequence[ContextPropagator] | None, +) -> tuple[ContextPropagator, ...]: + """Snapshot propagators so runtime behavior cannot change after setup.""" + return tuple(propagators or ()) + + +def header_to_context_fields(header: Header | None) -> dict[str, bytes]: + """Convert a Cadence Header to the cross-SDK ``str -> bytes`` carrier.""" + if header is None: + return {} + return {key: bytes(payload.data) for key, payload in header.fields.items()} + + +def inject_context_fields( + propagators: Sequence[ContextPropagator], +) -> dict[str, bytes]: + """Collect and validate fields emitted by configured propagators.""" + fields: dict[str, bytes] = {} + for propagator in propagators: + try: + emitted = propagator.inject() + except ContextPropagationError: + raise + except Exception as exc: + raise ContextPropagationError( + f"Context propagation inject failed for {type(propagator).__name__}" + ) from exc + if not isinstance(emitted, Mapping): + raise ContextPropagationError( + f"{type(propagator).__name__}.inject() must return a mapping" + ) + for key, value in emitted.items(): + if not isinstance(key, str): + raise ContextPropagationError( + f"{type(propagator).__name__}.inject() emitted a non-string key" + ) + if not isinstance(value, (bytes, bytearray, memoryview)): + raise ContextPropagationError( + f"{type(propagator).__name__}.inject() emitted a non-bytes value " + f"for header {key!r}" + ) + if key in fields: + raise ContextPropagationError( + f"Multiple context propagators emitted header {key!r}" + ) + fields[key] = bytes(value) + return fields + + +def context_header_from_propagators( + propagators: Sequence[ContextPropagator], +) -> Header | None: + """Build an independent protobuf Header for outbound propagation.""" + fields = inject_context_fields(propagators) + if not fields: + return None + header = Header() + for key, value in fields.items(): + header.fields[key].data = value + return header + + +@contextmanager +def context_propagation_scope( + propagators: Sequence[ContextPropagator], + header_or_fields: Header | Mapping[str, bytes] | None, +) -> Iterator[None]: + """Install incoming context for one logical workflow or activity invocation.""" + fields: Mapping[str, bytes] + if isinstance(header_or_fields, Header): + fields = header_to_context_fields(header_or_fields) + elif header_or_fields is None: + fields = {} + else: + fields = header_or_fields + + stack = ExitStack() + try: + for propagator in propagators: + try: + scope = propagator.extract(fields) + stack.enter_context(scope) + except ContextPropagationError: + raise + except Exception as exc: + raise ContextPropagationError( + f"Context propagation extract failed for {type(propagator).__name__}" + ) from exc + except Exception: + try: + stack.close() + except Exception as cleanup_error: + raise ContextPropagationError( + "Context propagation cleanup failed after extract error" + ) from cleanup_error + raise + + try: + yield + finally: + try: + stack.close() + except Exception as exc: + raise ContextPropagationError("Context propagation cleanup failed") from exc diff --git a/cadence/context.py b/cadence/context.py new file mode 100644 index 0000000..845af0a --- /dev/null +++ b/cadence/context.py @@ -0,0 +1,31 @@ +"""Context propagation interfaces for Cadence workflow boundaries. + +Propagators carry application-owned, request-scoped metadata through Cadence +headers. Implementations commonly use :mod:`contextvars` so values are scoped +to each client, workflow, or activity invocation. +""" + +from collections.abc import Mapping +from typing import ContextManager, Protocol + +from cadence.error import ContextPropagationError + +__all__ = ["ContextPropagationError", "ContextPropagator"] + + +class ContextPropagator(Protocol): + """Serialize and scope application context across Cadence boundaries. + + ``inject`` is called before the SDK writes a Cadence header. ``extract`` is + called for an incoming header and must return a context manager that restores + any ambient state when its scope exits. + + Workflow-side ``inject`` calls are replay-sensitive: implementations must + return identical bytes for identical persisted workflow state. + """ + + def inject(self) -> Mapping[str, bytes]: + """Return the current context as opaque header fields.""" + + def extract(self, headers: Mapping[str, bytes]) -> ContextManager[None]: + """Return a scope that installs context decoded from ``headers``.""" diff --git a/cadence/error.py b/cadence/error.py index c391e12..7307a7a 100644 --- a/cadence/error.py +++ b/cadence/error.py @@ -4,6 +4,10 @@ import grpc +class ContextPropagationError(RuntimeError): + """Raised when a context propagator cannot encode or install context.""" + + class ContinueAsNewError(Exception): def __init__( self, diff --git a/tests/cadence/test_context_propagation_primitives.py b/tests/cadence/test_context_propagation_primitives.py new file mode 100644 index 0000000..9f92f8d --- /dev/null +++ b/tests/cadence/test_context_propagation_primitives.py @@ -0,0 +1,159 @@ +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar +from typing import cast + +import pytest + +from cadence._internal.context_propagation import ( + context_header_from_propagators, + context_propagation_scope, + header_to_context_fields, + normalize_context_propagators, +) +from cadence.api.v1.common_pb2 import Header +from cadence.context import ContextPropagationError + + +class _ContextVarPropagator: + def __init__(self, key: str = "x-request-id") -> None: + self._key = key + self.value: ContextVar[str | None] = ContextVar(key, default=None) + + def inject(self) -> dict[str, bytes]: + value = self.value.get() + return {} if value is None else {self._key: value.encode()} + + @contextmanager + def extract(self, headers: Mapping[str, bytes]) -> Iterator[None]: + raw_value = headers.get(self._key) + token = self.value.set(raw_value.decode() if raw_value is not None else None) + try: + yield + finally: + self.value.reset(token) + + +class _StaticPropagator: + def __init__(self, fields: Mapping[str, bytes]) -> None: + self._fields = fields + + def inject(self) -> Mapping[str, bytes]: + return self._fields + + @contextmanager + def extract(self, _headers: Mapping[str, bytes]) -> Iterator[None]: + yield + + +class _BrokenPropagator: + def inject(self) -> Mapping[str, bytes]: + raise ValueError("cannot inject") + + @contextmanager + def extract(self, headers: Mapping[str, bytes]) -> Iterator[None]: + if headers: + yield + return + raise ValueError("cannot extract") + + +class _ContextPropagationErrorPropagator: + def inject(self) -> Mapping[str, bytes]: + raise ContextPropagationError("propagator-specific failure") + + @contextmanager + def extract(self, _headers: Mapping[str, bytes]) -> Iterator[None]: + yield + + +class _RecordingPropagator(_StaticPropagator): + def __init__(self, events: list[str], name: str) -> None: + super().__init__({}) + self._events = events + self._name = name + + @contextmanager + def extract(self, _headers: Mapping[str, bytes]) -> Iterator[None]: + self._events.append(f"enter:{self._name}") + try: + yield + finally: + self._events.append(f"exit:{self._name}") + + +class _CleanupBrokenPropagator(_StaticPropagator): + @contextmanager + def extract(self, _headers: Mapping[str, bytes]) -> Iterator[None]: + try: + yield + finally: + raise ValueError("cannot clean up") + + +def test_header_codec_preserves_raw_bytes_and_contextvar_scope() -> None: + propagator = _ContextVarPropagator() + header = Header() + header.fields["go.trace"].data = b"\x00\xffgo" + header.fields["x-request-id"].data = b"request-1" + header.fields["java.trace"].data = b"\x80java\x00" + + assert header_to_context_fields(header) == { + "go.trace": b"\x00\xffgo", + "x-request-id": b"request-1", + "java.trace": b"\x80java\x00", + } + with context_propagation_scope((propagator,), header): + assert propagator.value.get() == "request-1" + assert propagator.value.get() is None + + +def test_injection_rejects_collisions_and_invalid_values() -> None: + with pytest.raises(ContextPropagationError, match="Multiple context propagators"): + context_header_from_propagators( + ( + _StaticPropagator({"shared": b"one"}), + _StaticPropagator({"shared": b"two"}), + ) + ) + + invalid = _StaticPropagator(cast(Mapping[str, bytes], {"invalid": "value"})) + with pytest.raises(ContextPropagationError, match="non-bytes"): + context_header_from_propagators((invalid,)) + + with pytest.raises(ContextPropagationError, match="inject failed"): + context_header_from_propagators((_BrokenPropagator(),)) + + with pytest.raises( + ContextPropagationError, match="propagator-specific failure" + ) as error: + context_header_from_propagators((_ContextPropagationErrorPropagator(),)) + assert str(error.value) == "propagator-specific failure" + + +def test_scope_reverses_cleanup_and_wraps_propagator_errors() -> None: + events: list[str] = [] + with context_propagation_scope( + (_RecordingPropagator(events, "first"), _RecordingPropagator(events, "second")), + {}, + ): + assert events == ["enter:first", "enter:second"] + assert events == ["enter:first", "enter:second", "exit:second", "exit:first"] + + with pytest.raises(ContextPropagationError, match="extract failed"): + with context_propagation_scope((_BrokenPropagator(),), {}): + pass + with pytest.raises(ContextPropagationError, match="cleanup failed"): + with context_propagation_scope((_CleanupBrokenPropagator({}),), {}): + pass + + +def test_normalization_snapshots_configured_propagators() -> None: + propagator = _StaticPropagator({}) + configured = [propagator] + + normalized = normalize_context_propagators(configured) + configured.clear() + + assert normalized == (propagator,) + assert normalize_context_propagators(None) == ()