Repository: apache/airflow

Issue:
Many tasks updating dataset at once causes some of them to fail
### Apache Airflow version

main (development)

### What happened

I have 16 dags which all update the same dataset.  They're set to finish at the same time (when the seconds on the clock are 00). About three quarters of them behave as expected, but the other quarter fails with errors like:

```
[2022-07-21, 06:06:00 UTC] {standard_task_runner.py:97} ERROR - Failed to execute job 8 for task increment_source ((psycopg2.errors.UniqueViolation) duplicate key value violates unique constraint "dataset_dag_run_queue_pkey"
DETAIL:  Key (dataset_id, target_dag_id)=(1, simple_dataset_sink) already exists.

[SQL: INSERT INTO dataset_dag_run_queue (dataset_id, target_dag_id, created_at) VALUES (%(dataset_id)s, %(target_dag_id)s, %(created_at)s)]
[parameters: {'dataset_id': 1, 'target_dag_id': 'simple_dataset_sink', 'created_at': datetime.datetime(2022, 7, 21, 6, 6, 0, 131730, tzinfo=Timezone('UTC'))}]
(Background on this error at: https://sqlalche.me/e/14/gkpj); 375)
```

I've prepaired a gist with the details: https://gist.github.com/MatrixManAtYrService/b5e58be0949eab9180608d0760288d4d

### What you think should happen instead

All dags should succeed

### How to reproduce

See this gist:  https://gist.github.com/MatrixManAtYrService/b5e58be0949eab9180608d0760288d4d

Summary:  Unpause all of the dags which we expect to collide, wait two minutes.  Some will have collided.

### Operating System

docker/debian

### Versions of Apache Airflow Providers

n/a

### Deployment

Astronomer

### Deployment details

`astro dev start` targeting commit: cff7d9194f549d801947f47dfce4b5d6870bfaaa

be sure to have `pause` in requirements.txt

### Anything else

_No response_

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
This is described here:  https://github.com/apache/airflow/issues/24974

