From 419f7d0f0c3dadeb0bbecda763aaec73ed8da8dd Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Mon, 13 Feb 2023 16:47:22 +0100 Subject: feat(tvix/store): add NARCalculationService trait and impl This adds a NARCalculationService trait, which will take a root node, and provide a proto::CalculateNarResponse in return. It also adds a NonCachingNARCalculationService implementation, that will simply always render the NAR in memory to calculate the size and sha256, without any caching. Change-Id: Id1ffb18559212fa6001f70f2634bbc3dfd0aa343 Reviewed-on: https://cl.tvl.fyi/c/depot/+/8096 Reviewed-by: raitobezarius Tested-by: BuildkiteCI --- tvix/Cargo.lock | 8 ++++ tvix/Cargo.nix | 23 ++++++++++ tvix/store/Cargo.toml | 8 ++-- tvix/store/src/nar/mod.rs | 11 +++++ .../src/nar/non_caching_calculation_service.rs | 50 ++++++++++++++++++++++ 5 files changed, 97 insertions(+), 3 deletions(-) create mode 100644 tvix/store/src/nar/non_caching_calculation_service.rs (limited to 'tvix') diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock index 913d88b9fc12..106a1e38d0b4 100644 --- a/tvix/Cargo.lock +++ b/tvix/Cargo.lock @@ -422,6 +422,12 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f3ad85c1f65dc7b37604eb0e89748faf0b9653065f2a8ef69f96a687ec1e9279" +[[package]] +name = "count-write" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ced507ab50aa0123e2c54db8b5f44fdfee04b1c93744d69e924307945fe57a85" + [[package]] name = "countme" version = "3.0.1" @@ -2770,12 +2776,14 @@ dependencies = [ "anyhow", "blake3", "clap 4.1.4", + "count-write", "data-encoding", "fastcdc", "lazy_static", "nix-compat", "prost", "prost-build", + "sha2 0.10.6", "sled", "tempfile", "test-case", diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix index 126edecdff17..fc3f945df342 100644 --- a/tvix/Cargo.nix +++ b/tvix/Cargo.nix @@ -1246,6 +1246,21 @@ rec { ]; }; + "count-write" = rec { + crateName = "count-write"; + version = "0.1.0"; + edition = "2018"; + sha256 = "11bswmgr81s3jagdci1pr6qh9vnz9zsbbf2dqpi260daa2mhgmff"; + authors = [ + "SOFe " + ]; + features = { + "futures" = [ "futures-io-preview" ]; + "futures-io-preview" = [ "dep:futures-io-preview" ]; + "tokio" = [ "tokio-io" ]; + "tokio-io" = [ "dep:tokio-io" ]; + }; + }; "countme" = rec { crateName = "countme"; version = "3.0.1"; @@ -8252,6 +8267,10 @@ rec { packageId = "clap 4.1.4"; features = [ "derive" "env" ]; } + { + name = "count-write"; + packageId = "count-write"; + } { name = "data-encoding"; packageId = "data-encoding"; @@ -8272,6 +8291,10 @@ rec { name = "prost"; packageId = "prost"; } + { + name = "sha2"; + packageId = "sha2 0.10.6"; + } { name = "sled"; packageId = "sled"; diff --git a/tvix/store/Cargo.toml b/tvix/store/Cargo.toml index 769b496235d6..196e986b89ab 100644 --- a/tvix/store/Cargo.toml +++ b/tvix/store/Cargo.toml @@ -9,19 +9,21 @@ name = "tvix_store" [dependencies] anyhow = "1.0.68" blake3 = { version = "1.3.1", features = ["rayon", "std"] } +clap = { version = "4.0", features = ["derive", "env"] } +count-write = "0.1.0" data-encoding = "2.3.3" +fastcdc = "2.0.0" lazy_static = "1.4.0" -clap = { version = "4.0", features = ["derive", "env"] } prost = "0.11.2" +sha2 = "0.10.6" sled = { version = "0.34.7", features = ["compression"] } nix-compat = { path = "../nix-compat" } thiserror = "1.0.38" tokio-stream = "0.1.11" tokio = { version = "1.23.0", features = ["rt-multi-thread"] } tonic = "0.8.2" -tracing-subscriber = "0.3.16" tracing = "0.1.37" -fastcdc = "2.0.0" +tracing-subscriber = "0.3.16" [dependencies.tonic-reflection] optional = true diff --git a/tvix/store/src/nar/mod.rs b/tvix/store/src/nar/mod.rs index a3a8677d92a6..d9e0c8143b5a 100644 --- a/tvix/store/src/nar/mod.rs +++ b/tvix/store/src/nar/mod.rs @@ -1,8 +1,11 @@ +use crate::proto; use data_encoding::BASE64; use thiserror::Error; +mod non_caching_calculation_service; mod renderer; +pub use non_caching_calculation_service::NonCachingNARCalculationService; pub use renderer::NARRenderer; /// Errors that can encounter while rendering NARs. @@ -23,3 +26,11 @@ pub enum RenderError { #[error("failure using the NAR writer: {0}")] NARWriterError(std::io::Error), } + +/// The base trait for something calculating NARs, and returning their size and sha256. +pub trait NARCalculationService { + fn calculate_nar( + &self, + root_node: proto::node::Node, + ) -> Result; +} diff --git a/tvix/store/src/nar/non_caching_calculation_service.rs b/tvix/store/src/nar/non_caching_calculation_service.rs new file mode 100644 index 000000000000..880d12636979 --- /dev/null +++ b/tvix/store/src/nar/non_caching_calculation_service.rs @@ -0,0 +1,50 @@ +use count_write::CountWrite; +use sha2::{Digest, Sha256}; + +use crate::blobservice::BlobService; +use crate::chunkservice::ChunkService; +use crate::directoryservice::DirectoryService; +use crate::proto; + +use super::renderer::NARRenderer; +use super::{NARCalculationService, RenderError}; + +/// A NAR calculation service which simply renders the whole NAR whenever +/// we ask for the calculation. +#[derive(Clone)] +pub struct NonCachingNARCalculationService< + BS: BlobService, + CS: ChunkService + Clone, + DS: DirectoryService, +> { + nar_renderer: NARRenderer, +} + +impl + NonCachingNARCalculationService +{ + pub fn new(blob_service: BS, chunk_service: CS, directory_service: DS) -> Self { + Self { + nar_renderer: NARRenderer::new(blob_service, chunk_service, directory_service), + } + } +} + +impl NARCalculationService + for NonCachingNARCalculationService +{ + fn calculate_nar( + &self, + root_node: proto::node::Node, + ) -> Result { + let h = Sha256::new(); + let mut cw = CountWrite::from(h); + + self.nar_renderer.write_nar(&mut cw, root_node)?; + + Ok(proto::CalculateNarResponse { + nar_size: cw.count() as u32, + nar_sha256: cw.into_inner().finalize().to_vec(), + }) + } +} -- cgit 1.4.1