_elementsources.py: Use `StageTree` to stage previous sources
This will use buildbox-fuse (via buildbox-casd) if available, improving
performance of both stage and capture operations.
diff --git a/src/buildstream/_cas/cascache.py b/src/buildstream/_cas/cascache.py
index c6a5809..6ac3da5 100644
--- a/src/buildstream/_cas/cascache.py
+++ b/src/buildstream/_cas/cascache.py
@@ -443,6 +443,35 @@
return utils._message_digest(root_directory)
+ @contextlib.contextmanager
+ def stage_directory(self, directory_digest):
+ local_cas = self.get_local_cas()
+
+ request = local_cas_pb2.StageTreeRequest()
+ request.root_digest.CopyFrom(directory_digest)
+
+ done_event = threading.Event()
+
+ def request_iterator():
+ yield request
+
+ # Wait until staged tree is no longer needed
+ done_event.wait()
+
+ # A second (empty) request indicates that the staging location can be cleaned up
+ yield local_cas_pb2.StageTreeRequest()
+
+ response_stream = local_cas.StageTree(request_iterator())
+
+ # Read primary response and yield staging location
+ response = next(response_stream)
+ yield response.path
+
+ # Staged tree is no longer needed
+ done_event.set()
+ # Wait for cleanup to complete
+ next(response_stream)
+
# missing_blobs_for_directory():
#
# Determine which blobs of a directory tree are missing on the remote.
diff --git a/src/buildstream/_elementsources.py b/src/buildstream/_elementsources.py
index f089d17..315e973 100644
--- a/src/buildstream/_elementsources.py
+++ b/src/buildstream/_elementsources.py
@@ -459,7 +459,8 @@
# stop (Source): Only stage sources listed before this source
#
def _stage(self, *, stop=None):
- vdir = CasBasedDirectory(self._context.get_cascache())
+ cas = self._context.get_cascache()
+ vdir = CasBasedDirectory(cas)
for source in self._sources:
if source == stop:
@@ -474,10 +475,9 @@
if source.BST_STAGE_VIRTUAL_DIRECTORY:
source._stage(vsubdir)
else:
- with utils._tempdir(dir=self._context.tmpdir, prefix="staging-temp") as tmpdir:
- # Stage previous sources
- vsubdir._export_files(tmpdir)
-
+ # Stage previous sources
+ with cas.stage_directory(vsubdir._get_digest()) as tmpdir:
+ # Stage current source
source._stage(tmpdir)
# Capture modified tree
