Repository: apache/airflow

Issue:
Scheduler deadlock error when mapping over empty list
### Apache Airflow version

2.3.0b1 (pre-release)

### What happened

manually triggered this dag:
```python
from datetime import datetime
from airflow import DAG

with DAG(
    dag_id="null_mapped_2",
    start_date=datetime(1970, 1, 1),
    schedule_interval=None,
) as dag:

    @dag.task
    def empty():
        return []

    @dag.task
    def print_it(thing):
        print(thing)

    print_it.expand(thing=empty())

```

scheduler logs (whitespace added for emphasis):
```
   ____________       _____________
  ____    |__( )_________  __/__  /________      __
 ____  /| |_  /__  ___/_  /_ __  /_  __ \_ | /| / /
 ___  ___ |  / _  /   _  __/ _  / / /_/ /_ |/ |/ /
  _/_/  |_/_/  /_/    /_/    /_/  \____/____/|__/
 [2022-04-20 20:46:24,760] {scheduler_job.py:697} INFO - Starting the scheduler
 [2022-04-20 20:46:24,761] {scheduler_job.py:702} INFO - Processing each file at most -1 times
 [2022-04-20 20:46:24 +0000] [28] [INFO] Starting gunicorn 20.1.0
 [2022-04-20 20:46:24 +0000] [28] [INFO] Listening at: http://0.0.0.0:8793 (28)
 [2022-04-20 20:46:24 +0000] [28] [INFO] Using worker: sync
 [2022-04-20 20:46:24 +0000] [29] [INFO] Booting worker with pid: 29
 [2022-04-20 20:46:24,782] {executor_loader.py:106} INFO - Loaded executor: LocalExecutor
 [2022-04-20 20:46:24 +0000] [78] [INFO] Booting worker with pid: 78
 [2022-04-20 20:46:24,953] {manager.py:156} INFO - Launched DagFileProcessorManager with pid: 166
 [2022-04-20 20:46:24,962] {settings.py:55} INFO - Configured default timezone Timezone('UTC')
 /usr/local/lib/python3.9/site-packages/airflow/configuration.py:466 DeprecationWarning: The sql_alchemy_conn option in [core] has been moved to the sql_alchemy_conn option in [database] - the old setting has been used, but please update your config.
 [2022-04-20 20:46:24,988] {scheduler_job.py:1231} INFO - Resetting orphaned tasks for active dag runs
 /usr/local/lib/python3.9/site-packages/airflow/configuration.py:466 DeprecationWarning: The sql_alchemy_conn option in [core] has been moved to the sql_alchemy_conn option in [database] - the old setting has been used, but please update your config.
 [2022-04-20 20:46:32,124] {update_checks.py:128} INFO - Checking for new version of Astronomer Certified Airflow, previous check was performed at None
 [2022-04-20 20:46:32,441] {update_checks.py:84} INFO - Check finished, next check in 86400.0 seconds
 /usr/local/lib/python3.9/site-packages/airflow/configuration.py:466 DeprecationWarning: The sql_alchemy_conn option in [core] has been moved to the sql_alchemy_conn option in [database] - the old setting has been used, but please update your config.
 [2022-04-20 20:47:45,296] {scheduler_job.py:354} INFO - 1 tasks up for execution:
 	<TaskInstance: null_mapped_2.empty manual__2022-04-20T20:47:45.075321+00:00 [scheduled]>
 [2022-04-20 20:47:45,297] {scheduler_job.py:419} INFO - DAG null_mapped_2 has 0/16 running and queued tasks
 [2022-04-20 20:47:45,297] {scheduler_job.py:505} INFO - Setting the following tasks to queued state:
 	<TaskInstance: null_mapped_2.empty manual__2022-04-20T20:47:45.075321+00:00 [scheduled]>
 [2022-04-20 20:47:45,300] {scheduler_job.py:547} INFO - Sending TaskInstanceKey(dag_id='null_mapped_2', task_id='empty', run_id='manual__2022-04-20T20:47:45.075321+00:00', try_number=1, map_index=-1) to executor with priority 2 and queue default
 [2022-04-20 20:47:45,300] {base_executor.py:88} INFO - Adding to queue: ['airflow', 'tasks', 'run', 'null_mapped_2', 'empty', 'manual__2022-04-20T20:47:45.075321+00:00', '--local', '--subdir', 'DAGS_FOLDER/null_mapped_2.py']
 [2022-04-20 20:47:45,303] {local_executor.py:79} INFO - QueuedLocalWorker running ['airflow', 'tasks', 'run', 'null_mapped_2', 'empty', 'manual__2022-04-20T20:47:45.075321+00:00', '--local', '--subdir', 'DAGS_FOLDER/null_mapped_2.py']
 [2022-04-20 20:47:45,340] {dagbag.py:507} INFO - Filling up the DagBag from /usr/local/airflow/dags/null_mapped_2.py
 /usr/local/lib/python3.9/site-packages/airflow/configuration.py:466 DeprecationWarning: The sql_alchemy_conn option in [core] has been moved to the sql_alchemy_conn option in [database] - the old setting has been used, but please update your config.
 [2022-04-20 20:47:45,432] {task_command.py:369} INFO - Running <TaskInstance: null_mapped_2.empty manual__2022-04-20T20:47:45.075321+00:00 [queued]> on host 7ec8e95b149d


 [2022-04-20 20:47:46,796] {dagrun.py:583} ERROR - Deadlock; marking run <DagRun null_mapped_2 @ 2022-04-20 20:47:45.075321+00:00: manual__2022-04-20T20:47:45.075321+00:00, externally triggered: True> failed


 [2022-04-20 20:47:46,797] {dagrun.py:607} INFO - DagRun Finished: dag_id=null_mapped_2, execution_date=2022-04-20 20:47:45.075321+00:00, run_id=manual__2022-04-20T20:47:45.075321+00:00, run_start_date=2022-04-20 20:47:45.254176+00:00, run_end_date=2022-04-20 20:47:46.797390+00:00, run_duration=1.543214, state=failed, external_trigger=True, run_type=manual, data_interval_start=2022-04-20 20:47:45.075321+00:00, data_interval_end=2022-04-20 20:47:45.075321+00:00, dag_hash=8476a887126cf6d52573ee41fa81c637
 [2022-04-20 20:47:46,801] {dag.py:2894} INFO - Setting next_dagrun for null_mapped_2 to None, run_after=None
 [2022-04-20 20:47:46,818] {scheduler_job.py:600} INFO - Executor reports execution of null_mapped_2.empty run_id=manual__2022-04-20T20:47:45.075321+00:00 exited with status success for try_number 1
 [2022-04-20 20:47:46,831] {scheduler_job.py:644} INFO - TaskInstance Finished: dag_id=null_mapped_2, task_id=empty, run_id=manual__2022-04-20T20:47:45.075321+00:00, map_index=-1, run_start_date=2022-04-20 20:47:45.491953+00:00, run_end_date=2022-04-20 20:47:45.760654+00:00, run_duration=0.268701, state=success, executor_state=success, try_number=1, max_tries=0, job_id=2, pool=default_pool, queue=default, priority_weight=2, operator=_PythonDecoratedOperator, queued_dttm=2022-04-20 20:47:45.297857+00:00, queued_by_job_id=1, pid=250
```

