From 6a22723140df024f16b0e4afbe8eccd0a0033b7e Mon Sep 17 00:00:00 2001
From: yaythomas <tgaigher@amazon.com>
Date: Tue, 15 Sep 2026 07:42:00 +0000
Subject: [PATCH] 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
---
 .../concurrency/executor.py                   |  58 +-
 .../aws_durable_execution_sdk_python/state.py | 452 +++++---
 .../tests/concurrency_test.py                 | 190 +++-
 .../e2e/map_with_concurrent_waits_int_test.py |  89 +-
 .../tests/state_test.py                       | 985 ++++++++++++++----
 5 files changed, 1338 insertions(+), 436 deletions(-)

diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py
index c45ca0bd..2a39c4ef 100644
--- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py
+++ b/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
diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py
index 0b6a5fcf..f5ce7214 100644
--- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py
+++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py
@@ -3,13 +3,15 @@
 from __future__ import annotations
 
 import functools
+import heapq
+import itertools
 import json
 import logging
 import queue
 import threading
 import time
 from concurrent.futures import ThreadPoolExecutor
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum
 from threading import Lock
 from typing import TYPE_CHECKING, Callable, NoReturn
@@ -47,6 +49,14 @@
 
 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:
@@ -54,7 +64,10 @@ class CheckpointBatcherConfig:
 
     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). The collector flushes as soon as a wait on the queue
+            finds it empty. So this value only matters while operations keep
+            arriving.
         max_batch_operations: Maximum number of operations per batch (default: 250)
     """
 
@@ -76,6 +89,82 @@ class QueuedOperation:
     completion_event: CompletionEvent | None = None
 
 
+class _Signal(Enum):
+    """Control items that travel through the checkpoint queue, never sent."""
+
+    # A refresh was scheduled. The collector re-reads the heap and its timeout.
+    REFRESH_WAKE = "refresh_wake"
+
+
+@dataclass(frozen=True)
+class ScheduledRefresh:
+    """Handle for a delayed empty checkpoint from ExecutionState.schedule_refresh.
+
+    Internal to the SDK. Not part of the public API.
+
+    A response fetched before earliest_check_time cannot show the wait complete,
+    so the collector holds the checkpoint until then. wait() blocks until it has
+    been sent and the operations reloaded, or raises the error that stopped
+    checkpointing. cancel() drops it if it has not been sent yet. One already in
+    flight completes normally.
+    """
+
+    earliest_check_time: float
+    completion_event: CompletionEvent
+    cancelled: threading.Event = field(default_factory=threading.Event)
+
+    def wait(self, timeout: float | None = None) -> bool:
+        return self.completion_event.wait(timeout)
+
+    def is_set(self) -> bool:
+        return self.completion_event.is_set()
+
+    def is_cancelled(self) -> bool:
+        return self.cancelled.is_set()
+
+    def cancel(self) -> None:
+        self.cancelled.set()
+
+
+class _BatchAccumulator:
+    """The operations of one checkpoint request, with the limit accounting.
+
+    Empty checkpoints carry no bytes. The first one counts toward the operation
+    limit and later ones do not, so a resume wave of any width fits one request.
+    """
+
+    def __init__(self, config: CheckpointBatcherConfig) -> None:
+        self._config = config
+        self.operations: list[QueuedOperation] = []
+        self.total_size = 0
+        self.effective_count = 0
+        self.has_empty = False
+        # A sync checkpoint's caller is blocked until the batch persists, so it
+        # cannot queue more work. Waiting long for more work is then pointless.
+        self.has_blocked_caller = False
+
+    def is_empty(self) -> bool:
+        return not self.operations
+
+    def is_full(self) -> bool:
+        return self.effective_count >= self._config.max_batch_operations
+
+    def fits(self, size: int) -> bool:
+        return self.total_size + size <= self._config.max_batch_size_bytes
+
+    def add(self, op: QueuedOperation, size: int = 0) -> None:
+        self.operations.append(op)
+        if op.completion_event is not None:
+            self.has_blocked_caller = True
+        if op.operation_update is None:
+            if not self.has_empty:
+                self.effective_count += 1
+                self.has_empty = True
+            return
+        self.total_size += size
+        self.effective_count += 1
+
+
 # Statuses indicating an operation has finished and will not change on a later
 # replay. Includes TIMED_OUT/CANCELLED/STOPPED in addition to SUCCEEDED/FAILED
 _TERMINAL_OPERATION_STATUSES: frozenset[OperationStatus] = frozenset(
@@ -266,7 +355,10 @@ class ReplayStatus(Enum):
 
 
 class ExecutionState:
-    """Get, set and maintain execution state. This is mutable. Create and check checkpoints."""
+    """Get, set and maintain execution state. This is mutable. Create and check checkpoints.
+
+    Internal to the SDK. Not part of the public API.
+    """
 
     def __init__(
         self,
@@ -291,8 +383,14 @@ def __init__(
         )
 
         # Checkpoint batching components
-        self._checkpoint_queue: queue.Queue[QueuedOperation] = queue.Queue()
+        self._checkpoint_queue: queue.Queue[QueuedOperation | _Signal] = queue.Queue()
         self._overflow_queue: queue.Queue[QueuedOperation] = queue.Queue()
+        # Refreshes not yet sent, ordered by earliest_check_time. Producers push
+        # and the collector pops, so every access holds _completion_lock.
+        self._pending_refreshes: list[tuple[float, int, ScheduledRefresh]] = []
+        self._pending_refresh_seq = itertools.count()
+        # True while a REFRESH_WAKE is in the queue, so many refreshes make one.
+        self._refresh_wake_enqueued = False
         self._checkpointing_stopped: threading.Event = threading.Event()
         self._checkpointing_failed: CompletionEvent = CompletionEvent()
         # Set once the service confirms the execution has completed (a checkpoint
@@ -720,6 +818,77 @@ def create_checkpoint(
         else:
             logger.debug("Enqueued checkpoint operation for asynchronous processing")
 
+    def schedule_refresh(self, earliest_check_time: float) -> ScheduledRefresh:
+        """Enqueue a delayed empty checkpoint and return at once.
+
+        The collector holds it until earliest_check_time, then sends it with
+        every other one due at that time in one request, however far apart they
+        were requested. One requested after its time joins the next batch. The
+        caller waits on the handle when it needs the refreshed operations, and
+        cancels it if it stops needing them.
+        """
+        refresh = ScheduledRefresh(earliest_check_time, CompletionEvent())
+        with self._completion_lock:
+            if self._checkpointing_failed.is_set():
+                self._checkpointing_failed.wait()
+            self._reject_if_execution_completed(None)
+            if self._checkpointing_stopped.is_set():
+                raise OrphanedChildException(
+                    "Checkpointing stopped. The refresh will not be processed.",
+                    operation_id="",
+                )
+            heapq.heappush(
+                self._pending_refreshes,
+                (earliest_check_time, next(self._pending_refresh_seq), refresh),
+            )
+            if not self._refresh_wake_enqueued:
+                self._refresh_wake_enqueued = True
+                self._checkpoint_queue.put(_Signal.REFRESH_WAKE)
+        return refresh
+
+    def _consume_refresh_wake(self) -> None:
+        with self._completion_lock:
+            self._refresh_wake_enqueued = False
+
+    @staticmethod
+    def _settle_cancelled_refresh(refresh: ScheduledRefresh) -> None:
+        # Nobody waits on a cancelled refresh by contract. Setting the event
+        # keeps a waiter that broke the contract from blocking forever.
+        refresh.completion_event.set()
+
+    def _seconds_until_next_refresh(self, now: float) -> float | None:
+        # A cancelled refresh must not set the wake time. Each needless wake
+        # would extend the collection by one more read.
+        with self._completion_lock:
+            while (
+                self._pending_refreshes and self._pending_refreshes[0][2].is_cancelled()
+            ):
+                self._settle_cancelled_refresh(
+                    heapq.heappop(self._pending_refreshes)[2]
+                )
+            if not self._pending_refreshes:
+                return None
+            return self._pending_refreshes[0][0] - now
+
+    def _add_due_refreshes(self, batch: _BatchAccumulator, now: float) -> int:
+        """Move every refresh whose time has come into the batch. Returns how many."""
+        added = 0
+        with self._completion_lock:
+            while self._pending_refreshes and self._pending_refreshes[0][0] <= now:
+                refresh = heapq.heappop(self._pending_refreshes)[2]
+                if refresh.is_cancelled():
+                    self._settle_cancelled_refresh(refresh)
+                    continue
+                batch.add(QueuedOperation(None, refresh.completion_event))
+                added += 1
+        return added
+
+    def _drain_pending_refreshes(self) -> list[ScheduledRefresh]:
+        """Take every unsent refresh. The caller holds _completion_lock."""
+        drained = [entry[2] for entry in self._pending_refreshes]
+        self._pending_refreshes.clear()
+        return drained
+
     def create_checkpoint_sync(
         self,
         operation_update: OperationUpdate | None = None,
@@ -924,18 +1093,24 @@ def checkpoint_batches_forever(self) -> None:
                         while not self._overflow_queue.empty():
                             try:
                                 item = self._overflow_queue.get_nowait()
-                                if item.completion_event:
-                                    item.completion_event.set(bg_error)
                             except queue.Empty:
                                 break
+                            if item.completion_event:
+                                item.completion_event.set(bg_error)
 
                         while not self._checkpoint_queue.empty():
                             try:
-                                item = self._checkpoint_queue.get_nowait()
-                                if item.completion_event:
-                                    item.completion_event.set(bg_error)
+                                queued = self._checkpoint_queue.get_nowait()
                             except queue.Empty:
                                 break
+                            if (
+                                isinstance(queued, QueuedOperation)
+                                and queued.completion_event
+                            ):
+                                queued.completion_event.set(bg_error)
+
+                        for refresh in self._drain_pending_refreshes():
+                            refresh.completion_event.set(bg_error)
 
                         # Future checkpoint attempts fail immediately.
                         self._checkpointing_failed.set(bg_error)
@@ -943,6 +1118,19 @@ def checkpoint_batches_forever(self) -> None:
                     # Exit the loop - error has been signaled to main thread via completion events
                     break
 
+        # A refresh still pending at shutdown will never be sent. Its caller, if
+        # one is waiting, must not block forever. Settle it as orphaned. The lock
+        # orders this against schedule_refresh, which refuses once stopped.
+        with self._completion_lock:
+            unsent_refreshes = self._drain_pending_refreshes()
+        for refresh in unsent_refreshes:
+            refresh.completion_event.set(
+                OrphanedChildException(
+                    "Checkpointing stopped before the refresh time.",
+                    operation_id="",
+                )
+            )
+
         logger.debug("Background checkpoint processing stopped")
 
     def _settle_after_execution_completed(self) -> None:
@@ -959,24 +1147,34 @@ def _settle_after_execution_completed(self) -> None:
             self._execution_completed.set()
             self._checkpointing_stopped.set()
 
+            orphaned = OrphanedChildException(
+                "Execution already completed; checkpoint will not be processed.",
+                operation_id="",
+            )
+            for refresh in self._drain_pending_refreshes():
+                refresh.completion_event.set(orphaned)
+            unsent: list[QueuedOperation] = []
             for pending_queue in (self._overflow_queue, self._checkpoint_queue):
                 while not pending_queue.empty():
                     try:
-                        queued_op: QueuedOperation = pending_queue.get_nowait()
+                        item = pending_queue.get_nowait()
                     except queue.Empty:
                         break
-                    if queued_op.completion_event is not None:
-                        operation_id: str = (
-                            queued_op.operation_update.operation_id
-                            if queued_op.operation_update is not None
-                            else ""
-                        )
-                        queued_op.completion_event.set(
-                            OrphanedChildException(
-                                "Execution already completed; checkpoint will not be processed.",
-                                operation_id=operation_id,
-                            )
+                    if isinstance(item, QueuedOperation):
+                        unsent.append(item)
+            for queued_op in unsent:
+                if queued_op.completion_event is not None:
+                    operation_id: str = (
+                        queued_op.operation_update.operation_id
+                        if queued_op.operation_update is not None
+                        else ""
+                    )
+                    queued_op.completion_event.set(
+                        OrphanedChildException(
+                            "Execution already completed; checkpoint will not be processed.",
+                            operation_id=operation_id,
                         )
+                    )
 
     def stop_checkpointing(self) -> None:
         """Signal background thread to stop checkpointing.
