fix: remove idle waits from checkpoint batching

After collecting a synchronous checkpoint, the collector waited up to
100 ms on an empty queue. The caller waits until the batch persists, so
it cannot add work. The wait only delayed it. Sequential steps paid it
once per step.

Once a batch holds a synchronous checkpoint, wait 1 ms on an empty
queue instead. A batch with no blocked caller keeps the full window, so
a step's asynchronous START still shares a request with its SUCCEED.

With a 1 ms window, refreshes from independent coordinators can split
into separate requests when they arrive more than 1 ms apart. A refresh
is the empty checkpoint a coordinator sends to see that a wait has
ended. The coordinator knows the end time when the branch suspends, so
it now requests the refresh then, with that time attached, through
ExecutionState.schedule_refresh. The collector holds refreshes until
their time and sends all refreshes due at one time in one request. A
refresh scheduled before a batch is sealed joins it. Failure, completion
and shutdown settle every pending refresh.

On Lambda at 1024 MB, 1000 sequential steps went from 142 ms to 41 ms
per step. Ten nested coordinators resumed a wave 89 to 104 ms after its
end time instead of 288 to 303 ms.

Fixes #710

Author: yaythomas

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

@@ -8,6 +8,7 @@
 import time
 from collections import deque
 from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Generic, TypeVar, cast
 
 from aws_durable_execution_sdk_python.concurrency.models import (
@@ -56,6 +57,7 @@
     from aws_durable_execution_sdk_python.state import (
         CheckpointedResult,
         ExecutionState,
+        ScheduledRefresh,
     )
 
 
@@ -86,6 +88,20 @@ def _branch_error_object(err: Exception) -> ErrorObject:
     return ErrorObject.from_exception(err)
 
 
+@dataclass
+class ResumeWave(Generic[CallableType, ResultType]):
+    """Branches suspended until one time, and the refresh that resumes them.
+
+    The refresh is a delayed empty checkpoint requested when the first branch
+    suspends. Its response shows the waits complete, so the wave resumes one
+    round trip after its time.
+    """
+
+    resume_at: float
+    refresh: ScheduledRefresh
+    branches: list[Branch[CallableType, ResultType]] = field(default_factory=list)
+
+
 class ConcurrentExecutor(Generic[CallableType, ResultType]):
     """Execute durable operations concurrently. This contains the execution logic for Map and Parallel.
 
@@ -252,7 +268,9 @@ def execute(
 
         events: queue.Queue[BranchEvent[ResultType]] = queue.Queue()
         pending: deque[Branch[CallableType, ResultType]] = deque(self.branches)
-        timed_resumes: list[tuple[float, int]] = []
+        # The heap orders resume times. The dict groups the branches under each.
+        resume_times: list[float] = []
+        waves: dict[float, ResumeWave[CallableType, ResultType]] = {}
         branch_by_index: dict[int, Branch[CallableType, ResultType]] = {
             branch.index: branch for branch in self.branches
         }
@@ -317,19 +335,20 @@ def submit(branch: Branch[CallableType, ResultType]) -> None:
                     running += 1
                     needs_snapshot_rebuild = True
 
-                # Resume due timed suspends in-process. One checkpoint
-                # refresh serves the whole due wave; a failure is terminal
+                # Resume due waves in-process. A refresh failure is terminal
                 # for the execution and propagates from this thread.
                 now: float = time.time()
-                due: list[Branch[CallableType, ResultType]] = []
-                while timed_resumes and timed_resumes[0][0] <= now:
-                    _, index = heapq.heappop(timed_resumes)
-                    due.append(branch_by_index[index])
-                if due:
-                    execution_state.create_checkpoint()
-                    for branch in due:
+                resumed = False
+                while resume_times and resume_times[0] <= now:
+                    wave = waves.pop(heapq.heappop(resume_times))
+                    wave.refresh.wait()
+                    # Branches joined the wave in event order. Resume in index
+                    # order so scheduling stays deterministic.
+                    for branch in sorted(wave.branches, key=lambda b: b.index):
                         submit(branch)
                         running += 1
+                    resumed = True
+                if resumed:
                     continue
 
                 if running == 0:
@@ -340,18 +359,18 @@ def submit(branch: Branch[CallableType, ResultType]) -> None:
                         raise retryable_error
                     # Every in-flight branch is suspended and no slot is
                     # free (or no work remains): suspend the parent.
-                    if timed_resumes:
+                    if resume_times:
                         raise TimedSuspendExecution(
                             "All concurrent work complete or suspended pending retry.",
-                            timed_resumes[0][0],
+                            resume_times[0],
                         )
                     raise SuspendExecution(
                         "All concurrent work complete or suspended and pending external callback."
                     )
 
                 timeout: float | None = None
-                if timed_resumes:
-                    timeout = max(timed_resumes[0][0] - time.time(), 0)
+                if resume_times:
+                    timeout = max(resume_times[0] - time.time(), 0)
                 try:
                     event: BranchEvent[ResultType] = events.get(timeout=timeout)
                 except queue.Empty:
@@ -383,7 +402,13 @@ def submit(branch: Branch[CallableType, ResultType]) -> None:
                         needs_snapshot_rebuild = True
                     case BranchEventKind.SUSPENDED_UNTIL if event.resume_at is not None:
                         applied.suspend_until(event.resume_at)
-                        heapq.heappush(timed_resumes, (event.resume_at, event.index))
+                        if event.resume_at not in waves:
+                            waves[event.resume_at] = ResumeWave(
+                                event.resume_at,
+                                execution_state.schedule_refresh(event.resume_at),
+                            )
+                            heapq.heappush(resume_times, event.resume_at)
+                        waves[event.resume_at].branches.append(applied)
                         running -= 1
                         needs_snapshot_rebuild = True
                     case BranchEventKind.ORPHANED:
@@ -402,6 +427,9 @@ def submit(branch: Branch[CallableType, ResultType]) -> None:
                         msg = f"Unhandled branch event: {event}"
                         raise InvalidStateError(msg)
         finally:
+            # Nothing will wait for a refresh whose wave never resumed.
+            for wave in waves.values():
+                wave.refresh.cancel()
             # Shutdown without waiting for running threads for early return
             # when completion criteria are met (e.g., min_successful).
             # Running threads continue in the background of this invocation