The dagrun fails, even though none of its tasks get set to failed:

<img width="319" alt="Screen Shot 2022-04-20 at 2 43 15 PM" src="https://user-images.githubusercontent.com/5834582/164320377-663fcc0a-0bc4-4edc-8cc3-91884b84748d.png">



### What you think should happen instead

Mapping over a 0 length list should create no map_index'ed tasks, but the parent task should succeed because none of the 0 tasks failed.

### How to reproduce

Run the dag above

### Operating System

debian (docker)

### Versions of Apache Airflow Providers

n/a

### Deployment

Astronomer

### Deployment details

`astrocloud dev start`, image contains version 2.3.0.dev20220414

### 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:

Resolving patch:
diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -655,7 +655,18 @@ def task_instance_scheduling_decisions(self, session: Session = NEW_SESSION) ->
         if unfinished_tis:
             schedulable_tis = [ut for ut in unfinished_tis if ut.state in SCHEDULEABLE_STATES]
             self.log.debug("number of scheduleable tasks for %s: %s task(s)", self, len(schedulable_tis))
-            schedulable_tis, changed_tis = self._get_ready_tis(schedulable_tis, finished_tis, session)
+            schedulable_tis, changed_tis, expansion_happened = self._get_ready_tis(
+                schedulable_tis,
+                finished_tis,
+                session=session,
+            )
+
+            # During expansion we may change some tis into non-schedulable
+            # states, so we need to re-compute.
+            if expansion_happened:
+                new_unfinished_tis = [t for t in unfinished_tis if t.state in State.unfinished]
+                finished_tis.extend(t for t in unfinished_tis if t.state in State.finished)
+                unfinished_tis = new_unfinished_tis
 
         return TISchedulingDecision(
             tis=tis,
@@ -670,57 +681,51 @@ def _get_ready_tis(
         schedulable_tis: List[TI],
         finished_tis: List[TI],
         session: Session,
-    ) -> Tuple[List[TI], bool]:
+    ) -> Tuple[List[TI], bool, bool]:
         old_states = {}
         ready_tis: List[TI] = []
         changed_tis = False
 
         if not schedulable_tis:
