feat(tvix/store/directorysvc): add put_multiple_start r/6073

This provides a handle to upload multiple proto::Directory as part of
the same closure.

Change-Id: I9213dde257a260c8622239918ea541064b270484
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8356
Tested-by: BuildkiteCI
Reviewed-by: tazjin <tazjin@tvl.su>
Autosubmit: flokli <flokli@flokli.de>

1 files changed, 38 insertions, 0 deletions

diff --git a/tvix/store/src/directoryservice/utils.rs b/tvix/store/src/directoryservice/utils.rs
index 10edb9ba74dc..7d41e8d3ba01 100644
--- a/tvix/store/src/directoryservice/utils.rs
+++ b/tvix/store/src/directoryservice/utils.rs
@@ -1,3 +1,4 @@
+use super::DirectoryPutter;
 use super::DirectoryService;
 use crate::proto;
 use crate::Error;
@@ -104,3 +105,40 @@ impl<DS: DirectoryService> Iterator for DirectoryTraverser<DS> {
         }
     }
 }
+
+/// This is a simple implementation of a Directory uploader.
+/// TODO: verify connectivity? Factor out these checks into generic helpers?
+pub struct SimplePutter<DS: DirectoryService> {
+    directory_service: DS,
+    last_directory_digest: Option<[u8; 32]>,
+}
+
+impl<DS: DirectoryService> SimplePutter<DS> {
+    pub fn new(directory_service: DS) -> Self {
+        Self {
+            directory_service,
+            last_directory_digest: None,
+        }
+    }
+}
+
+impl<DS: DirectoryService> DirectoryPutter for SimplePutter<DS> {
+    fn put(&mut self, directory: proto::Directory) -> Result<(), Error> {
+        let digest = self.directory_service.put(directory)?;
+
+        // track the last directory digest
+        self.last_directory_digest = Some(digest);
+
+        Ok(())
+    }
+
+    /// We need to be mutable here, as that's the signature of the trait.
+    fn close(&mut self) -> Result<[u8; 32], Error> {
+        match self.last_directory_digest {
+            Some(last_digest) => Ok(last_digest),
+            None => Err(Error::InvalidRequest(
+                "no directories sent, can't show root digest".to_string(),
+            )),
+        }
+    }
+}