From 9b3228959a99e81ab6f1ce9c3a21136dfee82c62 Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Mon, 13 Feb 2023 18:55:40 +0100 Subject: feat(tvix/store): add pathinfoservice This adds a PathInfoService trait, and an implementation for it using sled, and one using a HashMap. Change-Id: I85fe4c9b6105808d7b0c095441326424ffc2d2bb Reviewed-on: https://cl.tvl.fyi/c/depot/+/8099 Tested-by: BuildkiteCI Reviewed-by: raitobezarius --- tvix/store/src/lib.rs | 1 + tvix/store/src/pathinfoservice/memory.rs | 62 ++++++++++++++++++++++++ tvix/store/src/pathinfoservice/mod.rs | 21 +++++++++ tvix/store/src/pathinfoservice/sled.rs | 81 ++++++++++++++++++++++++++++++++ 4 files changed, 165 insertions(+) create mode 100644 tvix/store/src/pathinfoservice/memory.rs create mode 100644 tvix/store/src/pathinfoservice/mod.rs create mode 100644 tvix/store/src/pathinfoservice/sled.rs (limited to 'tvix/store') diff --git a/tvix/store/src/lib.rs b/tvix/store/src/lib.rs index 0279d13f9b..c275be5ff0 100644 --- a/tvix/store/src/lib.rs +++ b/tvix/store/src/lib.rs @@ -6,6 +6,7 @@ mod errors; pub mod blobservice; pub mod chunkservice; pub mod directoryservice; +pub mod pathinfoservice; pub mod proto; pub use blobreader::BlobReader; diff --git a/tvix/store/src/pathinfoservice/memory.rs b/tvix/store/src/pathinfoservice/memory.rs new file mode 100644 index 0000000000..db31a2d771 --- /dev/null +++ b/tvix/store/src/pathinfoservice/memory.rs @@ -0,0 +1,62 @@ +use std::{ + collections::HashMap, + sync::{Arc, RwLock}, +}; + +use crate::{proto, Error}; +use nix_compat::store_path::DIGEST_SIZE; + +use super::PathInfoService; + +pub struct MemoryPathInfoService { + db: Arc, proto::PathInfo>>>, +} + +impl MemoryPathInfoService { + pub fn new() -> Self { + let db = Arc::new(RwLock::new(HashMap::default())); + + Self { db } + } +} + +impl PathInfoService for MemoryPathInfoService { + fn get( + &self, + by_what: proto::get_path_info_request::ByWhat, + ) -> Result, Error> { + match by_what { + proto::get_path_info_request::ByWhat::ByOutputHash(digest) => { + if digest.len() != DIGEST_SIZE { + return Err(Error::InvalidRequest("invalid digest length".to_string())); + } + + let db = self.db.read().unwrap(); + + match db.get(&digest) { + None => Ok(None), + Some(path_info) => Ok(Some(path_info.clone())), + } + } + } + } + + fn put(&self, path_info: proto::PathInfo) -> Result { + // Call validate on the received PathInfo message. + match path_info.validate() { + Err(e) => Err(Error::InvalidRequest(format!( + "failed to validate PathInfo: {}", + e + ))), + + // In case the PathInfo is valid, and we were able to extract a NixPath, store it in the database. + // This overwrites existing PathInfo objects. + Ok(nix_path) => { + let mut db = self.db.write().unwrap(); + db.insert(nix_path.digest.to_vec(), path_info.clone()); + + Ok(path_info) + } + } + } +} diff --git a/tvix/store/src/pathinfoservice/mod.rs b/tvix/store/src/pathinfoservice/mod.rs new file mode 100644 index 0000000000..b20985dd80 --- /dev/null +++ b/tvix/store/src/pathinfoservice/mod.rs @@ -0,0 +1,21 @@ +mod memory; +mod sled; + +use crate::{proto, Error}; + +pub use self::memory::MemoryPathInfoService; +pub use self::sled::SledPathInfoService; + +/// The base trait all PathInfo services need to implement. +/// This is a simple get and put of [proto::Directory], returning their digest. +pub trait PathInfoService { + /// Retrieve a PathInfo message. + fn get( + &self, + by_what: proto::get_path_info_request::ByWhat, + ) -> Result, Error>; + + /// Store a PathInfo message. Implementations MUST call validate and reject + /// invalid messages. + fn put(&self, path_info: proto::PathInfo) -> Result; +} diff --git a/tvix/store/src/pathinfoservice/sled.rs b/tvix/store/src/pathinfoservice/sled.rs new file mode 100644 index 0000000000..dc3c95b510 --- /dev/null +++ b/tvix/store/src/pathinfoservice/sled.rs @@ -0,0 +1,81 @@ +use crate::{proto, Error}; +use nix_compat::store_path::DIGEST_SIZE; +use prost::Message; +use std::path::PathBuf; +use tracing::warn; + +use super::PathInfoService; + +/// SledPathInfoService stores PathInfo in a [sled](https://github.com/spacejam/sled). +/// +/// The PathInfo messages are stored as encoded protos, and keyed by their output hash, +/// as that's currently the only request type available. +pub struct SledPathInfoService { + db: sled::Db, +} + +impl SledPathInfoService { + pub fn new(p: PathBuf) -> Result { + let config = sled::Config::default().use_compression(true).path(p); + let db = config.open()?; + + Ok(Self { db }) + } +} + +impl PathInfoService for SledPathInfoService { + fn get( + &self, + by_what: proto::get_path_info_request::ByWhat, + ) -> Result, Error> { + match by_what { + proto::get_path_info_request::ByWhat::ByOutputHash(digest) => { + if digest.len() != DIGEST_SIZE { + return Err(Error::InvalidRequest("invalid digest length".to_string())); + } + + match self.db.get(digest) { + Ok(None) => Ok(None), + Ok(Some(data)) => match proto::PathInfo::decode(&*data) { + Ok(path_info) => Ok(Some(path_info)), + Err(e) => { + warn!("failed to decode stored PathInfo: {}", e); + Err(Error::StorageError(format!( + "failed to decode stored PathInfo: {}", + e + ))) + } + }, + Err(e) => { + warn!("failed to retrieve PathInfo: {}", e); + Err(Error::StorageError(format!( + "failed to retrieve PathInfo: {}", + e + ))) + } + } + } + } + } + + fn put(&self, path_info: proto::PathInfo) -> Result { + // Call validate on the received PathInfo message. + match path_info.validate() { + Err(e) => Err(Error::InvalidRequest(format!( + "failed to validate PathInfo: {}", + e + ))), + // In case the PathInfo is valid, and we were able to extract a NixPath, store it in the database. + // This overwrites existing PathInfo objects. + Ok(nix_path) => match self.db.insert(nix_path.digest, path_info.encode_to_vec()) { + Ok(_) => Ok(path_info), + Err(e) => { + warn!("failed to insert PathInfo: {}", e); + Err(Error::StorageError(format! { + "failed to insert PathInfo: {}", e + })) + } + }, + } + } +} -- cgit 1.4.1