AIP-97 POC: classify a missed heartbeat as infra (executor-agnostic)

[1fanwang]

Why

AIP-97's answer to "does this only help Kubernetes" is that there are two tiers: a universal
floor every executor gets for free, and an opt-in ceiling for executors that can read a precise
cause (Kubernetes in apache#66405, Celery in #15, ECS in #28).

This is the floor, refreshed onto the current no-DB design. When a worker is killed outright it
reports nothing, so no executor can name the cause — but the scheduler already notices the missing
heartbeat and fails the task. That purge is the one classification point that needs no pod, no
backend API, and no executor code at all.

What changed

The heartbeat-timeout purge passes failure_kind=INFRA and reason="HeartbeatTimeout" to
handle_failure. Nothing else moves.

This is the same idea as #14, re-cut on the current head:
that POC predates the no-DB reconciliation and still passes the dropped infra_reason= argument.

Refs

• Parent: AIP-97: Disruption Readiness (Transparent Infra Failure Retries through Context Capture and Propagation) apache/airflow#66405
• Supersedes: AIP-97 POC: classify a missed heartbeat as infra (executor-agnostic) #14
• Ceiling tier on the same contract: AIP-97 POC: classify a lost Celery worker as infra #15 (Celery), AIP-97 POC: classify ECS task failures as infra #28 (ECS)

Testing Done

#	Scenario	Result
1	Missed heartbeat, refund on, retries=0	up_for_retry, max_tries 0->1
2	Missed heartbeat, refund off	failed, max_tries 0->0

Live Postgres 16, real scheduler purge

Drives the real _find_task_instances_without_heartbeats() and
_purge_task_instances_without_heartbeats() pair against a real database with a real registered
listener. The executor is a MagicMock for the callback sink only, mirroring the existing upstream
purge tests; it takes no part in the classification, which is the point.

$ docker run -d --name aip97-hb-pg -e POSTGRES_PASSWORD=airflow -e POSTGRES_USER=airflow \
    -e POSTGRES_DB=airflow -p 55434:5432 postgres:16
$ export AIRFLOW_HOME=/tmp/aip97-hb-e2e
$ export AIRFLOW__DATABASE__SQL_ALCHEMY_CONN="postgresql+psycopg2://airflow:airflow@localhost:55434/airflow"
$ export AIRFLOW__CORE__UNIT_TEST_MODE=False
$ uv run --project airflow-core airflow db migrate
Database migration done!

$ uv run --project airflow-core pytest /tmp/aip97_hb_e2e.py -q -s
[scheduler] _find_task_instances_without_heartbeats -> ['op']
[scheduler] purged a task with no heartbeat, no executor involvement
[db]        refund=True state=up_for_retry max_tries 0->1 (expected 1)
[listener]  received=[(<TaskFailureKind.INFRA: 'infra'>, 'HeartbeatTimeout')]

[scheduler] _find_task_instances_without_heartbeats -> ['op']
[scheduler] purged a task with no heartbeat, no executor involvement
[db]        refund=False state=failed max_tries 0->0 (expected 0)
[listener]  received=[(<TaskFailureKind.INFRA: 'infra'>, 'HeartbeatTimeout')]

2 passed

Scenario 1 uses retries=0, so the refund is what rescues it at all. With the flag off the kind
still reaches the listener, which is the alerting case working with no behavior change.

Committed unit test

$ uv run --project airflow-core pytest airflow-core/tests/unit/jobs/test_scheduler_job.py \
    -k "purge_without_heartbeat or task_instances_without_heartbeats" -q
5 passed, 400 deselected

Evidence script (inline, so it can be re-run)

from __future__ import annotations

import os
import sys
from datetime import timedelta
from unittest.mock import MagicMock

import pytest

from airflow._shared.state import TaskFailureKind
from airflow.jobs.job import Job
from airflow.jobs.scheduler_job_runner import HEARTBEAT_TIMEOUT_REASON, SchedulerJobRunner
from airflow.listeners import hookimpl
from airflow.listeners.listener import get_listener_manager
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.utils import timezone
from airflow.utils.state import State, TaskInstanceState
from tests_common.test_utils.config import conf_vars

received: list[tuple] = []


class _Listener:
    @hookimpl
    def on_task_instance_failed(self, previous_state, task_instance, error, failure_kind, reason):
        received.append((failure_kind, reason))


@pytest.mark.db_test
@pytest.mark.parametrize(
    ("refund_enabled", "expect_state", "expect_max_tries"),
    [
        ("True", TaskInstanceState.UP_FOR_RETRY, 1),
        ("False", TaskInstanceState.FAILED, 0),
    ],
    ids=["heartbeat_infra_refunded", "heartbeat_infra_flag_off"],
)
def test_heartbeat_timeout_is_infra(
    session, dag_maker, refund_enabled, expect_state, expect_max_tries
):
    """A worker that stopped heartbeating is infra, through the real scheduler purge, no executor."""
    received.clear()
    lm = get_listener_manager()
    lm.clear()
    lm.add_listener(_Listener())

    with dag_maker(dag_id=f"aip97_hb_{refund_enabled}", session=session):
        EmptyOperator(task_id="op", retries=0)
    dr = dag_maker.create_dagrun()
    ti = dr.task_instances[0]
    ti.task = dag_maker.dag.get_task("op")

    scheduler_job = Job()
    session.add(scheduler_job)
    session.commit()

    ti.state = State.RUNNING
    ti.start_date = timezone.utcnow() - timedelta(minutes=10)
    ti.last_heartbeat_at = timezone.utcnow() - timedelta(minutes=6)
    ti.queued_by_job_id = scheduler_job.id
    session.merge(ti)
    session.commit()

    pre_max_tries = ti.max_tries
    # A stand-in only for the callback sink; the classification below comes from the scheduler,
    # not from any executor. That is the point of this path.
    mock_executor = MagicMock()
    runner = SchedulerJobRunner(scheduler_job, executors=[mock_executor])

    with conf_vars({("core", "infra_failure_refund_retries"): refund_enabled}):
        found = runner._find_task_instances_without_heartbeats(session=session)
        print(f"\n[scheduler] _find_task_instances_without_heartbeats -> {[t.task_id for t in found]}")
        runner._purge_task_instances_without_heartbeats(found, session=session)
    session.commit()
    session.refresh(ti)

    print(
        f"\n[scheduler] purged a task with no heartbeat, no executor involvement"
        f"\n[db]        refund={refund_enabled} state={ti.state} "
        f"max_tries {pre_max_tries}->{ti.max_tries} (expected {expect_max_tries})"
        f"\n[listener]  received={received}"
    )

    assert ti.state == expect_state
    assert ti.max_tries == expect_max_tries
    assert received
    assert received[0] == (TaskFailureKind.INFRA, HEARTBEAT_TIMEOUT_REASON)
    lm.clear()


if __name__ == "__main__":
    sys.exit(pytest.main([os.path.abspath(__file__), "-p", "no:cacheprovider", "-q", "-s"]))