Skip to content

refactor(python): share one task lifecycle across both dispatch paths#438

Open
pratyush618 wants to merge 13 commits into
masterfrom
refactor/shared-task-lifecycle
Open

refactor(python): share one task lifecycle across both dispatch paths#438
pratyush618 wants to merge 13 commits into
masterfrom
refactor/shared-task-lifecycle

Conversation

@pratyush618

@pratyush618 pratyush618 commented Jul 16, 2026

Copy link
Copy Markdown
Collaborator

Why

Native async dispatch has never run: the pool reads _taskito_is_async off the registry entry, but
the decorator sets it on the TaskWrapper it returns, so every async task takes the blocking branch.
Activating it was a two-line change gated behind one problem — _wrap_task and
AsyncTaskExecutor._run_job were two independent implementations of one lifecycle, and the native
copy had silently lost the queue hooks, the saga compensation context, soft_timeout and per-item
batch handling.

This PR removes that gate by making the two paths share one body, and fixes the bugs that surfaced
while doing it. It does not activate native dispatch — see below.

What changed

One lifecycle. New taskito/task_lifecycle.py::run_lifecycle holds everything between "we have
the arguments" and "we have a result". _wrap_task drives it with run_maybe_async; the executor
awaits it. Each caller keeps only what genuinely differs — how arguments arrive, and how the result
is delivered. executor.py loses 101 lines and no longer owns a single lifecycle concern, so a new
lifecycle rule can't reach one path and miss the other.

Real bugs, each fixed with a regression test watched failing first:

  • A torn-down task was reported as completed. except Exception doesn't catch a BaseException,
    so the cleanup saw no error and emitted JOB_COMPLETED — and handed error=None to
    after_task/middleware — for a task whose body never returned. Latent on the blocking path;
    asyncio.CancelledError is a BaseException and is routine once dispatch goes native.
  • A failed task's exception outlived its lifecycle. It was held in a plain frame local (Python
    auto-clears an except ... as name for exactly this reason; error is not one) and passed as a
    logging argument, which a LogRecord keeps. Either one retains the exception, its traceback, and
    every frame that raised.
  • A sync task returning a coroutine broke — caught in review. It works on master and
    dead-lettered here: the lifecycle is itself a coroutine now, so the inner run_maybe_async nested
    inside its own event loop and raised. Awaiting whatever is awaitable covers both kinds and made
    the is_async parameter redundant.
  • Setup could strand what it acquired. Resources and proxies are taken before the body is
    entered, but only the body's teardown gave them back — so a raising before_task hook leaked
    them. Setup now unwinds explicitly.
  • Native dispatch dropped a torn-down job. A BaseException that is not a cancellation slipped
    past every clause in the executor, so the job reported nothing and sat Running until the reaper
    called it a timeout.

Tests: 22 new, parametrised over def/async def where the guarantee is shared and driven through a real @queue.task() and
a real worker — hooks, saga context, soft_timeout (which had no coverage at all), per-item batch,
logging, middleware-on-cancel. Each was watched failing with its guarantee removed. The async cases
route sync while dispatch is dormant; they are the net for activating it.

tests/worker/test_native_async.py grows a _fake_queue() helper, replacing ten hand-listed mock
setups. That mattered immediately: the lifecycle branches on is not None and truthiness, so a
MagicMock's auto-made _task_batch_configs reads as a per-item batch config and every task fails the
return-type contract. One list, not ten.

Why activation isn't here

Flipping both flags works — coroutines run on taskito-async-executor and all 18 parity tests pass
on the native path. Then measurement showed a failed async task stalls worker shutdown for the full
30s drain timeout
: 30.13s with nothing more exotic than an on_failure hook that collects errors.
That is past a typical container's grace period, so in-flight jobs get SIGKILLed.

The worker's drain loop ends only when the result channel disconnects, which needs the executor's
PyResultSender to drop — and a failed task's traceback pins the frame that owns it. The two
retention sources above are fixed, but a user hook that keeps the exception can't be fixed from
Python: every frame in that traceback can reach the sender. The drain has to stop depending on Python
object lifetime, which is a design change worth its own PR. ROADMAP.md records the findings for it,
including why Scheduler.in_flight is a tempting but unsafe drain signal.

