From b59df53774acc654ea4b23f02ccf5529587bceff Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Fri, 12 Jan 2024 12:18:43 +0200 Subject: refactor(tvix/store/pathinfoservice): make more generic We don't need Arcs in most of the cases, we're fine with some container. Change-Id: Ic4f8acb5b9d93e2b0923bb607463fb91e9d0e4fe Reviewed-on: https://cl.tvl.fyi/c/depot/+/10606 Autosubmit: flokli Tested-by: BuildkiteCI Reviewed-by: raitobezarius --- tvix/store/src/pathinfoservice/grpc.rs | 4 +-- tvix/store/src/pathinfoservice/memory.rs | 29 ++++++++---------- tvix/store/src/pathinfoservice/nix_http.rs | 21 +++++++------ tvix/store/src/pathinfoservice/sled.rs | 35 ++++++++++------------ tvix/store/src/proto/tests/grpc_pathinfoservice.rs | 4 +-- tvix/store/src/tests/utils.rs | 12 +++++--- 6 files changed, 51 insertions(+), 54 deletions(-) (limited to 'tvix') diff --git a/tvix/store/src/pathinfoservice/grpc.rs b/tvix/store/src/pathinfoservice/grpc.rs index 4ec1467525f9..575491333142 100644 --- a/tvix/store/src/pathinfoservice/grpc.rs +++ b/tvix/store/src/pathinfoservice/grpc.rs @@ -159,8 +159,8 @@ mod tests { let router = server.add_service( crate::proto::path_info_service_server::PathInfoServiceServer::new( GRPCPathInfoServiceWrapper::new(Box::new(MemoryPathInfoService::new( - gen_blob_service().into(), - gen_directory_service().into(), + gen_blob_service(), + gen_directory_service(), )) as Box), ), diff --git a/tvix/store/src/pathinfoservice/memory.rs b/tvix/store/src/pathinfoservice/memory.rs index 6f96c083476c..9f657a9c625b 100644 --- a/tvix/store/src/pathinfoservice/memory.rs +++ b/tvix/store/src/pathinfoservice/memory.rs @@ -11,18 +11,15 @@ use tvix_castore::proto as castorepb; use tvix_castore::Error; use tvix_castore::{blobservice::BlobService, directoryservice::DirectoryService}; -pub struct MemoryPathInfoService { +pub struct MemoryPathInfoService { db: Arc>>, - blob_service: Arc, - directory_service: Arc, + blob_service: BS, + directory_service: DS, } -impl MemoryPathInfoService { - pub fn new( - blob_service: Arc, - directory_service: Arc, - ) -> Self { +impl MemoryPathInfoService { + pub fn new(blob_service: BS, directory_service: DS) -> Self { Self { db: Default::default(), blob_service, @@ -32,7 +29,11 @@ impl MemoryPathInfoService { } #[async_trait] -impl PathInfoService for MemoryPathInfoService { +impl PathInfoService for MemoryPathInfoService +where + BS: AsRef + Send + Sync, + DS: AsRef + Send + Sync, +{ async fn get(&self, digest: [u8; 20]) -> Result, Error> { let db = self.db.read().unwrap(); @@ -65,13 +66,9 @@ impl PathInfoService for MemoryPathInfoService { &self, root_node: &castorepb::node::Node, ) -> Result<(u64, [u8; 32]), Error> { - calculate_size_and_sha256( - root_node, - self.blob_service.clone(), - self.directory_service.clone(), - ) - .await - .map_err(|e| Error::StorageError(e.to_string())) + calculate_size_and_sha256(root_node, &self.blob_service, &self.directory_service) + .await + .map_err(|e| Error::StorageError(e.to_string())) } fn list(&self) -> Pin> + Send>> { diff --git a/tvix/store/src/pathinfoservice/nix_http.rs b/tvix/store/src/pathinfoservice/nix_http.rs index ddca35c67a0e..4929355c6dde 100644 --- a/tvix/store/src/pathinfoservice/nix_http.rs +++ b/tvix/store/src/pathinfoservice/nix_http.rs @@ -1,7 +1,6 @@ use std::{ io::{self, BufRead, Read, Write}, pin::Pin, - sync::Arc, }; use data_encoding::BASE64; @@ -38,24 +37,20 @@ use super::PathInfoService; /// [PathInfoService::put] and [PathInfoService::calculate_nar] are not /// implemented and return an error if called. /// TODO: what about reading from nix-cache-info? -pub struct NixHTTPPathInfoService { +pub struct NixHTTPPathInfoService { base_url: url::Url, http_client: reqwest::Client, - blob_service: Arc, - directory_service: Arc, + blob_service: BS, + directory_service: DS, /// An optional list of [narinfo::PubKey]. /// If set, the .narinfo files received need to have correct signature by at least one of these. public_keys: Option>, } -impl NixHTTPPathInfoService { - pub fn new( - base_url: url::Url, - blob_service: Arc, - directory_service: Arc, - ) -> Self { +impl NixHTTPPathInfoService { + pub fn new(base_url: url::Url, blob_service: BS, directory_service: DS) -> Self { Self { base_url, http_client: reqwest::Client::new(), @@ -73,7 +68,11 @@ impl NixHTTPPathInfoService { } #[async_trait] -impl PathInfoService for NixHTTPPathInfoService { +impl PathInfoService for NixHTTPPathInfoService +where + BS: AsRef + Send + Sync + Clone + 'static, + DS: AsRef + Send + Sync + Clone + 'static, +{ #[instrument(skip_all, err, fields(path.digest=BASE64.encode(&digest)))] async fn get(&self, digest: [u8; 20]) -> Result, Error> { let narinfo_url = self diff --git a/tvix/store/src/pathinfoservice/sled.rs b/tvix/store/src/pathinfoservice/sled.rs index ff5ebee476b2..d4d2dedd0061 100644 --- a/tvix/store/src/pathinfoservice/sled.rs +++ b/tvix/store/src/pathinfoservice/sled.rs @@ -3,7 +3,7 @@ use crate::nar::calculate_size_and_sha256; use crate::proto::PathInfo; use futures::{stream::iter, Stream}; use prost::Message; -use std::{path::Path, pin::Pin, sync::Arc}; +use std::{path::Path, pin::Pin}; use tonic::async_trait; use tracing::warn; use tvix_castore::proto as castorepb; @@ -13,18 +13,18 @@ use tvix_castore::{blobservice::BlobService, directoryservice::DirectoryService, /// /// The PathInfo messages are stored as encoded protos, and keyed by their output hash, /// as that's currently the only request type available. -pub struct SledPathInfoService { +pub struct SledPathInfoService { db: sled::Db, - blob_service: Arc, - directory_service: Arc, + blob_service: BS, + directory_service: DS, } -impl SledPathInfoService { +impl SledPathInfoService { pub fn new>( p: P, - blob_service: Arc, - directory_service: Arc, + blob_service: BS, + directory_service: DS, ) -> Result { let config = sled::Config::default() .use_compression(false) // is a required parameter @@ -38,10 +38,7 @@ impl SledPathInfoService { }) } - pub fn new_temporary( - blob_service: Arc, - directory_service: Arc, - ) -> Result { + pub fn new_temporary(blob_service: BS, directory_service: DS) -> Result { let config = sled::Config::default().temporary(true); let db = config.open()?; @@ -54,7 +51,11 @@ impl SledPathInfoService { } #[async_trait] -impl PathInfoService for SledPathInfoService { +impl PathInfoService for SledPathInfoService +where + BS: AsRef + Send + Sync, + DS: AsRef + Send + Sync, +{ async fn get(&self, digest: [u8; 20]) -> Result, Error> { match self.db.get(digest) { Ok(None) => Ok(None), @@ -106,13 +107,9 @@ impl PathInfoService for SledPathInfoService { &self, root_node: &castorepb::node::Node, ) -> Result<(u64, [u8; 32]), Error> { - calculate_size_and_sha256( - root_node, - self.blob_service.clone(), - self.directory_service.clone(), - ) - .await - .map_err(|e| Error::StorageError(e.to_string())) + calculate_size_and_sha256(root_node, &self.blob_service, &self.directory_service) + .await + .map_err(|e| Error::StorageError(e.to_string())) } fn list(&self) -> Pin> + Send>> { diff --git a/tvix/store/src/proto/tests/grpc_pathinfoservice.rs b/tvix/store/src/proto/tests/grpc_pathinfoservice.rs index 66c0f0147615..e8da7792cdb1 100644 --- a/tvix/store/src/proto/tests/grpc_pathinfoservice.rs +++ b/tvix/store/src/proto/tests/grpc_pathinfoservice.rs @@ -22,8 +22,8 @@ fn gen_grpc_service( let blob_service = gen_blob_service(); let directory_service = gen_directory_service(); Arc::new(GRPCPathInfoServiceWrapper::new(gen_pathinfo_service( - blob_service.into(), - directory_service.into(), + blob_service, + directory_service, ))) } diff --git a/tvix/store/src/tests/utils.rs b/tvix/store/src/tests/utils.rs index 961be6e7ac07..040b7ee7f51f 100644 --- a/tvix/store/src/tests/utils.rs +++ b/tvix/store/src/tests/utils.rs @@ -4,9 +4,13 @@ use tvix_castore::{blobservice::BlobService, directoryservice::DirectoryService} pub use tvix_castore::utils::*; -pub fn gen_pathinfo_service( - blob_service: Arc, - directory_service: Arc, -) -> Arc { +pub fn gen_pathinfo_service( + blob_service: BS, + directory_service: DS, +) -> Arc +where + BS: AsRef + Send + Sync + 'static, + DS: AsRef + Send + Sync + 'static, +{ Arc::new(MemoryPathInfoService::new(blob_service, directory_service)) } -- cgit 1.4.1