fix: shorten checkpoint wait for blocked caller

The checkpoint collector waited on an empty queue after it had already
collected a synchronous checkpoint. A synchronous checkpoint has a
caller blocked until the batch is persisted, and a blocked caller cannot
enqueue more work. So that wait could not enlarge the batch. It only
delayed the caller by the shutdown-signal poll interval of 100 ms.

Sequential steps paid that cost once per step. Measured in us-east-1 on
python3.13 at 1024 MB over 1000 steps, warm, comparing this change
against d61985e:

  median interval between StepSucceeded   142 ms  ->  41 ms
  p95 interval                            163 ms  ->  62 ms
  invocation duration                    144.4 s  ->  43.1 s
  steps per second                          6.92  ->  23.22

The step body performs no I/O, so nearly all of the removed time was the
idle wait.

Changes

- Track whether a batch holds a synchronous checkpoint. Set the flag on
  all three insertion paths: the overflow drain, the first blocking
  read, and the additional read.
- Once set, wait only _BLOCKED_CALLER_WAIT_SECONDS on an empty queue
  instead of the full poll interval. Work already queued still joins the
  batch, so an asynchronous step START still travels in the same request
  as its synchronous SUCCEED. Sequential steps keep one checkpoint call
  of two updates per step, unchanged.
- Wait 1 ms rather than not at all. Branch threads start one at a time,
  so gaps in checkpoint arrivals occur while a parallel or map thread
  pool starts. Not waiting cut extra small batches during that startup.
  On Lambda, waiting 1 ms produced the same checkpoint call count as
  d61985e at both 200 and 500 branches, and ran 1.21x to 1.28x faster.
- Include empty synchronous checkpoints. ConcurrentExecutor collects
  every due branch, calls create_checkpoint() once for the whole resume
  wave, then resubmits. So a wide map resume enqueues one empty
  checkpoint, not one per branch, and the full window would delay that
  single caller with nothing to coalesce.
- Leave asynchronous-only batching unchanged. With no blocked caller the
  wait costs no latency and still reduces the API call count.
- Name the existing 0.1 second literal _STOP_SIGNAL_POLL_SECONDS. Both
  waits in the collector are now named, so the min() compares two named
  intervals rather than one named value and one literal.
- Correct the max_batch_time_seconds docstring. It bounds the total
  collection window, not the idle wait, which the poll interval caps
  regardless of the configured value.

Neither constant is reachable by an end user, so this changes no public
signature. CheckpointBatcherConfig is absent from the package exports,
no caller outside state.py passes batcher_config, and config.py has no
batcher settings.

Tests

Add six collector tests that assert the timeout values passed to the
queue read rather than elapsed time. The collector's waiting policy is
fully described by those timeouts, so the tests need no sleeps and no
wall-clock thresholds. Five of the six fail when the shortened wait is
removed. The sixth pins the asynchronous-only path this change leaves
alone.

Rewrite test_create_checkpoint_multiple_sync_calls_all_block. Its old
assertion required each caller to block for at least 0.15 s, but each
caller started its clock after the processor began its 0.15 s sleep, so
its own elapsed time was always under that. The assertion held only
because the collector added 100 ms on top, which is the defect. The test
now proves the invariant by observation order and uses no delay at all.
It waits until all three operations are queued, asserts no caller has
returned, collects the batch, asserts no caller has returned, then sets
the completion events and asserts every caller returned. Settling the
waiters before collection fails it.

Correct the description in map_with_concurrent_waits_int_test. It said
the test mirrors the resubmitter issuing one empty checkpoint per
branch. ConcurrentExecutor issues one per resume wave, so the test
exercises a synthetic burst. It still guards the batch-limit behaviour
that #325 changed.

Fixes #710

Author: yaythomas

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

@@ -47,14 +47,25 @@
 
 logger = logging.getLogger(__name__)
 
+# Longest wait on an empty queue before re-checking the shutdown signal.
+_STOP_SIGNAL_POLL_SECONDS = 0.1
+
+# Longest wait on an empty queue when a caller is already blocked on the batch.
+# Long enough for sibling branch threads to arrive. Short enough that the
+# blocked caller stays fast.
+_BLOCKED_CALLER_WAIT_SECONDS = 0.001
+
 
 @dataclass(frozen=True)
 class CheckpointBatcherConfig:
     """Configuration for checkpoint batching behavior.
 
     Attributes:
         max_batch_size_bytes: Maximum batch size in bytes (default: 750KB)
-        max_batch_time_seconds: Maximum time to wait before flushing batch (default: 1.0 second)
+        max_batch_time_seconds: Longest a batch keeps accumulating (default:
+            1.0 second). This caps total collection, not the idle wait. The
+            collector flushes as soon as one poll slice finds the queue empty.
+            So a larger value only matters while operations keep arriving.
         max_batch_operations: Maximum number of operations per batch (default: 250)
     """
 
@@ -1025,6 +1036,14 @@ def _collect_checkpoint_batch(self) -> list[QueuedOperation]:
         operation if queues are empty, then collects additional operations within the time
         window.
 
+        How long the collector waits on an empty queue depends on the batch. A batch
+        holding a sync checkpoint waits _BLOCKED_CALLER_WAIT_SECONDS, because a caller is
+        blocked on it. An async-only batch waits the full window, because no caller is
+        blocked and a larger batch means fewer API calls.
+
+        Both cases drain the queue first. So a step's async START and its sync SUCCEED go
+        in one request.
+
         Empty checkpoints (operation_update=None) are coalesced: the first empty checkpoint
         counts toward the batch operation limit, but subsequent empty checkpoints do not.
         All empty checkpoints remain in the batch so their completion events are signaled.