So the flags stay off and the note explaining why is updated to name the real blocker.

Verification

  • 1239 passed, 4 skipped — the 1217 baseline plus these 22, no regressions.
  • ruff check / ruff format / mypy clean on taskito/ and tests/.
  • cargo check --workspace, and with --features native-async (no Rust changed).
  • Every fix watched failing with the fix removed, and the shutdown numbers measured directly rather
    than inferred (0.17s vs 30.13s).

Summary by CodeRabbit

  • Bug Fixes
    • Improved consistency between synchronous and asynchronous task execution for hooks, middleware, timeouts, batch failures, compensation context, logging, and lifecycle events.
    • Fixed cancellation and fatal-interruption handling to avoid incorrect completion reporting and to ensure worker shutdown doesn’t stall after failures.
    • Strengthened per-task soft-timeout behavior (including clearing on re-registration) and ensured coroutine returns from sync tasks are awaited.
  • Tests
    • Expanded lifecycle parity and regression coverage across sync/async execution paths, including middleware error handling, saga context visibility, and cancellation/batch edge cases.

The blocking wrapper and the async executor each implement the lifecycle a task
runs through, and they have drifted: the native copy silently lost the queue
hooks, the saga compensation context, soft_timeout and per-item batch handling.
Patching the gaps into the executor would make a second copy that drifts again.

Move the body to taskito/task_lifecycle.py and have the wrapper drive it. Pure
move — the callers keep only what genuinely differs, which is where arguments
come from and where the result goes. current_job needs no special handling: its
context lookup already reads the contextvar before the thread-local.

soft_timeout moves into a _task_soft_timeouts dict, mirroring the eleven other
per-task options. It was the only one reachable solely through the wrapper's
closure, which is why the executor could never honour it.

_clear_context stays with the wrapper: it is edge state, and in the shared body
it would null out the thread-local of whichever thread the executor loop runs on.

No behaviour change — 1217 tests unchanged.
Drives each lifecycle guarantee through a real worker, parametrised over def
and async def. The async cases route sync while native dispatch is dormant;
they are the net for activating it.
`except Exception` misses a BaseException, so the cleanup saw no error and
emitted JOB_COMPLETED for a task whose body never returned. Latent today;
asyncio.CancelledError is a BaseException and is routine once async tasks
dispatch natively.
The executor reimplemented the lifecycle and had drifted, silently missing the
queue hooks, saga context, soft_timeout and per-item batch handling. It now
awaits run_lifecycle and keeps only what differs: decoding the payload and
answering the result sender.
The lifecycle held the exception in a plain local and passed it as a log
argument, so both the frame and any retained LogRecord kept it — and through its
traceback, every frame that raised. Latent on the blocking path; on native
dispatch those frames own the result sender the worker's drain loop waits on.
The lifecycle gaps the old note cites are closed; shutdown is what is left.
@coderabbitai

coderabbitai Bot commented Jul 16, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 6fcc9a9e-4ab2-4c69-8fdf-41dcb5b160c3

📥 Commits

Reviewing files that changed from the base of the PR and between 3741c95 and 06cafb0.

📒 Files selected for processing (3)
  • sdks/python/taskito/task_lifecycle.py
  • sdks/python/tests/worker/test_lifecycle_parity.py
  • sdks/python/tests/worker/test_native_async.py
🚧 Files skipped from review as they are similar to previous changes (3)
  • sdks/python/taskito/task_lifecycle.py
  • sdks/python/tests/worker/test_native_async.py
  • sdks/python/tests/worker/test_lifecycle_parity.py

📝 Walkthrough

Walkthrough

The PR centralizes task lifecycle execution in run_lifecycle, routes blocking and native async dispatch through it, stores per-task soft timeouts on queues, and adds sync/async parity and native executor test coverage.

Changes

Shared Task Lifecycle

