Add an ImageTree class and change to searching *all applicable* branches when looking for the best cache tag.

buildman/jobutil/buildjob.py

@@ -4,6 +4,7 @@ import logging
 from cachetools import lru_cache
 from endpoints.notificationhelper import spawn_notification
 from data import model
+from util.imagetree import ImageTree
 
 logger = logging.getLogger(__name__)
 
@@ -91,31 +92,31 @@ class BuildJob(object):
     repo_namespace = repo_build.repository.namespace_user.username
     repo_name = repo_build.repository.name
 
-    current_image = model.get_image(repo_build.repository, base_image_id)
-    next_image = None
-    if current_image is None:
+    base_image = model.get_image(repo_build.repository, base_image_id)
+    if base_image is None:
       return None
 
-    # For each cache comment, find a child image that matches the command.
-    for cache_command in cache_commands:
-      full_command = '["/bin/sh", "-c", "%s"]' % cache_command
-      next_image = model.find_child_image(repo_build.repository, current_image, full_command)
-      if next_image is None:
-        break
+    # Build an in-memory tree of the full heirarchy of images in the repository.
+    all_images = model.get_repository_images(repo_namespace, repo_name)
+    all_tags = model.list_repository_tags(repo_namespace, repo_name)
+    tree = ImageTree(all_images, all_tags, base_filter=base_image.id)
 
-      current_image = next_image
-      logger.debug('Found cached image %s for comment %s', current_image.id, full_command)
+    # Find a path in the tree, starting at the base image, that matches the cache comments
+    # or some subset thereof.
+    def checker(step, image):
+      if step >= len(cache_commands):
+        return False
 
-    # Find a tag associated with the image, if any.
-    # TODO(jschorr): We should just return the image ID instead of a parent tag, OR we should
-    # make this more efficient.
-    for tag in model.list_repository_tags(repo_namespace, repo_name):
-      tag_image = tag.image
-      ancestor_index = '/%s/' % current_image.id
-      if ancestor_index in tag_image.ancestors or tag_image.id == current_image.id:
-        return tag.name
+      full_command = '["/bin/sh", "-c", "%s"]' % cache_commands[step]
+      return image.storage.comment == full_command
+
+    path = tree.find_longest_path(base_image.id, checker)
+    if not path:
+      return None
+
+    # Find any tag associated with the last image in the path.
+    return tree.tag_containing_images(path[-1])
 
-    return None
 
   def _determine_cached_tag_by_tag(self):
     """ Determines the cached tag by looking for one of the tags being built, and seeing if it