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>

6 files changed, 51 insertions, 54 deletions

diff --git a/tvix/store/src/pathinfoservice/grpc.rs b/tvix/store/src/pathinfoservice/grpc.rs
index 4ec1467525..5754913331 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 6f96c08347..9f657a9c62 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 ddca35c67a..4929355c6d 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 ff5ebee476..d4d2dedd00 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 66c0f01476..e8da7792cd 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 961be6e7ac..040b7ee7f5 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))
 }