Layer / File(s) Summary
Central lifecycle runner
sdks/python/taskito/task_lifecycle.py, sdks/python/tests/observability/test_task_lifecycle_logs.py
Adds shared lifecycle handling for setup, execution, resources, middleware, hooks, batching, saga context, errors, cleanup, logging, and lifecycle events.
Blocking dispatch integration
sdks/python/taskito/app.py, sdks/python/taskito/mixins/decorators.py
Stores task soft timeouts on the queue and delegates blocking task execution to run_lifecycle.
Native async integration
sdks/python/taskito/async_support/executor.py
Delegates native async jobs to run_lifecycle, retaining concurrency gating, deserialization, failure hand-off, and async context cleanup.
Sync and async lifecycle parity tests
sdks/python/tests/worker/test_lifecycle_parity.py
Adds worker-backed coverage for hooks, failures, timeouts, batching, saga context, logging, middleware, teardown events, and shutdown.
Native async executor test support
sdks/python/tests/worker/test_native_async.py
Adds shared fake queue setup, preserves serializer and backpressure coverage, and verifies cancellation emits JOB_FAILED without JOB_COMPLETED.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Worker
  participant AsyncTaskExecutor
  participant run_lifecycle
  participant TaskFunction
  Worker->>AsyncTaskExecutor: Deserialize payload and start job
  AsyncTaskExecutor->>run_lifecycle: Await lifecycle execution
  run_lifecycle->>TaskFunction: Execute task and await awaitable result
  TaskFunction-->>run_lifecycle: Return result or raise exception
  run_lifecycle-->>AsyncTaskExecutor: Return result or re-raise exception
Loading

Possibly related PRs

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 45.28% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly summarizes the main refactor: unifying task lifecycle handling across sync and async dispatch paths.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch refactor/shared-task-lifecycle

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 7

🧹 Nitpick comments (1)
sdks/python/tests/worker/test_native_async.py (1)

24-52: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Make _fake_queue reject undeclared lifecycle attributes.

A bare MagicMock still fabricates newly accessed attributes, recreating the silent wrong-branch behavior this helper is intended to prevent. Use spec_set or a concrete stub so new lifecycle dependencies fail until explicitly configured.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@sdks/python/tests/worker/test_native_async.py` around lines 24 - 52, Update
_fake_queue to construct the queue mock with spec_set or a concrete stub that
declares every lifecycle attribute configured and accessed by the helper.
Preserve the existing defaults and overrides, while ensuring any newly accessed
undeclared attribute raises immediately instead of being auto-created.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@sdks/python/taskito/mixins/decorators.py`:
- Around line 419-421: The task registration logic around _wrap_task must remove
any existing _task_soft_timeouts entry when soft_timeout is None, while
continuing to store a provided timeout. Ensure replacing a task by task_name
cannot inherit a previous timeout.

In `@sdks/python/taskito/task_lifecycle.py`:
- Around line 269-314: Wrap the teardown operations in the finally
block—including compensation reset, resource release, proxy cleanup, hooks,
middleware, and event emission—in a nested try/finally, and move the existing
`error = None` into that nested finally. Ensure exception clearing runs even
when any teardown callback or cleanup operation raises, while preserving the
current cleanup order and event behavior.
- Around line 117-187: Move the cleanup guard to begin before proxy
reconstruction and resource acquisition in the task execution flow. Enclose
proxy cleanup, resource release callbacks, middleware setup, and before_task
hooks in an outer try/finally so exceptions during reconstruction, acquisition,
middleware lookup, or hooks still run all registered cleanup; preserve the
existing task execution error handling inside that guard.
- Around line 188-191: The run_lifecycle execution path must await awaitable
results returned by synchronous tasks instead of calling run_maybe_async inside
its active event loop. Update task_lifecycle.py lines 188-191 around
run_lifecycle to detect and await returned awaitables while preserving direct
handling for non-awaitable values; retain run_maybe_async in
mixins/decorators.py lines 157-166 as the blocking entrypoint, requiring no
direct change there.
- Around line 240-256: Handle non-cancellation BaseException cases in
async_support/executor.py around the async executor flow by emitting a terminal
non-retryable failure or equivalent disposition before re-raising
KeyboardInterrupt, SystemExit, GeneratorExit, and similar exceptions. In
task_lifecycle.py around the existing BaseException handler, preserve the
teardown-only behavior and do not add failure hooks or disposition reporting
there; the executor must own terminal reporting before propagation.

