diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py index 9d328a44..27511809 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -174,8 +174,24 @@ def _attempt_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str: # Context scope helpers # ------------------------------------------------------------------ def _attach_context(self, key: str, new_context: Context) -> None: - """Attach a context and remember its token under ``key``.""" + """Attach a context and remember its token under ``key``. + + A token already stored under ``key`` means the previous scope for that + operation was never released: its user function did not reach + ``on_user_function_end``, because it suspended and a timed in-process + resume re-entered the same operation. Release it before attaching, so + the new token does not bury one that can never be detached -- otherwise + releasing the new scope would restore the abandoned span and leave it + current for later work on this thread. + """ with self._lock: + if key in self._context_tokens: + logger.debug( + "Releasing an unreleased context scope for %s before " + "re-attaching; its user function did not report an end.", + key, + ) + self._detach_context(key) self._context_tokens[key] = ( threading.get_ident(), otel_context.attach(new_context), diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index f702fc7a..9e31b0df 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -176,8 +176,24 @@ def _attempt_span_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str: # Context scope helpers # ------------------------------------------------------------------ def _attach_context(self, key: str, new_context: Context) -> None: - """Attach a context and remember its token under ``key``.""" + """Attach a context and remember its token under ``key``. + + A token already stored under ``key`` means the previous scope for that + operation was never released: its user function did not reach + ``on_user_function_end``, because it suspended and a timed in-process + resume re-entered the same operation. Release it before attaching, so + the new token does not bury one that can never be detached -- otherwise + releasing the new scope would restore the abandoned span and leave it + current for later work on this thread. + """ with self._operation_spans_lock: + if key in self._context_tokens: + logger.debug( + "Releasing an unreleased context scope for %s before " + "re-attaching; its user function did not report an end.", + key, + ) + self._detach_context(key) self._context_tokens[key] = ( threading.get_ident(), context.attach(new_context), diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py index f03cb337..1fb2c9db 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -2,6 +2,7 @@ from __future__ import annotations +import threading from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime from types import SimpleNamespace @@ -594,13 +595,15 @@ def _step_end_info( ) -def _context_start_info(operation_id: str) -> UserFunctionStartInfo: +def _context_start_info( + operation_id: str, parent_id: str | None = None +) -> UserFunctionStartInfo: return UserFunctionStartInfo( operation_id=operation_id, operation_type=OperationType.CONTEXT, sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, name=operation_id, - parent_id=None, + parent_id=parent_id, start_time=START_TIME, is_replayed=False, status=OperationStatus.STARTED, @@ -609,13 +612,15 @@ def _context_start_info(operation_id: str) -> UserFunctionStartInfo: ) -def _context_end_info(operation_id: str) -> UserFunctionEndInfo: +def _context_end_info( + operation_id: str, parent_id: str | None = None +) -> UserFunctionEndInfo: return UserFunctionEndInfo( operation_id=operation_id, operation_type=OperationType.CONTEXT, sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, name=operation_id, - parent_id=None, + parent_id=parent_id, start_time=START_TIME, is_replayed=False, status=OperationStatus.STARTED, @@ -803,3 +808,164 @@ def test_detach_ignores_token_attached_on_another_thread(): assert otel_context.get_current() == before_context plugin.on_invocation_end(_invocation_end_info()) + + +# --------------------------------------------------------------------------- +# Re-entering an operation whose scope was never released +# --------------------------------------------------------------------------- +def test_reentered_child_context_does_not_leave_abandoned_span_current(): + """Verify a timed in-process resume unwinds the abandoned scope. + + A suspended child context never reaches on_user_function_end, and the + map/parallel coordinator can resume that branch in-process, re-entering the + same operation ID on the same thread. Without releasing the first scope, the + second scope's detach would restore the abandoned span and leave it current. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before_context = otel_context.get_current() + context_id = "ctx-1" + + # First run: the child context suspends, so no end hook fires. + plugin.on_user_function_start(_context_start_info(context_id)) + suspended_span = plugin._get_span(context_id) + assert suspended_span is not None + + # Timed in-process resume re-enters the same operation. + plugin.on_user_function_start(_context_start_info(context_id)) + assert len([key for key in plugin._context_tokens if key == context_id]) == 1 + + plugin.on_user_function_end(_context_end_info(context_id)) + + assert otel_context.get_current() == before_context + assert context_id not in plugin._context_tokens + assert ( + trace.get_current_span().get_span_context().span_id + != suspended_span.get_span_context().span_id + ) + + plugin.on_invocation_end(_invocation_end_info()) + + +def test_reentered_step_attempt_releases_the_previous_scope(): + """Verify re-entering the same attempt key unwinds the previous scope.""" + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before_context = otel_context.get_current() + + plugin.on_user_function_start(_step_start_info("step-1")) + plugin.on_user_function_start(_step_start_info("step-1")) + plugin.on_user_function_end(_step_end_info("step-1")) + + assert otel_context.get_current() == before_context + assert set(plugin._context_tokens) == {"__invocation_context__"} + + plugin.on_invocation_end(_invocation_end_info()) + assert plugin._context_tokens == {} + + +def test_reentry_on_another_thread_leaves_the_originating_worker_dirty(): + """Pin what re-entry can and cannot clean up across threads. + + A resumed branch can land on a different pool thread than the one that + suspended. Re-entry drops the foreign token instead of resetting it, because + a context token can only be reset on its own thread, and it unwinds cleanly + on the thread that re-entered. The worker that suspended keeps the abandoned + span current: releasing it needs a hook invoked on that thread when the user + function fails to complete, which the SDK does not provide. The worker is + kept alive here so this limitation is asserted rather than hidden by pool + shutdown; the assertion flips once such a hook exists. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before_context = otel_context.get_current() + span_key = "step-1:attempt:1" + + with ThreadPoolExecutor(max_workers=1) as worker: + # The suspending run happens on the worker and never reports an end. + worker.submit( + plugin.on_user_function_start, _step_start_info("step-1") + ).result() + abandoned_span = plugin._get_span(span_key) + assert abandoned_span is not None + foreign_thread_ident, _foreign_token = plugin._context_tokens[span_key] + assert foreign_thread_ident != threading.get_ident() + + # The timed resume lands on this thread instead. + plugin.on_user_function_start(_step_start_info("step-1")) + assert plugin._context_tokens[span_key][0] == threading.get_ident() + plugin.on_user_function_end(_step_end_info("step-1")) + + # This thread unwound to where it started. + assert otel_context.get_current() == before_context + + # The originating worker is still carrying the abandoned span. + worker_span_id = worker.submit( + lambda: trace.get_current_span().get_span_context().span_id + ).result() + assert worker_span_id == abandoned_span.get_span_context().span_id + + plugin.on_invocation_end(_invocation_end_info()) + + +def test_nested_reentry_restores_the_abandoned_outer_scope(): + """Pin nested re-entry: correct ids, but the abandoned outer span object. + + When an outer child context and an inner one both suspend, re-entry releases + each scope in the order the operations are replayed, which is not the reverse + of the order they were attached. Ending the inner operation therefore + restores the scope captured for the abandoned outer span rather than the + resumed one. Deterministic CONTEXT span ids make the two indistinguishable + downstream -- same trace id and span id, so parenting and log correlation are + unaffected -- but the current span object is one that is never exported, so + anything an instrumentation library records on it is lost. Reverse-order + unwinding needs the SDK to report the suspension; this test documents the + current behaviour and flips when that lands. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before_context = otel_context.get_current() + + # Both contexts suspend, so neither reports an end. + plugin.on_user_function_start(_context_start_info("ctx-outer")) + abandoned_outer = plugin._get_span("ctx-outer") + plugin.on_user_function_start( + _context_start_info("ctx-inner", parent_id="ctx-outer") + ) + assert abandoned_outer is not None + + # The timed in-process resume replays both contexts, outer first. + plugin.on_user_function_start(_context_start_info("ctx-outer")) + resumed_outer = plugin._get_span("ctx-outer") + plugin.on_user_function_start( + _context_start_info("ctx-inner", parent_id="ctx-outer") + ) + resumed_inner = plugin._get_span("ctx-inner") + assert resumed_outer is not None + assert resumed_inner is not None + assert resumed_outer is not abandoned_outer + + # Resumed inner code runs under the resumed inner span. + assert trace.get_current_span() is resumed_inner + + plugin.on_user_function_end(_context_end_info("ctx-inner", parent_id="ctx-outer")) + + # The restored scope carries the abandoned outer span, whose ids match the + # resumed one because CONTEXT span ids are derived from the operation id. + assert trace.get_current_span() is abandoned_outer + assert ( + abandoned_outer.get_span_context().span_id + == resumed_outer.get_span_context().span_id + ) + assert ( + abandoned_outer.get_span_context().trace_id + == resumed_outer.get_span_context().trace_id + ) + + # Leaving the outer context still unwinds to where the invocation started. + plugin.on_user_function_end(_context_end_info("ctx-outer")) + assert otel_context.get_current() == before_context + assert set(plugin._context_tokens) == {"__invocation_context__"} + + plugin.on_invocation_end(_invocation_end_info()) + assert plugin._context_tokens == {} diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index 34bd8bc8..6e7e6742 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -2,6 +2,7 @@ from __future__ import annotations +import threading import time from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime @@ -1319,3 +1320,184 @@ def test_detach_ignores_token_attached_on_another_thread(): assert otel_context.get_current() == before_context plugin.on_invocation_end(_invocation_end_info()) + + +# ---------------------------------------------------------------------- +# Re-entering an operation whose scope was never released +# ---------------------------------------------------------------------- +def test_reentered_child_context_does_not_leave_abandoned_span_current(): + """Verify a timed in-process resume unwinds the abandoned scope. + + A suspended child context never reaches on_user_function_end, and the + map/parallel coordinator can resume that branch in-process, re-entering the + same operation ID on the same thread. Without releasing the first scope, the + second scope's detach would restore the abandoned span and leave it current. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before_context = otel_context.get_current() + context_id = "ctx-1" + + # First run: the child context suspends, so no end hook fires. + plugin.on_user_function_start( + _user_function_start_info(context_id, operation_type=OperationType.CONTEXT) + ) + suspended_span = plugin._get_span(context_id) + assert suspended_span is not None + + # Timed in-process resume re-enters the same operation. + plugin.on_user_function_start( + _user_function_start_info(context_id, operation_type=OperationType.CONTEXT) + ) + assert len([key for key in plugin._context_tokens if key == context_id]) == 1 + + plugin.on_user_function_end( + _user_function_end_info(context_id, operation_type=OperationType.CONTEXT) + ) + + assert otel_context.get_current() == before_context + assert context_id not in plugin._context_tokens + assert ( + trace.get_current_span().get_span_context().span_id + != suspended_span.get_span_context().span_id + ) + + plugin.on_invocation_end(_invocation_end_info()) + + +def test_reentered_step_attempt_releases_the_previous_scope(): + """Verify re-entering the same attempt key unwinds the previous scope.""" + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before_context = otel_context.get_current() + operation_id = "step-1" + + plugin.on_user_function_start(_user_function_start_info(operation_id)) + plugin.on_user_function_start(_user_function_start_info(operation_id)) + plugin.on_user_function_end(_user_function_end_info(operation_id)) + + assert otel_context.get_current() == before_context + assert plugin._context_tokens == {} + + plugin.on_invocation_end(_invocation_end_info()) + + +def test_reentry_on_another_thread_leaves_the_originating_worker_dirty(): + """Pin what re-entry can and cannot clean up across threads. + + A resumed branch can land on a different pool thread than the one that + suspended. Re-entry drops the foreign token instead of resetting it, because + a context token can only be reset on its own thread, and it unwinds cleanly + on the thread that re-entered. The worker that suspended keeps the abandoned + span current: releasing it needs a hook invoked on that thread when the user + function fails to complete, which the SDK does not provide. The worker is + kept alive here so this limitation is asserted rather than hidden by pool + shutdown; the assertion flips once such a hook exists. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before_context = otel_context.get_current() + operation_id = "step-1" + span_key = "step-1:attempt:1" + + with ThreadPoolExecutor(max_workers=1) as worker: + # The suspending run happens on the worker and never reports an end. + worker.submit( + plugin.on_user_function_start, _user_function_start_info(operation_id) + ).result() + abandoned_span = plugin._get_span(span_key) + assert abandoned_span is not None + foreign_thread_ident, _foreign_token = plugin._context_tokens[span_key] + assert foreign_thread_ident != threading.get_ident() + + # The timed resume lands on this thread instead. + plugin.on_user_function_start(_user_function_start_info(operation_id)) + assert plugin._context_tokens[span_key][0] == threading.get_ident() + plugin.on_user_function_end(_user_function_end_info(operation_id)) + + # This thread unwound to where it started. + assert otel_context.get_current() == before_context + + # The originating worker is still carrying the abandoned span. + worker_span_id = worker.submit( + lambda: trace.get_current_span().get_span_context().span_id + ).result() + assert worker_span_id == abandoned_span.get_span_context().span_id + + plugin.on_invocation_end(_invocation_end_info()) + + +def test_nested_reentry_restores_the_abandoned_outer_scope(): + """Pin nested re-entry: correct ids, but the abandoned outer span object. + + When an outer child context and an inner one both suspend, re-entry releases + each scope in the order the operations are replayed, which is not the reverse + of the order they were attached. Ending the inner operation therefore + restores the scope captured for the abandoned outer span rather than the + resumed one. Deterministic CONTEXT span ids make the two indistinguishable + downstream -- same trace id and span id, so parenting and log correlation are + unaffected -- but the current span object is one that is never exported, so + anything an instrumentation library records on it is lost. Reverse-order + unwinding needs the SDK to report the suspension; this test documents the + current behaviour and flips when that lands. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before_context = otel_context.get_current() + + # Both contexts suspend, so neither reports an end. + plugin.on_user_function_start( + _user_function_start_info("ctx-outer", operation_type=OperationType.CONTEXT) + ) + abandoned_outer = plugin._get_span("ctx-outer") + plugin.on_user_function_start( + _user_function_start_info( + "ctx-inner", parent_id="ctx-outer", operation_type=OperationType.CONTEXT + ) + ) + assert abandoned_outer is not None + + # The timed in-process resume replays both contexts, outer first. + plugin.on_user_function_start( + _user_function_start_info("ctx-outer", operation_type=OperationType.CONTEXT) + ) + resumed_outer = plugin._get_span("ctx-outer") + plugin.on_user_function_start( + _user_function_start_info( + "ctx-inner", parent_id="ctx-outer", operation_type=OperationType.CONTEXT + ) + ) + resumed_inner = plugin._get_span("ctx-inner") + assert resumed_outer is not None + assert resumed_inner is not None + assert resumed_outer is not abandoned_outer + + # Resumed inner code runs under the resumed inner span. + assert trace.get_current_span() is resumed_inner + + plugin.on_user_function_end( + _user_function_end_info( + "ctx-inner", parent_id="ctx-outer", operation_type=OperationType.CONTEXT + ) + ) + + # The restored scope carries the abandoned outer span, whose ids match the + # resumed one because CONTEXT span ids are derived from the operation id. + assert trace.get_current_span() is abandoned_outer + assert ( + abandoned_outer.get_span_context().span_id + == resumed_outer.get_span_context().span_id + ) + assert ( + abandoned_outer.get_span_context().trace_id + == resumed_outer.get_span_context().trace_id + ) + + # Leaving the outer context still unwinds to where the invocation started. + plugin.on_user_function_end( + _user_function_end_info("ctx-outer", operation_type=OperationType.CONTEXT) + ) + assert otel_context.get_current() == before_context + assert plugin._context_tokens == {} + + plugin.on_invocation_end(_invocation_end_info())