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",