fix(plugin): give each invocation its own plugin session and log correlation

Review findings on the per-invocation plugin contract. The first one meant the
contract change did not actually deliver what it promised under concurrency.

`durable_execution` built one `PluginExecutor` at decoration time and every
invocation shared it, while the executor stored that invocation's plugin
instances, invocation metadata and operations provider on itself. So the SDK
created one plugin per invocation and then parked it in a shared mutable slot.
Measured with two invocations held inside their user function at once,
invocation A's operation hook was delivered to B's plugin and A's plugin never
received its own operation or end hook; in a second scenario A's scope exit
cleared the shared state while B was live and B lost `on_invocation_end`
entirely.

A new handler-lifetime `PluginHost` now holds nothing but the resolved factory
list and hands out a fresh `PluginExecutor` per invocation, which is the shape
JS and Java already had. The executor is reachable only from the frames of the
invocation that owns it, and a second `run()` on one executor raises, so reuse
fails loudly instead of as silent crosstalk. Its thread pool is per invocation
too: sharing one would let an ending invocation shut down a pool a concurrent
invocation was still submitting to.

The OTel log filter had the same class of bug. It held one mutable plugin
reference, so with two invocations open the one that started last won and every
other invocation's records carried its trace. Reading the active span from the
OTel context, as the Java plugin does, is not sufficient here: the invocation
span is never attached to the context, and the SDK runs the handler body on a
pool thread that does not inherit the context of the thread the plugin bound on.
The filter is now stateless and resolves per record, preferring the invocation
that claimed the emitting thread through a `ContextVar` and falling back to the
single open invocation when exactly one is open. With several open and an
unclaimed thread it leaves the record unstamped, because an unattributed record
is a smaller defect than one attributed to another customer's execution.
Installation is serialized so concurrent first-time callers cannot stack
filters.

Third, the Insight scheduler used `0` for both "no flush in flight" and "a flush
in flight covering zero exports", which are different facts. An invocation that
emitted no record could not tell that its flush was already running, so it
requested a second one that then ran after both invocations had returned —
customer exporter code past the invocation boundary, which the flush contract
forbids. Absence is now `None` and coverage an integer.

Author: ParidelPooya

packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py

@@ -101,11 +101,17 @@ def __init__(self, exporters: list[InsightExporter]) -> None:
         self._flushes_completed = 0
         self._flush_requested = False
         # Export counter coverage of the flush the worker is running right now, or
-        # 0 when no flush is in flight. Published when the worker commits to a
-        # flush, so a waiter woken while that flush runs -- before its coverage
-        # reaches _flushed_through -- can tell it is already covered instead of
-        # requesting a second flush that would run after its drain returned.
-        self._flush_in_flight = 0
+        # None when no flush is in flight. Presence and coverage are separate
+        # facts: a flush that covers zero exports is an ordinary flush -- it is
+        # what an invocation that emitted no record asks for -- and a single
+        # integer cannot say both "no flush is running" and "a flush covering
+        # nothing is running". Encoding the first as 0 made those two states
+        # identical, so a waiter needing zero coverage could not tell that its
+        # flush was already running and requested a second one that then ran
+        # after its invocation had returned. Published when the worker commits to
+        # a flush, so a waiter woken while that flush runs -- before its coverage
+        # reaches _flushed_through -- can tell it is already covered.
+        self._flush_in_flight: int | None = None
         # Value of the global schedule counter (_seq) when a flush was requested.
         # The worker defers the flush until no record scheduled at or before that
         # point is still pending. That is deliberately wider than the requester's
@@ -202,15 +208,15 @@ def drain(self, execution: _ExportState) -> None:
                     # made has already been consumed -- runs an extra flush after
                     # this drain, and the invocation, returned.
                     #
-                    # `_flush_in_flight` uses 0 as its "no flush is running"
-                    # sentinel, so the naive `self._flush_in_flight >= need`
-                    # reads as "already covered" when `need` is 0 -- precisely
-                    # when nothing is running at all. `need` is 0 for a drain
-                    # whose invocation emitted no record, so that form would let
-                    # such a drain skip its request and park until some other
-                    # execution happened to flush. Require a marker that is
-                    # actually set AND that reaches `need`.
-                    covered = 0 < self._flush_in_flight >= need
+                    # `_flush_in_flight` is None exactly while no flush is
+                    # running, so a flush that covers zero exports is still a
+                    # flush in flight. That case is the common one, not a corner:
+                    # `need` is 0 for a drain whose invocation emitted no record,
+                    # and the flush it asks for covers 0 exports when nothing has
+                    # ever been exported. Two such drains at once both see the
+                    # other's flush and neither asks for a second.
+                    in_flight = self._flush_in_flight
+                    covered = in_flight is not None and in_flight >= need
                     if not self._flush_requested and not covered:
                         self._flush_requested = True
                         self._flush_barrier = max(self._flush_barrier, self._seq)
@@ -260,7 +266,7 @@ def _ensure_worker_locked(self) -> tuple[_Dropped | None, Exception | None]:
             self._pending = {}
             self._flush_requested = False
             self._flush_barrier = 0
