This is an automated email from the ASF dual-hosted git repository.
ephraimbuddy 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 d4a4d853e9d Fix duplicate heartbeat-timeout task callbacks (#68008)
d4a4d853e9d is described below
commit d4a4d853e9d6f32b212d85d514666bd5848927cb
Author: Hemkumar Chheda <[email protected]>
AuthorDate: Wed Jul 22 12:04:46 2026 +0530
Fix duplicate heartbeat-timeout task callbacks (#68008)
* Fix duplicate heartbeat-timeout task callbacks
* Set task_callback_type on heartbeat-timeout callback based on retry
eligibility
Address review feedback: the heartbeat-timeout TaskCallbackRequest didn't
set
task_callback_type, so a timed-out task with retries left could dispatch
on_failure_callback instead of on_retry_callback. Mirrors the existing
externally-killed-task pattern.
* Load task and send email in heartbeat-timeout purge path
The purge path called handle_failure() without ever loading ti.task,
so fail_fast (ti.task.dag.fail_fast) silently no-opped and the
external-kill email path in process_executor_events never ran for
these TIs since the TI already left RUNNING/RESTARTING by the time
any executor event for it was processed. It also computed
task_callback_type with a max_tries > 0 guard that could disagree
with what handle_failure() actually persists for a RESTARTING TI
with max_tries=0.
Load ti.task before calling handle_failure(), decide the callback
type with plain is_eligible_to_retry() (matching
fetch_handle_failure_context exactly), and send an EmailRequest
alongside the existing TaskCallbackRequest.
closes: #42553
* Lock and revalidate TI state in heartbeat-timeout purge
The heartbeat scan selected RUNNING/RESTARTING TIs unlocked, and the same
session flowed into the purge. Between the scan and ti.handle_failure(), a
worker could commit the task as SUCCESS; handle_failure() refreshes the TI
but sets FAILED unconditionally, clobbering that terminal state after having
already enqueued a failure callback.
Lock the scan query with with_row_locks(of=TI, skip_locked=True) so a worker
can't change the row between the scan and the handle_failure() that follows
in
the same transaction, matching how process_executor_events and the other
TI-mutating paths in this file lock. As defense in depth, revalidate each
row's
committed state at the top of the purge loop and skip it (no callback, no
email,
no handle_failure) if it is no longer RUNNING/RESTARTING.
closes: #42553
* Address review nits on heartbeat-timeout purge
Use a real MockExecutor(do_update=False) instead of a bare MagicMock() in
test_heartbeat_timeout_converges_ti_state_before_next_scan, matching the
other
heartbeat-timeout tests, and assert on callback_sink.send.
Reword the _resolve_ti_callback_bundle_info docstring: it is used by the
heartbeat-timeout purge path; process_executor_events inlines the same
resolution rather than sharing this helper.
---
.../src/airflow/jobs/scheduler_job_runner.py | 158 ++++++++----
airflow-core/tests/unit/jobs/test_scheduler_job.py | 271 ++++++++++++++++++++-
2 files changed, 384 insertions(+), 45 deletions(-)
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 9b9e85c07e3..cb4c8f64552 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -187,6 +187,26 @@ def _eager_load_dag_run_for_validation() ->
tuple[LoaderOption, LoaderOption]:
)
+def _resolve_ti_callback_bundle_info(ti: TaskInstance) -> tuple[str, str |
None, Any]:
+ """
+ Resolve the bundle name/version/version-data needed to build a
TaskCallbackRequest or EmailRequest.
+
+ Used by the heartbeat-timeout purge path. Encapsulates the bundle-pinning
semantics: fall back
+ to ``dag_model`` for legacy tasks with no ``dag_version`` (pre-AIP-66
migrations), and leave the
+ bundle version unpinned when the dag run itself wasn't pinned
(``disable_bundle_versioning``),
+ so the callback runs against the same code as the task did.
``process_executor_events`` inlines
+ the same resolution for its externally-killed-task path.
+ """
+ bundle_name = ti.dag_version.bundle_name if ti.dag_version else
ti.dag_model.bundle_name
+ bundle_version = (
+ ti.dag_version.bundle_version
+ if ti.dag_version and ti.dag_run.bundle_version is not None
+ else ti.dag_run.bundle_version
+ )
+ version_data = _resolve_version_data(ti.dag_version,
ti.dag_run.bundle_version)
+ return bundle_name, bundle_version, version_data
+
+
def _ensure_ti_has_dag_version_id(ti: TaskInstance, session: Session, log:
Logger) -> bool:
"""
Ensure a TaskInstance has a valid dag_version_id for Pydantic
serialisation.
@@ -3548,22 +3568,27 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
self.log.debug("Finding 'running' jobs without a recent heartbeat")
limit_dttm = timezone.utcnow() -
timedelta(seconds=self._task_instance_heartbeat_timeout_secs)
asset_loader, alias_loader = _eager_load_dag_run_for_validation()
- task_instances_without_heartbeats = list(
- session.scalars(
- select(TI)
- .options(selectinload(TI.dag_model))
- .options(asset_loader)
- .options(alias_loader)
- .options(selectinload(TI.dag_version))
- .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
- .join(DM, TI.dag_id == DM.dag_id)
- .where(
- TI.state.in_((TaskInstanceState.RUNNING,
TaskInstanceState.RESTARTING)),
- TI.last_heartbeat_at < limit_dttm,
- )
- .where(TI.queued_by_job_id == self.job.id)
+ query = (
+ select(TI)
+ .options(selectinload(TI.dag_model))
+ .options(asset_loader)
+ .options(alias_loader)
+ .options(selectinload(TI.dag_version))
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(DM, TI.dag_id == DM.dag_id)
+ .where(
+ TI.state.in_((TaskInstanceState.RUNNING,
TaskInstanceState.RESTARTING)),
+ TI.last_heartbeat_at < limit_dttm,
)
+ .where(TI.queued_by_job_id == self.job.id)
)
+ # 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
+ # schedulers from blocking on each other.
_purge_task_instances_without_heartbeats still
+ # revalidates each row's state before acting, as defense in depth.
+ query = with_row_locks(query, of=TI, session=session, skip_locked=True)
+ task_instances_without_heartbeats = list(session.scalars(query))
if task_instances_without_heartbeats:
self.log.warning(
"Failing %s TIs without heartbeat after %s",
@@ -3582,46 +3607,72 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
dag_id_to_team_name = {}
for ti in task_instances_without_heartbeats:
+ # The scan locked this row (FOR UPDATE / skip_locked), but
revalidate against the
+ # committed state before emitting any side effect: a worker can
commit a terminal state
+ # (e.g. SUCCESS) around the same time the scan runs. Failing the
TI here would clobber
+ # that terminal state and emit a spurious failure callback.
Mirrors the lock-then-recheck
+ # guard in process_executor_events.
+ ti.refresh_from_db(session=session)
+ if ti.state not in (TaskInstanceState.RUNNING,
TaskInstanceState.RESTARTING):
+ self.log.info(
+ "Task instance %s is no longer running (state=%s);
skipping heartbeat-timeout purge",
+ ti,
+ ti.state,
+ )
+ continue
+
task_instance_heartbeat_timeout_message_details = (
self._generate_task_instance_heartbeat_timeout_message_details(ti)
)
- # Safely extract bundle info with fallback for legacy tasks
- # (dag_version may be None after Airflow 2 → 3 migration).
- _hb_bundle_name = ti.dag_version.bundle_name if ti.dag_version
else ti.dag_model.bundle_name
- # Mirror dag_run pinning: if the run wasn't pinned (e.g.
dag.disable_bundle_versioning=True),
- # leave the callback unpinned so it runs against the same code as
the task.
- _hb_bundle_version = (
- ti.dag_version.bundle_version
- if ti.dag_version and ti.dag_run.bundle_version is not None
- else ti.dag_run.bundle_version
+ msg = str(task_instance_heartbeat_timeout_message_details)
+
+ # Load the serialized task, mirroring how process_executor_events'
external-kill path
+ # loads it, so handle_failure() below can see fail_fast
(ti.task.dag.fail_fast) instead
+ # of silently skipping it, and so email/callback gating below can
check the real task
+ # definition. Unlike that path, there's no executor-reported state
to fall back to here,
+ # so a load failure still falls through to fail the TI below, just
without task context.
+ try:
+ dag =
self.scheduler_dag_bag.get_dag_for_run(dag_run=ti.dag_run, session=session)
+ if not dag:
+ raise DagNotFound(f"DAG '{ti.dag_id}' not found in
serialized_dag table")
+ task = dag.get_task(ti.task_id)
+ except Exception:
+ self.log.exception(
+ "Could not load task for heartbeat-timed-out task instance
%s; "
+ "continuing without fail_fast/email context",
+ ti,
+ )
+ task = None
+ ti.task = task
+
+ # Single source of truth for the retry decision, matching
+ # TaskInstance.fetch_handle_failure_context exactly, so the
callback type sent here can
+ # never disagree with the state handle_failure() actually persists
below (this previously
+ # diverged for RESTARTING task instances with max_tries=0).
+ task_callback_type = (
+ TaskInstanceState.UP_FOR_RETRY if ti.is_eligible_to_retry()
else TaskInstanceState.FAILED
)
- _hb_version_data = _resolve_version_data(ti.dag_version,
ti.dag_run.bundle_version)
+
+ bundle_name, bundle_version, version_data =
_resolve_ti_callback_bundle_info(ti)
# Backfill dag_version_id for legacy tasks (Pydantic requires
uuid.UUID).
if not _ensure_ti_has_dag_version_id(ti, session, self.log):
continue
- # ti.task isn't loaded in this purge path, so
is_eligible_to_retry() uses its
- # no-task fallback (``try_number <= max_tries``), which skips the
retries-configured
- # check its task-loaded branch applies; guard with ``max_tries >
0`` so a task
- # declared with retries=0 isn't treated as retry-eligible here.
- if ti.max_tries > 0 and ti.is_eligible_to_retry():
- task_callback_type = TaskInstanceState.UP_FOR_RETRY
- else:
- task_callback_type = TaskInstanceState.FAILED
+ context_from_server = TIRunContext(
+ dag_run=DRDataModel.model_validate(ti.dag_run,
from_attributes=True),
+ max_tries=ti.max_tries,
+ variables=[],
+ connections=[],
+ xcom_keys_to_clear=[],
+ )
request = TaskCallbackRequest(
filepath=ti.dag_model.relative_fileloc or "",
- bundle_name=_hb_bundle_name,
- bundle_version=_hb_bundle_version,
- version_data=_hb_version_data,
+ bundle_name=bundle_name,
+ bundle_version=bundle_version,
+ version_data=version_data,
ti=ti,
- msg=str(task_instance_heartbeat_timeout_message_details),
+ msg=msg,
task_callback_type=task_callback_type,
- context_from_server=TIRunContext(
- dag_run=DRDataModel.model_validate(ti.dag_run,
from_attributes=True),
- max_tries=ti.max_tries,
- variables=[],
- connections=[],
- xcom_keys_to_clear=[],
- ),
+ context_from_server=context_from_server,
)
session.add(
Log(
@@ -3642,6 +3693,27 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
request,
)
self.executor.send_callback(request)
+
+ # This purge path leaves the executor's own "task finished but TI
still looked queued"
+ # handling in process_executor_events unreachable for this TI once
handle_failure() below
+ # moves it out of RUNNING, so the email notification has to be
sent from here directly.
+ if task is not None and task.email and (task.email_on_failure or
task.email_on_retry):
+ self.executor.send_callback(
+ EmailRequest(
+ filepath=ti.dag_model.relative_fileloc or "",
+ bundle_name=bundle_name,
+ bundle_version=bundle_version,
+ version_data=version_data,
+ ti=ti,
+ msg=msg,
+ email_type=(
+ "retry" if task_callback_type ==
TaskInstanceState.UP_FOR_RETRY else "failure"
+ ),
+ context_from_server=context_from_server,
+ )
+ )
+
+ ti.handle_failure(error=msg, session=session)
executor = self._try_to_load_executor(
ti, session, team_name=dag_id_to_team_name.get(ti.dag_id,
NOTSET)
)
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index dff3adea48b..782bcf81c4a 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -46,7 +46,12 @@ from airflow._shared.observability.metrics.base_stats_logger
import StatsLogger
from airflow._shared.timezones import timezone
from airflow.api_fastapi.auth.tokens import JWTGenerator
from airflow.assets.manager import AssetManager
-from airflow.callbacks.callback_requests import DagCallbackRequest,
DagRunContext, TaskCallbackRequest
+from airflow.callbacks.callback_requests import (
+ DagCallbackRequest,
+ DagRunContext,
+ EmailRequest,
+ TaskCallbackRequest,
+)
from airflow.callbacks.database_callback_sink import DatabaseCallbackSink
from airflow.dag_processing.collection import AssetModelOperation,
DagModelOperation
from airflow.dag_processing.dagbag import DagBag, sync_bag_to_db
@@ -8720,6 +8725,82 @@ class TestSchedulerJob:
assert callback_request.context_from_server.dag_run.logical_date ==
dag_run.logical_date
assert callback_request.context_from_server.max_tries == ti.max_tries
+ def test_heartbeat_timeout_converges_ti_state_before_next_scan(self,
dag_maker, session):
+ """A heartbeat-timed-out TI should not be found again on the next
scheduler scan."""
+ with dag_maker(dag_id="test_heartbeat_timeout_dedupe",
session=session):
+ EmptyOperator(task_id="test_task", on_failure_callback=lambda
context: None)
+
+ dag_run = dag_maker.create_dagrun(run_id="test_run",
state=DagRunState.RUNNING)
+
+ executor = MockExecutor(do_update=False)
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[executor])
+
+ ti = dag_run.get_task_instance(task_id="test_task")
+ ti.state = TaskInstanceState.RUNNING
+ ti.try_number = 1
+ ti.max_tries = 0
+ ti.queued_by_job_id = scheduler_job.id
+ ti.start_date = timezone.utcnow() - timedelta(seconds=900)
+ ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600)
+ session.merge(ti)
+ session.commit()
+
+ self.job_runner._find_and_purge_task_instances_without_heartbeats()
+
+ session.expire_all()
+ ti.refresh_from_db(session=session)
+ assert ti.state != TaskInstanceState.RUNNING
+ assert
self.job_runner._find_task_instances_without_heartbeats(session=session) == []
+
+ self.job_runner._find_and_purge_task_instances_without_heartbeats()
+
+ self.job_runner.executor.callback_sink.send.assert_called_once()
+
+ @pytest.mark.parametrize(
+ ("retries", "callback_kind", "expected"),
+ [
+ (1, "retry", TaskInstanceState.UP_FOR_RETRY),
+ (0, "failure", TaskInstanceState.FAILED),
+ ],
+ )
+ def test_heartbeat_timeout_sets_callback_type_param(
+ self, dag_maker, session, retries, callback_kind, expected
+ ):
+ """Heartbeat timeout should mark callback type based on retry
eligibility."""
+ with dag_maker(dag_id=f"heartbeat_timeout_{callback_kind}",
session=session):
+ if callback_kind == "retry":
+ EmptyOperator(task_id="t1", retries=retries,
on_retry_callback=lambda ctx: None)
+ else:
+ EmptyOperator(task_id="t1", retries=retries,
on_failure_callback=lambda ctx: None)
+
+ dag_run = dag_maker.create_dagrun(run_id="test_run",
state=DagRunState.RUNNING)
+
+ executor = MockExecutor(do_update=False)
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[executor])
+
+ ti = dag_run.get_task_instance(task_id="t1")
+ ti.state = TaskInstanceState.RUNNING
+ ti.try_number = 1
+ ti.max_tries = retries
+ ti.queued_by_job_id = scheduler_job.id
+ ti.start_date = timezone.utcnow() - timedelta(seconds=900)
+ ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600)
+ session.merge(ti)
+ session.commit()
+
+ self.job_runner._find_and_purge_task_instances_without_heartbeats()
+
+ self.job_runner.executor.callback_sink.send.assert_called_once()
+ request = self.job_runner.executor.callback_sink.send.call_args[0][0]
+ assert isinstance(request, TaskCallbackRequest)
+ assert request.task_callback_type == expected
+
+ session.expire_all()
+ ti.refresh_from_db(session=session)
+ assert ti.state == expected
+
@pytest.mark.parametrize(
("state", "retries", "try_number", "expected_callback_type",
"expected_dispatched_callback"),
[
@@ -8948,6 +9029,192 @@ class TestSchedulerJob:
assert isinstance(request, TaskCallbackRequest)
assert request.bundle_version is None
+ @time_machine.travel(DEFAULT_DATE, tick=False)
+ def test_heartbeat_timeout_preserves_failure_email(self, dag_maker,
session):
+ """
+ The purge path moves the TI out of RUNNING itself, so
process_executor_events'
+ external-kill email path never sees this TI. The purge path must send
the failure
+ email directly instead of silently dropping it.
+ """
+ with dag_maker(dag_id="heartbeat_timeout_email", session=session):
+ EmptyOperator(
+ task_id="t1",
+ email="[email protected]",
+ email_on_failure=True,
+ )
+
+ dag_run = dag_maker.create_dagrun(run_id="test_run",
state=DagRunState.RUNNING)
+
+ executor = MockExecutor(do_update=False)
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[executor])
+
+ ti = dag_run.get_task_instance(task_id="t1")
+ ti.state = TaskInstanceState.RUNNING
+ ti.try_number = 1
+ ti.max_tries = 0
+ ti.queued_by_job_id = scheduler_job.id
+ ti.start_date = timezone.utcnow() - timedelta(seconds=900)
+ ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600)
+ session.merge(ti)
+ session.commit()
+
+ self.job_runner._find_and_purge_task_instances_without_heartbeats()
+
+ sent_requests = [c.args[0] for c in
self.job_runner.executor.callback_sink.send.call_args_list]
+ email_requests = [r for r in sent_requests if isinstance(r,
EmailRequest)]
+ assert len(email_requests) == 1
+ assert email_requests[0].email_type == "failure"
+
+ def
test_heartbeat_timeout_restarting_zero_max_tries_matches_final_state(self,
dag_maker, session):
+ """
+ is_eligible_to_retry() always returns True for a RESTARTING TI,
independent of
+ max_tries. The task_callback_type sent to the Dag processor must match
the state
+ handle_failure() actually persists -- these previously diverged for a
RESTARTING TI
+ with max_tries=0, where the callback was typed FAILED but the TI still
ended up
+ UP_FOR_RETRY.
+ """
+ with dag_maker(dag_id="hb_timeout_restarting_zero_max_tries",
session=session):
+ EmptyOperator(task_id="t1", retries=0)
+
+ dag_run = dag_maker.create_dagrun(run_id="test_run",
state=DagRunState.RUNNING)
+
+ executor = MockExecutor(do_update=False)
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[executor])
+
+ ti = dag_run.get_task_instance(task_id="t1")
+ ti.state = TaskInstanceState.RESTARTING
+ ti.try_number = 1
+ ti.max_tries = 0
+ ti.queued_by_job_id = scheduler_job.id
+ ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600)
+ session.merge(ti)
+ session.commit()
+
+ self.job_runner._find_and_purge_task_instances_without_heartbeats()
+
+ self.job_runner.executor.callback_sink.send.assert_called_once()
+ request = self.job_runner.executor.callback_sink.send.call_args[0][0]
+ assert isinstance(request, TaskCallbackRequest)
+ assert request.task_callback_type == TaskInstanceState.UP_FOR_RETRY
+
+ session.expire_all()
+ ti.refresh_from_db(session=session)
+ assert ti.state == TaskInstanceState.UP_FOR_RETRY
+
+ def test_heartbeat_timeout_honors_fail_fast(self, dag_maker, session):
+ """
+ handle_failure() only stops sibling tasks when ti.task.dag.fail_fast
is True, which
+ requires ti.task to be loaded. Before the fix, this purge path never
loaded ti.task,
+ so fail_fast silently no-opped and a sibling task kept running instead
of being
+ stopped.
+ """
+ with dag_maker(dag_id="hb_timeout_fail_fast", fail_fast=True):
+ EmptyOperator(task_id="t1")
+ EmptyOperator(task_id="sibling")
+
+ dag_run = dag_maker.create_dagrun(run_id="test_run",
state=DagRunState.RUNNING)
+
+ executor = MockExecutor(do_update=False)
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[executor])
+
+ ti = dag_run.get_task_instance(task_id="t1")
+ ti.state = TaskInstanceState.RUNNING
+ ti.try_number = 1
+ ti.max_tries = 0
+ ti.queued_by_job_id = scheduler_job.id
+ ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600)
+
+ sibling = dag_run.get_task_instance(task_id="sibling")
+ sibling.state = TaskInstanceState.RUNNING
+
+ session.merge(ti)
+ session.merge(sibling)
+ session.commit()
+
+ self.job_runner._find_and_purge_task_instances_without_heartbeats()
+
+ session.expire_all()
+ sibling.refresh_from_db(session=session)
+ assert sibling.state == TaskInstanceState.FAILED
+
+ def test_heartbeat_timeout_skips_ti_completed_concurrently(self,
dag_maker, session):
+ """
+ The heartbeat scan and a worker can race: the worker can commit a
terminal state (SUCCESS)
+ around the same time the scan picks the TI up. The purge must
revalidate the committed state
+ and skip such a TI, so it neither clobbers the terminal state with
FAILED nor emits a
+ spurious failure callback.
+ """
+ with dag_maker(dag_id="hb_timeout_concurrent_success",
session=session):
+ EmptyOperator(task_id="t1", on_failure_callback=lambda ctx: None)
+
+ dag_run = dag_maker.create_dagrun(run_id="test_run",
state=DagRunState.RUNNING)
+
+ executor = MockExecutor(do_update=False)
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[executor])
+
+ ti = dag_run.get_task_instance(task_id="t1")
+ ti.state = TaskInstanceState.RUNNING
+ ti.try_number = 1
+ ti.max_tries = 0
+ ti.queued_by_job_id = scheduler_job.id
+ ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600)
+ session.merge(ti)
+ session.commit()
+
+ # Simulate the worker winning the race: the DB row is now SUCCESS, but
the in-memory ``ti``
+ # still reads RUNNING, exactly as it would after an unlocked scan
handed it to the purge.
+ session.execute(
+ update(TaskInstance)
+ .where(TaskInstance.id == ti.id)
+ .values(state=TaskInstanceState.SUCCESS,
end_date=timezone.utcnow())
+ )
+ session.commit()
+ assert ti.state == TaskInstanceState.RUNNING
+
+ self.job_runner._purge_task_instances_without_heartbeats([ti],
session=session)
+
+ self.job_runner.executor.callback_sink.send.assert_not_called()
+ session.expire_all()
+ ti.refresh_from_db(session=session)
+ assert ti.state == TaskInstanceState.SUCCESS
+
+ def test_heartbeat_timeout_scan_locks_rows(self, dag_maker, session):
+ """
+ The heartbeat scan must lock the TI rows (``with_row_locks``,
``of=TI``, ``skip_locked=True``)
+ so a worker cannot commit a terminal state on the same TI between the
scan and the
+ handle_failure() that follows in the same transaction.
+ """
+ with dag_maker(dag_id="hb_timeout_scan_locks", session=session):
+ EmptyOperator(task_id="t1")
+
+ dag_run = dag_maker.create_dagrun(run_id="test_run",
state=DagRunState.RUNNING)
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[MockExecutor(do_update=False)])
+
+ ti = dag_run.get_task_instance(task_id="t1")
+ ti.state = TaskInstanceState.RUNNING
+ ti.queued_by_job_id = scheduler_job.id
+ ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600)
+ session.merge(ti)
+ session.commit()
+
+ with mock.patch(
+ "airflow.jobs.scheduler_job_runner.with_row_locks",
+ wraps=with_row_locks,
+ ) as wrapped:
+ found =
self.job_runner._find_task_instances_without_heartbeats(session=session)
+
+ assert [t.id for t in found] == [ti.id]
+ ti_lock_calls = [call for call in wrapped.mock_calls if
call.kwargs.get("of") is TaskInstance]
+ assert len(ti_lock_calls) == 1, f"Expected one with_row_locks call for
TI, got {ti_lock_calls}"
+ assert ti_lock_calls[0].kwargs["skip_locked"] is True
+ assert ti_lock_calls[0].kwargs["session"] is session
+
@conf_vars({("scheduler", "num_stuck_in_queued_retries"): "1"})
def test_stuck_in_queued_callback_bundle_version_follows_dag_run(
self, dag_maker, session, mock_executors
@@ -12205,7 +12472,7 @@ def _extract_bundle_version(ti):
class TestSchedulerCallbackBundleInfoDagVersionNullable:
"""
- Verify the bundle_name / bundle_version extraction logic used at all four
+ Verify the bundle_name / bundle_version extraction logic used at all five
TaskCallbackRequest / EmailRequest creation sites in
scheduler_job_runner.py.
When dag_version is present -> use dag_version.bundle_name /
bundle_version.
