Reduce task success asset registration lock contention

airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst

@@ -51,6 +51,21 @@ In addition to scheduling Dags based on time, you can also schedule Dags to run
 
     :ref:`asset_definitions` for how to declare assets.
 
+.. note::
+
+    Asset events are registered asynchronously. When a task succeeds, Airflow durably records its
+    asset events and the scheduler registers them shortly afterwards rather than inside the
+    task-completion request, usually within one scheduler loop and longer under heavy load.
+    Downstream asset-scheduled Dags are therefore triggered after a brief, eventually consistent
+    delay. The scheduler retries a queued registration up to ``[scheduler]
+    asset_event_queue_max_attempts`` times, after which the row is parked and reported via the
+    ``asset.event_queue.failures`` metric. Use the ``[scheduler] asset_event_queue_drain_interval``
+    option to tune how often the scheduler processes queued asset events.
+
+    ``dag.test`` runs without a scheduler, so it registers each task's asset events itself once
+    the task finishes; the events stay visible during the test, though downstream asset-scheduled
+    Dags are still only triggered by a running scheduler.
+
 Schedule Dags with assets
 -------------------------
 

airflow-core/docs/migrations-ref.rst

@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | Revision ID             | Revises ID       | Airflow Version   | Description                                                  |
 +=========================+==================+===================+==============================================================+
-| ``7a98f1b7dbd3`` (head) | ``c4e7a1f9b2d0`` | ``3.4.0``         | Add index on asset_event (asset_id, partition_key).          |
+| ``a63ae7baae97`` (head) | ``7a98f1b7dbd3`` | ``3.4.0``         | Add asset_event_queue table.                                 |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``7a98f1b7dbd3``        | ``c4e7a1f9b2d0`` | ``3.4.0``         | Add index on asset_event (asset_id, partition_key).          |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``c4e7a1f9b2d0``        | ``436dc127462c`` | ``3.4.0``         | Add index on asset.uri.                                      |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+

airflow-core/newsfragments/66854.significant.rst

@@ -0,0 +1,34 @@
+Asset events from successful tasks are registered asynchronously
+
+When a task succeeds, its asset events (and the resulting scheduling of downstream
+asset-triggered Dags) are no longer written synchronously inside the task state-update
+request. The execution API now records the events as a durable marker in a new
+``asset_event_queue`` table, committed in the same transaction as the task state, and the
+scheduler drains that queue and registers the events. This removes the asset registration
+work from the task-completion request path, which previously held the ``task_instance`` row
+lock while it ran and could pile up requests and exhaust the API server under high fan-out.
+
+As a result there can be a short delay between a task succeeding and its asset events being
+visible (typically one scheduler loop, longer under sustained load). A queued registration is
+retried up to ``[scheduler] asset_event_queue_max_attempts`` times; a row that still fails
+after that is parked in ``asset_event_queue`` and surfaced via the ``asset.event_queue.failures``
+metric rather than retried forever. ``dag.test`` runs without a scheduler, so it drains the queue
+itself after each task, keeping asset events visible during a test run. The drain cadence and
+batch size are tunable via ``[scheduler] asset_event_queue_drain_interval``,
+``asset_event_queue_batch_size`` and ``asset_event_queue_max_attempts``.
+
+* Types of change
+
+  * [ ] Dag changes
+  * [x] Config changes
+  * [ ] API changes
+  * [ ] CLI changes
+  * [x] Behaviour changes
+  * [ ] Plugin changes
+  * [ ] Dependency changes
+  * [ ] Code interface changes
+
+* Migration rules needed
+
+  * None - a new ``asset_event_queue`` table is created by the schema migration; no user
+    action is required.

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

@@ -40,6 +40,7 @@
 from sqlalchemy.sql import select
 from structlog.contextvars import bind_contextvars
 
+from airflow._shared.observability.metrics import stats
 from airflow._shared.observability.traces import override_ids
 from airflow._shared.state import TaskScope
 from airflow._shared.timezones import timezone
@@ -77,7 +78,7 @@
 )
 from airflow.configuration import conf
 from airflow.exceptions import InvalidPartitionKeyError, TaskNotFound
