From ee23220564987771c8e7909ded6fb9853f1d1b0d Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Thu, 16 Mar 2023 00:01:30 +0100 Subject: refactor(tvix/store/directorysvc): use [u8; 32] instead of Vec Also, simplify the trait interface, only allowing lookups of Directory objects by their digest. Change-Id: I6eec28a8cb0557bed9b69df8b8ff99a5e0f8fe35 Reviewed-on: https://cl.tvl.fyi/c/depot/+/8313 Tested-by: BuildkiteCI Autosubmit: flokli Reviewed-by: tazjin --- tvix/store/src/directoryservice/memory.rs | 56 ++++++++++++--------------- tvix/store/src/directoryservice/mod.rs | 7 +--- tvix/store/src/directoryservice/sled.rs | 64 ++++++++++++++----------------- 3 files changed, 54 insertions(+), 73 deletions(-) (limited to 'tvix/store/src/directoryservice') diff --git a/tvix/store/src/directoryservice/memory.rs b/tvix/store/src/directoryservice/memory.rs index 87cd038888e6..7440be112cfd 100644 --- a/tvix/store/src/directoryservice/memory.rs +++ b/tvix/store/src/directoryservice/memory.rs @@ -8,46 +8,38 @@ use super::DirectoryService; #[derive(Clone, Default)] pub struct MemoryDirectoryService { - db: Arc, proto::Directory>>>, + db: Arc>>, } impl DirectoryService for MemoryDirectoryService { - // TODO: change api to only be by digest - #[instrument(skip(self, by_what))] - fn get( - &self, - by_what: &proto::get_directory_request::ByWhat, - ) -> Result, Error> { - match by_what { - proto::get_directory_request::ByWhat::Digest(digest) => { - let db = self.db.read()?; - - match db.get(digest) { - // The directory was not found, return - None => Ok(None), - - // The directory was found, try to parse the data as Directory message - Some(directory) => { - // Validate the retrieved Directory indeed has the - // digest we expect it to have, to detect corruptions. - let actual_digest = directory.digest(); - if actual_digest.as_slice() != digest { - return Err(Error::StorageError(format!( - "requested directory with digest {}, but got {}", - BASE64.encode(digest), - BASE64.encode(&actual_digest) - ))); - } - - Ok(Some(directory.clone())) - } + #[instrument(skip(self, digest), fields(directory.digest = BASE64.encode(digest)))] + fn get(&self, digest: &[u8; 32]) -> Result, Error> { + let db = self.db.read()?; + + match db.get(digest) { + // The directory was not found, return + None => Ok(None), + + // The directory was found, try to parse the data as Directory message + Some(directory) => { + // Validate the retrieved Directory indeed has the + // digest we expect it to have, to detect corruptions. + let actual_digest = directory.digest(); + if actual_digest.as_slice() != digest { + return Err(Error::StorageError(format!( + "requested directory with digest {}, but got {}", + BASE64.encode(digest), + BASE64.encode(&actual_digest) + ))); } + + Ok(Some(directory.clone())) } } } #[instrument(skip(self, directory), fields(directory.digest = BASE64.encode(&directory.digest())))] - fn put(&self, directory: proto::Directory) -> Result, Error> { + fn put(&self, directory: proto::Directory) -> Result<[u8; 32], Error> { let digest = directory.digest(); // validate the directory itself. @@ -61,7 +53,7 @@ impl DirectoryService for MemoryDirectoryService { // store it let mut db = self.db.write()?; - db.insert(digest.clone(), directory); + db.insert(digest, directory); Ok(digest) } diff --git a/tvix/store/src/directoryservice/mod.rs b/tvix/store/src/directoryservice/mod.rs index 4abf823b23b4..a27f6745332a 100644 --- a/tvix/store/src/directoryservice/mod.rs +++ b/tvix/store/src/directoryservice/mod.rs @@ -11,11 +11,8 @@ pub use self::sled::SledDirectoryService; pub trait DirectoryService { /// Get looks up a single Directory message by its digest. /// In case the directory is not found, Ok(None) is returned. - fn get( - &self, - by_what: &proto::get_directory_request::ByWhat, - ) -> Result, Error>; + fn get(&self, digest: &[u8; 32]) -> Result, Error>; /// Get uploads a single Directory message, and returns the calculated /// digest, or an error. - fn put(&self, directory: proto::Directory) -> Result, Error>; + fn put(&self, directory: proto::Directory) -> Result<[u8; 32], Error>; } diff --git a/tvix/store/src/directoryservice/sled.rs b/tvix/store/src/directoryservice/sled.rs index 064f86ceaa6e..1f729a594c19 100644 --- a/tvix/store/src/directoryservice/sled.rs +++ b/tvix/store/src/directoryservice/sled.rs @@ -29,48 +29,40 @@ impl SledDirectoryService { } impl DirectoryService for SledDirectoryService { - // TODO: change api to only be by digest - #[instrument(name = "SledDirectoryService::get", skip(self, by_what))] - fn get( - &self, - by_what: &proto::get_directory_request::ByWhat, - ) -> Result, Error> { - match by_what { - proto::get_directory_request::ByWhat::Digest(digest) => { - match self.db.get(digest) { - // The directory was not found, return - Ok(None) => Ok(None), + #[instrument(name = "SledDirectoryService::get", skip(self, digest), fields(directory.digest = BASE64.encode(digest)))] + fn get(&self, digest: &[u8; 32]) -> Result, Error> { + match self.db.get(digest) { + // The directory was not found, return + Ok(None) => Ok(None), - // The directory was found, try to parse the data as Directory message - Ok(Some(data)) => match Directory::decode(&*data) { - Ok(directory) => { - // Validate the retrieved Directory indeed has the - // digest we expect it to have, to detect corruptions. - let actual_digest = directory.digest(); - if actual_digest.as_slice() != digest { - return Err(Error::StorageError(format!( - "requested directory with digest {}, but got {}", - BASE64.encode(digest), - BASE64.encode(&actual_digest) - ))); - } + // The directory was found, try to parse the data as Directory message + Ok(Some(data)) => match Directory::decode(&*data) { + Ok(directory) => { + // Validate the retrieved Directory indeed has the + // digest we expect it to have, to detect corruptions. + let actual_digest = directory.digest(); + if actual_digest.as_slice() != digest { + return Err(Error::StorageError(format!( + "requested directory with digest {}, but got {}", + BASE64.encode(digest), + BASE64.encode(&actual_digest) + ))); + } - Ok(Some(directory)) - } - Err(e) => { - warn!("unable to parse directory {}: {}", BASE64.encode(digest), e); - Err(Error::StorageError(e.to_string())) - } - }, - // some storage error? - Err(e) => Err(Error::StorageError(e.to_string())), + Ok(Some(directory)) } - } + Err(e) => { + warn!("unable to parse directory {}: {}", BASE64.encode(digest), e); + Err(Error::StorageError(e.to_string())) + } + }, + // some storage error? + Err(e) => Err(Error::StorageError(e.to_string())), } } #[instrument(name = "SledDirectoryService::put", skip(self, directory), fields(directory.digest = BASE64.encode(&directory.digest())))] - fn put(&self, directory: proto::Directory) -> Result, Error> { + fn put(&self, directory: proto::Directory) -> Result<[u8; 32], Error> { let digest = directory.digest(); // validate the directory itself. @@ -82,7 +74,7 @@ impl DirectoryService for SledDirectoryService { ))); } // store it - let result = self.db.insert(&digest, directory.encode_to_vec()); + let result = self.db.insert(digest, directory.encode_to_vec()); if let Err(e) = result { return Err(Error::StorageError(e.to_string())); } -- cgit 1.4.1