Reduce task success asset registration lock contention

[hkc-8010]

Closes #66853

Summary

Under high fan-out (for example a dbt/Cosmos DAG where many tasks emit assets at once), ti_update_state committed the task-instance state and then registered asset events while still holding the task_instance row lock. Holding that lock across the registration work (scanning scheduled DAGs and inserting AssetEvent/AssetDagRunQueue rows) piled up concurrent completions and OOMKilled the API server.

This reworks the task-success path so asset registration never holds the task_instance lock and a failed registration is never lost:

• Commit the task-instance state first, releasing the row lock.
• Then register asset events opportunistically, under a short lock timeout. Uncontended writes (the common case) still happen inline and are unaffected by the timeout.
• If the write cannot acquire its locks in time, or fails for any reason, the events are written to a new asset_event_queue table. The scheduler drains that queue and retries until the events are registered, so nothing is silently dropped.
• Replace the alias-event ORM append with a direct association-table insert, so a long-lived alias does not lazy-load its entire asset_events collection to add one row.

Asset-event creation is asynchronous when it cannot complete inline: there is a short, eventually consistent delay (usually one scheduler loop) before a deferred event and its downstream asset-triggered DAG run appear. This is documented and called out in a significant-change newsfragment. It replaces the earlier version of this PR, which returned success and only logged a dropped registration; the queue makes delivery at-least-once instead.

Changes

• airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py: commit state before registration; two-phase register-or-enqueue.
• airflow-core/src/airflow/models/asset.py and migration 0127_3_4_0_add_asset_event_queue.py: new asset_event_queue table.
• airflow-core/src/airflow/jobs/scheduler_job_runner.py: _drain_asset_event_queue sweep (claims with SKIP LOCKED, retries until success, parks a row after asset_event_queue_max_attempts).
• airflow-core/src/airflow/assets/manager.py: direct alias-event association insert.
• airflow-core/src/airflow/config_templates/config.yml: [scheduler] asset_event_queue_drain_interval, asset_event_queue_batch_size, asset_event_queue_max_attempts.
• Metrics asset.event_queue.enqueued / processed / failures / pending; docs note in asset-scheduling.rst; significant newsfragment.

Validation

• prek (ruff, ruff-format, mypy-airflow-core, newsfragment and migration checks) on the changed files.
• New unit tests for the endpoint (sync register, enqueue on failure, enqueue on lock contention, forced-failure skip, no-outlets short-circuit) and the scheduler drain (register and delete, retry then park, missing-TI purge, empty no-op).
• breeze testing core-tests on Python 3.10 with both --backend sqlite and --backend postgres.
• Alembic upgrade/downgrade round-trip for the new table.

PR Checklist

• My PR is targeted at the main branch
• Tests added
• Docs updated (asset-scheduling.rst) and significant newsfragment added
• prek passes on changed files
• Breeze core-tests pass (sqlite + postgres, Python 3.10)

---

Developed with AI assistance; I reviewed the design, code, and tests before submitting.

[hkc-8010]

The failing CI job (provider distributions tests / Compat 3.0.6:P3.10) is pre-existing on main and unrelated to this PR. It fails with:

ImportError: cannot import name 'Options' from 'jwt.types'

This is a flask_jwt_extended/PyJWT version incompatibility in providers/fab tests running against Airflow 3.0.6. Confirmed failing on main in run https://github.com/apache/airflow/actions/runs/25789005777 before this PR was opened.

All other CI checks pass.

[hkc-8010]

