Clarify a bit the code re updating next_dag run and next run after (#59724)

Currently it's a bit complicated and confusing.  It somewhat remains so after this update, but hopefully a little improved.

For example we have an argument "last automated dagrun" that takes not a dag run but a data interval.

And, when we call it, in places in the scheduler, we may call it with a run that is not actually "the last automated dagrun" -- it could be any dag run the scheduler touches, which of course could be a very old dag run.

I don't solve that problem here.

But I do try to make things a little clearer and a little simpler at the scheduler call sites.

In particular, there were two things that bothered me.

1. in the function `_should_update_dag_next_dagruns`, we did not only *check* whether we should update.  In some branches we actually *did* the update then said "no, don't update."  I pull this update higher up where it's more visible.
2. it's not obvious that the `calculate_dagrun_date_fields` function actually mutates the dag. It just says "calculate," right? So, to improve this slightly, and to reduce noise a bit at call site, I wrap this (along with the update mentioned above) in a function called `_update_next_dagrun_fields`.  At least it's clear we are (at least probably) updating the dag fields.  This has the side benefit bringing all invocations of `calculate_dagrun_date_fields` into one function in the scheduler (which is invoked in a few places).

So hopefully modest clarity improvement, and hopefully no behavior change.

Also -- discovered something very not obvious about this logic.

If you passed dag run to it after dag run created, then it would never update the fields. but this is probably the most important time to do the update, so teh scheduler doesn't try to create the same dagrun again.

So to prevent ogainst this confusing scenario, I just pull that logic out of this function and move it to the specific call sites where it is required -- only in the scheduling (not dagrun creation) context.

Author: dstandish

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -111,6 +111,7 @@
     from airflow.executors.executor_utils import ExecutorName
     from airflow.models.taskinstance import TaskInstanceKey
     from airflow.serialization.definitions.dag import SerializedDAG
+    from airflow.timetables.base import DataInterval
     from airflow.utils.sqlalchemy import CommitProhibitorGuard
 
 TI = TaskInstance
@@ -1846,14 +1847,14 @@ def _create_dag_runs(self, dag_models: Collection[DagModel], session: Session) -
                     #  https://github.com/apache/airflow/issues/59120
                     continue
 
-            if self._should_update_dag_next_dagruns(
-                serdag,
-                dag_model,
-                last_dag_run=None,
-                active_non_backfill_runs=active_runs_of_dags[serdag.dag_id],
+            self._update_next_dagrun_fields(
+                serdag=serdag,
+                dag_model=dag_model,
                 session=session,
-            ):
-                dag_model.calculate_dagrun_date_fields(serdag, data_interval)
+                active_non_backfill_runs=active_runs_of_dags[serdag.dag_id],
+                data_interval=data_interval,
+            )
+
         # TODO[HA]: Should we do a session.flush() so we don't have to keep lots of state/object in
         #  memory for larger dags? or expunge_all()
 
@@ -1933,45 +1934,6 @@ def _create_dag_runs_asset_triggered(
             dag_run.consumed_asset_events.extend(asset_events)
             session.execute(delete(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id == dag_run.dag_id))
 
-    def _should_update_dag_next_dagruns(
-        self,
-        dag: SerializedDAG,
-        dag_model: DagModel,
-        *,
-        last_dag_run: DagRun | None = None,
-        active_non_backfill_runs: int | None = None,
-        session: Session,
-    ) -> bool:
-        """Check if the dag's next_dagruns_create_after should be updated."""
-        # If last_dag_run is defined, the update was triggered by a scheduling decision in this DAG run.
-        # In such case, schedule next only if last_dag_run is finished and was an automated run.
-        if last_dag_run and not (
-            last_dag_run.state in State.finished_dr_states and last_dag_run.run_type == DagRunType.SCHEDULED
-        ):
-            return False
-        # If the DAG never schedules skip save runtime
-        if not dag.timetable.can_be_scheduled:
-            return False
-
-        if active_non_backfill_runs is None:
-            runs_dict = DagRun.active_runs_of_dags(
-                dag_ids=[dag.dag_id],
-                exclude_backfill=True,
-                session=session,
-            )
-            active_non_backfill_runs = runs_dict.get(dag.dag_id, 0)
-
-        if active_non_backfill_runs >= dag.max_active_runs:
-            self.log.info(
-                "DAG %s is at (or above) max_active_runs (%d of %d), not creating any more runs",
-                dag_model.dag_id,
-                active_non_backfill_runs,
-                dag.max_active_runs,
-            )
-            dag_model.next_dagrun_create_after = None
-            return False
-        return True
-
     def _lock_backfills(self, dag_runs: Collection[DagRun], session: Session) -> dict[int, Backfill]:
         """
         Lock Backfill rows to prevent race conditions when multiple schedulers run concurrently.
@@ -2147,16 +2109,13 @@ def _schedule_dag_run(
 
             dag = dag_run.dag = self.scheduler_dag_bag.get_dag_for_run(dag_run=dag_run, session=session)
             dag_model = DM.get_dagmodel(dag_run.dag_id, session)
-            if dag_model is None:
+            if not dag_model:
                 self.log.error("Couldn't find DAG model %s in database!", dag_run.dag_id)
                 return callback
 
             if not dag:
                 self.log.error("Couldn't find DAG %s in DAG bag!", dag_run.dag_id)
                 return callback
-            if not dag_model:
-                self.log.error("Couldn't find DAG model %s in database!", dag_run.dag_id)
-                return callback
 
             if (
                 dag_run.start_date
@@ -2176,10 +2135,17 @@ def _schedule_dag_run(
                 session.flush()
                 self.log.info("Run %s of %s has timed-out", dag_run.run_id, dag_run.dag_id)
 
-                if self._should_update_dag_next_dagruns(
-                    dag, dag_model, last_dag_run=dag_run, session=session
-                ):
-                    dag_model.calculate_dagrun_date_fields(dag, get_run_data_interval(dag.timetable, dag_run))
+                # TODO: questionable that this logic does what it is trying to do
+                #  I think its intent is, in part, to do this when it's the latest scheduled run
+                #  but it does not know that it is the latest. I think it could probably check that
+                #  logical date is equal to or greater than DagModel.next_dagrun, or something
+                if dag_run.state in State.finished_dr_states and dag_run.run_type == DagRunType.SCHEDULED:
+                    self._update_next_dagrun_fields(
+                        serdag=dag,
+                        dag_model=dag_model,
+                        session=session,
+                        data_interval=get_run_data_interval(dag.timetable, dag_run),
+                    )
 
                 dag_run_reloaded = session.scalar(
                     select(DagRun)
@@ -2252,8 +2218,18 @@ def _schedule_dag_run(
             # TODO[HA]: Rename update_state -> schedule_dag_run, ?? something else?
             schedulable_tis, callback_to_run = dag_run.update_state(session=session, execute_callbacks=False)
 
-            if self._should_update_dag_next_dagruns(dag, dag_model, last_dag_run=dag_run, session=session):
-                dag_model.calculate_dagrun_date_fields(dag, get_run_data_interval(dag.timetable, dag_run))
+            # TODO: questionable that this logic does what it is trying to do
+            #  I think its intent is, in part, to do this when it's the latest scheduled run
+            #  but it does not know that it is the latest. I think it could probably check that
+            #  logical date is equal to or greater than DagModel.next_dagrun, or something
+            if dag_run.state in State.finished_dr_states and dag_run.run_type == DagRunType.SCHEDULED:
+                self._update_next_dagrun_fields(
+                    serdag=dag,
+                    dag_model=dag_model,
+                    session=session,
+                    data_interval=get_run_data_interval(dag.timetable, dag_run),
+                )
+
             # This will do one query per dag run. We "could" build up a complex
             # query to update all the TIs across all the logical dates and dag
             # IDs in a single query, but it turns out that can be _very very slow_
@@ -2270,6 +2246,47 @@ def _schedule_dag_run(
 
             return callback_to_run
 
+    def _update_next_dagrun_fields(
+        self,
+        *,
+        serdag: SerializedDAG,
+        dag_model: DagModel,
+        session: Session,
+        active_non_backfill_runs: int | None = None,
+        data_interval: DataInterval | None,
+    ):
+        """
+        Conditionally update fields next_dagrun and next_dagrun_create_after on dag table.
+
+        If dag exceeds max active runs, set to None.
+
+        If dag's timetable not schedulable, don't update.
+
+        Otherwise, update via ``DagModel.calculate_dagrun_date_fields``.
+        """
+        exceeds_max, active_runs = self._exceeds_max_active_runs(
+            dag_model=dag_model,
+            active_non_backfill_runs=active_non_backfill_runs,
+            session=session,
+        )
+        if exceeds_max:
+            self.log.info(
+                "Dag exceeds max_active_runs; not creating any more runs",
+                dag_id=dag_model.dag_id,
+                active_runs=active_runs,
+                max_active_runs=dag_model.max_active_runs,
+            )
+            # null out next_dagrun_create_after so scheduler will not examine this dag
+            # this is periodically reconsidered in the scheduler and dag processor.
+            dag_model.next_dagrun_create_after = None
+            return
+
+        # If the DAG never schedules skip save runtime
+        if not serdag.timetable.can_be_scheduled:
+            return
+
+        dag_model.calculate_dagrun_date_fields(dag=serdag, last_automated_dag_run=data_interval)
+
     def _verify_integrity_if_dag_changed(self, dag_run: DagRun, session: Session) -> bool:
         """
         Only run DagRun.verify integrity if Serialized DAG has changed since it is slow.
@@ -3001,6 +3018,23 @@ def _try_to_load_executor(self, ti: TaskInstance, session, team_name=NOTSET) ->
 
         return executor
 
+    def _exceeds_max_active_runs(
+        self,
+        *,
+        dag_model: DagModel,
+        active_non_backfill_runs: int | None = None,
+        session: Session,
+    ):
+        if active_non_backfill_runs is None:
+            runs_dict = DagRun.active_runs_of_dags(
+                dag_ids=[dag_model.dag_id],
+                exclude_backfill=True,
+                session=session,
+            )
+            active_non_backfill_runs = runs_dict.get(dag_model.dag_id, 0)
+        exceeds = active_non_backfill_runs >= dag_model.max_active_runs
+        return exceeds, active_non_backfill_runs
+
 
 # Backcompat for older versions of task sdk import SchedulerDagBag from here
 SchedulerDagBag = DBDagBag

airflow-core/tests/unit/jobs/test_scheduler_job.py

@@ -63,7 +63,7 @@
     PartitionedAssetKeyLog,
 )
 from airflow.models.backfill import Backfill, _create_backfill
-from airflow.models.dag import DagModel, get_last_dagrun, infer_automated_data_interval
+from airflow.models.dag import DagModel, get_last_dagrun, get_run_data_interval, infer_automated_data_interval
 from airflow.models.dag_version import DagVersion
 from airflow.models.dagbundle import DagBundleModel
 from airflow.models.dagrun import DagRun
@@ -4421,8 +4421,9 @@ def test_should_update_dag_next_dagruns(self, provide_run_count: bool, kwargs: d
             EmptyOperator(task_id="dummy")
 
         index = 0
+        dr: DagRun
         for index in range(other_runs):
-            dag_maker.create_dagrun(
+            dr = dag_maker.create_dagrun(
                 run_id=f"run_{index}",
                 logical_date=(DEFAULT_DATE + timedelta(days=index)),
                 start_date=timezone.utcnow(),
@@ -4444,13 +4445,15 @@ def test_should_update_dag_next_dagruns(self, provide_run_count: bool, kwargs: d
         scheduler_job = Job(executor=self.null_exec)
         self.job_runner = SchedulerJobRunner(job=scheduler_job)
 
-        actual = self.job_runner._should_update_dag_next_dagruns(
-            dag=dag,
-            dag_model=dag_maker.dag_model,
-            active_non_backfill_runs=other_runs if provide_run_count else None,  # exclude backfill here
-            session=session,
-        )
-        assert actual == should_update
+        with patch("airflow.models.dag.DagModel.calculate_dagrun_date_fields") as mock_calc:
+            self.job_runner._update_next_dagrun_fields(
+                serdag=dag,
+                dag_model=dag_maker.dag_model,
+                active_non_backfill_runs=other_runs if provide_run_count else None,  # exclude backfill here
+                session=session,
+                data_interval=get_run_data_interval(dag.timetable, dr),
+            )
+        assert mock_calc.called == should_update
 
     @pytest.mark.parametrize(
         ("run_type", "expected"),
@@ -4472,11 +4475,9 @@ def test_should_update_dag_next_dagruns_after_run_type(self, run_type, expected,
         with dag_maker(
             schedule="*/1 * * * *",
             max_active_runs=3,
-        ) as dag:
+        ):
             EmptyOperator(task_id="dummy")
 
-        dag_model = dag_maker.dag_model
-
         run = dag_maker.create_dagrun(
             run_id="run",
             run_type=run_type,
@@ -4490,13 +4491,19 @@ def test_should_update_dag_next_dagruns_after_run_type(self, run_type, expected,
         scheduler_job = Job(executor=self.null_exec)
         self.job_runner = SchedulerJobRunner(job=scheduler_job)
 
-        actual = self.job_runner._should_update_dag_next_dagruns(
-            dag=dag,
-            dag_model=dag_model,
-            last_dag_run=run,
-            session=session,
-        )
-        assert actual == expected
+        # ensure terminal state; otherwise the run type is moot
+        run.state = DagRunState.FAILED
+        for ti in run.get_task_instances(session=session):
+            ti.state = "failed"
+        session.flush()
+
+        with patch("airflow.models.dag.DagModel.calculate_dagrun_date_fields") as mock_calc:
+            self.job_runner._schedule_dag_run(
+                dag_run=run,
+                session=session,
+            )
+
+        assert mock_calc.called == expected
 
     def test_create_dag_runs(self, dag_maker):
         """