This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch bound-heartbeat-purge-batch in repository https://gitbox.apache.org/repos/asf/airflow.git
commit d3f0e5e3db93b3ece31ee2d70e7d6afaee2c7ac5 Author: Rahul Vats <[email protected]> AuthorDate: Sun Oct 18 21:44:26 2026 +0530 Bound the heartbeat-timeout purge to max_tis_per_query per loop The scheduler finds every task instance past its heartbeat timeout and fails them in a single transaction, doing per-TI work (state revalidation, serialized-DAG load, failure callbacks) and holding a FOR UPDATE lock on each row. On a mass heartbeat-loss event -- a node dying with hundreds or thousands of RUNNING task instances -- that unbounded set turns the scheduler loop into one long, widely-locking transaction that blocks writers. Order the scan by oldest heartbeat and cap it at max_tis_per_query (the existing per-loop sizing knob) so each loop clears a bounded batch and the rest are handled on later loops, matching the batched-cleanup pattern used elsewhere in the scheduler. --- .../src/airflow/jobs/scheduler_job_runner.py | 7 +++++++ airflow-core/tests/unit/jobs/test_scheduler_job.py | 24 ++++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 19662ae5a07..49e8db2d074 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -3994,6 +3994,13 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin): session.execute(with_row_locks(run_query, of=DagRun, session=session, skip_locked=True)) ) query = query.where(or_(~regional_expression(TI), tuple_(TI.dag_id, TI.run_id).in_(locked_runs))) + # Bound the batch so a mass-heartbeat-loss event (e.g. a node dying with many RUNNING TIs) + # does not purge every stale row in one transaction: the purge below does per-TI work and + # holds FOR UPDATE on each locked row, so an unbounded set stalls the scheduler loop and + # blocks writers. Oldest heartbeat first, and the remainder are picked up on later loops. + query = query.order_by(TI.last_heartbeat_at) + if self.job.max_tis_per_query: + query = query.limit(self.job.max_tis_per_query) # Lock the rows (FOR UPDATE, of=TI so the FOR UPDATE isn't applied to the joined dag_model) # so a worker can't commit a terminal state on the same TI between this scan and the # handle_failure() in the purge that follows in the same transaction. skip_locked keeps HA diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 576cb59a898..c87e1baf0f8 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -918,6 +918,30 @@ class TestSchedulerJob: assert [row.id for row in runner._find_task_instances_without_heartbeats(session=session)] == [ti.id] + def test_find_task_instances_without_heartbeats_is_bounded_by_max_tis_per_query(self, dag_maker, session): + """A mass heartbeat-loss event must be purged in bounded batches, oldest heartbeat first.""" + with dag_maker(serialized=True): + EmptyOperator(task_id="task") + runner = SchedulerJobRunner(Job(), executors=[MockExecutor(do_update=False)]) + session.add(runner.job) + session.flush() + runner.job.max_tis_per_query = 2 + + ids_oldest_first = [] + for i in range(3): + dr = dag_maker.create_dagrun(run_id=f"run_{i}", logical_date=DEFAULT_DATE + timedelta(days=i)) + ti = dr.get_task_instance("task", session=session) + ti.state = State.RUNNING + ti.region_index = -1 + ti.last_heartbeat_at = DEFAULT_DATE + timedelta(minutes=i) + ti.queued_by_job_id = runner.job.id + ids_oldest_first.append(ti.id) + session.commit() + + found = runner._find_task_instances_without_heartbeats(session=session) + # Capped at max_tis_per_query, returning the two oldest-heartbeat rows; the third waits for a later loop. + assert [row.id for row in found] == ids_oldest_first[:2] + @pytest.mark.parametrize("event_state", [State.QUEUED, State.RUNNING]) def test_restarting_waits_for_terminal_executor_event(self, dag_maker, session, event_state): with dag_maker(dag_id="restart_waits_for_termination"):
