rjgoyln commented on code in PR #75962:
URL: https://github.com/apache/airflow/pull/75962#discussion_r4417285446
##########
airflow-core/newsfragments/75962.bugfix.rst:
##########
@@ -0,0 +1 @@
+Fix the scheduler deadlocking against the triggerer when it fails deferred
tasks whose trigger timeout has passed. The sweep now claims the timed-out rows
with ``SKIP LOCKED`` and fails at most 500 of them per tick, leaving the rest
for later ticks.
Review Comment:
OK! I have removed the newsfragment.
##########
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(wraps=session)
+ # Not wrapped: rolling the real session back would discard the
fixture's uncommitted rows.
+ might_fail_session.rollback = MagicMock()
Review Comment:
Thanks for the review!
I used the session instance as the spec rather than the Session class, since
bind is initialized on the instance. This also avoids issues in the MySQL CI
job.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]