fix(plugin): pair the end hook with the start hook

A start hook raising one of the three exceptions that instruct the calling
thread to stop propagates out of the dispatch loop, so plugins later in the
list never receive their start hook. The invocation-end hook that the
propagating exception then triggers reached them anyway, leaving a plugin to
tear down state it had never been told to build.

Hooks after the start hook are now dispatched to the plugins that received
the start hook, which makes the pairing an invariant rather than a
coincidence. A plugin counts as started before its hook is dispatched,
because a hook that begins and then fails may already have allocated what
its end hook releases.

A refused drain also requests a flush now. An exporter that re-enters a
plugin hook can queue a record from the export worker, and drain() refuses
the wait that hook makes. The worker exits its loop once nothing is pending
and no flush is requested, so the record reached the exporters and the
worker stopped, leaving a buffering exporter holding an execution's terminal
telemetry when Lambda froze the environment. The request is made without
waiting, so the worker stays unblocked, and its barrier covers the record
just queued.

Author: ParidelPooya

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

@@ -201,8 +201,8 @@ def drain(self, execution: _ExportState) -> None:
         exporter re-enters a plugin hook and the hook reaches an invocation end.
         The call is refused and reported rather than deadlocking the invocation;
         the record stays queued and this same worker exports it once it resumes
-        its loop. (Mirrors the Java
-        ``ExportScheduler.refuseWaitThatWouldBlockThePump``.)
+        its loop, and a flush covering it is requested on the way out. (Mirrors
+        the Java ``ExportScheduler.refuseWaitThatWouldBlockThePump``.)
         """
         if self._is_export_worker():
             _logger.warning(
@@ -211,6 +211,14 @@ def drain(self, execution: _ExportState) -> None:
                 "refused rather than deadlocking the invocation; an exporter "
                 "re-entered a plugin hook"
             )
+            # The refused call still leaves a flush behind. The hook that
+            # re-entered may have queued a record, and this worker exits its loop
+            # once nothing is pending and no flush is requested -- so without the
+            # request the record would be handed to the exporters and the worker
+            # would stop, leaving a buffering exporter holding an execution's
+            # terminal telemetry when Lambda freezes the environment. Requesting
+            # it rather than waiting for it is what keeps the worker unblocked.
+            self._request_flush()
             return
         failed_pending: _Dropped | None = None
         start_error: Exception | None = None
@@ -286,6 +294,20 @@ def _is_export_worker(self) -> bool:
         with self._condition:
             return self._worker is threading.current_thread()
 
+    def _request_flush(self) -> None:
+        """Ask the worker for a flush covering everything scheduled so far.
+
+        Returns without waiting, so it is safe to call from the worker itself. The
+        barrier is raised to the current schedule counter, which is what makes the
+        flush cover a record queued moments ago rather than running before it.
+        """
+        with self._condition:
+            if self._disabled:
+                return
+            self._flush_requested = True
+            self._flush_barrier = max(self._flush_barrier, self._seq)
+            self._condition.notify_all()
+
     def _disable_locked(self) -> _Dropped:
         """Latch asynchronous export off for good and surrender everything queued.
 

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

