feat(plugin): report user functions that do not complete - #659
Conversation
wrap_user_function reported an outcome only when the user function returned or raised an ordinary Exception. SuspendExecution, TimedSuspendExecution, OrphanedChildException, BackgroundThreadError and SystemExit all subclass BaseException and bypassed both handlers, so a plugin that bound state to the user-function thread in on_user_function_start was never told to release it -- and never told on the thread that owns it. - Add on_user_function_incomplete, dispatched synchronously like the start and end hooks, with a UserFunctionIncompleteInfo payload - Fire it from a finally block in wrap_user_function, gated on whether an outcome was already reported, so exactly one of end/incomplete follows every start - Release the OTel context scope from the new hook in both plugins, leaving the span open because the operation may still resume - Document the hook and its threading contract in the core README Nested scopes now unwind in reverse order: the hook fires as the exception propagates outward, so an inner context is released before its enclosing one, and a resume no longer restores a scope captured for the suspended run. A suspended branch also releases its scope on the worker that ran it, which the plugins could not do from another thread. The two limitation tests added with the earlier re-entry mitigation now assert the corrected behaviour. Resolves #658
| OperationStartInfo, | ||
| OperationType, | ||
| UserFunctionEndInfo, | ||
| UserFunctionIncompleteInfo, |
There was a problem hiding this comment.
Codex AI review
[P1] Raise the OTel package's minimum core SDK version. UserFunctionIncompleteInfo is absent from the currently supported core 1.8.0, while OTel still declares aws-durable-execution-sdk-python>=1.8.0. Environments pinned to 1.8.0 will resolve the new OTel package but fail during import and plugin entry-point loading. Update both the package dependency and test-pypi-otel environment to the first core release containing this hook, or provide a compatibility fallback.
| def suspend_on_worker() -> tuple[int, bool]: | ||
| plugin.on_user_function_start(_step_start_info("step-1")) | ||
| attached_span_id = trace.get_current_span().get_span_context().span_id | ||
| plugin.on_user_function_incomplete(_step_incomplete_info("step-1")) |
There was a problem hiding this comment.
Codex AI review
[P2] Add end-to-end coverage for automatic core-to-OTel dispatch. This test manually invokes the new callback, while the core tests use only a recording plugin, so both suites would pass if a real suspended branch never delivered this hook correctly. Add an integration test using the local runner and a real map/parallel suspension that verifies the originating worker is clean and resumed span parenting is correct, as required for this new public cross-component behavior.
Codex AI reviewFound a release compatibility break and a required cross-component test gap. Review was static; repository code was not executed. Reviewed commit |
Resolves #658. Root fix for the OTel context leak that #650 and #654 could only mitigate from inside the plugins.
Problem
ExecutionState.wrap_user_functionreported an outcome only when the user function returned or raised an ordinaryException. Five exception types subclassBaseExceptionand bypassed both handlers:SuspendExecution/TimedSuspendExecution(exceptions.py:468,479)OrphanedChildException(exceptions.py:549)BackgroundThreadError(exceptions.py:447)SystemExit/KeyboardInterruptA plugin that bound state to the user-function thread in
on_user_function_startwas therefore never told to release it, and — critically — never told on the thread that owns it. For the OTel plugins that state is an OpenTelemetry context scope, and acontextvarstoken can only be reset on its creating thread, so no amount of plugin-side bookkeeping could clean it up.Change
Core. New
on_user_function_incompletehook with aUserFunctionIncompleteInfopayload, dispatched synchronously like the start and end hooks. It fires from afinallyinwrap_user_function, gated on anoutcome_reportedflag, so exactly one of end/incomplete follows every start. One call site covers step, child-context, wait-for-condition and map/parallel branch user functions.OTel plugins. Both release the attached context scope from the new hook and leave the span open and registered — the operation is not finished, a suspended one resumes and ends later. Also factored the duplicated CONTEXT-vs-STEP key computation into a
_user_function_key/_user_function_span_keyhelper now that three hooks need it.Docs. New "User-function hooks and thread-bound state" section in the core README, stating the one-of-two contract and that an incomplete user function is not a finished operation.
What this fixes that the mitigations could not
OrphanedChildExceptionhad no re-entry to trigger the_attach_contextguard, so nothing released its scope.The guard added in #654 stays as a belt-and-braces measure: plugin hook exceptions are swallowed by
PluginExecutor._dispatch_plugin, so a plugin whose hook throws before releasing still recovers on the next attach for that key.API compatibility
Additive. The base class provides a default no-op, so existing plugins are unaffected and
DURABLE_INSTRUMENTATION_PLUGIN_API_VERSIONstays at 1 — happy to bump it if the team would rather signal the new surface. The alternative of reusingon_user_function_endwith a newSUSPENDEDoutcome was rejected in #658: it would fire an end hook where none fired before, so a plugin treating "not FAILED" as success would export a span for an attempt that never completed.Tests
Core (
tests/state_test.py):OTel (both plugins):
test_suspension_releases_the_scope_on_the_originating_workerandtest_nested_suspension_unwinds_scopes_in_reverse_orderresumed_outerVerification
hatch run test:all— 3252 passed, 2 skippedhatch run dev-core:test/dev-otel:test— greenhatch run types:checkanddev-core:typecheck— cleanhatch fmt --checkin both changed packages — cleanFollow-up not in this PR
The end-to-end runner test asked for in review on #654 — a real parallel branch that waits while a sibling works, asserting resumed instrumentation parenting — is still outstanding. It needs
aws-durable-execution-sdk-python-testingadded to thedev-otelhatch env (the roottestenv already has it) and is worth doing as its own PR now that this hook is what it would be validating.