Stop failing reschedule-mode sensors on a stale executor success

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -26,7 +26,7 @@
 import sys
 import time
 from collections import Counter, defaultdict, deque
-from collections.abc import Callable, Collection, Iterable, Iterator
+from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
 from contextlib import ExitStack
 from datetime import datetime, timedelta
 from functools import lru_cache, partial
@@ -107,6 +107,7 @@
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.models.taskinstance import TaskInstance
 from airflow.models.taskinstancekey import TaskInstanceKey
+from airflow.models.taskreschedule import TaskReschedule
 from airflow.models.team import Team
 from airflow.models.trigger import TRIGGER_FAIL_REPR, Trigger, TriggerFailureReason, handle_event_submit
 from airflow.observability.metrics import stats_utils
@@ -1319,6 +1320,23 @@ def _emit_executor_events_batch_metrics(num_events: int) -> None:
         stats.gauge("scheduler.executor_events.batch_size", num_events)
         stats.incr("scheduler.executor_events.processed", num_events)
 
+    @staticmethod
+    def _find_ti_ids_rescheduled_this_try(tis: Sequence[TI], *, session: Session) -> set[UUID]:
+        """Find which of the given task instances have already been rescheduled during their current try."""
+        candidate_ids = [
+            ti.id
+            for ti in tis
+            if ti.state in (TaskInstanceState.SCHEDULED, TaskInstanceState.QUEUED) and ti.next_method is None
+        ]
+        if not candidate_ids:
+            return set()
+        # Resolved once for the whole batch to keep the reschedule lookup off the per-event hot path.
+        return set(
+            session.scalars(
+                select(TaskReschedule.ti_id).where(TaskReschedule.ti_id.in_(candidate_ids)).distinct()
+            )
+        )
+
     @classmethod
     def process_executor_events(
         cls,
@@ -1452,7 +1470,8 @@ def process_executor_events(
         # row lock this entire set of taskinstances to make sure the scheduler doesn't fail when we have
         # multi-schedulers
         locked_query = with_row_locks(query, of=TI, session=session, skip_locked=True)
-        tis: Iterator[TI] = session.scalars(locked_query)
+        tis: Sequence[TI] = session.scalars(locked_query).all()
+        ti_ids_rescheduled_this_try = cls._find_ti_ids_rescheduled_this_try(tis, session=session)
         for ti in tis:
             try_number = ti_primary_key_to_try_number_map[ti.key.primary]
             buffer_key = ti.key.with_try_number(try_number)
@@ -1516,6 +1535,8 @@ def process_executor_events(
             # from the worker exit after defer() has not been processed yet - should not fail it.
             # 4) the trigger already put the TI back to queued (resume after defer) but the executor success
             # from the worker exit after defer() has not been processed yet - should not fail it.
+            # 5) a sensor in reschedule mode was put back to scheduled or queued for its next poke but the
+            # executor success from the previous poke has not been processed yet - should not fail it.
 
             # All of this could also happen if the state is "running",
             # but that is handled by the scheduler detecting task instances without heartbeats.
@@ -1530,11 +1551,18 @@ def process_executor_events(
                 ti.queued_by_job_id != job_id  # Another scheduler has queued this task again
                 or executor.has_task(ti)  # This scheduler has this task already
                 or (
-                    # Resume-after-defer: trigger moved TI to scheduled or queued (next_method set)
-                    # before we saw the executor success from the defer exit for the same try_number.
                     ti.state in (TaskInstanceState.SCHEDULED, TaskInstanceState.QUEUED)
                     and state == TaskInstanceState.SUCCESS
-                    and ti.next_method is not None
+                    and (
+                        # Resume-after-defer: trigger moved TI to scheduled or queued (next_method set)
+                        # before we saw the executor success from the defer exit for the same try_number.
+                        ti.next_method is not None
+                        # Sensor in reschedule mode: the TI was put back for its next poke before we saw
+                        # the executor success from the previous poke. A reschedule exit leaves
+                        # next_method unset and keeps the same try_number, so an earlier reschedule of
+                        # this try is the only thing that tells this apart from an externally killed task.
+                        or ti.id in ti_ids_rescheduled_this_try
+                    )
                 )
             )
 

airflow-core/tests/unit/jobs/test_scheduler_job.py

@@ -95,6 +95,7 @@
 from airflow.models.pool import Pool
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.models.taskinstance import TaskInstance
+from airflow.models.taskreschedule import TaskReschedule
 from airflow.models.team import Team
 from airflow.models.trigger import Trigger
 from airflow.partition_mappers.base import (
@@ -1006,6 +1007,55 @@ def test_process_executor_events_stale_success_when_queued_after_defer(
             tags={"dag_id": dag_id, "task_id": ti1.task_id},
         )
 
+    @pytest.mark.parametrize("ti_state", [State.SCHEDULED, State.QUEUED])
+    @mock.patch("airflow._shared.observability.metrics.stats._get_backend")
+    def test_process_executor_events_stale_success_between_reschedule_pokes(
+        self, mock_get_backend, ti_state, dag_maker
+    ):
+        """Stale success from an earlier poke of a reschedule-mode sensor must not read as a state mismatch."""
+        mock_stats = mock.MagicMock(spec=StatsLogger)
+        mock_get_backend.return_value = mock_stats
+        dag_id = "test_process_executor_events_stale_success_between_reschedule_pokes"
+
+        session = settings.Session()
+        with dag_maker(dag_id=dag_id, fileloc="/test_path1/"):
+            task1 = EmptyOperator(task_id="dummy_task")
+        ti1 = dag_maker.create_dagrun().get_task_instance(task1.task_id)
+
+        executor = MockExecutor(do_update=False)
+        scheduler_job = Job()
+        session.add(scheduler_job)
+        session.flush()
+        self.job_runner = SchedulerJobRunner(scheduler_job, executors=[executor])
+
+        # A reschedule exit leaves next_method unset and keeps the same try_number, so the reschedule
+        # row is the only thing standing between this and the externally-killed branch.
+        ti1.state = ti_state
+        ti1.next_method = None
+        ti1.queued_by_job_id = scheduler_job.id
+        ti1.try_number = 1
+        session.merge(ti1)
+        poke_end = DEFAULT_DATE + datetime.timedelta(seconds=5)
+        session.add(
+            TaskReschedule(
+                ti_id=ti1.id,
+                start_date=DEFAULT_DATE,
+                end_date=poke_end,
+                reschedule_date=poke_end + datetime.timedelta(seconds=60),
+            )
+        )
+        session.commit()
+
+        executor.event_buffer[ti1.key] = State.SUCCESS, None
+        executor.has_task = mock.MagicMock(spec=executor.has_task, return_value=False)
+        mock_stats.incr.reset_mock()
+
+        self.job_runner._process_executor_events(executor=executor, session=session)
+        ti1.refresh_from_db(session=session)
+        assert ti1.state == ti_state
+        self.job_runner.executor.callback_sink.send.assert_not_called()
+        mock_stats.incr.assert_called_once_with("scheduler.executor_events.processed", count=1)
+
     @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest")
     @mock.patch("airflow._shared.observability.metrics.stats._get_backend")
     def test_process_executor_events_multiple_try_numbers_warns(