In `@sdks/python/tests/worker/test_lifecycle_parity.py`:
- Around line 124-150: Update the failure capture in on_fail and the
corresponding hook handlers at
sdks/python/tests/worker/test_lifecycle_parity.py:124-150, 320-349, and 366-404
to store only each exception’s type and message (or equivalent derived values),
never the exception object or traceback; adjust the assertions to validate those
stored values while preserving the existing failure checks before _stop_worker
runs.
- Around line 65-68: Update _stop_worker to verify that the worker thread has
terminated after thread.join(timeout=10), and fail teardown if thread.is_alive()
remains true. Preserve the existing shutdown request and timeout while ensuring
a still-running daemon thread cannot allow the test to pass.

---

Nitpick comments:
In `@sdks/python/tests/worker/test_native_async.py`:
- Around line 24-52: Update _fake_queue to construct the queue mock with
spec_set or a concrete stub that declares every lifecycle attribute configured
and accessed by the helper. Preserve the existing defaults and overrides, while
ensuring any newly accessed undeclared attribute raises immediately instead of
being auto-created.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 0a52a6f4-2f29-4aa4-8f2f-0cb3cd52f4c1

📥 Commits

Reviewing files that changed from the base of the PR and between 7623ba3 and 7dd99ac.

📒 Files selected for processing (7)
  • sdks/python/taskito/app.py
  • sdks/python/taskito/async_support/executor.py
  • sdks/python/taskito/mixins/decorators.py
  • sdks/python/taskito/task_lifecycle.py
  • sdks/python/tests/observability/test_task_lifecycle_logs.py
  • sdks/python/tests/worker/test_lifecycle_parity.py
  • sdks/python/tests/worker/test_native_async.py

Comment thread sdks/python/taskito/mixins/decorators.py Outdated
Comment thread sdks/python/taskito/task_lifecycle.py
Comment thread sdks/python/taskito/task_lifecycle.py Outdated
Comment thread sdks/python/taskito/task_lifecycle.py
Comment thread sdks/python/taskito/task_lifecycle.py Outdated
Comment thread sdks/python/tests/worker/test_lifecycle_parity.py
Comment thread sdks/python/tests/worker/test_lifecycle_parity.py Outdated
A plain `def` handing back a coroutine worked before the lifecycle became one:
run_maybe_async now nests inside the lifecycle's own loop and raises. Deciding
on the returned value covers both kinds and drops the is_async parameter.
Resources and proxies are acquired before the body is entered, so a raising
before_task hook stranded them. The teardown could also skip clearing the
exception if a hook raised, which was the point of clearing it.
Registering a name again replaces the task, but left the old timeout in the map
for the new one to inherit.
A BaseException that is not a cancellation slipped past every clause, so the job
reported nothing and sat Running until the reaper called it a timeout.

The queue stub is now concrete: a bare MagicMock invents whatever the lifecycle
asks for, which is how a missing attribute takes the wrong branch in silence.
The stop helper returned on a live worker, which is how the shutdown stall in
this series nearly escaped: assertions had already passed, so only the clock
showed it. The hooks also held live exceptions — the stall's own trigger.
@pratyush618

Copy link
Copy Markdown
Collaborator Author

Thanks — all seven addressed, and one of them was a genuine regression I'd introduced. Details, since a couple needed digging:

Awaitable returns (task_lifecycle.py:191) — real regression, good catch. A plain def returning a coroutine worked on master and dead-lettered on this branch: the lifecycle is now a coroutine driven by asyncio.run, so the inner run_maybe_async saw a running loop and raised. Verified against a master worktree (returns 42) versus this branch (RuntimeError: Cannot run an async ... already has a running event loop).

Fixed as suggested — await what is awaitable — which also made the is_async parameter redundant, so it's gone from run_lifecycle and both callers. Deciding on the returned value rather than the declared kind is the more honest test anyway. Regression test: test_sync_task_returning_a_coroutine_is_awaited.

Cleanup guard before acquiring (:187). Correct, and it survived a mutation check: with the guard removed, a before_task hook that raises leaves a request-scoped resource unreleased (_release_acquired never called); with it, the resource comes back. Setup now unwinds explicitly via a shared _release_acquired. Kept the semantics narrow on purpose — a task that never started doesn't fire after_task or emit lifecycle events, so the failure path gives back only what it took. Test: test_setup_failure_releases_what_it_acquired (watched failing without the fix).

