[buildstream] 02/02: _elementsources.py: Use `StageTree` to stage previous sources

[ju...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

juergbi pushed a commit to branch juerg/opt
in repository https://gitbox.apache.org/repos/asf/buildstream.git

commit cbc8e791fc93b1f4028996595856fc0a670c9244
Author: Jürg Billeter <j...@bitron.ch>
AuthorDate: Fri Jul 8 14:32:09 2022 +0200

    _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.
---
 src/buildstream/_cas/cascache.py   | 29 +++++++++++++++++++++++++++++
 src/buildstream/_elementsources.py | 10 +++++-----
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/src/buildstream/_cas/cascache.py b/src/buildstream/_cas/cascache.py
index c6a58098f..6ac3da557 100644
--- a/src/buildstream/_cas/cascache.py
+++ b/src/buildstream/_cas/cascache.py
@@ -443,6 +443,35 @@ class CASCache:
 
         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 f089d17b1..315e973c0 100644
--- a/src/buildstream/_elementsources.py
+++ b/src/buildstream/_elementsources.py
@@ -459,7 +459,8 @@ class ElementSources:
     #   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 @@ class ElementSources:
                 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