Coverage for async-durable-execution/src/async_durable_execution/runner/local/executor.py: 98%
511 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 16:54 +0000
1"""Execution life-cycle logic."""
3from __future__ import annotations
5import logging
6import uuid
7from datetime import datetime, timezone
8from typing import TYPE_CHECKING
10from ...execution import (
11 DurableExecutionInvocationInput,
12 DurableExecutionInvocationOutput,
13 InvocationStatus,
14)
15from ...models import (
16 CallbackOptions,
17 CallbackTimeoutType,
18 CheckpointUpdatedExecutionState,
19 ErrorObject,
20 Operation,
21 OperationStatus,
22 OperationType,
23 OperationUpdate,
24)
25from ..exceptions import (
26 IllegalStateException,
27 InvalidParameterValueException,
28 ResourceNotFoundException,
29)
30from ..model import (
31 TERMINAL_STATUSES,
32 EventCreationContext,
33 GetDurableExecutionHistoryResponse,
34)
35from ..model import (
36 Event as HistoryEvent,
37)
38from .model import (
39 CallbackToken,
40 CheckpointDurableExecutionResponse,
41 GetDurableExecutionStateResponse,
42 Invoker,
43 SendDurableExecutionCallbackFailureResponse,
44 SendDurableExecutionCallbackHeartbeatResponse,
45 SendDurableExecutionCallbackSuccessResponse,
46 StartDurableExecutionInput,
47 StartDurableExecutionOutput,
48)
49from .time_scale import scale_delay
50from .execution import Execution
52if TYPE_CHECKING:
53 from collections.abc import Awaitable, Callable
54 from asyncio import Future
56 from . import InMemoryServiceClient
57 from .scheduler import Event, Scheduler
59logger = logging.getLogger(__name__)
61_CALLBACK_TIMEOUT_MINIMUM_DELAY_SECONDS = 5.0
64class Executor:
65 MAX_CONSECUTIVE_FAILED_ATTEMPTS: int = 5
66 RETRY_BACKOFF_SECONDS: int = 5
68 def __init__(
69 self,
70 scheduler: Scheduler,
71 invoker: Invoker,
72 service_client: InMemoryServiceClient,
73 ):
74 self._scheduler = scheduler
75 self._invoker = invoker
76 self._service_client = service_client
77 self._execution: Execution | None = None
78 self._execution_arn: str | None = None
79 self._completion_event: Event | None = None
80 self._callback_timeouts: dict[str, Future] = {}
81 self._callback_heartbeats: dict[str, Future] = {}
82 self._execution_timeout: Future | None = None
83 self._active_invocation: bool = False
84 self._scheduled_resume: bool = False
85 self._pending_resume: bool = False
86 self._deferred_wait_resumes: set[str] = set()
87 self._deferred_retry_resumes: set[str] = set()
88 self._deferred_callback_timeouts: set[tuple[str, CallbackTimeoutType]] = set()
90 def start_execution(
91 self,
92 input: StartDurableExecutionInput, # noqa: A002
93 ) -> StartDurableExecutionOutput:
94 self._ensure_can_start_execution()
96 # Generate invocation_id if not provided
97 if input.invocation_id is None:
98 input = StartDurableExecutionInput(
99 account_id=input.account_id,
100 function_name=input.function_name,
101 function_qualifier=input.function_qualifier,
102 execution_name=input.execution_name,
103 execution_timeout_seconds=input.execution_timeout_seconds,
104 execution_retention_period_days=input.execution_retention_period_days,
105 invocation_id=str(uuid.uuid4()),
106 trace_fields=input.trace_fields,
107 tenant_id=input.tenant_id,
108 input=input.input,
109 lambda_endpoint=input.lambda_endpoint,
110 )
112 execution = Execution.new(input=input)
113 execution.start()
114 self._save_execution(execution)
115 logger.debug("Created execution with ARN: %s", execution.durable_execution_arn)
117 completion_event = self._scheduler.create_event()
118 self._execution_arn = execution.durable_execution_arn
119 self._completion_event = completion_event
120 self._execution_timeout = None
121 self._active_invocation = False
122 self._scheduled_resume = False
123 self._pending_resume = False
124 self._deferred_wait_resumes.clear()
125 self._deferred_retry_resumes.clear()
126 self._deferred_callback_timeouts.clear()
128 # Schedule execution timeout
129 if input.execution_timeout_seconds > 0:
131 async def timeout_handler() -> None:
132 error = ErrorObject.from_message(
133 f"Execution timed out after {input.execution_timeout_seconds} seconds."
134 )
135 self.timeout_execution(execution.durable_execution_arn, error)
137 self._execution_timeout = self._scheduler.call_later(
138 timeout_handler,
139 delay=input.execution_timeout_seconds,
140 completion_event=completion_event,
141 )
143 # Schedule initial invocation to run immediately
144 self._invoke_execution(execution.durable_execution_arn)
146 return StartDurableExecutionOutput(
147 execution_arn=execution.durable_execution_arn
148 )
150 def _ensure_can_start_execution(self) -> None:
151 """Reject overlapping executions in the local runner."""
152 if self._execution is None:
153 return
155 if not self._execution.is_complete:
156 msg = "Local runner supports only one execution at a time."
157 raise IllegalStateException(msg)
159 def _save_execution(self, execution: Execution) -> None:
160 self._execution = execution
161 self._execution_arn = execution.durable_execution_arn
163 def _update_execution(self, execution: Execution) -> None:
164 self._execution = execution
165 self._execution_arn = execution.durable_execution_arn
167 def set_execution(self, execution: Execution) -> None:
168 """Replace the current local execution object."""
169 self._update_execution(execution)
171 def _validate_current_execution(self, execution_arn: str) -> None:
172 if self._execution_arn is None:
173 return
175 if self._execution_arn != execution_arn:
176 msg = f"Execution {execution_arn} is not the active local execution."
177 raise ResourceNotFoundException(msg)
179 def get_execution(self, execution_arn: str) -> Execution:
180 """Get execution by ARN.
182 Args:
183 execution_arn: The execution ARN to retrieve
185 Returns:
186 Execution: The execution object
188 Raises:
189 ResourceNotFoundException: If execution does not exist
190 """
191 if (
192 self._execution is None
193 or self._execution.durable_execution_arn != execution_arn
194 ):
195 msg: str = f"Execution {execution_arn} not found"
196 raise ResourceNotFoundException(msg)
197 return self._execution
199 def get_execution_state(
200 self,
201 execution_arn: str,
202 checkpoint_token: str | None = None,
203 marker: str | None = None,
204 max_items: int | None = None,
205 ) -> GetDurableExecutionStateResponse:
206 """Get execution state with operations.
208 Args:
209 execution_arn: The execution ARN
210 checkpoint_token: Checkpoint token for state consistency
211 marker: Pagination marker
212 max_items: Maximum items to return
214 Returns:
215 GetDurableExecutionStateResponse: Execution state with operations
217 Raises:
218 ResourceNotFoundException: If execution does not exist
219 InvalidParameterValueException: If checkpoint token is invalid
220 """
221 execution = self.get_execution(execution_arn)
223 # TODO: Validate checkpoint token if provided
224 if checkpoint_token and checkpoint_token not in execution.used_tokens:
225 msg: str = f"Invalid checkpoint token: {checkpoint_token}"
226 raise InvalidParameterValueException(msg)
228 # Get operations (excluding the initial EXECUTION operation for state)
229 operations = execution.get_assertable_operations()
231 # Apply pagination
232 if max_items is None:
233 max_items = 100
235 # Simple pagination - in real implementation would need proper marker handling
236 start_index = 0
237 if marker:
238 try:
239 start_index = int(marker)
240 except ValueError:
241 start_index = 0
243 end_index = start_index + max_items
244 paginated_operations = operations[start_index:end_index]
246 next_marker = None
247 if end_index < len(operations):
248 next_marker = str(end_index)
250 return GetDurableExecutionStateResponse(
251 operations=paginated_operations, next_marker=next_marker
252 )
254 def get_execution_history(
255 self,
256 execution_arn: str,
257 include_execution_data: bool = False, # noqa: FBT001, FBT002
258 reverse_order: bool = False, # noqa: FBT001, FBT002
259 marker: str | None = None,
260 max_items: int | None = None,
261 ) -> GetDurableExecutionHistoryResponse:
262 """Get execution history with events.
264 Args:
265 execution_arn: The execution ARN
266 include_execution_data: Whether to include execution data in events
267 reverse_order: Return events in reverse chronological order
268 marker: Pagination marker (event_id)
269 max_items: Maximum items to return
271 Returns:
272 GetDurableExecutionHistoryResponse: Execution history with events
274 Raises:
275 ResourceNotFoundException: If execution does not exist
276 """
277 execution: Execution = self.get_execution(execution_arn)
279 # Generate events
280 all_events: list[HistoryEvent] = []
281 ops: list[Operation] = execution.operations
282 updates: list[OperationUpdate] = execution.updates
283 updates_dict: dict[str, OperationUpdate] = {u.operation_id: u for u in updates}
284 durable_execution_arn: str = execution.durable_execution_arn
286 # Add InvocationCompleted events
287 for completion in execution.invocation_completions:
288 invocation_event = HistoryEvent.create_invocation_completed(
289 event_id=0, # Temporary, will be reassigned
290 event_timestamp=completion.end_timestamp,
291 start_timestamp=completion.start_timestamp,
292 end_timestamp=completion.end_timestamp,
293 request_id=completion.request_id,
294 )
295 all_events.append(invocation_event)
297 # Generate all events first (without final event IDs)
298 for op in ops:
299 operation_update: OperationUpdate | None = updates_dict.get(op.operation_id)
301 if op.status is OperationStatus.PENDING:
302 if (
303 op.operation_type is not OperationType.CHAINED_INVOKE
304 or op.start_timestamp is None
305 ):
306 continue
307 context: EventCreationContext = EventCreationContext(
308 op,
309 0, # Temporary event_id, will be reassigned after sorting
310 durable_execution_arn,
311 execution.start_input,
312 execution.result,
313 operation_update,
314 include_execution_data,
315 )
316 pending = HistoryEvent.create_chained_invoke_event_pending(context)
317 all_events.append(pending)
318 if op.start_timestamp is not None:
319 context = EventCreationContext(
320 op,
321 0, # Temporary event_id, will be reassigned after sorting
322 durable_execution_arn,
323 execution.start_input,
324 execution.result,
325 operation_update,
326 include_execution_data,
327 )
328 started = HistoryEvent.create_event_started(context)
329 all_events.append(started)
330 if op.end_timestamp is not None and op.status in TERMINAL_STATUSES:
331 context = EventCreationContext(
332 op,
333 0, # Temporary event_id, will be reassigned after sorting
334 durable_execution_arn,
335 execution.start_input,
336 execution.result,
337 operation_update,
338 include_execution_data,
339 )
340 finished = HistoryEvent.create_event_terminated(context)
341 all_events.append(finished)
343 # Sort events by timestamp to get correct chronological order
344 all_events.sort(key=lambda event: event.event_timestamp)
346 # Reassign event IDs based on chronological order
347 all_events = [
348 HistoryEvent.from_event_with_id(event, i)
349 for i, event in enumerate(all_events, 1)
350 ]
352 # Apply cursor-based pagination
353 if max_items is None:
354 max_items = 100
356 # Handle pagination marker
357 if reverse_order:
358 all_events.reverse()
359 start_index: int = 0
360 if marker:
361 try:
362 marker_event_id: int = int(marker)
363 # Find the index of the first event with event_id >= marker
364 start_index = len(all_events)
365 for i, e in enumerate(all_events):
366 is_valid_page_start: bool = (
367 e.event_id < marker_event_id
368 if reverse_order
369 else e.event_id >= marker_event_id
370 )
371 if is_valid_page_start:
372 start_index = i
373 break
374 except ValueError:
375 start_index = 0
377 # Get paginated events
378 end_index: int = start_index + max_items
379 paginated_events: list[HistoryEvent] = all_events[start_index:end_index]
381 # Generate next marker
382 next_marker: str | None = None
383 if end_index < len(all_events):
384 if reverse_order:
385 # Next marker is the event_id of the last returned event
386 next_marker = (
387 str(paginated_events[-1].event_id) if paginated_events else None
388 )
389 else:
390 # Next marker is the event_id of the next event after the last returned
391 next_marker = (
392 str(all_events[end_index].event_id)
393 if end_index < len(all_events)
394 else None
395 )
397 return GetDurableExecutionHistoryResponse(
398 events=paginated_events, next_marker=next_marker
399 )
401 def checkpoint_execution(
402 self,
403 execution_arn: str,
404 checkpoint_token: str,
405 updates: list[OperationUpdate] | None = None,
406 client_token: str | None = None,
407 ) -> CheckpointDurableExecutionResponse:
408 """Process checkpoint for an execution.
410 Args:
411 execution_arn: The execution ARN
412 checkpoint_token: Current checkpoint token
413 updates: List of operation updates to process
414 client_token: Client token for idempotency
416 Returns:
417 CheckpointDurableExecutionResponse: Updated checkpoint token and state
419 Raises:
420 ResourceNotFoundException: If execution does not exist
421 InvalidParameterValueException: If checkpoint token is invalid
422 """
423 execution = self.get_execution(execution_arn)
425 # Validate checkpoint token
426 if checkpoint_token not in execution.used_tokens:
427 msg: str = f"Invalid checkpoint token: {checkpoint_token}"
428 raise InvalidParameterValueException(msg)
430 if updates:
431 checkpoint_output = self._service_client.process_checkpoint(
432 checkpoint_token=checkpoint_token,
433 updates=updates,
434 client_token=client_token,
435 )
437 new_execution_state = None
438 if checkpoint_output.new_execution_state:
439 new_execution_state = CheckpointUpdatedExecutionState(
440 operations=checkpoint_output.new_execution_state.operations,
441 next_marker=checkpoint_output.new_execution_state.next_marker,
442 )
444 return CheckpointDurableExecutionResponse(
445 checkpoint_token=checkpoint_output.checkpoint_token,
446 new_execution_state=new_execution_state,
447 )
449 # Save execution state after generating new token
450 new_checkpoint_token = execution.get_new_checkpoint_token()
451 self._update_execution(execution)
453 return CheckpointDurableExecutionResponse(
454 checkpoint_token=new_checkpoint_token,
455 new_execution_state=None,
456 )
458 def send_callback_success(
459 self,
460 callback_id: str,
461 result: bytes | None = None,
462 ) -> SendDurableExecutionCallbackSuccessResponse:
463 """Send callback success response.
465 Args:
466 callback_id: The callback ID to respond to
467 result: Optional result data for the callback
469 Returns:
470 SendDurableExecutionCallbackSuccessResponse: Empty response
472 Raises:
473 InvalidParameterValueException: If callback_id is invalid
474 ResourceNotFoundException: If callback does not exist
475 """
476 if not callback_id:
477 msg: str = "callback_id is required"
478 raise InvalidParameterValueException(msg)
480 try:
481 callback_token = CallbackToken.from_str(callback_id)
482 execution = self.get_execution(callback_token.execution_arn)
483 execution.complete_callback_success(callback_id, result)
484 self._update_execution(execution)
485 self._cleanup_callback_timeouts(callback_id)
486 self._schedule_resume(callback_token.execution_arn)
487 logger.info("Callback success completed for callback_id: %s", callback_id)
488 except Exception as e:
489 msg = f"Failed to process callback success: {e}"
490 raise ResourceNotFoundException(msg) from e
492 return SendDurableExecutionCallbackSuccessResponse()
494 def send_callback_failure(
495 self,
496 callback_id: str,
497 error: ErrorObject | None = None,
498 ) -> SendDurableExecutionCallbackFailureResponse:
499 """Send callback failure response.
501 Args:
502 callback_id: The callback ID to respond to
503 error: Optional error object for the callback failure
505 Returns:
506 SendDurableExecutionCallbackFailureResponse: Empty response
508 Raises:
509 InvalidParameterValueException: If callback_id is invalid
510 ResourceNotFoundException: If callback does not exist
511 """
512 if not callback_id:
513 msg: str = "callback_id is required"
514 raise InvalidParameterValueException(msg)
516 callback_error: ErrorObject = error or ErrorObject.from_message("")
518 try:
519 callback_token: CallbackToken = CallbackToken.from_str(callback_id)
520 execution: Execution = self.get_execution(callback_token.execution_arn)
521 execution.complete_callback_failure(callback_id, callback_error)
522 self._update_execution(execution)
523 self._cleanup_callback_timeouts(callback_id)
524 self._schedule_resume(callback_token.execution_arn)
525 logger.info("Callback failure completed for callback_id: %s", callback_id)
526 except Exception as e:
527 msg = f"Failed to process callback failure: {e}"
528 raise ResourceNotFoundException(msg) from e
530 return SendDurableExecutionCallbackFailureResponse()
532 def send_callback_heartbeat(
533 self, callback_id: str
534 ) -> SendDurableExecutionCallbackHeartbeatResponse:
535 """Send callback heartbeat to keep callback alive.
537 Args:
538 callback_id: The callback ID to send heartbeat for
540 Returns:
541 SendDurableExecutionCallbackHeartbeatResponse: Empty response
543 Raises:
544 InvalidParameterValueException: If callback_id is invalid
545 ResourceNotFoundException: If callback does not exist
546 """
547 if not callback_id:
548 msg: str = "callback_id is required"
549 raise InvalidParameterValueException(msg)
551 try:
552 callback_token: CallbackToken = CallbackToken.from_str(callback_id)
553 execution: Execution = self.get_execution(callback_token.execution_arn)
555 # Find callback operation to verify it exists and is active
556 _, operation = execution.find_callback_operation(callback_id)
557 if operation.status != OperationStatus.STARTED:
558 msg = f"Callback {callback_id} is not active"
559 raise ResourceNotFoundException(msg)
561 # Reset heartbeat timeout if configured
562 self._reset_callback_heartbeat_timeout(
563 callback_id, execution.durable_execution_arn
564 )
565 logger.info("Callback heartbeat processed for callback_id: %s", callback_id)
566 except Exception as e:
567 msg = f"Failed to process callback heartbeat: {e}"
568 raise ResourceNotFoundException(msg) from e
570 return SendDurableExecutionCallbackHeartbeatResponse()
572 def _validate_invocation_response(
573 self,
574 execution_arn: str,
575 response: DurableExecutionInvocationOutput,
576 execution: Execution,
577 ):
578 """Validate response status and apply the resulting execution changes.
580 Raises:
581 InvalidParameterValueException: If the response status is invalid.
582 IllegalStateException: If the response status is valid but the execution is already completed.
583 """
584 if execution.is_complete:
585 msg_already_complete: str = "Execution already completed, ignoring result"
587 raise IllegalStateException(msg_already_complete)
589 if response.status is None:
590 msg_status_required: str = "Response status is required"
592 raise InvalidParameterValueException(msg_status_required)
594 match response.status:
595 case InvocationStatus.FAILED:
596 if response.result is not None:
597 msg_failed_result: str = (
598 "Cannot provide a Result for FAILED status."
599 )
600 raise InvalidParameterValueException(msg_failed_result)
601 logger.info("[%s] Execution failed", execution_arn)
602 self._complete_workflow(
603 execution_arn, result=None, error=response.error
604 )
606 case InvocationStatus.SUCCEEDED:
607 if response.error is not None:
608 msg_success_error: str = (
609 "Cannot provide an Error for SUCCEEDED status."
610 )
611 raise InvalidParameterValueException(msg_success_error)
612 logger.info("[%s] Execution succeeded", execution_arn)
613 self._complete_workflow(
614 execution_arn, result=response.result, error=None
615 )
617 case InvocationStatus.PENDING:
618 if not execution.has_pending_operations():
619 msg_pending_ops: str = (
620 "Cannot return PENDING status with no pending operations."
621 )
622 raise InvalidParameterValueException(msg_pending_ops)
623 logger.info("[%s] Execution pending async work", execution_arn)
625 case _:
626 msg_unexpected_status: str = (
627 f"Unexpected invocation status: {response.status}"
628 )
629 raise IllegalStateException(msg_unexpected_status)
631 def _invoke_handler(self, execution_arn: str) -> Callable[[], Awaitable[None]]:
632 """Create a parameterless callable that captures execution arn for the scheduler."""
634 async def invoke() -> None:
635 self._mark_invocation_started(execution_arn)
636 execution: Execution = self.get_execution(execution_arn)
638 # Early exit if execution is already completed - like Java's COMPLETED check
639 if execution.is_complete:
640 logger.info(
641 "[%s] Execution already completed, ignoring result", execution_arn
642 )
643 return
645 try:
646 checkpoint_token = execution.get_new_checkpoint_token()
647 invocation_input: DurableExecutionInvocationInput = (
648 self._invoker.create_invocation_input(
649 start_input=execution.start_input,
650 durable_execution_arn=execution.durable_execution_arn,
651 checkpoint_token=checkpoint_token,
652 operations=execution.operations,
653 )
654 )
656 self._save_execution(execution)
658 invocation_start = datetime.now(timezone.utc)
659 invoke_response = await self._invoker.invoke(
660 execution.start_input.function_name,
661 invocation_input,
662 execution.start_input.lambda_endpoint,
663 )
664 invocation_end = datetime.now(timezone.utc)
666 # Reload execution after invocation in case it was completed via checkpoint
667 execution = self.get_execution(execution_arn)
669 # Record invocation completion and save immediately
670 execution.record_invocation_completion(
671 invocation_start, invocation_end, invoke_response.request_id
672 )
673 self._save_execution(execution)
675 if execution.is_complete:
676 logger.info(
677 "[%s] Execution completed during invocation, ignoring result",
678 execution_arn,
679 )
680 return
682 # Process successful received response - validate status and handle accordingly
683 response = invoke_response.invocation_output
684 try:
685 self._validate_invocation_response(
686 execution_arn, response, execution
687 )
688 except (InvalidParameterValueException, IllegalStateException) as e:
689 logger.warning(
690 "[%s] Lambda output validation failure: %s", execution_arn, e
691 )
692 error_obj = ErrorObject.from_exception(e)
693 self._retry_invocation(execution, error_obj)
695 except ResourceNotFoundException:
696 logger.warning(
697 "[%s] Function No longer exists: %s",
698 execution_arn,
699 execution.start_input.function_name,
700 )
701 error_obj = ErrorObject.from_message(
702 message=f"Function not found: {execution.start_input.function_name}"
703 )
704 self._fail_workflow(execution_arn, error_obj)
706 except Exception as e: # noqa: BLE001
707 # Handle invocation errors (network, function not found, etc.)
708 logger.warning("[%s] Invocation failed: %s", execution_arn, e)
709 error_obj = ErrorObject.from_exception(e)
710 self._retry_invocation(execution, error_obj)
711 finally:
712 self._mark_invocation_finished(execution_arn)
714 return invoke
716 def _schedule_callback_resume(self, execution_arn: str) -> None:
717 """Schedule a callback-triggered resume.
719 Kept as a named wrapper for tests and callers that exercise callback behavior.
720 """
721 self._schedule_resume(execution_arn)
723 def _schedule_resume(self, execution_arn: str) -> None:
724 """Coalesce external resumes to avoid overlapping replays."""
725 self._validate_current_execution(execution_arn)
726 if self._active_invocation:
727 self._pending_resume = True
728 return
730 if self._scheduled_resume:
731 return
733 self._scheduled_resume = True
734 self._invoke_execution(execution_arn)
736 def _mark_invocation_started(self, execution_arn: str) -> None:
737 self._validate_current_execution(execution_arn)
738 self._scheduled_resume = False
739 self._active_invocation = True
741 def _mark_invocation_finished(self, execution_arn: str) -> None:
742 self._validate_current_execution(execution_arn)
743 self._active_invocation = False
744 should_resume = self._apply_deferred_resume_events(execution_arn)
745 if self._pending_resume:
746 self._pending_resume = False
747 should_resume = True
749 if should_resume:
750 execution = self.get_execution(execution_arn)
751 if not execution.is_complete:
752 self._scheduled_resume = True
753 self._invoke_execution(execution_arn)
755 def _invoke_execution(self, execution_arn: str, delay: float = 0) -> None:
756 """Invoke execution after delay in seconds."""
757 self._validate_current_execution(execution_arn)
758 self._scheduler.call_later(
759 self._invoke_handler(execution_arn),
760 delay=delay,
761 completion_event=self._completion_event,
762 )
764 def _complete_workflow(
765 self, execution_arn: str, result: str | None, error: ErrorObject | None
766 ):
767 """Complete workflow - handles both success and failure with terminal state validation."""
768 execution = self.get_execution(execution_arn)
770 if execution.is_complete:
771 msg: str = "Cannot make multiple close workflow decisions."
773 raise IllegalStateException(msg)
775 if error is not None:
776 self.fail_execution(execution_arn, error)
777 else:
778 self.complete_execution(execution_arn, result)
780 def _fail_workflow(self, execution_arn: str, error: ErrorObject):
781 """Fail workflow with terminal state validation."""
782 execution = self.get_execution(execution_arn)
784 if execution.is_complete:
785 msg: str = "Cannot make multiple close workflow decisions."
787 raise IllegalStateException(msg)
789 self.fail_execution(execution_arn, error)
791 def _retry_invocation(self, execution: Execution, error: ErrorObject):
792 """Handle retry logic or fail execution if retries exhausted."""
793 if (
794 execution.consecutive_failed_invocation_attempts
795 > self.MAX_CONSECUTIVE_FAILED_ATTEMPTS
796 ):
797 # Exhausted retries - fail the execution
798 self._fail_workflow(
799 execution_arn=execution.durable_execution_arn, error=error
800 )
801 else:
802 # Schedule retry with backoff
803 execution.consecutive_failed_invocation_attempts += 1
804 self._save_execution(execution)
805 self._invoke_execution(
806 execution_arn=execution.durable_execution_arn,
807 delay=self.RETRY_BACKOFF_SECONDS,
808 )
810 def _complete_events(self, execution_arn: str):
811 self._validate_current_execution(execution_arn)
812 # complete doesn't actually checkpoint explicitly
813 if self._completion_event:
814 self._completion_event.set()
815 if execution_timeout := self._execution_timeout:
816 execution_timeout.cancel()
817 self._execution_timeout = None
819 async def wait_until_complete(
820 self, execution_arn: str, timeout: float | None = None
821 ) -> bool:
822 """Wait until execution completion.
824 Args
825 timeout (int|float|None): Wait for event to set until this timeout.
827 Returns:
828 True when set. False if the event timed out without being set.
829 """
830 self._validate_current_execution(execution_arn)
831 if event := self._completion_event:
832 return await event.wait_async(timeout)
834 # this really shouldn't happen - implies execution timed out?
835 msg: str = "execution does not exist."
837 raise ResourceNotFoundException(msg)
839 def complete_execution(self, execution_arn: str, result: str | None = None) -> None:
840 """Complete execution successfully (COMPLETE_WORKFLOW_EXECUTION decision)."""
841 logger.debug("[%s] Completing execution with result: %s", execution_arn, result)
842 execution: Execution = self.get_execution(execution_arn)
843 execution.complete_success(result=result) # Sets CloseStatus.COMPLETED
844 self._update_execution(execution)
845 if execution.result is None:
846 msg: str = "Execution result is required"
847 raise IllegalStateException(msg)
848 self._complete_events(execution_arn=execution_arn)
850 def fail_execution(self, execution_arn: str, error: ErrorObject) -> None:
851 """Fail execution with error (FAIL_WORKFLOW_EXECUTION decision)."""
852 logger.error("[%s] Completing execution with error: %s", execution_arn, error)
853 execution: Execution = self.get_execution(execution_arn)
854 execution.complete_fail(error=error) # Sets CloseStatus.FAILED
855 self._update_execution(execution)
856 # set by complete_fail
857 if execution.result is None:
858 msg: str = "Execution result is required"
859 raise IllegalStateException(msg)
860 self._complete_events(execution_arn=execution_arn)
862 def _on_wait_succeeded(self, execution_arn: str, operation_id: str) -> bool:
863 """Private method - called when a wait operation completes successfully."""
864 execution = self.get_execution(execution_arn)
866 if execution.is_complete:
867 logger.info(
868 "[%s] Execution already completed, ignoring wait succeeded event",
869 execution_arn,
870 )
871 return False
873 try:
874 execution.complete_wait(operation_id=operation_id)
875 self._update_execution(execution)
876 logger.debug(
877 "[%s] Wait succeeded for operation %s", execution_arn, operation_id
878 )
879 return True
880 except Exception:
881 logger.exception("[%s] Error processing wait succeeded.", execution_arn)
882 return False
884 def _on_retry_ready(self, execution_arn: str, operation_id: str) -> bool:
885 """Private method - called when a retry delay has elapsed and retry is ready."""
886 execution = self.get_execution(execution_arn)
888 if execution.is_complete:
889 logger.info(
890 "[%s] Execution already completed, ignoring retry", execution_arn
891 )
892 return False
894 try:
895 execution.complete_retry(operation_id=operation_id)
896 self._update_execution(execution)
897 logger.debug(
898 "[%s] Retry ready for operation %s", execution_arn, operation_id
899 )
900 return True
901 except Exception:
902 logger.exception("[%s] Error processing retry ready.", execution_arn)
903 return False
905 def _defer_resume_event_if_active(
906 self,
907 execution_arn: str,
908 operation_id: str,
909 deferred_events: set[str],
910 ) -> bool:
911 self._validate_current_execution(execution_arn)
912 if not self._active_invocation:
913 return False
915 deferred_events.add(operation_id)
916 return True
918 def _defer_callback_timeout_if_active(
919 self,
920 execution_arn: str,
921 callback_id: str,
922 timeout_type: CallbackTimeoutType,
923 ) -> bool:
924 self._validate_current_execution(execution_arn)
925 if not self._active_invocation:
926 return False
928 self._deferred_callback_timeouts.add((callback_id, timeout_type))
929 return True
931 def _apply_deferred_resume_events(self, execution_arn: str) -> bool:
932 self._validate_current_execution(execution_arn)
933 should_resume = False
935 for operation_id in list(self._deferred_wait_resumes):
936 self._deferred_wait_resumes.discard(operation_id)
937 should_resume = (
938 self._on_wait_succeeded(execution_arn, operation_id) or should_resume
939 )
941 for operation_id in list(self._deferred_retry_resumes):
942 self._deferred_retry_resumes.discard(operation_id)
943 should_resume = (
944 self._on_retry_ready(execution_arn, operation_id) or should_resume
945 )
947 for callback_event in list(self._deferred_callback_timeouts):
948 callback_id, timeout_type = callback_event
949 self._deferred_callback_timeouts.discard(callback_event)
950 should_resume = (
951 self._complete_callback_timeout(
952 execution_arn, callback_id, timeout_type
953 )
954 or should_resume
955 )
957 return should_resume
959 def timeout_execution(self, execution_arn: str, error: ErrorObject) -> None:
960 """Handle execution timeout."""
961 logger.exception("[%s] Execution timed out.", execution_arn)
962 execution: Execution = self.get_execution(execution_arn)
963 execution.complete_timeout(error=error) # Sets CloseStatus.TIMED_OUT
964 self._update_execution(execution)
965 self._complete_events(execution_arn=execution_arn)
967 def stop_execution(self, execution_arn: str, error: ErrorObject) -> None:
968 """Handle execution stop."""
969 self.fail_execution(execution_arn, error)
971 def schedule_wait_timer(
972 self, execution_arn: str, operation_id: str, delay: float
973 ) -> None:
974 """Schedule a wait operation."""
975 logger.debug("[%s] scheduling wait with delay: %d", execution_arn, delay)
977 async def wait_handler() -> None:
978 if self._defer_resume_event_if_active(
979 execution_arn, operation_id, self._deferred_wait_resumes
980 ):
981 return
982 if self._on_wait_succeeded(execution_arn, operation_id):
983 self._schedule_resume(execution_arn)
985 self._scheduler.call_later(
986 wait_handler, delay=delay, completion_event=self._completion_event
987 )
989 def schedule_step_retry(
990 self, execution_arn: str, operation_id: str, delay: float
991 ) -> None:
992 """Schedule a step retry."""
993 logger.debug(
994 "[%s] scheduling retry for %s with delay: %d",
995 execution_arn,
996 operation_id,
997 delay,
998 )
1000 async def retry_handler() -> None:
1001 if self._defer_resume_event_if_active(
1002 execution_arn, operation_id, self._deferred_retry_resumes
1003 ):
1004 return
1005 if self._on_retry_ready(execution_arn, operation_id):
1006 self._schedule_resume(execution_arn)
1008 self._scheduler.call_later(
1009 retry_handler, delay=delay, completion_event=self._completion_event
1010 )
1012 def schedule_callback_timeouts(
1013 self,
1014 execution_arn: str,
1015 callback_options: CallbackOptions | None,
1016 callback_id: str,
1017 ) -> None:
1018 """Schedule callback timeout and heartbeat timeout if configured."""
1019 self._schedule_callback_timeouts(execution_arn, callback_options, callback_id)
1021 def _schedule_callback_timeouts(
1022 self,
1023 execution_arn: str,
1024 callback_options: CallbackOptions | None,
1025 callback_id: str,
1026 ) -> None:
1027 """Schedule callback timeout and heartbeat timeout if configured."""
1028 try:
1029 if not callback_options:
1030 return
1032 # Schedule main timeout if configured
1033 if callback_options.timeout_seconds > 0:
1034 timeout_delay = scale_delay(
1035 callback_options.timeout_seconds,
1036 minimum=_CALLBACK_TIMEOUT_MINIMUM_DELAY_SECONDS,
1037 )
1039 async def timeout_handler() -> None:
1040 self._on_callback_timeout(execution_arn, callback_id)
1042 timeout_future = self._scheduler.call_later(
1043 timeout_handler,
1044 delay=timeout_delay,
1045 completion_event=self._completion_event,
1046 )
1047 self._callback_timeouts[callback_id] = timeout_future
1049 # Schedule heartbeat timeout if configured
1050 if callback_options.heartbeat_timeout_seconds > 0:
1051 heartbeat_delay = scale_delay(
1052 callback_options.heartbeat_timeout_seconds,
1053 minimum=_CALLBACK_TIMEOUT_MINIMUM_DELAY_SECONDS,
1054 )
1056 async def heartbeat_timeout_handler() -> None:
1057 self._on_callback_heartbeat_timeout(execution_arn, callback_id)
1059 heartbeat_future = self._scheduler.call_later(
1060 heartbeat_timeout_handler,
1061 delay=heartbeat_delay,
1062 completion_event=self._completion_event,
1063 )
1064 self._callback_heartbeats[callback_id] = heartbeat_future
1066 except Exception:
1067 logger.exception(
1068 "[%s] Error scheduling callback timeouts for %s",
1069 execution_arn,
1070 callback_id,
1071 )
1073 def _reset_callback_heartbeat_timeout(
1074 self, callback_id: str, execution_arn: str
1075 ) -> None:
1076 """Reset the heartbeat timeout for a callback."""
1077 # Cancel existing heartbeat timeout
1078 if heartbeat_future := self._callback_heartbeats.pop(callback_id, None):
1079 heartbeat_future.cancel()
1081 # Find callback options to reschedule heartbeat timeout
1082 try:
1083 callback_token = CallbackToken.from_str(callback_id)
1084 execution = self.get_execution(callback_token.execution_arn)
1086 callback_options = None
1087 for update in execution.updates:
1088 if (
1089 update.operation_id == callback_token.operation_id
1090 and update.callback_options
1091 and update.action.value == "START"
1092 ):
1093 callback_options = update.callback_options
1094 break
1096 if callback_options and callback_options.heartbeat_timeout_seconds > 0:
1097 heartbeat_delay = scale_delay(
1098 callback_options.heartbeat_timeout_seconds,
1099 minimum=_CALLBACK_TIMEOUT_MINIMUM_DELAY_SECONDS,
1100 )
1102 async def heartbeat_timeout_handler() -> None:
1103 self._on_callback_heartbeat_timeout(execution_arn, callback_id)
1105 heartbeat_future = self._scheduler.call_later(
1106 heartbeat_timeout_handler,
1107 delay=heartbeat_delay,
1108 completion_event=self._completion_event,
1109 )
1110 self._callback_heartbeats[callback_id] = heartbeat_future
1112 except Exception:
1113 logger.exception(
1114 "[%s] Error resetting callback heartbeat timeout for %s",
1115 execution_arn,
1116 callback_id,
1117 )
1119 def _cleanup_callback_timeouts(self, callback_id: str) -> None:
1120 """Clean up timeout events for a completed callback."""
1121 # Clean up main timeout
1122 if timeout_future := self._callback_timeouts.pop(callback_id, None):
1123 timeout_future.cancel()
1125 # Clean up heartbeat timeout
1126 if heartbeat_future := self._callback_heartbeats.pop(callback_id, None):
1127 heartbeat_future.cancel()
1129 def _complete_callback_timeout(
1130 self,
1131 execution_arn: str,
1132 callback_id: str,
1133 timeout_type: CallbackTimeoutType,
1134 ) -> bool:
1135 """Complete a callback with a timeout if the execution is still active."""
1136 try:
1137 callback_token = CallbackToken.from_str(callback_id)
1138 execution = self.get_execution(callback_token.execution_arn)
1140 if execution.is_complete:
1141 return False
1143 timeout_label = (
1144 "Callback heartbeat timed out"
1145 if timeout_type is CallbackTimeoutType.HEARTBEAT
1146 else "Callback timed out"
1147 )
1148 timeout_error = ErrorObject.from_message(
1149 f"{timeout_label}: {timeout_type.value}"
1150 )
1151 execution.complete_callback_timeout(callback_id, timeout_error)
1152 self._update_execution(execution)
1153 logger.warning("[%s] %s %s", execution_arn, timeout_label, callback_id)
1154 return True
1155 except Exception:
1156 logger.exception(
1157 "[%s] Error processing callback timeout for %s",
1158 execution_arn,
1159 callback_id,
1160 )
1161 return False
1163 def _on_callback_timeout(self, execution_arn: str, callback_id: str) -> None:
1164 """Handle callback timeout."""
1165 if self._defer_callback_timeout_if_active(
1166 execution_arn, callback_id, CallbackTimeoutType.TIMEOUT
1167 ):
1168 return
1169 if self._complete_callback_timeout(
1170 execution_arn, callback_id, CallbackTimeoutType.TIMEOUT
1171 ):
1172 self._schedule_resume(execution_arn)
1174 def _on_callback_heartbeat_timeout(
1175 self, execution_arn: str, callback_id: str
1176 ) -> None:
1177 """Handle callback heartbeat timeout."""
1178 if self._defer_callback_timeout_if_active(
1179 execution_arn, callback_id, CallbackTimeoutType.HEARTBEAT
1180 ):
1181 return
1182 if self._complete_callback_timeout(
1183 execution_arn, callback_id, CallbackTimeoutType.HEARTBEAT
1184 ):
1185 self._schedule_resume(execution_arn)