fix(plugin): dispatch on one thread, keep the real failure

Two findings from a fresh review pass.

The failure exit fired the end hook inside its except block, so an
exception out of the hook left that block before the handler's own failure
was re-raised. The caller saw the plugin's exception and the real failure
survived only as __context__. The end-hook dispatch does raise: it re-raises
the control exceptions it holds through the fan-out. Instrumentation does
not decide what an execution failed with, so the hook's exception is now
contained and logged there and the original is re-raised unchanged.

execute_plugins also had a sync parameter no caller ever passed. Its
asynchronous branch submitted the hook to a per-invocation thread pool,
which skipped the end-hook fan-out rule and swallowed a control exception in
a Future nobody reads -- a way to opt out of the invariants this PR
establishes, on a path with no production caller. The parameter and the pool
are removed: hooks are dispatched on the calling thread, which is what lets
a plugin set thread-affine state the SDK's own logging reads, and the gate
that used to test the pool now tests the factory list. The pool created no
thread either way, because ThreadPoolExecutor spawns lazily on first submit.

Author: ParidelPooya

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

@@ -6,7 +6,6 @@
 import functools
 import logging
 from collections.abc import Iterator, Mapping, Sequence
-from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any, Callable, MutableMapping, Protocol, cast
@@ -603,7 +602,6 @@ def __init__(self, plugins: list[DurableInstrumentationPluginFactory] | None):
         # 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
         self._run_entered = False
@@ -625,25 +623,14 @@ class exists to prevent; failing loudly here keeps the bug from
             )
             raise RuntimeError(msg)
         self._run_entered = True
-        if self._plugin_factories:
-            self._executor = ThreadPoolExecutor(
-                max_workers=1,
-                thread_name_prefix="plugin-executor",
-            )
         try:
             yield
         finally:
             self._invocation_status = None
             self._operations_provider = None
-            # Shut down the thread pool, waiting for pending tasks to complete.
-            # The pool belongs to this invocation, so this drains only this
-            # invocation's queued dispatches and cannot cut short a concurrent
-            # invocation's.
-            if self._executor:
-                self._executor.shutdown(wait=True)
-            # Drop this invocation's plugin instances. After the pool has
-            # drained, so no queued dispatch still holds one: nothing outlives
-            # the invocation.
+            # Drop this invocation's plugin instances: nothing outlives the
+            # invocation. Every dispatch is synchronous, so there is no queued
+            # work still holding one.
             self._plugins = []
             self._started = []
 
@@ -739,7 +726,7 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None:
             if control is not None:
                 raise control from None
 
-    def execute_plugins(self, info, sync):
+    def execute_plugins(self, info):
         """Dispatch one hook to this invocation's plugins.
 
         A plugin receives a hook only once it has received the invocation-start
@@ -770,20 +757,25 @@ def execute_plugins(self, info, sync):
         ``BaseException`` rather than naming the three again. Naming them would
         miss a :class:`BaseExceptionGroup` carrying one, which is what
         :func:`_contain_plugin_failure` hands back.
+
+        Every hook is dispatched on the calling thread. That is what lets a
+        plugin set a ``ThreadLocal`` or an MDC key the SDK's own logging then
+        reads, and it is what makes the pairing and re-raise rules above
+        enforceable: a hook dispatched to a pool would land in a
+        :class:`~concurrent.futures.Future` nobody reads, so a control exception
+        raised there would be swallowed and the end-hook fan-out could not hold
+        it. An earlier ``sync`` parameter offered the pool path; no caller ever
+        passed it, and it is removed rather than left as a way to opt out of
+        those rules.
         """
-        if not self._executor:
+        if not self._plugin_factories:
             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 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
