jason810496 commented on code in PR #66854:
URL: https://github.com/apache/airflow/pull/66854#discussion_r3620724237
##########
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:
```suggestion
```
##########
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:
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.
```suggestion
```
##########
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:
Additionally, we should remove all the explicit `commit` and `rollback` in
the route to make this atomic at session injection level.
```suggestion
except Exception:
```
##########
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:
After revisiting this part. I think we should not switch to the new
generated TI.id .
We should use the old TI.id instead. Since the AssetEvent table use
source_task_id without FK, if we re-map to a new generated TI.id, we change the
write sementic. IIUC, we should drop the FK of AssetEventQueue with TI.id and
check if the TI and TIH existed in the scheduler.
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -5712,6 +5713,150 @@ def dict_from_obj(obj):
assert created_run.creating_job_id == scheduler_job.id
+ def _seed_asset_producer_and_queue_row(self, session, dag_maker, *,
outlet_events=None):
+ """Create a producer DAG whose task has an asset outlet, a downstream
DAG scheduled on that
+ asset, a successful producer TI, and a matching ``AssetEventQueue``
row. Returns the TI."""
+ asset = Asset(uri="test://drain-asset", name="drain_asset",
group="test_group")
+
+ with dag_maker(dag_id="drain-producer", start_date=timezone.utcnow(),
session=session) as dag:
+ BashOperator(task_id="task", bash_command="echo 1",
outlets=[asset])
+ dr = dag_maker.create_dagrun(run_id="producer-run")
+ [ti] = dr.get_task_instances(session=session)
+ ti.state = State.SUCCESS
+ # Capture the outlet profiles before creating other DAGs staled the
serialized reference.
+ task_outlets = [o.asprofile().model_dump(mode="json") for o in
dag.get_task("task").outlets]
+ session.flush()
+
+ # Downstream DAG scheduled on the asset -> registration should enqueue
it (ADRQ row).
+ with dag_maker(dag_id="drain-consumer", schedule=[asset],
session=session):
+ EmptyOperator(task_id="downstream")
+ session.add(
+ AssetEventQueue(
+ ti_id=ti.id,
+ payload={"task_outlets": task_outlets, "outlet_events":
outlet_events or []},
+ attempts=0,
+ )
+ )
+ session.commit()
+ return ti
+
+ @pytest.mark.need_serialized_dag
+ def test_drain_asset_event_queue_registers_and_deletes(self, session,
dag_maker):
+ """A pending queue row for a successful producer TI is registered
(AssetEvent + downstream
+ AssetDagRunQueue), the queue row is deleted, and the processed metric
is emitted."""
+ ti = self._seed_asset_producer_and_queue_row(session, dag_maker)
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job,
executors=[self.null_exec])
+
+ with mock.patch("airflow.jobs.scheduler_job_runner.stats") as
mock_stats:
+ self.job_runner._drain_asset_event_queue(session=session)
+
+ # The queue row was consumed.
+ assert session.scalars(select(AssetEventQueue)).all() == []
+ # An AssetEvent was produced by the (now drained) TI.
+ event =
session.scalar(select(AssetEvent).where(AssetEvent.source_run_id == ti.run_id))
+ assert event is not None
+ # The downstream asset-scheduled DAG was queued.
+ adrq = session.scalars(
+ select(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id ==
"drain-consumer")
+ ).all()
+ assert len(adrq) == 1
+ mock_stats.incr.assert_any_call("asset.event_queue.processed", 1)
+
+ @pytest.mark.need_serialized_dag
+ def test_drain_asset_event_queue_retries_then_parks(self, session,
dag_maker):
+ """A row whose registration keeps failing is retried across drain
passes, one attempt per
+ pass, and once it hits max_attempts it is parked and the failures
metric is emitted."""
+ ti = self._seed_asset_producer_and_queue_row(session, dag_maker)
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job,
executors=[self.null_exec])
+
+ max_attempts = self.job_runner._asset_event_queue_max_attempts
+
+ with mock.patch(
+
"airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db",
+ side_effect=RuntimeError("boom"),
+ ):
+ # First pass: the single pending row is claimed once, its
registration fails, so attempts
+ # becomes 1 and it is not yet parked.
+ with mock.patch("airflow.jobs.scheduler_job_runner.stats") as
mock_stats:
+ self.job_runner._drain_asset_event_queue(session=session)
+ row = session.scalars(select(AssetEventQueue)).one()
+ assert row.ti_id == ti.id
+ assert row.attempts == 1
+ assert mock.call("asset.event_queue.failures") not in
mock_stats.incr.mock_calls
+
+ # Run further passes; on the pass that pushes attempts to
max_attempts the row is parked
+ # and the failures metric fires.
+ for _ in range(max_attempts - 1):
Review Comment:
Actually, we can make the test like the following to exercise more:
```python
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.
```
##########
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:
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.
```suggestion
```
##########
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:
Ditto. We can remove
`test_ti_update_state_asset_registration_failure_enqueues` and
`test_ti_update_state_lock_contention_enqueues`.
##########
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:
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.
##########
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:
IIUC, We didn't add the actual statement to point the `AssetEventQueue` row
to the new generated id here.
--
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]
