hkc-8010 commented on code in PR #66854:
URL: https://github.com/apache/airflow/pull/66854#discussion_r3619807152
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -525,14 +535,28 @@ 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,
)
+ # Register asset events now that the TI state is committed and the row
lock released, so
+ # the registration work never holds the task_instance lock. Only a
genuine RUNNING->SUCCESS
+ # transition reaches here (a duplicate SUCCESS->SUCCESS short-circuits
earlier), so a
+ # completion is never registered twice.
+ if isinstance(ti_patch_payload, TISuccessStatePayload):
+ _register_asset_changes_or_enqueue(
Review Comment:
Closed by the outbox rework. The queue marker is now committed in the same
transaction as the SUCCESS state, before the row lock is released, so the death
window you described can no longer drop the event. The worker retry seeing
SUCCESS is fine because the pending events live in the queue. Added a test that
neutralizes the inline path and asserts the marker survives.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -499,6 +506,9 @@ def ti_update_state(
extra=json.dumps({"host_name": hostname}) if hostname else
None,
)
)
+ # Commit now to release the task_instance row lock before asset
registration runs, so
+ # concurrent task completions don't pile up waiting on it.
+ session.commit()
Review Comment:
You're right, that was the real gap. Reworked it into a proper transactional
outbox: the `asset_event_queue` marker is written in the same transaction as
the state update and commits with it, before the row lock is released. The
synchronous registration is now a fast-path that runs after that commit and
deletes the marker on success. If the process dies in between, the marker is
already durable and the scheduler drain finishes it. Regression test:
`test_ti_update_state_success_marker_committed_atomically`.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -525,14 +535,28 @@ def ti_update_state(
task_id=task_id,
map_index=map_index,
)
+ session.commit()
except Exception:
+ session.rollback()
log.warning(
Review Comment:
Done. The enqueue now happens inside the state-update transaction and
commits atomically with the SUCCESS state. The append-only `asset_event_queue`
write keeps the lock time short compared to running
`register_asset_changes_in_db` under the request.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -622,6 +646,95 @@ def _validate_outlet_event_partition_keys(outlet_events:
list[dict[str, Any]]) -
)
+# Lock timeout (seconds) for the opportunistic synchronous asset-event write.
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.
+_ASSET_SYNC_LOCK_TIMEOUT = 2
+
+
+def _register_asset_changes_or_enqueue(
Review Comment:
I looked hard at going queue-only and decided to keep the inline fast-path.
Two reasons. First, this was the shape agreed with the maintainers:
opportunistic inline write, queue as the fallback, so the common uncontended
completion stays synchronous. Second, as @Andrushika noted here, queue-only
moves all registration onto the scheduler loop and adds a drain interval of
latency to every asset event even without contention. With the outbox rework
the durability no longer depends on the inline write (the marker is already
committed), so the inline path is a pure optimization with no consistency cost.
Happy to revisit if you still prefer queue-only.
##########
airflow-core/src/airflow/models/asset.py:
##########
@@ -788,6 +789,45 @@ def __repr__(self):
return f"{self.__class__.__name__}({', '.join(args)})"
+class AssetEventQueue(Base):
+ """
+ Pending asset-event registrations awaiting processing by the scheduler.
+
+ A row is enqueued on the task-success path in the execution API when
synchronous
+ asset-event registration is skipped or fails (for example under database
lock
+ contention). The scheduler drains this table, (re)runs
+
:meth:`~airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db`
to
+ create the ``AssetEvent`` and ``AssetDagRunQueue`` rows, and deletes the
queue row
+ once that write commits. This keeps the long-running asset-registration
work off the
+ synchronous ``ti_update_state`` request and guarantees the events are
eventually
+ delivered instead of being dropped when the write cannot complete inline.
+
+ ``ti_id`` is the primary key: at most one pending registration exists per
task
+ instance, and the row is cascade-deleted if the task instance is removed.
+ """
+
+ ti_id: Mapped[UUID] = mapped_column(sa.Uuid(), primary_key=True,
nullable=False)
+ task_outlets: Mapped[list] = mapped_column(sa.JSON(), nullable=False,
default=list)
+ outlet_events: Mapped[list] = mapped_column(sa.JSON(), nullable=False,
default=list)
Review Comment:
Done. The table now has a single `payload` JSON column holding
`{"task_outlets": ..., "outlet_events": ...}` instead of two columns, and the
enqueue is one `merge`. The table is new in 3.4.0 so there is no back-compat
concern. Updated the model, the migration, and the drain read path.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -622,6 +646,95 @@ def _validate_outlet_event_partition_keys(outlet_events:
list[dict[str, Any]]) -
)
+# Lock timeout (seconds) for the opportunistic synchronous asset-event write.
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.
+_ASSET_SYNC_LOCK_TIMEOUT = 2
+
+
+def _register_asset_changes_or_enqueue(
+ *,
+ task_instance_id: UUID,
+ task_outlets: list[AssetProfile],
+ outlet_events: list[dict[str, Any]],
+ session: SessionDep,
+) -> None:
+ """
+ Register a successful task's asset events, falling back to the durable
queue on contention.
+
+ Runs after the task-instance state has been committed and its row lock
released. The write is
+ attempted synchronously under a short lock timeout so the common,
uncontended case stays fast.
+ If the write cannot acquire its locks in time (or fails for any other
reason) the events are
+ written to ``asset_event_queue`` and the scheduler drains them, retrying
until they land. The
+ events are therefore never dropped, and the request never holds the
task_instance lock for the
+ registration work.
+ """
+ if not task_outlets and not outlet_events:
+ return
+
+ 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)
+ # Supersede any row a previous failed attempt for this TI left in
the queue so the
+ # scheduler drain does not re-register the same events after a
synchronous success.
+
session.execute(delete(AssetEventQueue).where(AssetEventQueue.ti_id ==
task_instance_id))
+ session.commit()
+ return
+ except OperationalError as e:
+ session.rollback()
+ if not is_lock_not_available_error(e):
+ log.warning(
+ "Asset registration failed with a database error; deferring to
the asset event "
+ "queue for the scheduler to retry.",
+ task_instance_id=str(task_instance_id),
+ exc_info=True,
+ )
+ except Exception:
+ session.rollback()
+ log.warning(
+ "Asset registration failed; deferring to the asset event queue for
the scheduler to retry.",
+ task_instance_id=str(task_instance_id),
+ exc_info=True,
+ )
+
+ _enqueue_asset_events(
+ task_instance_id=task_instance_id,
+ task_outlets=task_outlets,
+ outlet_events=outlet_events,
+ session=session,
+ )
+
+
+def _enqueue_asset_events(
+ *,
+ task_instance_id: UUID,
+ task_outlets: list[AssetProfile],
+ outlet_events: list[dict[str, Any]],
+ session: SessionDep,
+) -> None:
+ """
+ Durably record a completed task's asset events for the scheduler to
register later.
+
+ Uses ``merge`` so a retried completion overwrites the pending row rather
than conflicting on
+ the ``ti_id`` primary key.
+ """
+ session.merge(
+ AssetEventQueue(
+ ti_id=task_instance_id,
+ task_outlets=[outlet.model_dump(mode="json") for outlet in
task_outlets],
+ outlet_events=outlet_events,
+ attempts=0,
+ )
Review Comment:
Addressed by the same change. The enqueue is a single `merge` of one JSON
`payload` column now, so there is no per-column serialization on the
task-success path.
##########
airflow-core/docs/authoring-and-scheduling/asset-scheduling.rst:
##########
@@ -51,6 +51,17 @@ In addition to scheduling Dags based on time, you can also
schedule Dags to run
:ref:`asset_definitions` for how to declare assets.
+.. note::
+
+ Asset events are registered asynchronously. When a task succeeds, Airflow
registers its
+ asset events shortly afterwards rather than strictly inside the
task-completion request,
+ usually within one scheduler loop and longer under heavy load or database
contention.
+ Downstream asset-scheduled Dags are therefore triggered after a brief,
eventually
+ consistent delay. The events are never dropped: if the registration cannot
complete
+ inline it is handed to the scheduler, which retries until it succeeds. Use
the
+ ``[scheduler] asset_event_queue_drain_interval`` option to tune how often
the scheduler
+ processes deferred asset events.
+
Review Comment:
Added a note to the asset-scheduling doc and the newsfragment: under
`dag.test` there is no scheduler to drain the queue, so an event that cannot
register inline will not trigger its downstream runs during the test.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3676,6 +3682,88 @@ def _remove_unreferenced_triggers(self, *, session:
Session = NEW_SESSION) -> No
.execution_options(synchronize_session="fetch")
)
+ @provide_session
+ def _drain_asset_event_queue(self, *, session: Session = NEW_SESSION) ->
None:
+ """
+ Process pending asset-event registrations enqueued by the execution
API.
+
+ The task-success path enqueues an
:class:`~airflow.models.asset.AssetEventQueue` row when
+ synchronous asset-event registration is skipped or fails under
database lock contention.
+ Here we claim pending rows one at a time (``FOR UPDATE SKIP LOCKED``,
so multiple schedulers
+ do not collide), (re)run ``register_asset_changes_in_db`` to create
the ``AssetEvent`` and
+ ``AssetDagRunQueue`` rows, and delete the queue row in the same
transaction so it is
+ processed exactly once. A row that keeps failing is retried on later
passes until it exceeds
+ ``asset_event_queue_max_attempts``, after which it is left in place
(parked) and surfaced via
+ the ``asset.event_queue.failures`` metric for an operator to inspect.
Review Comment:
Reworked along your suggestion. The drain claims the whole batch in a single
`FOR UPDATE SKIP LOCKED` query, processes each row inside a
`session.begin_nested()` savepoint, and commits once at the end instead of
per-row commit/rollback. Dropped the per-pass `seen` set since a single batched
select already avoids re-claiming a row within the pass.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3676,6 +3682,88 @@ def _remove_unreferenced_triggers(self, *, session:
Session = NEW_SESSION) -> No
.execution_options(synchronize_session="fetch")
)
+ @provide_session
+ def _drain_asset_event_queue(self, *, session: Session = NEW_SESSION) ->
None:
+ """
+ Process pending asset-event registrations enqueued by the execution
API.
+
+ The task-success path enqueues an
:class:`~airflow.models.asset.AssetEventQueue` row when
+ synchronous asset-event registration is skipped or fails under
database lock contention.
+ Here we claim pending rows one at a time (``FOR UPDATE SKIP LOCKED``,
so multiple schedulers
+ do not collide), (re)run ``register_asset_changes_in_db`` to create
the ``AssetEvent`` and
+ ``AssetDagRunQueue`` rows, and delete the queue row in the same
transaction so it is
+ processed exactly once. A row that keeps failing is retried on later
passes until it exceeds
+ ``asset_event_queue_max_attempts``, after which it is left in place
(parked) and surfaced via
+ the ``asset.event_queue.failures`` metric for an operator to inspect.
+ """
+ from airflow.api_fastapi.execution_api.datamodels.asset import
AssetProfile
+
+ batch_size = conf.getint("scheduler", "asset_event_queue_batch_size",
fallback=100)
+ max_attempts = conf.getint("scheduler",
"asset_event_queue_max_attempts", fallback=5)
Review Comment:
Moved `batch_size`, `max_attempts` and `drain_interval` into
`SchedulerJobRunner.__init__`, next to the other scheduler config reads that
are done there on purpose to keep the loop free of conf lookups (the existing
comment in `__init__` calls out the secret-backend / HA-lock hazard you
mentioned).
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3676,6 +3682,88 @@ def _remove_unreferenced_triggers(self, *, session:
Session = NEW_SESSION) -> No
.execution_options(synchronize_session="fetch")
)
+ @provide_session
+ def _drain_asset_event_queue(self, *, session: Session = NEW_SESSION) ->
None:
+ """
+ Process pending asset-event registrations enqueued by the execution
API.
+
+ The task-success path enqueues an
:class:`~airflow.models.asset.AssetEventQueue` row when
+ synchronous asset-event registration is skipped or fails under
database lock contention.
+ Here we claim pending rows one at a time (``FOR UPDATE SKIP LOCKED``,
so multiple schedulers
+ do not collide), (re)run ``register_asset_changes_in_db`` to create
the ``AssetEvent`` and
+ ``AssetDagRunQueue`` rows, and delete the queue row in the same
transaction so it is
+ processed exactly once. A row that keeps failing is retried on later
passes until it exceeds
+ ``asset_event_queue_max_attempts``, after which it is left in place
(parked) and surfaced via
+ the ``asset.event_queue.failures`` metric for an operator to inspect.
+ """
+ from airflow.api_fastapi.execution_api.datamodels.asset import
AssetProfile
+
+ batch_size = conf.getint("scheduler", "asset_event_queue_batch_size",
fallback=100)
+ max_attempts = conf.getint("scheduler",
"asset_event_queue_max_attempts", fallback=5)
+
+ processed = 0
+ # Track rows attempted this pass so a freshly-failed row (whose
attempts were just bumped)
+ # is not re-claimed and hammered within the same drain; it waits for
the next pass instead.
+ seen: set[UUID] = set()
+ for _ in range(batch_size):
+ conditions = [AssetEventQueue.attempts < max_attempts]
+ if seen:
+ conditions.append(AssetEventQueue.ti_id.notin_(seen))
+ row = session.scalars(
+ with_row_locks(
+
select(AssetEventQueue).where(*conditions).order_by(AssetEventQueue.created_at).limit(1),
+ of=AssetEventQueue,
+ session=session,
+ skip_locked=True,
+ )
+ ).first()
+ if row is None:
+ break
+
+ ti_id = row.ti_id
+ seen.add(ti_id)
+ ti = session.get(TaskInstance, ti_id)
+ try:
Review Comment:
Real bug, thanks. I fixed it with `ON UPDATE CASCADE` on the `aeq_ti_fkey`
foreign key rather than a manual `UPDATE` in `prepare_db_for_next_try`. The
manual update runs into FK ordering: the constraint is not deferrable, so
updating `asset_event_queue.ti_id` to the new id before `task_instance` has
that id fails, and doing it after the id change orphans the row while the
parent update runs. `ON UPDATE CASCADE` lets the database re-point the queue
row atomically with the id change, and it covers every path that reassigns the
id, not just this method. Regression test:
`test_clearing_ti_repoints_asset_event_queue`.
--
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]
