diff options
Diffstat (limited to 'tvix/store/src/nar/renderer.rs')
| -rw-r--r-- | tvix/store/src/nar/renderer.rs | 126 |
1 files changed, 72 insertions, 54 deletions
diff --git a/tvix/store/src/nar/renderer.rs b/tvix/store/src/nar/renderer.rs index 0816b8e973c7..07cdc4b1e31f 100644 --- a/tvix/store/src/nar/renderer.rs +++ b/tvix/store/src/nar/renderer.rs @@ -1,21 +1,53 @@ use crate::utils::AsyncIoBridge; -use super::RenderError; -use async_recursion::async_recursion; +use super::{NarCalculationService, RenderError}; use count_write::CountWrite; use nix_compat::nar::writer::r#async as nar_writer; use sha2::{Digest, Sha256}; use tokio::io::{self, AsyncWrite, BufReader}; -use tvix_castore::{ - blobservice::BlobService, - directoryservice::DirectoryService, - proto::{self as castorepb, NamedNode}, -}; +use tonic::async_trait; +use tracing::instrument; +use tvix_castore::{blobservice::BlobService, directoryservice::DirectoryService, Node}; + +pub struct SimpleRenderer<BS, DS> { + blob_service: BS, + directory_service: DS, +} + +impl<BS, DS> SimpleRenderer<BS, DS> { + pub fn new(blob_service: BS, directory_service: DS) -> Self { + Self { + blob_service, + directory_service, + } + } +} + +#[async_trait] +impl<BS, DS> NarCalculationService for SimpleRenderer<BS, DS> +where + BS: BlobService + Clone, + DS: DirectoryService + Clone, +{ + async fn calculate_nar( + &self, + root_node: &Node, + ) -> Result<(u64, [u8; 32]), tvix_castore::Error> { + calculate_size_and_sha256( + root_node, + self.blob_service.clone(), + self.directory_service.clone(), + ) + .await + .map_err(|e| tvix_castore::Error::StorageError(format!("failed rendering nar: {}", e))) + } +} /// Invoke [write_nar], and return the size and sha256 digest of the produced /// NAR output. +#[instrument(skip_all)] pub async fn calculate_size_and_sha256<BS, DS>( - root_node: &castorepb::node::Node, + root_node: &Node, blob_service: BS, directory_service: DS, ) -> Result<(u64, [u8; 32]), RenderError> @@ -39,13 +71,13 @@ where Ok((cw.count(), h.finalize().into())) } -/// Accepts a [castorepb::node::Node] pointing to the root of a (store) path, +/// Accepts a [Node] pointing to the root of a (store) path, /// and uses the passed blob_service and directory_service to perform the /// necessary lookups as it traverses the structure. /// The contents in NAR serialization are writen to the passed [AsyncWrite]. pub async fn write_nar<W, BS, DS>( mut w: W, - proto_root_node: &castorepb::node::Node, + root_node: &Node, blob_service: BS, directory_service: DS, ) -> Result<(), RenderError> @@ -61,7 +93,8 @@ where walk_node( nar_root_node, - proto_root_node, + root_node, + b"", blob_service, directory_service, ) @@ -72,10 +105,10 @@ where /// Process an intermediate node in the structure. /// This consumes the node. -#[async_recursion] async fn walk_node<BS, DS>( - nar_node: nar_writer::Node<'async_recursion, '_>, - proto_node: &castorepb::node::Node, + nar_node: nar_writer::Node<'_, '_>, + castore_node: &Node, + name: &[u8], blob_service: BS, directory_service: DS, ) -> Result<(BS, DS), RenderError> @@ -83,24 +116,20 @@ where BS: BlobService + Send, DS: DirectoryService + Send, { - match proto_node { - castorepb::node::Node::Symlink(proto_symlink_node) => { + match castore_node { + Node::Symlink { target, .. } => { nar_node - .symlink(&proto_symlink_node.target) + .symlink(target.as_ref()) .await .map_err(RenderError::NARWriterError)?; } - castorepb::node::Node::File(proto_file_node) => { - let digest_len = proto_file_node.digest.len(); - let digest = proto_file_node.digest.clone().try_into().map_err(|_| { - RenderError::StoreError(io::Error::new( - io::ErrorKind::Other, - format!("invalid digest len {} in file node", digest_len), - )) - })?; - + Node::File { + digest, + size, + executable, + } => { let mut blob_reader = match blob_service - .open_read(&digest) + .open_read(digest) .await .map_err(RenderError::StoreError)? { @@ -112,39 +141,23 @@ where }?; nar_node - .file( - proto_file_node.executable, - proto_file_node.size, - &mut blob_reader, - ) + .file(*executable, *size, &mut blob_reader) .await .map_err(RenderError::NARWriterError)?; } - castorepb::node::Node::Directory(proto_directory_node) => { - let digest_len = proto_directory_node.digest.len(); - let digest = proto_directory_node - .digest - .clone() - .try_into() - .map_err(|_| { - RenderError::StoreError(io::Error::new( - io::ErrorKind::InvalidData, - format!("invalid digest len {} in directory node", digest_len), - )) - })?; - + Node::Directory { digest, .. } => { // look it up with the directory service match directory_service - .get(&digest) + .get(digest) .await .map_err(|e| RenderError::StoreError(e.into()))? { // if it's None, that's an error! None => Err(RenderError::DirectoryNotFound( - digest, - proto_directory_node.name.clone(), + digest.clone(), + bytes::Bytes::copy_from_slice(name), ))?, - Some(proto_directory) => { + Some(directory) => { // start a directory node let mut nar_node_directory = nar_node .directory() @@ -158,15 +171,20 @@ where // for each node in the directory, create a new entry with its name, // and then recurse on that entry. - for proto_node in proto_directory.nodes() { + for (name, node) in directory.nodes() { let child_node = nar_node_directory - .entry(proto_node.get_name()) + .entry(name.as_ref()) .await .map_err(RenderError::NARWriterError)?; - (blob_service, directory_service) = - walk_node(child_node, &proto_node, blob_service, directory_service) - .await?; + (blob_service, directory_service) = Box::pin(walk_node( + child_node, + node, + name.as_ref(), + blob_service, + directory_service, + )) + .await?; } // close the directory |