-from airflow.models.asset import AssetActive
+from airflow.models.asset import AssetActive, AssetEventQueue
 from airflow.models.base import ID_LEN
 from airflow.models.dag import DagModel
 from airflow.models.dagrun import DagRun as DR
@@ -97,6 +98,8 @@
 if TYPE_CHECKING:
     from sqlalchemy.sql.dml import Update
 
+    from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile
+
 router = VersionedAPIRouter()
 
 ti_id_router = VersionedAPIRouter(
@@ -499,6 +502,30 @@ def ti_update_state(
                 extra=json.dumps({"host_name": hostname}) if hostname else None,
             )
         )
+        # Durably record the successful task's asset events in the SAME transaction that commits the
+        # state, so the marker can never be lost in a crash between this commit and their
+        # registration. The scheduler drain is the single writer that registers them; the request
+        # never runs asset registration itself, which is what keeps the task_instance row lock short
+        # under high fan-out. Only a genuine RUNNING->SUCCESS transition reaches here (a duplicate
+        # SUCCESS->SUCCESS short-circuits earlier), so a completion is never enqueued twice.
+        if (
+            updated_state == TaskInstanceState.SUCCESS
+            and isinstance(ti_patch_payload, TISuccessStatePayload)
+            and (ti_patch_payload.task_outlets or ti_patch_payload.outlet_events)
+        ):
+            _enqueue_asset_events(
+                task_instance_id=task_instance_id,
+                dag_id=dag_id,
+                run_id=run_id,
+                task_id=task_id,
+                map_index=map_index,
+                task_outlets=ti_patch_payload.task_outlets,
+                outlet_events=ti_patch_payload.outlet_events,
+                session=session,
+            )
+        # Commit the state, log entry and durable queue marker together so they land atomically and
+        # the row lock is released promptly.
+        session.commit()
     except DataError:
         # Let DataErrorHandler return a 422 (not the opaque 500 below).
         raise
@@ -525,7 +552,9 @@ def ti_update_state(
                     task_id=task_id,
                     map_index=map_index,
                 )
+                session.commit()
             except Exception:
+                session.rollback()
                 log.warning(
                     "Failed to clear task state on success",
                     dag_id=dag_id,
@@ -622,6 +651,58 @@ def _validate_outlet_event_partition_keys(outlet_events: list[dict[str, Any]]) -
             )
 
 
