From ffb6588429be9017c8cdce30809614f01d479b8a Mon Sep 17 00:00:00 2001 From: warren Date: Sun, 30 Aug 2026 17:06:22 +0800 Subject: [PATCH] fix: reject orphan checkpoint race Recheck parent completion while holding the queue insertion locks so an already-known branch cannot enqueue after becoming orphaned. Add a deterministic event-coordinated regression test. --- .../aws_durable_execution_sdk_python/state.py | 45 +++++++---- .../tests/state_test.py | 81 +++++++++++++++++++ 2 files changed, 110 insertions(+), 16 deletions(-) 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 9f2391ce..8642c0ff 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 @@ -308,7 +308,8 @@ def __init__( # Operations whose parent has completed self._parent_done: set[str] = set() - # Protects parent_to_children and parent_done + # Protects parent_to_children and parent_done. When both state locks are + # required, acquire _completion_lock before _parent_done_lock. self._parent_done_lock: Lock = Lock() # Branch thread pools created by concurrency coordinators. A pool @@ -538,6 +539,23 @@ def _reject_if_execution_completed( operation_id=operation_id, ) + def _reject_if_parent_done(self, operation_update: OperationUpdate) -> None: + """Raise OrphanedChildException when the operation is orphaned. + + Must be called while holding _parent_done_lock so the check can be + linearized with parent completion and, at the final call site, enqueue. + """ + if operation_update.operation_id not in self._parent_done: + return + logger.debug( + "Rejecting checkpoint for operation %s - parent is done", + operation_update.operation_id, + ) + raise OrphanedChildException( + "Parent context completed, child operation cannot checkpoint", + operation_id=operation_update.operation_id, + ) + def create_checkpoint( self, operation_update: OperationUpdate | None = None, @@ -631,18 +649,7 @@ def create_checkpoint( self._mark_orphans(operation_update.operation_id) # Check if this operation's parent is done - if operation_update.operation_id in self._parent_done: - logger.debug( - "Rejecting checkpoint for operation %s - parent is done", - operation_update.operation_id, - ) - error_msg = ( - "Parent context completed, child operation cannot checkpoint" - ) - raise OrphanedChildException( - error_msg, - operation_id=operation_update.operation_id, - ) + self._reject_if_parent_done(operation_update) # Check if background checkpointing has failed if self._checkpointing_failed.is_set(): @@ -675,14 +682,20 @@ def create_checkpoint( # 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 a waiter that blocks forever. Acquire _parent_done_lock inside + # _completion_lock and hold both through queue insertion so parent completion + # cannot mark the operation orphaned between the final check and enqueue. 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) + if operation_update is None: + self._checkpoint_queue.put(queued_op) + else: + with self._parent_done_lock: + self._reject_if_parent_done(operation_update) + self._checkpoint_queue.put(queued_op) # Conditionally wait for completion based on is_sync parameter if is_sync: 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 7c33d3b0..bb5f841d 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -1426,6 +1426,87 @@ def test_rejection_of_operations_from_completed_parents(): assert exc_info.value.operation_id == "child_1" +def test_parent_done_recheck_rejects_checkpoint_racing_parent_completion(): + """A checkpoint that passed validation is rejected if its parent completes. + + The child pauses in the operation hook after releasing _parent_done_lock and + before acquiring _completion_lock. Completing the parent during that pause + deterministically reproduces the validation-to-enqueue race without sleeps. + """ + mock_lambda_client = Mock(spec=LambdaClient) + mock_plugin_executor = Mock(spec=PluginExecutor) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=mock_plugin_executor, + ) + + child_reached_hook = threading.Event() + release_child_hook = threading.Event() + + def block_child_completion(operation_update: OperationUpdate, **_: object) -> None: + if ( + operation_update.operation_id == "child_1" + and operation_update.action == OperationAction.SUCCEED + ): + child_reached_hook.set() + assert release_child_hook.wait(timeout=2.0), ( + "child checkpoint was not released" + ) + + mock_plugin_executor.on_operation_action.side_effect = block_child_completion + + state.create_checkpoint( + OperationUpdate( + operation_id="parent_1", + operation_type=OperationType.CONTEXT, + action=OperationAction.START, + ), + is_sync=False, + ) + state.create_checkpoint( + OperationUpdate( + operation_id="child_1", + operation_type=OperationType.CONTEXT, + action=OperationAction.START, + parent_id="parent_1", + ), + is_sync=False, + ) + child_complete = OperationUpdate( + operation_id="child_1", + operation_type=OperationType.CONTEXT, + action=OperationAction.SUCCEED, + parent_id="parent_1", + ) + + with ThreadPoolExecutor(max_workers=1) as executor: + child_future = executor.submit(state.create_checkpoint, child_complete, False) + assert child_reached_hook.wait(timeout=2.0), ( + "child checkpoint did not reach the hook" + ) + + try: + state.create_checkpoint( + OperationUpdate( + operation_id="parent_1", + operation_type=OperationType.CONTEXT, + action=OperationAction.SUCCEED, + ), + is_sync=False, + ) + assert "child_1" in state._parent_done + finally: + release_child_hook.set() + + with pytest.raises(OrphanedChildException) as exc_info: + child_future.result(timeout=2.0) + + assert exc_info.value.operation_id == "child_1" + + def test_nested_parallel_operations_deep_hierarchy(): """Test that nested parallel operations handle deep hierarchies correctly.""" mock_lambda_client = Mock(spec=LambdaClient)