diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 6a5e74bcc0a22..2d5922dc34f99 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/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( diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index c2636b3276410..a26b6776952cc 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/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()