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 that carries an
  operation update. 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.
- Leave empty synchronous checkpoints on the full window. The branch
  resubmitter issues one per branch on resume, so a wide map resuming at
  once produces hundreds simultaneously, and flushing early would split
  them across requests. That is the coalescing added in #325 and guarded
  by map_with_concurrent_waits_int_test. The exclusion costs one delay
  per resume wave, not per branch.
- 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. Four of the six fail when the shortened wait is
removed; the other two pin the paths this change leaves alone.

Replace the wall-clock assertion in
test_create_checkpoint_multiple_sync_calls_all_block. Each caller
started its clock after the processor began its 150 ms sleep, so its
own elapsed time was always under 150 ms. The assertion held only
because the collector added 100 ms on top, which is the defect. It now
asserts the real invariant: no caller returns before its batch is
collected, and every caller stays blocked across the processor delay.

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 with an operation update waits
+        _BLOCKED_CALLER_WAIT_SECONDS, because a caller is blocked on it. Any other batch
+        waits the full window, because no caller is blocked.
+
+        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,6 +1056,11 @@ def _collect_checkpoint_batch(self) -> list[QueuedOperation]:
         """
         batch: list[QueuedOperation] = []
         has_empty_checkpoint = False
+        # Set once the batch holds a sync checkpoint with an operation update. That
+        # caller is blocked until the batch persists, so it cannot queue more work.
+        # Empty sync checkpoints do not count because many arrive together when a
+        # map or parallel resumes, and they must coalesce.
+        has_blocked_caller = False
         total_size = 0
         effective_operation_count = 0  # Operations that count toward batch limit
 
@@ -1060,6 +1084,8 @@ def _collect_checkpoint_batch(self) -> list[QueuedOperation]:
                         self._overflow_queue.put(overflow_op)
                         break
                     batch.append(overflow_op)
+                    if overflow_op.completion_event is not None:
+                        has_blocked_caller = True
                     total_size += op_size
                     effective_operation_count += 1
         except queue.Empty:
@@ -1071,15 +1097,17 @@ 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.operation_update is None:
                         has_empty_checkpoint = True
                     else:
                         total_size += self._calculate_operation_size(first_op)
+                        if first_op.completion_event is not None:
+                            has_blocked_caller = True
 
                     effective_operation_count = 1
                     break
@@ -1101,14 +1129,19 @@ 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.operation_update is None:  # Empty checkpoint
@@ -1130,6 +1163,8 @@ def _collect_checkpoint_batch(self) -> list[QueuedOperation]:
                         )
                         break
                     batch.append(additional_op)
+                    if additional_op.completion_event is not None:
+                        has_blocked_caller = True
                     total_size += op_size
                     effective_operation_count += 1