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

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