-            return ready_tis, changed_tis
+            return ready_tis, changed_tis, False
 
         # If we expand TIs, we need a new list so that we iterate over them too. (We can't alter
         # `schedulable_tis` in place and have the `for` loop pick them up
-        expanded_tis: List[TI] = []
+        additional_tis: List[TI] = []
         dep_context = DepContext(
             flag_upstream_failed=True,
             ignore_unmapped_tasks=True,  # Ignore this Dep, as we will expand it if we can.
             finished_tis=finished_tis,
         )
 
-        # Check dependencies
-        for schedulable in itertools.chain(schedulable_tis, expanded_tis):
-
+        # Check dependencies.
+        expansion_happened = False
+        for schedulable in itertools.chain(schedulable_tis, additional_tis):
             old_state = schedulable.state
-            if schedulable.are_dependencies_met(session=session, dep_context=dep_context):
-                ready_tis.append(schedulable)
-            else:
+            if not schedulable.are_dependencies_met(session=session, dep_context=dep_context):
                 old_states[schedulable.key] = old_state
                 continue
-
-            # This is called in two places: First (and ideally) is in the mini scheduler at the end of
-            # LocalTaskJob, and then as an "expansion of last resort" in the scheduler
-            # to ensure that the mapped task is correctly expanded before we try to execute it.
-            if schedulable.map_index < 0 and schedulable.task.is_mapped:
-                # HACK. This needs a better way, one that copes with multiple upstreams!
-                for ti in finished_tis:
-                    if schedulable.task_id in ti.task.downstream_task_ids:
-
-                        assert isinstance(schedulable.task, MappedOperator)
-                        new_tis, _ = schedulable.task.expand_mapped_task(self.run_id, session=session)
-                        if schedulable.state == TaskInstanceState.SKIPPED:
-                            # Task is now skipped (likely cos upstream returned 0 tasks
-                            continue
-                        assert new_tis[0] is schedulable
-                        expanded_tis.extend(new_tis[1:])
-                        break
+            # If schedulable is from a mapped task, but not yet expanded, do it
+            # now. This is called in two places: First and ideally in the mini
+            # scheduler at the end of LocalTaskJob, and then as an "expansion of
+            # last resort" in the scheduler to ensure that the mapped task is
+            # correctly expanded before executed.
+            if schedulable.map_index < 0 and isinstance(schedulable.task, MappedOperator):
+                expanded_tis, _ = schedulable.task.expand_mapped_task(self.run_id, session=session)
+                if expanded_tis:
+                    assert expanded_tis[0] is schedulable
+                    additional_tis.extend(expanded_tis[1:])
+                expansion_happened = True
+            if schedulable.state in SCHEDULEABLE_STATES:
+                ready_tis.append(schedulable)
 
         # Check if any ti changed state
-        tis_filter = TI.filter_for_tis(old_states.keys())
+        tis_filter = TI.filter_for_tis(old_states)
         if tis_filter is not None:
             fresh_tis = session.query(TI).filter(tis_filter).all()
             changed_tis = any(ti.state != old_states[ti.key] for ti in fresh_tis)
 
