From 8f2729ea07f3c180247a75820c39f57aa1e4cf27 Mon Sep 17 00:00:00 2001 From: Harry Nguyen Date: Wed, 28 Oct 2026 10:56:15 -0700 Subject: [PATCH 1/5] test: cover an oversized operation in the overflow queue A step result above max_batch_size_bytes (750 KB) that reaches the overflow queue is never sent. Both tests fail on the current code with [] == [large_op]. Refs #871 --- .../tests/state_test.py | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) 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 128aa33cd..8bd5e6a96 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -56,6 +56,7 @@ from aws_durable_execution_sdk_python.state import ( _BLOCKED_CALLER_WAIT_SECONDS, _STOP_SIGNAL_POLL_SECONDS, + _BatchAccumulator, _Signal, CheckpointBatcherConfig, CheckpointedResult, @@ -1367,6 +1368,96 @@ def test_collect_checkpoint_batch_uses_overflow_queue(): assert batch[1].operation_update.operation_id == "overflow_2" +def _oversized_queued_operation( + operation_id: str, max_batch_size_bytes: int +) -> QueuedOperation: + """A sync operation whose serialized size alone exceeds the batch limit.""" + return QueuedOperation( + OperationUpdate( + operation_id=operation_id, + operation_type=OperationType.STEP, + action=OperationAction.SUCCEED, + payload="y" * (2 * max_batch_size_bytes), + ), + CompletionEvent(), + ) + + +def test_collect_checkpoint_batch_sends_oversized_operation_from_overflow(): + """An operation above max_batch_size_bytes that reached the overflow queue is + still sent, alone, in the next batch (issue #871). + + This is the default AT_LEAST_ONCE_PER_RETRY step shape: a small START is + already in the batch when the large SUCCEED arrives, so SUCCEED does not fit + and moves to the overflow queue. The next batch must send it rather than + return it to the queue forever. + """ + mock_lambda_client = Mock(spec=LambdaClient) + config = CheckpointBatcherConfig(max_batch_size_bytes=500) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + batcher_config=config, + ) + + small_op = QueuedOperation( + OperationUpdate( + operation_id="big_step", + operation_type=OperationType.STEP, + action=OperationAction.START, + ), + None, + ) + large_op = _oversized_queued_operation("big_step", config.max_batch_size_bytes) + assert state._calculate_operation_size(large_op) > config.max_batch_size_bytes + state._checkpoint_queue.put(small_op) + state._checkpoint_queue.put(large_op) + + first_batch = state._collect_checkpoint_batch() + assert first_batch == [small_op] + assert state._overflow_queue.qsize() == 1 + + # The handler thread is blocked on large_op, so no further operation will + # arrive. Stopping the batcher makes a regression return [] instead of + # blocking the test forever in _wait_for_first_operation. + state._checkpointing_stopped.set() + second_batch = state._collect_checkpoint_batch() + assert second_batch == [large_op] + assert state._overflow_queue.empty() + + +def test_drain_overflow_takes_oversized_operation_when_batch_is_empty(): + """_drain_overflow adds an operation above the size limit to an empty batch, + and refuses it once the batch holds another operation. + """ + mock_lambda_client = Mock(spec=LambdaClient) + config = CheckpointBatcherConfig(max_batch_size_bytes=500) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_lambda_client, + plugin_executor=PluginExecutor(plugins=None), + batcher_config=config, + ) + first_large = _oversized_queued_operation("large_1", config.max_batch_size_bytes) + second_large = _oversized_queued_operation("large_2", config.max_batch_size_bytes) + state._overflow_queue.put(first_large) + state._overflow_queue.put(second_large) + + batch = _BatchAccumulator(config) + state._drain_overflow(batch) + + # The first operation is taken despite its size. The second does not fit + # beside it and stays queued for the next batch. + assert batch.operations == [first_large] + assert state._overflow_queue.qsize() == 1 + assert state._overflow_queue.get_nowait() is second_large + + def test_collect_checkpoint_batch_handles_empty_checkpoint(): """Test batch collection with empty checkpoints (None operation_update).""" mock_lambda_client = Mock(spec=LambdaClient) From 079fa503822188e6af27c38379954015aefe155f Mon Sep 17 00:00:00 2001 From: Harry Nguyen Date: Wed, 28 Oct 2026 10:56:15 -0700 Subject: [PATCH 2/5] fix: send an oversized first operation in a batch _drain_overflow refused an operation above max_batch_size_bytes even when the batch was empty, so the operation returned to the overflow queue on every pass and was never sent. Its sync caller then blocked forever, and the execution hung until the Lambda timeout. _wait_for_first_operation already accepted an oversized first operation. The rule now lives in _BatchAccumulator.fits: an empty batch accepts an operation of any size. The service enforces the per-operation limit. Fixes #871 --- .../src/aws_durable_execution_sdk_python/state.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 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 e9a7588d2..a9883b3ea 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 @@ -154,6 +154,12 @@ def is_full(self) -> bool: return self.effective_count >= self._config.max_batch_operations def fits(self, size: int) -> bool: + # An empty batch accepts an operation of any size. An operation above + # the limit can only ever be sent alone, so refusing it here would + # return it to its queue on every pass, and its sync caller would block + # forever. The service, not the batcher, enforces the per-operation limit. + if self.is_empty(): + return True return self.total_size + size <= self._config.max_batch_size_bytes def add(self, op: QueuedOperation, size: int = 0) -> None: @@ -1487,8 +1493,7 @@ def _wait_for_first_operation(self, batch: _BatchAccumulator) -> bool: if isinstance(item, _Signal): self._consume_refresh_wake() continue - # The first operation is never refused for size. It could not be - # sent otherwise. + # The first operation is never refused for size; see fits(). batch.add(item, self._calculate_operation_size(item)) return True return False From 5c77fe1dae99d05cbfc1b2281f99c21f0eac026c Mon Sep 17 00:00:00 2001 From: Harry Nguyen Date: Wed, 28 Oct 2026 11:46:02 -0700 Subject: [PATCH 3/5] test: run an 800 KB step result end to end Runs the #871 reproduction through DurableFunctionTestRunner. The result is above max_batch_size_bytes (750 KB) and below the service's 1 MB limit. On the code before the fix, run() raises TimeoutError after 30 s. Refs #871 --- .../tests/e2e/large_step_result_int_test.py | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/large_step_result_int_test.py diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/large_step_result_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/large_step_result_int_test.py new file mode 100644 index 000000000..69ecc12de --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/large_step_result_int_test.py @@ -0,0 +1,44 @@ +"""A step result above the checkpoint batch size limit completes (issue #871). + +The step result is 800 KB. That is above ``max_batch_size_bytes`` (750 KB) and +below the service's 1 MB per-operation limit. With the default +``AT_LEAST_ONCE_PER_RETRY`` semantics the START checkpoint is async, so the +batcher still holds it when the sync SUCCEED arrives, and SUCCEED takes the +overflow path. Before the fix the SUCCEED operation was never sent and the +handler thread blocked forever, so ``run()`` raised ``TimeoutError``. + +The emulated Lambda timeout stays at its 900 s default on purpose. A shorter +one would re-invoke the handler, and the re-run SUCCEED would open a fresh +batch and be sent, which would hide the bug. + +On a regression the test fails after ``execution_timeout`` (30 s). In local +runs the pytest process then did not exit on its own. The likely cause is the +handler thread, which stays blocked on the SUCCEED checkpoint. +""" + +from __future__ import annotations + +import json + +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.execution import InvocationStatus +from aws_durable_execution_sdk_python_testing.runner import ( + DurableFunctionTestResult, + DurableFunctionTestRunner, +) + +RESULT_SIZE_BYTES = 800 * 1024 + + +@durable_execution +def _large_step_handler(_event: dict[str, str], context: DurableContext) -> str: + return context.step(lambda _sc: "y" * RESULT_SIZE_BYTES, name="big_step") + + +def test_step_result_above_batch_size_limit_completes() -> None: + with DurableFunctionTestRunner(handler=_large_step_handler) as runner: + result: DurableFunctionTestResult = runner.run(input="{}", execution_timeout=30) + + assert result.status is InvocationStatus.SUCCEEDED + assert result.result is not None + assert len(json.loads(result.result)) == RESULT_SIZE_BYTES From 1de039475d9c44b361ec7b46b07b6d2cee1e4625 Mon Sep 17 00:00:00 2001 From: Harry Nguyen Date: Wed, 28 Oct 2026 11:50:17 -0700 Subject: [PATCH 4/5] test: trim the large step result docstring --- .../tests/e2e/large_step_result_int_test.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/large_step_result_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/large_step_result_int_test.py index 69ecc12de..0d7eb3ac3 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/large_step_result_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/large_step_result_int_test.py @@ -6,14 +6,6 @@ batcher still holds it when the sync SUCCEED arrives, and SUCCEED takes the overflow path. Before the fix the SUCCEED operation was never sent and the handler thread blocked forever, so ``run()`` raised ``TimeoutError``. - -The emulated Lambda timeout stays at its 900 s default on purpose. A shorter -one would re-invoke the handler, and the re-run SUCCEED would open a fresh -batch and be sent, which would hide the bug. - -On a regression the test fails after ``execution_timeout`` (30 s). In local -runs the pytest process then did not exit on its own. The likely cause is the -handler thread, which stays blocked on the SUCCEED checkpoint. """ from __future__ import annotations From 028501308ec116cf9e4e4e08bd18ccfb3b49cdae Mon Sep 17 00:00:00 2001 From: Harry Nguyen Date: Wed, 28 Oct 2026 11:52:20 -0700 Subject: [PATCH 5/5] test: annotate the new overflow tests with -> None CONTRIBUTING asks for return annotations so that mypy checks the function bodies. --- packages/aws-durable-execution-sdk-python/tests/state_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 8bd5e6a96..45c41e53c 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -1383,7 +1383,7 @@ def _oversized_queued_operation( ) -def test_collect_checkpoint_batch_sends_oversized_operation_from_overflow(): +def test_collect_checkpoint_batch_sends_oversized_operation_from_overflow() -> None: """An operation above max_batch_size_bytes that reached the overflow queue is still sent, alone, in the next batch (issue #871). @@ -1429,7 +1429,7 @@ def test_collect_checkpoint_batch_sends_oversized_operation_from_overflow(): assert state._overflow_queue.empty() -def test_drain_overflow_takes_oversized_operation_when_batch_is_empty(): +def test_drain_overflow_takes_oversized_operation_when_batch_is_empty() -> None: """_drain_overflow adds an operation above the size limit to an empty batch, and refuses it once the batch holds another operation. """