_sandboxremote.py: Fetch blobs on artifact creation, not after command

We need to ensure artifact files are available locally also for elements
that don't run any commands (e.g., compose elements). It's also not
necessary to fetch file blobs after every command. This moves file blob
fetching from SandboxRemote.process_job_output() (part of Sandbox.run())
to a separate fetch_missing_blobs() method, which is invoked by
Element._cache_artifact().

3 files changed, 23 insertions, 19 deletions

diff --git a/src/buildstream/element.py b/src/buildstream/element.py
index 21c38bc1a..0d80113ab 100644
--- a/src/buildstream/element.py
+++ b/src/buildstream/element.py
@@ -1703,6 +1703,7 @@ class Element(Plugin):
             try:
                 sandbox_build_dir = sandbox_vroot.descend(
                     *self.get_variable('build-root').lstrip(os.sep).split(os.sep))
+                sandbox._fetch_missing_blobs(sandbox_build_dir)
             except VirtualDirectoryError:
                 # Directory could not be found. Pre-virtual
                 # directory behaviour was to continue silently
@@ -1712,6 +1713,7 @@ class Element(Plugin):
         if collect is not None:
             try:
                 collectvdir = sandbox_vroot.descend(*collect.lstrip(os.sep).split(os.sep))
+                sandbox._fetch_missing_blobs(collectvdir)
             except VirtualDirectoryError:
                 pass
 
diff --git a/src/buildstream/sandbox/_sandboxremote.py b/src/buildstream/sandbox/_sandboxremote.py
index c84bfa4ef..8c9515497 100644
--- a/src/buildstream/sandbox/_sandboxremote.py
+++ b/src/buildstream/sandbox/_sandboxremote.py
@@ -30,7 +30,6 @@ from .. import utils
 from ..node import Node
 from .._message import Message, MessageType
 from .sandbox import Sandbox, SandboxCommandError, _SandboxBatch
-from ..storage.directory import VirtualDirectoryError
 from ..storage._casbaseddirectory import CasBasedDirectory
 from .. import _signals
 from .._protos.build.bazel.remote.execution.v2 import remote_execution_pb2, remote_execution_pb2_grpc
@@ -280,7 +279,6 @@ class SandboxRemote(Sandbox):
             raise SandboxError("Output directory structure had no digest attached.")
 
         context = self._get_context()
-        project = self._get_project()
         cascache = context.get_cascache()
         artifactcache = context.artifactcache
         casremote = CASRemote(self.storage_remote_spec)
@@ -294,27 +292,20 @@ class SandboxRemote(Sandbox):
         # to replace the sandbox's virtual directory with that. Creating a new virtual directory object
         # from another hash will be interesting, though...
 
-        new_dir = CasBasedDirectory(context.artifactcache.cas, digest=dir_digest)
+        new_dir = CasBasedDirectory(artifactcache.cas, digest=dir_digest)
         self._set_virtual_directory(new_dir)
 
+    def _fetch_missing_blobs(self, vdir):
+        context = self._get_context()
+        project = self._get_project()
+        cascache = context.get_cascache()
+        artifactcache = context.artifactcache
+        casremote = CASRemote(self.storage_remote_spec)
+
         # Fetch the file blobs if needed
         if self._output_files_required or artifactcache.has_push_remotes():
-            required_blobs = []
-            directories = []
-
-            directories.append(self._output_directory)
-            if self._build_directory and (self._build_directory_always or failure):
-                directories.append(self._build_directory)
-
-            for directory in directories:
-                try:
-                    vdir = new_dir.descend(*directory.strip(os.sep).split(os.sep))
-                    dir_digest = vdir._get_digest()
-                    required_blobs += cascache.required_blobs_for_directory(dir_digest)
-                except VirtualDirectoryError:
-                    # If the directory does not exist, there is no need to
-                    # download file blobs.
-                    pass
+            dir_digest = vdir._get_digest()
+            required_blobs = cascache.required_blobs_for_directory(dir_digest)
 
             local_missing_blobs = cascache.local_missing_blobs(required_blobs)
             if local_missing_blobs:
diff --git a/src/buildstream/sandbox/sandbox.py b/src/buildstream/sandbox/sandbox.py
index ece15c949..8ff0aafa2 100644
--- a/src/buildstream/sandbox/sandbox.py
+++ b/src/buildstream/sandbox/sandbox.py
@@ -410,6 +410,17 @@ class Sandbox():
 
         return 'BST_CAS_DIRECTORIES' in os.environ
 
+    # _fetch_missing_blobs()
+    #
+    # Fetch required file blobs missing from the local cache for sandboxes using
+    # remote execution. This is a no-op for local sandboxes.
+    #
+    # Args:
+    #    vdir (Directory): The virtual directory whose blobs to fetch
+    #
+    def _fetch_missing_blobs(self, vdir):
+        pass
+
     ################################################
     #               Private methods                #
     ################################################