[AIRFLOW-4797] Use same zombies in all DAG file processors

(cherry picked from commit cb0dbe3)

Author: KevinYang21

airflow/models/dagbag.py

@@ -30,8 +30,8 @@
 from collections import namedtuple
 from datetime import datetime
 
+from croniter import CroniterBadCronError, CroniterBadDateError, CroniterNotAlphaError, croniter
 import six
-from croniter import croniter, CroniterBadCronError, CroniterBadDateError, CroniterNotAlphaError
 
 from airflow import settings
 from airflow.configuration import conf

airflow/utils/dag_processing.py

@@ -40,6 +40,7 @@
 from setproctitle import setproctitle
 import six
 from six.moves import reload_module
+from sqlalchemy import or_
 from tabulate import tabulate
 
 # To avoid circular imports
@@ -54,7 +55,6 @@
 from airflow.utils.db import provide_session
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.state import State
-from sqlalchemy import or_
 
 if six.PY2:
     ConnectionError = IOError
@@ -776,13 +776,14 @@ def __init__(self,
         # Map from file path to stats about the file
         self._file_stats = {}  # type: dict(str, DagFileStat)
 
-        self._last_zombie_query_time = timezone.utcnow()
+        self._last_zombie_query_time = None
         # Last time that the DAG dir was traversed to look for files
         self.last_dag_dir_refresh_time = timezone.utcnow()
         # Last time stats were printed
         self.last_stat_print_time = timezone.datetime(2000, 1, 1)
         # TODO: Remove magic number
         self._zombie_query_interval = 10
+        self._zombies = []
         # How long to wait before timing out a process to parse a DAG file
         self._processor_timeout = processor_timeout
 
@@ -852,6 +853,7 @@ def start(self):
                 continue
 
             self._refresh_dag_dir()
+            self._find_zombies()
 
             simple_dags = self.heartbeat()
             for simple_dag in simple_dags:
@@ -1247,13 +1249,11 @@ def heartbeat(self):
 
             self._file_path_queue.extend(files_paths_to_queue)
 
-        zombies = self._find_zombies()
-
         # Start more processors if we have enough slots and files to process
         while (self._parallelism - len(self._processors) > 0 and
                len(self._file_path_queue) > 0):
             file_path = self._file_path_queue.pop(0)
-            processor = self._processor_factory(file_path, zombies)
+            processor = self._processor_factory(file_path, self._zombies)
             Stats.incr('dag_processing.processes')
 
             processor.start()
@@ -1271,13 +1271,13 @@ def heartbeat(self):
     @provide_session
     def _find_zombies(self, session):
         """
-        Find zombie task instances, which are tasks haven't heartbeated for too long.
-        :return: Zombie task instances in SimpleTaskInstance format.
+        Find zombie task instances, which are tasks haven't heartbeated for too long
+        and update the current zombie list.
         """
         now = timezone.utcnow()
         zombies = []
-        if (now - self._last_zombie_query_time).total_seconds() \
-                > self._zombie_query_interval:
+        if not self._last_zombie_query_time or \
+                (now - self._last_zombie_query_time).total_seconds() > self._zombie_query_interval:
             # to avoid circular imports
             from airflow.jobs import LocalTaskJob as LJ
             self.log.info("Finding 'running' jobs without a recent heartbeat")
@@ -1305,7 +1305,7 @@ def _find_zombies(self, session):
                     sti.dag_id, sti.task_id, sti.execution_date.isoformat())
                 zombies.append(sti)
 
