Fix duplicate heartbeat-timeout task callbacks

[hkc-8010]

closes: #42553

Description

When the scheduler detects a running task instance that has stopped heartbeating, it currently builds and enqueues a TaskCallbackRequest, but the task instance can remain in RUNNING until the callback is processed asynchronously.

That leaves a window where the next scheduler heartbeat-timeout scan can find the same task instance attempt again and enqueue another callback for the same try.

This changes the heartbeat-timeout purge path to enqueue the existing callback request, then immediately run the established task-instance failure transition path so the task instance leaves RUNNING in the same scheduler pass.

Executor cleanup via executor.change_state(..., FAILED, remove_running=True) is unchanged.

Related: #66767 handles callback type metadata for heartbeat-timed-out retries. This PR is scoped to preventing duplicate same-try callback emission.

Tests

• uv run --frozen --no-sync pytest --with-db-init airflow-core/tests/unit/jobs/test_scheduler_job.py::TestSchedulerJob::test_heartbeat_timeout_converges_ti_state_before_next_scan -q
  • Failed before the fix as expected.
• uv run --frozen --no-sync pytest airflow-core/tests/unit/jobs/test_scheduler_job.py::TestSchedulerJob::test_heartbeat_timeout_converges_ti_state_before_next_scan -q
• uv run --frozen --no-sync pytest airflow-core/tests/unit/jobs/test_scheduler_job.py -k "heartbeat_timeout" -q
• uv run --frozen --no-sync pytest airflow-core/tests/unit/jobs/test_scheduler_job.py -q
• prek run --files airflow-core/src/airflow/jobs/scheduler_job_runner.py airflow-core/tests/unit/jobs/test_scheduler_job.py
• .venv/bin/python scripts/ci/prek/run_mypy_full_dist_local_venv_or_breeze_in_ci.py airflow-core
• breeze testing core-tests --backend postgres --python 3.10 --db-reset -- airflow-core/tests/unit/jobs/test_scheduler_job.py -k "heartbeat_timeout" -q
• breeze testing core-tests --backend sqlite --python 3.10 --force-lowest-dependencies --db-reset -- airflow-core/tests/unit/jobs/test_scheduler_job.py -k "heartbeat_timeout" -q

[potiuk]

@hkc-8010 A few things need addressing before review — see our Pull Request quality criteria.

• ❌ Pre-commit / static checks. See docs.

No rush.

Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you.

[hkc-8010]

Thanks @potiuk. I addressed the static/pre-commit failure in ca401fa6ba9.

The failed check was caused by generated/provider_dependencies.json.sha256sum being generated without a trailing newline, which made end-of-file-fixer modify it during CI. The follow-up commit writes that checksum file with a trailing newline and adds a focused Breeze regression test for the cache-miss path.

Validated locally:

• uv run --project dev/breeze pytest dev/breeze/tests/test_provider_dependencies.py -q
• prek run --files dev/breeze/src/airflow_breeze/utils/provider_dependencies.py dev/breeze/tests/test_provider_dependencies.py
• prek run end-of-file-fixer --files generated/provider_dependencies.json.sha256sum

Could you please take another look when you get a chance?

[eladkal]

@hkc-8010 there are conflicts to resolve

[hkc-8010]

@hkc-8010 there are conflicts to resolve

Thanks. I've resolved the conflicts.

[jscheffl]

Looks good to me. But as I am not an expert on callbacks - is it OK from "state model" if callback is executed async that the state of the task is already set to failed? Or will this be achange to previous executions (that callbacks assumed the tasks are still in "running" state and will this harm callback logic?

[hkc-8010]

Thanks @jscheffl, good question.

I think this is OK from the state-model side and does not need an extra code change. This makes the heartbeat-timeout path consistent with the existing scheduler path for externally killed tasks: the scheduler creates/sends the TaskCallbackRequest, then calls ti.handle_failure(...) to converge the metastore state.

The callback itself is still executed asynchronously from the callback request payload. That payload is built before handle_failure() and uses the simplified execution API TI datamodel, so it is not passing a live ORM object whose state mutates underneath the callback. If callback code explicitly queries the metastore, it will now see the task as failed, which seems like the correct state once the scheduler has determined the task heartbeat timed out.

The main issue this PR fixes is that leaving the DB row as running after queueing the callback lets the next scheduler heartbeat find the same TI again and enqueue another same-try failure callback. Moving through handle_failure() in the same purge pass closes that gap while keeping the callback asynchronous.

[Vamsi-klu]

This closes the duplicate-scan window, but it now makes the heartbeat-timeout request the durable callback path for retryable tasks. Since the request is built before handle_failure() and does not set task_callback_type, a timed-out task with retries left can dispatch on_failure_callback instead of on_retry_callback. The regression test only covers max_tries = 0. Please set task_callback_type from retry eligibility before send_callback() and add a retryable heartbeat-timeout case.

[hkc-8010]

@Vamsi-klu Good catch, fixed. _purge_task_instances_without_heartbeats now sets task_callback_type the same way _process_executor_events does for the external-kill path: UP_FOR_RETRY if ti.is_eligible_to_retry() else FAILED, computed before handle_failure(). Added test_heartbeat_timeout_sets_callback_type_param, parametrised over a retry-eligible case (max_tries=1) and a non-eligible case (max_tries=0), asserting both the callback request's task_callback_type and the TI's final state.

[ephraimbuddy]

Can you add this test I suggested?

[hkc-8010]

@ephraimbuddy Addressed all three comments in 8f9eb4b — loads ti.task before handle_failure(), sends the email directly from this path, and unifies the retry-type decision with is_eligible_to_retry(). Added tests for each. Would you mind taking another look?

[hkc-8010]

@ephraimbuddy Fixed the race in 3c9eef3: the scan now locks the rows (with_row_locks, skip_locked) and the purge revalidates each TI's committed state before failing it, so a concurrently-completed task isn't clobbered and gets no failure callback. Tests added for both the race and the lock. Mind taking another look?

[Copilot]

Pull request overview

This PR updates the scheduler’s heartbeat-timeout cleanup so that, after enqueuing the callback request for a heartbeat-timed-out task instance, it immediately runs the normal task-instance failure transition to move the task instance out of RUNNING in the same scheduler pass—preventing duplicate same-try callback emission on subsequent scans.

Changes:

• Add row-locking to the heartbeat-timeout scan and re-check state in the purge loop to reduce race windows.
• In the heartbeat-timeout purge path, load the serialized task (when available), enqueue the callback request (and email request when configured), then immediately call ti.handle_failure(...) to persist the non-RUNNING state.
• Add/extend regression tests around deduplication, callback typing, email behavior, fail_fast, concurrent completion, and row-locking semantics.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
airflow-core/src/airflow/jobs/scheduler_job_runner.py	Updates zombie/heartbeat-timeout detection + purge to enqueue callbacks and immediately transition TI state out of RUNNING, with added locking and context loading.
airflow-core/tests/unit/jobs/test_scheduler_job.py	Adds regression coverage to ensure heartbeat-timeout cleanup converges TI state before the next scan and preserves expected callback/email/fail_fast behaviors.

[ephraimbuddy]

Nit.
Also fix the copilot review