diff options
Diffstat (limited to '')
-rw-r--r-- | tvix/store/src/nar/renderer.rs | 61 |
1 files changed, 38 insertions, 23 deletions
diff --git a/tvix/store/src/nar/renderer.rs b/tvix/store/src/nar/renderer.rs index 4255148fc5e8..f1392472a50e 100644 --- a/tvix/store/src/nar/renderer.rs +++ b/tvix/store/src/nar/renderer.rs @@ -8,20 +8,20 @@ use count_write::CountWrite; use nix_compat::nar; use sha2::{Digest, Sha256}; use std::{io, sync::Arc}; -use tokio::io::BufReader; +use tokio::{io::BufReader, task::spawn_blocking}; use tracing::warn; /// Invoke [write_nar], and return the size and sha256 digest of the produced /// NAR output. -pub fn calculate_size_and_sha256( +pub async fn calculate_size_and_sha256( root_node: &proto::node::Node, blob_service: Arc<dyn BlobService>, directory_service: Arc<dyn DirectoryService>, ) -> Result<(u64, [u8; 32]), RenderError> { let h = Sha256::new(); - let mut cw = CountWrite::from(h); + let cw = CountWrite::from(h); - write_nar(&mut cw, root_node, blob_service, directory_service)?; + let cw = write_nar(cw, root_node, blob_service, directory_service).await?; Ok((cw.count(), cw.into_inner().finalize().into())) } @@ -30,26 +30,44 @@ pub fn calculate_size_and_sha256( /// 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 [std::io::Write]. -pub fn write_nar<W: std::io::Write>( - w: &mut W, +/// +/// The writer is passed back in the return value. This is done because async Rust +/// lacks scoped blocking tasks, so we need to transfer ownership of the writer +/// internally. +/// +/// # Panics +/// This will panic if called outside the context of a Tokio runtime. +pub async fn write_nar<W: std::io::Write + Send + 'static>( + mut w: W, proto_root_node: &proto::node::Node, blob_service: Arc<dyn BlobService>, directory_service: Arc<dyn DirectoryService>, -) -> Result<(), RenderError> { - // Initialize NAR writer - let nar_root_node = nar::writer::open(w).map_err(RenderError::NARWriterError)?; +) -> Result<W, RenderError> { + let tokio_handle = tokio::runtime::Handle::current(); + let proto_root_node = proto_root_node.clone(); + + spawn_blocking(move || { + // Initialize NAR writer + let nar_root_node = nar::writer::open(&mut w).map_err(RenderError::NARWriterError)?; - walk_node( - nar_root_node, - proto_root_node, - blob_service, - directory_service, - ) + walk_node( + tokio_handle, + nar_root_node, + &proto_root_node, + blob_service, + directory_service, + )?; + + Ok(w) + }) + .await + .unwrap() } /// Process an intermediate node in the structure. /// This consumes the node. fn walk_node( + tokio_handle: tokio::runtime::Handle, nar_node: nar::writer::Node, proto_node: &proto::node::Node, blob_service: Arc<dyn BlobService>, @@ -73,9 +91,6 @@ fn walk_node( )) })?; - // HACK: blob_service is async, but this function isn't async yet.. - let tokio_handle = tokio::runtime::Handle::current(); - let blob_reader = match tokio_handle .block_on(async { blob_service.open_read(&digest).await }) .map_err(RenderError::StoreError)? @@ -107,11 +122,10 @@ fn walk_node( })?; // look it up with the directory service - let resp = directory_service - .get(&digest) - .map_err(RenderError::StoreError)?; - - match resp { + match tokio_handle + .block_on(async { directory_service.get(&digest).await }) + .map_err(RenderError::StoreError)? + { // if it's None, that's an error! None => { return Err(RenderError::DirectoryNotFound( @@ -131,6 +145,7 @@ fn walk_node( .entry(proto_node.get_name()) .map_err(RenderError::NARWriterError)?; walk_node( + tokio_handle.clone(), child_node, &proto_node, blob_service.clone(), |