Coverage for async-durable-execution/src/async_durable_execution/execution.py: 96%

1from __future__ import annotations 

2 

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 

10 

11from .context import bind_current_context 

12from .primitive.child import DurableContext 

13from .exceptions import ( 

14 CheckpointError, 

15 ExecutionError, 

16 InvocationError, 

17 SuspendExecution, 

18) 

19from .models import ( 

20 DurableExecutionInvocationOutput, 

21 ErrorObject, 

22 InvocationStatus, 

23 Operation, 

24 OperationUpdate, 

25 OperationIdentifier, 

26 SerializableModel, 

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 

36 

37 

38if TYPE_CHECKING: 

39 from collections.abc import MutableMapping 

40 

41 from .client import AsyncLambdaApiClient, DurableServiceClient, LambdaApiClient 

42 from .models import LambdaContext 

43 

44configure_durable_logger(logging.getLogger()) 

45logger = logging.getLogger(__name__) 

46 

47# 6MB in bytes, minus 50 bytes for envelope 

48LAMBDA_RESPONSE_SIZE_LIMIT = 6 * 1024 * 1024 - 50 

49 

50T = TypeVar("T") 

51Params = ParamSpec("Params") 

52 

53 

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. 

58 

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. 

62 

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__)) 

70 

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 

78 

79 return wrapper 

80 

81 

82@dataclass(frozen=True) 

83class InitialExecutionState(SerializableModel): 

84 """Initial page of operation history included with an invocation event.""" 

85 

86 operations: list[Operation] = field( 

87 default_factory=list, 

88 metadata={"alias": "Operations"}, 

89 ) 

90 next_marker: str = field(default="", metadata={"alias": "NextMarker"}) 

91 

92 

93@dataclass(frozen=True) 

94class DurableExecutionInvocationInput(SerializableModel): 

95 """Event payload delivered to a durable Lambda invocation.""" 

96 

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 ) 

103 

104 

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.""" 

110 

111 handler_attrs = getattr(handler, "__dict__", {}) 

112 original_handler = handler_attrs.get("_durable_execution_original") 

113 if original_handler is None: 

114 return handler 

115 

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 ) 

122 

123 

124@dataclass(frozen=True) 

125class DurableConfig: 

126 boto3_client: LambdaApiClient | AsyncLambdaApiClient | None = None 

127 service_client: DurableServiceClient | None = None 

128 

129 

130def durable_execution( 

131 func: Callable[..., Awaitable[Any]] | None = None, 

132 /, 

133 **kwargs, 

134) -> Callable[[Any, LambdaContext], Any]: 

135 """ 

136 Decorator to create a durable execution handler. 

137 

138 Args: 

139 func: The user function to decorate 

140 boto3_client: Optional sync or async Lambda API client to use 

141 service_client: Optional durable service client to use. Intended for 

142 testing and local execution tooling. 

