diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index c1bac7960234d..ad8918241567c 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/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, diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 542ce7eaaf15b..5499e641e776f 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/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"), [