Stop the scheduler deadlocking with the triggerer on trigger timeouts

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

@@ -173,6 +173,10 @@
 # safety bound, not a behavioural knob operators need to tune.
 MAX_PARTITION_DAG_RUNS_PER_LOOP = 500
 
+# Per-tick cap on timed-out deferred task instances the scheduler fails, for the
+# same reason as above; the rest are failed on subsequent ticks.
+MAX_TIMED_OUT_TASK_INSTANCES_PER_LOOP = 500
+
 MAX_TI_REFILL_QUERIES_PER_LOOP = 3
 """
 Max extra task instance queries per critical section to fill a partially selected batch.
@@ -3862,23 +3866,49 @@ def check_trigger_timeouts(
         """Mark any "deferred" task as failed if the trigger or execution timeout has passed."""
         for attempt in run_with_db_retries(max_retries, logger=self.log):
             with attempt:
-                result = session.execute(
-                    update(TI)
-                    .where(
-                        TI.state == TaskInstanceState.DEFERRED,
-                        TI.trigger_timeout < timezone.utcnow(),
+                try:
+                    now = timezone.utcnow()
+                    query = (
+                        select(TI.id)
+                        .where(
+                            TI.state == TaskInstanceState.DEFERRED,
+                            TI.trigger_timeout < now,
+                        )
+                        .order_by(TI.id)
+                        .limit(MAX_TIMED_OUT_TASK_INSTANCES_PER_LOOP)
                     )
-                    .values(
-                        state=TaskInstanceState.SCHEDULED,
-                        next_method=TRIGGER_FAIL_REPR,
-                        next_kwargs={"error": TriggerFailureReason.TRIGGER_TIMEOUT},
-                        scheduled_dttm=timezone.utcnow(),
-                        trigger_id=None,
+                    # Claim the rows rather than letting a blind bulk UPDATE wait on them: a sweep
+                    # that waits can be one side of a lock cycle with the triggerer, which writes
+                    # these same rows as tasks defer and fire. SKIP LOCKED leaves the rows a writer
+                    # already holds for the next tick.
+                    query = with_row_locks(query, of=TI, session=session, skip_locked=True)
+                    timed_out_ids = session.scalars(query).all()
+                    if not timed_out_ids:
+                        return
+
+                    result = session.execute(
+                        update(TI)
+                        .where(
+                            TI.id.in_(timed_out_ids),
+                            TI.state == TaskInstanceState.DEFERRED,
+                            TI.trigger_timeout < now,
+                        )
+                        .values(
+                            state=TaskInstanceState.SCHEDULED,
+                            next_method=TRIGGER_FAIL_REPR,
+                            next_kwargs={"error": TriggerFailureReason.TRIGGER_TIMEOUT},
+                            scheduled_dttm=timezone.utcnow(),
+                            trigger_id=None,
+                        )
                     )
-                )
-                num_timed_out_tasks = getattr(result, "rowcount", 0)
-                if num_timed_out_tasks:
-                    self.log.info("Timed out %i deferred tasks without fired triggers", num_timed_out_tasks)
+                    num_timed_out_tasks = getattr(result, "rowcount", 0)
+                    if num_timed_out_tasks:
+                        self.log.info(
+                            "Timed out %i deferred tasks without fired triggers", num_timed_out_tasks
+                        )
+                except OperationalError:
+                    session.rollback()
+                    raise
 
     @provide_session
     def check_awaiting_input_timeouts(

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

@@ -9540,6 +9540,95 @@ def test_timeout_triggers(self, dag_maker):
         assert ti1.next_method == "__fail__"
         assert ti2.state == State.DEFERRED
 
+    def _make_timed_out_deferred_tis(self, dag_maker, session, count: int) -> list[TaskInstance]:
+        with dag_maker(
+            dag_id="test_timeout_triggers_sweep",
+            start_date=DEFAULT_DATE,
+            schedule="@once",
+            max_active_runs=count,
+            session=session,
+        ):
+            EmptyOperator(task_id="dummy1")
+
+        tis = []
+        for i in range(count):
+            dag_run = dag_maker.create_dagrun(
+                run_id=f"test_sweep_{i}",
+                logical_date=DEFAULT_DATE + datetime.timedelta(seconds=i),
+            )
+            ti = dag_run.get_task_instance("dummy1", session=session)
+            ti.state = State.DEFERRED
+            ti.trigger_timeout = timezone.utcnow() - datetime.timedelta(seconds=60)
+            tis.append(ti)
+        session.flush()
+        return tis
+
+    @mock.patch("airflow.jobs.scheduler_job_runner.with_row_locks", wraps=with_row_locks)
+    def test_timeout_triggers_locks_rows_with_skip_locked(self, wrapped, dag_maker, session):
+        """
+        The timeout sweep must claim the TI rows with ``skip_locked`` rather than letting a blind
+        bulk UPDATE wait on rows the triggerer holds, which is what deadlocks it on MySQL.
+        """
+        self._make_timed_out_deferred_tis(dag_maker, session, 1)
+        self.job_runner = SchedulerJobRunner(job=Job())
+
+        self.job_runner.check_trigger_timeouts(session=session)
+
+        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
+
+    def test_timeout_triggers_roll_back_before_retrying(self, dag_maker, session):
+        """
+        A failed attempt must roll back before the next one, the way ``adopt_or_reset_orphaned_tasks``
+        does, so the retry starts from a clean transaction instead of the half-claimed rows.
+        """
+        from sqlalchemy.exc import OperationalError
+
+        tis = self._make_timed_out_deferred_tis(dag_maker, session, 1)
+        self.job_runner = SchedulerJobRunner(job=Job())
+
+        might_fail_session = MagicMock(spec=session, wraps=session)
+        # Not wrapped: rolling the real session back would discard the fixture's uncommitted rows.
+        might_fail_session.rollback = MagicMock(spec=session.rollback)
+        claim_failed = False
+
+        def claim(*args, **kwargs):
+            nonlocal claim_failed
+            if not claim_failed:
+                claim_failed = True
+                raise OperationalError("any_statement", "any_params", "any_orig")
+            return session.scalars(*args, **kwargs)
+
+        might_fail_session.scalars.side_effect = claim
+
+        self.job_runner.check_trigger_timeouts(max_retries=2, session=might_fail_session)
+
+        might_fail_session.rollback.assert_called_once()
+        session.refresh(tis[0])
+        assert tis[0].state == State.SCHEDULED
+        assert tis[0].next_method == "__fail__"
+
+    @mock.patch("airflow.jobs.scheduler_job_runner.MAX_TIMED_OUT_TASK_INSTANCES_PER_LOOP", 2)
+    def test_timeout_triggers_sweep_is_bounded(self, dag_maker, session):
+        """Each sweep fails at most one batch; the rest are picked up on later ticks."""
+        tis = self._make_timed_out_deferred_tis(dag_maker, session, 3)
+        self.job_runner = SchedulerJobRunner(job=Job())
+
+        self.job_runner.check_trigger_timeouts(session=session)
+
+        for ti in tis:
+            session.refresh(ti)
+        assert sorted(ti.state for ti in tis) == [State.DEFERRED, State.SCHEDULED, State.SCHEDULED]
+
+        self.job_runner.check_trigger_timeouts(session=session)
+
+        for ti in tis:
+            session.refresh(ti)
+        assert all(ti.state == State.SCHEDULED for ti in tis)
+        assert all(ti.next_method == "__fail__" for ti in tis)
+
     def test_awaiting_input_timeout_with_defaults_resumes(self, dag_maker):
         """
         A parked ``awaiting_input`` task past its deadline with defaults is resumed to SCHEDULED by
@@ -9734,7 +9823,8 @@ def side_effect(*args, **kwargs):
                 ti1.trigger_timeout = timezone.utcnow() - datetime.timedelta(seconds=60)
                 ti2.state = State.DEFERRED
                 ti2.trigger_timeout = timezone.utcnow() + datetime.timedelta(seconds=60)
-                session.flush()
+                # Committed, not just flushed: a failed attempt now rolls back before retrying.
+                session.commit()
 
                 # Boot up the scheduler and make it check timeouts
                 scheduler_job = Job()