Coverage for async_durable_execution/_runner/local/executor.py: 96%
492 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +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 ..._core import (
11 CallbackOptions,
12 CallbackTimeoutType,
13 DurableExecutionInvocationInput,
14 DurableExecutionInvocationOutput,
15 ErrorObject,
16 InvocationStatus,
17 Operation,
18 OperationStatus,
19 OperationType,
20 OperationUpdate,
21)
22from ..exceptions import (
23 IllegalStateException,
24 InvalidParameterValueException,
25 ResourceNotFoundException,
26)
27from ..model import (
28 TERMINAL_STATUSES,
29 EventCreationContext,
30 GetDurableExecutionHistoryResponse,
31)
32from ..model import (
33 Event as HistoryEvent,
34)
35from .model import (
36 CallbackToken,
37 GetDurableExecutionStateResponse,
38 Invoker,
39 SendDurableExecutionCallbackFailureResponse,
40 SendDurableExecutionCallbackHeartbeatResponse,
41 SendDurableExecutionCallbackSuccessResponse,
42 StartDurableExecutionInput,
43 StartDurableExecutionOutput,
44)
45from .time_scale import scale_delay
46from .execution import Execution
48if TYPE_CHECKING:
49 from collections.abc import Awaitable, Callable
50 from asyncio import Future
52 from . import InMemoryServiceClient
53 from .scheduler import Event, Scheduler
55logger = logging.getLogger(__name__)
57_CALLBACK_TIMEOUT_MINIMUM_DELAY_SECONDS = 5.0
60class Executor:
61 MAX_CONSECUTIVE_FAILED_ATTEMPTS: int = 5
62 RETRY_BACKOFF_SECONDS: int = 5
64 def __init__(
65 self,
66 scheduler: Scheduler,
67 invoker: Invoker,
68 service_client: InMemoryServiceClient,
69 ) -> None:
70 self._scheduler = scheduler
71 self._invoker = invoker
72 self._service_client = service_client
73 self._execution: Execution | None = None
74 self._execution_arn: str | None = None
75 self._completion_event: Event | None = None
76 self._callback_timeouts: dict[str, Future] = {}
77 self._callback_heartbeats: dict[str, Future] = {}
78 self._execution_timeout: Future | None = None
79 self._active_invocation: bool = False
80 self._scheduled_resume: bool = False
81 self._pending_resume: bool = False
82 self._deferred_wait_resumes: set[str] = set()
83 self._deferred_retry_resumes: set[str] = set()
84 self._deferred_callback_timeouts: set[tuple[str, CallbackTimeoutType]] = set()
86 def start_execution(
87 self,
88 input: StartDurableExecutionInput, # noqa: A002
89 ) -> StartDurableExecutionOutput:
90 self._ensure_can_start_execution()
92 # Generate invocation_id if not provided
93 if input.invocation_id is None:
94 input = StartDurableExecutionInput(
95 account_id=input.account_id,
96 function_name=input.function_name,
97 function_qualifier=input.function_qualifier,
98 execution_name=input.execution_name,
99 execution_timeout_seconds=input.execution_timeout_seconds,
100 execution_retention_period_days=input.execution_retention_period_days,
101 invocation_id=str(uuid.uuid4()),
102 trace_fields=input.trace_fields,
103 tenant_id=input.tenant_id,
104 input=input.input,
105 lambda_endpoint=input.lambda_endpoint,
106 )
108 execution = Execution.new(input=input)
109 execution.start()
110 self._save_execution(execution)
111 logger.debug("Created execution with ARN: %s", execution.durable_execution_arn)
113 completion_event = self._scheduler.create_event()
114 self._execution_arn = execution.durable_execution_arn
115 self._completion_event = completion_event
116 self._execution_timeout = None
117 self._active_invocation = False
118 self._scheduled_resume = False
119 self._pending_resume = False
120 self._deferred_wait_resumes.clear()
121 self._deferred_retry_resumes.clear()
122 self._deferred_callback_timeouts.clear()
124 # Schedule execution timeout
125 if input.execution_timeout_seconds > 0:
127 async def timeout_handler() -> None:
128 error = ErrorObject.from_message(
129 f"Execution timed out after {input.execution_timeout_seconds} seconds."
130 )
131 self.timeout_execution(execution.durable_execution_arn, error)
133 self._execution_timeout = self._scheduler.call_later(
134 timeout_handler,
135 delay=input.execution_timeout_seconds,
136 completion_event=completion_event,
137 )
139 # Schedule initial invocation to run immediately
140 self._invoke_execution(execution.durable_execution_arn)
142 return StartDurableExecutionOutput(
143 execution_arn=execution.durable_execution_arn
144 )
146 def _ensure_can_start_execution(self) -> None:
147 """Reject overlapping executions in the local runner."""
148 if self._execution is None:
149 return
151 if not self._execution.is_complete: 151 ↛ exitline 151 didn't return from function '_ensure_can_start_execution' because the condition on line 151 was always true
152 msg = "Local runner supports only one execution at a time."
153 raise IllegalStateException(msg)
155 def _save_execution(self, execution: Execution) -> None:
156 self._execution = execution
157 self._execution_arn = execution.durable_execution_arn
159 def _update_execution(self, execution: Execution) -> None:
160 self._execution = execution
161 self._execution_arn = execution.durable_execution_arn
163 def set_execution(self, execution: Execution) -> None:
164 """Replace the current local execution object."""
165 self._update_execution(execution)
167 def _validate_current_execution(self, execution_arn: str) -> None:
168 if self._execution_arn is None:
169 return
171 if self._execution_arn != execution_arn: 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true
172 msg = f"Execution {execution_arn} is not the active local execution."
173 raise ResourceNotFoundException(msg)
175 def get_execution(self, execution_arn: str) -> Execution:
176 """Get execution by ARN.
178 Args:
179 execution_arn: The execution ARN to retrieve
181 Returns:
182 Execution: The execution object
184 Raises:
185 ResourceNotFoundException: If execution does not exist
186 """
187 if ( 187 ↛ 191line 187 didn't jump to line 191 because the condition on line 187 was never true
188 self._execution is None
189 or self._execution.durable_execution_arn != execution_arn
190 ):
191 msg: str = f"Execution {execution_arn} not found"
192 raise ResourceNotFoundException(msg)
193 return self._execution
195 def get_execution_state(
196 self,
197 execution_arn: str,
198 checkpoint_token: str | None = None,
199 marker: str | None = None,
200 max_items: int | None = None,
201 ) -> GetDurableExecutionStateResponse:
202 """Get execution state with operations.
204 Args:
205 execution_arn: The execution ARN
206 checkpoint_token: Checkpoint token for state consistency
207 marker: Pagination marker
208 max_items: Maximum items to return
210 Returns:
211 GetDurableExecutionStateResponse: Execution state with operations
213 Raises:
214 ResourceNotFoundException: If execution does not exist
215 InvalidParameterValueException: If checkpoint token is invalid
216 """
217 execution = self.get_execution(execution_arn)
219 # TODO: Validate checkpoint token if provided
220 if checkpoint_token and checkpoint_token not in execution.used_tokens:
221 msg: str = f"Invalid checkpoint token: {checkpoint_token}"
222 raise InvalidParameterValueException(msg)
224 # Get operations (excluding the initial EXECUTION operation for state)
225 operations = execution.get_assertable_operations()
227 # Apply pagination
228 if max_items is None:
229 max_items = 100
231 # Simple pagination - in real implementation would need proper marker handling
232 start_index = 0
233 if marker:
234 try:
235 start_index = int(marker)
236 except ValueError:
237 start_index = 0
239 end_index = start_index + max_items
240 paginated_operations = operations[start_index:end_index]
242 next_marker = None
243 if end_index < len(operations):
244 next_marker = str(end_index)
246 return GetDurableExecutionStateResponse(
247 operations=paginated_operations, next_marker=next_marker
248 )
250 def get_execution_history(
251 self,
252 execution_arn: str,
253 include_execution_data: bool = False, # noqa: FBT001, FBT002
254 reverse_order: bool = False, # noqa: FBT001, FBT002
255 marker: str | None = None,
256 max_items: int | None = None,
257 ) -> GetDurableExecutionHistoryResponse:
258 """Get execution history with events.
260 Args:
261 execution_arn: The execution ARN
262 include_execution_data: Whether to include execution data in events
263 reverse_order: Return events in reverse chronological order
264 marker: Pagination marker (event_id)
265 max_items: Maximum items to return
267 Returns:
268 GetDurableExecutionHistoryResponse: Execution history with events
270 Raises:
271 ResourceNotFoundException: If execution does not exist
272 """
273 execution: Execution = self.get_execution(execution_arn)
275 # Generate events
276 all_events: list[HistoryEvent] = []
277 ops: list[Operation] = execution.operations
278 updates: list[OperationUpdate] = execution.updates
279 updates_dict: dict[str, OperationUpdate] = {u.operation_id: u for u in updates}
280 durable_execution_arn: str = execution.durable_execution_arn
282 # Add InvocationCompleted events
283 for completion in execution.invocation_completions:
284 invocation_event = HistoryEvent.create_invocation_completed(
285 event_id=0, # Temporary, will be reassigned
286 event_timestamp=completion.end_timestamp,
287 start_timestamp=completion.start_timestamp,
288 end_timestamp=completion.end_timestamp,
289 request_id=completion.request_id,
290 )
291 all_events.append(invocation_event)
293 # Generate all events first (without final event IDs)
294 for op in ops:
295 operation_update: OperationUpdate | None = updates_dict.get(op.operation_id)
297 if op.status is OperationStatus.PENDING:
298 if ( 298 ↛ 302line 298 didn't jump to line 302 because the condition on line 298 was never true
299 op.operation_type is not OperationType.CHAINED_INVOKE
300 or op.start_timestamp is None
301 ):
302 continue
303 context: EventCreationContext = EventCreationContext(
304 op,
305 0, # Temporary event_id, will be reassigned after sorting
306 durable_execution_arn,
307 execution.start_input,
308 execution.result,
309 operation_update,
310 include_execution_data,
311 )
312 pending = HistoryEvent.create_chained_invoke_event_pending(context)
313 all_events.append(pending)
314 if op.start_timestamp is not None: 314 ↛ 326line 314 didn't jump to line 326 because the condition on line 314 was always true
315 context = EventCreationContext(
316 op,
317 0, # Temporary event_id, will be reassigned after sorting
318 durable_execution_arn,
319 execution.start_input,
320 execution.result,
321 operation_update,
322 include_execution_data,
323 )
324 started = HistoryEvent.create_event_started(context)
325 all_events.append(started)
326 if op.end_timestamp is not None and op.status in TERMINAL_STATUSES:
327 context = EventCreationContext(
328 op,
329 0, # Temporary event_id, will be reassigned after sorting
330 durable_execution_arn,
331 execution.start_input,
332 execution.result,
333 operation_update,
334 include_execution_data,
335 )
336 finished = HistoryEvent.create_event_terminated(context)
337 all_events.append(finished)
339 # Sort events by timestamp to get correct chronological order
340 all_events.sort(key=lambda event: event.event_timestamp)
342 # Reassign event IDs based on chronological order
343 all_events = [
344 HistoryEvent.from_event_with_id(event, i)
345 for i, event in enumerate(all_events, 1)
346 ]
348 # Apply cursor-based pagination
349 if max_items is None:
350 max_items = 100
352 # Handle pagination marker
353 if reverse_order:
354 all_events.reverse()
355 start_index: int = 0
356 if marker:
357 try:
358 marker_event_id: int = int(marker)
359 # Find the index of the first event with event_id >= marker
360 start_index = len(all_events)
361 for i, e in enumerate(all_events): 361 ↛ 374line 361 didn't jump to line 374 because the loop on line 361 didn't complete
362 is_valid_page_start: bool = (
363 e.event_id < marker_event_id
364 if reverse_order
365 else e.event_id >= marker_event_id
366 )
367 if is_valid_page_start:
368 start_index = i
369 break
370 except ValueError:
371 start_index = 0
373 # Get paginated events
374 end_index: int = start_index + max_items
375 paginated_events: list[HistoryEvent] = all_events[start_index:end_index]
377 # Generate next marker
378 next_marker: str | None = None
379 if end_index < len(all_events):
380 if reverse_order:
381 # Next marker is the event_id of the last returned event
382 next_marker = (
383 str(paginated_events[-1].event_id) if paginated_events else None
384 )
385 else:
386 # Next marker is the event_id of the next event after the last returned
387 next_marker = (
388 str(all_events[end_index].event_id)
389 if end_index < len(all_events)
390 else None
391 )
393 return GetDurableExecutionHistoryResponse(
394 events=paginated_events, next_marker=next_marker
395 )
397 def send_callback_success(
398 self,
399 callback_id: str,
400 result: bytes | None = None,
401 ) -> SendDurableExecutionCallbackSuccessResponse:
402 """Send callback success response.
404 Args:
405 callback_id: The callback ID to respond to
406 result: Optional result data for the callback
408 Returns:
409 SendDurableExecutionCallbackSuccessResponse: Empty response
411 Raises:
412 InvalidParameterValueException: If callback_id is invalid
413 ResourceNotFoundException: If callback does not exist
414 """
415 if not callback_id:
416 msg: str = "callback_id is required"
417 raise InvalidParameterValueException(msg)
419 try:
420 callback_token = CallbackToken.from_str(callback_id)
421 execution = self.get_execution(callback_token.execution_arn)
422 execution.complete_callback_success(callback_id, result)
423 self._update_execution(execution)
424 self._cleanup_callback_timeouts(callback_id)
425 self._schedule_resume(callback_token.execution_arn)
426 logger.info("Callback success completed for callback_id: %s", callback_id)
427 except Exception as e:
428 msg = f"Failed to process callback success: {e}"
429 raise ResourceNotFoundException(msg) from e
431 return SendDurableExecutionCallbackSuccessResponse()
433 def send_callback_failure(
434 self,
435 callback_id: str,
436 error: ErrorObject | None = None,
437 ) -> SendDurableExecutionCallbackFailureResponse:
438 """Send callback failure response.
440 Args:
441 callback_id: The callback ID to respond to
442 error: Optional error object for the callback failure
444 Returns:
445 SendDurableExecutionCallbackFailureResponse: Empty response
447 Raises:
448 InvalidParameterValueException: If callback_id is invalid
449 ResourceNotFoundException: If callback does not exist
450 """
451 if not callback_id:
452 msg: str = "callback_id is required"
453 raise InvalidParameterValueException(msg)
455 callback_error: ErrorObject = error or ErrorObject.from_message("")
457 try:
458 callback_token: CallbackToken = CallbackToken.from_str(callback_id)
459 execution: Execution = self.get_execution(callback_token.execution_arn)
460 execution.complete_callback_failure(callback_id, callback_error)
461 self._update_execution(execution)
462 self._cleanup_callback_timeouts(callback_id)
463 self._schedule_resume(callback_token.execution_arn)
464 logger.info("Callback failure completed for callback_id: %s", callback_id)
465 except Exception as e:
466 msg = f"Failed to process callback failure: {e}"
467 raise ResourceNotFoundException(msg) from e
469 return SendDurableExecutionCallbackFailureResponse()
471 def send_callback_heartbeat(
472 self, callback_id: str
473 ) -> SendDurableExecutionCallbackHeartbeatResponse:
474 """Send callback heartbeat to keep callback alive.
476 Args:
477 callback_id: The callback ID to send heartbeat for
479 Returns:
480 SendDurableExecutionCallbackHeartbeatResponse: Empty response
482 Raises:
483 InvalidParameterValueException: If callback_id is invalid
484 ResourceNotFoundException: If callback does not exist
485 """
486 if not callback_id:
487 msg: str = "callback_id is required"
488 raise InvalidParameterValueException(msg)
490 try:
491 callback_token: CallbackToken = CallbackToken.from_str(callback_id)
492 execution: Execution = self.get_execution(callback_token.execution_arn)
494 # Find callback operation to verify it exists and is active
495 _, operation = execution.find_callback_operation(callback_id)
496 if operation.status != OperationStatus.STARTED:
497 msg = f"Callback {callback_id} is not active"
498 raise ResourceNotFoundException(msg)
500 # Reset heartbeat timeout if configured
501 self._reset_callback_heartbeat_timeout(
502 callback_id, execution.durable_execution_arn
503 )
504 logger.info("Callback heartbeat processed for callback_id: %s", callback_id)
505 except Exception as e:
506 msg = f"Failed to process callback heartbeat: {e}"
507 raise ResourceNotFoundException(msg) from e
509 return SendDurableExecutionCallbackHeartbeatResponse()
511 def _validate_invocation_response(
512 self,
513 execution_arn: str,
514 response: DurableExecutionInvocationOutput,
515 execution: Execution,
516 ) -> None:
517 """Validate response status and apply the resulting execution changes.
519 Raises:
520 InvalidParameterValueException: If the response status is invalid.
521 IllegalStateException: If the response status is valid but the execution is already completed.
522 """
523 if execution.is_complete:
524 msg_already_complete: str = "Execution already completed, ignoring result"
526 raise IllegalStateException(msg_already_complete)
528 if response.status is None:
529 msg_status_required: str = "Response status is required"
531 raise InvalidParameterValueException(msg_status_required)
533 match response.status:
534 case InvocationStatus.FAILED:
535 if response.result is not None:
536 msg_failed_result: str = (
537 "Cannot provide a Result for FAILED status."
538 )
539 raise InvalidParameterValueException(msg_failed_result)
540 logger.info("[%s] Execution failed", execution_arn)
541 self._complete_workflow(
542 execution_arn, result=None, error=response.error
543 )
545 case InvocationStatus.SUCCEEDED:
546 if response.error is not None:
547 msg_success_error: str = (
548 "Cannot provide an Error for SUCCEEDED status."
549 )
550 raise InvalidParameterValueException(msg_success_error)
551 logger.info("[%s] Execution succeeded", execution_arn)
552 self._complete_workflow(
553 execution_arn, result=response.result, error=None
554 )
556 case InvocationStatus.PENDING:
557 if not execution.has_pending_operations():
558 msg_pending_ops: str = (
559 "Cannot return PENDING status with no pending operations."
560 )
561 raise InvalidParameterValueException(msg_pending_ops)
562 logger.info("[%s] Execution pending async work", execution_arn)
564 case _:
565 msg_unexpected_status: str = (
566 f"Unexpected invocation status: {response.status}"
567 )
568 raise IllegalStateException(msg_unexpected_status)
570 def _invoke_handler(self, execution_arn: str) -> Callable[[], Awaitable[None]]:
571 """Create a parameterless callable that captures execution arn for the scheduler."""
573 async def invoke() -> None:
574 self._mark_invocation_started(execution_arn)
575 execution: Execution = self.get_execution(execution_arn)
577 # Early exit if execution is already completed - like Java's COMPLETED check
578 if execution.is_complete:
579 logger.info(
580 "[%s] Execution already completed, ignoring result", execution_arn
581 )
582 return
584 try:
585 checkpoint_token = execution.get_new_checkpoint_token()
586 invocation_input: DurableExecutionInvocationInput = (
587 self._invoker.create_invocation_input(
588 start_input=execution.start_input,
589 durable_execution_arn=execution.durable_execution_arn,
590 checkpoint_token=checkpoint_token,
591 operations=execution.operations,
592 )
593 )
595 self._save_execution(execution)
597 invocation_start = datetime.now(timezone.utc)
598 invoke_response = await self._invoker.invoke(
599 execution.start_input.function_name,
600 invocation_input,
601 execution.start_input.lambda_endpoint,
602 )
603 invocation_end = datetime.now(timezone.utc)
605 # Reload execution after invocation in case it was completed via checkpoint
606 execution = self.get_execution(execution_arn)
608 # Record invocation completion and save immediately
609 execution.record_invocation_completion(
610 invocation_start, invocation_end, invoke_response.request_id
611 )
612 self._save_execution(execution)
614 if execution.is_complete:
615 logger.info(
616 "[%s] Execution completed during invocation, ignoring result",
617 execution_arn,
618 )
619 return
621 # Process successful received response - validate status and handle accordingly
622 response = invoke_response.invocation_output
623 try:
624 self._validate_invocation_response(
625 execution_arn, response, execution
626 )
627 except (InvalidParameterValueException, IllegalStateException) as e:
628 logger.warning(
629 "[%s] Lambda output validation failure: %s", execution_arn, e
630 )
631 error_obj = ErrorObject.from_exception(e)
632 self._retry_invocation(execution, error_obj)
634 except ResourceNotFoundException:
635 logger.warning(
636 "[%s] Function No longer exists: %s",
637 execution_arn,
638 execution.start_input.function_name,
639 )
640 error_obj = ErrorObject.from_message(
641 message=f"Function not found: {execution.start_input.function_name}"
642 )
643 self._fail_workflow(execution_arn, error_obj)
645 except Exception as e: # noqa: BLE001
646 # Handle invocation errors (network, function not found, etc.)
647 logger.warning("[%s] Invocation failed: %s", execution_arn, e)
648 error_obj = ErrorObject.from_exception(e)
649 self._retry_invocation(execution, error_obj)
650 finally:
651 self._mark_invocation_finished(execution_arn)
653 return invoke
655 def _schedule_resume(self, execution_arn: str) -> None:
656 """Coalesce external resumes to avoid overlapping replays."""
657 self._validate_current_execution(execution_arn)
658 if self._active_invocation:
659 self._pending_resume = True
660 return
662 if self._scheduled_resume:
663 return
665 self._scheduled_resume = True
666 self._invoke_execution(execution_arn)
668 def _mark_invocation_started(self, execution_arn: str) -> None:
669 self._validate_current_execution(execution_arn)
670 self._scheduled_resume = False
671 self._active_invocation = True
673 def _mark_invocation_finished(self, execution_arn: str) -> None:
674 self._validate_current_execution(execution_arn)
675 self._active_invocation = False
676 should_resume = self._apply_deferred_resume_events(execution_arn)
677 if self._pending_resume:
678 self._pending_resume = False
679 should_resume = True
681 if should_resume:
682 execution = self.get_execution(execution_arn)
683 if not execution.is_complete:
684 self._scheduled_resume = True
685 self._invoke_execution(execution_arn)
687 def _invoke_execution(self, execution_arn: str, delay: float = 0) -> None:
688 """Invoke execution after delay in seconds."""
689 self._validate_current_execution(execution_arn)
690 self._scheduler.call_later(
691 self._invoke_handler(execution_arn),
692 delay=delay,
693 completion_event=self._completion_event,
694 )
696 def _complete_workflow(
697 self, execution_arn: str, result: str | None, error: ErrorObject | None
698 ) -> None:
699 """Complete workflow - handles both success and failure with terminal state validation."""
700 execution = self.get_execution(execution_arn)
702 if execution.is_complete:
703 msg: str = "Cannot make multiple close workflow decisions."
705 raise IllegalStateException(msg)
707 if error is not None:
708 self.fail_execution(execution_arn, error)
709 else:
710 self.complete_execution(execution_arn, result)
712 def _fail_workflow(self, execution_arn: str, error: ErrorObject) -> None:
713 """Fail workflow with terminal state validation."""
714 execution = self.get_execution(execution_arn)
716 if execution.is_complete:
717 msg: str = "Cannot make multiple close workflow decisions."
719 raise IllegalStateException(msg)
721 self.fail_execution(execution_arn, error)
723 def _retry_invocation(self, execution: Execution, error: ErrorObject) -> None:
724 """Handle retry logic or fail execution if retries exhausted."""
725 if (
726 execution.consecutive_failed_invocation_attempts
727 > self.MAX_CONSECUTIVE_FAILED_ATTEMPTS
728 ):
729 # Exhausted retries - fail the execution
730 self._fail_workflow(
731 execution_arn=execution.durable_execution_arn, error=error
732 )
733 else:
734 # Schedule retry with backoff
735 execution.consecutive_failed_invocation_attempts += 1
736 self._save_execution(execution)
737 self._invoke_execution(
738 execution_arn=execution.durable_execution_arn,
739 delay=self.RETRY_BACKOFF_SECONDS,
740 )
742 def _complete_events(self, execution_arn: str) -> None:
743 self._validate_current_execution(execution_arn)
744 # complete doesn't actually checkpoint explicitly
745 if self._completion_event:
746 self._completion_event.set()
747 if execution_timeout := self._execution_timeout:
748 execution_timeout.cancel()
749 self._execution_timeout = None
751 async def wait_until_complete(
752 self, execution_arn: str, timeout: float | None = None
753 ) -> bool:
754 """Wait until execution completion.
756 Args
757 timeout (int|float|None): Wait for event to set until this timeout.
759 Returns:
760 True when set. False if the event timed out without being set.
761 """
762 self._validate_current_execution(execution_arn)
763 if event := self._completion_event:
764 return await event.wait_async(timeout)
766 # this really shouldn't happen - implies execution timed out?
767 msg: str = "execution does not exist."
769 raise ResourceNotFoundException(msg)
771 def complete_execution(self, execution_arn: str, result: str | None = None) -> None:
772 """Complete execution successfully (COMPLETE_WORKFLOW_EXECUTION decision)."""
773 logger.debug("[%s] Completing execution with result: %s", execution_arn, result)
774 execution: Execution = self.get_execution(execution_arn)
775 execution.complete_success(result=result) # Sets CloseStatus.COMPLETED
776 self._update_execution(execution)
777 if execution.result is None:
778 msg: str = "Execution result is required"
779 raise IllegalStateException(msg)
780 self._complete_events(execution_arn=execution_arn)
782 def fail_execution(self, execution_arn: str, error: ErrorObject) -> None:
783 """Fail execution with error (FAIL_WORKFLOW_EXECUTION decision)."""
784 logger.error("[%s] Completing execution with error: %s", execution_arn, error)
785 execution: Execution = self.get_execution(execution_arn)
786 execution.complete_fail(error=error) # Sets CloseStatus.FAILED
787 self._update_execution(execution)
788 # set by complete_fail
789 if execution.result is None:
790 msg: str = "Execution result is required"
791 raise IllegalStateException(msg)
792 self._complete_events(execution_arn=execution_arn)
794 def _on_wait_succeeded(self, execution_arn: str, operation_id: str) -> bool:
795 """Private method - called when a wait operation completes successfully."""
796 execution = self.get_execution(execution_arn)
798 if execution.is_complete:
799 logger.info(
800 "[%s] Execution already completed, ignoring wait succeeded event",
801 execution_arn,
802 )
803 return False
805 try:
806 execution.complete_wait(operation_id=operation_id)
807 self._update_execution(execution)
808 logger.debug(
809 "[%s] Wait succeeded for operation %s", execution_arn, operation_id
810 )
811 return True
812 except Exception:
813 logger.exception("[%s] Error processing wait succeeded.", execution_arn)
814 return False
816 def _on_retry_ready(self, execution_arn: str, operation_id: str) -> bool:
817 """Private method - called when a retry delay has elapsed and retry is ready."""
818 execution = self.get_execution(execution_arn)
820 if execution.is_complete:
821 logger.info(
822 "[%s] Execution already completed, ignoring retry", execution_arn
823 )
824 return False
826 try:
827 execution.complete_retry(operation_id=operation_id)
828 self._update_execution(execution)
829 logger.debug(
830 "[%s] Retry ready for operation %s", execution_arn, operation_id
831 )
832 return True
833 except Exception:
834 logger.exception("[%s] Error processing retry ready.", execution_arn)
835 return False
837 def _defer_resume_event_if_active(
838 self,
839 execution_arn: str,
840 operation_id: str,
841 deferred_events: set[str],
842 ) -> bool:
843 self._validate_current_execution(execution_arn)
844 if not self._active_invocation:
845 return False
847 deferred_events.add(operation_id)
848 return True
850 def _defer_callback_timeout_if_active(
851 self,
852 execution_arn: str,
853 callback_id: str,
854 timeout_type: CallbackTimeoutType,
855 ) -> bool:
856 self._validate_current_execution(execution_arn)
857 if not self._active_invocation: 857 ↛ 860line 857 didn't jump to line 860 because the condition on line 857 was always true
858 return False
860 self._deferred_callback_timeouts.add((callback_id, timeout_type))
861 return True
863 def _apply_deferred_resume_events(self, execution_arn: str) -> bool:
864 self._validate_current_execution(execution_arn)
865 should_resume = False
867 for operation_id in list(self._deferred_wait_resumes):
868 self._deferred_wait_resumes.discard(operation_id)
869 should_resume = (
870 self._on_wait_succeeded(execution_arn, operation_id) or should_resume
871 )
873 for operation_id in list(self._deferred_retry_resumes):
874 self._deferred_retry_resumes.discard(operation_id)
875 should_resume = (
876 self._on_retry_ready(execution_arn, operation_id) or should_resume
877 )
879 for callback_event in list(self._deferred_callback_timeouts): 879 ↛ 880line 879 didn't jump to line 880 because the loop on line 879 never started
880 callback_id, timeout_type = callback_event
881 self._deferred_callback_timeouts.discard(callback_event)
882 should_resume = (
883 self._complete_callback_timeout(
884 execution_arn, callback_id, timeout_type
885 )
886 or should_resume
887 )
889 return should_resume
891 def timeout_execution(self, execution_arn: str, error: ErrorObject) -> None:
892 """Handle execution timeout."""
893 logger.exception("[%s] Execution timed out.", execution_arn)
894 execution: Execution = self.get_execution(execution_arn)
895 execution.complete_timeout(error=error) # Sets CloseStatus.TIMED_OUT
896 self._update_execution(execution)
897 self._complete_events(execution_arn=execution_arn)
899 def schedule_wait_timer(
900 self, execution_arn: str, operation_id: str, delay: float
901 ) -> None:
902 """Schedule a wait operation."""
903 logger.debug("[%s] scheduling wait with delay: %d", execution_arn, delay)
905 async def wait_handler() -> None:
906 if self._defer_resume_event_if_active(
907 execution_arn, operation_id, self._deferred_wait_resumes
908 ):
909 return
910 if self._on_wait_succeeded(execution_arn, operation_id): 910 ↛ exitline 910 didn't return from function 'wait_handler' because the condition on line 910 was always true
911 self._schedule_resume(execution_arn)
913 self._scheduler.call_later(
914 wait_handler, delay=delay, completion_event=self._completion_event
915 )
917 def schedule_step_retry(
918 self, execution_arn: str, operation_id: str, delay: float
919 ) -> None:
920 """Schedule a step retry."""
921 logger.debug(
922 "[%s] scheduling retry for %s with delay: %d",
923 execution_arn,
924 operation_id,
925 delay,
926 )
928 async def retry_handler() -> None:
929 if self._defer_resume_event_if_active(
930 execution_arn, operation_id, self._deferred_retry_resumes
931 ):
932 return
933 if self._on_retry_ready(execution_arn, operation_id):
934 self._schedule_resume(execution_arn)
936 self._scheduler.call_later(
937 retry_handler, delay=delay, completion_event=self._completion_event
938 )
940 def schedule_callback_timeouts(
941 self,
942 execution_arn: str,
943 callback_options: CallbackOptions | None,
944 callback_id: str,
945 ) -> None:
946 """Schedule callback timeout and heartbeat timeout if configured."""
947 self._schedule_callback_timeouts(execution_arn, callback_options, callback_id)
949 def _schedule_callback_timeouts(
950 self,
951 execution_arn: str,
952 callback_options: CallbackOptions | None,
953 callback_id: str,
954 ) -> None:
955 """Schedule callback timeout and heartbeat timeout if configured."""
956 try:
957 if not callback_options:
958 return
960 # Schedule main timeout if configured
961 if callback_options.timeout_seconds > 0:
962 timeout_delay = scale_delay(
963 callback_options.timeout_seconds,
964 minimum=_CALLBACK_TIMEOUT_MINIMUM_DELAY_SECONDS,
965 )
967 async def timeout_handler() -> None:
968 self._on_callback_timeout(execution_arn, callback_id)
970 timeout_future = self._scheduler.call_later(
971 timeout_handler,
972 delay=timeout_delay,
973 completion_event=self._completion_event,
974 )
975 self._callback_timeouts[callback_id] = timeout_future
977 # Schedule heartbeat timeout if configured
978 if callback_options.heartbeat_timeout_seconds > 0:
979 heartbeat_delay = scale_delay(
980 callback_options.heartbeat_timeout_seconds,
981 minimum=_CALLBACK_TIMEOUT_MINIMUM_DELAY_SECONDS,
982 )
984 async def heartbeat_timeout_handler() -> None:
985 self._on_callback_heartbeat_timeout(execution_arn, callback_id)
987 heartbeat_future = self._scheduler.call_later(
988 heartbeat_timeout_handler,
989 delay=heartbeat_delay,
990 completion_event=self._completion_event,
991 )
992 self._callback_heartbeats[callback_id] = heartbeat_future
994 except Exception:
995 logger.exception(
996 "[%s] Error scheduling callback timeouts for %s",
997 execution_arn,
998 callback_id,
999 )
1001 def _reset_callback_heartbeat_timeout(
1002 self, callback_id: str, execution_arn: str
1003 ) -> None:
1004 """Reset the heartbeat timeout for a callback."""
1005 # Cancel existing heartbeat timeout
1006 if heartbeat_future := self._callback_heartbeats.pop(callback_id, None):
1007 heartbeat_future.cancel()
1009 # Find callback options to reschedule heartbeat timeout
1010 try:
1011 callback_token = CallbackToken.from_str(callback_id)
1012 execution = self.get_execution(callback_token.execution_arn)
1014 callback_options = None
1015 for update in execution.updates:
1016 if (
1017 update.operation_id == callback_token.operation_id
1018 and update.callback_options
1019 and update.action.value == "START"
1020 ):
1021 callback_options = update.callback_options
1022 break
1024 if callback_options and callback_options.heartbeat_timeout_seconds > 0:
1025 heartbeat_delay = scale_delay(
1026 callback_options.heartbeat_timeout_seconds,
1027 minimum=_CALLBACK_TIMEOUT_MINIMUM_DELAY_SECONDS,
1028 )
1030 async def heartbeat_timeout_handler() -> None:
1031 self._on_callback_heartbeat_timeout(execution_arn, callback_id)
1033 heartbeat_future = self._scheduler.call_later(
1034 heartbeat_timeout_handler,
1035 delay=heartbeat_delay,
1036 completion_event=self._completion_event,
1037 )
1038 self._callback_heartbeats[callback_id] = heartbeat_future
1040 except Exception:
1041 logger.exception(
1042 "[%s] Error resetting callback heartbeat timeout for %s",
1043 execution_arn,
1044 callback_id,
1045 )
1047 def _cleanup_callback_timeouts(self, callback_id: str) -> None:
1048 """Clean up timeout events for a completed callback."""
1049 # Clean up main timeout
1050 if timeout_future := self._callback_timeouts.pop(callback_id, None):
1051 timeout_future.cancel()
1053 # Clean up heartbeat timeout
1054 if heartbeat_future := self._callback_heartbeats.pop(callback_id, None):
1055 heartbeat_future.cancel()
1057 def _complete_callback_timeout(
1058 self,
1059 execution_arn: str,
1060 callback_id: str,
1061 timeout_type: CallbackTimeoutType,
1062 ) -> bool:
1063 """Complete a callback with a timeout if the execution is still active."""
1064 try:
1065 callback_token = CallbackToken.from_str(callback_id)
1066 execution = self.get_execution(callback_token.execution_arn)
1068 if execution.is_complete:
1069 return False
1071 timeout_label = (
1072 "Callback heartbeat timed out"
1073 if timeout_type is CallbackTimeoutType.HEARTBEAT
1074 else "Callback timed out"
1075 )
1076 timeout_error = ErrorObject.from_message(
1077 f"{timeout_label}: {timeout_type.value}"
1078 )
1079 execution.complete_callback_timeout(callback_id, timeout_error)
1080 self._update_execution(execution)
1081 logger.warning("[%s] %s %s", execution_arn, timeout_label, callback_id)
1082 return True
1083 except Exception:
1084 logger.exception(
1085 "[%s] Error processing callback timeout for %s",
1086 execution_arn,
1087 callback_id,
1088 )
1089 return False
1091 def _on_callback_timeout(self, execution_arn: str, callback_id: str) -> None:
1092 """Handle callback timeout."""
1093 if self._defer_callback_timeout_if_active( 1093 ↛ 1096line 1093 didn't jump to line 1096 because the condition on line 1093 was never true
1094 execution_arn, callback_id, CallbackTimeoutType.TIMEOUT
1095 ):
1096 return
1097 if self._complete_callback_timeout(
1098 execution_arn, callback_id, CallbackTimeoutType.TIMEOUT
1099 ):
1100 self._schedule_resume(execution_arn)
1102 def _on_callback_heartbeat_timeout(
1103 self, execution_arn: str, callback_id: str
1104 ) -> None:
1105 """Handle callback heartbeat timeout."""
1106 if self._defer_callback_timeout_if_active( 1106 ↛ 1109line 1106 didn't jump to line 1109 because the condition on line 1106 was never true
1107 execution_arn, callback_id, CallbackTimeoutType.HEARTBEAT
1108 ):
1109 return
1110 if self._complete_callback_timeout(
1111 execution_arn, callback_id, CallbackTimeoutType.HEARTBEAT
1112 ):
1113 self._schedule_resume(execution_arn)