Skip to content

[plugin] Notify plugins when a user function does not complete #658

Description

@wangyb-A

Summary

ExecutionState.wrap_user_function reports a user-function end only on a normal return and on except Exception. Several exception types subclass BaseException and therefore bypass it, so a plugin that acquired per-attempt state in on_user_function_start is never told to release it — and never told on the thread that owns it.

# state.py:1163-1178
def wrapper(*args, **kwargs):
    start_info = self._plugin_executor.on_user_function_start(...)
    try:
        result = user_function(*args, **kwargs)
        self._plugin_executor.on_user_function_end(start_info, None)
        return result
    except SuspendExecution:
        raise                      # no end hook
    except Exception as e:
        self._plugin_executor.on_user_function_end(start_info, ErrorObject.from_exception(e))
        raise

Paths that skip the hook today:

  • SuspendExecution / TimedSuspendExecutionexceptions.py:468,479
  • OrphanedChildExceptionexceptions.py:549, raised at state.py:536, state.py:642, state.py:952
  • BackgroundThreadErrorexceptions.py:447
  • SystemExit / KeyboardInterrupt

Why this matters beyond bookkeeping

The OTel plugins attach an OpenTelemetry context scope in on_user_function_start and release it in on_user_function_end. When the end hook never fires, the scope stays attached to the worker thread that ran the user function, and a contextvars token can only be reset on its creating thread — so nothing else can clean it up.

Concrete consequences observed while fixing #643:

  1. In-process timed resume. The map/parallel coordinator resumes a timed suspend inside the same invocation (concurrency/executor.py, the while timed_resumes and timed_resumes[0][0] <= now: ... submit(branch) loop), re-entering the same operation ID. fix(otel): release unreleased scope on operation re-entry #654 added a guard that releases the previous scope on re-entry, which fixes the token loss but only on the re-entering thread.
  2. Nested suspends unwind out of order. With an outer child context and an inner one both suspended, re-entry releases scopes in replay order rather than reverse attach order, so ending the resumed inner operation restores the scope captured for the abandoned outer span. Deterministic CONTEXT span ids make the ids identical, so parenting and log correlation are unaffected, but the current span object is one that is never exported.
  3. Resume on a different worker. If the resumed branch lands on another pool thread, the original worker keeps the abandoned span current for the windows between attached scopes.
  4. Orphaned branches. OrphanedChildException abandons a scope with no re-entry at all, so no plugin-side guard can reach it.

Cases 2-4 are pinned by tests in #654 that assert current behaviour and are expected to fail once this issue is fixed:

  • test_nested_reentry_restores_the_abandoned_outer_scope
  • test_reentry_on_another_thread_leaves_the_originating_worker_dirty

both in packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py and .../test_invocation_plugin.py.

Proposed change

Fire a new notification from a finally block, gated on whether an end hook already ran:

completed = False
try:
    result = user_function(*args, **kwargs)
    self._plugin_executor.on_user_function_end(start_info, None)
    completed = True
    return result
except SuspendExecution:
    raise
except Exception as e:
    self._plugin_executor.on_user_function_end(start_info, ErrorObject.from_exception(e))
    completed = True
    raise
finally:
    if not completed:
        self._plugin_executor.on_user_function_complete(start_info)

wrap_user_function is the single choke point for step, child-context, wait-for-condition and map/parallel branch user functions, so one call site covers every path, and the hook runs on the thread that executed the user code — which is what makes cleanup possible at all.

Alternative considered and not recommended

Adding UserFunctionOutcome.SUSPENDED and reusing on_user_function_end is smaller, but it fires an end hook where none fired before. Any existing plugin that treats "not FAILED" as success would end and export a span for an attempt that never completed. An additive hook with a default no-op avoids that.

Acceptance criteria

  • A plugin is notified, on the thread that ran the user function, whenever that function does not report an end: suspension (timed and untimed), OrphanedChildException, BackgroundThreadError, and SystemExit/KeyboardInterrupt.
  • The notification never double-fires alongside on_user_function_end.
  • The hook's documented contract states that it runs on the user-code thread and that implementations must release scoped state without ending spans, because the operation may still resume.
  • Both OTel plugins release the attached context scope from the new hook, so nested suspends unwind in reverse order and a suspended scope is released on its originating thread.
  • The two tripwire tests above are updated to assert the corrected behaviour.
  • Plugin exceptions remain swallowed and logged, as with the existing hooks (PluginExecutor._dispatch_plugin).

Work breakdown and estimate

Item Files Size
New info dataclass, base hook, dispatch case, executor method plugin.py (OperationInfo at :90, UserFunctionEndInfo.from_start_info at :216, base hooks :378-436, _dispatch_plugin :478-500, executor end hook :653) ~45 lines
finally wiring state.py:1163-1178 ~6 lines
Release the scope from the new hook both OTel plugins ~40 lines
Flip the tripwire assertions 2 OTel test files ~30 lines
New tests: fires on suspend / orphan / background error, no double-fire, nested reverse-order unwind, originating-thread release tests/plugin_test.py, tests/state_test.py, both OTel test files ~250 lines
Docs "Dynamic instrumentation plugins" in packages/aws-durable-execution-sdk-python/README.md (line 30+) ~20 lines

Roughly 400 lines, 1-2 days including review. An end-to-end test driving a real parallel branch that waits while a sibling works is also needed to prove the hook fires on the right thread in the real coordinator; note [tool.hatch.envs.dev-otel] in the root pyproject.toml does not depend on aws-durable-execution-sdk-python-testing (the root test env does), so that dev-only dependency has to be added for such a test to live under the OTel package.

Open questions

References

Metadata

Metadata

Assignees

Labels

bugpkg:otelPackage: aws-durable-execution-sdk-python-otelpkg:sdkPackage: aws-durable-execution-sdk-python

Type

No type

Projects

Status
In progress

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions