From 48a7ed44a306a82772dad8dfc9eeb69d3f1e08a4 Mon Sep 17 00:00:00 2001 From: Andrew Chang Date: Sat, 18 Jul 2026 01:07:55 +0800 Subject: [PATCH] Register asset events before taking the task_instance row lock Under high fan-out, ti_update_state held the task_instance row lock for the whole asset-event registration (asset lookups, event inserts, dag-run queueing). Concurrent completions piled up on the lock, each occupying an API-server thread, until the server was OOMKilled (#66853). Register the events before acquiring the lock instead, still inside the same transaction. The locked SELECT then re-checks the state: on the normal path the state flip and the events commit atomically; if a duplicate or a concurrent failure wins the race, the transaction rolls back and the events are discarded with it. Nothing is ever registered without its SUCCESS state, so no queue or reconciliation is needed. --- .../execution_api/routes/task_instances.py | 39 ++++- .../versions/head/test_task_instances.py | 133 ++++++++++++++++++ 2 files changed, 171 insertions(+), 1 deletion(-) 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"), [