This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 8699a93e659 Bound job loading for component health checks (#76751)
8699a93e659 is described below
commit 8699a93e659038ae7726446a87d4857d75aa534f
Author: Vlad Vladov <[email protected]>
AuthorDate: Tue Nov 10 18:43:11 2026 +0200
Bound job loading for component health checks (#76751)
Uncooperative shutdowns leave unfinished jobs indefinitely. Hydrating
every abandoned row makes health probes increasingly expensive.
MySQL cannot semijoin an IN over a UNION ALL, so it ran the combined
health query as a dependent subquery per job row and the cost still grew
with job history. Two plain indexed queries stay flat.
The alive rule now lives in one SQL helper next to Job._is_alive, so the
health query and the jobs API filter cannot drift apart.
Amp-Thread-ID:
https://ampcode.com/threads/T-01a1c4c7-1193-77df-941c-935af7861fbe
---
.../src/airflow/api/common/airflow_health.py | 34 ++++++---
.../airflow/api_fastapi/common/parameters/misc.py | 30 ++------
airflow-core/src/airflow/jobs/job.py | 26 ++++++-
.../tests/unit/api/common/test_airflow_health.py | 80 +++++++++++++++-------
airflow-core/tests/unit/jobs/test_base_job.py | 29 ++++++++
5 files changed, 141 insertions(+), 58 deletions(-)
diff --git a/airflow-core/src/airflow/api/common/airflow_health.py
b/airflow-core/src/airflow/api/common/airflow_health.py
index 728fba5a2a0..7b05a03f959 100644
--- a/airflow-core/src/airflow/api/common/airflow_health.py
+++ b/airflow-core/src/airflow/api/common/airflow_health.py
@@ -17,11 +17,13 @@
from __future__ import annotations
import logging
+from datetime import datetime
from enum import Enum
from typing import TYPE_CHECKING, Any
from sqlalchemy import select
+from airflow._shared.timezones import timezone
from airflow.configuration import conf
from airflow.dag_processing.bundles.manager import
get_configured_bundle_team_names
from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner
@@ -53,16 +55,28 @@ class DetailedHealthStatus(str, Enum):
@provide_session
def get_jobs_health(job_runner_class, *, session: Session = NEW_SESSION) ->
list[Job]:
- """Return unfinished jobs for the runner class, ordered by latest
heartbeat."""
- return list(
- session.scalars(
- select(Job)
- .where(
- Job.job_type == job_runner_class.job_type,
- Job.end_date.is_(None),
- )
- .order_by(Job.latest_heartbeat.desc())
- )
+ """
+ Return live unfinished jobs and the latest unfinished job, ordered by
heartbeat.
+
+ Keep the latest job even when dead so health responses retain its
heartbeat.
+ """
+ job_type = job_runner_class.job_type
+ unfinished = Job.end_date.is_(None)
+ # Two separate queries rather than one ``id IN (... UNION ALL ...)``:
MySQL cannot semijoin that
+ # form and runs it as a dependent subquery per ``job`` row, which scales
with job history.
+ live_jobs = session.scalars(
+ select(Job).where(Job.job_type == job_type, unfinished,
Job.is_alive_clause([job_type]))
+ )
+ latest_job = session.scalar(
+ select(Job).where(Job.job_type == job_type,
unfinished).order_by(Job.latest_heartbeat.desc()).limit(1)
+ )
+ jobs = {job.id: job for job in live_jobs}
+ if latest_job is not None:
+ jobs.setdefault(latest_job.id, latest_job)
+ return sorted(
+ jobs.values(),
+ key=lambda job: job.latest_heartbeat or
datetime.min.replace(tzinfo=timezone.utc),
+ reverse=True,
)
diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters/misc.py
b/airflow-core/src/airflow/api_fastapi/common/parameters/misc.py
index e25cd99cbf3..b6130a77f47 100644
--- a/airflow-core/src/airflow/api_fastapi/common/parameters/misc.py
+++ b/airflow-core/src/airflow/api_fastapi/common/parameters/misc.py
@@ -88,30 +88,14 @@ class _IsAliveFilter(BaseParam[bool]):
if self.value is None and self.skip_none:
return select
- from datetime import timedelta
-
- from sqlalchemy import and_, case, or_
-
- from airflow._shared.timezones import timezone
- from airflow.jobs.job import Job, JobState, health_check_threshold
-
- now = timezone.utcnow()
- # Only the concrete core job types are persisted here and their
thresholds come from config
- # rather than the per-row heartrate, so the cutoff is a bound
timestamp and the predicate stays
- # portable. Kept in sync with Job._is_alive: alive means RUNNING with
a heartbeat newer than the
- # job type's health-check threshold. The case wraps the result so it
is always True/False (never
- # NULL), keeping the negated is_alive=False query exact.
- fresh_per_type = or_(
- *(
- and_(
- Job.job_type == job_type,
- Job.latest_heartbeat > now -
timedelta(seconds=health_check_threshold(job_type, 0)),
- )
- for job_type in ("SchedulerJob", "TriggererJob",
"DagProcessorJob")
- )
- )
+ from sqlalchemy import case
+
+ from airflow.jobs.job import Job
+
+ # The case wraps the result so it is always True/False (never NULL),
keeping the negated
+ # is_alive=False query exact.
alive = case(
- (and_(Job.state == JobState.RUNNING,
Job.latest_heartbeat.is_not(None), fresh_per_type), True),
+ (Job.is_alive_clause(("SchedulerJob", "TriggererJob",
"DagProcessorJob")), True),
else_=False,
)
return select.where(alive if self.value else ~alive)
diff --git a/airflow-core/src/airflow/jobs/job.py
b/airflow-core/src/airflow/jobs/job.py
index f08b2fd1cb8..43862eb5a2b 100644
--- a/airflow-core/src/airflow/jobs/job.py
+++ b/airflow-core/src/airflow/jobs/job.py
@@ -18,14 +18,14 @@
from __future__ import annotations
from collections.abc import Callable, Iterable
-from datetime import datetime
+from datetime import datetime, timedelta
from enum import Enum
from functools import cached_property, lru_cache
from time import sleep
from typing import TYPE_CHECKING, NoReturn
from uuid import UUID
-from sqlalchemy import Index, Integer, String, Uuid, case, select
+from sqlalchemy import Index, Integer, String, Uuid, and_, case, or_, select
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import Mapped, backref, foreign, mapped_column,
relationship
from sqlalchemy.orm.session import make_transient
@@ -59,6 +59,7 @@ class JobState(str, Enum):
if TYPE_CHECKING:
from sqlalchemy.orm.session import Session
+ from sqlalchemy.sql.elements import ColumnElement
def _resolve_dagrun_model():
@@ -334,6 +335,27 @@ class Job(Base, LoggingMixin):
and (timezone.utcnow() - latest_heartbeat).total_seconds() <
health_check_threshold_value
)
+ @staticmethod
+ def is_alive_clause(job_types: Iterable[str]) -> ColumnElement[bool]:
+ """
+ Return the SQL form of :meth:`_is_alive` for the given job types.
+
+ Only the core job types are supported: their thresholds come from
config rather than the per-row
+ heartrate, so the cutoff is a bound timestamp and the predicate stays
portable. Keep in sync
+ with :meth:`_is_alive`.
+ """
+ now = timezone.utcnow()
+ fresh_per_type = or_(
+ *(
+ and_(
+ Job.job_type == job_type,
+ Job.latest_heartbeat > now -
timedelta(seconds=health_check_threshold(job_type, 0)),
+ )
+ for job_type in job_types
+ )
+ )
+ return and_(Job.state == JobState.RUNNING,
Job.latest_heartbeat.is_not(None), fresh_per_type)
+
@provide_session
def most_recent_job(job_type: str, *, session: Session = NEW_SESSION) -> Job |
None:
diff --git a/airflow-core/tests/unit/api/common/test_airflow_health.py
b/airflow-core/tests/unit/api/common/test_airflow_health.py
index 9affe9faab6..649ad465c5c 100644
--- a/airflow-core/tests/unit/api/common/test_airflow_health.py
+++ b/airflow-core/tests/unit/api/common/test_airflow_health.py
@@ -20,6 +20,7 @@ from datetime import datetime, timedelta
from unittest.mock import MagicMock, patch
import pytest
+import time_machine
from airflow._shared.timezones import timezone
from airflow.api.common.airflow_health import (
@@ -31,7 +32,8 @@ from airflow.api.common.airflow_health import (
get_airflow_health,
get_jobs_health,
)
-from airflow.jobs.job import Job, JobState
+from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner
+from airflow.jobs.job import Job, JobState, health_check_threshold
from airflow.jobs.scheduler_job_runner import SchedulerJobRunner
from airflow.jobs.triggerer_job_runner import TriggererJobRunner
from airflow.utils.session import provide_session
@@ -83,13 +85,14 @@ def _create_job(
runner_class,
*,
hostname: str,
- heartbeat: datetime,
+ heartbeat: datetime | None,
state: JobState = JobState.RUNNING,
end_date: datetime | None = None,
team_names: list[str] | None = None,
bundle_names: list[str] | None = None,
) -> Job:
job = Job(
+ job_type=runner_class.job_type,
state=state,
latest_heartbeat=heartbeat,
hostname=hostname,
@@ -97,10 +100,6 @@ def _create_job(
team_names=team_names or [],
bundle_names=bundle_names,
)
- if runner_class is TriggererJobRunner:
- runner_class(job=job, capacity=1)
- else:
- runner_class(job=job)
session.add(job)
return job
@@ -482,39 +481,71 @@ class TestAirflowHealthFromDb:
@pytest.fixture(autouse=True)
def cleanup_jobs(self):
clear_db_jobs()
+ health_check_threshold.cache_clear()
yield
clear_db_jobs()
+ health_check_threshold.cache_clear()
+ @pytest.mark.parametrize(
+ ("runner_class", "threshold"),
+ [(SchedulerJobRunner, 17), (TriggererJobRunner, 41.5),
(DagProcessorJobRunner, 73)],
+ )
+ @pytest.mark.parametrize("latest_failed", [False, True])
+ @time_machine.travel("2026-11-09T09:00:00Z", tick=False)
+ @conf_vars(
+ {
+ ("scheduler", "scheduler_health_check_threshold"): "17",
+ ("triggerer", "triggerer_health_check_threshold"): "41.5",
+ ("dag_processor", "health_check_threshold"): "73",
+ }
+ )
@provide_session
- def
test_get_jobs_health_returns_unfinished_jobs_ordered_by_heartbeat(self, *,
session):
- older = timezone.utcnow() - timedelta(minutes=5)
- newer = timezone.utcnow()
- older_job = _create_job(session, SchedulerJobRunner, hostname="older",
heartbeat=older)
- newer_job = _create_job(session, SchedulerJobRunner, hostname="newer",
heartbeat=newer)
+ def test_get_jobs_health_returns_live_jobs_and_latest_unfinished(
+ self, runner_class, threshold, latest_failed, *, session
+ ):
+ now = timezone.utcnow()
+ cutoff = now - timedelta(seconds=threshold)
+ _create_job(session, runner_class, hostname="live-newer",
heartbeat=now - timedelta(seconds=1))
+ _create_job(
+ session, runner_class, hostname="live-older", heartbeat=cutoff +
timedelta(microseconds=1)
+ )
+ _create_job(session, runner_class, hostname="at-threshold",
heartbeat=cutoff)
+ _create_job(session, runner_class, hostname="stale", heartbeat=cutoff
- timedelta(microseconds=1))
_create_job(
session,
- SchedulerJobRunner,
+ runner_class,
hostname="ended",
- heartbeat=newer,
- end_date=newer,
+ heartbeat=now + timedelta(seconds=1),
+ end_date=now,
)
- unfinished_failed = _create_job(
+ _create_job(
session,
- SchedulerJobRunner,
+ runner_class,
hostname="failed-unfinished",
- heartbeat=older - timedelta(minutes=1),
+ heartbeat=now if latest_failed else now - timedelta(seconds=2),
state=JobState.FAILED,
)
- _create_job(session, TriggererJobRunner, hostname="triggerer",
heartbeat=newer)
+ other_runner = SchedulerJobRunner if runner_class is
TriggererJobRunner else TriggererJobRunner
+ _create_job(session, other_runner, hostname="other-component",
heartbeat=now)
+ session.commit()
+
+ jobs = get_jobs_health(runner_class, session=session)
+
+ expected = ["live-newer", "live-older"]
+ if latest_failed:
+ expected.insert(0, "failed-unfinished")
+ assert [job.hostname for job in jobs] == expected
+
+ @provide_session
+ def test_get_jobs_health_keeps_one_job_with_null_heartbeat(self, *,
session):
+ _create_job(session, SchedulerJobRunner, hostname="no-heartbeat-1",
heartbeat=None)
+ _create_job(session, SchedulerJobRunner, hostname="no-heartbeat-2",
heartbeat=None)
session.commit()
jobs = get_jobs_health(SchedulerJobRunner, session=session)
- assert [job.hostname for job in jobs] == [
- newer_job.hostname,
- older_job.hostname,
- unfinished_failed.hostname,
- ]
+ assert len(jobs) == 1
+ assert jobs[0].latest_heartbeat is None
def test_get_airflow_health_no_jobs(self):
health_status = get_airflow_health()
@@ -590,6 +621,9 @@ class TestAirflowHealthFromDb:
_create_job(session, SchedulerJobRunner, hostname="stale-3",
heartbeat=third)
session.commit()
+ jobs = get_jobs_health(SchedulerJobRunner, session=session)
+ assert [job.hostname for job in jobs] == ["stale-1"]
+
health_status = get_airflow_health()
assert health_status["scheduler"]["status"] == HealthStatus.UNHEALTHY
diff --git a/airflow-core/tests/unit/jobs/test_base_job.py
b/airflow-core/tests/unit/jobs/test_base_job.py
index eda85fa270b..8443f30985d 100644
--- a/airflow-core/tests/unit/jobs/test_base_job.py
+++ b/airflow-core/tests/unit/jobs/test_base_job.py
@@ -33,6 +33,7 @@ from airflow.utils.session import create_session
from airflow.utils.state import State
from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.db import clear_db_jobs
from unit.listeners import lifecycle_listener
from unit.utils.test_helpers import MockJobRunner, SchedulerJobRunner,
TriggererJobRunner
@@ -214,6 +215,34 @@ class TestJob:
job.latest_heartbeat = timezone.utcnow() -
datetime.timedelta(seconds=10)
assert job.is_alive() is False, "Completed jobs even with recent
heartbeat should not be alive"
+ def test_is_alive_clause_matches_is_alive(self, session):
+ clear_db_jobs()
+ core_types = ("SchedulerJob", "TriggererJob", "DagProcessorJob")
+ now = timezone.utcnow()
+ for job_type in core_types:
+ threshold = health_check_threshold(job_type, 0)
+ for state in (State.RUNNING, State.SUCCESS, State.FAILED,
State.RESTARTING):
+ for age in (None, 0, threshold - 1, threshold + 1, 86400):
+ session.add(
+ Job(
+ job_type=job_type,
+ state=state,
+ hostname=f"{job_type}|{state}|{age}",
+ latest_heartbeat=None if age is None else now -
datetime.timedelta(seconds=age),
+ )
+ )
+ session.commit()
+
+ jobs = session.scalars(select(Job)).all()
+ alive_in_sql = {
+ job.hostname for job in
session.scalars(select(Job).where(Job.is_alive_clause(core_types)))
+ }
+
+ assert alive_in_sql == {job.hostname for job in jobs if job.is_alive()}
+ assert alive_in_sql, "the scenarios must include live jobs"
+ assert len(alive_in_sql) < len(jobs)
+ clear_db_jobs()
+
@patch("airflow.jobs.job.create_session")
def test_heartbeat_failed(self, mock_create_session, caplog):
when = timezone.utcnow() - datetime.timedelta(seconds=60)