-        return zombies
+            self._zombies = zombies
 
     def _kill_timed_out_processors(self):
         """

tests/models/test_dagbag.py

@@ -17,13 +17,13 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from datetime import datetime
 import inspect
 import os
 import shutil
 import textwrap
 import unittest
-from tempfile import mkdtemp, NamedTemporaryFile
+from datetime import datetime
+from tempfile import NamedTemporaryFile, mkdtemp
 
 from mock import patch, ANY
 

tests/utils/test_dag_processing.py

@@ -27,14 +27,16 @@
 from mock import MagicMock, PropertyMock
 
 from airflow.configuration import conf, mkdir_p
-from airflow.jobs import DagFileProcessor
-from airflow.jobs import LocalTaskJob as LJ
+from airflow.jobs import DagFileProcessor, LocalTaskJob as LJ
 from airflow.models import DagBag, TaskInstance as TI
 from airflow.utils import timezone
-from airflow.utils.dag_processing import (DagFileProcessorAgent, DagFileProcessorManager,
-                                          DagFileStat, SimpleTaskInstance, correct_maybe_zipped)
+from airflow.utils.dag_processing import (
+    DagFileProcessorAgent, DagFileProcessorManager, DagFileStat, SimpleTaskInstance, correct_maybe_zipped,
+)
 from airflow.utils.db import create_session
 from airflow.utils.state import State
+from tests.test_utils.config import conf_vars
+from tests.test_utils.db import clear_db_runs
 
 TEST_DAG_FOLDER = os.path.join(
     os.path.dirname(os.path.realpath(__file__)), os.pardir, 'dags')
@@ -133,6 +135,9 @@ def __exit__(self, *exc_info):
 
 
 class TestDagFileProcessorManager(unittest.TestCase):
+    def setUp(self):
+        clear_db_runs()
+
     def test_set_file_paths_when_processor_file_path_not_in_new_file_paths(self):
         manager = DagFileProcessorManager(
             dag_directory='directory',
@@ -202,7 +207,8 @@ def test_find_zombies(self):
 
             manager._last_zombie_query_time = timezone.utcnow() - timedelta(
                 seconds=manager._zombie_threshold_secs + 1)
-            zombies = manager._find_zombies()
+            manager._find_zombies()
+            zombies = manager._zombies
             self.assertEqual(1, len(zombies))
             self.assertIsInstance(zombies[0], SimpleTaskInstance)
             self.assertEqual(ti.dag_id, zombies[0].dag_id)
@@ -212,6 +218,87 @@ def test_find_zombies(self):
             session.query(TI).delete()
             session.query(LJ).delete()
 
+    def test_zombies_are_correctly_passed_to_dag_file_processor(self):
+        """
+        Check that the same set of zombies are passed to the dag
+        file processors until the next zombie detection logic is invoked.
+        """
+        with conf_vars({('scheduler', 'max_threads'): '1',
+                        ('core', 'load_examples'): 'False'}):
+            dagbag = DagBag(os.path.join(TEST_DAG_FOLDER, 'test_example_bash_operator.py'))
+            with create_session() as session:
+                session.query(LJ).delete()
+                dag = dagbag.get_dag('test_example_bash_operator')
+                task = dag.get_task(task_id='run_this_last')
+
+                ti = TI(task, DEFAULT_DATE, State.RUNNING)
+                lj = LJ(ti)
+                lj.state = State.SHUTDOWN
+                lj.id = 1
+                ti.job_id = lj.id
+
+                session.add(lj)
+                session.add(ti)
+                session.commit()
+                fake_zombies = [SimpleTaskInstance(ti)]
+
+            class FakeDagFIleProcessor(DagFileProcessor):
+                # This fake processor will return the zombies it received in constructor
+                # as its processing result w/o actually parsing anything.
+                def __init__(self, file_path, pickle_dags, dag_id_white_list, zombies):
+                    super(FakeDagFIleProcessor, self).__init__(
+                        file_path, pickle_dags, dag_id_white_list, zombies
+                    )
+
+                    self._result = zombies, 0
+
+                def start(self):
+                    pass
+
+                @property
+                def start_time(self):
+                    return DEFAULT_DATE
+
+                @property
+                def pid(self):
+                    return 1234
+
+                @property
+                def done(self):
+                    return True
+
+                @property
+                def result(self):
+                    return self._result
+
+            def processor_factory(file_path, zombies):
+                return FakeDagFIleProcessor(file_path,
+                                            False,
+                                            [],
+                                            zombies)
+
+            test_dag_path = os.path.join(TEST_DAG_FOLDER,
+                                         'test_example_bash_operator.py')
+            async_mode = 'sqlite' not in conf.get('core', 'sql_alchemy_conn')
+            processor_agent = DagFileProcessorAgent(test_dag_path,
+                                                    [],
+                                                    1,
+                                                    processor_factory,
+                                                    timedelta.max,
+                                                    async_mode)
+            processor_agent.start()
+            parsing_result = []
+            if not async_mode:
+                processor_agent.heartbeat()
+            while not processor_agent.done:
+                if not async_mode:
+                    processor_agent.wait_until_finished()
+                parsing_result.extend(processor_agent.harvest_simple_dags())
+
+            self.assertEqual(len(fake_zombies), len(parsing_result))
+            self.assertEqual(set([zombie.key for zombie in fake_zombies]),
+                             set([result.key for result in parsing_result]))
+
     @mock.patch("airflow.jobs.DagFileProcessor.pid", new_callable=PropertyMock)
     @mock.patch("airflow.jobs.DagFileProcessor.kill")
     def test_kill_timed_out_processors_kill(self, mock_kill, mock_pid):