diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py index 1db42ec3..339a0e9f 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/lambda_service.py @@ -1086,7 +1086,8 @@ def from_dict( class CheckpointOutput: """Representation of the CheckpointDurableExecutionOutput structure of the DEX CheckpointDurableExecution API.""" - checkpoint_token: str + # None on the terminal checkpoint that ends the execution. + checkpoint_token: str | None new_execution_state: CheckpointUpdatedExecutionState @classmethod @@ -1109,8 +1110,7 @@ def from_dict(cls, data: MutableMapping[str, Any]) -> CheckpointOutput: new_execution_state = CheckpointUpdatedExecutionState() return cls( - # TODO: maybe should throw if empty? - checkpoint_token=data.get("CheckpointToken", ""), + checkpoint_token=data.get("CheckpointToken"), new_execution_state=new_execution_state, ) @@ -1201,6 +1201,11 @@ def checkpoint( updates: list[OperationUpdate], client_token: str | None, ) -> CheckpointOutput: + # A checkpoint token is required. Raise a clear, retryable error (so the + # invocation re-drives) rather than letting the client reject an empty + # value with an opaque validation error. + if not checkpoint_token: + raise CheckpointError("Cannot checkpoint without a checkpoint token.") try: optional_params: dict[str, str] = {} if client_token is not None: @@ -1230,6 +1235,13 @@ def get_execution_state( next_marker: str, max_items: int = 1000, ) -> StateOutput: + # A checkpoint token is required. Raise a clear, retryable error (so the + # invocation re-drives) rather than letting the client reject an empty + # value with an opaque validation error. + if not checkpoint_token: + raise GetExecutionStateError( + "Cannot get execution state without a checkpoint token." + ) try: result: GetDurableExecutionStateResponseTypeDef = ( self.client.get_durable_execution_state( 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 cb174441..fb66bc05 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 @@ -16,6 +16,7 @@ from aws_durable_execution_sdk_python.exceptions import ( BackgroundThreadError, + CheckpointError, DurableExecutionsError, DurableOperationError, GetExecutionStateError, @@ -37,7 +38,7 @@ from aws_durable_execution_sdk_python.plugin import ( PluginExecutor, ) -from aws_durable_execution_sdk_python.threading import CompletionEvent, OrderedLock +from aws_durable_execution_sdk_python.threading import CompletionEvent if TYPE_CHECKING: @@ -282,7 +283,6 @@ def __init__( self._operations: dict[str, Operation] = dict(operations) self._service_client: DurableServiceClient = service_client self._plugin_executor: PluginExecutor = plugin_executor - self._ordered_checkpoint_lock: OrderedLock = OrderedLock() self._operations_lock: Lock = Lock() # Checkpoint batching configuration @@ -295,6 +295,12 @@ def __init__( self._overflow_queue: queue.Queue[QueuedOperation] = queue.Queue() self._checkpointing_stopped: threading.Event = threading.Event() self._checkpointing_failed: CompletionEvent = CompletionEvent() + # Set once the service confirms the execution has completed (a checkpoint + # response with no token). No further checkpoint can succeed afterward. + self._execution_completed: threading.Event = threading.Event() + # Serializes the execution-completed check with enqueueing so a checkpoint + # is never enqueued after _settle_after_execution_completed drains the queue. + self._completion_lock: Lock = Lock() # Concurrency management for parallel operations: parent_id -> {child_operation_ids} self._parent_to_children: dict[str, set[str]] = {} @@ -496,6 +502,26 @@ def emit_operation_update_hook(self, operation: Operation) -> None: self._plugin_executor.on_operation_update(operation) + def _reject_if_execution_completed( + self, operation_update: OperationUpdate | None + ) -> None: + """Raise OrphanedChildException when the execution has already completed. + + Called before dispatching the START hook and again inside + _completion_lock: the first keeps an orphaned operation from emitting a + START with no matching completion, the second closes the race with a + concurrent completion. + """ + if not self._execution_completed.is_set(): + return + operation_id: str = ( + operation_update.operation_id if operation_update is not None else "" + ) + raise OrphanedChildException( + "Execution already completed; checkpoint will not be processed.", + operation_id=operation_id, + ) + def create_checkpoint( self, operation_update: OperationUpdate | None = None, @@ -544,8 +570,13 @@ def create_checkpoint( If False, returns immediately without blocking for performance. Raises: - Any exception from the background checkpoint processing will propagate - through the ThreadPoolExecutor to the main thread, terminating the Lambda. + OrphanedChildException: If the operation's parent context has already + completed, or the execution itself has already completed, so the + checkpoint cannot be processed. + BackgroundThreadError: If background checkpoint processing has failed; + the stored failure is re-raised to the caller. For a synchronous + checkpoint, a later background failure also surfaces here through + the completion event. Examples: # Synchronous checkpoint (default, safe) @@ -602,6 +633,12 @@ def create_checkpoint( # This will raise the stored BackgroundThreadError self._checkpointing_failed.wait() + # Reject a late checkpoint before dispatching the START hook, so an + # orphaned operation does not emit a START with no matching completion + # (mirrors the parent-done check above). The _completion_lock block below + # re-checks to close the race with a concurrent completion. + self._reject_if_execution_completed(operation_update) + # Conditionally create completion event based on is_sync parameter completion_event: CompletionEvent | None = ( CompletionEvent() if is_sync else None @@ -618,8 +655,18 @@ def create_checkpoint( # Create wrapper object for queue queued_op = QueuedOperation(operation_update, completion_event) - # Enqueue the wrapper object (operation_update can be None for empty checkpoints) - self._checkpoint_queue.put(queued_op) + # Enqueue under the same lock the background loop holds while it stops and + # drains the queue - on completion via _settle_after_execution_completed, or + # on failure in the exception handler. Re-check both terminal conditions + # inside the lock so a checkpoint is never enqueued after a drain and left + # with a waiter that blocks forever. + with self._completion_lock: + if self._checkpointing_failed.is_set(): + # Raises the stored BackgroundThreadError. + self._checkpointing_failed.wait() + self._reject_if_execution_completed(operation_update) + # Enqueue the wrapper object (operation_update can be None for empty checkpoints) + self._checkpoint_queue.put(queued_op) # Conditionally wait for completion based on is_sync parameter if is_sync: @@ -713,22 +760,20 @@ def _mark_orphans(self, context_id: str) -> None: def checkpoint_batches_forever(self) -> None: """Single background thread that batches operations and processes results. - Runs until shutdown is signaled. This method processes checkpoint operations - in batches, makes API calls to persist them, and updates the execution state - with the results. - - The method maintains the checkpoint token locally and updates it after each - successful batch processing. It continues running until stop_checkpointing() - is called. - - Note: When shutdown is signaled, only non-essential async checkpoints may remain - in the queue. All critical synchronous checkpoints (SUCCEED, FAIL, etc.) will - have already completed because the main thread blocks on them. Therefore, we - don't need to drain the queue - the Lambda timeout will handle cleanup. - - Raises: - Any exception from the service client checkpoint call will propagate naturally, - terminating the background thread and signaling an error to the main thread. + Collects queued operations into batches, persists each batch with one API + call, refreshes execution state from the response, and wakes the batch's + waiters. The checkpoint token is held locally and advanced after each + successful batch. + + The loop stops on any of three conditions: shutdown is signaled + (stop_checkpointing), the service reports the execution has completed (a + terminal batch whose response omits the token), or a batch fails. On + completion or failure the remaining queued operations are drained and their + waiters settled - with OrphanedChildException on completion, or the wrapped + BackgroundThreadError on failure - so no waiter blocks forever. On a plain + shutdown signal the queue is not drained: only non-essential asynchronous + checkpoints can remain, since the main thread blocks on every synchronous + one. """ # Keep checkpoint token as local variable in the loop current_checkpoint_token: str = self._current_checkpoint_token @@ -766,16 +811,43 @@ def checkpoint_batches_forever(self) -> None: logger.debug("Checkpoint batch processed successfully") - # Update local token for next iteration - current_checkpoint_token = output.checkpoint_token + # The service omits the token only when it completes the + # execution. Advance on a returned token; otherwise treat a + # terminal batch as completion and any other omission as an + # invalid response. + execution_completed: bool = False + if output.checkpoint_token: + current_checkpoint_token = output.checkpoint_token + elif any( + update.operation_type is OperationType.EXECUTION + and ( + update.action is OperationAction.SUCCEED + or update.action is OperationAction.FAIL + ) + for update in updates + ): + execution_completed = True + else: + raise CheckpointError( + "Checkpoint response omitted the token outside of " + "execution completion." + ) previous_operations = self.operations - # Fetch new operations from the API before unblocking sync waiters + # Fetch new operations from the API before unblocking sync + # waiters. On completion the token is consumed, so skip + # pagination (which would reuse the spent token) and record + # only the operations the terminal response carries inline. + fetch_marker: str | None = ( + None + if execution_completed + else output.new_execution_state.next_marker + ) updated_operations = self.fetch_paginated_operations( output.new_execution_state.operations, - output.checkpoint_token, - output.new_execution_state.next_marker, + current_checkpoint_token, + fetch_marker, ) self._plugin_executor.on_operation_update( updated_operations, @@ -787,6 +859,15 @@ def checkpoint_batches_forever(self) -> None: for queued_op in batch: if queued_op.completion_event is not None: queued_op.completion_event.set() + + if execution_completed: + # The execution is complete; no further checkpoint can + # succeed. Stop the loop and settle any still-queued + # operations (e.g. from orphaned concurrent branches) so + # their waiters do not block and no further checkpoint + # request is issued with the consumed token. + self._settle_after_execution_completed() + break except Exception as e: # Checkpoint failed - wake all blocked threads so they can raise error # Drain both queues and signal all completion events @@ -795,38 +876,74 @@ def checkpoint_batches_forever(self) -> None: "Checkpoint creation failed", e ) - # FIFO: although at this point order not really import any anymore - # Signal completion events for the failed batch + # Signal completion events for the failed batch (already dequeued, + # so no producer can race these). for queued_op in batch: if queued_op.completion_event is not None: queued_op.completion_event.set(bg_error) - # overflow 1st: although at this point order not really import any anymore - 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 - - # finally Wake all blocked threads in main queue - while not self._checkpoint_queue.empty(): - try: - item = self._checkpoint_queue.get_nowait() - if item.completion_event: - item.completion_event.set(bg_error) - except queue.Empty: - break - - # Set the failure event so future checkpoint attempts fail immediately - self._checkpointing_failed.set(bg_error) + # Drain the queues and set the failure flag under the completion + # lock so a concurrent create_checkpoint either observes the + # failure and raises, or has its operation drained here - never + # enqueued after the drain and left blocked. + with self._completion_lock: + 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 + + while not self._checkpoint_queue.empty(): + try: + item = self._checkpoint_queue.get_nowait() + if item.completion_event: + item.completion_event.set(bg_error) + except queue.Empty: + break + + # Future checkpoint attempts fail immediately. + self._checkpointing_failed.set(bg_error) # Exit the loop - error has been signaled to main thread via completion events break logger.debug("Background checkpoint processing stopped") + def _settle_after_execution_completed(self) -> None: + """Stop checkpointing and settle queued operations after the execution ends. + + Called when a checkpoint response omits the token on a terminal batch, + which the service does only when it completes the execution. Any operation + still queued belongs to work that can no longer be checkpointed (typically + an orphaned concurrent branch), so its waiter is settled with + OrphanedChildException rather than left blocking or sent with the consumed + token. + """ + with self._completion_lock: + self._execution_completed.set() + self._checkpointing_stopped.set() + + for pending_queue in (self._overflow_queue, self._checkpoint_queue): + while not pending_queue.empty(): + try: + queued_op: QueuedOperation = 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, + ) + ) + def stop_checkpointing(self) -> None: """Signal background thread to stop checkpointing. diff --git a/packages/aws-durable-execution-sdk-python/tests/lambda_service_test.py b/packages/aws-durable-execution-sdk-python/tests/lambda_service_test.py index 2d0ff829..55f1ebce 100644 --- a/packages/aws-durable-execution-sdk-python/tests/lambda_service_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/lambda_service_test.py @@ -1641,10 +1641,14 @@ def test_checkpoint_output_from_dict(): def test_checkpoint_output_from_dict_empty(): - """Test CheckpointOutput.from_dict with empty data.""" + """Test CheckpointOutput.from_dict with empty data. + + The service omits CheckpointToken on the terminal checkpoint; from_dict + surfaces that as None rather than an empty string. + """ data = {} output = CheckpointOutput.from_dict(data) - assert not output.checkpoint_token + assert output.checkpoint_token is None assert len(output.new_execution_state.operations) == 0 assert output.new_execution_state.next_marker is None @@ -1825,6 +1829,37 @@ def test_lambda_client_checkpoint_with_explicit_none_client_token(): assert result.checkpoint_token == "new_token" # noqa: S105 +@pytest.mark.parametrize("token", ["", None]) +def test_lambda_client_checkpoint_empty_token_raises(token): + """An empty/None checkpoint token raises a retryable CheckpointError, not a client error.""" + mock_client = Mock() + lambda_client = LambdaClient(mock_client) + update = OperationUpdate( + operation_id="op1", + operation_type=OperationType.STEP, + action=OperationAction.START, + ) + + with pytest.raises(CheckpointError) as exc_info: + lambda_client.checkpoint("arn123", token, [update], None) + + assert exc_info.value.is_retryable() + mock_client.checkpoint_durable_execution.assert_not_called() + + +@pytest.mark.parametrize("token", ["", None]) +def test_lambda_client_get_execution_state_empty_token_raises(token): + """An empty/None checkpoint token raises a retryable GetExecutionStateError, not a client error.""" + mock_client = Mock() + lambda_client = LambdaClient(mock_client) + + with pytest.raises(GetExecutionStateError) as exc_info: + lambda_client.get_execution_state("arn123", token, "marker") + + assert exc_info.value.is_retryable() + mock_client.get_durable_execution_state.assert_not_called() + + def test_lambda_client_checkpoint_with_empty_string_client_token(): """Test LambdaClient.checkpoint method with empty string client_token - should pass empty string.""" mock_client = Mock() 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 bb9a141f..7cea2b0c 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -15,6 +15,7 @@ from aws_durable_execution_sdk_python.exceptions import ( BackgroundThreadError, + CheckpointError, DurableApiErrorCategory, GetExecutionStateError, OrphanedChildException, @@ -1884,6 +1885,409 @@ def run_batching(): pass # Expected +def test_checkpoint_missing_token_on_non_terminal_batch_fails(): + """A non-terminal checkpoint response without a token fails the checkpoint. + + The service omits the token only when it completes the execution, so a + missing token on any other checkpoint is an invalid response. + """ + mock_lambda_client = Mock(spec=LambdaClient) + mock_lambda_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token=None, + new_execution_state=CheckpointUpdatedExecutionState( + operations=[], next_marker=None + ), + ) + + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + + completion_event = CompletionEvent() + state._checkpoint_queue.put( + QueuedOperation( + OperationUpdate( + operation_id="op1", + operation_type=OperationType.STEP, + action=OperationAction.START, + ), + completion_event, + ) + ) + + thread = threading.Thread(daemon=True, target=state.checkpoint_batches_forever) + thread.start() + thread.join(timeout=2.0) + + assert completion_event.is_set() + try: + completion_event.wait() + pytest.fail("Should have raised BackgroundThreadError") + except BackgroundThreadError as bg_error: + assert isinstance(bg_error.source_exception, CheckpointError) + assert bg_error.source_exception.is_retryable() + + +@pytest.mark.parametrize( + "terminal_update", + [ + OperationUpdate.create_execution_succeed(payload="{}"), + OperationUpdate.create_execution_fail( + error=ErrorObject(message="boom", type="Error", data=None, stack_trace=None) + ), + ], +) +def test_checkpoint_missing_token_on_terminal_batch_succeeds(terminal_update): + """A terminal checkpoint response without a token completes normally. + + The service omits the token when it ends the execution (SUCCEED or FAIL); + that is expected and must not fail the checkpoint. + """ + mock_lambda_client = Mock(spec=LambdaClient) + mock_lambda_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token=None, + new_execution_state=CheckpointUpdatedExecutionState( + operations=[], next_marker=None + ), + ) + + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + + completion_event = CompletionEvent() + state._checkpoint_queue.put(QueuedOperation(terminal_update, completion_event)) + + thread = threading.Thread(daemon=True, target=state.checkpoint_batches_forever) + thread.start() + try: + # Returns without raising once the terminal checkpoint is processed. + completion_event.wait() + finally: + state.stop_checkpointing() + thread.join(timeout=2.0) + + assert completion_event.is_set() + assert state._execution_completed.is_set() + + +def test_settle_after_execution_completed_orphans_pending_operations(): + """When the execution completes, still-queued operations are settled as orphaned. + + Their waiters must be released with OrphanedChildException rather than left + blocking or sent with the consumed token, and checkpointing must stop. + """ + mock_lambda_client = Mock(spec=LambdaClient) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + + completion_event = CompletionEvent() + state._checkpoint_queue.put( + QueuedOperation( + OperationUpdate( + operation_id="orphan_op", + operation_type=OperationType.STEP, + action=OperationAction.START, + ), + completion_event, + ) + ) + + state._settle_after_execution_completed() + + assert state._execution_completed.is_set() + assert state._checkpointing_stopped.is_set() + with pytest.raises(OrphanedChildException): + completion_event.wait() + + +def _arm_execution_completed(state: ExecutionState) -> None: + state._execution_completed.set() + + +def _arm_checkpointing_failed(state: ExecutionState) -> None: + state._checkpointing_failed.set( + BackgroundThreadError("Checkpoint creation failed", RuntimeError("boom")) + ) + + +@pytest.mark.parametrize( + ("arm_terminal_state", "expected_error"), + [ + pytest.param( + _arm_execution_completed, OrphanedChildException, id="execution-completed" + ), + pytest.param( + _arm_checkpointing_failed, BackgroundThreadError, id="checkpointing-failed" + ), + ], +) +def test_create_checkpoint_rejected_in_terminal_state( + arm_terminal_state, expected_error +): + """A checkpoint attempted in a terminal state is rejected, not enqueued. + + Whether the background loop has completed the execution (_execution_completed + -> OrphanedChildException) or failed (_checkpointing_failed -> + BackgroundThreadError), create_checkpoint surfaces the terminal condition + rather than enqueueing an operation whose waiter would block forever. + """ + mock_lambda_client = Mock(spec=LambdaClient) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + arm_terminal_state(state) + + with pytest.raises(expected_error): + state.create_checkpoint( + OperationUpdate( + operation_id="late_op", + operation_type=OperationType.STEP, + action=OperationAction.START, + ), + is_sync=True, + ) + + assert state._checkpoint_queue.empty() + + +def test_create_checkpoint_after_completion_skips_start_hook(): + """A checkpoint rejected after completion must not emit a START plugin hook. + + The orphan check runs before the START dispatch, matching the parent-done + path, so a rejected operation never fires a START with no completion. + """ + mock_lambda_client = Mock(spec=LambdaClient) + mock_plugin_executor = Mock() + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=mock_plugin_executor, + ) + state._execution_completed.set() + + with pytest.raises(OrphanedChildException): + state.create_checkpoint( + OperationUpdate( + operation_id="late_op", + operation_type=OperationType.STEP, + action=OperationAction.START, + ), + is_sync=True, + ) + + mock_plugin_executor.on_operation_action.assert_not_called() + + +def test_checkpoint_completion_skips_pagination_with_consumed_token(): + """On completion the consumed token must not be reused to paginate. + + The terminal response carries its operations inline; if it also reports a + next_marker, the loop must not call get_execution_state with the spent + token (which would turn a clean completion into a failure). + """ + mock_lambda_client = Mock(spec=LambdaClient) + mock_lambda_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token=None, + new_execution_state=CheckpointUpdatedExecutionState( + operations=[], next_marker="more-pages" + ), + ) + + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + + completion_event: CompletionEvent = CompletionEvent() + state._checkpoint_queue.put( + QueuedOperation( + OperationUpdate.create_execution_succeed(payload="{}"), + completion_event, + ) + ) + + thread = threading.Thread(daemon=True, target=state.checkpoint_batches_forever) + thread.start() + try: + completion_event.wait() + finally: + state.stop_checkpointing() + thread.join(timeout=2.0) + + assert state._execution_completed.is_set() + mock_lambda_client.get_execution_state.assert_not_called() + + +def test_completion_lock_rejects_producer_racing_completion(): + """The in-lock re-check rejects a producer that raced a concurrent completion. + + A producer can pass the pre-lock orphan check and then block acquiring + _completion_lock while the execution completes. When it finally takes the + lock it must observe _execution_completed and raise, rather than enqueue an + operation into an already-drained queue where its waiter would block forever. + """ + mock_lambda_client = Mock(spec=LambdaClient) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + + outcome: dict[str, str] = {} + started: threading.Event = threading.Event() + + def producer() -> None: + started.set() + try: + state.create_checkpoint( + OperationUpdate( + operation_id="racing_op", + operation_type=OperationType.STEP, + action=OperationAction.START, + ), + is_sync=False, + ) + outcome["result"] = "enqueued" + except OrphanedChildException: + outcome["result"] = "orphaned" + + thread = threading.Thread(target=producer, daemon=True) + + # Hold the lock to stand in for a settle/drain in progress. The producer + # passes the pre-lock check (execution not yet completed) and then blocks + # acquiring the lock; completion happens while it waits. The lock is always + # released so the producer can never block the test. + state._completion_lock.acquire() + try: + thread.start() + assert started.wait(timeout=2.0), "producer thread did not start" + # Give the producer time to pass the pre-lock check and block on the lock. + time.sleep(0.1) + state._execution_completed.set() + finally: + state._completion_lock.release() + + thread.join(timeout=2.0) + + assert outcome["result"] == "orphaned" + assert state._checkpoint_queue.empty() + + +def test_checkpoint_missing_token_on_mixed_terminal_batch_succeeds(): + """A batch mixing a terminal update with non-terminal ones completes. + + The terminal probe scans the sent updates, so a batch that carries both a + non-terminal START and a terminal EXECUTION update is recognized as + completion when the response omits the token. + """ + mock_lambda_client = Mock(spec=LambdaClient) + mock_lambda_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token=None, + new_execution_state=CheckpointUpdatedExecutionState( + operations=[], next_marker=None + ), + ) + + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + + non_terminal_event: CompletionEvent = CompletionEvent() + terminal_event: CompletionEvent = CompletionEvent() + state._checkpoint_queue.put( + QueuedOperation( + OperationUpdate( + operation_id="op1", + operation_type=OperationType.STEP, + action=OperationAction.START, + ), + non_terminal_event, + ) + ) + state._checkpoint_queue.put( + QueuedOperation( + OperationUpdate.create_execution_succeed(payload="{}"), terminal_event + ) + ) + + thread = threading.Thread(daemon=True, target=state.checkpoint_batches_forever) + thread.start() + try: + non_terminal_event.wait() + terminal_event.wait() + finally: + state.stop_checkpointing() + thread.join(timeout=2.0) + + assert state._execution_completed.is_set() + + +def test_checkpoint_missing_token_on_empty_only_batch_fails(): + """An empty-only batch (no updates) returning no token is invalid. + + With nothing sent, a missing token cannot mean completion, so it must fail + the checkpoint rather than be treated as terminal. + """ + mock_lambda_client = Mock(spec=LambdaClient) + mock_lambda_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token=None, + new_execution_state=CheckpointUpdatedExecutionState( + operations=[], next_marker=None + ), + ) + + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + ) + + empty_event: CompletionEvent = CompletionEvent() + state._checkpoint_queue.put(QueuedOperation(None, empty_event)) + + thread = threading.Thread(daemon=True, target=state.checkpoint_batches_forever) + thread.start() + thread.join(timeout=2.0) + + assert empty_event.is_set() + try: + empty_event.wait() + pytest.fail("Should have raised BackgroundThreadError") + except BackgroundThreadError as bg_error: + assert isinstance(bg_error.source_exception, CheckpointError) + + def test_collect_checkpoint_batch_shutdown_path(): """Test _collect_checkpoint_batch during shutdown with operations in queue.