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