Coverage for async_durable_execution/_core/state.py: 97%
284 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +0000
1"""Model for execution state."""
3from __future__ import annotations
5import asyncio
6import json
7import logging
8from collections import deque
9import time
10from dataclasses import dataclass
11from typing import TYPE_CHECKING, Any, cast
13from .exceptions import (
14 CheckpointError,
15 DurableExecutionsError,
16 GetExecutionStateError,
17 OrphanedChildException,
18)
19from .models import (
20 Operation,
21 OperationAction,
22 OperationType,
23 OperationUpdate,
24)
25from .client import DurableServiceClient
28if TYPE_CHECKING:
29 from collections.abc import MutableMapping
30 from .models import LambdaContext
32logger = logging.getLogger(__name__)
34RECURSIVE_LEVEL_INPUT_FIELD = "__recursive_level"
37@dataclass(frozen=True)
38class CheckpointBatcherConfig:
39 """Configuration for checkpoint batching behavior.
41 Attributes:
42 max_batch_size_bytes: Maximum batch size in bytes (default: 750KB)
43 max_batch_time_seconds: Maximum time to wait before flushing batch (default: 1.0 second)
44 max_batch_operations: Maximum number of operations per batch (default: 250)
45 """
47 max_batch_size_bytes: int = 750 * 1024 # 750KB
48 max_batch_time_seconds: float = 1.0
49 max_batch_operations: int = 250
52@dataclass(frozen=True)
53class QueuedOperation:
54 """Wrapper for operations in the checkpoint queue.
56 Attributes:
57 operation_update: The operation update to be checkpointed, or None for empty checkpoints
58 completion_future: Completion future for synchronous operations, or None for async operations
59 """
61 operation_update: OperationUpdate | None
62 completion_future: asyncio.Future[Operation | None] | None = None
65def _completion_done(completion) -> bool:
66 return completion is None or completion.done()
69def _completion_set_result(completion, operation: Operation | None = None) -> None:
70 if completion is None: 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true
71 return
72 completion.set_result(operation)
75def _completion_set_exception(completion, error: Exception) -> None:
76 if completion is None: 76 ↛ 77line 76 didn't jump to line 77 because the condition on line 76 was never true
77 return
78 completion.set_exception(error)
81class ExecutionState:
82 """Get, set and maintain execution state. This is mutable. Create and check checkpoints."""
84 # Keep operations visible on the class for Mock(spec=ExecutionState).
85 operations: MutableMapping[str, Operation] = cast(
86 "MutableMapping[str, Operation]", None
87 )
89 def __init__(
90 self,
91 durable_execution_arn: str,
92 initial_checkpoint_token: str,
93 service_client: DurableServiceClient,
94 lambda_context: LambdaContext | None = None,
95 batcher_config: CheckpointBatcherConfig | None = None,
96 operations: MutableMapping[str, Operation] | None = None,
97 ) -> None:
98 self.operations: MutableMapping[str, Operation] = dict(operations or {})
99 self.durable_execution_arn: str = durable_execution_arn
100 self.lambda_context: LambdaContext | None = lambda_context
101 self._current_checkpoint_token: str = initial_checkpoint_token
102 self._service_client: DurableServiceClient = service_client
104 # Checkpoint batching configuration
105 self._batcher_config = batcher_config or CheckpointBatcherConfig()
107 # Checkpoint batching components
108 self._checkpoint_queue: asyncio.Queue[QueuedOperation | None] = asyncio.Queue()
109 self._overflow_queue: deque[QueuedOperation] = deque()
110 self._checkpointing_stopped = asyncio.Event()
111 self._checkpointing_failed = asyncio.Event()
112 self._checkpointing_failure: Exception | None = None
113 self._checkpointing_task: asyncio.Task[None] | None = None
115 # Concurrency management for parallel operations: parent_id -> {child_operation_ids}
116 self._parent_to_children: dict[str, set[str]] = {}
118 # Operations whose parent has completed
119 self._parent_done: set[str] = set()
121 async def initialize(self, invocation_input) -> None:
122 await self.fetch_paginated_operations(
123 invocation_input.initial_execution_state.operations,
124 invocation_input.checkpoint_token,
125 invocation_input.initial_execution_state.next_marker,
126 )
128 async def fetch_paginated_operations(
129 self,
130 initial_operations: list[Operation],
131 checkpoint_token: str,
132 next_marker: str | None,
133 ) -> list[Operation]:
134 """Add initial operations and fetch all paginated operations from the Durable Functions API.
136 The checkpoint_token is passed explicitly as a parameter rather than using the
137 instance variable so pagination continues from the correct checkpoint state.
139 Args:
140 initial_operations: initial operations to be added to ExecutionState
141 checkpoint_token: checkpoint token used to call Durable Functions API.
142 next_marker: a marker indicates that there are paginated operations.
143 Returns:
144 List of all operations fetched from the Durable Functions API
146 Raises:
147 GetExecutionStateError: If the API call fails. The error is logged
148 with structured extras before re-raising. Callers are responsible
149 for deciding whether to fail the execution or allow Lambda retry
150 based on is_retryable().
151 """
152 all_operations: list[Operation] = (
153 initial_operations.copy() if initial_operations else []
154 )
155 try:
156 while next_marker:
157 output = await self._service_client.get_execution_state(
158 durable_execution_arn=self.durable_execution_arn,
159 checkpoint_token=checkpoint_token,
160 next_marker=next_marker,
161 )
162 all_operations.extend(output.operations)
163 next_marker = output.next_marker
164 except GetExecutionStateError as e:
165 logger.exception(
166 "Durable API error during state fetch.",
167 extra=e.build_logger_extras(),
168 )
169 raise
170 finally:
171 # Always store whatever operations we successfully fetched
172 if all_operations:
173 self.operations.update({op.operation_id: op for op in all_operations})
174 return all_operations
176 def get_raw_input_payload(self) -> str | None:
177 # It is possible that backend will not provide an execution operation
178 # for the initial page of results.
179 if not (operations := self.get_execution_operation()):
180 return None
181 if not (execution_details := operations.execution_details):
182 return None
183 return execution_details.input_payload
185 def get_input_event(self) -> Any:
186 # Python RIC LambdaMarshaller just uses standard json deserialization for event
187 # https://github.com/aws/aws-lambda-python-runtime-interface-client/blob/main/awslambdaric/lambda_runtime_marshaller.py#L46
188 raw_input_payload: str | None = self.get_raw_input_payload()
189 input_event: Any = {}
190 if raw_input_payload and raw_input_payload.strip():
191 try:
192 input_event = json.loads(raw_input_payload)
193 except json.JSONDecodeError:
194 logger.exception(
195 "Failed to parse input payload as JSON: payload: %r",
196 raw_input_payload,
197 )
198 raise
199 return input_event
201 @property
202 def recursive_level(self) -> int:
203 input_event = self.get_input_event()
204 if not isinstance(input_event, dict):
205 return 0
207 value = input_event.get(RECURSIVE_LEVEL_INPUT_FIELD, 0)
208 if isinstance(value, bool):
209 return 0
210 if isinstance(value, int):
211 return value
213 try:
214 return int(value)
215 except (TypeError, ValueError):
216 return 0
218 def get_execution_operation(self) -> Operation | None:
219 # invocation id is id of execution operation
220 invocation_id = self.durable_execution_arn.split("/")[-1]
221 candidate = self.operations.get(invocation_id)
222 if not candidate:
223 # Due to payload size limitations we may have an empty operations list.
224 # This will only happen when loading the initial page of results and is
225 # expected behaviour. We don't fail, but instead return None
226 # as the execution operation does not exist
227 msg: str = "No durable operations found in execution state."
228 logger.debug(msg)
229 return None
230 if candidate.operation_type is not OperationType.EXECUTION:
231 msg = f"The execution operation in execution state does not have EXECUTION type: {candidate.operation_type}"
232 raise DurableExecutionsError(msg)
234 return candidate
236 def has_prior_operations(self) -> bool:
237 """Return True if loaded state contains any non-execution operation."""
238 return any(
239 op.operation_type is not OperationType.EXECUTION
240 for op in tuple(self.operations.values())
241 )
243 async def create_checkpoint(
244 self,
245 operation_update: OperationUpdate | None = None,
246 is_sync: bool = True, # noqa: FBT001, FBT002
247 ) -> Operation | None:
248 """Create a checkpoint with optional synchronous behavior.
250 This method enqueues a checkpoint operation for processing by the background
251 batching thread. By default, the operation is synchronous (blocking) to ensure
252 the checkpoint is persisted before continuing. For performance-critical paths
253 where immediate confirmation is not required, set is_sync=False.
255 Synchronous checkpoints (is_sync=True, default):
256 - Block the caller until the checkpoint is processed by the background thread
257 - Ensure the checkpoint is persisted before continuing
258 - Safe default for correctness
259 - Use cases: Most operations requiring confirmation before proceeding
261 Asynchronous checkpoints (is_sync=False, opt-in):
262 - Return immediately without waiting for the checkpoint to complete
263 - Performance optimization for specific use cases
264 - Use cases: observability checkpoints, fire-and-forget operations
266 When to use synchronous checkpoints (is_sync=True, default):
267 1. Step START with AtMostOncePerRetry semantics - prevents duplicate execution
268 2. Operation completion (SUCCEED/FAIL) - ensures state persisted before returning
269 3. Retry operations - ensures retry state recorded before continuing
270 4. Callback START - must wait for API to generate callback ID
271 5. Invoke START - ensures chained invoke recorded before proceeding
272 6. Child context results - ensures results persisted before returning
273 7. Large results - ensures results saved before returning to caller
274 8. Wait for condition completion - ensures state recorded before proceeding
275 9. Most operations - safe default
277 When to use asynchronous checkpoints (is_sync=False, opt-in):
278 1. Step START with AtLeastOncePerRetry semantics - performance optimization
279 2. Child context START - fire-and-forget for performance
280 3. Wait for condition START - observability only, no blocking needed
281 4. Any checkpoint where immediate confirmation is not required AND performance matters
283 Args:
284 operation_update: The checkpoint to create. If None, creates an empty
285 checkpoint to get a fresh checkpoint token and updated
286 operations list.
287 is_sync: If True (default), blocks until the checkpoint is processed.
288 If False, returns immediately without blocking for performance.
290 Returns:
291 The updated operation for synchronous checkpoints with an operation update.
292 Returns None for asynchronous checkpoints and empty checkpoints.
294 Raises:
295 Any exception from checkpoint processing will propagate back to the
296 awaiting coroutine, terminating the Lambda invocation.
298 Examples:
299 # Synchronous checkpoint (default, safe)
300 execution_state.create_checkpoint(operation_update)
302 # Explicit synchronous checkpoint
303 execution_state.create_checkpoint(operation_update, is_sync=True)
305 # Asynchronous checkpoint (opt-in for performance)
306 execution_state.create_checkpoint(operation_update, is_sync=False)
308 # Empty checkpoint (sync by default)
309 execution_state.create_checkpoint()
311 # Empty checkpoint (async for performance)
312 execution_state.create_checkpoint(is_sync=False)
313 """
314 # if this is CONTEXT complete, mark incomplete descendants as orphans so the children can't complete after the parent
315 if operation_update is not None:
316 if operation_update.parent_id:
317 if operation_update.parent_id not in self._parent_to_children:
318 self._parent_to_children[operation_update.parent_id] = set()
319 self._parent_to_children[operation_update.parent_id].add(
320 operation_update.operation_id
321 )
323 if (
324 operation_update.operation_type == OperationType.CONTEXT
325 and operation_update.action
326 in {OperationAction.SUCCEED, OperationAction.FAIL}
327 ):
328 self._mark_orphans(operation_update.operation_id)
330 if operation_update.operation_id in self._parent_done:
331 logger.debug(
332 "Rejecting checkpoint for operation %s - parent is done",
333 operation_update.operation_id,
334 )
335 error_msg = (
336 "Parent context completed, child operation cannot checkpoint"
337 )
338 raise OrphanedChildException(
339 error_msg,
340 operation_id=operation_update.operation_id,
341 )
343 if self._checkpointing_failure is not None:
344 raise self._checkpointing_failure
346 if self._checkpointing_task is None or self._checkpointing_task.done():
347 self.start_checkpointing()
349 completion_future: asyncio.Future[Operation | None] | None = None
350 if is_sync:
351 completion_future = asyncio.get_running_loop().create_future()
353 # Create wrapper object for queue
354 queued_op = QueuedOperation(operation_update, completion_future)
356 # Enqueue the wrapper object (operation_update can be None for empty checkpoints)
357 await self._checkpoint_queue.put(queued_op)
359 # Conditionally wait for completion based on is_sync parameter
360 if is_sync:
361 logger.debug("Enqueued checkpoint operation for synchronous processing")
362 if completion_future is None: 362 ↛ 364line 362 didn't jump to line 364 because the condition on line 362 was never true
363 # this shouldn't ever be possible
364 msg: str = "completion_future must be set for synchronous execution"
365 raise DurableExecutionsError(msg)
366 return await completion_future
367 else:
368 logger.debug("Enqueued checkpoint operation for asynchronous processing")
369 return None
371 def _mark_orphans(self, context_id: str) -> None:
372 """Mark all descendants (direct and transitive) as orphaned.
374 This method uses BFS (Breadth-First Search) to recursively collect all
375 descendants of the given context operation and marks them as orphaned.
376 Once marked, these operations will be rejected if they attempt to checkpoint.
378 Args:
379 context_id: The operation ID of the CONTEXT that has completed
380 """
381 # Collect all descendants recursively using BFS
382 all_descendants = set()
383 # Start with root
384 to_process: set[str] = {context_id}
386 while to_process:
387 current_id = to_process.pop()
389 # Skip if already processed (avoid cycles, though shouldn't happen)
390 if current_id in all_descendants:
391 continue
393 all_descendants.add(current_id)
395 # Add all direct children to processing queue
396 direct_children = self._parent_to_children.get(current_id, set())
397 to_process.update(direct_children)
399 # Remove the root itself (we only want descendants)
400 all_descendants.discard(context_id)
402 # Mark all descendants as orphaned
403 self._parent_done.update(all_descendants)
404 logger.debug(
405 "Marked %d descendants as parent-done for context %s",
406 len(all_descendants),
407 context_id,
408 )
410 def start_checkpointing(self) -> None:
411 """Start the checkpoint processor on the current event loop."""
412 if self._checkpointing_task is None or self._checkpointing_task.done(): 412 ↛ exitline 412 didn't return from function 'start_checkpointing' because the condition on line 412 was always true
413 self._checkpointing_stopped.clear()
414 self._checkpointing_failed.clear()
415 self._checkpointing_task = asyncio.create_task(
416 self.checkpoint_batches_forever()
417 )
419 async def checkpoint_batches_forever(self) -> None:
420 """Background coroutine that batches operations and processes results.
422 Runs until shutdown is signaled. This method processes checkpoint operations
423 in batches, makes API calls to persist them, and updates the execution state
424 with the results.
426 The method maintains the checkpoint token locally and updates it after each
427 successful batch processing. It continues running until stop_checkpointing()
428 is called.
430 Note: When shutdown is signaled, only non-essential async checkpoints may remain
431 in the queue. All critical synchronous checkpoints (SUCCEED, FAIL, etc.) will
432 have already completed because the main thread blocks on them. Therefore, we
433 don't need to drain the queue - the Lambda timeout will handle cleanup.
435 Raises:
436 Any exception from the service client checkpoint call will propagate naturally,
437 terminating the background thread and signaling an error to the main thread.
438 """
439 await asyncio.sleep(0)
441 # Keep checkpoint token as local variable in the loop
442 current_checkpoint_token: str = self._current_checkpoint_token
444 while not self._checkpointing_stopped.is_set():
445 # Collect operations into a batch
446 batch: list[QueuedOperation] = await self._collect_checkpoint_batch()
448 if batch:
449 # Extract OperationUpdates, excluding empty checkpoints from API call
450 updates: list[OperationUpdate] = []
451 empty_count = 0
453 for q in batch:
454 if q.operation_update is not None:
455 updates.append(q.operation_update)
456 else:
457 empty_count += 1
459 logger.debug(
460 "Sending %d OperationUpdates out of %d operations, excluding %d empty checkpoints",
461 len(updates),
462 len(batch),
463 empty_count,
464 )
466 try:
467 # Make API call with batched operations
468 output = await self._service_client.checkpoint(
469 durable_execution_arn=self.durable_execution_arn,
470 checkpoint_token=current_checkpoint_token,
471 updates=updates,
472 client_token=None,
473 )
475 logger.debug("Checkpoint batch processed successfully")
477 if output.checkpoint_token:
478 current_checkpoint_token = output.checkpoint_token
479 else:
480 batch_is_terminal = any(
481 update.operation_type is OperationType.EXECUTION
482 and update.action
483 in {OperationAction.SUCCEED, OperationAction.FAIL}
484 for update in updates
485 )
486 if not batch_is_terminal:
487 msg = (
488 "Checkpoint response omitted the token outside of "
489 "execution completion."
490 )
491 raise CheckpointError(msg)
493 # Fetch new operations from the API before unblocking sync waiters
494 updated_operations = await self.fetch_paginated_operations(
495 output.new_execution_state.operations,
496 current_checkpoint_token,
497 output.new_execution_state.next_marker,
498 )
499 updated_operations_by_id = {
500 operation.operation_id: operation
501 for operation in updated_operations
502 }
504 # Signal completion for any synchronous operations
505 for queued_op in batch:
506 if not _completion_done(queued_op.completion_future):
507 operation = None
508 if queued_op.operation_update is not None:
509 operation = updated_operations_by_id.get(
510 queued_op.operation_update.operation_id
511 )
512 _completion_set_result(
513 queued_op.completion_future, operation
514 )
515 except Exception as e:
516 # Checkpoint failed - wake blocked coroutines so they can raise error
517 logger.exception("Checkpoint batch processing failed")
518 self._checkpointing_failure = e
519 self._checkpointing_failed.set()
521 # Signal completion futures for the failed batch
522 for queued_op in batch:
523 if not _completion_done(queued_op.completion_future):
524 _completion_set_exception(queued_op.completion_future, e)
526 while self._overflow_queue:
527 overflow_item = self._overflow_queue.popleft()
528 if not _completion_done(overflow_item.completion_future): 528 ↛ 526line 528 didn't jump to line 526 because the condition on line 528 was always true
529 _completion_set_exception(
530 overflow_item.completion_future, e
531 )
533 while not self._checkpoint_queue.empty():
534 queued_item: QueuedOperation | None = (
535 self._checkpoint_queue.get_nowait()
536 )
537 if queued_item is not None and not _completion_done( 537 ↛ 533line 537 didn't jump to line 533 because the condition on line 537 was always true
538 queued_item.completion_future
539 ):
540 _completion_set_exception(queued_item.completion_future, e)
541 break
543 logger.debug("Background checkpoint processing stopped")
545 def stop_checkpointing(self) -> None:
546 """Signal the checkpoint processor to stop.
548 This method sets the checkpointing stopped event, which signals the background
549 thread to exit. Any remaining async checkpoints in the queue are non-essential
550 (observability only) and will be abandoned. All critical synchronous checkpoints
551 will have already completed before this is called.
552 """
553 logger.debug("Signaling checkpoint processor to stop")
554 self._checkpointing_stopped.set()
555 if self._checkpointing_task is not None and not self._checkpoint_queue.full():
556 self._checkpoint_queue.put_nowait(None)
558 async def _collect_checkpoint_batch(self) -> list[QueuedOperation]:
559 """Collect multiple checkpoint operations into a batch for API efficiency.
561 Processes overflow queue first to maintain FIFO order, then collects from main queue.
562 Respects configured size, time, and operation count limits. Blocks for the first
563 operation if queues are empty, then collects additional operations within the time
564 window.
566 Empty checkpoints (operation_update=None) are coalesced: the first empty checkpoint
567 counts toward the batch operation limit, but subsequent empty checkpoints do not.
568 All empty checkpoints remain in the batch so their completion events are signaled.
569 This avoids unnecessary batches when many concurrent map/parallel branches resume
570 simultaneously and each queues an empty checkpoint.
572 Returns:
573 List of QueuedOperation objects ready for batch processing. Returns empty list
574 if no operations are available.
575 """
576 batch: list[QueuedOperation] = []
577 has_empty_checkpoint = False
578 total_size = 0
579 effective_operation_count = 0 # Operations that count toward batch limit
581 # First, drain overflow queue (FIFO order preserved)
582 while (
583 self._overflow_queue
584 and effective_operation_count < self._batcher_config.max_batch_operations
585 ):
586 overflow_op = self._overflow_queue.popleft()
588 if overflow_op.operation_update is None: # Empty checkpoint
589 batch.append(overflow_op)
590 if not has_empty_checkpoint:
591 effective_operation_count += 1
592 has_empty_checkpoint = True
593 else:
594 op_size = self._calculate_operation_size(overflow_op)
595 if total_size + op_size > self._batcher_config.max_batch_size_bytes:
596 self._overflow_queue.appendleft(overflow_op)
597 break
598 batch.append(overflow_op)
599 total_size += op_size
600 effective_operation_count += 1
602 # If batch is empty, get first operation from main queue
603 if not batch:
604 while not self._checkpointing_stopped.is_set():
605 try:
606 first_op = await asyncio.wait_for(
607 self._checkpoint_queue.get(), timeout=0.1
608 )
609 if first_op is None:
610 continue
611 batch.append(first_op)
613 if first_op.operation_update is None:
614 has_empty_checkpoint = True
615 else:
616 total_size += self._calculate_operation_size(first_op)
618 effective_operation_count = 1
619 break
620 except asyncio.TimeoutError:
621 continue
623 # If stopped and no operation retrieved, return empty batch
624 if not batch:
625 return batch
627 # Start batching window using configured time
628 batch_deadline = time.time() + self._batcher_config.max_batch_time_seconds
630 # Collect additional operations within the time window. Once the batch
631 # reaches the real-operation limit, keep accepting empty checkpoints so
632 # concurrent resubmits can coalesce instead of spilling into a new API call.
633 while time.time() < batch_deadline and not self._checkpointing_stopped.is_set():
634 remaining_time = min(
635 batch_deadline - time.time(),
636 0.1, # Check stop signal every 100ms
637 )
639 if remaining_time <= 0:
640 break
642 try:
643 additional_op = await asyncio.wait_for(
644 self._checkpoint_queue.get(), timeout=remaining_time
645 )
646 if additional_op is None:
647 continue
649 if additional_op.operation_update is None: # Empty checkpoint
650 batch.append(additional_op)
651 if not has_empty_checkpoint:
652 effective_operation_count += (
653 1 # First empty counts toward limit
654 )
655 has_empty_checkpoint = True
656 # Subsequent empties don't count toward limit
657 else:
658 if (
659 effective_operation_count
660 >= self._batcher_config.max_batch_operations
661 ):
662 self._overflow_queue.append(additional_op)
663 logger.debug(
664 "Batch operation limit reached, moving operation to overflow queue"
665 )
666 break
668 op_size = self._calculate_operation_size(additional_op)
669 # Check if adding this operation would exceed size limit
670 if total_size + op_size > self._batcher_config.max_batch_size_bytes:
671 # Put in overflow queue for next batch
672 self._overflow_queue.append(additional_op)
673 logger.debug(
674 "Batch size limit reached, moving operation to overflow queue"
675 )
676 break
677 batch.append(additional_op)
678 total_size += op_size
679 effective_operation_count += 1
681 except asyncio.TimeoutError:
682 break
684 empty_count = sum(1 for q in batch if q.operation_update is None)
685 logger.debug(
686 "Collected batch of %d operations (%d effective, %d non-empty, %d empty), total size: %d bytes",
687 len(batch),
688 effective_operation_count,
689 len(batch) - empty_count,
690 empty_count,
691 total_size,
692 )
693 return batch
695 @staticmethod
696 def _calculate_operation_size(queued_op: QueuedOperation) -> int:
697 """Calculate the serialized size of a queued operation for batching limits.
699 Uses JSON serialization to estimate the size of the operation update. Empty
700 checkpoints (None operation_update) have zero size.
702 Args:
703 queued_op: The queued operation to calculate size for
705 Returns:
706 Size in bytes of the serialized operation, or 0 for empty checkpoints
707 """
708 # Empty checkpoints have no size
709 if queued_op.operation_update is None:
710 return 0
712 # Use JSON serialization to estimate size
713 serialized = json.dumps(queued_op.operation_update.to_dict()).encode("utf-8")
714 return len(serialized)
716 async def aclose(self) -> None:
717 self.stop_checkpointing()
718 if self._checkpointing_task is not None:
719 await self._checkpointing_task
721 def close(self) -> None:
722 self.stop_checkpointing()