-        return ready_tis, changed_tis
+        return ready_tis, changed_tis, expansion_happened
 
     def _are_premature_tis(
         self,
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -622,7 +622,7 @@ def expand_mapped_task(self, run_id: str, *, session: Session) -> Tuple[Sequence
             .one_or_none()
         )
 
-        ret: List[TaskInstance] = []
+        all_expanded_tis: List[TaskInstance] = []
 
         if unmapped_ti:
             # The unmapped task instance still exists and is unfinished, i.e. we
@@ -636,14 +636,13 @@ def expand_mapped_task(self, run_id: str, *, session: Session) -> Tuple[Sequence
                     total_length,
                 )
                 unmapped_ti.state = TaskInstanceState.SKIPPED
-                session.flush()
-                return ret, 0
-            # Otherwise convert this into the first mapped index, and create
-            # TaskInstance for other indexes.
-            unmapped_ti.map_index = 0
+            else:
+                # Otherwise convert this into the first mapped index, and create
+                # TaskInstance for other indexes.
+                unmapped_ti.map_index = 0
+                self.log.debug("Updated in place to become %s", unmapped_ti)
+                all_expanded_tis.append(unmapped_ti)
             state = unmapped_ti.state
-            self.log.debug("Updated in place to become %s", unmapped_ti)
-            ret.append(unmapped_ti)
             indexes_to_map = range(1, total_length)
         else:
             # Only create "missing" ones.
@@ -660,13 +659,12 @@ def expand_mapped_task(self, run_id: str, *, session: Session) -> Tuple[Sequence
 
         for index in indexes_to_map:
             # TODO: Make more efficient with bulk_insert_mappings/bulk_save_mappings.
-            # TODO: Change `TaskInstance` ctor to take Operator, not BaseOperator
-            ti = TaskInstance(self, run_id=run_id, map_index=index, state=state)  # type: ignore
+            ti = TaskInstance(self, run_id=run_id, map_index=index, state=state)
             self.log.debug("Expanding TIs upserted %s", ti)
             task_instance_mutation_hook(ti)
             ti = session.merge(ti)
-            ti.task = self
-            ret.append(ti)
+            ti.refresh_from_task(self)  # session.merge() loses task information.
+            all_expanded_tis.append(ti)
 
         # Set to "REMOVED" any (old) TaskInstances with map indices greater
         # than the current map value
@@ -678,8 +676,7 @@ def expand_mapped_task(self, run_id: str, *, session: Session) -> Tuple[Sequence
         ).update({TaskInstance.state: TaskInstanceState.REMOVED})
 
         session.flush()
-
-        return ret, total_length
+        return all_expanded_tis, total_length
 
     def prepare_for_execution(self) -> "MappedOperator":
         # Since a mapped operator cannot be used for execution, and an unmapped


Test patch:
diff --git a/tests/models/test_dagrun.py b/tests/models/test_dagrun.py
--- a/tests/models/test_dagrun.py
+++ b/tests/models/test_dagrun.py
@@ -30,6 +30,7 @@
 from airflow.decorators import task
 from airflow.models import DAG, DagBag, DagModel, DagRun, TaskInstance as TI, clear_task_instances
 from airflow.models.baseoperator import BaseOperator
+from airflow.models.dagrun import TISchedulingDecision
 from airflow.models.taskmap import TaskMap
 from airflow.models.xcom_arg import XComArg
 from airflow.operators.empty import EmptyOperator
@@ -1055,14 +1056,24 @@ def test_ti_scheduling_mapped_zero_length(dag_maker, session):
         mapped = MockOperator.partial(task_id='task_2').expand(arg2=XComArg(task))
 
     dr: DagRun = dag_maker.create_dagrun()
-    ti1, _ = sorted(dr.task_instances, key=lambda ti: ti.task_id)
+    ti1, ti2 = sorted(dr.task_instances, key=lambda ti: ti.task_id)
     ti1.state = TaskInstanceState.SUCCESS
     session.add(
         TaskMap(dag_id=dr.dag_id, task_id=ti1.task_id, run_id=dr.run_id, map_index=-1, length=0, keys=None)
     )
     session.flush()
 
-    dr.task_instance_scheduling_decisions(session=session)
+    decision = dr.task_instance_scheduling_decisions(session=session)
+
+    # ti1 finished execution. ti2 goes directly to finished state because it's
+    # expanded against a zero-length XCom.
+    assert decision == TISchedulingDecision(
+        tis=[ti1, ti2],
+        schedulable_tis=[],
+        changed_tis=False,
+        unfinished_tis=[],
+        finished_tis=[ti1, ti2],
+    )
 
     indices = (
         session.query(TI.map_index, TI.state)