feat(tvix/store/pathinfosvc): add calculate_nar method r/6271

Putting this in the PathInfoService trait makes much more sense, we can
have direct control over where/how to cache the results in the
implementation.

This now requires each PathInfoService to hold pointers to BlobService
and DirectoryService.

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

4 files changed, 102 insertions, 14 deletions

diff --git a/tvix/store/src/pathinfoservice/grpc.rs b/tvix/store/src/pathinfoservice/grpc.rs
index 6bb774c668a3..871c3b592256 100644
--- a/tvix/store/src/pathinfoservice/grpc.rs
+++ b/tvix/store/src/pathinfoservice/grpc.rs
@@ -78,4 +78,36 @@ impl PathInfoService for GRPCPathInfoService {
             .block_on(task)?
             .map_err(|e| crate::Error::StorageError(e.to_string()))
     }
+
+    fn calculate_nar(
+        &self,
+        root_node: &proto::node::Node,
+    ) -> Result<(u64, [u8; 32]), crate::Error> {
+        // Get a new handle to the gRPC client.
+        let mut grpc_client = self.grpc_client.clone();
+        let root_node = root_node.clone();
+
+        let task: tokio::task::JoinHandle<Result<_, Status>> =
+            self.tokio_handle.spawn(async move {
+                let path_info = grpc_client
+                    .calculate_nar(proto::Node {
+                        node: Some(root_node),
+                    })
+                    .await?
+                    .into_inner();
+                Ok(path_info)
+            });
+
+        let resp = self
+            .tokio_handle
+            .block_on(task)?
+            .map_err(|e| crate::Error::StorageError(e.to_string()))?;
+
+        let nar_sha256: [u8; 32] = resp
+            .nar_sha256
+            .try_into()
+            .map_err(|_e| crate::Error::StorageError("invalid digest length".to_string()))?;
+
+        Ok((resp.nar_size, nar_sha256))
+    }
 }
diff --git a/tvix/store/src/pathinfoservice/memory.rs b/tvix/store/src/pathinfoservice/memory.rs
index d0ff1976efab..5b48ed9efa34 100644
--- a/tvix/store/src/pathinfoservice/memory.rs
+++ b/tvix/store/src/pathinfoservice/memory.rs
@@ -1,16 +1,31 @@
 use super::PathInfoService;
-use crate::{proto, Error};
+use crate::{
+    blobservice::BlobService, directoryservice::DirectoryService, nar::calculate_size_and_sha256,
+    proto, Error,
+};
 use std::{
     collections::HashMap,
     sync::{Arc, RwLock},
 };
 
