Merge pull request #1966 from apache/juerg/uploadtree

cascache.py: Use `UploadTree` with `storage-service`, if available

src/buildstream/_artifact.py

@@ -339,7 +339,7 @@ def cached_buildroot(self):
 
         buildroot_digest = self._get_field_digest("buildroot")
         if buildroot_digest:
-            return self._cas.contains_directory(buildroot_digest, with_files=True)
+            return self._cas.contains_directory(buildroot_digest)
         else:
             return False
 
@@ -371,7 +371,7 @@ def cached_buildtree(self):
 
         buildtree_digest = self._get_field_digest("buildtree")
         if buildtree_digest:
-            return self._cas.contains_directory(buildtree_digest, with_files=True)
+            return self._cas.contains_directory(buildtree_digest)
         else:
             return False
 
@@ -400,7 +400,7 @@ def cached_sources(self):
 
         sources_digest = self._get_field_digest("sources")
         if sources_digest:
-            return self._cas.contains_directory(sources_digest, with_files=True)
+            return self._cas.contains_directory(sources_digest)
         else:
             return False
 
@@ -602,7 +602,7 @@ def query_cache(self):
             return False
 
         # Check whether 'files' subdirectory is available, with or without file contents
-        if str(artifact.files) and not self._cas.contains_directory(artifact.files, with_files=True):
+        if str(artifact.files) and not self._cas.contains_directory(artifact.files):
             self._cached = False
             return False
 

src/buildstream/_cas/cascache.py

@@ -124,24 +124,23 @@ def contains_files(self, digests):
 
     # contains_directory():
     #
-    # Check whether the specified directory and subdirectories are in the cache,
-    # i.e non dangling.
+    # Check whether the specified directory, subdirectories and files are in the
+    # cache, i.e non dangling.
     #
     # Args:
     #     digest (Digest): The directory digest to check
-    #     with_files (bool): Whether to check files as well
     #
     # Returns: True if the directory is available in the local cache
     #
-    def contains_directory(self, digest, *, with_files):
+    def contains_directory(self, digest):
         local_cas = self.get_local_cas()
 
         # Without a remote cache, `FetchTree` simply checks the local cache.
         request = local_cas_pb2.FetchTreeRequest()
         request.root_digest.CopyFrom(digest)
         # Always fetch Directory protos as they are needed to enumerate subdirectories and files.
         # Don't implicitly fetch file blobs from the remote cache as we don't need them.
-        request.fetch_file_blobs = with_files and not self._remote_cache
+        request.fetch_file_blobs = not self._remote_cache
 
         try:
             local_cas.FetchTree(request)
@@ -154,9 +153,20 @@ def contains_directory(self, digest, *, with_files):
                 raise CASCacheError("Unsupported buildbox-casd version: FetchTree unimplemented") from e
             raise
 
-        # Check whether everything is available in the remote cache.
-        missing_blobs = self.missing_blobs_for_directory(digest, remote=self._default_remote)
-        return not missing_blobs
+        # Make sure everything is available in the remote cache (storage-service)
+        request = local_cas_pb2.UploadTreeRequest()
+        request.root_digest.CopyFrom(digest)
+        try:
+            local_cas.UploadTree(request)
+            return True
+        except grpc.RpcError as e:
+            if e.code() == grpc.StatusCode.NOT_FOUND:
+                return False
+            if e.code() == grpc.StatusCode.UNIMPLEMENTED:
+                # Fallback path if buildbox-casd is too old to support UploadTree
+                missing_blobs = self.missing_blobs_for_directory(digest, remote=self._default_remote)
+                return not missing_blobs
+            raise
 
     # checkout():
     #

src/buildstream/_elementsources.py

@@ -308,7 +308,7 @@ def query_cache(self):
             self._cached = False
             return False
 
-        if not cas.contains_directory(source_proto.files, with_files=True):
+        if not cas.contains_directory(source_proto.files):
             self._cached = False
             return False
 

src/buildstream/_sourcecache.py

@@ -58,7 +58,7 @@ def contains(self, source):
 
         # check files
         source_proto = self._get_source(ref)
-        return self.cas.contains_directory(source_proto.files, with_files=True)
+        return self.cas.contains_directory(source_proto.files)
 
     # commit()
     #