@@ -870,7 +862,7 @@ def on_invocation_start(
         # Build this invocation's plugin instances from the very info their first
         # hook receives, and before that hook is dispatched.
         self._create_plugins(self._invocation_status)
-        self.execute_plugins(self._invocation_status, sync=True)
+        self.execute_plugins(self._invocation_status)
 
     def _snapshot_execution_input(self, execution_input: Any) -> Any:
         """Deep-copy the execution input so the plugin view is isolated.
@@ -918,7 +910,7 @@ def on_invocation_end(
                 operations=self._snapshot_operation_infos(self._operations_provider),
             )
         )
-        self.execute_plugins(invocation_end_info, sync=True)
+        self.execute_plugins(invocation_end_info)
 
     def on_user_function_start(
         self,
@@ -939,7 +931,7 @@ def on_user_function_start(
             is_replay_children=is_replay_children,
             attempt=attempt,
         )
-        self.execute_plugins(start_info, sync=True)
+        self.execute_plugins(start_info)
         return start_info
 
     def on_user_function_end(
@@ -952,7 +944,6 @@ def on_user_function_end(
         """Execute plugins when a user function returns, fails, or is incomplete."""
         self.execute_plugins(
             UserFunctionEndInfo.from_start_info(start_info, error, outcome=outcome),
-            sync=True,
         )
 
     def on_operation_action(
@@ -982,7 +973,6 @@ def on_operation_action(
                     is_replayed=previous_operation is not None,
                     status=OperationStatus.STARTED,
                 ),
-                sync=True,
             )
 
     def on_operation_replay(self, operation: Operation) -> None:
@@ -1000,7 +990,7 @@ def on_operation_replay(self, operation: Operation) -> None:
             is_replayed=True,
             status=operation.status,
         )
-        self.execute_plugins(start_info, sync=True)
+        self.execute_plugins(start_info)
 
     def on_child_context_end(
         self,
@@ -1025,7 +1015,6 @@ def on_child_context_end(
                 error=error,
                 is_replayed=is_replayed,
             ),
-            sync=True,
         )
 
     def on_operation_update(
@@ -1075,7 +1064,6 @@ def on_operation_update(
                         ),
                         is_replayed=False,
                     ),
-                    sync=True,
                 )
 
         if (
@@ -1103,7 +1091,6 @@ def on_operation_update(
                 },
                 operations=_to_operation_info_map(operations),
             ),
-            sync=True,
         )
 
     @staticmethod
@@ -1204,13 +1191,30 @@ def wrapper(event: Any, context: LambdaContext):
                         # they also fire the hook. The hook is what a plugin needs
                         # to flush, and a process being torn down is when flushing
                         # matters; the cost is the same bounded work any
-                        # invocation end does. The exception itself is re-raised
-                        # unchanged, so what the caller sees is untouched.
-                        plugin_executor.on_invocation_end(
-                            output=DurableExecutionInvocationOutput.create_retry(
-                                ErrorObject.from_exception(e)
-                            ),
-                        )
+                        # invocation end does.
+                        #
+                        # The handler's exception is what the caller sees,
+                        # whatever the hook does. The end-hook dispatch can raise
+                        # -- it re-raises the control exceptions it holds through
+                        # the fan-out -- and letting that replace the handler's
+                        # failure would report an instrumentation problem as the
+                        # execution's outcome and leave the real failure reachable
+                        # only as __context__. Instrumentation does not decide
+                        # what an execution failed with, so the hook's exception
+                        # is contained here and the original is re-raised
+                        # unchanged.
+                        try:
+                            plugin_executor.on_invocation_end(
+                                output=DurableExecutionInvocationOutput.create_retry(
+                                    ErrorObject.from_exception(e)
+                                ),
+                            )
+                        except BaseException:  # noqa: BLE001 - the handler's failure wins
+                            logger.exception(
+                                "Plugin invocation-end hook failed while the "
+                                "invocation was already failing; the original "
+                                "failure is raised"
+                            )
                         raise
                     plugin_executor.on_invocation_end(output=completed)
                     return output

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

@@ -3,6 +3,7 @@
 import datetime
 import logging
 import pickle
+import threading
 import unittest
 from collections.abc import Iterator
 from copy import deepcopy
@@ -705,6 +706,38 @@ def test_a_factory_raising_a_group_without_a_control_exception_is_contained(self
         self.assertIn("boom", "\n".join(logs.output))
         self.assertEqual(surviving.calls, ["invocation_start:req-1"])
 
+    def test_a_failing_end_hook_does_not_replace_the_handlers_failure(self):
+        """Instrumentation does not decide what an execution failed with.
+
+        The end-hook dispatch can raise: it re-raises the control exceptions it
+        holds through the fan-out. On the failure exit that raise used to leave
+        the ``except`` block before the handler's own exception was re-raised, so
+        the caller saw the plugin's exception and the real failure survived only
+        as ``__context__``.
+        """
+        plugin = _ControlOnEndPlugin()
+        host = PluginHost(plugins=[plugin_factory(plugin)])
+
+        @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 ValueError("the real failure")
+
+        with self.assertLogs(
+            "aws_durable_execution_sdk_python.plugin", level=logging.ERROR
+        ):
+            with self.assertRaises(ValueError) as raised:
+                handler({}, LAMBDA_CTX)
+
+        self.assertEqual(str(raised.exception), "the real failure")
+        # The hook still ran, and still saw the failing outcome.
+        self.assertEqual(plugin.end_statuses, [InvocationStatus.RETRY])
+
     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.
 
@@ -1136,26 +1169,50 @@ def test_a_hook_raising_cancellation_is_contained(self):
             "aws_durable_execution_sdk_python.plugin", level=logging.ERROR
         ) as logs:
             with _invocation(executor, tracking):
-                executor.execute_plugins(OPERATION_START_INFO, sync=True)
+                executor.execute_plugins(OPERATION_START_INFO)
 
         self.assertIn("hook cancelled", "\n".join(logs.output))
         self.assertIn("operation_start:op-2", tracking.calls)
 
 
 class TestPluginExecutor(unittest.TestCase):
-    def test_no_thread_pool_when_plugins_is_none(self):
-        """Tests that PluginExecutor does not create a thread pool when plugins is empty."""
-        executor = PluginExecutor(plugins=None)
-        self.assertIsNone(executor._executor)
-
-    def test_no_thread_pool_when_plugins_is_empty_list(self):
-        executor = PluginExecutor(plugins=[])
-        self.assertIsNone(executor._executor)
+    def test_dispatch_is_a_no_op_when_no_factory_is_registered(self):
+        """An executor with no factories dispatches nothing and needs no thread.
+
+        Hooks are dispatched on the calling thread, which is what lets a plugin
+        set thread-affine state the SDK's own logging then reads, and what makes
+        the end-hook pairing and re-raise rules enforceable. There is therefore no
+        pool to create, and an executor with nothing registered returns before it
+        touches anything.
+        """
+        for plugins in (None, []):
+            with self.subTest(plugins=plugins):
+                executor = PluginExecutor(plugins=plugins)
+                self.assertEqual(executor._plugin_factories, [])
+                with executor.run():
+                    # Nothing registered, so no hook reaches a plugin and no
+                    # dispatch raises on the way through.
+                    executor.execute_plugins(INVOCATION_START_INFO)
+                    executor.execute_plugins(OPERATION_START_INFO)
+
+    def test_hooks_are_dispatched_on_the_calling_thread(self):
+        """Thread affinity is the contract, so it is asserted rather than assumed."""
+        seen: list[str] = []
+
+        class _ThreadRecordingPlugin(DurableInstrumentationPlugin):
+            def on_invocation_start(self, info: InvocationStartInfo) -> None:
+                seen.append(threading.current_thread().name)
 
-    def test_thread_pool_created_when_plugins_provided(self):
-        executor = PluginExecutor(plugins=[plugin_factory(_NoOpPlugin())])
+        executor = PluginExecutor(plugins=[plugin_factory(_ThreadRecordingPlugin())])
         with executor.run():
-            self.assertIsNotNone(executor._executor)
+            executor.on_invocation_start(
+                execution_arn="arn:exec",
+                lambda_context=LAMBDA_CTX,
+                execution_start_time=START_TS,
+                is_first_invocation=True,
+            )
+
+        self.assertEqual(seen, [threading.current_thread().name])
 
     def test_start_is_noop_when_empty(self):
         executor = PluginExecutor(plugins=[])
@@ -1234,37 +1291,37 @@ def setUp(self):
 
     def test_dispatch_invocation_start_info(self):
         with _invocation(self.executor, self.plugin):
-            self.executor.execute_plugins(INVOCATION_START_INFO, sync=True)
+            self.executor.execute_plugins(INVOCATION_START_INFO)
         self.assertIn("invocation_start:req-1", self.plugin.calls)
 
     def test_dispatch_invocation_end_info(self):
         with _invocation(self.executor, self.plugin):
-            self.executor.execute_plugins(INVOCATION_END_INFO, sync=True)
+            self.executor.execute_plugins(INVOCATION_END_INFO)
         self.assertIn("invocation_end:req-1", self.plugin.calls)
 
     def test_dispatch_operation_end_info(self):
         with _invocation(self.executor, self.plugin):
-            self.executor.execute_plugins(OPERATION_END_INFO, sync=False)
+            self.executor.execute_plugins(OPERATION_END_INFO)
         self.assertIn("operation_end:op-1", self.plugin.calls)
 
     def test_dispatch_operation_start_info(self):
         with _invocation(self.executor, self.plugin):
-            self.executor.execute_plugins(OPERATION_START_INFO, sync=False)
+            self.executor.execute_plugins(OPERATION_START_INFO)
         self.assertIn("operation_start:op-2", self.plugin.calls)
 
     def test_dispatch_operation_change_info(self):
         with _invocation(self.executor, self.plugin):
-            self.executor.execute_plugins(OPERATION_CHANGE_INFO, sync=False)
+            self.executor.execute_plugins(OPERATION_CHANGE_INFO)
         self.assertIn("operation_change:op-1", self.plugin.calls)
 
     def test_dispatch_user_function_start_info(self):
         with _invocation(self.executor, self.plugin):
-            self.executor.execute_plugins(USER_FUNCTION_START_INFO, sync=True)
+            self.executor.execute_plugins(USER_FUNCTION_START_INFO)
         self.assertIn("user_function_start:op-1", self.plugin.calls)
 
     def test_dispatch_user_function_end_info(self):
         with _invocation(self.executor, self.plugin):
-            self.executor.execute_plugins(USER_FUNCTION_END_INFO, sync=True)
+            self.executor.execute_plugins(USER_FUNCTION_END_INFO)
         self.assertIn("user_function_end:op-1", self.plugin.calls)
 
     def test_dispatch_unknown_type_logs_exception(self):
@@ -1273,7 +1330,7 @@ def test_dispatch_unknown_type_logs_exception(self):
             "aws_durable_execution_sdk_python.plugin", level=logging.ERROR
         ):
             with _invocation(self.executor, self.plugin):
-                self.executor.execute_plugins("not a valid info type", sync=True)
+                self.executor.execute_plugins("not a valid info type")
 
     def test_plugin_exception_is_swallowed(self):
         """If a plugin raises, the exception is logged and execution continues."""
@@ -1287,7 +1344,7 @@ def test_plugin_exception_is_swallowed(self):
             "aws_durable_execution_sdk_python.plugin", level=logging.ERROR
         ):
             with _invocation(executor, tracking_plugin):
-                executor.execute_plugins(OPERATION_START_INFO, sync=True)
+                executor.execute_plugins(OPERATION_START_INFO)
 
         # The second plugin should still have been called
         self.assertIn("operation_start:op-2", tracking_plugin.calls)
@@ -1298,7 +1355,7 @@ def test_multiple_plugins_all_called(self):
         executor = PluginExecutor(plugins=[plugin_factory(p1), plugin_factory(p2)])
 
         with _invocation(executor, p1, p2):
-            executor.execute_plugins(OPERATION_START_INFO, sync=True)
+            executor.execute_plugins(OPERATION_START_INFO)
 
         self.assertIn("operation_start:op-2", p1.calls)
         self.assertIn("operation_start:op-2", p2.calls)