-#[derive(Default)]
-pub struct MemoryPathInfoService {
+pub struct MemoryPathInfoService<DS: DirectoryService> {
     db: Arc<RwLock<HashMap<[u8; 20], proto::PathInfo>>>,
+
+    blob_service: Box<dyn BlobService>,
+    directory_service: DS,
+}
+
+impl<DS: DirectoryService> MemoryPathInfoService<DS> {
+    pub fn new(blob_service: Box<dyn BlobService>, directory_service: DS) -> Self {
+        Self {
+            db: Default::default(),
+            blob_service,
+            directory_service,
+        }
+    }
 }
 
-impl PathInfoService for MemoryPathInfoService {
+impl<DS: DirectoryService + Clone> PathInfoService for MemoryPathInfoService<DS> {
     fn get(&self, digest: [u8; 20]) -> Result<Option<proto::PathInfo>, Error> {
         let db = self.db.read().unwrap();
 
@@ -38,4 +53,13 @@ impl PathInfoService for MemoryPathInfoService {
             }
         }
     }
+
+    fn calculate_nar(&self, root_node: &proto::node::Node) -> Result<(u64, [u8; 32]), Error> {
+        calculate_size_and_sha256(
+            root_node,
+            &self.blob_service,
+            self.directory_service.clone(),
+        )
+        .map_err(|e| Error::StorageError(e.to_string()))
+    }
 }
diff --git a/tvix/store/src/pathinfoservice/mod.rs b/tvix/store/src/pathinfoservice/mod.rs
index ddede5851575..2483909a1190 100644
--- a/tvix/store/src/pathinfoservice/mod.rs
+++ b/tvix/store/src/pathinfoservice/mod.rs
@@ -18,5 +18,8 @@ pub trait PathInfoService {
     /// invalid messages.
     fn put(&self, path_info: proto::PathInfo) -> Result<proto::PathInfo, Error>;
 
-    // TODO: add default impl for nar calculation here, and override from GRPC client!
+    /// Return the nar size and nar sha256 digest for a given root node.
+    /// This can be used to calculate NAR-based output paths,
+    /// and implementations are encouraged to cache it.
+    fn calculate_nar(&self, root_node: &proto::node::Node) -> Result<(u64, [u8; 32]), Error>;
 }
diff --git a/tvix/store/src/pathinfoservice/sled.rs b/tvix/store/src/pathinfoservice/sled.rs
index 8776ebcbc106..98ea60ff4440 100644
--- a/tvix/store/src/pathinfoservice/sled.rs
+++ b/tvix/store/src/pathinfoservice/sled.rs
@@ -1,5 +1,8 @@
 use super::PathInfoService;
-use crate::{proto, Error};
+use crate::{
+    blobservice::BlobService, directoryservice::DirectoryService, nar::calculate_size_and_sha256,
+    proto, Error,
+};
 use prost::Message;
 use std::path::PathBuf;
 use tracing::warn;
@@ -8,28 +11,45 @@ use tracing::warn;
 ///
 /// The PathInfo messages are stored as encoded protos, and keyed by their output hash,
 /// as that's currently the only request type available.
-#[derive(Clone)]
-pub struct SledPathInfoService {
+pub struct SledPathInfoService<DS: DirectoryService> {
     db: sled::Db,
+
+    blob_service: Box<dyn BlobService>,
+    directory_service: DS,
 }
 
-impl SledPathInfoService {
-    pub fn new(p: PathBuf) -> Result<Self, sled::Error> {
+impl<DS: DirectoryService> SledPathInfoService<DS> {
+    pub fn new(
+        p: PathBuf,
+        blob_service: Box<dyn BlobService>,
+        directory_service: DS,
+    ) -> Result<Self, sled::Error> {
         let config = sled::Config::default().use_compression(true).path(p);
         let db = config.open()?;
 
-        Ok(Self { db })
+        Ok(Self {
+            db,
+            blob_service,
+            directory_service,
+        })
     }
 
-    pub fn new_temporary() -> Result<Self, sled::Error> {
+    pub fn new_temporary(
+        blob_service: Box<dyn BlobService>,
+        directory_service: DS,
+    ) -> Result<Self, sled::Error> {
         let config = sled::Config::default().temporary(true);
         let db = config.open()?;
 
-        Ok(Self { db })
+        Ok(Self {
+            db,
+            blob_service,
+            directory_service,
+        })
     }
 }
 
-impl PathInfoService for SledPathInfoService {
+impl<DS: DirectoryService + Clone> PathInfoService for SledPathInfoService<DS> {
     fn get(&self, digest: [u8; 20]) -> Result<Option<proto::PathInfo>, Error> {
         match self.db.get(digest) {
             Ok(None) => Ok(None),
@@ -73,4 +93,13 @@ impl PathInfoService for SledPathInfoService {
             },
         }
     }
+
+    fn calculate_nar(&self, root_node: &proto::node::Node) -> Result<(u64, [u8; 32]), Error> {
+        calculate_size_and_sha256(
+            root_node,
+            &self.blob_service,
+            self.directory_service.clone(),
+        )
+        .map_err(|e| Error::StorageError(e.to_string()))
+    }
 }