Coverage for async_durable_execution/_core/execution.py: 95%
174 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +0000
1from __future__ import annotations
3import asyncio
4import functools
5import json
6import logging
7from collections.abc import Awaitable, Callable
8from dataclasses import dataclass, field
9from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast
11from .context import DurableContext, bind_current_context
12from .exceptions import (
13 CheckpointError,
14 ExecutionError,
15 InvocationError,
16 SuspendExecution,
17 _sdk_error_type_name,
18)
19from .models import (
20 DurableExecutionInvocationOutput,
21 ErrorObject,
22 InvocationStatus,
23 Operation,
24 OperationUpdate,
25 OperationIdentifier,
26 JsonSerializableModel,
27)
28from .client import (
29 AsyncLambdaClient,
30 ThreadedSyncLambdaClient,
31 create_default_client,
32 lambda_api_client_is_async,
33)
34from .logger import configure_durable_logger
35from .state import ExecutionState
38if TYPE_CHECKING:
39 from collections.abc import MutableMapping
41 from .client import AsyncLambdaApiClient, DurableServiceClient, LambdaApiClient
42 from .models import LambdaContext
44configure_durable_logger(logging.getLogger())
45logger = logging.getLogger(__name__)
47# 6MB in bytes, minus 50 bytes for envelope
48LAMBDA_RESPONSE_SIZE_LIMIT = 6 * 1024 * 1024 - 50
50T = TypeVar("T")
51Params = ParamSpec("Params")
54def durable_callable(
55 func: Callable[Params, Awaitable[T]],
56) -> Callable[Params, Callable[[], Awaitable[T]]]:
57 """Wrap an async function so calling it returns a zero-argument durable callable.
59 The returned callable can be passed to durable operations such as `step()`
60 and `run_in_child_context()`, keeping durable operation creation explicit
61 while avoiding manual `functools.partial(...)` wrapping at the callsite.
63 Class and static methods are supported with either decorator order:
64 `@classmethod`/`@staticmethod` may appear above or below `@durable_callable`.
65 """
66 if isinstance(func, classmethod):
67 return classmethod(durable_callable(func.__func__))
68 if isinstance(func, staticmethod):
69 return staticmethod(durable_callable(func.__func__))
71 @functools.wraps(func)
72 def wrapper(
73 *args: Params.args, **kwargs: Params.kwargs
74 ) -> Callable[[], Awaitable[T]]:
75 bound = functools.partial(func, *args, **kwargs)
76 setattr(bound, "__name__", func.__name__)
77 return bound
79 return wrapper
82@dataclass(frozen=True)
83class InitialExecutionState(JsonSerializableModel):
84 """Initial page of operation history included with an invocation event."""
86 operations: list[Operation] = field(
87 default_factory=list,
88 metadata={"alias": "Operations"},
89 )
90 next_marker: str = field(default="", metadata={"alias": "NextMarker"})
93@dataclass(frozen=True)
94class DurableExecutionInvocationInput(JsonSerializableModel):
95 """Event payload delivered to a durable Lambda invocation."""
97 durable_execution_arn: str = field(metadata={"alias": "DurableExecutionArn"})
98 checkpoint_token: str = field(metadata={"alias": "CheckpointToken"})
99 initial_execution_state: InitialExecutionState = field(
100 default_factory=InitialExecutionState,
101 metadata={"alias": "InitialExecutionState"},
102 )
105def _bind_service_client_to_handler(
106 handler: Callable[[Any, LambdaContext], Any],
107 service_client: DurableServiceClient,
108) -> Callable[[Any, LambdaContext], Any]:
109 """Recreate a durable handler with a specific service client bound."""
111 handler_attrs = getattr(handler, "__dict__", {})
112 original_handler = handler_attrs.get("_durable_execution_original")
113 if original_handler is None:
114 return handler
116 boto3_client = handler_attrs.get("_durable_execution_boto3_client")
117 return durable_execution(
118 original_handler,
119 boto3_client=boto3_client,
120 service_client=service_client,
121 )
124@dataclass(frozen=True)
125class DurableConfig:
126 boto3_client: LambdaApiClient | AsyncLambdaApiClient | None = None
127 service_client: DurableServiceClient | None = None
130def durable_execution(
131 func: Callable[..., Awaitable[Any]] | None = None,
132 /,
133 *,
134 boto3_client: LambdaApiClient | AsyncLambdaApiClient | None = None,
135 service_client: DurableServiceClient | None = None,
136) -> Callable[[Any, LambdaContext], Any]:
137 """
138 Decorator to create a durable execution handler.
140 Args:
141 func: The user function to decorate
142 boto3_client: Optional sync or async Lambda API client to use
143 service_client: Optional durable service client to use. Intended for
144 testing and local execution tooling.
145 """
146 # Decorator called with parameters
147 if func is None:
148 logger.debug("Decorator called with parameters")
149 return functools.partial(
150 durable_execution,
151 boto3_client=boto3_client,
152 service_client=service_client,
153 )
154 config = DurableConfig(
155 boto3_client=boto3_client,
156 service_client=service_client,
157 )
158 logger.debug("Starting durable execution handler...")
160 # Use the explicitly provided durable client when present. Otherwise, delay
161 # Lambda API client construction until invocation so importing decorated handlers
162 # does not require AWS environment configuration.
163 active_service_client = config.service_client
164 active_service_client_loop: asyncio.AbstractEventLoop | None = None
165 handler_loop: asyncio.AbstractEventLoop | None = None
167 def get_or_create_handler_loop() -> asyncio.AbstractEventLoop:
168 nonlocal handler_loop
169 if handler_loop is None or handler_loop.is_closed():
170 handler_loop = asyncio.new_event_loop()
171 return handler_loop
173 async def get_active_service_client() -> DurableServiceClient:
174 nonlocal active_service_client, active_service_client_loop
175 current_loop = asyncio.get_running_loop()
176 if active_service_client is not None:
177 if ( 177 ↛ 185line 177 didn't jump to line 185 because the condition on line 177 was always true
178 active_service_client_loop is None
179 or active_service_client_loop is current_loop
180 ):
181 return active_service_client
183 # The SDK-owned async client is bound to the loop that first used it.
184 # If test code calls the async handler on another loop, rebuild it there.
185 active_service_client = None
186 active_service_client_loop = None
188 if config.boto3_client is not None:
189 if lambda_api_client_is_async(config.boto3_client): 189 ↛ 190line 189 didn't jump to line 190 because the condition on line 189 was never true
190 active_service_client = AsyncLambdaClient(
191 cast("AsyncLambdaApiClient", config.boto3_client)
192 )
193 else:
194 active_service_client = ThreadedSyncLambdaClient(
195 client=cast("LambdaApiClient", config.boto3_client)
196 )
197 return active_service_client
199 lambda_client = create_default_client()
200 if lambda_api_client_is_async(lambda_client):
201 active_service_client = AsyncLambdaClient(
202 cast("AsyncLambdaApiClient", lambda_client)
203 )
204 active_service_client_loop = current_loop
205 return active_service_client
207 active_service_client = ThreadedSyncLambdaClient(
208 client=cast("LambdaApiClient", lambda_client)
209 )
210 return active_service_client
212 async def async_wrapper(
213 event: Any, context: LambdaContext
214 ) -> MutableMapping[str, Any]:
215 service_client = await get_active_service_client()
216 return (await _wrapper_async(func, event, context, service_client)).to_dict()
218 @functools.wraps(func)
219 def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
220 return _run_on_event_loop(
221 get_or_create_handler_loop(), async_wrapper, event, context
222 )
224 setattr(wrapper, "_async_handler", async_wrapper)
225 setattr(wrapper, "_durable_execution_original", func)
226 setattr(wrapper, "_durable_execution_boto3_client", config.boto3_client)
228 return wrapper
231def _run_on_event_loop(
232 loop: asyncio.AbstractEventLoop,
233 async_func: Callable[..., Awaitable[MutableMapping[str, Any]]],
234 *args: Any,
235) -> MutableMapping[str, Any]:
236 try:
237 asyncio.get_running_loop()
238 except RuntimeError:
239 pass
240 else:
241 msg = (
242 "durable_execution sync handlers cannot be called from a running "
243 "event loop. Use the handler's _async_handler attribute instead."
244 )
245 raise RuntimeError(msg)
247 previous_loop: asyncio.AbstractEventLoop | None = None
248 had_previous_loop = True
249 try:
250 previous_loop = asyncio.get_event_loop()
251 except RuntimeError:
252 had_previous_loop = False
254 asyncio.set_event_loop(loop)
255 try:
256 return loop.run_until_complete(async_func(*args))
257 finally:
258 asyncio.set_event_loop(previous_loop if had_previous_loop else None)
261def deserialize_input(event: Any) -> DurableExecutionInvocationInput:
262 try:
263 logger.debug("durableExecutionArn: %s", event.get("DurableExecutionArn"))
264 return DurableExecutionInvocationInput.from_dict(event)
265 except (KeyError, TypeError, AttributeError) as e:
266 msg = (
267 "Unexpected payload provided to start the durable execution. "
268 "Check your resource configurations to confirm the durability is set."
269 )
270 raise ExecutionError(msg) from e
273async def _wrapper_async(
274 user_func: Callable[[Any], Any],
275 event: Any,
276 context: LambdaContext,
277 service_client: DurableServiceClient,
278) -> DurableExecutionInvocationOutput:
279 invocation_input = deserialize_input(event)
280 execution_state: ExecutionState = ExecutionState(
281 durable_execution_arn=invocation_input.durable_execution_arn,
282 initial_checkpoint_token=invocation_input.checkpoint_token,
283 service_client=service_client,
284 lambda_context=context,
285 )
287 try:
288 await execution_state.initialize(invocation_input)
290 input_event = execution_state.get_input_event()
292 root_context = DurableContext(
293 execution_state=execution_state,
294 operation_identifier=OperationIdentifier.create_execution_op(),
295 replaying=execution_state.has_prior_operations(),
296 )
298 if execution_state.get_execution_operation() is None: 298 ↛ 299line 298 didn't jump to line 299 because the condition on line 298 was never true
299 msg = "Execution state is missing the root execution operation."
300 raise RuntimeError(msg)
301 execution_state.start_checkpointing()
303 logger.debug("execution arn: %s", invocation_input.durable_execution_arn)
305 with bind_current_context(root_context):
306 result = await user_func(input_event)
307 return await handle_user_function_result(execution_state, result)
309 except SuspendExecution:
310 logger.debug("Suspending execution...")
311 return DurableExecutionInvocationOutput(status=InvocationStatus.PENDING)
312 except Exception as e:
313 return await handle_user_function_exception(execution_state, e)
314 finally:
315 await execution_state.aclose()
318async def handle_user_function_result(
319 execution_state, result
320) -> DurableExecutionInvocationOutput:
321 # done with userland
322 serialized_result = json.dumps(result)
323 # large response handling here. Remember if checkpointing to complete, NOT to include
324 # payload in response
325 if serialized_result and len(serialized_result) > LAMBDA_RESPONSE_SIZE_LIMIT:
326 logger.debug(
327 "Response size (%s bytes) exceeds Lambda limit (%s) bytes). Checkpointing result.",
328 len(serialized_result),
329 LAMBDA_RESPONSE_SIZE_LIMIT,
330 )
331 success_operation = OperationUpdate.create_execution_succeed(
332 payload=serialized_result
333 )
334 # Checkpoint large result with blocking (is_sync=True, default).
335 # Must ensure the result is persisted before returning to Lambda.
336 # Large results exceed Lambda response limits and must be stored durably
337 # before the execution completes.
338 await execution_state.create_checkpoint(success_operation, is_sync=True)
340 return DurableExecutionInvocationOutput.create_succeeded(result="")
341 return DurableExecutionInvocationOutput.create_succeeded(result=serialized_result)
344async def handle_user_function_exception(
345 execution_state, e: Exception
346) -> DurableExecutionInvocationOutput:
347 if isinstance(e, CheckpointError):
348 return handle_checkpoint_error(e)
349 if isinstance(e, InvocationError):
350 # Non-retryable Durable API errors (e.g., customer configuration issues,
351 # 4xx client errors) will never succeed on retry — fail the execution immediately.
352 if not e.is_retryable():
353 logger.exception(
354 "Non-retryable Durable API error. Must fail execution without retry.",
355 extra=e.build_logger_extras(),
356 )
357 return DurableExecutionInvocationOutput(
358 status=InvocationStatus.FAILED,
359 error=ErrorObject(
360 message=str(e),
361 type=_sdk_error_type_name(e),
362 ),
363 )
364 logger.exception("Invocation error. Must terminate.")
365 # Throw the error to trigger Lambda retry
366 raise
367 if isinstance(e, ExecutionError):
368 logger.exception("Execution error. Must fail execution without retry.")
369 return DurableExecutionInvocationOutput(
370 status=InvocationStatus.FAILED,
371 error=ErrorObject(
372 message=str(e),
373 type=_sdk_error_type_name(e),
374 ),
375 )
377 # all user-space errors go here
378 logger.exception("Execution failed")
380 result = DurableExecutionInvocationOutput(
381 status=InvocationStatus.FAILED, error=ErrorObject.from_exception(e)
382 )
384 serialized_result = json.dumps(result.to_dict())
386 if serialized_result and len(serialized_result) > LAMBDA_RESPONSE_SIZE_LIMIT:
387 logger.debug(
388 "Response size (%s bytes) exceeds Lambda limit (%s) bytes). Checkpointing result.",
389 len(serialized_result),
390 LAMBDA_RESPONSE_SIZE_LIMIT,
391 )
392 failed_operation = OperationUpdate.create_execution_fail(
393 error=ErrorObject.from_exception(e)
394 )
396 # Checkpoint large result with blocking (is_sync=True, default).
397 # Must ensure the result is persisted before returning to Lambda.
398 # Large results exceed Lambda response limits and must be stored durably
399 # before the execution completes.
400 try:
401 await execution_state.create_checkpoint(failed_operation, is_sync=True)
402 except CheckpointError as e:
403 return handle_checkpoint_error(e)
404 return DurableExecutionInvocationOutput(status=InvocationStatus.FAILED)
405 return result
408def handle_checkpoint_error(error: CheckpointError) -> DurableExecutionInvocationOutput:
409 """Convert checkpoint failures into a final result or retry trigger."""
410 # Checkpoint system is broken - stop background thread and exit immediately
411 logger.exception("Checkpoint system failed", extra=error.build_logger_extras())
412 if error.is_retryable():
413 raise error from None # Terminate Lambda immediately and have it be retried
414 return DurableExecutionInvocationOutput(
415 status=InvocationStatus.FAILED, error=ErrorObject.from_exception(error)
416 )