villebro opened a new pull request, #43658:
URL: https://github.com/apache/superset/pull/43658
### SUMMARY
Follow-up to #43627 (merged), closing the one gap it explicitly deferred and
tightening the reaper's schedule.
**1. Cancel the orphaned warehouse query.** #43627 made async chart-data
query tasks cancellable *while the worker is alive* (an `on_abort` handler runs
`db_engine_spec.cancel_query` over a fresh connection). But when a **worker
dies**, there's no live thread to run that handler, so the reaper only marked
the `Task` `FAILURE` and revoked the Celery job — the abandoned warehouse query
kept running until the warehouse's own statement/idle timeout. Now the reaper
cancels it out-of-band.
**2. Dedicated reaper beat job.** Reaping was folded into `prune_tasks`.
Reaping wants a *short* cadence (detect a dead worker — and cancel its query —
within ~a minute); retention pruning is a heavy, infrequent bulk DELETE. One
shared cadence can't serve both, so reaping now runs on its own
`reap_orphaned_tasks` beat schedule.
#### How the orphan-cancel works
The reaper runs in a different process from the dead worker, so the engine
cancel-id must be on the `Task` row:
- **Persist at capture** — when the capture sink obtains a cancel-id,
`TaskContext.set_cancellation(database_id, cancel_id)` merges
`{cancel_database_id, cancel_query_id}` into the property cache. It does
**not** write: the task registers its abort handler immediately after, and that
`is_abortable` write already flushes the whole cache — so the handle is
persisted in an existing write, **no extra DB write** on the hot path. Only
engines that expose a cancel-id before execution
(Postgres/MySQL/Snowflake/Redshift, as in the live path) persist a handle.
- **Reaper cancels** — `ReapOrphanedTasksCommand._reap` reads the handle,
loads the `Database`, and reuses the existing `cancel_chart_query(database,
cancel_id)` helper (fresh connection, id validation, `gtf.query.cancel` /
`gtf.query.cancel_failed` metrics) **before** the revoke + CAS→`FAILURE`.
Best-effort: no handle, a missing database, or a cancel failure never blocks
the FAILURE transition. A dead worker holds no lock on the row, so this is
race-free.
No new helper, config key, or migration — `cancel_chart_query` and the
reaper already exist; this reuses them. The live abort path is unchanged (it
still uses the in-memory closure).
#### Scheduling change
- `reap_orphaned_tasks` is a new Celery beat task; `prune_tasks` reverts to
retention-only. `CELERY_CONFIG.beat_schedule` gains a commented
`reap_orphaned_tasks` template on a short interval (e.g. every minute),
separate from the `prune_tasks` retention template. Docs/UPDATING updated.
### TESTING INSTRUCTIONS
- Unit (`tests/unit_tests/tasks/`): `TaskContext.set_cancellation` merges
the handle into the cache (persisted by the ensuing `is_abortable` write,
existing keys preserved); the capture sink calls `set_cancellation` before
`on_abort` (and neither without a cancel-id); `_reap` calls
`cancel_chart_query` with the persisted `(database, cancel_id)` when present,
skips it when absent or the database is gone, and reaps regardless;
`reap_orphaned_tasks` delegates to the command.
- Integration (`tests/integration_tests/tasks/commands/test_reap.py`): an
orphan with a persisted handle → the reaper invokes `cancel_chart_query` with
the loaded `Database` and still ends the row `FAILURE`.
- Manual (Celery worker + `reap_orphaned_tasks` beat +
`DISTRIBUTED_COORDINATION_CONFIG`, a cancellable engine like Postgres): run an
async chart query, `kill -9` the worker mid-query → within
`GTF_ORPHAN_TASK_TIMEOUT` the task is `FAILURE` and the backend query is gone
from `pg_stat_activity`.
### ADDITIONAL INFORMATION
- [ ] Has associated issue:
- [x] Required feature flags: `GLOBAL_TASK_FRAMEWORK` (async chart data also
needs `GLOBAL_ASYNC_QUERIES`)
- [ ] Changes UI
- [ ] Includes DB Migration (the `cancel_*` handle lives in the existing
`properties` JSON blob — no schema change)
- [x] Introduces new feature or API
- [ ] Removes existing feature or API
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