@@ -1037,13 +1056,18 @@ def _collect_checkpoint_batch(self) -> list[QueuedOperation]:
         """
         batch: list[QueuedOperation] = []
         has_empty_checkpoint = False
+        # Set once the batch holds any sync checkpoint. That caller is blocked until
+        # the batch persists, so it cannot queue more work.
+        has_blocked_caller = False
         total_size = 0
         effective_operation_count = 0  # Operations that count toward batch limit
 
         # First, drain overflow queue (FIFO order preserved)
         try:
             while effective_operation_count < self._batcher_config.max_batch_operations:
                 overflow_op = self._overflow_queue.get_nowait()
+                if overflow_op.completion_event is not None:
+                    has_blocked_caller = True
 
                 if overflow_op.operation_update is None:  # Empty checkpoint
                     batch.append(overflow_op)
@@ -1071,10 +1095,12 @@ def _collect_checkpoint_batch(self) -> list[QueuedOperation]:
             while not self._checkpointing_stopped.is_set():
                 try:
                     first_op = self._checkpoint_queue.get(
-                        timeout=0.1
-                    )  # Check stop signal every 100ms
+                        timeout=_STOP_SIGNAL_POLL_SECONDS
+                    )
                     self._checkpoint_queue.task_done()
                     batch.append(first_op)
+                    if first_op.completion_event is not None:
+                        has_blocked_caller = True
 
                     if first_op.operation_update is None:
                         has_empty_checkpoint = True
@@ -1101,15 +1127,22 @@ def _collect_checkpoint_batch(self) -> list[QueuedOperation]:
         ):
             remaining_time = min(
                 batch_deadline - time.time(),
-                0.1,  # Check stop signal every 100ms
+                _STOP_SIGNAL_POLL_SECONDS,
             )
 
             if remaining_time <= 0:
                 break
 
             try:
-                additional_op = self._checkpoint_queue.get(timeout=remaining_time)
+                if has_blocked_caller:
+                    additional_op = self._checkpoint_queue.get(
+                        timeout=min(_BLOCKED_CALLER_WAIT_SECONDS, remaining_time)
+                    )
+                else:
+                    additional_op = self._checkpoint_queue.get(timeout=remaining_time)
                 self._checkpoint_queue.task_done()
+                if additional_op.completion_event is not None:
+                    has_blocked_caller = True
 
                 if additional_op.operation_update is None:  # Empty checkpoint
                     batch.append(additional_op)

packages/aws-durable-execution-sdk-python/tests/e2e/map_with_concurrent_waits_int_test.py

@@ -8,24 +8,23 @@
 
 Background
 ----------
-When a map branch suspends via TimedSuspendExecution and later resumes, the
-ConcurrentExecutor resubmitter calls::
+NOTE ON PRODUCTION FIDELITY. ConcurrentExecutor no longer issues one empty
+checkpoint per branch. It collects every due branch, calls
+execution_state.create_checkpoint() once for the whole resume wave, then
+resubmits (executor.py, "One checkpoint refresh serves the whole due wave").
+So a wide map resume enqueues ONE empty checkpoint, not hundreds.
 
-    execution_state.create_checkpoint()  # empty checkpoint
-
-before resubmitting the branch. In high-concurrency scenarios (300+ branches)
-all resuming at the same time, 300+ empty checkpoints flood the checkpoint
-queue.
+These tests therefore exercise a synthetic burst of concurrent empty
+checkpoints rather than the current resubmitter. They still guard the batch
+limit behaviour below, which is what issue #325 changed.
 
 Without the coalescing optimization (issue #325), the 250-operation batch limit
 causes these to be split across multiple batches → multiple API calls.
 With the optimization, all subsequent empty checkpoints beyond the first do
 NOT count toward the batch limit, so they are coalesced into a single batch
 and a single API call.
 
-These tests directly simulate that concurrent-checkpoint pattern by launching
-many threads that each call ``create_checkpoint()`` simultaneously, mirroring
-what the map resubmitter does when all branches resume at once.
+They launch many threads that each call ``create_checkpoint()`` at once.
 """
 
 from __future__ import annotations
@@ -92,16 +91,15 @@ def _checkpoint(
 
 
 def test_map_with_concurrent_waits_coalesces_empty_checkpoints():
-    """300 concurrent branches all create empty checkpoints simultaneously.
+    """300 concurrent empty checkpoints must collect into one batch.
 
-    Simulates the Java MapWithConditionAndCallbackExample scenario: 300 map
-    branches all resuming from a wait operation at the same time, each calling
-    the resubmitter which enqueues an empty checkpoint.
+    A synthetic burst, not the current resubmitter, which issues one empty
+    checkpoint per resume wave. See the module docstring.
 
-    Without the coalescing optimization, the 250-op batch limit splits 300
-    empty checkpoints into 2 batches (250 + 50) → 2 API calls.
-    With the optimization (effective_operation_count stays 1 for empties),
-    all 300 are collected in a single batch → 1 API call.
+    Without the coalescing optimization the 250-op batch limit splits 300 empty
+    checkpoints into 2 batches of 250 and 50, so 2 API calls. With it,
+    effective_operation_count stays 1 for empties, so all 300 collect into one
+    batch and one API call.
     """
     mock_client, calls = _make_tracking_client()
     state = _make_state(mock_client, batch_time=5.0, max_ops=250)