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]

Reply via email to