@@ -1018,134 +1216,116 @@ def record_branch_fatal_error(
             return True
 
     def _collect_checkpoint_batch(self) -> list[QueuedOperation]:
-        """Collect multiple checkpoint operations into a batch for API efficiency.
+        """Collect the operations for one checkpoint request.
 
-        Processes overflow queue first to maintain FIFO order, then collects from main queue.
-        Respects configured size, time, and operation count limits. Blocks for the first
-        operation if queues are empty, then collects additional operations within the time
-        window.
+        Intake order is the overflow queue, then refreshes whose time has come,
+        then the main queue. The first read blocks until an operation arrives, a
+        refresh comes due, or shutdown is signalled. Later reads wait
+        _BLOCKED_CALLER_WAIT_SECONDS once the batch holds a sync checkpoint,
+        because that caller is blocked and cannot add work, and otherwise up to
+        max_batch_time_seconds.
 
-        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.
-        This avoids unnecessary batches when many concurrent map/parallel branches resume
-        simultaneously and each queues an empty checkpoint.
+        A refresh waits in the heap until its earliest_check_time, so refreshes
+        due at the same time share one request. The final due check seals the
+        batch. A refresh scheduled after it joins the next batch.
 
-        Returns:
-            List of QueuedOperation objects ready for batch processing. Returns empty list
-            if no operations are available.
+        Returns the operations for the request, or an empty list at shutdown.
         """
-        batch: list[QueuedOperation] = []
-        has_empty_checkpoint = 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.operation_update is None:  # Empty checkpoint
-                    batch.append(overflow_op)
-                    if not has_empty_checkpoint:
-                        effective_operation_count += (
-                            1  # First empty counts toward limit
-                        )
-                        has_empty_checkpoint = True
-                    # Subsequent empties don't count toward limit
-                else:
-                    op_size = self._calculate_operation_size(overflow_op)
-                    if total_size + op_size > self._batcher_config.max_batch_size_bytes:
-                        # Put back and stop
-                        self._overflow_queue.put(overflow_op)
-                        break
-                    batch.append(overflow_op)
-                    total_size += op_size
-                    effective_operation_count += 1
-        except queue.Empty:
-            pass
-
-        # If batch is empty, get first operation from main queue
-        if not batch:
-            # Block for first operation, checking stop signal periodically
-            while not self._checkpointing_stopped.is_set():
-                try:
-                    first_op = self._checkpoint_queue.get(
-                        timeout=0.1
-                    )  # Check stop signal every 100ms
-                    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)
-
-                    effective_operation_count = 1
-                    break
-                except queue.Empty:
-                    continue
-
-            # If stopped and no operation retrieved, return empty batch
-            if not batch:
-                return batch
-
-        # Start batching window using configured time
-        batch_deadline = time.time() + self._batcher_config.max_batch_time_seconds
-
-        # Collect additional operations within the time window
-        while (
-            time.time() < batch_deadline
-            and effective_operation_count < self._batcher_config.max_batch_operations
-            and not self._checkpointing_stopped.is_set()
-        ):
-            remaining_time = min(
-                batch_deadline - time.time(),
-                0.1,  # Check stop signal every 100ms
-            )
-
-            if remaining_time <= 0:
+        batch = _BatchAccumulator(self._batcher_config)
+        self._drain_overflow(batch)
+        self._add_due_refreshes(batch, time.time())
+
+        if batch.is_empty() and not self._wait_for_first_operation(batch):
+            return []
+
+        deadline = time.time() + self._batcher_config.max_batch_time_seconds
+        while not batch.is_full() and not self._checkpointing_stopped.is_set():
+            now = time.time()
+            self._add_due_refreshes(batch, now)
+            timeout = self._read_timeout(batch, now, deadline)
+            if timeout <= 0:
                 break
-
             try:
-                additional_op = self._checkpoint_queue.get(timeout=remaining_time)
-                self._checkpoint_queue.task_done()
-
-                if additional_op.operation_update is None:  # Empty checkpoint
-                    batch.append(additional_op)
-                    if not has_empty_checkpoint:
-                        effective_operation_count += (
-                            1  # First empty counts toward limit
-                        )
-                        has_empty_checkpoint = True
-                    # Subsequent empties don't count toward limit
-                else:
-                    op_size = self._calculate_operation_size(additional_op)
-                    # Check if adding this operation would exceed size limit
-                    if total_size + op_size > self._batcher_config.max_batch_size_bytes:
-                        # Put in overflow queue for next batch
-                        self._overflow_queue.put(additional_op)
-                        logger.debug(
-                            "Batch size limit reached, moving operation to overflow queue"
-                        )
-                        break
-                    batch.append(additional_op)
-                    total_size += op_size
-                    effective_operation_count += 1
-
+                item = self._checkpoint_queue.get(timeout=timeout)
             except queue.Empty:
+                # A due refresh can end the read early. Keep collecting only if one
+                # joined, so refreshes due within one wait share the batch.
+                if self._add_due_refreshes(batch, time.time()):
+                    continue
+                break
+            self._checkpoint_queue.task_done()
+            if isinstance(item, _Signal):
+                self._consume_refresh_wake()
+                continue
+            size = self._calculate_operation_size(item)
+            if not batch.fits(size):
+                self._overflow_queue.put(item)
+                logger.debug(
+                    "Batch size limit reached, moving operation to overflow queue"
+                )
                 break
+            batch.add(item, size)
+
+        self._add_due_refreshes(batch, time.time())
 
-        empty_count = sum(1 for q in batch if q.operation_update is None)
+        empty_count = sum(1 for q in batch.operations if q.operation_update is None)
         logger.debug(
             "Collected batch of %d operations (%d effective, %d non-empty, %d empty), total size: %d bytes",
-            len(batch),
-            effective_operation_count,
-            len(batch) - empty_count,
+            len(batch.operations),
+            batch.effective_count,
+            len(batch.operations) - empty_count,
             empty_count,
-            total_size,
+            batch.total_size,
         )
-        return batch
+        return batch.operations
+
+    def _drain_overflow(self, batch: _BatchAccumulator) -> None:
+        """Take operations left over from earlier batches, oldest first."""
+        try:
+            while not batch.is_full():
+                op = self._overflow_queue.get_nowait()
+                size = self._calculate_operation_size(op)
+                if not batch.fits(size):
+                    self._overflow_queue.put(op)
+                    break
+                batch.add(op, size)
+        except queue.Empty:
+            pass
+
+    def _wait_for_first_operation(self, batch: _BatchAccumulator) -> bool:
+        """Block until the batch has an operation. Returns False at shutdown."""
+        while not self._checkpointing_stopped.is_set():
+            now = time.time()
+            timeout = _STOP_SIGNAL_POLL_SECONDS
+            until_refresh = self._seconds_until_next_refresh(now)
+            if until_refresh is not None:
+                timeout = max(0.0, min(timeout, until_refresh))
+            try:
+                item = self._checkpoint_queue.get(timeout=timeout)
+            except queue.Empty:
+                if self._add_due_refreshes(batch, time.time()):
+                    return True
+                continue
+            self._checkpoint_queue.task_done()
+            if isinstance(item, _Signal):
+                self._consume_refresh_wake()
+                continue
+            # The first operation is never refused for size. It could not be
+            # sent otherwise.
+            batch.add(item, self._calculate_operation_size(item))
+            return True
+        return False
+
+    def _read_timeout(
+        self, batch: _BatchAccumulator, now: float, deadline: float
+    ) -> float:
+        timeout = min(deadline - now, _STOP_SIGNAL_POLL_SECONDS)
+        until_refresh = self._seconds_until_next_refresh(now)
+        if until_refresh is not None:
+            timeout = min(timeout, until_refresh)
+        if batch.has_blocked_caller:
+            timeout = min(timeout, _BLOCKED_CALLER_WAIT_SECONDS)
+        return timeout
 
     @staticmethod
     def _calculate_operation_size(queued_op: QueuedOperation) -> int:
diff --git a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py
index 2a6154cd..a9313566 100644
--- a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py
+++ b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py
@@ -1410,14 +1410,12 @@ def execute_item(self, child_context, executable):
 
     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)
+    # The coordinator requests the refresh when task 1 suspends and waits on
+    # the returned event when task 1 is due. Fail the wait. The branches' own
+    # checkpoints stay no-ops.
+    failing_refresh = Mock()
+    failing_refresh.wait = Mock(side_effect=RuntimeError("resume refresh failed"))
+    execution_state.schedule_refresh = Mock(return_value=failing_refresh)
 
     executor_context = Mock()
     executor_context._create_step_id_for_logical_step = lambda *args: "1"
@@ -3500,8 +3498,180 @@ def slow_sibling():
     assert call_counts[0] == 2
     assert result.success_count == 2
     assert result.completion_reason is CompletionReason.ALL_COMPLETED
-    # A resume wave refreshes state once before resubmitting.
-    execution_state.create_checkpoint.assert_called()
+    # One refresh is requested for the resume time and waited on before the
+    # wave is resubmitted.
+    execution_state.schedule_refresh.assert_called_once()
+    execution_state.schedule_refresh.return_value.wait.assert_called_once()
+
+
+def test_branches_sharing_a_resume_time_request_one_refresh():
+    """Two branches suspend until the same time while a sibling runs. The
+    coordinator requests one refresh for that time and waits on it once."""
+    # In the future, so both suspensions are recorded before the wave is due.
+    resume_at = time.time() + 0.3
+    release = threading.Event()
+    call_counts: dict[int, int] = {}
+
+    class TestExecutor(ConcurrentExecutor):
+        def execute_item(self, child_context, executable):
+            index = executable.index
+            call_counts[index] = call_counts.get(index, 0) + 1
+            if index == 0:
+                assert release.wait(timeout=5)
+                return "long"
+            if call_counts[index] == 1:
+                msg = "wait"
+                raise TimedSuspendExecution(msg, resume_at)
+            if all(call_counts.get(i) == 2 for i in (1, 2)):
+                release.set()
+            return f"resumed{index}"
+
+    executor = TestExecutor(
+        executables=[Executable(i, lambda: None) for i in range(3)],
+        max_concurrency=3,
+        completion_config=CompletionConfig(
+            min_successful=3,
+            tolerated_failure_count=None,
+            tolerated_failure_percentage=None,
+        ),
+        sub_type_top="TOP",
+        sub_type_iteration="ITER",
+        name_prefix="test_",
+        serdes=None,
+        operation_id_namespace=_StubNamespace(),
+    )
+    execution_state = Mock()
+    executor_context = Mock()
+    executor_context._create_step_id_for_logical_step = lambda *args: "1"
+    child_context = Mock()
+    child_context.state.wrap_user_function = lambda func, *args, **kwargs: func
+    executor_context.create_child_context = lambda *args, **kwargs: child_context
+
+    result = executor.execute(execution_state, executor_context)
+
+    assert result.success_count == 3
+    execution_state.schedule_refresh.assert_called_once_with(resume_at)
+    execution_state.schedule_refresh.return_value.wait.assert_called_once()
+
+
+def test_wave_resumes_branches_in_index_order():
+    """Branch 2 suspends before branch 1, both until the same time. The wave
+    resumes them as 1 then 2, not in the order their suspensions arrived."""
+    resume_at = time.time() + 0.3
+    release_long = threading.Event()
+    resumed: list[int] = []  # start() calls on suspended branches, coordinator order
+    call_counts: dict[int, int] = {}
+    original_start = Branch.start
+
+    def recording_start(branch):
+        if branch.status is BranchStatus.SUSPENDED_WITH_TIMEOUT:
+            resumed.append(branch.index)
+        original_start(branch)
+
+    class TestExecutor(ConcurrentExecutor):
+        def execute_item(self, child_context, executable):
+            index = executable.index
+            call_counts[index] = call_counts.get(index, 0) + 1
+            if index == 0:
+                assert release_long.wait(timeout=5)
+                return "long"
+            if call_counts[index] == 1:
+                if index == 1:
+                    # Raise only once the coordinator has recorded branch 2's
+                    # suspension, so the two events arrive as 2 then 1.
+                    deadline = time.monotonic() + 5
+                    while (
+                        self.branches[2].status
+                        is not BranchStatus.SUSPENDED_WITH_TIMEOUT
+                    ):
+                        assert time.monotonic() < deadline
+                        time.sleep(0.001)
+                msg = "wait"
+                raise TimedSuspendExecution(msg, resume_at)
+            if call_counts.get(1) == 2 and call_counts.get(2) == 2:
+                release_long.set()
+            return f"resumed{index}"
+
+    executor = TestExecutor(
+        executables=[Executable(i, lambda: None) for i in range(3)],
+        max_concurrency=3,
+        completion_config=CompletionConfig(
+            min_successful=3,
+            tolerated_failure_count=None,
+            tolerated_failure_percentage=None,
+        ),
+        sub_type_top="TOP",
+        sub_type_iteration="ITER",
+        name_prefix="test_",
+        serdes=None,
+        operation_id_namespace=_StubNamespace(),
+    )
+    execution_state = Mock()
+    executor_context = Mock()
+    executor_context._create_step_id_for_logical_step = lambda *args: "1"
+    child_context = Mock()
+    child_context.state.wrap_user_function = lambda func, *args, **kwargs: func
+    executor_context.create_child_context = lambda *args, **kwargs: child_context
+
+    with patch.object(Branch, "start", recording_start):
+        result = executor.execute(execution_state, executor_context)
+
+    assert result.success_count == 3
+    assert resumed == [1, 2]
+
+
+def test_early_completion_cancels_unused_refresh():
+    """min_successful is met while another branch waits on a far timed resume.
+    The coordinator will never wait for that refresh, so it cancels it. Sending
+    it later would waste a request."""
+
+    refresh_requested = threading.Event()
+
+    class TestExecutor(ConcurrentExecutor):
+        def execute_item(self, child_context, executable):
+            if executable.index == 0:
+                # Complete only after the coordinator has requested the refresh
+                # for branch 1, so the cancel path is always exercised.
+                assert refresh_requested.wait(timeout=5)
+                return "fast"
+            msg = "far away"
+            raise TimedSuspendExecution(msg, time.time() + 3600)
+
+    executor = TestExecutor(
+        executables=[Executable(0, lambda: "a"), Executable(1, lambda: "b")],
+        max_concurrency=2,
+        completion_config=CompletionConfig(
+            min_successful=1,
+            tolerated_failure_count=None,
+            tolerated_failure_percentage=None,
+        ),
+        sub_type_top="TOP",
+        sub_type_iteration="ITER",
+        name_prefix="test_",
+        serdes=None,
+        operation_id_namespace=_StubNamespace(),
+    )
+
+    execution_state = Mock()
+    refresh = Mock()
+
+    def schedule_refresh(_resume_at):
+        refresh_requested.set()
+        return refresh
+
+    execution_state.schedule_refresh = Mock(side_effect=schedule_refresh)
+    executor_context = Mock()
+    executor_context._create_step_id_for_logical_step = lambda *args: "1"
+    child_context = Mock()
+    child_context.state.wrap_user_function = lambda func, *args, **kwargs: func
+    executor_context.create_child_context = lambda *args, **kwargs: child_context
+
+    result = executor.execute(execution_state, executor_context)
+
+    assert result.success_count == 1
+    execution_state.schedule_refresh.assert_called_once()
+    refresh.cancel.assert_called_once()
+    refresh.wait.assert_not_called()
 
 
 def test_all_timed_suspended_parent_suspends_with_earliest_timestamp():
diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/map_with_concurrent_waits_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/map_with_concurrent_waits_int_test.py
index 62ad7c2b..62cd640d 100644
--- a/packages/aws-durable-execution-sdk-python/tests/e2e/map_with_concurrent_waits_int_test.py
+++ b/packages/aws-durable-execution-sdk-python/tests/e2e/map_with_concurrent_waits_int_test.py
@@ -1,36 +1,24 @@
-"""Integration test: empty checkpoint coalescing with concurrent map + wait.
-
-Python equivalent of the Java MapWithConditionAndCallbackExample referenced in
-issue #325. Verifies that when many concurrent map branches resume from timed
-wait operations simultaneously, the empty checkpoints produced by the
-resubmitter (executor.py) are coalesced into minimal API calls instead of
-being split across multiple batches.
-
-Background
-----------
-When a map branch suspends via TimedSuspendExecution and later resumes, the
-ConcurrentExecutor resubmitter calls::
-
-    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.
-
-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.
+"""Integration tests for coalescing the refreshes of concurrent resume waves.
+
+A coordinator that resumes timed waits in-process needs a refresh, an empty
+checkpoint whose response shows the waits complete. It requests one refresh
+per distinct resume time through ExecutionState.schedule_refresh(resume_at).
+Independent coordinators, such as nested maps, each request their own.
+
+The batcher holds a refresh until resume_at, then sends every refresh due at
+that time in one request. So the number of requests depends on the resume
+times, not on how far apart the requests were made or how the threads were
+scheduled.
+
+The batch operation limit still applies. The first empty checkpoint counts
+toward the 250-operation limit and the rest do not, so 300 refreshes fit in
+one batch. These tests verify both rules.
 """
 
 from __future__ import annotations
 
 import threading
+import time
 from concurrent.futures import ThreadPoolExecutor
 
 
@@ -92,33 +80,28 @@ def _checkpoint(
 
 
 def test_map_with_concurrent_waits_coalesces_empty_checkpoints():
-    """300 concurrent branches all create empty checkpoints simultaneously.
+    """300 due refreshes from 300 independent callers must make one API call.
 
-    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.
-
-    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.
+    All 300 threads request their refreshes before the batcher starts, so the
+    result does not depend on how fast the scheduler runs them. The check time
+    is in the past: deferral of future refreshes is covered by unit tests, and
+    here every refresh is due when the batcher first looks. Without the
+    batch-limit optimization the 250-op limit would split them into 2 requests.
     """
     mock_client, calls = _make_tracking_client()
     state = _make_state(mock_client, batch_time=5.0, max_ops=250)
 
-    batcher = ThreadPoolExecutor(max_workers=1)
-    batcher.submit(state.checkpoint_batches_forever)
-
-    # 300 branches all call create_checkpoint() concurrently, each blocking
-    # until the batch is processed — mirrors the resubmitter pattern.
     branch_count = 300
-    start_barrier = threading.Barrier(branch_count)
+    check_time = time.time() - 1.0
     errors: list[Exception] = []
+    handles = []
+    handles_lock = threading.Lock()
 
     def branch_work():
         try:
-            start_barrier.wait()  # all start simultaneously
-            state.create_checkpoint()  # empty checkpoint, synchronous
+            handle = state.schedule_refresh(check_time)
+            with handles_lock:
+                handles.append(handle)
         except Exception as e:  # noqa: BLE001
             errors.append(e)
 
@@ -127,17 +110,19 @@ def branch_work():
         t.start()
     for t in threads:
         t.join(timeout=30)
+    assert not errors, f"Branch errors: {errors}"
+    assert len(handles) == branch_count, "every caller must have enqueued first"
 
+    batcher = ThreadPoolExecutor(max_workers=1)
+    batcher.submit(state.checkpoint_batches_forever)
     try:
-        assert not errors, f"Branch errors: {errors}"
-
-        # All 300 empty checkpoints should be batched into 1 API call.
-        # Without the fix, 300 > 250 limit would produce 2 calls.
+        for handle in handles:
+            assert handle.wait(timeout=30)
         assert len(calls) == 1, (
-            f"Expected 1 coalesced API call for {branch_count} concurrent empty "
-            f"checkpoints, got {len(calls)}. The 250-op limit must not split empties."
+            f"Expected 1 coalesced API call for {branch_count} due refreshes, "
+            f"got {len(calls)}."
         )
-        assert calls[0] == [], "Empty checkpoints should produce an empty updates list"
+        assert calls[0] == [], "Refreshes should produce an empty updates list"
     finally:
         state.stop_checkpointing()
         batcher.shutdown(wait=True)
diff --git a/packages/aws-durable-execution-sdk-python/tests/state_test.py b/packages/aws-durable-execution-sdk-python/tests/state_test.py
index 92b07f58..713f63fb 100644
--- a/packages/aws-durable-execution-sdk-python/tests/state_test.py
+++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py
@@ -2,9 +2,12 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
+
 import contextlib
 import datetime
 import json
+import queue
 import threading
 import time
 import unittest.mock
@@ -51,6 +54,9 @@
     UserFunctionOutcome,
 )
 from aws_durable_execution_sdk_python.state import (
+    _BLOCKED_CALLER_WAIT_SECONDS,
+    _STOP_SIGNAL_POLL_SECONDS,
+    _Signal,
     CheckpointBatcherConfig,
     CheckpointedResult,
     ExecutionState,
@@ -1665,6 +1671,39 @@ def test_nested_parallel_operations_deep_hierarchy():
 
 
 # Test 8.4: Thread safety and synchronous operations
+def _assert_sync_call_blocks_until_event(
+    state: ExecutionState, call: Callable[[], None]
+) -> None:
+    """Assert that ``call`` stays inside create_checkpoint until its event is set.
+
+    No sleeps and no clock. The caller is observed queued and not returned,
+    then collected and still not returned, then released by its event.
+    """
+    returned = threading.Event()
+
+    def call_checkpoint() -> None:
+        call()
+        returned.set()
+
+    caller = threading.Thread(daemon=True, target=call_checkpoint)
+    caller.start()
+
+    deadline = time.monotonic() + 10.0
+    while state._checkpoint_queue.qsize() < 1:
+        assert time.monotonic() < deadline, "the operation was never enqueued"
+        time.sleep(0.001)
+    assert not returned.is_set(), "the caller returned while still queued"
+
+    batch = state._collect_checkpoint_batch()
+    assert len(batch) == 1
+    assert not returned.is_set(), "the caller returned before its event was set"
+
+    assert batch[0].completion_event is not None
+    batch[0].completion_event.set()
+    caller.join(timeout=10.0)
+    assert returned.is_set(), "the caller never returned"
+
+
 def test_synchronous_checkpoint_blocks_until_complete():
     """Test that create_checkpoint_sync blocks until checkpoint is processed."""
     mock_lambda_client = Mock(spec=LambdaClient)
@@ -1690,35 +1729,9 @@ def test_synchronous_checkpoint_blocks_until_complete():
         action=OperationAction.START,
     )
 
-    # Track if operation completed
-    completed = threading.Event()
-
-    def background_processor():
-        """Simulate background processing."""
-        time.sleep(0.1)  # Small delay
-        batch = state._collect_checkpoint_batch()
-        if batch:
-            # Signal completion events
-            for queued_op in batch:
-                if queued_op.completion_event:
-                    queued_op.completion_event.set()
-        completed.set()
-
-    # Start background processor
-    processor_thread = threading.Thread(daemon=True, target=background_processor)
-    processor_thread.start()
-
-    # Call synchronous checkpoint (should block)
-    start_time = time.time()
-    state.create_checkpoint(operation_update, is_sync=True)
-    elapsed = time.time() - start_time
-
-    # Verify it blocked for at least the delay time
-    assert elapsed >= 0.1
-
-    # Wait for background thread
-    processor_thread.join(timeout=1.0)
-    assert completed.is_set()
+    _assert_sync_call_blocks_until_event(
+        state, lambda: state.create_checkpoint(operation_update, is_sync=True)
+    )
 
 
 def test_concurrent_access_to_operations_dictionary():
@@ -3500,11 +3513,7 @@ def test_collect_checkpoint_batch_overflow_queue_size_limit_final():
 
 
 def test_create_checkpoint_blocks_until_completion_default():
-    """Test that create_checkpoint() blocks until completion when is_sync=True (default).
-
-    Verifies that calling create_checkpoint without specifying is_sync results in
-    synchronous blocking behavior until the background thread processes the checkpoint.
-    """
+    """create_checkpoint() blocks until completion when is_sync is left at its default."""
     mock_lambda_client = Mock(spec=LambdaClient)
     mock_lambda_client.checkpoint.return_value = CheckpointOutput(
         checkpoint_token="new_token",  # noqa: S106
@@ -3528,55 +3537,13 @@ def test_create_checkpoint_blocks_until_completion_default():
         action=OperationAction.START,
     )
 
-    # Track timing and completion
-    call_completed = threading.Event()
-    start_time = None
-    end_time = None
-
-    def call_checkpoint():
-        nonlocal start_time, end_time
-        start_time = time.time()
-        # Call without is_sync parameter (defaults to True)
-        state.create_checkpoint(operation_update)
-        end_time = time.time()
-        call_completed.set()
-
-    def background_processor():
-        """Simulate background processing with delay."""
-        time.sleep(0.15)  # Delay to verify blocking
-        batch = state._collect_checkpoint_batch()
-        if batch:
-            # Signal completion events
-            for queued_op in batch:
-                if queued_op.completion_event:
-                    queued_op.completion_event.set()
-
-    # Start background processor
-    processor_thread = threading.Thread(daemon=True, target=background_processor)
-    processor_thread.start()
-
-    # Start checkpoint call
-    caller_thread = threading.Thread(daemon=True, target=call_checkpoint)
-    caller_thread.start()
-
-    # Wait for both threads
-    caller_thread.join(timeout=2.0)
-    processor_thread.join(timeout=1.0)
-
-    # Verify call completed
-    assert call_completed.is_set()
-
-    # Verify it blocked for at least the delay time
-    elapsed = end_time - start_time
-    assert elapsed >= 0.15, f"Expected blocking for at least 0.15s, got {elapsed}s"
+    _assert_sync_call_blocks_until_event(
+        state, lambda: state.create_checkpoint(operation_update)
+    )
 
 
 def test_create_checkpoint_blocks_until_completion_explicit_true():
-    """Test that create_checkpoint(is_sync=True) blocks until completion.
-
-    Verifies that explicitly setting is_sync=True results in synchronous blocking
-    behavior until the background thread processes the checkpoint.
-    """
+    """create_checkpoint(is_sync=True) blocks until completion."""
     mock_lambda_client = Mock(spec=LambdaClient)
     mock_lambda_client.checkpoint.return_value = CheckpointOutput(
         checkpoint_token="new_token",  # noqa: S106
@@ -3600,47 +3567,9 @@ def test_create_checkpoint_blocks_until_completion_explicit_true():
         action=OperationAction.START,
     )
 
-    # Track timing and completion
-    call_completed = threading.Event()
-    start_time = None
-    end_time = None
-
-    def call_checkpoint():
-        nonlocal start_time, end_time
-        start_time = time.time()
-        # Call with explicit is_sync=True
-        state.create_checkpoint(operation_update, is_sync=True)
-        end_time = time.time()
-        call_completed.set()
-
-    def background_processor():
-        """Simulate background processing with delay."""
-        time.sleep(0.15)  # Delay to verify blocking
-        batch = state._collect_checkpoint_batch()
-        if batch:
-            # Signal completion events
-            for queued_op in batch:
-                if queued_op.completion_event:
-                    queued_op.completion_event.set()
-
-    # Start background processor
-    processor_thread = threading.Thread(daemon=True, target=background_processor)
-    processor_thread.start()
-
-    # Start checkpoint call
-    caller_thread = threading.Thread(daemon=True, target=call_checkpoint)
-    caller_thread.start()
-
-    # Wait for both threads
-    caller_thread.join(timeout=2.0)
-    processor_thread.join(timeout=1.0)
-
-    # Verify call completed
-    assert call_completed.is_set()
-
-    # Verify it blocked for at least the delay time
-    elapsed = end_time - start_time
-    assert elapsed >= 0.15, f"Expected blocking for at least 0.15s, got {elapsed}s"
+    _assert_sync_call_blocks_until_event(
+        state, lambda: state.create_checkpoint(operation_update, is_sync=True)
+    )
 
 
 def test_create_checkpoint_completion_event_created_and_signaled():
@@ -3879,17 +3808,12 @@ def run_background_processor():
 
 
 def test_create_checkpoint_multiple_sync_calls_all_block():
-    """Test that multiple synchronous checkpoint calls all block correctly.
-
-    Verifies that when multiple threads call create_checkpoint synchronously,
-    they all block until their respective completion events are signaled.
-    """
+    """No sync caller returns until its batch is collected and its event is set."""
     mock_lambda_client = Mock(spec=LambdaClient)
     mock_lambda_client.checkpoint.return_value = CheckpointOutput(
         checkpoint_token="new_token",  # noqa: S106
         new_execution_state=CheckpointUpdatedExecutionState(
-            operations=[],
-            next_marker=None,
+            operations=[], next_marker=None
         ),
     )
 
@@ -3902,65 +3826,226 @@ def test_create_checkpoint_multiple_sync_calls_all_block():
     )
 
     num_callers = 3
-    completion_events = [CompletionEvent() for _ in range(num_callers)]
-    start_times = [None] * num_callers
-    end_times = [None] * num_callers
+    returned = [False] * num_callers
 
-    def call_checkpoint(index):
-        """Call synchronous checkpoint."""
-        operation_update = OperationUpdate(
-            operation_id=f"test_op_{index}",
-            operation_type=OperationType.STEP,
-            action=OperationAction.START,
+    def call_checkpoint(index: int) -> None:
+        state.create_checkpoint(
+            OperationUpdate(
+                operation_id=f"test_op_{index}",
+                operation_type=OperationType.STEP,
+                action=OperationAction.START,
+            ),
+            is_sync=True,
         )
-        start_times[index] = time.time()
-        state.create_checkpoint(operation_update, is_sync=True)
-        end_times[index] = time.time()
-        completion_events[index].set()
+        returned[index] = True
 
-    def background_processor():
-        """Process all checkpoints with delay."""
-        time.sleep(0.15)  # Delay to verify blocking
-        batch = state._collect_checkpoint_batch()
-        if batch:
-            # Signal all completion events
-            for queued_op in batch:
-                if queued_op.completion_event:
-                    queued_op.completion_event.set()
+    caller_threads = [
+        threading.Thread(daemon=True, target=call_checkpoint, args=(i,))
+        for i in range(num_callers)
+    ]
+    for thread in caller_threads:
+        thread.start()
 
-    # Start background processor
-    processor_thread = threading.Thread(daemon=True, target=background_processor)
-    processor_thread.start()
+    # Wait until every caller has enqueued its operation, not for a fixed time.
+    # The loop is bounded so a genuine hang fails instead of blocking.
+    deadline = time.monotonic() + 10.0
+    while state._checkpoint_queue.qsize() < num_callers:
+        assert time.monotonic() < deadline, (
+            f"only {state._checkpoint_queue.qsize()} of {num_callers} operations "
+            f"were enqueued"
+        )
+        time.sleep(0.001)
 
-    # Start multiple caller threads
-    caller_threads = []
-    for i in range(num_callers):
-        thread = threading.Thread(daemon=True, target=call_checkpoint, args=(i,))
-        thread.start()
-        caller_threads.append(thread)
+    # Every caller is queued and none has returned. So each is blocked inside
+    # create_checkpoint, which is what is_sync=True promises.
+    assert returned == [False] * num_callers, (
+        f"a caller returned while still queued: {returned}"
+    )
+
+    batch = state._collect_checkpoint_batch()
+
+    # One batch carries all three, so the collector drains queued work rather
+    # than flushing one operation at a time.
+    assert len(batch) == num_callers
+
+    # Collection alone does not release a caller. Only the completion event does.
+    assert returned == [False] * num_callers, (
+        f"a caller returned before its event was set: {returned}"
+    )
+
+    for queued_op in batch:
+        assert queued_op.completion_event is not None
+        queued_op.completion_event.set()
 
-    # Wait for all threads
     for thread in caller_threads:
-        thread.join(timeout=2.0)
-    processor_thread.join(timeout=1.0)
+        thread.join(timeout=10.0)
 
-    # Verify all calls completed
-    for i, event in enumerate(completion_events):
-        assert event.is_set(), f"Caller {i} did not complete"
+    assert returned == [True] * num_callers, f"a caller never returned: {returned}"
 
-    # Verify all calls blocked for at least the delay time
-    for i in range(num_callers):
-        elapsed = end_times[i] - start_times[i]
-        assert elapsed >= 0.15, (
-            f"Caller {i} expected blocking for at least 0.15s, got {elapsed}s"
-        )
 
+class _RecordingQueue(queue.Queue):
+    """Records each timed get's timeout, so tests assert the policy without a clock."""
 
-def test_create_checkpoint_sync_with_empty_checkpoint():
-    """Test synchronous behavior with empty checkpoint (None operation_update).
+    def __init__(self) -> None:
+        super().__init__()
+        self.timeouts: list[float] = []
+        self.expired_timeouts: list[float] = []
 
-    Verifies that empty checkpoints also block correctly when is_sync=True.
-    """
+    on_get = None  # optional callable run inside each timed get, for boundary tests
+
+    def get(self, block=True, timeout=None):  # noqa: FBT002
+        if timeout is not None:
+            self.timeouts.append(timeout)
+            if self.on_get is not None:
+                self.on_get()
+        try:
+            return super().get(block, timeout)
+        except queue.Empty:
+            if timeout is not None:
+                self.expired_timeouts.append(timeout)
+            raise
+
+
+def _batching_state(max_ops: int = 250) -> ExecutionState:
+    return ExecutionState(
+        durable_execution_arn="test_arn",
+        initial_checkpoint_token="token123",  # noqa: S106
+        operations={},
+        service_client=Mock(),
+        plugin_executor=PluginExecutor(plugins=None),
+        batcher_config=CheckpointBatcherConfig(max_batch_operations=max_ops),
+    )
+
+
+def _step_update(action, op_id: str = "step"):
+    return OperationUpdate(
+        operation_id=op_id, operation_type=OperationType.STEP, action=action
+    )
+
+
+def _record_queue_reads(state) -> _RecordingQueue:
+    recording = _RecordingQueue()
+    state._checkpoint_queue = recording
+    return recording
+
+
+def _assert_short_wait(recording: _RecordingQueue) -> None:
+    """Assert the collector stopped waiting after the blocked-caller interval."""
+    assert recording.expired_timeouts, "collector never reached an empty queue"
+    assert max(recording.expired_timeouts) <= _BLOCKED_CALLER_WAIT_SECONDS, (
+        f"collector waited {max(recording.expired_timeouts)}s on an empty queue "
+        f"while a caller was blocked; expected at most "
+        f"{_BLOCKED_CALLER_WAIT_SECONDS}s"
+    )
+
+
+def _assert_full_wait(recording: _RecordingQueue) -> None:
+    """Assert the collector kept the full batching window."""
+    assert recording.expired_timeouts == [_STOP_SIGNAL_POLL_SECONDS], (
+        f"expected one expired wait of {_STOP_SIGNAL_POLL_SECONDS}s, got "
+        f"{recording.expired_timeouts}"
+    )
+
+
+def test_collect_batch_shortens_wait_once_it_holds_a_sync_checkpoint():
+    """A caller blocks until the batch persists, so the collector must not hold it."""
+    state = _batching_state()
+    recording = _record_queue_reads(state)
+
+    waiter = CompletionEvent()
+    recording.put(QueuedOperation(_step_update(OperationAction.SUCCEED), waiter))
+
+    batch = state._collect_checkpoint_batch()
+
+    assert [q.operation_update.action for q in batch] == [OperationAction.SUCCEED]
+    assert batch[0].completion_event is waiter
+    assert not waiter.is_set(), "the collector must not settle the waiter"
+    _assert_short_wait(recording)
+
+
+def test_collect_batch_coalesces_async_start_with_sync_succeed():
+    """START and SUCCEED must share one request, or the call count doubles."""
+    state = _batching_state()
+    recording = _record_queue_reads(state)
+
+    waiter = CompletionEvent()
+    recording.put(QueuedOperation(_step_update(OperationAction.START), None))
+    recording.put(QueuedOperation(_step_update(OperationAction.SUCCEED), waiter))
+
+    batch = state._collect_checkpoint_batch()
+
+    assert [q.operation_update.action for q in batch] == [
+        OperationAction.START,
+        OperationAction.SUCCEED,
+    ]
+    assert batch[-1].completion_event is waiter
+    _assert_short_wait(recording)
+
+
+def test_collect_batch_drains_queued_work_before_flushing():
+    """Everything already queued still joins the batch. Only future work is skipped."""
+    state = _batching_state()
+    recording = _record_queue_reads(state)
+
+    waiter = CompletionEvent()
+    recording.put(QueuedOperation(_step_update(OperationAction.SUCCEED, "a"), waiter))
+    for op_id in ("b", "c", "d"):
+        recording.put(QueuedOperation(_step_update(OperationAction.START, op_id), None))
+
+    batch = state._collect_checkpoint_batch()
+
+    assert [q.operation_update.operation_id for q in batch] == ["a", "b", "c", "d"]
+    _assert_short_wait(recording)
+
+
+def test_collect_batch_shortens_wait_for_sync_checkpoint_from_overflow_queue():
+    """A size-limited batch defers to overflow, so that path must set the flag too."""
+    state = _batching_state()
+    recording = _record_queue_reads(state)
+
+    waiter = CompletionEvent()
+    state._overflow_queue.put(
+        QueuedOperation(_step_update(OperationAction.SUCCEED), waiter)
+    )
+
+    batch = state._collect_checkpoint_batch()
+
+    assert len(batch) == 1
+    assert batch[0].completion_event is waiter
+    _assert_short_wait(recording)
+
+
+def test_collect_batch_shortens_wait_for_sync_empty_checkpoint():
+    """An empty sync checkpoint blocks a caller too, so it gets the short wait."""
+    state = _batching_state()
+    recording = _record_queue_reads(state)
+
+    waiter = CompletionEvent()
+    recording.put(QueuedOperation(None, waiter))
+
+    batch = state._collect_checkpoint_batch()
+
+    assert len(batch) == 1
+    assert batch[0].operation_update is None
+    assert batch[0].completion_event is waiter
+    _assert_short_wait(recording)
+
+
+def test_collect_batch_keeps_full_window_for_async_only_batch():
+    """No caller blocks, so holding the window costs nothing and cuts API calls."""
+    state = _batching_state()
+    recording = _record_queue_reads(state)
+
+    recording.put(QueuedOperation(_step_update(OperationAction.START), None))
+
+    batch = state._collect_checkpoint_batch()
+
+    assert [q.operation_update.action for q in batch] == [OperationAction.START]
+    _assert_full_wait(recording)
+
+
+def test_create_checkpoint_sync_with_empty_checkpoint():
+    """An empty checkpoint (None operation_update) also blocks when is_sync=True."""
     mock_lambda_client = Mock(spec=LambdaClient)
     mock_lambda_client.checkpoint.return_value = CheckpointOutput(
         checkpoint_token="new_token",  # noqa: S106
@@ -3978,47 +4063,9 @@ def test_create_checkpoint_sync_with_empty_checkpoint():
         plugin_executor=PluginExecutor(plugins=None),
     )
 
-    # Track timing and completion
-    call_completed = threading.Event()
-    start_time = None
-    end_time = None
-
-    def call_checkpoint():
-        nonlocal start_time, end_time
-        start_time = time.time()
-        # Call with None (empty checkpoint) and is_sync=True
-        state.create_checkpoint(None, is_sync=True)
-        end_time = time.time()
-        call_completed.set()
-
-    def background_processor():
-        """Simulate background processing with delay."""
-        time.sleep(0.15)  # Delay to verify blocking
-        batch = state._collect_checkpoint_batch()
-        if batch:
-            # Signal completion events
-            for queued_op in batch:
-                if queued_op.completion_event:
-                    queued_op.completion_event.set()
-
-    # Start background processor
-    processor_thread = threading.Thread(daemon=True, target=background_processor)
-    processor_thread.start()
-
-    # Start checkpoint call
-    caller_thread = threading.Thread(daemon=True, target=call_checkpoint)
-    caller_thread.start()
-
-    # Wait for both threads
-    caller_thread.join(timeout=2.0)
-    processor_thread.join(timeout=1.0)
-
-    # Verify call completed
-    assert call_completed.is_set()
-
-    # Verify it blocked for at least the delay time
-    elapsed = end_time - start_time
-    assert elapsed >= 0.15, f"Expected blocking for at least 0.15s, got {elapsed}s"
+    _assert_sync_call_blocks_until_event(
+        state, lambda: state.create_checkpoint(None, is_sync=True)
+    )
 
 
 def test_create_checkpoint_sync_success():
@@ -4338,8 +4385,8 @@ def test_checkpoint_batches_forever_single_api_call_for_many_empty_checkpoints()
 def test_collect_checkpoint_batch_first_empty_counts_toward_limit():
     """Test that only the first empty checkpoint counts toward the batch operation limit.
 
-    With limit=2: an empty op (effective=1) + a real op (effective=2) exactly fills the
-    batch. The loop exits after the limit is hit; items after the limit stay in the queue.
+    With limit=2, an empty op (effective=1) and a real op (effective=2) fill the
+    batch. Items after the limit go in the next batch.
     """
     mock_lambda_client = Mock(spec=LambdaClient)
 
@@ -4387,11 +4434,17 @@ def test_collect_checkpoint_batch_first_empty_counts_toward_limit():
     # The batch contains exactly: 1 leading empty + op_1 (limit=2 effective ops)
     assert len(real_in_batch) == 1
     assert real_in_batch[0].operation_update.operation_id == "op_1"
-    assert (
-        len(empty_in_batch) == 1
-    )  # Only the leading empty; trailing deferred to next batch
-    # op_2 and trailing empties remain in the queue
-    assert state._checkpoint_queue.qsize() == 51
+    assert len(empty_in_batch) == 1  # only the leading empty
+
+    # The rest follows in order. op_2 and the first trailing empty fill the
+    # next batch. The remaining empties share one batch, since only the first
+    # empty in a batch counts.
+    second = state._collect_checkpoint_batch()
+    assert second[0].operation_update.operation_id == "op_2"
+    assert len(second) == 2
+    third = state._collect_checkpoint_batch()
+    assert len(third) == 49
+    assert all(q.operation_update is None for q in third)
 
 
 def test_execution_state_get_execution_operation_no_operations():
@@ -5239,6 +5292,492 @@ def on_operation_start(self, info):
 # endregion Plugin Executor Integration Tests
 
 
+# --- Refresh scheduling. An empty checkpoint sent no earlier than a clock time ---
+
+
+def _refresh_client(calls: list[list]) -> Mock:
+    """A client that records each request's updates and answers with a fresh token."""
+    client = Mock(spec=LambdaClient)
+
+    def _checkpoint(
+        durable_execution_arn, checkpoint_token, updates, client_token=None
+    ):
+        calls.append(list(updates))
+        return CheckpointOutput(
+            checkpoint_token=f"token_{len(calls)}",
+            new_execution_state=CheckpointUpdatedExecutionState(),
+        )
+
+    client.checkpoint = _checkpoint
+    return client
+
+
+def _refresh_state(client=None) -> ExecutionState:
+    return ExecutionState(
+        durable_execution_arn="test_arn",
+        initial_checkpoint_token="token123",  # noqa: S106
+        operations={},
+        service_client=client if client is not None else Mock(),
+        plugin_executor=PluginExecutor(plugins=None),
+    )
+
+
+def _pending(state) -> list:
+    return [entry[2] for entry in state._pending_refreshes]
+
+
+def _queued_operations(state) -> list:
+    items = list(state._checkpoint_queue.queue)
+    return [item for item in items if isinstance(item, QueuedOperation)]
+
+
+def test_schedule_refresh_enqueues_and_returns_without_blocking():
+    state = _refresh_state()
+    check_time = time.time() + 60
+
+    started = time.monotonic()
+    refresh = state.schedule_refresh(check_time)
+
+    assert time.monotonic() - started < 0.1
+    assert not refresh.is_set()
+    assert _pending(state) == [refresh]
+    assert list(state._checkpoint_queue.queue) == [_Signal.REFRESH_WAKE]
+
+
+def test_many_refreshes_make_one_wake_signal():
+    state = _refresh_state()
+    for i in range(100):
+        state.schedule_refresh(time.time() + 60 + i)
+    assert list(state._checkpoint_queue.queue) == [_Signal.REFRESH_WAKE]
+
+    state._checkpoint_queue.get_nowait()
+    state._consume_refresh_wake()
+    state.schedule_refresh(time.time() + 200)
+    assert list(state._checkpoint_queue.queue) == [_Signal.REFRESH_WAKE]
+
+
+def test_collect_batch_defers_refresh_whose_time_has_not_come():
+    """An early refresh must not travel with unrelated work. Its response could
+    not show the wait complete, so sending it early wastes the request."""
+    state = _refresh_state()
+    recording = _record_queue_reads(state)
+    refresh = state.schedule_refresh(time.time() + 60)
+    recording.put(
+        QueuedOperation(_step_update(OperationAction.SUCCEED), CompletionEvent())
+    )
+
+    batch = state._collect_checkpoint_batch()
+
+    assert [q.operation_update.action for q in batch] == [OperationAction.SUCCEED]
+    assert _pending(state) == [refresh]
+    assert not refresh.is_set()
+
+
+def test_collect_batch_sends_due_refresh_with_short_wait():
+    """A due refresh has a blocked coordinator. It is treated like any sync op."""
+    state = _refresh_state()
+    recording = _record_queue_reads(state)
+    refresh = state.schedule_refresh(time.time() - 1)
+
+    batch = state._collect_checkpoint_batch()
+
+    assert [q.completion_event for q in batch] == [refresh.completion_event]
+    assert _pending(state) == []
+    _assert_short_wait(recording)
+
+
+def test_collect_batch_wakes_for_refresh_before_poll_interval():
+    """With nothing else queued, the collector sleeps only until the refresh is
+    due, not for the full stop-signal poll interval."""
+    state = _refresh_state()
+    recording = _record_queue_reads(state)
+    refresh = state.schedule_refresh(time.time() + 0.03)
+
+    started = time.monotonic()
+    batch = state._collect_checkpoint_batch()
+    elapsed = time.monotonic() - started
+
+    assert [q.completion_event for q in batch] == [refresh.completion_event]
+    assert elapsed < _STOP_SIGNAL_POLL_SECONDS
+    assert all(timeout <= 0.031 for timeout in recording.timeouts)
+
+
+def test_collect_batch_groups_refreshes_by_check_time_not_arrival():
+    """Three refreshes with one check time go in one batch. Two with other
+    times stay pending. Arrival order plays no part."""
+    state = _refresh_state()
+    _record_queue_reads(state)
+    check_time = time.time() + 0.02
+    later = state.schedule_refresh(check_time + 60)
+    first = state.schedule_refresh(check_time)
+    much_later = state.schedule_refresh(check_time + 120)
+    second = state.schedule_refresh(check_time)
+    third = state.schedule_refresh(check_time)
+
+    batch = state._collect_checkpoint_batch()
+
+    assert [q.completion_event for q in batch] == [
+        r.completion_event for r in (first, second, third)
+    ]
+    assert _pending(state) == [later, much_later]
+
+
+def test_collect_batch_groups_refreshes_due_within_one_short_wait():
+    """Refreshes whose check times differ by less than the blocked-caller wait
+    must share one batch. A read cut short by the next refresh coming due is
+    not an expired idle window, so the collector keeps collecting."""
+    state = _refresh_state()
+    _record_queue_reads(state)
+    base = time.time() + 0.02
+    refreshes = [state.schedule_refresh(base + i * 0.0008) for i in range(5)]
+
+    batch = state._collect_checkpoint_batch()
+
+    assert [q.completion_event for q in batch] == [
+        r.completion_event for r in refreshes
+    ]
+    assert _pending(state) == []
+
+
+def test_collect_batch_adds_refresh_that_comes_due_during_async_window():
+    """An async-only batch keeps its window, but a refresh coming due inside
+    that window closes it and joins the batch."""
+    state = _refresh_state()
+    recording = _record_queue_reads(state)
+    recording.put(QueuedOperation(_step_update(OperationAction.START), None))
+    refresh = state.schedule_refresh(time.time() + 0.02)
+
+    started = time.monotonic()
+    batch = state._collect_checkpoint_batch()
+    elapsed = time.monotonic() - started
+
+    assert [q.operation_update is None for q in batch] == [False, True]
+    assert batch[1].completion_event is refresh.completion_event
+    assert elapsed < _STOP_SIGNAL_POLL_SECONDS
+
+
+def test_refreshes_with_one_check_time_share_one_request_despite_spacing():
+    """Refreshes requested 5 ms apart, wider than the blocked-caller wait,
+    still make one request because they share a check time."""
+    calls: list[list] = []
+    state = _refresh_state(_refresh_client(calls))
+    batcher = ThreadPoolExecutor(max_workers=1)
+    batcher.submit(state.checkpoint_batches_forever)
+    try:
+        check_time = time.time() + 0.4
+        events = []
+        for _ in range(5):
+            events.append(state.schedule_refresh(check_time))
+            time.sleep(0.005)
+        for event in events:
+            assert event.wait(timeout=5)
+        assert calls == [[]], f"expected one empty request, got {calls}"
+    finally:
+        state.stop_checkpointing()
+        batcher.shutdown(wait=True)
+
+
+def test_refresh_is_not_sent_with_an_earlier_unrelated_request():
+    """A step's checkpoint flushes at once. The pending refresh stays behind
+    and is sent later, alone, at its own time."""
+    calls: list[list] = []
+    state = _refresh_state(_refresh_client(calls))
+    batcher = ThreadPoolExecutor(max_workers=1)
+    batcher.submit(state.checkpoint_batches_forever)
+    try:
+        refresh = state.schedule_refresh(time.time() + 0.3)
+        state._checkpoint_queue.join()  # the collector has deferred it
+        state.create_checkpoint(_step_update(OperationAction.SUCCEED))
+        assert len(calls) == 1
+        assert [u.action for u in calls[0]] == [OperationAction.SUCCEED]
+        assert not refresh.is_set()
+        assert refresh.wait(timeout=5)
+        assert calls[1] == []
+        assert len(calls) == 2
+    finally:
+        state.stop_checkpointing()
+        batcher.shutdown(wait=True)
+
+
+def test_batch_failure_settles_pending_refresh():
+    """A coordinator waiting on a deferred refresh must see the failure, not
+    block until its check time."""
+    client = Mock(spec=LambdaClient)
+    client.checkpoint.side_effect = RuntimeError("service down")
+    state = _refresh_state(client)
+    batcher = ThreadPoolExecutor(max_workers=1)
+    batcher.submit(state.checkpoint_batches_forever)
+    try:
+        refreshes = [state.schedule_refresh(time.time() + 60 * n) for n in (1, 2)]
+        state._checkpoint_queue.join()
+        with pytest.raises(BackgroundThreadError):
+            state.create_checkpoint(_step_update(OperationAction.SUCCEED))
+        for refresh in refreshes:
+            with pytest.raises(BackgroundThreadError):
+                refresh.wait(timeout=5)
+        with pytest.raises(BackgroundThreadError):
+            state.schedule_refresh(time.time() + 60)
+    finally:
+        state.stop_checkpointing()
+        batcher.shutdown(wait=True)
+
+
+def test_execution_completion_settles_pending_refresh():
+    """When the service ends the execution, a deferred refresh can never be
+    sent. Its waiter is settled as orphaned."""
+    client = Mock(spec=LambdaClient)
+    client.checkpoint.return_value = CheckpointOutput(
+        checkpoint_token="",
+        new_execution_state=CheckpointUpdatedExecutionState(),
+    )
+    state = _refresh_state(client)
+    batcher = ThreadPoolExecutor(max_workers=1)
+    batcher.submit(state.checkpoint_batches_forever)
+    try:
+        refreshes = [state.schedule_refresh(time.time() + 60 * n) for n in (1, 2)]
+        state._checkpoint_queue.join()
+        state.create_checkpoint(_step_update(OperationAction.SUCCEED))
+        for refresh in refreshes:
+            with pytest.raises(OrphanedChildException):
+                refresh.wait(timeout=5)
+    finally:
+        state.stop_checkpointing()
+        batcher.shutdown(wait=True)
+
+
+def test_stop_settles_every_pending_refresh():
+    """Refreshes still deferred or still queued at shutdown are settled, so no
+    thread can wait on an event that is never set."""
+    state = _refresh_state(_refresh_client([]))
+    batcher = ThreadPoolExecutor(max_workers=1)
+    batcher.submit(state.checkpoint_batches_forever)
+    deferred = state.schedule_refresh(time.time() + 60)
+    state._checkpoint_queue.join()  # deferred to the heap
+    still_queued = state.schedule_refresh(time.time() + 120)  # may not be dequeued
+
+    state.stop_checkpointing()
+    batcher.shutdown(wait=True)
+
+    for refresh in (deferred, still_queued):
+        with pytest.raises(OrphanedChildException):
+            refresh.wait(timeout=5)
+    with pytest.raises(OrphanedChildException):
+        state.schedule_refresh(time.time() + 60)
+
+
+def test_stop_settles_pending_refresh_and_keeps_queued_items():
+    """At shutdown a pending refresh is settled. Any queued operation stays in
+    the queue."""
+    state = _refresh_state(_refresh_client([]))
+    refresh = state.schedule_refresh(time.time() + 60)
+    async_op = QueuedOperation(_step_update(OperationAction.START), None)
+    state._checkpoint_queue.put(async_op)
+    state.stop_checkpointing()
+
+    state.checkpoint_batches_forever()  # exits at once, then settles
+
+    with pytest.raises(OrphanedChildException):
+        refresh.wait(timeout=1)
+    assert _queued_operations(state) == [async_op]
+
+
+def test_cancelled_refresh_deferred_then_due_is_dropped():
+    """A refresh cancelled while pending is dropped when its time comes, and
+    the batch goes on without it."""
+    state = _refresh_state()
+    recording = _record_queue_reads(state)
+    handle = state.schedule_refresh(time.time() + 0.02)
+    recording.put(
+        QueuedOperation(_step_update(OperationAction.SUCCEED), CompletionEvent())
+    )
+    first = state._collect_checkpoint_batch()
+    assert [q.operation_update.action for q in first] == [OperationAction.SUCCEED]
+    assert _pending(state) == [handle]
+
+    handle.cancel()
+    time.sleep(0.03)
+    recording.put(
+        QueuedOperation(_step_update(OperationAction.SUCCEED), CompletionEvent())
+    )
+    second = state._collect_checkpoint_batch()
+
+    assert [q.operation_update.action for q in second] == [OperationAction.SUCCEED]
+    assert _pending(state) == []
+    assert handle.is_set()
+
+
+def test_collect_batch_opens_with_refresh_that_came_due_between_batches():
+    """A pending refresh whose time passed while the previous batch was in
+    flight opens the next batch, ahead of the first queue read."""
+    state = _refresh_state()
+    recording = _record_queue_reads(state)
+    refresh = state.schedule_refresh(time.time() + 0.02)
+    recording.put(
+        QueuedOperation(_step_update(OperationAction.SUCCEED), CompletionEvent())
+    )
+    state._collect_checkpoint_batch()
+    assert _pending(state) == [refresh]
+
+    time.sleep(0.03)
+    recording.put(
+        QueuedOperation(_step_update(OperationAction.SUCCEED), CompletionEvent())
+    )
+    batch = state._collect_checkpoint_batch()
+
+    assert batch[0].completion_event is refresh.completion_event
+    assert [q.operation_update is None for q in batch] == [True, False]
+
+
+def test_collect_batch_counts_first_plain_empty_checkpoint_read_after_other_work():
+    """A plain empty checkpoint, with no check time, read after another item
+    counts once toward the operation limit like any first empty."""
+    state = _refresh_state()
+    recording = _record_queue_reads(state)
+    recording.put(
+        QueuedOperation(_step_update(OperationAction.SUCCEED), CompletionEvent())
+    )
+    empties = [QueuedOperation(None, CompletionEvent()) for _ in range(3)]
+    for op in empties:
+        recording.put(op)
+
+    batch = state._collect_checkpoint_batch()
+
+    assert batch[1:] == empties
+    assert len(batch) == 4
+
+
+def test_cancelled_refreshes_do_not_delay_a_batch():
+    """A hundred pending refreshes, cancelled, coming due 0.8 ms apart during a
+    sync batch. They must not hold the batch open. Each would otherwise cost
+    one more read."""
+    state = _refresh_state()
+    recording = _record_queue_reads(state)
+    now = time.time()
+    handles = [state.schedule_refresh(now + 0.0005 + i * 0.0008) for i in range(100)]
+    for handle in handles:
+        handle.cancel()
+    recording.put(
+        QueuedOperation(_step_update(OperationAction.SUCCEED), CompletionEvent())
+    )
+
+    started = time.monotonic()
+    batch = state._collect_checkpoint_batch()
+    elapsed = time.monotonic() - started
+
+    assert [q.operation_update for q in batch if q.operation_update] == [
+        batch[-1].operation_update
+    ]
+    assert elapsed < 0.03, f"batch held open for {elapsed * 1000:.0f} ms"
+    assert _pending(state) == []
+    assert all(handle.is_set() for handle in handles)
+
+
+def test_same_time_refreshes_share_a_batch_when_the_operation_limit_closes_it():
+    """Two refreshes for one past time, the second scheduled while the batch is
+    being collected, with room for two operations. Both go in this batch. The
+    step that did not fit goes in the next."""
+    state = _batching_state(max_ops=2)
+    recording = _record_queue_reads(state)
+    check_time = time.time() - 1
+    first = state.schedule_refresh(check_time)
+    assert (
+        recording.get_nowait() is _Signal.REFRESH_WAKE
+    )  # consumed, as a prior read would
+    state._consume_refresh_wake()
+    recording.put(
+        QueuedOperation(_step_update(OperationAction.SUCCEED, "a"), CompletionEvent())
+    )
+    recording.put(
+        QueuedOperation(_step_update(OperationAction.SUCCEED, "b"), CompletionEvent())
+    )
+    second: list = []
+
+    def schedule_second_during_collection():
+        if not second:
+            second.append(state.schedule_refresh(check_time))
+
+    recording.on_get = schedule_second_during_collection
+
+    batch = state._collect_checkpoint_batch()
+
+    assert [q.completion_event for q in batch if q.operation_update is None] == [
+        first.completion_event,
+        second[0].completion_event,
+    ]
+    assert [q.operation_update.operation_id for q in batch if q.operation_update] == [
+        "a"
+    ]
+    recording.on_get = None
+    next_batch = state._collect_checkpoint_batch()
+    assert [
+        q.operation_update.operation_id for q in next_batch if q.operation_update
+    ] == ["b"]
+
+
+def test_collect_batch_reads_no_more_than_one_batch_from_a_backlog():
+    """Sealing a batch must not depend on the size of the queue behind it."""
+    state = _batching_state(max_ops=250)
+    for i in range(5000):
+        state._checkpoint_queue.put(
+            QueuedOperation(_step_update(OperationAction.START, f"op{i}"), None)
+        )
+
+    batch = state._collect_checkpoint_batch()
+
+    assert len(batch) == 250
+    assert state._checkpoint_queue.qsize() == 4750
+    assert state._overflow_queue.empty()
+
+
+def test_cancelled_refresh_is_not_sent():
+    """A coordinator that stops needing a refresh cancels it. No request is made
+    when its time comes."""
+    calls: list[list] = []
+    state = _refresh_state(_refresh_client(calls))
+    batcher = ThreadPoolExecutor(max_workers=1)
+    batcher.submit(state.checkpoint_batches_forever)
+    try:
+        check_time = time.time() + 0.15
+        cancelled = state.schedule_refresh(check_time)
+        kept = state.schedule_refresh(check_time)
+        cancelled.cancel()
+        assert kept.wait(timeout=5)
+        time.sleep(0.05)
+        assert calls == [[]], "the kept refresh makes one request, the cancelled none"
+
+        alone = state.schedule_refresh(time.time() + 0.1)
+        alone.cancel()
+        time.sleep(0.25)
+        assert calls == [[]], "a cancelled refresh with no companion makes no request"
+    finally:
+        state.stop_checkpointing()
+        batcher.shutdown(wait=True)
+
+
+def test_refresh_requested_after_its_time_was_sent_makes_its_own_request():
+    """Two independent coordinators, one check time. The first request is sent
+    at that time. A refresh for the same time requested afterwards is due at
+    once and goes in the next request. Coalescing covers only refreshes
+    requested before their time."""
+    calls: list[list] = []
+    state = _refresh_state(_refresh_client(calls))
+    batcher = ThreadPoolExecutor(max_workers=1)
+    batcher.submit(state.checkpoint_batches_forever)
+    try:
+        check_time = time.time() + 0.1
+        first = state.schedule_refresh(check_time)
+        second = state.schedule_refresh(check_time)
+        assert first.wait(timeout=5) and second.wait(timeout=5)
+        assert calls == [[]]
+
+        late = state.schedule_refresh(check_time)
+        assert late.wait(timeout=5)
+        assert calls == [[], []]
+    finally:
+        state.stop_checkpointing()
+        batcher.shutdown(wait=True)
+
+
 def _make_execution_state_for_operations(mock_lambda_client, *, operations=None):
     return ExecutionState(
         durable_execution_arn="test_arn",