Fix API server OOMKill: release row lock before asset event emission under high concurrency

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

@@ -454,11 +454,26 @@ def ti_update_state(
                 extra=json.dumps({"host_name": hostname}) if hostname else None,
             )
         )
+        session.commit()
     except SQLAlchemyError as e:
         log.error("Error updating Task Instance state", error=str(e))
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Database error occurred"
         )
+
+    if isinstance(ti_patch_payload, TISuccessStatePayload):
+        try:
+            ti = session.get(TI, task_instance_id)
+            if ti is not None:
+                TI.register_asset_changes_in_db(
+                    ti,
+                    ti_patch_payload.task_outlets,
+                    ti_patch_payload.outlet_events,
+                    session,
+                )
+                session.commit()
+        except Exception:
+            log.exception("Error registering asset changes. The task state is already saved as SUCCESS.")
 
 
 def _emit_task_span(ti, state):
@@ -547,14 +562,6 @@ def _create_ti_state_update_query_and_update_state(
                 retry_delay_override=ti_patch_payload.retry_delay_seconds,
                 retry_reason=(ti_patch_payload.retry_reason[:500] if ti_patch_payload.retry_reason else None),
             )
-        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,
-                )
         _emit_task_span(ti, state=updated_state)
     elif isinstance(ti_patch_payload, TIDeferredStatePayload):
         # Calculate timeout if it was passed

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

@@ -322,10 +322,17 @@ 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)
+            if asset_alias_models:
+                from sqlalchemy import insert
+                from airflow.models.asset import asset_alias_asset_event_association_table
+
+                session.execute(
+                    insert(asset_alias_asset_event_association_table).values(
+                        [{"alias_id": aam.id, "event_id": asset_event.id} for aam in asset_alias_models]
+                    )
+                )
 
+            for asset_alias_model in asset_alias_models:
                 dags_to_queue_from_asset_alias |= {
                     alias_ref.dag
                     for alias_ref in asset_alias_model.scheduled_dags

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py

@@ -1328,6 +1328,35 @@ def test_ti_update_state_to_success_with_asset_alias_events(
             assert events[0].asset == AssetModel(name="my-task", uri="s3://bucket/my-task", extra={})
             assert events[0].extra == expected_extra
 
+    def test_ti_update_state_to_success_asset_registration_failure_returns_204(
+        self, client, session, create_task_instance
+    ):
+        ti = create_task_instance(
+            task_id="test_ti_update_state_to_success_asset_registration_failure_returns_204",
+            start_date=DEFAULT_START_DATE,
+            state=State.RUNNING,
+        )
+        session.commit()
+
+        with mock.patch("airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db") as mock_register:
+            mock_register.side_effect = Exception("Asset registration failed")
+
+            response = client.patch(
+                f"/execution/task-instances/{ti.id}/state",
+                json={
+                    "state": "success",
+                    "end_date": DEFAULT_END_DATE.isoformat(),
+                    "task_outlets": [{"name": "my-task", "uri": "s3://bucket/my-task", "type": "Asset"}],
+                },
+            )
+
+        assert response.status_code == 204
+        assert response.text == ""
+
+        session.expire_all()
+        ti = session.get(TaskInstance, ti.id)
+        assert ti.state == State.SUCCESS
+
     def test_ti_update_state_not_found(self, client, session):
         """
         Test that a 404 error is returned when the Task Instance does not exist.

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

@@ -162,6 +162,86 @@ def test_register_asset_change_with_alias(
         )
         assert session.scalar(select(func.count()).select_from(AssetDagRunQueue)) == 2
 
+    @pytest.mark.usefixtures("clear_assets")
+    def test_register_asset_change_with_alias_no_lazy_load(
+        self, session, dag_maker, mock_task_instance, testing_dag_bundle
+    ):
+        bundle_name = "testing"
+
+        consumer_dag_1 = DagModel(
+            dag_id="conumser_1", bundle_name=bundle_name, is_stale=False, fileloc="dag1.py"
+        )
+        session.add(consumer_dag_1)
+
+        asm = AssetModel(uri="test://asset1/", name="test_asset_uri", group="asset")
+        session.add(asm)
+
+        asam = AssetAliasModel(name="test_alias_name", group="test")
+        session.add(asam)
+        asam.scheduled_dags = [
+            DagScheduleAssetAliasReference(alias_id=asam.id, dag_id=consumer_dag_1.dag_id)
+        ]
+        session.execute(delete(AssetDagRunQueue))
+        session.flush()
+
+        # Add a bunch of preexisting asset events for the alias
+        from airflow.models.asset import AssetEvent, asset_alias_asset_event_association_table
+        from sqlalchemy import insert
+
+        pre_existing_events = []
+        for _ in range(5):
+            event = AssetEvent(asset_id=asm.id)
+            session.add(event)
+            pre_existing_events.append(event)
+        session.flush()
+
+        session.execute(
+            insert(asset_alias_asset_event_association_table).values(
+                [{"alias_id": asam.id, "event_id": ev.id} for ev in pre_existing_events]
+            )
+        )
+        session.commit()
+
+        asset = Asset(uri="test://asset1", name="test_asset_uri")
+        asset_manager = AssetManager()
+
+        # Test that no queries are run against asset_alias_asset_event during registration
+        with mock.patch("airflow.assets.manager.Session.scalars") as mock_scalars:
+            # We don't want to actually mock the scalars call since it would break the normal flow
+            # Instead we use the SQLAlchemy event listener to count queries to asset_alias_asset_event
+            pass
+
+        from sqlalchemy import event
+
+        query_count = 0
+        def before_cursor_execute(conn, cursor, statement, parameters, context, executemany):
+            nonlocal query_count
+            if "asset_alias_asset_event" in statement and "SELECT" in statement.upper():
+                query_count += 1
+
+        event.listen(session.bind, "before_cursor_execute", before_cursor_execute)
+
+        asset_manager.register_asset_change(
+            task_instance=mock_task_instance,
+            asset=asset,
+            source_alias_names=["test_alias_name"],
+            session=session,
+        )
+        session.flush()
+
+        event.remove(session.bind, "before_cursor_execute", before_cursor_execute)
+
+        # Ensure no SELECTs were made to the association table (no lazy loading)
+        assert query_count == 0
+
+        # Ensure we've created the new asset event and it is associated
+        # Total events should be 5 pre-existing + 1 new = 6
+        assert (
+            session.scalar(select(func.count()).select_from(AssetEvent).where(AssetEvent.asset_id == asm.id))
+            == 6
+        )
+        assert session.scalar(select(func.count()).select_from(AssetDagRunQueue)) == 1
+
     def test_register_asset_change_no_downstreams(self, session, mock_task_instance):
         asset_manager = AssetManager()