Register asset events before taking the task_instance row lock

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

@@ -358,6 +358,33 @@ def ti_update_state(
     bind_contextvars(ti_id=str(task_instance_id))
     log.debug("Updating task instance state", new_state=ti_patch_payload.state)
 
+    # For a success payload, register outlet asset events *before* acquiring the task_instance
+    # row lock below, so the registration work (asset lookups, event inserts, dag-run queueing)
+    # never runs while holding that lock and concurrent completions do not pile up behind it.
+    # The locked SELECT below re-checks the state: if the transaction does not commit (duplicate
+    # or invalid transition), these writes roll back with it, so the registered events are only
+    # ever observable together with the committed SUCCESS state.
+    assets_registered = False
+    if isinstance(ti_patch_payload, TISuccessStatePayload) and (
+        ti_patch_payload.task_outlets or ti_patch_payload.outlet_events
+    ):
+        ti = session.get(TI, task_instance_id)
+        if ti is not None and ti.state == TaskInstanceState.RUNNING:
+            try:
+                _validate_outlet_event_partition_keys(ti_patch_payload.outlet_events)
+            except InvalidPartitionKeyError as e:
+                raise HTTPException(
+                    status_code=HTTP_422_UNPROCESSABLE_CONTENT,
+                    detail={"reason": "invalid_partition_key", "message": str(e)},
+                ) from e
+            TI.register_asset_changes_in_db(
+                ti,
+                ti_patch_payload.task_outlets,
+                ti_patch_payload.outlet_events,
+                session=session,
+            )
+            assets_registered = True
+
     old = (
         select(
             TI.state,
@@ -411,6 +438,10 @@ def ti_update_state(
     # SUCCESS or DEFERRED -> DEFERRED), including duplicates that would not pass the RUNNING
     # transition check below.
     if ti_patch_payload.state.value == previous_state:
+        if assets_registered:
+            # A concurrent completion won the race and committed its own registration;
+            # discard ours so the events are not recorded twice.
+            session.rollback()
         log.info(
             "Duplicate state update request received; state already set",
             requested_state=ti_patch_payload.state.value,
@@ -460,6 +491,7 @@ def ti_update_state(
             query=query,
             dag_id=dag_id,
             dag_bag=dag_bag,
+            assets_registered=assets_registered,
         )
     except DataError:
         # Let DataErrorHandler return a 422 instead of silently marking the TI FAILED below.
@@ -630,6 +662,7 @@ def _create_ti_state_update_query_and_update_state(
     session: SessionDep,
     dag_bag: DagBagDep,
     dag_id: str,
+    assets_registered: bool = False,
 ) -> tuple[Update, TaskInstanceState]:
     if isinstance(ti_patch_payload, (TITerminalStatePayload, TIRetryStatePayload, TISuccessStatePayload)):
         ti = session.get(TI, task_instance_id, with_for_update={"of": TI})
@@ -660,7 +693,11 @@ def _create_ti_state_update_query_and_update_state(
             # 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:
+            # Normally the events were registered before the row lock was taken (see
+            # ti_update_state). Falling in here means the state was not RUNNING at the
+            # unlocked pre-read but is now: a new try started in between. Register under
+            # the lock, preserving the previous behaviour for that race.
+            if ti is not None and not assets_registered:
                 TI.register_asset_changes_in_db(
                     ti,
                     ti_patch_payload.task_outlets,

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

@@ -1294,6 +1294,139 @@ def test_ti_update_state_to_success_with_asset_events(
         assert event[0].asset == AssetModel(name="my-task", uri="s3://bucket/my-task", extra={})
         assert event[0].extra == expected_extra
 
+    def test_ti_update_state_success_registers_assets_before_taking_ti_row_lock(
+        self, client, session, create_task_instance
+    ):
+        import sqlalchemy.event
+
+        asset = AssetModel(id=1, name="my-task", uri="s3://bucket/my-task", group="asset", extra={})
+        session.add_all([asset, AssetActive.for_asset(asset)])
+        ti = create_task_instance(
+            task_id="test_register_assets_before_ti_row_lock",
+            start_date=DEFAULT_START_DATE,
+            state=State.RUNNING,
+        )
+        session.commit()
+
+        statement_order: list[str] = []
+
+        def _before_execute(conn, clauseelement, multiparams, params, execution_options):
+            if getattr(clauseelement, "_for_update_arg", None) is not None:
+                statement_order.append("locked_select")
+            elif getattr(getattr(clauseelement, "table", None), "name", None) == "asset_event":
+                statement_order.append("asset_event_insert")
+
+        engine = session.get_bind()
+        sqlalchemy.event.listen(engine, "before_execute", _before_execute)
+        try:
+            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"}],
+                    "outlet_events": [],
+                },
+            )
+        finally:
+            sqlalchemy.event.remove(engine, "before_execute", _before_execute)
+
+        assert response.status_code == 204
+        assert "asset_event_insert" in statement_order
+        assert "locked_select" in statement_order
+        assert statement_order.index("asset_event_insert") < statement_order.index("locked_select")
+
+    @pytest.mark.parametrize(
+        ("concurrent_state", "expected_status"),
+        [
+            pytest.param(TaskInstanceState.SUCCESS, 200, id="duplicate-success"),
+            pytest.param(TaskInstanceState.FAILED, 409, id="failed-during-registration"),
+        ],
+    )
+    @mock.patch.object(TaskInstance, "register_asset_changes_in_db")
+    def test_ti_update_state_success_rolls_back_asset_events_when_state_changes_during_registration(
+        self,
+        mock_register,
+        client,
+        session,
+        create_task_instance,
+        concurrent_state,
+        expected_status,
+    ):
+        real_register = mock_register.get_original()[0]
+
+        asset = AssetModel(id=1, name="my-task", uri="s3://bucket/my-task", group="asset", extra={})
+        session.add_all([asset, AssetActive.for_asset(asset)])
+        ti = create_task_instance(
+            task_id="test_rollback_asset_events_on_concurrent_state_change",
+            start_date=DEFAULT_START_DATE,
+            state=State.RUNNING,
+        )
+        session.commit()
+
+        def _register_then_concurrent_transition(ti_arg, task_outlets, outlet_events, *, session):
+            real_register(ti_arg, task_outlets, outlet_events, session=session)
+            # Simulate a concurrent transition landing between the unlocked registration
+            # and the locked state re-check.
+            session.execute(
+                update(TaskInstance).where(TaskInstance.id == ti_arg.id).values(state=concurrent_state)
+            )
+
+        mock_register.side_effect = _register_then_concurrent_transition
+
+        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"}],
+                "outlet_events": [],
+            },
+        )
+
+        assert response.status_code == expected_status
+        session.expire_all()
+        assert session.scalars(select(AssetEvent)).all() == []
+
+    def test_ti_update_state_success_registers_assets_under_lock_when_pre_read_skips(
+        self, client, session, create_task_instance
+    ):
+        """A success whose unlocked pre-read cannot confirm RUNNING must still register its events."""
+        asset = AssetModel(id=1, name="my-task", uri="s3://bucket/my-task", group="asset", extra={})
+        session.add_all([asset, AssetActive.for_asset(asset)])
+        ti = create_task_instance(
+            task_id="test_fallback_registration_under_lock",
+            start_date=DEFAULT_START_DATE,
+            state=State.RUNNING,
+        )
+        session.commit()
+
+        real_get = Session.get
+        pre_read_skipped = False
+
+        def _none_on_first_ti_get(session_self, entity, *args, **kwargs):
+            nonlocal pre_read_skipped
+            if not pre_read_skipped and entity is TaskInstance:
+                pre_read_skipped = True
+                return None
+            return real_get(session_self, entity, *args, **kwargs)
+
+        with mock.patch.object(Session, "get", autospec=True, side_effect=_none_on_first_ti_get):
+            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"}],
+                    "outlet_events": [],
+                },
+            )
+
+        assert pre_read_skipped
+        assert response.status_code == 204
+        session.expire_all()
+        assert len(session.scalars(select(AssetEvent)).all()) == 1
+
     @pytest.mark.parametrize(
         ("outlet_events", "expected_extra"),
         [