refactor(tvix/glue): simplify TvixStoreIO further r/7352

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 <tvl@lahfa.xyz>
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI

1 files changed, 43 insertions, 45 deletions

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<dyn BlobService>,
-    directory_service: Arc<dyn DirectoryService>,
-    path_info_service: Arc<dyn PathInfoService>,
+pub struct TvixStoreIO<BS, DS, PS> {
+    blob_service: BS,
+    directory_service: DS,
+    path_info_service: PS,
     std_io: StdIO,
     tokio_handle: tokio::runtime::Handle,
 }
 
-impl TvixStoreIO {
+impl<BS, DS, PS> TvixStoreIO<BS, DS, PS>
+where
+    DS: Deref<Target = dyn DirectoryService>,
+    PS: Deref<Target = dyn PathInfoService>,
+{
     pub fn new(
-        blob_service: Arc<dyn BlobService>,
-        directory_service: Arc<dyn DirectoryService>,
-        path_info_service: Arc<dyn PathInfoService>,
+        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<BS, DS, PS> EvalIO for TvixStoreIO<BS, DS, PS>
+where
+    BS: Deref<Target = dyn BlobService> + Clone,
+    DS: Deref<Target = dyn DirectoryService>,
+    PS: Deref<Target = dyn PathInfoService>,
+{
     #[instrument(skip(self), ret, err)]
     fn path_exists(&self, path: &Path) -> io::Result<bool> {
         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<PathBuf> {
-        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<dyn BlobService>;
+        let directory_service =
+            Arc::new(MemoryDirectoryService::default()) as Arc<dyn DirectoryService>;
+        let path_info_service = Box::new(MemoryPathInfoService::new(
             blob_service.clone(),
             directory_service.clone(),
-        ));
+        )) as Box<dyn PathInfoService>;
         let runtime = tokio::runtime::Runtime::new().unwrap();
 
         eval.io_handle = Box::new(TvixStoreIO::new(