Fix schedule_tis HA race on try_number/state transitions

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
 

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):