[buildstream] 01/02: cascache.py: Add stage_directory() method

[ju...@apache.org]

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

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

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

    cascache.py: Add stage_directory() method
---
 src/buildstream/_cas/cascache.py | 42 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/src/buildstream/_cas/cascache.py b/src/buildstream/_cas/cascache.py
index c6a58098f..f9574802c 100644
--- a/src/buildstream/_cas/cascache.py
+++ b/src/buildstream/_cas/cascache.py
@@ -443,6 +443,48 @@ class CASCache:
 
         return utils._message_digest(root_directory)
 
+    # stage_directory():
+    #
+    # A contextmanager to stage a CAS directory tree in the local filesystem.
+    #
+    # This makes the specified directory tree temporarily available for local
+    # filesystem access. This may use FUSE or hardlinking.
+    #
+    # Args:
+    #     directory_digest (Digest): The digest of a directory
+    #
+    # Yields:
+    #     (str): The local filesystem path
+    #
+    @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.