@@ -903,9 +903,15 @@ def test_drain_from_the_export_worker_is_refused_rather_than_deadlocking(
         scheduler.schedule(ARN_A, _record("r1"))
 
         assert exporter.returned.wait(timeout=10), "the refused drain must return"
+        assert _wait_until(lambda: "refused rather than deadlocking" in caplog.text)
 
     assert ("export", "r1") in exporter.calls
-    assert "refused rather than deadlocking" in caplog.text
+    # The refused drain still leaves a flush behind, with no external drain to ask
+    # for one. The re-entering hook may have queued a record, and the worker exits
+    # once nothing is pending and no flush is requested, so without the request a
+    # buffering exporter would be holding that record when the environment froze.
+    assert _wait_until(lambda: ("flush", None) in exporter.calls), (
+        "a refused drain must request a flush on its way out"
+    )
     # The worker is still serving: a drain from any other thread completes.
     scheduler.drain(ARN_B)
-    assert ("flush", None) in exporter.calls

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py

@@ -569,6 +569,10 @@ def __init__(self, plugins: list[DurableInstrumentationPluginFactory] | None):
         # on_invocation_start and emptied when the invocation scope exits.
         self._plugin_factories = list(plugins or [])
         self._plugins: list[DurableInstrumentationPlugin] = []
+        # The subset of _plugins whose invocation-start hook has been dispatched.
+        # Every later hook is dispatched to this list, so a plugin that never
+        # received its start hook never receives its end hook.
+        self._started: list[DurableInstrumentationPlugin] = []
         self._executor: ThreadPoolExecutor | None = None
         self._invocation_status: InvocationStartInfo | None = None
         self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None
@@ -611,6 +615,7 @@ class exists to prevent; failing loudly here keeps the bug from
             # drained, so no queued dispatch still holds one: nothing outlives
             # the invocation.
             self._plugins = []
+            self._started = []
 
     def _create_plugins(self, info: InvocationStartInfo) -> None:
         """Build this invocation's plugin instances from its start info.
@@ -704,9 +709,27 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None:
             logger.exception("Plugin %s exception ignored", plugin.__class__.__name__)
 
     def execute_plugins(self, info, sync):
+        """Dispatch one hook to this invocation's plugins.
+
+        A plugin receives a hook only once it has received the invocation-start
+        hook, which makes the pairing an invariant rather than a coincidence.
+        Without it one dispatch order breaks the pairing: a start hook that raises
+        one of the three exceptions :data:`_PLUGIN_THREAD_CONTROL_EXCEPTIONS`
+        names propagates out of this loop, so plugins later in the list never
+        receive their start hook -- and the invocation-end hook that the
+        propagating exception then triggers used to reach them anyway, leaving a
+        plugin to tear down state it had never been told to build.
+
+        A plugin is counted as started before its start hook is dispatched rather
+        than after, because a hook that begins and then fails may already have
+        allocated what its end hook releases.
+        """
         if not self._executor:
             return
-        for plugin in self._plugins:
+        starting = isinstance(info, InvocationStartInfo)
+        for plugin in self._plugins if starting else self._started:
+            if starting:
+                self._started.append(plugin)
             if sync:
                 # this is called synchronously, so plugins will be able to manipulate thread local objects
                 self._dispatch_plugin(plugin, info)

packages/aws-durable-execution-sdk-python/tests/plugin_test.py

@@ -647,6 +647,35 @@ def create_plugin(self, info: InvocationStartInfo) -> _TrackingPlugin:
         self.assertEqual(built[0].calls, ["invocation_start:req-1"])
         self.assertEqual(built[1].calls, ["invocation_start:req-2"])
 
+    def test_a_start_hook_that_stops_the_thread_leaves_later_plugins_unpaired(self):
+        """A plugin that never received the start hook never receives the end hook.
+
+        A start hook raising one of the three control exceptions propagates out of
+        the dispatch loop, so plugins after it in the list never receive their
+        start hook. The invocation-end hook that the propagating exception then
+        triggers used to reach them anyway, leaving a plugin to tear down state it
+        had never been told to build.
+        """
+        later = _TrackingPlugin()
+        host = PluginHost(
+            plugins=[plugin_factory(_ControlOnStartPlugin()), plugin_factory(later)]
+        )
+
+        @host.handle_durable_output
+        def handler(event, context, plugin_executor):
+            plugin_executor.on_invocation_start(
+                execution_arn="arn:exec",
+                lambda_context=LAMBDA_CTX,
+                execution_start_time=START_TS,
+                is_first_invocation=True,
+            )
+            raise AssertionError("unreachable: the start hook stops the thread")
+
+        with self.assertRaises(KeyboardInterrupt):
+            handler({}, LAMBDA_CTX)
+
+        self.assertEqual(later.calls, [])
+
     def test_an_end_hook_that_stops_the_thread_is_not_reported_twice(self):
         """Exactly one end notification per invocation, whatever the hook does.
 
@@ -2335,6 +2364,13 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
         raise KeyboardInterrupt
 
 
+class _ControlOnStartPlugin(DurableInstrumentationPlugin):
+    """Plugin whose start hook stops the thread, so later plugins never start."""
+
+    def on_invocation_start(self, info: InvocationStartInfo) -> None:
+        raise KeyboardInterrupt
+
+
 class _FailingPlugin(DurableInstrumentationPlugin):
     """Plugin that raises on every hook call."""