Coverage for async_durable_execution/_core/context.py: 94%
216 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +0000
1from __future__ import annotations
3import functools
4import hashlib
5import logging
6from collections.abc import Iterator
7from contextlib import contextmanager
8from contextvars import ContextVar, Token
9from dataclasses import dataclass
10from typing import TYPE_CHECKING, Any, cast
12from .exceptions import InvalidStateError
13from .models import (
14 Operation,
15 OperationIdentifier,
16 OperationStatus,
17 OperationSubType,
18 OperationSubTypeValue,
19 OperationType,
20)
22if TYPE_CHECKING:
23 from .models import LambdaContext
24 from .state import ExecutionState
27logger = logging.getLogger(__name__)
30@dataclass(frozen=True)
31class SerDesContext:
32 """Context for serialization operations and composable SerDes stages.
34 The first three fields retain the original public construction order.
35 Pipeline stages additionally receive durable operation metadata and the
36 root value being serialized through ``original_value``. Deserialization
37 stages always receive ``original_value=None``.
38 """
40 operation_id: str = ""
41 durable_execution_arn: str = ""
42 recursive_level: int = 0
43 entity_id: str = ""
44 operation_name: str | None = None
45 parent_id: str | None = None
46 operation_type: OperationType | None = None
47 operation_sub_type: OperationSubTypeValue | None = None
48 attempt: int | None = None
49 original_value: Any = None
52class OperationIdGenerator:
53 """Generate deterministic operation ids within a durable execution scope."""
55 def __init__(self, prefix: str | None) -> None:
56 self._prefix = prefix
57 self._counter = 0
58 self._claimed_local_ids: set[str] = set()
59 self._unconsumed_reservations: dict[str, bool] = {}
60 self._unconsumed_checkpoint_count = 0
61 self._reservation_selection_started = False
62 self._replay_frontier_pending = False
64 def increment(self) -> int:
65 self._counter += 1
66 return self._counter
68 def get_current(self) -> int:
69 return self._counter
71 def _create_id(self, value: str) -> str:
72 """Hash one context-local identity value."""
73 prefix = self._prefix
74 step_id = f"{prefix}-{value}" if prefix else value
75 return hashlib.blake2b(step_id.encode()).hexdigest()[:64]
77 def _create_id_for_local_id(self, local_id: str) -> str:
78 """Generate an id in the caller-provided local-id namespace."""
79 return self._create_id(f"local:{local_id}")
81 def _create_step_id_for_logical_step(self, step: int) -> str:
82 """Generate the stable operation id for a logical step."""
83 return self._create_id(str(step))
85 def create_step_id(self) -> str:
86 """Generate an operation id and advance the logical step counter."""
87 return self._create_step_id_for_logical_step(self.increment())
89 def create_step_id_for_local_id(self, local_id: str) -> str:
90 """Generate an operation id from a stable caller-provided local id."""
91 if not isinstance(local_id, str): 91 ↛ 92line 91 didn't jump to line 92 because the condition on line 91 was never true
92 msg = "local_operation_id must be a string"
93 raise TypeError(msg)
94 if not local_id.strip(): 94 ↛ 95line 94 didn't jump to line 95 because the condition on line 94 was never true
95 msg = "local_operation_id must not be blank"
96 raise ValueError(msg)
97 if self._reservation_selection_started:
98 msg = (
99 "local_operation_id reservations must be created before any "
100 "reserved operation is selected"
101 )
102 raise RuntimeError(msg)
103 if local_id in self._claimed_local_ids: 103 ↛ 104line 103 didn't jump to line 104 because the condition on line 103 was never true
104 msg = f"local_operation_id is already reserved: {local_id}"
105 raise ValueError(msg)
107 self._claimed_local_ids.add(local_id)
108 return self._create_id_for_local_id(local_id)
110 def _register_reservation(
111 self,
112 operation_id: str,
113 *,
114 has_checkpoint: bool,
115 ) -> None:
116 """Track an allocated reservation until workflow code selects it."""
117 previous = self._unconsumed_reservations.get(operation_id)
118 if operation_id in self._unconsumed_reservations:
119 if previous == has_checkpoint: 119 ↛ 121line 119 didn't jump to line 121 because the condition on line 119 was always true
120 return
121 if previous:
122 self._unconsumed_checkpoint_count -= 1
123 self._unconsumed_reservations[operation_id] = has_checkpoint
124 if has_checkpoint:
125 self._unconsumed_checkpoint_count += 1
127 def _consume_reservation(self, operation_id: str) -> None:
128 """Discard a reservation after workflow code selects it."""
129 self._mark_reservation_selected()
130 has_checkpoint = self._unconsumed_reservations.pop(operation_id, False)
131 if has_checkpoint:
132 self._unconsumed_checkpoint_count -= 1
134 def _mark_reservation_selected(self) -> None:
135 """Prevent explicit local ids from being registered after selection."""
136 self._reservation_selection_started = True
138 def _has_unconsumed_checkpoint(self) -> bool:
139 """Return whether any allocated reservation still has replay history."""
140 return self._unconsumed_checkpoint_count > 0
142 def _mark_replay_frontier(self) -> None:
143 """Remember that replay ended immediately before a virtual scope."""
144 self._replay_frontier_pending = True
146 def _clear_replay_frontier(self) -> None:
147 """Clear a replay frontier after the next operation boundary is known."""
148 self._replay_frontier_pending = False
150 def _is_replay_frontier_pending(self) -> bool:
151 """Return whether a virtual scope may still contain flattened history."""
152 return self._replay_frontier_pending
155@dataclass(frozen=True)
156class OperationContext:
157 """Base context shared by durable execution scopes."""
159 execution_state: ExecutionState
160 operation_identifier: OperationIdentifier
162 @property
163 def lambda_context(self) -> LambdaContext | None:
164 """Return the Lambda context for the active invocation."""
165 return self.execution_state.lambda_context
167 @property
168 def durable_execution_arn(self) -> str:
169 """Return the ARN of the durable execution."""
170 return self.execution_state.durable_execution_arn
172 @property
173 def parent_id(self) -> str | None:
174 return self.operation_identifier.parent_id
176 @property
177 def operation_id(self) -> str | None:
178 return self.operation_identifier.operation_id
180 @property
181 def operation_name(self) -> str | None:
182 return self.operation_identifier.name
184 @property
185 def recursive_level(self) -> int:
186 """Return the recursion depth recorded on the execution input."""
187 return self.execution_state.recursive_level
189 def is_replaying(self) -> bool:
190 """Return whether the active context is replaying prior user code."""
191 return False
194@dataclass(frozen=True)
195class DurableContext(OperationContext):
196 """Runtime context available to a durable handler or child context."""
198 step_id_prefix: str | None = None
199 replaying: bool = False
201 @functools.cached_property
202 def step_counter(self) -> OperationIdGenerator:
203 return OperationIdGenerator(self.operation_id_generator_prefix)
205 @property
206 def is_virtual(self) -> bool:
207 return self.operation_identifier.parent_id != self.operation_id_generator_prefix
209 @property
210 def operation_id_generator_prefix(self) -> str | None:
211 return (
212 self.step_id_prefix
213 if self.step_id_prefix is not None
214 else self.operation_identifier.parent_id
215 )
217 def create_child_context(
218 self,
219 operation_id: str,
220 *,
221 is_virtual: bool = False,
222 replaying: bool | None = None,
223 ) -> DurableContext:
224 """Create a child context for the given operation."""
225 child_parent_id = self.parent_id if is_virtual else operation_id
226 logger.debug(
227 "Creating child context for operation %s (is_virtual=%s)",
228 operation_id,
229 is_virtual,
230 )
231 return DurableContext(
232 execution_state=self.execution_state,
233 operation_identifier=OperationIdentifier(
234 operation_id=None,
235 sub_type=OperationSubType.EXECUTION,
236 parent_id=child_parent_id,
237 ),
238 step_id_prefix=operation_id,
239 replaying=self.is_replaying() if replaying is None else replaying,
240 )
242 def is_replaying(self) -> bool:
243 """Return True while this context is replaying prior operations."""
244 return self.replaying
246 def _set_replay_status_new(self) -> None:
247 object.__setattr__(self, "replaying", False)
248 self.step_counter._clear_replay_frontier() # noqa: SLF001
250 def _set_replay_status_frontier(self) -> None:
251 """End parent replay while retaining a snapshot for a virtual child."""
252 object.__setattr__(self, "replaying", False)
253 self.step_counter._mark_replay_frontier() # noqa: SLF001
255 def _virtual_child_replay_snapshot(self) -> bool:
256 """Return replay state including flattened history past the frontier."""
257 return (
258 self.is_replaying() or self.step_counter._is_replay_frontier_pending() # noqa: SLF001
259 )
261 def _peek_next_operation_id(self) -> str:
262 return self.step_counter._create_step_id_for_logical_step( # noqa: SLF001
263 self.step_counter.get_current() + 1
264 )
266 def _next_operation_result(self) -> Operation | None:
267 return self._operation_result(self._peek_next_operation_id())
269 def _operation_result(self, operation_id: str) -> Operation | None:
270 return self.execution_state.operations.get(operation_id)
272 def _next_operation_exists(self) -> bool:
273 return self._next_operation_result() is not None
275 def _next_reserved_or_sequential_operation_exists(self) -> bool:
276 if self.step_counter._has_unconsumed_checkpoint(): # noqa: SLF001
277 return True
278 return self._next_operation_exists()
280 def _next_operation_is_terminal_checkpoint(self) -> bool:
281 return self._operation_is_terminal_checkpoint(self._peek_next_operation_id())
283 def _operation_is_terminal_checkpoint(self, operation_id: str) -> bool:
284 operation = self._operation_result(operation_id)
285 if operation is None:
286 return False
287 return operation.status in {
288 OperationStatus.SUCCEEDED,
289 OperationStatus.FAILED,
290 OperationStatus.CANCELLED,
291 OperationStatus.STOPPED,
292 OperationStatus.TIMED_OUT,
293 }
295 @contextmanager
296 def _replay_aware(
297 self,
298 *,
299 operation_id: str | None = None,
300 executes_user_code: bool = False,
301 consume_reservation: bool = True,
302 ) -> Iterator[None]:
303 """Update replay status around one durable operation.
305 `operation_id` identifies an operation that was allocated before entering
306 this scope. Pre-allocated reservations are tracked independently from the
307 sequential counter so launch order does not end replay prematurely.
308 """
309 was_replaying = self.is_replaying()
310 self.step_counter._clear_replay_frontier() # noqa: SLF001
311 current_operation_id = operation_id or self._peek_next_operation_id()
312 if operation_id is not None and consume_reservation:
313 self.step_counter._consume_reservation(operation_id) # noqa: SLF001
314 current_exists = was_replaying and (
315 self._operation_result(current_operation_id) is not None
316 )
317 current_terminal = was_replaying and self._operation_is_terminal_checkpoint(
318 current_operation_id
319 )
320 flip_after = (
321 was_replaying
322 and not executes_user_code
323 and current_exists
324 and not current_terminal
325 )
327 if was_replaying and (
328 not current_exists or (executes_user_code and not current_terminal)
329 ):
330 self._set_replay_status_new()
332 try:
333 yield
334 finally:
335 if flip_after:
336 self._set_replay_status_frontier()
337 elif self.is_replaying():
338 next_operation_exists = (
339 self._next_reserved_or_sequential_operation_exists()
340 if operation_id is not None
341 else self._next_operation_exists()
342 )
343 if not next_operation_exists:
344 self._set_replay_status_frontier()
347_current_context: ContextVar = ContextVar(
348 "async_durable_execution.current_context",
349 default=None,
350)
351_durable_definition_operation: ContextVar[str | None] = ContextVar(
352 "async_durable_execution.durable_definition_operation",
353 default=None,
354)
357def ensure_durable_operations_allowed(operation_name: str) -> None:
358 """Reject durable operation creation during a synchronous definition phase."""
359 definition_operation = _durable_definition_operation.get()
360 if definition_operation is None:
361 return
363 msg = (
364 f"{operation_name} cannot be used while defining a "
365 f"{definition_operation}. Durable operations may only run after the "
366 "definition has been validated."
367 )
368 raise InvalidStateError(msg)
371@contextmanager
372def bind_durable_definition(operation_name: str) -> Iterator[None]:
373 """Mark a synchronous definition phase in the current context."""
374 token = _durable_definition_operation.set(operation_name)
375 try:
376 yield
377 finally:
378 _durable_definition_operation.reset(token)
381def set_current_context(context) -> Token:
382 """Bind the active durable context for the current async task."""
383 return _current_context.set(context)
386def reset_current_context(token: Token) -> None:
387 """Restore the previous durable context after a temporary override."""
388 _current_context.reset(token)
391def get_current_context() -> OperationContext | SerDesContext:
392 """Return the currently active durable execution context.
394 Raises:
395 RuntimeError: If called outside supported durable user code.
396 """
397 ensure_durable_operations_allowed("Durable operations")
398 current_context = _current_context.get()
399 if current_context is None:
400 msg = (
401 "get_current_context() can only be used while a durable function, "
402 "step function, flow node, wait_for_callback submitter, or "
403 "wait_for_condition check, or SerDes operation is executing."
404 )
405 raise RuntimeError(msg)
406 return current_context
409def get_durable_context() -> DurableContext:
410 """Return the current context after validating durable operations are allowed."""
411 current_context = get_current_context()
412 if (
413 not hasattr(current_context, "execution_state")
414 or not hasattr(current_context, "operation_identifier")
415 or not hasattr(current_context, "step_counter")
416 or not hasattr(current_context, "create_child_context")
417 ):
418 operation_name = getattr(current_context, "operation_name", None)
419 msg = (
420 f"{operation_name or 'Durable operations'} can only be used while a "
421 "durable function or child context is executing."
422 )
423 raise RuntimeError(msg)
424 return cast(DurableContext, current_context)
427@contextmanager
428def bind_current_context(
429 context: OperationContext | SerDesContext,
430) -> Iterator[None]:
431 """Temporarily bind the supplied durable context while invoking user code."""
432 token = set_current_context(context)
433 try:
434 yield
435 finally:
436 reset_current_context(token)