aws-durable-execution-sdk-python 2.0.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

aws_durable_execution_sdk_python/__about__.py

@@ -1,4 +1,4 @@
 # SPDX-FileCopyrightText: 2025-present Amazon.com, Inc. or its affiliates.
 #
 # SPDX-License-Identifier: Apache-2.0
-__version__ = "2.0.0"
+__version__ = "2.0.1"

aws_durable_execution_sdk_python/concurrency/executor.py

@@ -8,6 +8,7 @@ import queue
 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 @@ if TYPE_CHECKING:
     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 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
 
         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 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
                     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 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
                         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 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
                         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 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
                         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

aws_durable_execution_sdk_python/exceptions.py

@@ -13,7 +13,7 @@ BAD_REQUEST_ERROR: int = 400
 TOO_MANY_REQUESTS_ERROR: int = 429
 SERVICE_ERROR: int = 500
 INVALID_PARAMETER_VALUE_EXCEPTION: str = "InvalidParameterValueException"
-INVALID_CHECKPOINT_TOKEN_PREFIX: str = "Invalid Checkpoint Token"
+INVALID_CHECKPOINT_TOKEN_PREFIX: str = "Invalid checkpoint token"
 
 # Non-retryable customer error codes that arrive as non-4xx (e.g. HTTP 502) from Lambda.
 # Unlike typical 5xx errors, these require customer intervention (e.g., fixing
@@ -162,7 +162,7 @@ class BotoClientError(InvocationError):
           These arrive as HTTP 502 but require customer intervention to fix.
         - 4xx errors → EXECUTION, except:
           - 429 (TooManyRequests) → INVOCATION (throttling is transient)
