hkc-8010 commented on code in PR #66854:
URL: https://github.com/apache/airflow/pull/66854#discussion_r3639431004
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -499,6 +507,26 @@ def ti_update_state(
extra=json.dumps({"host_name": hostname}) if hostname else
None,
)
)
+ # Durably record the successful task's asset events in the SAME
transaction that commits the
+ # state, so the marker can never be lost in a crash between this
commit and the fast-path
+ # registration below. Only a genuine RUNNING->SUCCESS transition
reaches here (a duplicate
+ # SUCCESS->SUCCESS short-circuits earlier), so a completion is never
enqueued twice.
+ if (
+ updated_state == TaskInstanceState.SUCCESS
+ and isinstance(ti_patch_payload, TISuccessStatePayload)
+ and (ti_patch_payload.task_outlets or
ti_patch_payload.outlet_events)
+ ):
+ _enqueue_asset_events(
+ task_instance_id=task_instance_id,
+ task_outlets=ti_patch_payload.task_outlets,
+ outlet_events=ti_patch_payload.outlet_events,
+ session=session,
+ )
+ enqueued_asset_events = True
+ # 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()
Review Comment:
Agreed, and done. The enqueue stays inside the state-update transaction and
commits atomically with the SUCCESS state, but the request no longer calls
`register_asset_changes_in_db` at all. The append to `asset_event_queue` is the
only asset work on the request path now, so the row lock is held only for the
state write plus one insert.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -622,6 +665,96 @@ def _validate_outlet_event_partition_keys(outlet_events:
list[dict[str, Any]]) -
)
+# Lock timeout (seconds) for the fast-path synchronous asset-event
registration. Kept short so a
+# task completion sheds to the durable queue quickly under asset-row
contention rather than holding
+# an API-server worker. This only bounds lock waits; an uncontended write does
not wait on locks and
+# so is unaffected by the timeout. The events are already durably queued
before this runs, so a
+# timeout drops nothing.
+_ASSET_SYNC_LOCK_TIMEOUT = 2
+
+
+def _asset_event_payload(
+ task_outlets: list[AssetProfile], outlet_events: list[dict[str, Any]]
+) -> dict[str, Any]:
+ """Serialize a task's outlets and outlet events into the queue row's JSON
payload."""
+ return {
+ "task_outlets": [outlet.model_dump(mode="json") for outlet in
task_outlets],
+ "outlet_events": outlet_events,
+ }
+
+
+def _enqueue_asset_events(
+ *,
+ task_instance_id: UUID,
+ task_outlets: list[AssetProfile],
+ outlet_events: list[dict[str, Any]],
+ session: SessionDep,
+) -> None:
+ """
+ Add a durable queue marker for a completed task's asset events, without
committing.
+
+ The caller commits this in the same transaction as the task-state update,
so the marker and the
+ ``SUCCESS`` state land atomically: a crash before the fast-path
registration below can never
+ leave the state committed with no record of the pending events. ``merge``
overwrites any pending
+ row a retried completion left behind rather than conflicting on the
``ti_id`` primary key.
+ """
+ session.merge(
+ AssetEventQueue(
+ ti_id=task_instance_id,
+ payload=_asset_event_payload(task_outlets, outlet_events),
+ attempts=0,
+ )
+ )
+ stats.incr("asset.event_queue.enqueued")
+
+
+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,
+ )
Review Comment:
Done. `_try_register_enqueued_asset_events` and the
`_ASSET_SYNC_LOCK_TIMEOUT` it used are gone, and the scheduler drain is the
only caller of `register_asset_changes_in_db`. The `AssetEventQueue` lifecycle
is now unambiguous: the execution API creates the row, the scheduler consumes
it. This also removes the inline lock-timeout path and the
`is_lock_not_available_error`/`with_db_lock_timeout` plumbing entirely.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -525,14 +553,29 @@ def ti_update_state(
task_id=task_id,
map_index=map_index,
)
+ session.commit()
except Exception:
+ session.rollback()
log.warning(
"Failed to clear task state on success",
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
)
+ # 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,
+ )
+
Review Comment:
Removed, along with the `enqueued_asset_events` flag it was gated on.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -525,14 +553,29 @@ def ti_update_state(
task_id=task_id,
map_index=map_index,
)
+ session.commit()
except Exception:
+ session.rollback()
Review Comment:
I removed the inline path's separate commit/rollback with it. Two commits
remain and I kept them on purpose:
1. The single `session.commit()` after the state update. That is the atomic
outbox boundary you asked for on the other thread: state, log entry and queue
marker land together, and the row lock is released right after. Deferring it to
the session dependency's end-of-request commit doesn't work because the
optional `clear_on_success` block runs between here and there and would fold
into the same transaction.
2. The `commit`/`rollback` inside the `clear_on_success` block. That block
is pre-existing best-effort cleanup: a failure there must not roll back an
already-durable success (that would fail a task that actually succeeded and
make the worker retry it). It rolls back only its own changes and leaves the
committed success intact, so I left it as is rather than couple cleanup failure
to task state.
Happy to revisit if you'd rather change the `clear_on_success` semantics,
but that felt out of scope for this PR.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py:
##########
@@ -1364,6 +1364,254 @@ 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_success_registers_assets_synchronously(
+ self, client, session, create_task_instance
+ ):
+ """On the default (sqlite) path a successful TI with outlets registers
assets inline: the
+ AssetEvent is written and no durable asset_event_queue row is left
behind."""
+ asset = AssetModel(id=1, name="my-task", uri="s3://bucket/my-task",
group="asset", extra={})
+ session.add_all([asset, AssetActive.for_asset(asset)])
Review Comment:
All three are gone (`_success_registers_assets_synchronously`,
`_asset_registration_failure_enqueues`, `_lock_contention_enqueues`), and there
is no sqlite-specific branch anywhere now. The success test is reframed as "a
successful task enqueues the marker; the drain registers later" and asserts the
payload and the natural key, with no `AssetEvent` created inline.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py:
##########
@@ -1364,6 +1364,254 @@ 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_success_registers_assets_synchronously(
+ self, client, session, create_task_instance
+ ):
+ """On the default (sqlite) path a successful TI with outlets registers
assets inline: the
+ AssetEvent is written and no durable asset_event_queue row is left
behind."""
+ 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_ti_update_state_success_registers_assets_synchronously",
+ start_date=DEFAULT_START_DATE,
+ state=State.RUNNING,
+ )
+ session.commit()
+
+ 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 == 204
+ assert response.text == ""
+ session.expire_all()
+
+ assert session.get(TaskInstance, ti.id).state == State.SUCCESS
+ events = session.scalars(select(AssetEvent)).all()
+ assert len(events) == 1
+ assert events[0].asset == AssetModel(name="my-task",
uri="s3://bucket/my-task", extra={})
+ # Synchronous path: nothing deferred to the durable queue.
+ assert session.scalars(select(AssetEventQueue)).all() == []
+
+ def test_ti_update_state_asset_registration_failure_enqueues(self, client,
session, create_task_instance):
+ """If synchronous asset registration raises, the TI success is still
durable (committed) and
+ the events are handed to the asset_event_queue for the scheduler to
retry."""
+ ti = create_task_instance(
+ task_id="test_ti_update_state_asset_registration_failure_enqueues",
+ start_date=DEFAULT_START_DATE,
+ state=State.RUNNING,
+ )
+ session.commit()
+
+ task_outlets = [{"name": "my-task", "uri": "s3://bucket/my-task",
"type": "Asset"}]
+ outlet_events = [
+ {"dest_asset_key": {"name": "my-task", "uri":
"s3://bucket/my-task"}, "extra": {"foo": 1}}
+ ]
+
+ with (
+ mock.patch(
+
"airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db",
+ side_effect=RuntimeError("boom"),
+ ),
+
mock.patch("airflow.api_fastapi.execution_api.routes.task_instances.stats") as
mock_stats,
+ ):
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/state",
+ json={
+ "state": "success",
+ "end_date": DEFAULT_END_DATE.isoformat(),
+ "task_outlets": task_outlets,
+ "outlet_events": outlet_events,
+ },
+ )
+
+ assert response.status_code == 204
+ assert response.text == ""
+ session.expire_all()
+
+ # TI success is durable even though asset registration failed.
+ assert session.get(TaskInstance, ti.id).state == State.SUCCESS
+
+ # Exactly one queue row, with the serialized outlets/events and the
right ti_id.
+ rows = session.scalars(select(AssetEventQueue)).all()
+ assert len(rows) == 1
+ assert rows[0].ti_id == ti.id
+ assert rows[0].payload == {"task_outlets": task_outlets,
"outlet_events": outlet_events}
+ assert rows[0].attempts == 0
+
+ # No AssetEvent was persisted (the write was rolled back), and enqueue
was metered.
+ assert session.scalars(select(AssetEvent)).all() == []
+ mock_stats.incr.assert_called_once_with("asset.event_queue.enqueued")
+
+ def test_ti_update_state_lock_contention_enqueues(self, client, session,
create_task_instance):
Review Comment:
Same here: `test_ti_update_state_asset_registration_failure_enqueues` and
`test_ti_update_state_lock_contention_enqueues` are both removed. They only
existed to cover the inline path's failure and lock-timeout fallbacks, which no
longer exist now that the request only appends the queue marker. Drain-side
retry and parking are covered by the scheduler tests instead.
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -1035,6 +1035,9 @@ def prepare_db_for_next_try(self, session: Session):
TaskInstanceHistory.record_ti(self, session=session)
session.execute(delete(TaskReschedule).filter_by(ti_id=self.id))
+ # Reassigning the id re-points any not-yet-drained asset_event_queue
rows via the FK's
+ # ON UPDATE CASCADE, so the scheduler drain still finds this TI
instead of dropping the
+ # pending asset events.
Review Comment:
Good catch, and my earlier `ON UPDATE CASCADE` answer was not actually
correct across backends. Airflow doesn't enable `PRAGMA foreign_keys` on sqlite
in the production engine (only inside migrations), so on sqlite the cascade
silently does nothing and the pending row keeps the old id. The manual `UPDATE`
you suggested is blocked on Postgres by the non-deferrable FK ordering. So
neither pure approach is right on both.
Instead I stopped depending on the id for lookup. The enqueue now stores the
task instance's natural key (`dag_id`, `run_id`, `task_id`, `map_index`) in the
payload, and the drain resolves the live task instance by that key rather than
by the surrogate `ti_id`. Clearing reassigns the id but not the natural key,
and the run is unchanged, so the drain still finds the current attempt and its
DagRun partition context. The regression test
(`test_clearing_ti_still_registers_via_natural_key`) runs on the default sqlite
engine, exactly the condition the cascade approach failed under, and asserts
the events still register after a clear.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