-            self._flush_in_flight = 0
+            self._flush_in_flight = None
             # Release every waiter; the permanent disable latch means no record
             # will ever be exported.
             self._condition.notify_all()
@@ -304,7 +310,10 @@ def _run_loop(self) -> None:
                         flush_covers = self._export_count
                         # Publish what this flush will cover before releasing the
                         # lock, so a waiter that wakes while it runs can see that
-                        # this flush releases it and skip asking for another.
+                        # this flush releases it and skip asking for another. A
+                        # coverage of 0 is published like any other: it means this
+                        # flush covers every export so far, of which there are
+                        # none, and it is still a flush in flight.
                         self._flush_in_flight = flush_covers
                         break
                     if self._pending:
@@ -353,7 +362,7 @@ def _run_loop(self) -> None:
                     # Retire the marker whatever happened: a stale one would park
                     # every later waiter that trusted this flush to cover it. Only
                     # a flush that ran to completion publishes its coverage.
-                    self._flush_in_flight = 0
+                    self._flush_in_flight = None
                     if flushed:
                         self._flushes_completed += 1
                         if flush_covers > self._flushed_through:

packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py

@@ -430,6 +430,92 @@ def drain() -> None:
     assert capture.calls == [("flush", None)]
 
 
+class GatedFlushExporter(CaptureExporter):
+    """Holds each ``flush()`` open until the test releases it, and counts them."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._lock = threading.Lock()
+        self.flush_count = 0
+        self.started: dict[int, threading.Event] = {}
+        self.release: dict[int, threading.Event] = {}
+        for index in (1, 2):
+            self.started[index] = threading.Event()
+            self.release[index] = threading.Event()
+
+    def flush(self) -> None:
+        with self._lock:
+            self.flush_count += 1
+            index = self.flush_count
+        if index in self.started:
+            self.started[index].set()
+            self.release[index].wait(10.0)
+        super().flush()
+
+
+def test_concurrent_drains_with_nothing_to_export_share_one_flush() -> None:
+    # Two invocations that emitted no record drain at the same time. Neither has
+    # an export to be covered, so both need a flush that covers zero exports.
+    # While `_flush_in_flight` used 0 for "no flush is running", the second drain
+    # could not tell that the flush it needs was already running, and asked for
+    # another. The completion of the first flush then released both drains, and
+    # the second flush ran after both invocations had already returned -- customer
+    # exporter code running past the invocation boundary, which is what the flush
+    # contract forbids. One flush must serve both, and whichever drain a flush
+    # belongs to must stay parked until that flush completes.
+    exporter = GatedFlushExporter()
+    scheduler = _ArnScheduler([exporter])
+    returned: list[str] = []
+    returned_lock = threading.Lock()
+
+    def drain(name: str, execution_arn: str) -> None:
+        scheduler.drain(execution_arn)
+        with returned_lock:
+            returned.append(name)
+
+    threads = [
+        threading.Thread(target=drain, args=("first", ARN_A), daemon=True),
+        threading.Thread(target=drain, args=("second", ARN_B), daemon=True),
+    ]
+    try:
+        threads[0].start()
+        assert exporter.started[1].wait(10.0), "the first drain never flushed"
+        threads[1].start()
+
+        first = scheduler.executions[ARN_A]
+
+        def both_parked() -> bool:
+            second = scheduler.executions.get(ARN_B)
+            if second is None:
+                return False
+            with scheduler._condition:
+                return first.waiters == 1 and second.waiters == 1
+
+        assert _wait_until(both_parked), "a drain raced past the flush it needs"
+        with returned_lock:
+            assert returned == [], "a drain returned before its flush completed"
+
+        exporter.release[1].set()
+        for thread in threads:
+            thread.join(10.0)
+        assert not any(thread.is_alive() for thread in threads)
+        with returned_lock:
+            assert sorted(returned) == ["first", "second"]
+
+        # The redundant request, if one was made, was recorded before either
+        # drain returned, so the worker starts that flush without further
+        # prompting. Nothing arriving here is what proves no second flush was
+        # requested.
+        assert not exporter.started[2].wait(0.75), (
+            "a second flush ran after both invocations had returned"
+        )
+        assert exporter.flush_count == 1
+    finally:
+        exporter.release[1].set()
+        exporter.release[2].set()
+        _wait_until(lambda: not scheduler._worker_alive())
+
+
 def test_drain_never_rides_on_a_flush_that_finished_before_it_started() -> None:
     # Export coverage alone would let the second drain return immediately: every
     # export is already covered by the first drain's flush. A drain must wait for
@@ -475,14 +561,15 @@ def fail_start(self) -> None:  # noqa: ARG001
             for execution in scheduler.executions.values()
         )
         assert scheduler._flush_requested is False
-        assert scheduler._flush_in_flight == 0
+        assert scheduler._flush_in_flight is None
 
 
 def test_disabled_latch_clears_a_published_flush_in_flight_marker(monkeypatch) -> None:
     # `_flush_in_flight` is the coverage of the flush the worker is running right
     # now, published so a waiter woken during that flush can tell it is already
-    # covered and skip requesting another. 0 means "no flush is running", so the
-    # marker is a claim that a flush is in flight and will complete.
+    # covered and skip requesting another. None means "no flush is running", so
+    # any integer -- 0 included -- is a claim that a flush is in flight and will
+    # complete.
     #
     # The _disabled latch makes that claim permanently false: no worker exists and
     # none will ever be started again, so the published flush can never complete.
@@ -503,7 +590,7 @@ def fail_start(self) -> None:  # noqa: ARG001
 
     with scheduler._condition:
         assert scheduler._disabled
-        assert scheduler._flush_in_flight == 0, (
+        assert scheduler._flush_in_flight is None, (
             "the _disabled latch left a flush-in-flight marker behind for a flush "
             "that can never run"
         )

packages/aws-durable-execution-sdk-python-otel/README.md

@@ -430,7 +430,10 @@ Structured trace context and sampling decision returned by context extractors.
 
 The logging filter (and its installer) used to stamp trace context onto log
 records. Installed automatically when `enrich_logger=True`; exported for manual
-setups.
+setups, where `install_log_filter(target_logger)` attaches it to a logger of your
+choice. The filter carries no invocation identity of its own: it resolves the
+invocation a record belongs to at emit time, so one filter serves every
+invocation the environment runs, including concurrent ones.
 
 ## Requirements
 

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py

@@ -90,7 +90,11 @@
     canonical_trace_id,
 )
 from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig
-from aws_durable_execution_sdk_python_otel.log_filter import install_log_filter
+from aws_durable_execution_sdk_python_otel.log_filter import (
+    bind_invocation,
+    install_log_filter,
+    unbind_invocation,
+)
 from aws_durable_execution_sdk_python_otel.provider import create_tracer_provider
 
 
@@ -177,9 +181,12 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None:
         self._tracing_enabled = False
 
         if self._config.enrich_logger:
-            # Install (or, on a warm environment, rebind) the root-logger filter
-            # so every log record is stamped with this invocation's span context.
-            install_log_filter(self)
+            # Install the root-logger filter so every log record is stamped with
+            # the active span context. On a warm environment the handler already
+            # carries the filter a previous invocation installed and it is reused
+            # as is: the filter holds no invocation identity, and this plugin
+            # claims the invocation in on_invocation_start instead.
+            install_log_filter()
 
     def _bind_sdk_tracer(self) -> bool:
         """Bind to an SDK tracer, retrying a deferred global provider."""
@@ -426,6 +433,12 @@ def _with_sampling(self, parent_context: Context) -> Context:
     # ------------------------------------------------------------------
     def on_invocation_start(self, info: InvocationStartInfo) -> None:
         logger.debug("Durable invocation started: %s", info)
+        # Claim log correlation for this invocation before anything can fail
+        # below: the claim is what keeps a concurrent invocation's records off
+        # this invocation's trace, and it is registered even when tracing turns
+        # out to be disabled so that the filter can still tell how many
+        # invocations are open.
+        bind_invocation(self)
         if info.execution_start_time is None:
             logger.warning(
                 "ExecutionOtelPlugin requires InvocationStartInfo.execution_start_time "
@@ -660,12 +673,17 @@ def _release_invocation_scope(self) -> None:
           plugin, so any scope this plugin attached and did not release must be
           detached here or it would stay current on a warm environment's thread
           after the invocation returns.
+        * The log filter's record of open invocations is process-global, so this
+          invocation must be removed from it. Until it is, a record emitted on an
+          unclaimed thread could still be correlated to this finished
+          invocation's spans.
         * ``_tracing_enabled`` is cleared so a hook that arrives after the
           invocation end -- one dispatched off the checkpointing path, for
           instance -- cannot start a span after the invocation span was ended and
           the provider flushed.
         """
         self._detach_remaining_contexts()