Maybe it's ok that the update failed--given that the downstream dag ran as desired--but I don't think it should cause the task to fail.
nice work!
Ran across this again (at least I think it's the same thing) with this dag:

```python3
from airflow import Dataset, DAG
from airflow.operators.python import PythonOperator
from datetime import datetime


fan_out = Dataset("fan_out")
fan_in = Dataset("fan_in")

# the leader
with DAG(
    dag_id="momma_duck", start_date=datetime(1970, 1, 1), schedule_interval=None
) as leader:

    PythonOperator(
        task_id="has_outlet", python_callable=lambda: None, outlets=[fan_out]
    )

# the many
for i in range(1, 1000):
    with DAG(
        dag_id=f"duckling_{i}", start_date=datetime(1970, 1, 1), schedule_on=[fan_out]
    ) as duck:

        PythonOperator(
            task_id="has_outlet", python_callable=lambda: None, outlets=[fan_in]
        )
    globals()[f"duck_{i}"] = duck


# the straggler
with DAG(
    dag_id="straggler_duck", start_date=datetime(1970, 1, 1), schedule_on=[fan_in]
) as straggler:

    PythonOperator(task_id="has_outlet", python_callable=lambda: None)

```

This time I noticed that, with the LocalExecutor at least, there are messages regarding zombines in the scheduler log:

```{dag.py:3178} INFO - Setting next_dagrun for duckling_897 to None, run_after=None
{dagrun.py:567} INFO - Marking run <DagRun duckling_855 @ 2022-08-08 21:01:33.571168+00:00: dataset_triggered__2022-08-08T21:01:33.568548+00:00, state:running, queued_at: 2022-08-08 21:01:33.569170+00:00. externally triggered: False> successful
{dagrun.py:612} INFO - DagRun Finished: dag_id=duckling_855, execution_date=2022-08-08 21:01:33.571168+00:00, run_id=dataset_triggered__2022-08-08T21:01:33.568548+00:00, run_start_date=2022-08-08 21:03:40.332173+00:00, run_end_date=2022-08-08 21:07:27.579980+00:00, run_duration=227.247807, state=success, external_trigger=False, run_type=dataset_triggered, data_interval_start=None, data_interval_end=None, dag_hash=ff772192a0a9c22e28dbebf2b14f1ebc
{dag.py:3178} INFO - Setting next_dagrun for duckling_855 to None, run_after=None
{scheduler_job.py:1364} WARNING - Failing (58) jobs without heartbeat after 2022-08-08 21:02:27.737896+00:00
{scheduler_job.py:1372} ERROR - Detected zombie job: {'full_filepath': '/home/matt/2022/08/07/dags/many_to_one.py', 'msg': 'Detected <TaskInstance: duckling_289.has_outlet dataset_triggered__2022-08-08T21:01:31.846596+00:00 [running]> as zombie', 'simple_task_instance': <airflow.models.taskinstance.SimpleTaskInstance object at 0x7f98ebffc850>, 'is_failure_callback': True}
{scheduler_job.py:1372} ERROR - Detected zombie job: {'full_filepath': '/home/matt/2022/08/07/dags/many_to_one.py', 'msg': 'Detected <TaskInstance: duckling_127.has_outlet dataset_triggered__2022-08-08T21:01:31.823203+00:00 [running]> as zombie', 'simple_task_instance': <airflow.models.taskinstance.SimpleTaskInstance object at 0x7f98ec93f7f0>, 'is_failure_callback': True}

```

Resolving patch:
diff --git a/airflow/datasets/manager.py b/airflow/datasets/manager.py
--- a/airflow/datasets/manager.py
+++ b/airflow/datasets/manager.py
@@ -17,6 +17,7 @@
 # under the License.
 from typing import TYPE_CHECKING
 
+from sqlalchemy import exc
 from sqlalchemy.orm.session import Session
 
 from airflow.configuration import conf
@@ -63,13 +64,47 @@ def register_dataset_change(
         self._queue_dagruns(dataset_model, session)
 
     def _queue_dagruns(self, dataset: DatasetModel, session: Session) -> None:
+        # Possible race condition: if multiple dags or multiple (usually
+        # mapped) tasks update the same dataset, this can fail with a unique
+        # constraint violation.
+        #
+        # If we support it, use ON CONFLICT to do nothing, otherwise
+        # "fallback" to running this in a nested transaction. This is needed
+        # so that the adding of these rows happens in the same transaction
+        # where `ti.state` is changed.
+
+        if session.bind.dialect.name == "postgresql":
+            return self._postgres_queue_dagruns(dataset, session)
+        return self._slow_path_queue_dagruns(dataset, session)
+
+    def _slow_path_queue_dagruns(self, dataset: DatasetModel, session: Session) -> None:
         consuming_dag_ids = [x.dag_id for x in dataset.consuming_dags]
         self.log.debug("consuming dag ids %s", consuming_dag_ids)
+
+        # Don't error whole transaction when a single RunQueue item conflicts.
+        # https://docs.sqlalchemy.org/en/14/orm/session_transaction.html#using-savepoint
         for dag_id in consuming_dag_ids:
-            session.merge(DatasetDagRunQueue(dataset_id=dataset.id, target_dag_id=dag_id))
+            item = DatasetDagRunQueue(target_dag_id=dag_id, dataset_id=dataset.id)
+            try:
+                with session.begin_nested():
+                    session.merge(item)
+            except exc.IntegrityError:
+                self.log.debug("Skipping record %s", item, exc_info=True)
+
+        session.flush()
+
+    def _postgres_queue_dagruns(self, dataset: DatasetModel, session: Session) -> None:
+        from sqlalchemy.dialects.postgresql import insert
+
+        stmt = insert(DatasetDagRunQueue).values(dataset_id=dataset.id).on_conflict_do_nothing()
+        session.execute(
+            stmt,
+            [{'target_dag_id': target_dag.dag_id} for target_dag in dataset.consuming_dags],
+        )
+        session.flush()
 
 
-def resolve_dataset_manager():
+def resolve_dataset_manager() -> "DatasetManager":
     _dataset_manager_class = conf.getimport(
         section='core',
         key='dataset_manager_class',
diff --git a/airflow/models/dataset.py b/airflow/models/dataset.py
--- a/airflow/models/dataset.py
+++ b/airflow/models/dataset.py
@@ -117,6 +117,11 @@ class DagScheduleDatasetReference(Base):
             DagScheduleDatasetReference.dag_id == foreign(DatasetDagRunQueue.target_dag_id),
         )""",
     )
+    dag = relationship(
+        "DagModel",
+        primaryjoin="foreign(DagModel.dag_id) == DagScheduleDatasetReference.dag_id",
+        uselist=False,
+    )
 
     __tablename__ = "dag_schedule_dataset_reference"
     __table_args__ = (


Test patch:
diff --git a/tests/datasets/test_manager.py b/tests/datasets/test_manager.py
--- a/tests/datasets/test_manager.py
+++ b/tests/datasets/test_manager.py
@@ -22,7 +22,8 @@
 
 from airflow.datasets import Dataset
 from airflow.datasets.manager import DatasetManager
-from airflow.models.dataset import DatasetModel
+from airflow.models.dag import DagModel
+from airflow.models.dataset import DagScheduleDatasetReference, DatasetDagRunQueue, DatasetEvent, DatasetModel
 
 
 @pytest.fixture()
@@ -61,26 +62,21 @@ def test_register_dataset_change_dataset_doesnt_exist(self, mock_task_instance):
         mock_session.add.assert_not_called()
         mock_session.merge.assert_not_called()
 
-    def test_register_dataset_change(self, mock_task_instance):
+    def test_register_dataset_change(self, session, dag_maker, mock_task_instance):
         dsem = DatasetManager()
 
-        mock_dag_1 = mock.MagicMock()
-        mock_dag_1.dag_id = 1
-        mock_dag_2 = mock.MagicMock()
-        mock_dag_2.dag_id = 2
-
         ds = Dataset(uri="test_dataset_uri")
+        dag1 = DagModel(dag_id="dag1")
+        dag2 = DagModel(dag_id="dag2")
+        session.add_all([dag1, dag2])
 
         dsm = DatasetModel(uri="test_dataset_uri")
-        dsm.consuming_dags = [mock_dag_1, mock_dag_2]
-
-        mock_session = mock.Mock()
-        # Gotta mock up the query results
-        mock_session.query.return_value.filter.return_value.one_or_none.return_value = dsm
+        session.add(dsm)
+        dsm.consuming_dags = [DagScheduleDatasetReference(dag_id=dag.dag_id) for dag in (dag1, dag2)]
+        session.flush()
 
-        dsem.register_dataset_change(task_instance=mock_task_instance, dataset=ds, session=mock_session)
+        dsem.register_dataset_change(task_instance=mock_task_instance, dataset=ds, session=session)
 
         # Ensure we've created a dataset
-        mock_session.add.assert_called_once()
-        # Ensure that we've created DatasetDagRunQueue rows
-        assert mock_session.merge.call_count == 2
+        assert session.query(DatasetEvent).filter_by(dataset_id=dsm.id).count() == 1
+        assert session.query(DatasetDagRunQueue).count() == 2
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -58,7 +58,7 @@
     Variable,
     XCom,
 )
-from airflow.models.dataset import DatasetDagRunQueue, DatasetEvent, DatasetModel, TaskOutletDatasetReference
+from airflow.models.dataset import DatasetDagRunQueue, DatasetEvent, DatasetModel
 from airflow.models.expandinput import EXPAND_INPUT_EMPTY
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.models.taskfail import TaskFail
@@ -1716,10 +1716,19 @@ def test_outlet_datasets(self, create_task_instance, clear_datasets):
         ti.refresh_from_db()
         assert ti.state == TaskInstanceState.SUCCESS
 
+        # check that no other dataset events recorded
+        event = (
+            session.query(DatasetEvent)
+            .join(DatasetEvent.dataset)
+            .filter(DatasetEvent.source_task_instance == ti)
+            .one()
+        )
+        assert event
+        assert event.dataset
+
         # check that one queue record created for each dag that depends on dataset 1
-        assert session.query(DatasetDagRunQueue.target_dag_id).filter(
-            TaskOutletDatasetReference.dag_id == dag1.dag_id,
-            TaskOutletDatasetReference.task_id == 'producing_task_1',
+        assert session.query(DatasetDagRunQueue.target_dag_id).filter_by(
+            dataset_id=event.dataset.id
         ).order_by(DatasetDagRunQueue.target_dag_id).all() == [
             ('dataset_consumes_1',),
             ('dataset_consumes_1_and_2',),
@@ -1731,14 +1740,6 @@ def test_outlet_datasets(self, create_task_instance, clear_datasets):
             DatasetEvent.source_task_instance == ti
         ).one() == ('s3://dag1/output_1.txt',)
 
-        # check that no other dataset events recorded
-        assert (
-            session.query(DatasetModel.uri)
-            .join(DatasetEvent.dataset)
-            .filter(DatasetEvent.source_task_instance == ti)
-            .count()
-        ) == 1
-
     def test_outlet_datasets_failed(self, create_task_instance, clear_datasets):
         """
         Verify that when we have an outlet dataset on a task, and the task