The CI failure in Integration and System Tests / Integration core otel (test_export_legacy_metric_names) is unrelated to this PR — it tests scheduler-side metric emission timing and has a history of flakiness (see #61070, #65867). This PR touches only task_instances.py and assets/manager.py; no OTEL or scheduler metric code is modified.

[potiuk]

@hkc-8010 Converting to draft — this PR doesn't yet meet our Pull Request quality criteria.

• ❌ Pre-commit / static checks. See docs.
• ❌ Unresolved review comments: 2 thread(s). See docs.

See the linked criteria for how to fix each item, then mark the PR "Ready for review". This is not a rejection — just an invitation to bring the PR up to standard. No rush.

---

Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you.

---

Drafted-by: Claude Code (Opus 4.7); reviewed by @potiuk before posting

[uranusjr]

Is it a good idea to make the endpoint succeed when lock contention happens? I feel it’s a consistency issue. We should either do retry in the endpoint, or report back to the client there’s something wrong.

[hkc-8010]

Retry/reconciliation for this path is intentionally out of scope for this PR (called out in the description) — the goal here is resolving the task_instance row lock contention causing idle-in-transaction pile-ups and apiserver OOMKills, not building a full delivery guarantee for asset events.

On why it still returns success: by the time register_asset_changes_in_db runs, the TI state is already committed and durable. Returning 500 here would make the task-sdk worker retry a state update for a task that already succeeded, which is a worse failure mode than a dropped asset event. The failure is now observable via the asset.registration_failures metric and a log line, so an operator can catch and manually reconcile if it happens.

Agreed it's not fully consistent. Happy to file a follow-up issue to track retry/reconciliation for dropped asset events as separate work if that addresses the concern for this PR.

[hkc-8010]

Renamed the PR title, check_no_conventional_commit_message was failing on the Conventional-Commits-style title.

Also looked at the Integration and System Tests / Integration core otel failure (test_dag_execution_succeeds[detail_spans], span-hierarchy assertion on sub_span1/_execute_task). This is unrelated to this PR: test_otel.py isn't touched by this diff, and git blame shows the failing expected-hierarchy block was added by a separate commit (2e9cbc25ebbe, "Add configurable task span detail level for OTel tracing") that landed on main independently. The test itself notes the gap ("Additional detail spans are deferred to follow-up PRs"), so this looks like a pre-existing baseline failure rather than something introduced here.

[uranusjr]

How do you plan to do the followup? The response is indistinguishable to when the asset registration is successful. You response is too hand-wavy to the inconsistency. A concrete plan should be layed out so the behavior is not just silently broken until your next submission is merged. This PR is simply not an improvement on its own.

[Andrushika]

Since this PR and the discussion mainly revolve around the early-commit approach and the resulting consistency trade-offs, I was wondering whether there might be another approach that avoids the early commit.

To be more specific, would it be possible to keep the TI state update and asset registration in the same transaction, but defer acquiring the TI row lock until immediately before the ti.state update?

My current understanding of the original flow is:

ti_update_state()
→ SELECT the TI row FOR UPDATE (lock acquire)
→ perform asset registration (long job)
→ update ti.state
→ commit and release the lock

This seems to mean that when asset registration takes a long time, the TI row lock is held for that entire period, causing later requests for the same TI to wait on SELECT ... FOR UPDATE.

I was wondering whether the ordering could instead be:

ti_update_state()
→ perform asset registration (long job)
→ acquire the TI row lock and re-check the state
→ update ti.state
→ commit both changes together

If the state is no longer valid at that point, the transaction could be rolled back together with the asset changes.

Would this potentially shorten the TI lock window while preserving atomicity between the state update and asset events? (Since #66853 identifies the long-held TI lock and waiter pile-up as the OOM path)
I have not audited the other TI writers or the lock-ordering implications, so there may be a constraint I have overlooked.

[hkc-8010]

@uranusjr @Andrushika you're right, and I've stopped defending the current shape. Returning 204 while a failed register_asset_changes_in_db is only logged and counted means a dropped asset event silently never fires its downstream DAG, and "follow up later" isn't good enough. As it stands the PR isn't an improvement on its own. Before I push more code I want to agree on the shape.

Proposal: a durable transactional outbox plus a scheduler reconciliation sweep, so a dropped registration is guaranteed to be delivered eventually rather than lost.

1. Durable marker committed atomically with the TI state. On the success path, alongside the existing early session.commit(), I insert one row into a new asset_registration_outbox table capturing ti_id and the (already-validated) task_outlets and outlet_events. Because it commits in the same transaction as the state update, the 204 is only returned once the pending-registration record is durable, so even a crash between the state commit and registration can't lose the event. That is the "not silently broken" guarantee.
2. Inline registration stays post-commit, as an optimization. The expensive work (scheduled-dag scan, AssetDagRunQueue inserts, alias association) still runs after the lock is released, so the API server OOMKill: task_instance row lock held during asset event emission under high concurrency #66853 contention fix is preserved. On success it deletes the outbox row in the same transaction as the asset writes, which keeps it exactly-once against that row; on failure it rolls the asset writes back and leaves the outbox row for the scheduler. The outbox insert itself touches no hot rows and runs no queries, so it doesn't reintroduce the task_instance lock contention.
3. Scheduler reconciles leftovers. A recurring sweep (claiming rows with FOR UPDATE SKIP LOCKED, HA-safe, modeled on the existing heartbeat/zombie purge) replays any leftover outbox rows through register_asset_changes_in_db and deletes them, so the events land as AssetEvent and AssetDagRunQueue rows and _create_dag_runs_asset_triggered picks them up. A stats.gauge exposes the pending backlog so operators can see if reconciliation falls behind.

@Andrushika thanks, your reorder (register before taking the TI lock, then commit state and assets together) is a clean way to shrink the lock window and I considered it seriously. I leaned toward the outbox for two reasons. The reorder still needs a full audit of every task_instance writer for lock-ordering and deadlock risk before I'd trust it, and it still loses the event if the process dies between the asset work and the commit. The outbox gives an explicit, testable at-least-once guarantee instead. I'm happy to go the reorder route if you'd both prefer it; I don't have a strong attachment beyond wanting the guarantee to be verifiable.

Two questions before I implement:

• Does this shape address the consistency concern for you?
• Should the outbox and reconciliation land in this PR (it adds a table, a migration, and a scheduler timer), or as a stacked PR that merges together with this one? I'd rather not split it in a way that leaves a window where the drop is unhandled on main.

[hkc-8010]

@uranusjr I reworked this along the lines we discussed. The set-state endpoint no longer registers asset events under the task_instance lock and no longer drops them on failure.

On success it now commits the TI state first (releasing the row lock), then tries the asset write once under a short lock timeout. If it cannot get the locks in time, or fails for any reason, it writes the events to a new asset_event_queue table. The scheduler drains that queue and retries until the events land, parking a row only after a bounded number of failures (with a metric). So the common case still registers inline, contention sheds to the queue, and nothing is silently dropped.

I also kept the alias direct-INSERT fix, documented the enqueue-to-registration window in the asset-scheduling docs, and added a significant newsfragment. Validation: new endpoint and scheduler unit tests, plus breeze core-tests on Python 3.10 against both sqlite and postgres.

Draining the TI state update itself through the same queue is left out, as we agreed. Would appreciate another look when you get a chance.

[Andrushika]

it still loses the event if the process dies between the asset work and the commit.

No, I think this part may be a misunderstanding.
The one possible solution I mentioned simply changes the lock-acquisition sequence while preserving the original pattern of a DB transaction. Asset registration and state updates are atomic here.

[Andrushika]

To make the comparison concrete, I prototyped it: #70078

It registers the events before taking the task_instance row lock, still within the same transaction, then rechecks the state under the lock and rolls back if it loses the race. About 40 lines of code plus tests. Test results and Postgres numbers are in the draft PR.

I opened it as a draft for comparison. Happy to move it forward if this direction is useful, or drop it if the queue direction is preferred :)

[jason810496]

It's a interesting direction to solve this critical issue, thanks!
#protm when this get merged.

[jason810496]

Wow this direction looks really clean!
But I traced the code, and I have one concern: if we skip the inline registration, registration only happens in _drain_asset_event_queue, which runs on the scheduler's main loop thread (am I reading this right?).

Yes, you're correct!

So all the registration work would move onto the scheduler loop, and registration can be a heavy job. Every event would also wait for the next drain tick, even when there is no contention.
Find another place to run the registration jobs? But I don't know where is appropriate.

This is the point: "registration can be a heavy job" so it would better not run within this API server hot path. The batch processing on scheduler side is a ideal location for me, mind that we will execute user defined "Airflow Listener" on asset change.

[jason810496]

Suggested change

	# Fast-path: now that the state and durable queue marker are committed and the row lock is
	# released, try to register the asset events synchronously and delete the marker. Under
	# contention (or any failure) the marker stays and the scheduler drain registers them later.
	# The isinstance check is redundant with ``enqueued_asset_events`` at runtime but narrows the
	# payload type for the attribute access below.
	if enqueued_asset_events and isinstance(ti_patch_payload, TISuccessStatePayload):
	_try_register_enqueued_asset_events(
	task_instance_id=task_instance_id,
	task_outlets=ti_patch_payload.task_outlets,
	outlet_events=ti_patch_payload.outlet_events,
	session=session,
	)

[jason810496]

I'd prefer to remove _try_register_enqueued_asset_events at all to make scheduler as the single entrypoint for calling the TI.register_asset_changes_in_db. So that the lifecycle of AssetEventQueue and Asset Events will be more clear.

Suggested change

	def _try_register_enqueued_asset_events(
	*,
	task_instance_id: UUID,
	task_outlets: list[AssetProfile],
	outlet_events: list[dict[str, Any]],
	session: SessionDep,
	) -> None:
	"""
	Fast-path synchronous registration of a completed task's already-enqueued asset events.
	Runs after the task-instance state (and the durable queue marker) have been committed and the
	row lock released, so this is only an optimization. On success it registers the events and
	deletes the queue row in one transaction; on lock contention or any error it rolls back and
	leaves the row for the scheduler drain. The task_instance row lock is never held for the
	registration work, and because the marker is already durable the events are registered exactly
	once and never dropped.
	"""
	ti = session.get(TI, task_instance_id)
	if ti is None:
	return
	try:
	with with_db_lock_timeout(session, lock_timeout=_ASSET_SYNC_LOCK_TIMEOUT):
	TI.register_asset_changes_in_db(ti, task_outlets, outlet_events, session=session)
	session.execute(delete(AssetEventQueue).where(AssetEventQueue.ti_id == task_instance_id))
	session.commit()
	stats.incr("asset.event_queue.registered_inline")
	return
	except OperationalError as e:
	session.rollback()
	if not is_lock_not_available_error(e):
	log.warning(
	"Synchronous asset registration failed with a database error; leaving the event on "
	"the queue for the scheduler to retry.",
	task_instance_id=str(task_instance_id),
	exc_info=True,
	)
	except Exception:
	session.rollback()
	log.warning(
	"Synchronous asset registration failed; leaving the event on the queue for the "
	"scheduler to retry.",
	task_instance_id=str(task_instance_id),
	exc_info=True,
	)

[jason810496]

It's fine to call _enqueue_asset_events in the single transaction (and we should for the atomicity) if we don't call the TI.register_asset_changes_in_db. Since adding new AssetEventQueue record take way less time.

Suggested change

	# Commit now to release the task_instance row lock before the asset registration runs, so
	# concurrent task completions don't pile up waiting on it. The state, log entry and queue
	# marker all commit together here.
	session.commit()

[jason810496]

Additionally, we should remove all the explicit commit and rollback in the route to make this atomic at session injection level.

Suggested change

	session.commit()
	except Exception:
	session.rollback()
	except Exception:

[jason810496]

IIUC, We didn't add the actual statement to point the AssetEventQueue row to the new generated id here.

[jason810496]

Then we should remove this test case after removing the _try_register_enqueued_asset_events call. I don't feel it worthwhile to add a specialize archciture for sqlite path. Having a single writer to register the asset changes should be better.

[jason810496]

Ditto. We can remove test_ti_update_state_asset_registration_failure_enqueues and test_ti_update_state_lock_contention_enqueues.

[jason810496]

Actually, we can make the test like the following to exercise more:

for cur_attempt in range(max_attempts+1):
    # the row.attempts will be same as cur_attempt+1 for till max_attempts
    # then row.attempts will stay as max_attempts and will not over max_attempts.

[jason810496]

If we really want to keep the fast path, we should create a nested session and don't explicit commit or rollback at route handler level. We should defer to DB level timeout under nested session, and let the SessionDep level session injection commit to ensure the correctness.

Somehow like (but definitely need more refinement)

def _try_register_inline_or_enqueue(
    *,
    session: Session,
    ti: TaskInstance,
    task_outlets: list[AssetProfile],
    outlet_events: list[dict[str, Any]],
) -> None:
    """
    Attempt asset-event registration inside the state transaction; enqueue on failure.

    Runs under the caller's transaction — the one that holds the task_instance FOR UPDATE
    lock and will be committed by the SessionDep teardown. The savepoint isolates the
    attempt so a failure rolls back the partial registration without losing the state
    change; the marker enqueued on failure commits atomically with the state.

    BOUNDS — read carefully:
      * lock_timeout / innodb_lock_wait_timeout bound each individual lock WAIT only.
      * statement_timeout (PG) bounds each individual statement only.
      * Nothing here bounds TOTAL attempt duration; a many-statement registration
        (large outlet fan-out) holds the task_instance row lock for its full runtime.
        This is the pre-PR contention shape; prefer commit-first if that matters.
    """
    dialect_name = get_dialect_name(session)
    old_mysql_timeout = None
    try:
        with session.begin_nested():  # SAVEPOINT; ctx mgr emits ROLLBACK TO on error
            if dialect_name == "postgresql":
                # SET LOCAL is transactional: rolled back with the savepoint on failure.
                # On success (RELEASE) it persists to end of transaction — reset below.
                session.execute(text(f"SET LOCAL lock_timeout = '{_ASSET_SYNC_LOCK_TIMEOUT}s'"))
                session.execute(text(f"SET LOCAL statement_timeout = '{_ASSET_SYNC_LOCK_TIMEOUT}s'"))
            elif dialect_name == "mysql":
                # Session-scoped regardless of savepoints; restored in finally.
                old_mysql_timeout = session.execute(
                    text("SELECT @@SESSION.innodb_lock_wait_timeout")
                ).scalar()
                session.execute(
                    text(f"SET SESSION innodb_lock_wait_timeout = {_ASSET_SYNC_LOCK_TIMEOUT}")
                )
            TaskInstance.register_asset_changes_in_db(
                ti, task_outlets, outlet_events, session=session
            )
        # Savepoint released: the SET LOCALs would otherwise apply to the rest of the
        # request transaction (including the teardown commit path).
        if dialect_name == "postgresql":
            session.execute(text("SET LOCAL lock_timeout = DEFAULT"))
            session.execute(text("SET LOCAL statement_timeout = DEFAULT"))
        stats.incr("asset.event_queue.registered_inline")
        return
    except OperationalError as e:
        # ROLLBACK TO SAVEPOINT has already run; the outer transaction (state change)
        # is intact and the session is usable without a full rollback.
        if not is_lock_not_available_error(e):
            log.warning(
                "Inline asset registration failed with a database error; enqueueing.",
                task_instance_id=str(ti.id),
                exc_info=True,
            )
    except Exception:
        log.warning(
            "Inline asset registration failed; enqueueing.",
            task_instance_id=str(ti.id),
            exc_info=True,
        )
    finally:
        if old_mysql_timeout is not None:
            # Safe here: any failure above was contained by the savepoint rollback.
            session.execute(
                text(f"SET SESSION innodb_lock_wait_timeout = {old_mysql_timeout}")
            )

    # Fallback: marker in the SAME transaction as the state — they commit together at
    # SessionDep teardown, so at-least-once holds (a crash rolls back both; the client
    # retries the state transition). merge, not add: overwrite a cleared try's leftover.
    session.merge(
        AssetEventQueue(
            ti_id=ti.id,
            payload=_asset_event_payload(task_outlets, outlet_events),
            attempts=0,
        )
    )
    stats.incr("asset.event_queue.enqueued")

[hkc-8010]

@uranusjr following up on the queue design we landed on in Slack. After @jason810496's review I simplified this to be queue-only: the execution API commits a durable asset_event_queue marker in the same transaction as the task-success state, and the scheduler drain is now the single writer that runs register_asset_changes_in_db. That drops the inline registration fast-path I had added on the request path, so there is no asset registration inside the state-update request at all now, which is the cleanest form of the scheduler-drained queue we discussed.

The tradeoff is that a task's asset events (and the downstream asset-triggered runs) become visible after a drain interval instead of immediately, and dag.test does not register them at all since it has no scheduler. Both are documented in the newsfragment and the asset-scheduling note.

Wanted to confirm you are good dropping the inline path before I treat the direction as settled. @jason810496 thanks for pushing on this.