+def _asset_event_payload(
+    task_outlets: list[AssetProfile],
+    outlet_events: list[dict[str, Any]],
+    ti_key: dict[str, Any],
+) -> dict[str, Any]:
+    """Serialize a task's outlets, outlet events, and natural key into the queue row's JSON payload."""
+    return {
+        "task_outlets": [outlet.model_dump(mode="json") for outlet in task_outlets],
+        "outlet_events": outlet_events,
+        # The scheduler drain resolves the live task instance by this natural key rather than by the
+        # surrogate ``ti_id``. Clearing a task reassigns its uuid7 id, and relying on the foreign
+        # key's ON UPDATE CASCADE to re-point the row only works on Postgres -- SQLite does not
+        # enforce foreign keys in Airflow's production engine -- so a natural-key lookup is what
+        # keeps the pending events reachable on every backend.
+        "ti_key": ti_key,
+    }
+
+
+def _enqueue_asset_events(
+    *,
+    task_instance_id: UUID,
+    dag_id: str,
+    run_id: str,
+    task_id: str,
+    map_index: int,
+    task_outlets: list[AssetProfile],
+    outlet_events: list[dict[str, Any]],
+    session: SessionDep,
+) -> None:
+    """
+    Add a durable queue marker for a completed task's asset events, without committing.
+
+    The caller commits this in the same transaction as the task-state update, so the marker and the
+    ``SUCCESS`` state land atomically: the events can never be lost in a crash between the state
+    commit and their registration. The scheduler drain is the single writer that later registers the
+    events and deletes the row. ``merge`` overwrites any pending row a retried completion left behind
+    rather than conflicting on the ``ti_id`` primary key.
+    """
+    session.merge(
+        AssetEventQueue(
+            ti_id=task_instance_id,
+            payload=_asset_event_payload(
+                task_outlets,
+                outlet_events,
+                {"dag_id": dag_id, "run_id": run_id, "task_id": task_id, "map_index": map_index},
+            ),
+            attempts=0,
+        )
+    )
+    stats.incr("asset.event_queue.enqueued")
+
+
 def _create_ti_state_update_query_and_update_state(
     *,
     ti_patch_payload: TIStateUpdate,
@@ -659,14 +740,6 @@ def _create_ti_state_update_query_and_update_state(
             # Store retry policy overrides so next_retry_datetime() can read them.
             # These are cleared when the task enters RUNNING (ti_run).
             query = query.values(retry_delay_override=retry_delay_override, retry_reason=retry_reason)
-        elif isinstance(ti_patch_payload, TISuccessStatePayload):
-            if ti is not None:
-                TI.register_asset_changes_in_db(
-                    ti,
-                    ti_patch_payload.task_outlets,
-                    ti_patch_payload.outlet_events,
-                    session=session,
-                )
         try:
             _emit_task_span(ti, state=updated_state)
         except Exception:

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

@@ -22,7 +22,7 @@
 from typing import TYPE_CHECKING
 
 import structlog
-from sqlalchemy import exc, or_, select
+from sqlalchemy import exc, insert, or_, select
 from sqlalchemy.orm import joinedload
 
 from airflow._shared.observability.metrics import stats
@@ -41,6 +41,7 @@
     DagScheduleAssetUriReference,
     PartitionedAssetKeyLog,
     TaskOutletAssetReference,
+    asset_alias_asset_event_association_table,
 )
 from airflow.models.log import Log
 from airflow.timetables.base import compute_rollup_fingerprint
@@ -359,8 +360,14 @@ def register_asset_change(
             ).unique()
 
             for asset_alias_model in asset_alias_models:
-                asset_alias_model.asset_events.append(asset_event)
-                session.add(asset_alias_model)
+                # Direct INSERT instead of asset_alias_model.asset_events.append(...), which
+                # would lazy-load the alias's entire asset_events collection just to add one row.
+                session.execute(
+                    insert(asset_alias_asset_event_association_table).values(
+                        alias_id=asset_alias_model.id,
+                        event_id=asset_event.id,
+                    )
+                )
 
                 dags_to_queue_from_asset_alias |= {
                     alias_ref.dag

airflow-core/src/airflow/config_templates/config.yml

@@ -2639,6 +2639,31 @@ scheduler:
       type: integer
       example: ~
       default: "60"
+    asset_event_queue_drain_interval:
+      description: |
+        How often (in seconds) the scheduler drains the asset event queue, registering asset
+        events that the execution API deferred to the queue (for example when the synchronous
+        write on task completion hit database lock contention).
+      version_added: 3.4.0
+      type: float
+      example: ~
+      default: "5.0"
+    asset_event_queue_batch_size:
+      description: |
+        The maximum number of pending asset event registrations the scheduler processes in a
+        single drain pass.
+      version_added: 3.4.0
+      type: integer
+      example: ~
+      default: "100"
+    asset_event_queue_max_attempts:
+      description: |
+        How many times the scheduler retries a failing asset event registration before parking
+        the row for inspection and emitting the ``asset.event_queue.failures`` metric.
+      version_added: 3.4.0
+      type: integer
+      example: ~
+      default: "5"
     pool_metrics_interval:
       description: |
         How often (in seconds) should pool usage stats be sent to StatsD (if statsd_on is enabled)

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -77,6 +77,7 @@
     AssetAliasModel,
     AssetDagRunQueue,
     AssetEvent,
+    AssetEventQueue,
     AssetModel,
     AssetPartitionDagRun,
     AssetWatcherModel,
@@ -85,6 +86,7 @@
     PartitionedAssetKeyLog,
     TaskInletAssetReference,
     TaskOutletAssetReference,
+    _register_queued_asset_event,
 )
 from airflow.models.asset_state_store import AssetStateStoreModel
 from airflow.models.backfill import Backfill, BackfillDagRun
@@ -335,6 +337,17 @@ def __init__(
         self._parallelism = conf.getint("core", "parallelism")
         self._multi_team = conf.getboolean("core", "multi_team")
         self._dag_tags_in_metrics = conf.getboolean("metrics", "dag_tags_in_metrics", fallback=False)
+        # Asset-event queue drain tuning. Read here (not per drain tick) so the drain loop stays free
+        # of conf lookups, which can hit the secret backend and break HA locks (see note above).
+        self._asset_event_queue_batch_size = conf.getint(
+            "scheduler", "asset_event_queue_batch_size", fallback=100
+        )
+        self._asset_event_queue_max_attempts = conf.getint(
+            "scheduler", "asset_event_queue_max_attempts", fallback=5
+        )
+        self._asset_event_queue_drain_interval = conf.getfloat(
+            "scheduler", "asset_event_queue_drain_interval", fallback=5.0
+        )
         self._max_partition_dag_runs_per_loop = MAX_PARTITION_DAG_RUNS_PER_LOOP
         self._dag_id_to_team_name: dict[str, str | None] = {}
 
@@ -1788,6 +1801,11 @@ def _run_scheduler_loop(self) -> None:
             self._remove_unreferenced_triggers,
         )
 
+        timers.call_regular_interval(
+            self._asset_event_queue_drain_interval,
+            self._drain_asset_event_queue,
+        )
+
         if any(x.is_local for x in self.executors):
             bundle_cleanup_mgr = BundleUsageTrackingManager()
             check_interval = conf.getint(
@@ -3698,6 +3716,66 @@ def _remove_unreferenced_triggers(self, *, session: Session = NEW_SESSION) -> No
             .execution_options(synchronize_session="fetch")
         )
 
+    @provide_session
+    def _drain_asset_event_queue(self, *, session: Session = NEW_SESSION) -> None:
+        """
+        Process pending asset-event registrations enqueued by the execution API.
+
+        The task-success path commits an :class:`~airflow.models.asset.AssetEventQueue` marker
+        atomically with the task state instead of registering the events inline, so the scheduler
+        is the sole writer that runs ``register_asset_changes_in_db`` in a live deployment. We claim
+        a batch of pending rows in a single ``FOR UPDATE SKIP LOCKED`` query (so multiple schedulers
+        do not collide), register each row's events and delete it in the same transaction so it is
+        processed exactly once. A row that keeps failing is retried on later passes until it exceeds
+        ``asset_event_queue_max_attempts``, after which it is left in place (parked) and surfaced via
+        the ``asset.event_queue.failures`` metric for an operator to inspect.
+        """
+        max_attempts = self._asset_event_queue_max_attempts
+        rows = session.scalars(
+            with_row_locks(
+                select(AssetEventQueue)
+                .where(AssetEventQueue.attempts < max_attempts)
+                .order_by(AssetEventQueue.created_at)
+                .limit(self._asset_event_queue_batch_size),
+                of=AssetEventQueue,
+                session=session,
+                skip_locked=True,
+            )
+        ).all()
+
+        processed = 0
+        for row in rows:
+            ti_id = row.ti_id
+            try:
+                # Each row is registered in its own savepoint so one poison row rolls back alone
+                # instead of discarding the whole batch; the outer transaction commits once at the end.
+                with session.begin_nested():
+                    _register_queued_asset_event(row, session=session)
+                processed += 1
+            except Exception:
+                with session.begin_nested():
+                    row.attempts += 1
+                if row.attempts >= max_attempts:
+                    self.log.error(
+                        "Giving up on asset-event registration for task instance %s after %s "
+                        "attempts; the row is left in asset_event_queue for inspection.",
+                        ti_id,
+                        row.attempts,
+                    )
+                    stats.incr("asset.event_queue.failures")
+                else:
+                    self.log.warning(
+                        "Asset-event registration for task instance %s failed; will retry (attempt %s).",
+                        ti_id,
+                        row.attempts,
+                    )
+
+        session.commit()
+        if processed:
+            stats.incr("asset.event_queue.processed", processed)
+        pending = session.scalar(select(func.count()).select_from(AssetEventQueue))
+        stats.gauge("asset.event_queue.pending", pending or 0)
+
     @provide_session
     def _update_asset_orphanage(self, *, session: Session = NEW_SESSION) -> None:
         """