Cache tasks by uuid in PlayIterator for O(1) lookupsfeature_improve_playiterator_task_lookup

Rather than repeatedly searching for tasks by uuid via iterating over
all known blocks, cache the tasks when they are added to the PlayIterator
so the lookup becomes a simple key check in a dict.

3 files changed, 25 insertions, 43 deletions

diff --git a/lib/ansible/executor/play_iterator.py b/lib/ansible/executor/play_iterator.py
index d1366bbbc5..a8c00833e7 100644
--- a/lib/ansible/executor/play_iterator.py
+++ b/lib/ansible/executor/play_iterator.py
@@ -154,6 +154,8 @@ class PlayIterator:
         self._play = play
         self._blocks = []
 
+        self._task_uuid_cache = dict()
+
         # Default options to gather
         gather_subset = C.DEFAULT_GATHER_SUBSET
         gather_timeout = C.DEFAULT_GATHER_TIMEOUT
@@ -179,12 +181,17 @@ class PlayIterator:
 
         setup_block = setup_block.filter_tagged_tasks(play_context, all_vars)
         self._blocks.append(setup_block)
+        self.cache_block_tasks(setup_block)
 
         for block in self._play.compile():
             new_block = block.filter_tagged_tasks(play_context, all_vars)
             if new_block.has_tasks():
+                self.cache_block_tasks(new_block)
                 self._blocks.append(new_block)
 
+        for handler_block in self._play.handlers:
+            self.cache_block_tasks(handler_block)
+
         self._host_states = {}
         start_at_matched = False
         for host in inventory.get_hosts(self._play.hosts):
@@ -227,6 +234,18 @@ class PlayIterator:
 
         return self._host_states[host.name].copy()
 
+    def cache_block_tasks(self, block):
+        def _cache_portion(p):
+            for t in p:
+                if isinstance(t, Block):
+                    self.cache_block_tasks(t)
+                elif t._uuid not in self._task_uuid_cache:
+                    self._task_uuid_cache[t._uuid] = t
+
+        for portion in (block.block, block.rescue, block.always):
+            if portion is not None:
+                _cache_portion(portion)
+
     def get_next_task_for_host(self, host, peek=False):
 
         display.debug("getting the next task for host %s" % host.name)
@@ -514,46 +533,7 @@ class PlayIterator:
         else:
             the_uuid = task
 
-        def _search_block(block):
-            '''
-            helper method to check a block's task lists (block/rescue/always)
-            for a given task uuid. If a Block is encountered in the place of a
-            task, it will be recursively searched (this happens when a task
-            include inserts one or more blocks into a task list).
-            '''
-            for b in (block.block, block.rescue, block.always):
-                for t in b:
-                    if isinstance(t, Block):
-                        res = _search_block(t)
-                        if res:
-                            return res
-                    elif t._uuid == the_uuid:
-                        return t
-            return None
-
-        def _search_state(state):
-            for block in state._blocks:
-                res = _search_block(block)
-                if res:
-                    return res
-            for child_state in (state.tasks_child_state, state.rescue_child_state, state.always_child_state):
-                if child_state is not None:
-                    res = _search_state(child_state)
-                    if res:
-                        return res
-            return None
-
-        s = self.get_host_state(host)
-        res = _search_state(s)
-        if res:
-            return res
-
-        for block in self._play.handlers:
-            res = _search_block(block)
-            if res:
-                return res
-
-        return None
+        return self._task_uuid_cache.get(the_uuid, None)
 
     def _insert_tasks_into_state(self, state, task_list):
         # if we've failed at all, or if the task list is empty, just return the current state
@@ -590,5 +570,7 @@ class PlayIterator:
         return state
 
     def add_tasks(self, host, task_list):
+        for b in task_list:
+            self.cache_block_tasks(b)
         self._host_states[host.name] = self._insert_tasks_into_state(self.get_host_state(host), task_list)
 
diff --git a/test/integration/Makefile b/test/integration/Makefile
index cc5d26a02e..bf861f641d 100644
--- a/test/integration/Makefile
+++ b/test/integration/Makefile
@@ -202,13 +202,13 @@ blocks: setup
 	# remove old output log
 	rm -f block_test.out
 	# run the test and check to make sure the right number of completions was logged
-	ansible-playbook -vv -e outputdir=$(TEST_DIR) test_blocks/main.yml | tee block_test.out
+	ansible-playbook -vv $(TEST_FLAGS) -e outputdir=$(TEST_DIR) test_blocks/main.yml | tee block_test.out
 	env python -c 'import sys, re; sys.stdout.write(re.sub("\x1B\[([0-9]{1,2}(;[0-9]{1,2})?)?[m|K]", "", sys.stdin.read()))' <block_test.out >block_test_wo_colors.out
 	[ "$$(grep 'TEST COMPLETE' block_test.out | wc -l | sed 's/ *//')" = "$$(egrep '^[0-9]+ plays in' block_test_wo_colors.out | cut -f1 -d' ')" ]
 	# cleanup the output log again, to make sure the test is clean
 	rm -f block_test.out block_test_wo_colors.out
 	# run test with free strategy and again count the completions
-	ansible-playbook -vv -e outputdir=$(TEST_DIR) test_blocks/main.yml -e test_strategy=free | tee block_test.out
+	ansible-playbook -vv $(TEST_FLAGS) -e outputdir=$(TEST_DIR) test_blocks/main.yml -e test_strategy=free | tee block_test.out
 	env python -c 'import sys, re; sys.stdout.write(re.sub("\x1B\[([0-9]{1,2}(;[0-9]{1,2})?)?[m|K]", "", sys.stdin.read()))' <block_test.out >block_test_wo_colors.out
 	[ "$$(grep 'TEST COMPLETE' block_test.out | wc -l | sed 's/ *//')" = "$$(egrep '^[0-9]+ plays in' block_test_wo_colors.out | cut -f1 -d' ')" ]
 
diff --git a/test/units/executor/test_play_iterator.py b/test/units/executor/test_play_iterator.py
index c0f23e531d..671deaf279 100644
--- a/test/units/executor/test_play_iterator.py
+++ b/test/units/executor/test_play_iterator.py
@@ -329,7 +329,7 @@ class TestPlayIterator(unittest.TestCase):
         # test the high-level add_tasks() method
         s = HostState(blocks=[0,1,2])
         itr._insert_tasks_into_state = MagicMock(return_value=s)
-        itr.add_tasks(hosts[0], [3,4,5])
+        itr.add_tasks(hosts[0], [MagicMock(), MagicMock(), MagicMock()])
         self.assertEqual(itr._host_states[hosts[0].name], s)
 
         # now actually test the lower-level method that does the work