-          - InvalidParameterValueException with "Invalid Checkpoint Token" → INVOCATION
+          - InvalidParameterValueException with "Invalid checkpoint token" (exact match) → INVOCATION
             (stale token from a concurrent checkpoint; next invocation gets a fresh token)
         - 5xx, network errors → INVOCATION
         """

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 @@ if TYPE_CHECKING:
 
 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 @@ class ExecutionState:
         )
 
         # 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 @@ class ExecutionState:
         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 @@ class ExecutionState:
                         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 @@ class ExecutionState:
                     # 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 @@ class ExecutionState:
             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 @@ class ExecutionState:
             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:

{aws_durable_execution_sdk_python-2.0.0.dist-info → aws_durable_execution_sdk_python-2.0.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: aws-durable-execution-sdk-python
-Version: 2.0.0
+Version: 2.0.1
 Summary: AWS Durable Execution SDK for Python
 Project-URL: Documentation, https://github.com/aws/aws-durable-execution-sdk-python#readme
 Project-URL: Issues, https://github.com/aws/aws-durable-execution-sdk-python/issues

{aws_durable_execution_sdk_python-2.0.0.dist-info → aws_durable_execution_sdk_python-2.0.1.dist-info}/RECORD

@@ -1,10 +1,10 @@
 aws_durable_execution_sdk_python/.gitignore,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aws_durable_execution_sdk_python/__about__.py,sha256=iMryfzT8zKHfkT0Q7FfA6pgZKIB24Mz7jR5gb7CxJGM,137
+aws_durable_execution_sdk_python/__about__.py,sha256=JTKxzhd0vBceCgP4Nsi3Wjz8K-nKVjiSzkQR_y8dkAY,137
 aws_durable_execution_sdk_python/__init__.py,sha256=K10hqr0JcYNhf41AkKVvboKK1AXLrZVQJ3Foy4h3taE,2488
 aws_durable_execution_sdk_python/config.py,sha256=lTV4g_LXiFcyjiXQ5Kn7XvtFDuybw8tvgigetY55u9I,29176
 aws_durable_execution_sdk_python/constants.py,sha256=Pdgfagm9wcWt5UZRv1tLXvCUqkalv89zrWoEoUSL4EA,259
 aws_durable_execution_sdk_python/context.py,sha256=D1Z-0pXzUTDXTE3PMK6Mw-XpPFV9q-XWtAy6k_nYLBI,37408
-aws_durable_execution_sdk_python/exceptions.py,sha256=pi-HAkp_hW-C_J7-5QLaofJqSxr5Tfaxttz-X6mf_Ag,21341
+aws_durable_execution_sdk_python/exceptions.py,sha256=nBg9aZxfKSYilmy-xdMIjR2gVGqKfOf2BZA9TJO4RlU,21355
 aws_durable_execution_sdk_python/execution.py,sha256=kLDGU0EvwuJdjxVvE9ZcGKUz8lBhWhS8pirGDjYz2Kc,19968
 aws_durable_execution_sdk_python/filesystem_serdes.py,sha256=wPURZb9bwfUqTUhnNp3kOBojMyFnXxjL-QPkO6Yk1aw,10685
 aws_durable_execution_sdk_python/identifier.py,sha256=0NQBir435WivGGZWlj6520u4UqT0yxIig2iUFLAzuVQ,3023
@@ -16,13 +16,13 @@ aws_durable_execution_sdk_python/preview.py,sha256=taizk3reZ-1SZYsM-0YGW5Oyze-tD
 aws_durable_execution_sdk_python/py.typed,sha256=DbCDcAirv769HS1ibzZLXq8VPazYf7JCkKSzj3ctudM,58
 aws_durable_execution_sdk_python/retries.py,sha256=X8hfFPnrNR_LLJdiQ9g-KX8aQN5BKZTscmlQBVTt9NU,13678
 aws_durable_execution_sdk_python/serdes.py,sha256=_AH09m5vLKzd_6vn8krnit7p9ogqV4eLvH0shyoog4M,17347
-aws_durable_execution_sdk_python/state.py,sha256=94PUbiCtoj2WxsZM9wq7O0tlf9QZb5Ha0kPqnEXGQ_w,53775
+aws_durable_execution_sdk_python/state.py,sha256=7C56NoRA44k9TRpISjUDQbA7YYLZeepp8tNkSZMK1wA,61123
 aws_durable_execution_sdk_python/suspend.py,sha256=PIgc6w5-diZIq58-tPv51dzBSHGW7e4oR0QFam8Vgi0,3323
 aws_durable_execution_sdk_python/threading.py,sha256=cU7vIEWHYco-qqkgoAMyPHSOQ8b_Q_r_rZGG7m1g7I4,7650
 aws_durable_execution_sdk_python/types.py,sha256=MFTaRL8XN2PoelJKfonMhNEgTZkOGYDkGRqx2WzQHNo,7022
 aws_durable_execution_sdk_python/waits.py,sha256=NIF77G5VtXYb4748IC8MEvbE_L1ipBPA_HSkSQE4PlA,4119
 aws_durable_execution_sdk_python/concurrency/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aws_durable_execution_sdk_python/concurrency/executor.py,sha256=61Lbm2ahp7VTT6U5-AtjeueaREbu0EX5wBK_EBRN9pw,36221
+aws_durable_execution_sdk_python/concurrency/executor.py,sha256=2EZiXHAJTX3zBnuji56zBiuS86AgXQRqD6W8OuZN5i0,37433
 aws_durable_execution_sdk_python/concurrency/models.py,sha256=Agi0Y_N8aKzWxQBeqhoiN97n-o0mkcw0n7-Fuw_FHt0,25135
 aws_durable_execution_sdk_python/operation/__init__.py,sha256=xjtSpukZb6sftgJ5oVPSc6TcgKgVQg4vXHeYIeKOPks,25
 aws_durable_execution_sdk_python/operation/base.py,sha256=kuUroQnJFo36rH3Ad0IJjuLi4XlL2d13tN9fmKuqw3o,7632
@@ -36,6 +36,6 @@ aws_durable_execution_sdk_python/operation/wait.py,sha256=YMVomY1QpPfyACBnN-kQk8
 aws_durable_execution_sdk_python/operation/wait_for_condition.py,sha256=BS8rFZGQqQoV0q3gG_1MZSZxXXdl8u8dyYTRGcguMS4,13082
 aws_durable_execution_sdk_python/LICENSE,sha256=CeipvOyAZxBGUsFoaFqwkx54aPnIKEtm9a5u2uXxEws,10142
 aws_durable_execution_sdk_python/NOTICE,sha256=1CkO1kwu3Q_OHYTj-d-yiBJA_lNN73a4zSntavaD4oc,67
-aws_durable_execution_sdk_python-2.0.0.dist-info/METADATA,sha256=ganlTgYTZS0rrZKMbNDbrjDZd_HSRgzM6Fyiq_7DxGY,6902
-aws_durable_execution_sdk_python-2.0.0.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
-aws_durable_execution_sdk_python-2.0.0.dist-info/RECORD,,
+aws_durable_execution_sdk_python-2.0.1.dist-info/METADATA,sha256=NMLsc77gVyCpm75B39WaUlEP9Hkob5NYhCy_5mNni-k,6902
+aws_durable_execution_sdk_python-2.0.1.dist-info/WHEEL,sha256=THafob7ofN-NsuMN7Mg4qZyHaQI7KkD-QlcQatYhXPo,87
+aws_durable_execution_sdk_python-2.0.1.dist-info/RECORD,,

{aws_durable_execution_sdk_python-2.0.0.dist-info → aws_durable_execution_sdk_python-2.0.1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.32.0
+Generator: hatchling 1.32.3
 Root-Is-Purelib: true
 Tag: py3-none-any