Guaranteed exception clearing (:314). Right — and it defeated the whole point of the fix, since a raising after_task hook would have skipped the clear and left the cycle. Teardown is now a nested try/finally with error = None in the inner finally.

Non-cancellation BaseException in the executor (:256). Fixed with a BaseException fallback that reports a failure via _hand_off_now (no awaiting on a loop that may be going down) and re-raises. Reported as a failure with the retry policy left to decide, matching what the blocking path does with the same exception. Test: test_base_exception_reports_failure.

soft_timeout on re-registration (decorators.py:421). Real bug in code this PR added; pop on None as suggested. Test: test_re_registering_a_task_clears_its_soft_timeout.

_stop_worker hiding a live thread. Accepted, and this one stung: it's exactly how the shutdown stall documented in the PR body nearly escaped me — the assertions had already passed, so the tests were green and only the wall clock (10s per failing async test) showed anything wrong. The helper now asserts the thread stopped.

Tests retaining exceptions. Accepted — they now record type and message. Retention is the documented stall condition, so holding live exceptions in the very tests meant to be the net for activation would have made each one slow while proving nothing extra.

Nitpick — spec_set on _fake_queue. Agreed on the principle, and a bare MagicMock did re-open the trap the helper exists to close. spec_set=Queue doesn't work (the attributes the lifecycle reads live on instances, not the class), so it's now a concrete _FakeQueue with __slots__ — an undeclared lifecycle read raises, and so does a typo'd override.

Verification: full suite green, ruff/mypy clean, and each fix watched failing with the fix removed.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@sdks/python/taskito/task_lifecycle.py`:
- Around line 178-210: Update the setup-failure handling around the lifecycle
orchestration so every middleware recorded in completed_mw receives its after()
teardown when a later before_task hook or saga setup step raises. Invoke
teardown for completed middleware before releasing resources, while preserving
compensation-context reset, cleanup, and exception propagation.

In `@sdks/python/tests/worker/test_native_async.py`:
- Around line 576-600: Update test_base_exception_reports_failure to verify
shutdown after executor.stop(), using the existing shutdown assertion helper if
available or asserting that the executor thread is no longer alive. Keep the
existing failure-reporting and KeyboardInterrupt behavior unchanged.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 4c0d22d4-546c-4ef2-ac38-dc5b44db0e6f

📥 Commits

Reviewing files that changed from the base of the PR and between 7dd99ac and 3741c95.

📒 Files selected for processing (5)
  • sdks/python/taskito/async_support/executor.py
  • sdks/python/taskito/mixins/decorators.py
  • sdks/python/taskito/task_lifecycle.py
  • sdks/python/tests/worker/test_lifecycle_parity.py
  • sdks/python/tests/worker/test_native_async.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • sdks/python/taskito/mixins/decorators.py

Comment thread sdks/python/taskito/task_lifecycle.py
Comment thread sdks/python/tests/worker/test_native_async.py
`completed_mw` exists so every before() gets an after(); the setup unwind
released resources but skipped it, leaving middleware-owned state behind.
`stop()` only logs when its join times out, so the test could pass on a leaked
thread — the same way the worker stop helper hid a stall.
@pratyush618

Copy link
Copy Markdown
Collaborator Author

Both accepted — each one a gap in my own fix from the last round.

Middleware teardown on setup failure. Right, and I'd drawn the line in the wrong place. I reasoned that a task which never started shouldn't fire after_task or emit lifecycle events, and kept the unwind to resources only — but middleware isn't in that category. completed_mw exists precisely to pair every before() with an after(), so a middleware that set something up is owed the chance to take it down whether or not the task it prepared for ever ran. The setup unwind now runs after(job, None, exc) for the completed ones, and test_setup_failure_undoes_what_setup_did covers it (verified failing with the pairing removed: "middleware that ran before() never got its after()").

The hooks stay as they were, deliberately: there's no completed_hooks equivalent, so partial-hook pairing would be a new semantic rather than an existing invariant being honoured.

Asserting the executor thread stopped. Fair, and the same trap as _stop_worker in the last round — stop() only logs on join timeout, so with the thread-exception warning filtered the test really could have gone green on a leaked thread. Now asserts executor._thread is not alive.

Full suite green, ruff/mypy clean.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant