From 590d384c0d7401daf6a70b1f04ff991ff8921e82 Mon Sep 17 00:00:00 2001 From: Siddharthan P S Date: Wed, 11 Mar 2026 13:51:06 -0400 Subject: [PATCH] Guard schedule_tis updates against HA race and add regression tests --- airflow-core/src/airflow/models/dagrun.py | 74 +++++- airflow-core/tests/unit/models/test_dagrun.py | 237 +++++++++++++++++- 2 files changed, 297 insertions(+), 14 deletions(-) diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index 61242e45390d6..2113efd9ca12b 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -1987,29 +1987,56 @@ def schedule_tis( empty_ti_ids.append(ti.id) count = 0 + # Guard updates by state as well as TI id so stale scheduler views do not + # re-schedule already transitioned rows in HA race windows. + non_null_schedulable_states = tuple(s for s in SCHEDULEABLE_STATES if s is not None) + schedulable_state_clause = or_( + TI.state.is_(None), + TI.state.in_(non_null_schedulable_states), + ) + non_reschedule_schedulable_states = tuple( + s for s in non_null_schedulable_states if s != TaskInstanceState.UP_FOR_RESCHEDULE + ) + incrementing_schedulable_state_clause = ( + or_( + TI.state.is_(None), + TI.state.in_(non_reschedule_schedulable_states), + ) + if non_reschedule_schedulable_states + else TI.state.is_(None) + ) if schedulable_ti_ids: schedulable_ti_ids_chunks = chunks( schedulable_ti_ids, max_tis_per_query or len(schedulable_ti_ids) ) for id_chunk in schedulable_ti_ids_chunks: - result = session.execute( + up_for_reschedule_result = session.execute( update(TI) - .where(TI.id.in_(id_chunk)) + .where( + TI.id.in_(id_chunk), + schedulable_state_clause, + TI.state == TaskInstanceState.UP_FOR_RESCHEDULE, + ) .values( + try_number=TI.try_number, + state=TaskInstanceState.SCHEDULED, + scheduled_dttm=timezone.utcnow(), + ) + .execution_options(synchronize_session=False) + ) + incrementing_result = session.execute( + update(TI) + .where(TI.id.in_(id_chunk), incrementing_schedulable_state_clause) + .values( + try_number=TI.try_number + 1, state=TaskInstanceState.SCHEDULED, scheduled_dttm=timezone.utcnow(), - try_number=case( - ( - or_(TI.state.is_(None), TI.state != TaskInstanceState.UP_FOR_RESCHEDULE), - TI.try_number + 1, - ), - else_=TI.try_number, - ), ) .execution_options(synchronize_session=False) ) - count += getattr(result, "rowcount", 0) + count += getattr(up_for_reschedule_result, "rowcount", 0) + count += getattr(incrementing_result, "rowcount", 0) if debug_try_number_check: rows = session.execute( select(TI.id, TI.try_number, TI.state).where(TI.id.in_(id_chunk)) @@ -2026,6 +2053,8 @@ def schedule_tis( continue expected_try_number, pre_update_try_number, pre_update_state = expected db_try_number, db_state = db_row + if db_state != TaskInstanceState.SCHEDULED: + continue if db_try_number != expected_try_number: self.log.warning( "schedule_tis: try_number mismatch after scheduling for ti_id=%s " @@ -2047,21 +2076,40 @@ def schedule_tis( if empty_ti_ids: dummy_ti_ids_chunks = chunks(empty_ti_ids, max_tis_per_query or len(empty_ti_ids)) for id_chunk in dummy_ti_ids_chunks: - result = session.execute( + up_for_reschedule_result = session.execute( update(TI) - .where(TI.id.in_(id_chunk)) + .where( + TI.id.in_(id_chunk), + schedulable_state_clause, + TI.state == TaskInstanceState.UP_FOR_RESCHEDULE, + ) .values( + try_number=TI.try_number, state=TaskInstanceState.SUCCESS, start_date=timezone.utcnow(), end_date=timezone.utcnow(), duration=0, + ) + .execution_options( + synchronize_session=False, + ) + ) + incrementing_result = session.execute( + update(TI) + .where(TI.id.in_(id_chunk), incrementing_schedulable_state_clause) + .values( try_number=TI.try_number + 1, + state=TaskInstanceState.SUCCESS, + start_date=timezone.utcnow(), + end_date=timezone.utcnow(), + duration=0, ) .execution_options( synchronize_session=False, ) ) - count += getattr(result, "rowcount", 0) + count += getattr(up_for_reschedule_result, "rowcount", 0) + count += getattr(incrementing_result, "rowcount", 0) return count diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 14722f83b0cce..160babfe35703 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -28,7 +28,7 @@ import pendulum import pytest from opentelemetry.sdk.trace import TracerProvider -from sqlalchemy import func, select +from sqlalchemy import func, select, update from sqlalchemy.orm import joinedload from airflow import settings @@ -2046,6 +2046,241 @@ def test_schedule_tis_map_index(dag_maker, session): assert ti2.state == TaskInstanceState.SUCCESS +def test_schedule_tis_does_not_increment_try_number_if_ti_already_queued_by_other_scheduler( + dag_maker, session +): + from airflow.utils.session import create_session + + with dag_maker(session=session) as dag: + BashOperator(task_id="task", bash_command="echo 1") + + dr = dag_maker.create_dagrun(session=session) + ti = dr.get_task_instance("task", session=session) + assert ti is not None + ti.refresh_from_task(dag.get_task("task")) + assert ti.state is None + + ti.try_number = 1 + session.flush() + session.commit() + + with create_session() as other_session: + filter_for_tis = TI.filter_for_tis([ti]) + assert filter_for_tis is not None + other_session.execute( + update(TI) + .where(filter_for_tis) + .values(state=TaskInstanceState.QUEUED, try_number=1) + .execution_options(synchronize_session=False) + ) + + assert dr.schedule_tis((ti,), session=session) == 0 + + refreshed_ti = session.scalar( + select(TI).where( + TI.dag_id == ti.dag_id, + TI.task_id == ti.task_id, + TI.run_id == ti.run_id, + TI.map_index == ti.map_index, + ) + ) + assert refreshed_ti is not None + assert refreshed_ti.state == TaskInstanceState.QUEUED + assert refreshed_ti.try_number == 1 + + +def test_schedule_tis_empty_operator_does_not_short_circuit_if_ti_already_queued(dag_maker, session): + from airflow.utils.session import create_session + + with dag_maker(session=session) as dag: + EmptyOperator(task_id="empty_task") + + dr = dag_maker.create_dagrun(session=session) + ti = dr.get_task_instance("empty_task", session=session) + assert ti is not None + ti.refresh_from_task(dag.get_task("empty_task")) + assert ti.state is None + + ti.try_number = 1 + session.flush() + session.commit() + + with create_session() as other_session: + filter_for_tis = TI.filter_for_tis([ti]) + assert filter_for_tis is not None + other_session.execute( + update(TI) + .where(filter_for_tis) + .values(state=TaskInstanceState.QUEUED, try_number=1) + .execution_options(synchronize_session=False) + ) + + assert dr.schedule_tis((ti,), session=session) == 0 + + refreshed_ti = session.scalar( + select(TI).where( + TI.dag_id == ti.dag_id, + TI.task_id == ti.task_id, + TI.run_id == ti.run_id, + TI.map_index == ti.map_index, + ) + ) + assert refreshed_ti is not None + assert refreshed_ti.state == TaskInstanceState.QUEUED + assert refreshed_ti.try_number == 1 + + +def test_schedule_tis_up_for_reschedule_does_not_increment_try_number(dag_maker, session): + with dag_maker(session=session) as dag: + BashOperator(task_id="task", bash_command="echo 1") + + dr = dag_maker.create_dagrun(session=session) + ti = dr.get_task_instance("task", session=session) + assert ti is not None + ti.refresh_from_task(dag.get_task("task")) + + ti.state = TaskInstanceState.UP_FOR_RESCHEDULE + ti.try_number = 3 + session.commit() + + assert dr.schedule_tis((ti,), session=session) == 1 + session.commit() + session.expire_all() + + refreshed_ti = session.scalar( + select(TI).where( + TI.dag_id == ti.dag_id, + TI.task_id == ti.task_id, + TI.run_id == ti.run_id, + TI.map_index == ti.map_index, + ) + ) + assert refreshed_ti is not None + assert refreshed_ti.state == TaskInstanceState.SCHEDULED + assert refreshed_ti.try_number == 3 + + +def test_schedule_tis_is_noop_if_ti_transitions_to_nonschedulable_state_before_update(dag_maker, session): + from airflow.utils.session import create_session + + with dag_maker(session=session) as dag: + BashOperator(task_id="task", bash_command="echo 1") + + dr = dag_maker.create_dagrun(session=session) + ti = dr.get_task_instance("task", session=session) + assert ti is not None + ti.refresh_from_task(dag.get_task("task")) + + ti.try_number = 1 + session.commit() + + with create_session() as other_session: + filter_for_tis = TI.filter_for_tis([ti]) + assert filter_for_tis is not None + other_session.execute( + update(TI) + .where(filter_for_tis) + .values(state=TaskInstanceState.QUEUED, try_number=1) + .execution_options(synchronize_session=False) + ) + + assert dr.schedule_tis((ti,), session=session) == 0 + + refreshed_ti = session.scalar( + select(TI).where( + TI.dag_id == ti.dag_id, + TI.task_id == ti.task_id, + TI.run_id == ti.run_id, + TI.map_index == ti.map_index, + ) + ) + assert refreshed_ti is not None + assert refreshed_ti.state == TaskInstanceState.QUEUED + assert refreshed_ti.try_number == 1 + + +def test_schedule_tis_empty_operator_is_noop_if_ti_already_running(dag_maker, session): + from airflow.utils.session import create_session + + with dag_maker(session=session) as dag: + EmptyOperator(task_id="empty_task") + + dr = dag_maker.create_dagrun(session=session) + ti = dr.get_task_instance("empty_task", session=session) + assert ti is not None + ti.refresh_from_task(dag.get_task("empty_task")) + + ti.try_number = 3 + session.commit() + + with create_session() as other_session: + filter_for_tis = TI.filter_for_tis([ti]) + assert filter_for_tis is not None + other_session.execute( + update(TI) + .where(filter_for_tis) + .values(state=TaskInstanceState.RUNNING, try_number=3) + .execution_options(synchronize_session=False) + ) + + assert dr.schedule_tis((ti,), session=session) == 0 + + refreshed_ti = session.scalar( + select(TI).where( + TI.dag_id == ti.dag_id, + TI.task_id == ti.task_id, + TI.run_id == ti.run_id, + TI.map_index == ti.map_index, + ) + ) + assert refreshed_ti is not None + assert refreshed_ti.state == TaskInstanceState.RUNNING + assert refreshed_ti.try_number == 3 + + +def test_schedule_tis_only_one_scheduler_update_succeeds_when_competing(dag_maker, session): + from airflow.utils.session import create_session + + with dag_maker(session=session) as dag: + BashOperator(task_id="task", bash_command="echo 1") + + dr = dag_maker.create_dagrun(session=session) + ti = dr.get_task_instance("task", session=session) + assert ti is not None + ti.refresh_from_task(dag.get_task("task")) + assert ti.state is None + + ti.try_number = 0 + session.commit() + + assert dr.schedule_tis((ti,), session=session) == 1 + session.commit() + + with create_session() as scheduler_b_session: + ti_b = scheduler_b_session.scalar( + select(TI).where( + TI.dag_id == ti.dag_id, + TI.task_id == ti.task_id, + TI.run_id == ti.run_id, + TI.map_index == ti.map_index, + ) + ) + assert ti_b is not None + assert dr.schedule_tis((ti_b,), session=scheduler_b_session) == 0 + + refreshed_ti = session.scalar( + select(TI).where( + TI.dag_id == ti.dag_id, + TI.task_id == ti.task_id, + TI.run_id == ti.run_id, + TI.map_index == ti.map_index, + ) + ) + assert refreshed_ti is not None + assert refreshed_ti.state == TaskInstanceState.SCHEDULED + assert refreshed_ti.try_number == 1 + + @pytest.mark.xfail(reason="We can't keep this behaviour with remote workers where scheduler can't reach xcom") @pytest.mark.need_serialized_dag def test_schedule_tis_start_trigger(dag_maker, session):