From 52992812ed4a621ab60f08a0674e724cf052e293 Mon Sep 17 00:00:00 2001
From: Abderrahim Kitouni <abderrahim.kitouni@codethink.co.uk>
Date: Thu, 1 Aug 2024 13:23:44 +0100
Subject: [PATCH 1/2] cascache: add a helper method to ensure a tree is cached
 locally

---
 src/buildstream/_cas/cascache.py | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/src/buildstream/_cas/cascache.py b/src/buildstream/_cas/cascache.py
index 627778c16..66d861aec 100644
--- a/src/buildstream/_cas/cascache.py
+++ b/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.

From f797f1870aff3f35a20fb8e1b543b85802119c71 Mon Sep 17 00:00:00 2001
From: Abderrahim Kitouni <abderrahim.kitouni@codethink.co.uk>
Date: Thu, 1 Aug 2024 14:28:06 +0100
Subject: [PATCH 2/2] storage: ensure blobs are available when using
 CasBasedDirectory with storage service

---
 src/buildstream/storage/_casbaseddirectory.py  | 4 ++++
 src/buildstream/storage/_filebaseddirectory.py | 3 +++
 src/buildstream/storage/directory.py           | 7 +++++++
 3 files changed, 14 insertions(+)

diff --git a/src/buildstream/storage/_casbaseddirectory.py b/src/buildstream/storage/_casbaseddirectory.py
index 9c2e83a7b..80e9b3dfc 100644
--- a/src/buildstream/storage/_casbaseddirectory.py
+++ b/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]
diff --git a/src/buildstream/storage/_filebaseddirectory.py b/src/buildstream/storage/_filebaseddirectory.py
index 93fc17768..87749fda2 100644
--- a/src/buildstream/storage/_filebaseddirectory.py
+++ b/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)
diff --git a/src/buildstream/storage/directory.py b/src/buildstream/storage/directory.py
index fd19393d0..3cd64a966 100644
--- a/src/buildstream/storage/directory.py
+++ b/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: