fix(plugin): finish the end dispatch, flush only new work

Two follow-ups to the previous commit, both from its own review.

An end hook raising one of the three control exceptions aborted the dispatch
loop, so plugins after it lost the hook that is their only chance to finish:
Insight would not drain and OTel would leave spans unended. Every plugin the
loop reaches has already started, so the first such exception is now held
and re-raised once every plugin has been called. No other hook defers,
because stopping a start-hook loop early leaves later plugins with nothing
to clean up: the pairing rule withholds their end hook too.

The flush a refused drain requests is now conditional on a pending record.
An exporter whose flush() re-enters a plugin hook arrives at the refused
drain from inside a flush, and an unconditional request asked for the next
one, which re-entered and requested again for as long as the environment
lived. A pending record is what distinguishes the record the re-entering
hook queued from that loop.

Author: ParidelPooya

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

@@ -211,14 +211,20 @@ 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()
+            # The refused call still leaves a flush behind, but only when there is
+            # something for it to cover. The hook that re-entered may have queued a
+            # record, and this worker exits its loop once nothing is pending and no
+            # flush is requested -- so without the request the record would be
+            # handed to the exporters and the worker would stop, leaving a
+            # buffering exporter holding an execution's terminal telemetry when
+            # Lambda freezes the environment.
+            #
+            # Requesting one unconditionally would livelock instead: an exporter
+            # whose flush() re-enters a hook arrives here from inside a flush, and
+            # an unconditional request would ask for the next one, which re-enters
+            # again, for as long as the environment lives. A pending record is what
+            # distinguishes new work from that loop.
+            self._request_flush_for_pending_records()
             return
         failed_pending: _Dropped | None = None
         start_error: Exception | None = None
@@ -294,15 +300,23 @@ 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.
+    def _request_flush_for_pending_records(self) -> None:
+        """Ask the worker for a flush, but only if a record is waiting for one.
 
         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.
+
+        A pending record is the condition, not a formality. This is called from a
+        drain refused on the worker thread, and one way to reach that is an
+        exporter whose ``flush()`` re-enters a plugin hook: the call then arrives
+        from inside a flush, and requesting the next one unconditionally would
+        produce a flush that re-enters, requests, and flushes again for as long as
+        the environment lives -- after the invocation has returned. Nothing is
+        pending in that case, so nothing is requested.
         """
         with self._condition:
-            if self._disabled:
+            if self._disabled or not self._pending:
                 return
             self._flush_requested = True
             self._flush_barrier = max(self._flush_barrier, self._seq)

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

@@ -866,22 +866,27 @@ def drain(arn: str) -> None:
 
 
 class ReentrantDrainExporter(CaptureExporter):
-    """Exporter that drains from inside export(), as a plugin hook would.
+    """Exporter that queues a record and drains from inside export().
 
-    An exporter that re-enters a plugin hook reaches ``on_invocation_end``, and
-    that hook drains. The re-entry therefore arrives on the export worker thread,
-    which is the one thread able to serve the wait.
+    This is what an exporter re-entering a plugin hook produces: the hook emits
+    its record and then, at an invocation end, drains. Both arrive on the export
+    worker, which is the one thread able to serve the wait.
     """
 
     def __init__(self) -> None:
         super().__init__()
         self.scheduler: _ArnScheduler | None = None
         self.returned = threading.Event()
+        self._reentered = False
 
     def export(self, record: dict[str, Any]) -> None:
         super().export(record)
         assert self.scheduler is not None
-        self.scheduler.drain(ARN_A)
+        if self._reentered:
+            return
+        self._reentered = True
+        self.scheduler.schedule(ARN_B, _record("r2"))
+        self.scheduler.drain(ARN_B)
         self.returned.set()
 
 
@@ -907,11 +912,59 @@ def test_drain_from_the_export_worker_is_refused_rather_than_deadlocking(
 
     assert ("export", "r1") in exporter.calls
     # The refused drain still leaves a flush behind, with no external drain to ask
-    # for one. The re-entering hook may have queued a record, and the worker exits
-    # once nothing is pending and no flush is requested, so without the request a
+    # for one. The re-entering hook queued a record, and the worker exits once
+    # nothing is pending and no flush is requested, so without the request a
     # buffering exporter would be holding that record when the environment froze.
-    assert _wait_until(lambda: ("flush", None) in exporter.calls), (
-        "a refused drain must request a flush on its way out"
-    )
+    assert _wait_until(lambda: ("export", "r2") in exporter.calls)
+    assert _wait_until(
+        lambda: exporter.calls.index(("flush", None))
+        > exporter.calls.index(("export", "r2"))
+    ), "a refused drain must request a flush that covers the record it queued"
     # The worker is still serving: a drain from any other thread completes.
     scheduler.drain(ARN_B)
+
+
+class ReentrantFlushExporter(CaptureExporter):
+    """Exporter whose flush() drains, as a hook re-entered from a flush would.
+
+    An exporter that re-enters a plugin hook from ``flush()`` reaches
+    ``on_invocation_end``, which drains. The drain arrives on the export worker
+    from inside a flush, with nothing pending.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.scheduler: _ArnScheduler | None = None
+        self.flushes = 0
+
+    def flush(self) -> None:
+        super().flush()
+        self.flushes += 1
+        assert self.scheduler is not None
+        self.scheduler.drain(ARN_A)
+
+
+def test_a_drain_refused_from_inside_a_flush_does_not_re_arm_it() -> None:
+    """A refused drain with nothing pending asks for no further flush.
+
+    Requesting one unconditionally would keep the worker flushing for as long as
+    the environment lived: the flush re-enters the hook, the hook drains, the
+    refused drain asks for the next flush. A pending record is what distinguishes
+    new work from that loop, so a drain refused from inside a flush leaves no
+    request behind.
+    """
+    exporter = ReentrantFlushExporter()
+    scheduler = _ArnScheduler([exporter])
+    exporter.scheduler = scheduler
+
+    scheduler.schedule(ARN_A, _record("r1"))
+    scheduler.drain(ARN_A)
+
+    flushes_after_drain = exporter.flushes
+    assert flushes_after_drain >= 1, "the drain must have flushed"
+
+    # Give a re-armed flush time to appear. The worker retires when nothing is
+    # pending and no flush is requested, so a bounded count here is the whole
+    # assertion: an unconditional request never settles.
+    assert not _wait_until(lambda: exporter.flushes > flushes_after_drain, timeout=1.0)
+    assert _wait_until(lambda: not scheduler._worker_alive())

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

@@ -723,19 +723,40 @@ def execute_plugins(self, info, sync):
         A plugin is counted as started before its start hook is dispatched rather
         than after, because a hook that begins and then fails may already have
         allocated what its end hook releases.
+
+        The invocation-end hook is the one hook that finishes dispatching even
+        when a plugin raises one of those three. Every plugin it reaches has
+        already started, so cutting the loop short costs a plugin its only chance
+        to finish: Insight would not drain, and OTel would leave spans unended.
+        The first such exception is held and re-raised once every plugin has been
+        called, so the thread still stops and nothing is swallowed. No other hook
+        defers: stopping a start-hook loop early leaves later plugins with nothing
+        to clean up, because the pairing rule above then withholds their end hook
+        too.
         """
         if not self._executor:
             return
         starting = isinstance(info, InvocationStartInfo)
+        ending = isinstance(info, InvocationEndInfo)
+        deferred_control: BaseException | None = None
         for plugin in self._plugins if starting else self._started:
             if starting:
                 self._started.append(plugin)
-            if sync:
-                # this is called synchronously, so plugins will be able to manipulate thread local objects
-                self._dispatch_plugin(plugin, info)
-            else:
+            if not sync:
                 # this is called asynchronously, so plugins cannot manipulate thread local objects
                 self._executor.submit(self._dispatch_plugin, plugin, info)
+                continue
+            # this is called synchronously, so plugins will be able to manipulate thread local objects
+            if not ending:
+                self._dispatch_plugin(plugin, info)
+                continue
+            try:
+                self._dispatch_plugin(plugin, info)
+            except _PLUGIN_THREAD_CONTROL_EXCEPTIONS as control:
+                if deferred_control is None:
+                    deferred_control = control
+        if deferred_control is not None:
+            raise deferred_control
 
     def _snapshot_operation_infos(
         self,

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

@@ -647,6 +647,40 @@ 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_an_end_hook_that_stops_the_thread_still_finishes_the_dispatch(self):
+        """Every started plugin receives the end hook, then the thread stops.
+
+        The end hook is a plugin's only chance to finish -- Insight drains there
+        and OTel ends its spans -- so a plugin raising one of the three control
+        exceptions must not cost the plugins after it in the list their own end
+        hook. The exception is held and re-raised once every plugin has been
+        called.
+        """
+        later = _TrackingPlugin()
+        host = PluginHost(
+            plugins=[plugin_factory(_ControlOnEndPlugin()), plugin_factory(later)]
+        )
+
+        @host.handle_durable_output
+        def handler(event, context, plugin_executor):
+            plugin_executor.on_invocation_start(
+                execution_arn="arn:exec",
+                lambda_context=LAMBDA_CTX,
+                execution_start_time=START_TS,
+                is_first_invocation=True,
+            )
+            return {
+                "Status": ServiceInvocationStatus.SUCCEEDED.value,
+                "Result": None,
+            }
+
+        with self.assertRaises(KeyboardInterrupt):
+            handler({}, LAMBDA_CTX)
+
+        self.assertEqual(
+            later.calls, ["invocation_start:req-1", "invocation_end:req-1"]
+        )
+
     def test_a_start_hook_that_stops_the_thread_leaves_later_plugins_unpaired(self):
         """A plugin that never received the start hook never receives the end hook.