Merge pull request #1938 from apache/abderrahim/ensure-blobs

Ensure blobs are available when needed (when using storage-service)

src/buildstream/_cas/cascache.py

@@ -169,15 +169,9 @@ def contains_directory(self, digest, *, with_files):
     #     can_link (bool): Whether we can create hard links in the destination
     #
     def checkout(self, dest, tree, *, can_link=False, _fetch=True):
-        if _fetch and self._remote_cache:
+        if _fetch:
             # We need the files in the local cache
-            local_cas = self.get_local_cas()
-
-            request = local_cas_pb2.FetchTreeRequest()
-            request.root_digest.CopyFrom(tree)
-            request.fetch_file_blobs = True
-
-            local_cas.FetchTree(request)
+            self.ensure_tree(tree)
 
         os.makedirs(dest, exist_ok=True)
 
@@ -223,6 +217,24 @@ def checkout(self, dest, tree, *, can_link=False, _fetch=True):
             fullpath = os.path.join(dest, symlinknode.name)
             os.symlink(symlinknode.target, fullpath)
 
+    # ensure_tree():
+    #
+    # Make sure all blobs referenced by the given directory tree are available
+    # in the local cache when using a remote cache.
+    #
+    # Args:
+    #     tree (Digest): The digest of the tree
+    #
+    def ensure_tree(self, tree):
+        if self._remote_cache:
+            local_cas = self.get_local_cas()
+
+            request = local_cas_pb2.FetchTreeRequest()
+            request.root_digest.CopyFrom(tree)
+            request.fetch_file_blobs = True
+
+            local_cas.FetchTree(request)
+
     # pull_tree():
     #
     # Pull a single Tree rather than a ref.

src/buildstream/storage/_casbaseddirectory.py

@@ -185,6 +185,7 @@ def import_single_file(self, external_pathspec: str) -> FileListResult:
         return result
 
     def export_to_tar(self, tarfile: TarFile, destination_dir: str, mtime: int = BST_ARBITRARY_TIMESTAMP) -> None:
+        self._ensure_local()
         for filename, entry in sorted(self.__index.items()):
             arcname = os.path.join(destination_dir, filename)
             if entry.type == FileType.DIRECTORY:
@@ -396,6 +397,9 @@ def _export_files(self, to_directory: str, *, can_link: bool = False, can_destro
     def _set_deterministic_user(self) -> None:
         pass
 
+    def _ensure_local(self):
+        self.__cas_cache.ensure_tree(self._get_digest())
+
     def _get_underlying_path(self, filename) -> str:
         try:
             entry = self.__index[filename]

src/buildstream/storage/_filebaseddirectory.py

@@ -267,6 +267,9 @@ def _import_files(
             #
             assert isinstance(external_pathspec, Directory)
 
+            # Ensure blobs are available locally
+            external_pathspec._ensure_local()
+
             def copy_action(src_path, dest_path, mtime, result):
                 utils.safe_copy(src_path, dest_path, result=result)
                 utils._set_file_mtime(dest_path, mtime)

src/buildstream/storage/directory.py

@@ -466,6 +466,13 @@ def _import_files(
     def _export_files(self, to_directory: str, *, can_link: bool = False, can_destroy: bool = False) -> None:
         raise NotImplementedError()
 
+    # _ensure_local()
+    #
+    # Makes sure the files for the directory are available locally. Should be called before
+    # using _get_underlying_path()
+    def _ensure_local(self):
+        pass
+
     # _get_underlying_path()
     #
     # Args: