Fix MySQL deadlock when queueing asset-triggered Dags

airflow-core/src/airflow/assets/manager.py

@@ -518,13 +518,21 @@ def _queue_dagruns(
         # mapped) tasks update the same asset, this can fail with a unique
         # constraint violation.
         #
-        # If we support it, use ON CONFLICT to do nothing, otherwise
+        # If we support it, use ON CONFLICT / ON DUPLICATE KEY to do nothing, otherwise
         # "fallback" to running this in a nested transaction. This is needed
         # so that the adding of these rows happens in the same transaction
         # where `ti.state` is changed.
-        if get_dialect_name(session) == "postgresql":
-            return cls._queue_dagruns_nonpartitioned_postgres(asset_id, non_partitioned_dags, session)
-        return cls._queue_dagruns_nonpartitioned_slow_path(asset_id, non_partitioned_dags, session)
+        #
+        # Rows are inserted in sorted order so concurrent fan-outs acquire
+        # row locks in a consistent order, which prevents deadlocks (a
+        # set's iteration order differs between processes).
+        dag_ids = sorted(dag.dag_id for dag in non_partitioned_dags)
+        dialect = get_dialect_name(session)
+        if dialect == "postgresql":
+            return cls._queue_dagruns_nonpartitioned_postgres(asset_id, dag_ids, session)
+        if dialect == "mysql":
+            return cls._queue_dagruns_nonpartitioned_mysql(asset_id, dag_ids, session)
+        return cls._queue_dagruns_nonpartitioned_slow_path(asset_id, dag_ids, session)
 
     @classmethod
     def _queue_partitioned_dags(
@@ -790,30 +798,37 @@ def _get_or_create_apdr(
 
     @classmethod
     def _queue_dagruns_nonpartitioned_slow_path(
-        cls, asset_id: int, dags_to_queue: set[DagModel], session: Session
+        cls, asset_id: int, dag_ids: list[str], session: Session
     ) -> None:
-        def _queue_dagrun_if_needed(dag: DagModel) -> str | None:
-            item = AssetDagRunQueue(target_dag_id=dag.dag_id, asset_id=asset_id)
+        for dag_id in dag_ids:
+            item = AssetDagRunQueue(target_dag_id=dag_id, asset_id=asset_id)
             # Don't error whole transaction when a single RunQueue item conflicts.
             # https://docs.sqlalchemy.org/en/20/orm/session_transaction.html#using-savepoint
             try:
                 with session.begin_nested():
                     session.merge(item)
             except exc.IntegrityError:
                 cls.logger().debug("Skipping record %s", item, exc_info=True)
-            return dag.dag_id
+        cls.logger().debug("consuming dag ids %s", dag_ids)
+
+    @classmethod
+    def _queue_dagruns_nonpartitioned_mysql(cls, asset_id: int, dag_ids: list[str], session: Session) -> None:
+        from sqlalchemy.dialects.mysql import insert
 
-        queued_results = (_queue_dagrun_if_needed(dag) for dag in dags_to_queue)
-        if queued_dag_ids := [r for r in queued_results if r is not None]:
-            cls.logger().debug("consuming dag ids %s", queued_dag_ids)
+        # The ON DUPLICATE KEY UPDATE skips duplicates while still throwing errors for other critical issues.
+        stmt = insert(AssetDagRunQueue).values(
+            [{"asset_id": asset_id, "target_dag_id": dag_id} for dag_id in dag_ids]
+        )
+        stmt = stmt.on_duplicate_key_update(target_dag_id=stmt.inserted.target_dag_id)
+        session.execute(stmt)
 
     @classmethod
     def _queue_dagruns_nonpartitioned_postgres(
-        cls, asset_id: int, dags_to_queue: set[DagModel], session: Session
+        cls, asset_id: int, dag_ids: list[str], session: Session
     ) -> None:
         from sqlalchemy.dialects.postgresql import insert
 
-        values = [{"target_dag_id": dag.dag_id} for dag in dags_to_queue]
+        values = [{"target_dag_id": dag_id} for dag_id in dag_ids]
         stmt = insert(AssetDagRunQueue).values(asset_id=asset_id).on_conflict_do_nothing()
         session.execute(stmt, values)
 

airflow-core/src/airflow/models/trigger.py

@@ -37,7 +37,7 @@
 from airflow.models.taskinstance import TaskInstance
 from airflow.serialization.enums import stringify_encoding_keys
 from airflow.triggers.base import BaseTaskEndEvent
-from airflow.utils.retries import run_with_db_retries
+from airflow.utils.retries import retry_db_transaction, run_with_db_retries
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.utils.sqlalchemy import UtcDateTime, get_dialect_name, with_row_locks
 from airflow.utils.state import TaskInstanceState
@@ -270,12 +270,17 @@ def clean_unused(cls, *, session: Session = NEW_SESSION) -> None:
 
     @classmethod
     @provide_session
+    @retry_db_transaction
     def submit_event(cls, trigger_id, event: TriggerEvent, *, session: Session = NEW_SESSION) -> None:
         """
         Fire an event.
 
         Resume all tasks that were in deferred state.
         Send an event to all assets associated to the trigger.
+
+        Retried as a whole on transient database errors, e.g. a deadlock
+        while fanning an asset event out to consumer Dags. Such errors abort
+        the entire transaction, so the retry must replay it from the start.
         """
         # Resume deferred tasks
         for task_instance in session.scalars(

airflow-core/tests/unit/assets/test_manager.py

@@ -26,6 +26,7 @@
 
 import pytest
 from sqlalchemy import delete, func, select
+from sqlalchemy.dialects import mysql
 from sqlalchemy.orm import Session
 
 from airflow import settings
@@ -149,6 +150,91 @@ def test_register_asset_change(self, session, mock_task_instance):
         )
         assert session.scalar(select(func.count()).select_from(AssetDagRunQueue)) == 2
 
+    @pytest.mark.usefixtures("dag_maker", "testing_dag_bundle")
+    def test_register_asset_change_twice_deduplicates_queue_records(self, session, mock_task_instance):
+        """Fanning out the same asset again must tolerate the already-queued rows, on every backend."""
+        asset_manager = AssetManager()
+
+        asset = Asset(uri="test://asset1", name="test_asset_uri", group="asset")
+        bundle_name = "testing"
+
+        dag1 = DagModel(dag_id="dag1", is_stale=False, bundle_name=bundle_name)
+        dag2 = DagModel(dag_id="dag2", is_stale=False, bundle_name=bundle_name)
+        session.add_all([dag1, dag2])
+
+        asm = AssetModel(uri="test://asset1/", name="test_asset_uri", group="asset")
+        session.add(asm)
+        asm.scheduled_dags = [DagScheduleAssetReference(dag_id=dag.dag_id) for dag in (dag1, dag2)]
+        session.execute(delete(AssetDagRunQueue))
+        session.flush()
+
+        for _ in range(2):
+            asset_manager.register_asset_change(
+                task_instance=mock_task_instance, asset=asset, session=session
+            )
+            session.flush()
+
+        assert (
+            session.scalar(select(func.count()).select_from(AssetEvent).where(AssetEvent.asset_id == asm.id))
+            == 2
+        )
+        assert session.scalar(select(func.count()).select_from(AssetDagRunQueue)) == 2
+
+    @pytest.mark.parametrize(
+        ("dialect", "expected_helper"),
+        [
+            ("postgresql", "_queue_dagruns_nonpartitioned_postgres"),
+            ("mysql", "_queue_dagruns_nonpartitioned_mysql"),
+            ("sqlite", "_queue_dagruns_nonpartitioned_slow_path"),
+        ],
+    )
+    @mock.patch.object(AssetManager, "_queue_dagruns_nonpartitioned_slow_path", autospec=True)
+    @mock.patch.object(AssetManager, "_queue_dagruns_nonpartitioned_mysql", autospec=True)
+    @mock.patch.object(AssetManager, "_queue_dagruns_nonpartitioned_postgres", autospec=True)
+    @mock.patch("airflow.assets.manager.get_dialect_name", autospec=True)
+    def test_queue_dagruns_dispatches_by_dialect_with_sorted_dag_ids(
+        self, mock_get_dialect_name, mock_postgres, mock_mysql, mock_slow_path, dialect, expected_helper
+    ):
+        mock_get_dialect_name.return_value = dialect
+        dags_to_queue = {DagModel(dag_id=dag_id) for dag_id in ("dag_b", "dag_a", "dag_c")}
+        session = mock.Mock(spec=Session)
+
+        AssetManager._queue_dagruns(
+            asset_id=1,
+            dags_to_queue=dags_to_queue,
+            partition_key=None,
+            partition_date=None,
+            event=mock.Mock(spec=AssetEvent),
+            task_instance=None,
+            session=session,
+        )
+
+        helper_mocks = {
+            "_queue_dagruns_nonpartitioned_postgres": mock_postgres,
+            "_queue_dagruns_nonpartitioned_mysql": mock_mysql,
+            "_queue_dagruns_nonpartitioned_slow_path": mock_slow_path,
+        }
+        for name, helper_mock in helper_mocks.items():
+            if name == expected_helper:
+                helper_mock.assert_called_once_with(1, ["dag_a", "dag_b", "dag_c"], session)
+            else:
+                helper_mock.assert_not_called()
+
+    def test_queue_dagruns_nonpartitioned_mysql_single_multi_row_upsert(self):
+        session = mock.Mock(spec=Session)
+
+        AssetManager._queue_dagruns_nonpartitioned_mysql(42, ["dag_a", "dag_b"], session)
+
+        session.execute.assert_called_once()
+        stmt = session.execute.call_args.args[0]
+        sql = str(stmt.compile(dialect=mysql.dialect(), compile_kwargs={"literal_binds": True}))
+        # NULL is a compile-time placeholder: the created_at default is evaluated per row at execution.
+        assert (
+            "INSERT INTO asset_dag_run_queue (asset_id, target_dag_id, created_at) "
+            "VALUES (42, 'dag_a', NULL), (42, 'dag_b', NULL) "
+            "ON DUPLICATE KEY UPDATE target_dag_id = VALUES(target_dag_id)"
+        ) in sql
+
     @pytest.mark.usefixtures("clear_assets")
     def test_register_asset_change_with_alias(
         self, session, dag_maker, mock_task_instance, testing_dag_bundle

airflow-core/tests/unit/models/test_trigger.py

@@ -27,6 +27,7 @@
 import pytz
 from cryptography.fernet import Fernet
 from sqlalchemy import delete, func, select
+from sqlalchemy.exc import OperationalError
 
 from airflow._shared.timezones import timezone
 from airflow.jobs.job import Job
@@ -256,6 +257,33 @@ def test_submit_event_no_n_plus_one_for_assets(_, session, asset_count, expected
         Trigger.submit_event(trigger_id, TriggerEvent("payload"), session=session)
 
 
+@patch("airflow.models.trigger.AssetManager.register_asset_change")
+def test_submit_event_retries_on_transient_db_error(mock_register_asset_change, session):
+    """
+    A transient database error during the asset-event fan-out rolls back the whole transaction,
+    so submit_event must replay it from the start.
+    """
+    trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={})
+    session.add(trigger)
+    session.flush()
+    asset = AssetModel("test")
+    asset.add_trigger(trigger, "test_asset_watcher")
+    session.add(asset)
+    session.commit()
+    trigger_id = trigger.id
+
+    deadlock = OperationalError(
+        "INSERT INTO asset_dag_run_queue",
+        {},
+        Exception("(1213, 'Deadlock found when trying to get lock; try restarting transaction')"),
+    )
+    mock_register_asset_change.side_effect = [deadlock, None]
+
+    Trigger.submit_event(trigger_id, TriggerEvent("payload"), session=session)
+
+    assert mock_register_asset_change.call_count == 2
+
+
 def test_submit_failure(session, create_task_instance):
     """
     Tests that failures submitted to a trigger fail their dependent