fix: concurrency drain timer resumes outside lock

A map or parallel stays in the current invocation while at least one
branch is still running. When one iteration keeps it alive and the
other branches wait on short timers (a wait, a wait_for_condition poll,
or a step retry backoff), those branches resume in-process as their
timers come due.

Before this commit, a single timer thread runs those resumes one at a
time and holds its lock across a blocking checkpoint that refreshes
state. Every resume costs one network round trip under the lock, and
every branch trying to register its next wait queues behind it. When
many timers come due together the timer thread falls behind, the
invocation reaches its function timeout, and the backend reinvokes, so
a map that should finish in seconds runs for minutes across several
timed-out invocations. Holding the lock across the submit also allowed
a latent self-deadlock, where a branch that finished inline reacquired
the same lock on the timer thread through its done-callback.

This commit holds the lock only long enough to take all due timers off
the queue and mark them pending, then releases it and runs one shared
refresh for the whole wave before handing the branches back to the
worker pool. One round trip now serves the whole wave instead of one
per resume, and new waits no longer queue behind a network call. The
take and the mark stay atomic so a branch never looks parked while it
is about to resume, which would otherwise suspend the whole operation
by mistake. If the refresh fails, which happens only when the
checkpoint subsystem has already failed and is terminal, the timer
thread records that one error and re-raises it from execute() so the
platform retries from the last checkpoint.

Closes #473

Author: yaythomas

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

@@ -59,7 +59,7 @@ class TimerScheduler:
     """Manage timed suspend tasks with a background timer thread."""
 
     def __init__(
-        self, resubmit_callback: Callable[[ExecutableWithState], None]
+        self, resubmit_callback: Callable[[list[ExecutableWithState]], None]
     ) -> None:
         self.resubmit_callback = resubmit_callback
         self._pending_resumes: list[tuple[float, int, ExecutableWithState]] = []
@@ -114,18 +114,31 @@ def _timer_loop(self) -> None:
 
             current_time = time.time()
             if current_time >= next_resume_time:
-                # Time to resume
+                # Drain every due resume under the lock, transitioning each to
+                # PENDING atomically with the pop. Keeping pop+reset_to_pending
+                # together is required: should_execution_suspend reads branch
+                # status without this lock, so an item that is removed from the
+                # heap but still SUSPENDED_WITH_TIMEOUT could trigger a spurious
+                # parent suspend.
+                ready: list[ExecutableWithState] = []
                 with self._lock:
-                    # no branch cover because hard to test reliably - this is a double-safety check if heap mutated
-                    # since the first peek on next_resume_time further up
-                    if (  # pragma: no branch
+                    while (
                         self._pending_resumes
                         and self._pending_resumes[0][0] <= current_time
                     ):
                         _, _, exe_state = heapq.heappop(self._pending_resumes)
                         if exe_state.can_resume:
                             exe_state.reset_to_pending()
-                            self.resubmit_callback(exe_state)
+                            ready.append(exe_state)
+                # Resubmit outside the lock. Only the heap pop and the PENDING
+                # transition need the lock. The checkpoint refresh is a blocking
+                # network call and the submit hands work to the pool, so running
+                # them off the lock keeps timed resumes from serializing behind
+                # the network round trip and keeps the timer thread from
+                # re-entering this non-reentrant lock when a submitted future
+                # completes inline and its done-callback calls schedule_resume.
+                if ready:
+                    self.resubmit_callback(ready)
             else:
                 # Wait until next resume time
                 wait_time = min(next_resume_time - current_time, 0.1)
@@ -169,6 +182,7 @@ def __init__(
         # Event-driven state tracking for when the executor is done
         self._completion_event = threading.Event()
         self._suspend_exception: SuspendExecution | None = None
+        self._resume_error: Exception | None = None
 
         # ExecutionCounters will keep track of completion criteria and on-going counters
         min_successful = self.completion_config.min_successful or len(self.executables)
@@ -222,11 +236,32 @@ def execute(
         ]
         self._completion_event.clear()
         self._suspend_exception = None
-
-        def resubmitter(executable_with_state: ExecutableWithState) -> None:
-            """Resubmit a timed suspended task."""
-            execution_state.create_checkpoint()
-            submit_task(executable_with_state)
+        self._resume_error = None
+
+        def resubmitter(ready: list[ExecutableWithState]) -> None:
+            """Resubmit a wave of timed-suspended tasks.
+
+            One checkpoint refresh serves the whole due wave: the fetch returns
+            all operations, so every resumed branch reads fresh state. The
+            refresh only raises when the background checkpoint subsystem has
+            failed, which is terminal for the whole execution, so record the
+            error and wake the parent to re-raise it. Catching here keeps the
+            single timer thread alive so a failure does not strand the other
+            pending resumes.
+            """
+            try:
+                execution_state.create_checkpoint()
+            except Exception as exc:  # noqa: BLE001
+                # resubmitter runs only on the single timer thread, so this
+                # check-then-set needs no lock. First error wins: keep the
+                # earliest failure if several waves fail before execute() reads
+                # it (they are the same terminal checkpoint failure anyway).
+                if self._resume_error is None:  # pragma: no branch
+                    self._resume_error = exc
+                self._completion_event.set()
+                return
+            for executable_with_state in ready:
+                submit_task(executable_with_state)
 
         thread_executor = ThreadPoolExecutor(max_workers=max_workers)
         try:
@@ -259,6 +294,12 @@ def on_done(future: Future) -> None:
                 for future in futures:
                     future.cancel()
 
+                # A timed resume failed to refresh state (terminal checkpoint
+                # subsystem failure). Re-raise so the invocation fails and the
+                # backend retries from the last durable checkpoint.
+                if self._resume_error is not None:
+                    raise self._resume_error
+
                 # Suspend execution if everything done and at least one of the tasks raised a suspend exception.
                 if self._suspend_exception:
                     raise self._suspend_exception

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

@@ -1370,6 +1370,74 @@ def execute_item(self, child_context, executable):
         executor.execute(execution_state, executor_context)
 
 
+def test_concurrent_executor_resume_checkpoint_failure_propagates():
+    """A resume-time checkpoint refresh failure propagates out of execute().
+
+    Regression guard: the timer resubmit does a blocking checkpoint refresh.
+    That refresh only raises when the checkpoint subsystem has failed, which
+    is terminal. execute() must re-raise it (so the invocation fails and the
+    backend retries from the last durable checkpoint) rather than leave the
+    wave PENDING forever - the completion wait has no timeout, so a stranded
+    PENDING branch would hang the whole map.
+    """
+
+    class TestExecutor(ConcurrentExecutor):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.calls: dict[int, int] = {}
+            self.long_runner_release = threading.Event()
+
+        def execute_item(self, child_context, executable):
+            task_id = executable.index
+            self.calls[task_id] = self.calls.get(task_id, 0) + 1
+            if task_id == 0:
+                # Long-runner keeps the map alive so task 1 resumes in-process.
+                self.long_runner_release.wait(timeout=5)
+                return "result_A"
+            # Task 1 suspends with a past timestamp -> immediate in-process resume.
+            msg = "resume-me"
+            raise TimedSuspendExecution(msg, time.time() - 1)
+
+    executables = [Executable(0, lambda: "task_A"), Executable(1, lambda: "task_B")]
+    completion_config = CompletionConfig(
+        min_successful=2,
+        tolerated_failure_count=None,
+        tolerated_failure_percentage=None,
+    )
+
+    executor = TestExecutor(
+        executables=executables,
+        max_concurrency=2,
+        completion_config=completion_config,
+        sub_type_top="TOP",
+        sub_type_iteration="ITER",
+        name_prefix="test_",
+        serdes=None,
+    )
+
+    execution_state = Mock()
+
+    def checkpoint(*args, **kwargs):
+        # The resume refresh calls create_checkpoint() with no arguments.
+        # Fail that call; leave the branches' own checkpoints as no-ops.
+        if not args and not kwargs:
+            msg = "resume refresh failed"
+            raise RuntimeError(msg)
+
+    execution_state.create_checkpoint = Mock(side_effect=checkpoint)
+
+    executor_context = Mock()
+    executor_context._create_step_id_for_logical_step = lambda *args: "1"  # noqa: SLF001
+    child_context = Mock()
+    child_context.state.wrap_user_function = lambda func, *args, **kwargs: func
+    executor_context.create_child_context = lambda *args, **kwargs: child_context
+
+    # Must re-raise (not hang): the resume failure surfaces as the original error.
+    with pytest.raises(RuntimeError, match="resume refresh failed"):
+        executor.execute(execution_state, executor_context)
+    executor.long_runner_release.set()
+
+
 def test_concurrent_executor_with_timed_resubmit_while_other_task_running():
     """Test timed resubmission while other tasks are still running."""
 
@@ -3200,7 +3268,9 @@ def test_timer_scheduler_fifo_ordering_with_same_timestamp():
     items synchronously, so callback order is deterministic.
     """
     results = []
-    resubmit_callback = Mock(side_effect=lambda exe: results.append(exe.index))
+    resubmit_callback = Mock(
+        side_effect=lambda batch: results.extend(exe.index for exe in batch)
+    )
 
     with TimerScheduler(resubmit_callback) as scheduler:
         # Use a past timestamp so they trigger immediately