diff options
author | Florian Klink <flokli@flokli.de> | 2024-01-12T10·18+0200 |
---|---|---|
committer | clbot <clbot@tvl.fyi> | 2024-01-12T20·37+0000 |
commit | b59df53774acc654ea4b23f02ccf5529587bceff (patch) | |
tree | fa35191cd12dbca2478884616de2527eb89066e2 /tvix | |
parent | 7d51193f7db8d6126ee2970eadf009a5d87b694f (diff) |
refactor(tvix/store/pathinfoservice): make more generic r/7371
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 <flokli@flokli.de> Tested-by: BuildkiteCI Reviewed-by: raitobezarius <tvl@lahfa.xyz>
Diffstat (limited to '')
-rw-r--r-- | tvix/store/src/pathinfoservice/grpc.rs | 4 | ||||
-rw-r--r-- | tvix/store/src/pathinfoservice/memory.rs | 29 | ||||
-rw-r--r-- | tvix/store/src/pathinfoservice/nix_http.rs | 21 | ||||
-rw-r--r-- | tvix/store/src/pathinfoservice/sled.rs | 35 | ||||
-rw-r--r-- | tvix/store/src/proto/tests/grpc_pathinfoservice.rs | 4 | ||||
-rw-r--r-- | tvix/store/src/tests/utils.rs | 12 |
6 files changed, 51 insertions, 54 deletions
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<dyn PathInfoService>), ), 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<BS, DS> { db: Arc<RwLock<HashMap<[u8; 20], PathInfo>>>, - blob_service: Arc<dyn BlobService>, - directory_service: Arc<dyn DirectoryService>, + blob_service: BS, + directory_service: DS, } -impl MemoryPathInfoService { - pub fn new( - blob_service: Arc<dyn BlobService>, - directory_service: Arc<dyn DirectoryService>, - ) -> Self { +impl<BS, DS> MemoryPathInfoService<BS, DS> { + 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<BS, DS> PathInfoService for MemoryPathInfoService<BS, DS> +where + BS: AsRef<dyn BlobService> + Send + Sync, + DS: AsRef<dyn DirectoryService> + Send + Sync, +{ async fn get(&self, digest: [u8; 20]) -> Result<Option<PathInfo>, 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<Box<dyn Stream<Item = Result<PathInfo, Error>> + 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<BS, DS> { base_url: url::Url, http_client: reqwest::Client, - blob_service: Arc<dyn BlobService>, - directory_service: Arc<dyn DirectoryService>, + 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<Vec<narinfo::PubKey>>, } -impl NixHTTPPathInfoService { - pub fn new( - base_url: url::Url, - blob_service: Arc<dyn BlobService>, - directory_service: Arc<dyn DirectoryService>, - ) -> Self { +impl<BS, DS> NixHTTPPathInfoService<BS, DS> { + 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<BS, DS> PathInfoService for NixHTTPPathInfoService<BS, DS> +where + BS: AsRef<dyn BlobService> + Send + Sync + Clone + 'static, + DS: AsRef<dyn DirectoryService> + Send + Sync + Clone + 'static, +{ #[instrument(skip_all, err, fields(path.digest=BASE64.encode(&digest)))] async fn get(&self, digest: [u8; 20]) -> Result<Option<PathInfo>, 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<BS, DS> { db: sled::Db, - blob_service: Arc<dyn BlobService>, - directory_service: Arc<dyn DirectoryService>, + blob_service: BS, + directory_service: DS, } -impl SledPathInfoService { +impl<BS, DS> SledPathInfoService<BS, DS> { pub fn new<P: AsRef<Path>>( p: P, - blob_service: Arc<dyn BlobService>, - directory_service: Arc<dyn DirectoryService>, + blob_service: BS, + directory_service: DS, ) -> Result<Self, sled::Error> { 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<dyn BlobService>, - directory_service: Arc<dyn DirectoryService>, - ) -> Result<Self, sled::Error> { + pub fn new_temporary(blob_service: BS, directory_service: DS) -> Result<Self, sled::Error> { let config = sled::Config::default().temporary(true); let db = config.open()?; @@ -54,7 +51,11 @@ impl SledPathInfoService { } #[async_trait] -impl PathInfoService for SledPathInfoService { +impl<BS, DS> PathInfoService for SledPathInfoService<BS, DS> +where + BS: AsRef<dyn BlobService> + Send + Sync, + DS: AsRef<dyn DirectoryService> + Send + Sync, +{ async fn get(&self, digest: [u8; 20]) -> Result<Option<PathInfo>, 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<Box<dyn Stream<Item = Result<PathInfo, Error>> + 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<dyn BlobService>, - directory_service: Arc<dyn DirectoryService>, -) -> Arc<dyn PathInfoService> { +pub fn gen_pathinfo_service<BS, DS>( + blob_service: BS, + directory_service: DS, +) -> Arc<dyn PathInfoService> +where + BS: AsRef<dyn BlobService> + Send + Sync + 'static, + DS: AsRef<dyn DirectoryService> + Send + Sync + 'static, +{ Arc::new(MemoryPathInfoService::new(blob_service, directory_service)) } |