From 6f85dbfc06c4fa96deb968cfeb7e98ba36e95043 Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Fri, 9 Jun 2023 10:28:02 +0300 Subject: feat(tvix/store/pathinfosvc): add calculate_nar method 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 Reviewed-by: tazjin --- tvix/store/src/pathinfoservice/grpc.rs | 32 ++++++++++++++++++++++ tvix/store/src/pathinfoservice/memory.rs | 32 +++++++++++++++++++--- tvix/store/src/pathinfoservice/mod.rs | 5 +++- tvix/store/src/pathinfoservice/sled.rs | 47 ++++++++++++++++++++++++++------ 4 files changed, 102 insertions(+), 14 deletions(-) (limited to 'tvix/store/src/pathinfoservice') 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> = + 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 { db: Arc>>, + + blob_service: Box, + directory_service: DS, +} + +impl MemoryPathInfoService { + pub fn new(blob_service: Box, directory_service: DS) -> Self { + Self { + db: Default::default(), + blob_service, + directory_service, + } + } } -impl PathInfoService for MemoryPathInfoService { +impl PathInfoService for MemoryPathInfoService { fn get(&self, digest: [u8; 20]) -> Result, 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; - // 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 { db: sled::Db, + + blob_service: Box, + directory_service: DS, } -impl SledPathInfoService { - pub fn new(p: PathBuf) -> Result { +impl SledPathInfoService { + pub fn new( + p: PathBuf, + blob_service: Box, + directory_service: DS, + ) -> Result { 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 { + pub fn new_temporary( + blob_service: Box, + directory_service: DS, + ) -> Result { 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 PathInfoService for SledPathInfoService { fn get(&self, digest: [u8; 20]) -> Result, 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())) + } } -- cgit 1.4.1