From 5f0360c566a8c3a6ebfea1c3b2ddb068ebba4859 Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Fri, 5 Jan 2024 17:03:13 +0200 Subject: refactor(tvix/glue): simplify TvixStoreIO further We don't need to spawn in all these places, we can just block_on directly, this is all IO bound. This also means, we don't need to clone any of the service handles (except preserving clone-ability of the BlobService). Change-Id: I7d90f4d6a263a98491caa071ada538a5197a5472 Reviewed-on: https://cl.tvl.fyi/c/depot/+/10540 Reviewed-by: raitobezarius Autosubmit: flokli Tested-by: BuildkiteCI --- tvix/glue/src/tvix_store_io.rs | 88 +++++++++++++++++++++--------------------- 1 file changed, 43 insertions(+), 45 deletions(-) (limited to 'tvix/glue/src') diff --git a/tvix/glue/src/tvix_store_io.rs b/tvix/glue/src/tvix_store_io.rs index f6d3d0792360..eba8cee7336f 100644 --- a/tvix/glue/src/tvix_store_io.rs +++ b/tvix/glue/src/tvix_store_io.rs @@ -5,7 +5,6 @@ use std::{ io, ops::Deref, path::{Path, PathBuf}, - sync::Arc, }; use tokio::io::AsyncReadExt; use tracing::{error, instrument, warn}; @@ -26,19 +25,23 @@ use tvix_store::pathinfoservice::PathInfoService; /// This is to both cover cases of syntactically valid store paths, that exist /// on the filesystem (still managed by Nix), as well as being able to read /// files outside store paths. -pub struct TvixStoreIO { - blob_service: Arc, - directory_service: Arc, - path_info_service: Arc, +pub struct TvixStoreIO { + blob_service: BS, + directory_service: DS, + path_info_service: PS, std_io: StdIO, tokio_handle: tokio::runtime::Handle, } -impl TvixStoreIO { +impl TvixStoreIO +where + DS: Deref, + PS: Deref, +{ pub fn new( - blob_service: Arc, - directory_service: Arc, - path_info_service: Arc, + blob_service: BS, + directory_service: DS, + path_info_service: PS, tokio_handle: tokio::runtime::Handle, ) -> Self { Self { @@ -97,7 +100,12 @@ impl TvixStoreIO { } } -impl EvalIO for TvixStoreIO { +impl EvalIO for TvixStoreIO +where + BS: Deref + Clone, + DS: Deref, + PS: Deref, +{ #[instrument(skip(self), ret, err)] fn path_exists(&self, path: &Path) -> io::Result { if let Ok((store_path, sub_path)) = @@ -144,11 +152,9 @@ impl EvalIO for TvixStoreIO { ) })?; - let blob_service = self.blob_service.clone(); - - let task = self.tokio_handle.spawn(async move { + self.tokio_handle.block_on(async { let mut reader = { - let resp = blob_service.open_read(&digest).await?; + let resp = self.blob_service.deref().open_read(&digest).await?; match resp { Some(blob_reader) => blob_reader, None => { @@ -168,9 +174,7 @@ impl EvalIO for TvixStoreIO { reader.read_to_string(&mut buf).await?; Ok(buf) - }); - - self.tokio_handle.block_on(task).unwrap() + }) } Node::Symlink(_symlink_node) => Err(io::Error::new( io::ErrorKind::Unsupported, @@ -208,12 +212,10 @@ impl EvalIO for TvixStoreIO { ) })?; - let directory_service = self.directory_service.clone(); - let digest_clone = digest.clone(); - let task = self + if let Some(directory) = self .tokio_handle - .spawn(async move { directory_service.get(&digest_clone).await }); - if let Some(directory) = self.tokio_handle.block_on(task).unwrap()? { + .block_on(async { self.directory_service.deref().get(&digest).await })? + { let mut children: Vec<(bytes::Bytes, FileType)> = Vec::new(); for node in directory.nodes() { children.push(match node { @@ -258,23 +260,15 @@ impl EvalIO for TvixStoreIO { #[instrument(skip(self), ret, err)] fn import_path(&self, path: &Path) -> io::Result { - let task = self.tokio_handle.spawn({ - let blob_service = self.blob_service.clone(); - let directory_service = self.directory_service.clone(); - let path_info_service = self.path_info_service.clone(); - let path = path.to_owned(); - async move { - tvix_store::utils::import_path( - path, - blob_service, - directory_service, - path_info_service, - ) - .await - } - }); - - let output_path = self.tokio_handle.block_on(task)??; + let output_path = self.tokio_handle.block_on(async { + tvix_store::utils::import_path( + path, + self.blob_service.deref(), + self.directory_service.deref(), + self.path_info_service.deref(), + ) + .await + })?; Ok(output_path.to_absolute_path().into()) } @@ -290,9 +284,12 @@ mod tests { use std::{cell::RefCell, path::Path, rc::Rc, sync::Arc}; use tempfile::TempDir; - use tvix_castore::{blobservice::MemoryBlobService, directoryservice::MemoryDirectoryService}; + use tvix_castore::{ + blobservice::{BlobService, MemoryBlobService}, + directoryservice::{DirectoryService, MemoryDirectoryService}, + }; use tvix_eval::EvaluationResult; - use tvix_store::pathinfoservice::MemoryPathInfoService; + use tvix_store::pathinfoservice::{MemoryPathInfoService, PathInfoService}; use crate::{builtins::add_derivation_builtins, known_paths::KnownPaths}; @@ -304,12 +301,13 @@ mod tests { fn eval(str: &str) -> EvaluationResult { let mut eval = tvix_eval::Evaluation::new_impure(); - let blob_service = Arc::new(MemoryBlobService::default()); - let directory_service = Arc::new(MemoryDirectoryService::default()); - let path_info_service = Arc::new(MemoryPathInfoService::new( + let blob_service = Arc::new(MemoryBlobService::default()) as Arc; + let directory_service = + Arc::new(MemoryDirectoryService::default()) as Arc; + let path_info_service = Box::new(MemoryPathInfoService::new( blob_service.clone(), directory_service.clone(), - )); + )) as Box; let runtime = tokio::runtime::Runtime::new().unwrap(); eval.io_handle = Box::new(TvixStoreIO::new( -- cgit 1.4.1