villebro opened a new pull request, #43627: URL: https://github.com/apache/superset/pull/43627
### SUMMARY Part of the GAQ→GTF epic (targets `gaq-to-gtf`, not `master`). Two resilience gaps remained in the Global Task Framework once async chart data moved onto it, both called out in the code itself: 1. **Orphaned tasks** — a worker killed mid-execution (OOM, SIGKILL, crash, lost broker message) left the `Task` row stuck `IN_PROGRESS` forever and its Celery job a zombie. `prune_tasks` only deleted *old terminal* rows; the only backstop was client-side (`GLOBAL_ASYNC_QUERIES_POLLING_STALE_TIMEOUT`), which just makes the browser give up while the task + Celery job + warehouse query keep running. 2. **Non-cancellable queries** — `execute_chart_query` deliberately set no timeout and no abort handler, because the chart-data path (`Database.get_df` → `db_engine_spec.execute`) has none of SQL Lab's cancellation plumbing, so a GTF abort/timeout only flipped the row terminal while the warehouse kept churning. This PR adds a worker **heartbeat + orphan reaper** (folded into the existing `prune_tasks` cron) and makes chart-data query tasks **cancellable** on engines that support it, then restores an (opt-in) per-query timeout. #### What GAQ had, what GTF has today, and what this adds | Capability | Legacy GAQ | GTF (today, pre-PR) | After this PR | |---|---|---|---| | Detect a worker that died mid-task | ❌ none | ❌ none (task stuck `IN_PROGRESS` forever) | ✅ liveness heartbeat + reaper marks it `FAILURE` | | Kill the Celery job | ✅ `revoke(terminate, SIGUSR1)` | ❌ cooperative DB abort only | ✅ tiered: cooperative first, forced `revoke` as escalation/reaper | | Reclaim a wedged worker slot | ✅ (blunt `revoke`) | ❌ (task must cooperate) | ✅ reaper force-`revoke`s a task wedged in `ABORTING` | | Graceful cleanup on cancel | ⚠️ can interrupt mid-cleanup | ✅ `on_abort` / `on_cleanup` run | ✅ preserved (cooperative path runs first) | | Cancel the actual warehouse query | ⚠️ SQL-Lab STOPPED flag only, not chart data | ❌ none for chart data | ✅ engine `cancel_query` over a fresh connection (supported engines) | | Per-query timeout for async chart data | ❌ | ❌ (removed pending cancellation) | ✅ opt-in `GLOBAL_ASYNC_QUERIES_QUERY_TIMEOUT` | | Terminal state after a cancelled query | ⚠️ marked failed/cancelled ad hoc | n/a | ✅ `ABORTED` / `TIMED_OUT` (not `FAILURE`) | | Source of truth | Redis job registry / flags | DB `Task` row (CAS) | DB `Task` row (CAS); `revoke`/`cancel` are best-effort accelerators | #### Design notes - **Heartbeat** (`superset/tasks/heartbeat.py`): while a worker holds a task it bumps `tasks.last_heartbeat` on a daemon thread every `GTF_TASK_HEARTBEAT_INTERVAL` (default 15s). The write is a raw `text()` UPDATE so it **never advances `changed_on`** — otherwise every heartbeat would resurface the task in the `status_changes` poll and reset client backoff. `execute_task` is split into a thin wrapper (records the Celery job id, wraps the body in the heartbeat) + `_execute_task_body`; the heartbeat starts **before** the DAG gate so a task blocked on prerequisites still counts as alive. - **Reaper** (`ReapOrphanedTasksCommand`, run from `prune_tasks` before deletion): a dead-worker orphan (stale heartbeat) is `revoke`d + CAS-transitioned to `FAILURE` + `publish_completion` (so waiters/DAG dependents/pollers unblock); a live worker wedged in `ABORTING` past the grace window is force-`revoke`d only, letting its own `finally` finalize. The CAS is authoritative, so a revived worker that already committed a terminal status just no-ops the reaper. - **Query cancellation** (`superset/tasks/query_cancel.py`): a contextvar cursor-capture seam (hooked into `Database._execute_sql_with_mutation_and_logging` before the blocking execute) lets the task capture an engine cancel id via `db_engine_spec.get_cancel_query_id`; the abort handler then kills the backend over a fresh connection (`cancel_chart_query`, mirroring `sql_lab.cancel_query`). Cancellation auto-enables **per engine** — Postgres/MySQL/Snowflake/Redshift are cancellable; engines without cancel support stay non-abortable and behave exactly as before. - **Metrics** on every op: `gtf.task.heartbeat`, `gtf.task.heartbeat_failure`, `gtf.task.orphan_reaped`, `gtf.task.revoke`, `gtf.task.revoke_failure`, `gtf.task.abort_escalated`, `gtf.query.cancel`, `gtf.query.cancel_failed`. #### Known follow-up (intentionally out of scope) Out-of-band cancellation of an **orphan's** warehouse query (from the reaper) is deferred — it needs the cancel handle persisted on the task, which races the wholesale `properties` writes. The reaper still revokes + fails the orphaned task; its abandoned query relies on warehouse-side timeouts until then. **Live** user-abort/timeout cancellation is fully implemented. ### TESTING INSTRUCTIONS - Unit + integration: `pytest tests/unit_tests/tasks/ tests/integration_tests/tasks/commands/test_reap.py` — 408 unit tests pass. New coverage: heartbeat lifecycle/failure, `find_orphaned` orphan + wedged-abort selection (DB-backed), all reaper branches (orphan / wedged / CAS-race / no-celery-id), the executor except-guard (a cancel-induced exception ends `TIMED_OUT`/`ABORTED`, not `FAILURE`), `touch_heartbeat` not advancing `changed_on`, and the full cancel seam (capture / notify / cancel success / decline / exception). - Manual (needs `DISTRIBUTED_COORDINATION_CONFIG` + a Celery worker + `celery beat` running `prune_tasks` on a short interval): submit an async chart query, `kill -9` the worker mid-query → within `GTF_ORPHAN_TASK_TIMEOUT` the task flips to `FAILURE` and its Celery job is revoked. Then, on a supported engine, abort a long-running query from the UI (or set `GLOBAL_ASYNC_QUERIES_QUERY_TIMEOUT`) and confirm the warehouse query is gone (e.g. `pg_stat_activity`) and the task ends `ABORTED`/`TIMED_OUT`. ### ADDITIONAL INFORMATION - [ ] Has associated issue: - [x] Required feature flags: `GLOBAL_TASK_FRAMEWORK` (async chart data also needs `GLOBAL_ASYNC_QUERIES`) - [ ] Changes UI - [x] Includes DB Migration (follow approval process in [SIP-59](https://github.com/apache/superset/issues/13351)) - [x] Migration is atomic, supports rollback & is backwards-compatible - [ ] Confirm DB migration upgrade and downgrade tested - [ ] Runtime estimates and downtime expectations provided - [x] Introduces new feature or API - [ ] Removes existing feature or API <sub>Note: the `tasks.last_heartbeat` column is folded into the branch's existing task-schema migration `7e2c9a4f1b83` (nullable add + index, reversible) rather than a new revision, matching how `guest_key` was folded — so there is no new migration file.</sub> -- 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]
