fat-catTW commented on code in PR #75962:
URL: https://github.com/apache/airflow/pull/75962#discussion_r4415821093
##########
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:
I think this newsfragment may not be needed. The PR looks like a bug fix.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -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)
Review Comment:
I think this is ineffective on MySQL: `with_row_locks()` returns before
applying `SKIP LOCKED` when `supports_for_update_of` is false, although MySQL 8
supports `FOR UPDATE SKIP LOCKED`.
Could we update the guard and test the compiled MySQL SQL?
##########
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 PR!
Maybe we can give this mock a `Session` spec, or use `autospec`?
This test verifies the rollback/retry interaction, and an unspecced
`MagicMock` would also accept a misspelled or removed Session API.
--
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]
