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 <tvl@lahfa.xyz>
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
This commit is contained in:
Florian Klink 2024-01-05 17:03:13 +02:00 committed by clbot
parent 4284cd82ef
commit 5f0360c566
2 changed files with 44 additions and 46 deletions

View file

@ -94,7 +94,7 @@ fn interpret(code: &str, path: Option<PathBuf>, args: &Args, explain: bool) -> b
eval.io_handle = Box::new(tvix_glue::tvix_io::TvixIO::new(TvixStoreIO::new(
blob_service,
directory_service,
path_info_service.into(), // we need an Arc<_> here.
path_info_service,
tokio_runtime.handle().clone(),
)));

View file

@ -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(