143 """ 

144 # Decorator called with parameters 

145 if func is None: 

146 logger.debug("Decorator called with parameters") 

147 return functools.partial( 

148 durable_execution, 

149 **kwargs, 

150 ) 

151 config = DurableConfig(**kwargs) 

152 logger.debug("Starting durable execution handler...") 

153 

154 # Use the explicitly provided durable client when present. Otherwise, delay 

155 # Lambda API client construction until invocation so importing decorated handlers 

156 # does not require AWS environment configuration. 

157 active_service_client = config.service_client 

158 active_service_client_loop: asyncio.AbstractEventLoop | None = None 

159 handler_loop: asyncio.AbstractEventLoop | None = None 

160 

161 def get_or_create_handler_loop() -> asyncio.AbstractEventLoop: 

162 nonlocal handler_loop 

163 if handler_loop is None or handler_loop.is_closed(): 

164 handler_loop = asyncio.new_event_loop() 

165 return handler_loop 

166 

167 async def get_active_service_client() -> DurableServiceClient: 

168 nonlocal active_service_client, active_service_client_loop 

169 current_loop = asyncio.get_running_loop() 

170 if active_service_client is not None: 

171 if ( 

172 active_service_client_loop is None 

173 or active_service_client_loop is current_loop 

174 ): 

175 return active_service_client 

176 

177 # The SDK-owned async client is bound to the loop that first used it. 

178 # If test code calls the async handler on another loop, rebuild it there. 

179 active_service_client = None 

180 active_service_client_loop = None 

181 

182 if config.boto3_client is not None: 

183 if lambda_api_client_is_async(config.boto3_client): 

184 active_service_client = AsyncLambdaClient( 

185 cast("AsyncLambdaApiClient", config.boto3_client) 

186 ) 

187 else: 

188 active_service_client = ThreadedSyncLambdaClient( 

189 client=cast("LambdaApiClient", config.boto3_client) 

190 ) 

191 return active_service_client 

192 

193 lambda_client = create_default_client() 

194 if lambda_api_client_is_async(lambda_client): 

195 active_service_client = AsyncLambdaClient( 

196 cast("AsyncLambdaApiClient", lambda_client) 

197 ) 

198 active_service_client_loop = current_loop 

199 return active_service_client 

200 

201 active_service_client = ThreadedSyncLambdaClient( 

202 client=cast("LambdaApiClient", lambda_client) 

203 ) 

204 return active_service_client 

205 

206 async def async_wrapper( 

207 event: Any, context: LambdaContext 

208 ) -> MutableMapping[str, Any]: 

209 service_client = await get_active_service_client() 

210 return (await _wrapper_async(func, event, context, service_client)).to_dict() 

211 

212 @functools.wraps(func) 

213 def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: 

214 return _run_on_event_loop( 

215 get_or_create_handler_loop(), async_wrapper, event, context 

216 ) 

217 

218 setattr(wrapper, "_async_handler", async_wrapper) 

219 setattr(wrapper, "_durable_execution_original", func) 

220 setattr(wrapper, "_durable_execution_boto3_client", config.boto3_client) 

221 

222 return wrapper 

223 

224 

225def _run_on_event_loop( 

226 loop: asyncio.AbstractEventLoop, 

227 async_func: Callable[..., Awaitable[MutableMapping[str, Any]]], 

228 *args: Any, 

229) -> MutableMapping[str, Any]: 

230 try: 

231 asyncio.get_running_loop() 

232 except RuntimeError: 

233 pass 

234 else: 

235 msg = ( 

236 "durable_execution sync handlers cannot be called from a running " 

237 "event loop. Use the handler's _async_handler attribute instead." 

238 ) 

239 raise RuntimeError(msg) 

240 

241 previous_loop: asyncio.AbstractEventLoop | None = None 

242 had_previous_loop = True 

243 try: 

244 previous_loop = asyncio.get_event_loop() 

245 except RuntimeError: 

246 had_previous_loop = False 

247 

248 asyncio.set_event_loop(loop) 

249 try: 

250 return loop.run_until_complete(async_func(*args)) 

251 finally: 

252 asyncio.set_event_loop(previous_loop if had_previous_loop else None) 

253 

254 

255def deserialize_input(event: Any) -> DurableExecutionInvocationInput: 

256 try: 

257 logger.debug("durableExecutionArn: %s", event.get("DurableExecutionArn")) 

258 return DurableExecutionInvocationInput.from_json_dict(event) 

259 except (KeyError, TypeError, AttributeError) as e: 

260 msg = ( 

261 "Unexpected payload provided to start the durable execution. " 

262 "Check your resource configurations to confirm the durability is set." 

263 ) 

264 raise ExecutionError(msg) from e 

265 

266 

267async def _wrapper_async( 

268 user_func: Callable[[Any], Any], 

269 event: Any, 

270 context: LambdaContext, 

271 service_client: DurableServiceClient, 

272) -> DurableExecutionInvocationOutput: 

273 invocation_input = deserialize_input(event) 

274 execution_state: ExecutionState = ExecutionState( 

275 durable_execution_arn=invocation_input.durable_execution_arn, 

276 initial_checkpoint_token=invocation_input.checkpoint_token, 

277 service_client=service_client, 

278 lambda_context=context, 

279 ) 

280 

281 try: 

282 await execution_state.initialize(invocation_input) 

283 

284 input_event = execution_state.get_input_event() 

285 

286 root_context = DurableContext( 

287 execution_state=execution_state, 

288 operation_identifier=OperationIdentifier.create_execution_op(), 

289 replaying=execution_state.has_prior_operations(), 

290 ) 

291 

292 if execution_state.get_execution_operation() is None: 

293 msg = "Execution state is missing the root execution operation." 

294 raise RuntimeError(msg) 

295 execution_state.start_checkpointing() 

296 

297 logger.debug("execution arn: %s", invocation_input.durable_execution_arn) 

298 

299 with bind_current_context(root_context): 

300 result = await user_func(input_event) 

301 return await handle_user_function_result(execution_state, result) 

302 

303 except SuspendExecution: 

304 logger.debug("Suspending execution...") 

305 return DurableExecutionInvocationOutput(status=InvocationStatus.PENDING) 

306 except Exception as e: 

307 return await handle_user_function_exception(execution_state, e) 

308 finally: 

309 await execution_state.aclose() 

310 

311 

312async def handle_user_function_result( 

313 execution_state, result 

314) -> DurableExecutionInvocationOutput: 

315 # done with userland 

316 serialized_result = json.dumps(result) 

317 # large response handling here. Remember if checkpointing to complete, NOT to include 

318 # payload in response 

319 if serialized_result and len(serialized_result) > LAMBDA_RESPONSE_SIZE_LIMIT: 

320 logger.debug( 

321 "Response size (%s bytes) exceeds Lambda limit (%s) bytes). Checkpointing result.", 

322 len(serialized_result), 

323 LAMBDA_RESPONSE_SIZE_LIMIT, 

324 ) 

325 success_operation = OperationUpdate.create_execution_succeed( 

326 payload=serialized_result 

327 ) 

328 # Checkpoint large result with blocking (is_sync=True, default). 

329 # Must ensure the result is persisted before returning to Lambda. 

330 # Large results exceed Lambda response limits and must be stored durably 

331 # before the execution completes. 

332 await execution_state.create_checkpoint(success_operation, is_sync=True) 

333 

334 return DurableExecutionInvocationOutput.create_succeeded(result="") 

335 return DurableExecutionInvocationOutput.create_succeeded(result=serialized_result) 

336 

337 

338async def handle_user_function_exception( 

339 execution_state, e: Exception 

340) -> DurableExecutionInvocationOutput: 

341 if isinstance(e, CheckpointError): 

342 return handle_checkpoint_error(e) 

343 if isinstance(e, InvocationError): 

344 # Non-retryable Durable API errors (e.g., customer configuration issues, 

345 # 4xx client errors) will never succeed on retry — fail the execution immediately. 

346 if not e.is_retryable(): 

347 logger.exception( 

348 "Non-retryable Durable API error. Must fail execution without retry.", 

349 extra=e.build_logger_extras(), 

350 ) 

351 return DurableExecutionInvocationOutput( 

352 status=InvocationStatus.FAILED, 

353 error=ErrorObject.from_exception(e), 

354 ) 

355 logger.exception("Invocation error. Must terminate.") 

356 # Throw the error to trigger Lambda retry 

357 raise 

358 if isinstance(e, ExecutionError): 

359 logger.exception("Execution error. Must fail execution without retry.") 

360 return DurableExecutionInvocationOutput( 

361 status=InvocationStatus.FAILED, 

362 error=ErrorObject.from_exception(e), 

363 ) 

364 

365 # all user-space errors go here 

366 logger.exception("Execution failed") 

367 

368 result = DurableExecutionInvocationOutput( 

369 status=InvocationStatus.FAILED, error=ErrorObject.from_exception(e) 

370 ) 

371 

372 serialized_result = json.dumps(result.to_dict()) 

373 

374 if serialized_result and len(serialized_result) > LAMBDA_RESPONSE_SIZE_LIMIT: 

375 logger.debug( 

376 "Response size (%s bytes) exceeds Lambda limit (%s) bytes). Checkpointing result.", 

377 len(serialized_result), 

378 LAMBDA_RESPONSE_SIZE_LIMIT, 

379 ) 

380 failed_operation = OperationUpdate.create_execution_fail( 

381 error=ErrorObject.from_exception(e) 

382 ) 

383 

384 # Checkpoint large result with blocking (is_sync=True, default). 

385 # Must ensure the result is persisted before returning to Lambda. 

386 # Large results exceed Lambda response limits and must be stored durably 

387 # before the execution completes. 

388 try: 

389 await execution_state.create_checkpoint(failed_operation, is_sync=True) 

390 except CheckpointError as e: 

391 return handle_checkpoint_error(e) 

392 return DurableExecutionInvocationOutput(status=InvocationStatus.FAILED) 

393 return result 

394 

395 

396def handle_checkpoint_error(error: CheckpointError) -> DurableExecutionInvocationOutput: 

397 """Convert checkpoint failures into a final result or retry trigger.""" 

398 # Checkpoint system is broken - stop background thread and exit immediately 

399 logger.exception("Checkpoint system failed", extra=error.build_logger_extras()) 

400 if error.is_retryable(): 

401 raise error from None # Terminate Lambda immediately and have it be retried 

402 return DurableExecutionInvocationOutput( 

403 status=InvocationStatus.FAILED, error=ErrorObject.from_exception(error) 

404 )