+        unbind_invocation(self)
         self._tracing_enabled = False
 
     # ------------------------------------------------------------------
@@ -675,6 +693,14 @@ def on_operation_start(self, info: OperationStartInfo) -> None:
         logger.debug("Durable operation started: %s", info)
         if not self._tracing_enabled:
             return
+        # Runs on the thread that drives the durable operation, which is the
+        # thread running the handler body and not the thread the
+        # invocation-start hook claimed. Claim it too, so records emitted from
+        # top-level handler code are correlated to this invocation even while
+        # another invocation is open in the same process. Claimed after the
+        # tracing-enabled gate, so a hook arriving after the invocation ended
+        # cannot re-register a finished invocation.
+        bind_invocation(self)
         if info.operation_type is OperationType.CONTEXT:
             with self._lock:
                 self._checkpointed_context_ids.add(info.operation_id)
@@ -807,6 +833,9 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
         logger.debug("Durable user function started: %s", info)
         if not self._tracing_enabled:
             return
+        # Runs on the thread executing user code -- a parallel branch runs on its
+        # own thread -- so claim that thread for this invocation as well.
+        bind_invocation(self)
         if info.operation_type not in (OperationType.CONTEXT, OperationType.STEP):
             raise RuntimeError(
                 "on_user_function_start only supports CONTEXT and STEP operations"