Add a way to build a rpc handler that still has the S type parameter

Author: rklaehn

src/net_protocol.rs

@@ -25,7 +25,7 @@ use crate::{
         Stats,
     },
     provider::EventSender,
-    rpc::{client::blobs::MemClient, RpcHandler},
+    rpc::{client::blobs::MemClient, MemRpcHandler, RpcHandler},
     store::{GcConfig, Store},
     util::{
         local_pool::{self, LocalPoolHandle},
@@ -63,7 +63,7 @@ pub(crate) struct BlobsInner<S> {
     endpoint: Endpoint,
     gc_state: Arc<std::sync::Mutex<GcState>>,
     #[cfg(feature = "rpc")]
-    pub(crate) rpc_handler: Arc<OnceLock<crate::rpc::RpcHandler>>,
+    pub(crate) rpc_handler: Arc<OnceLock<crate::rpc::MemRpcHandler>>,
 }
 
 /// Name used for logging when new node addresses are added from gossip.
@@ -137,15 +137,26 @@ impl<S: crate::store::Store> Builder<S> {
     /// Build the Blobs protocol handler.
     /// You need to provide a local pool handle and an endpoint.
     pub fn build(self, rt: &LocalPoolHandle, endpoint: &Endpoint) -> Arc<Blobs> {
+        let inner = self.build_inner(rt, endpoint);
+        Arc::new(Blobs { inner })
+    }
+
+    pub fn build_rpc_handler(self, rt: &LocalPoolHandle, endpoint: &Endpoint) -> RpcHandler<S> {
+        let inner = self.build_inner(rt, endpoint);
+        RpcHandler::from_blobs(inner)
+    }
+
+    /// Build the Blobs protocol handler.
+    /// You need to provide a local pool handle and an endpoint.
+    fn build_inner(self, rt: &LocalPoolHandle, endpoint: &Endpoint) -> Arc<BlobsInner<S>> {
         let downloader = Downloader::new(self.store.clone(), endpoint.clone(), rt.clone());
-        let inner = Arc::new(BlobsInner::new(
+        Arc::new(BlobsInner::new(
             self.store,
             rt.clone(),
             self.events.unwrap_or_default(),
             downloader,
             endpoint.clone(),
-        ));
-        Arc::new(Blobs { inner })
+        ))
     }
 }
 
@@ -394,6 +405,10 @@ pub struct Blobs {
 }
 
 impl Blobs {
+    pub(crate) fn from_inner<S: Store>(inner: Arc<BlobsInner<S>>) -> Self {
+        Self { inner }
+    }
+
     pub fn client(&self) -> MemClient {
         self.inner.clone().client()
     }
@@ -466,7 +481,7 @@ impl<S: crate::store::Store> DynBlobs for BlobsInner<S> {
     fn client(self: Arc<Self>) -> MemClient {
         let client = self
             .rpc_handler
-            .get_or_init(|| RpcHandler::new(&self))
+            .get_or_init(|| MemRpcHandler::new(&self))
             .client
             .clone();
         MemClient::new(client)

src/rpc.rs

@@ -43,7 +43,7 @@ use crate::{
     export::ExportProgress,
     format::collection::Collection,
     get::db::DownloadProgress,
-    net_protocol::{BlobDownloadRequest, BlobsInner},
+    net_protocol::{BlobDownloadRequest, Blobs, BlobsInner},
     provider::{AddProgress, BatchAddPathProgress},
     store::{ConsistencyCheckProgress, ImportProgress, MapEntry, ValidateProgress},
     util::{
@@ -71,7 +71,7 @@ impl<D: crate::store::Store> BlobsInner<D> {
     where
         C: ChannelTypes<RpcService>,
     {
-        Handler {
+        RpcHandler {
             blobs: self.clone(),
         }
         .handle_rpc_request(msg, chan)
@@ -81,11 +81,11 @@ impl<D: crate::store::Store> BlobsInner<D> {
 
 /// RPC handler for the blobs protocol
 #[derive(Debug, Clone)]
-pub struct Handler<S> {
+pub struct RpcHandler<S> {
     blobs: Arc<BlobsInner<S>>,
 }
 
-impl<D: crate::store::Store> Handler<D> {
+impl<D: crate::store::Store> RpcHandler<D> {
     fn store(&self) -> &D {
         &self.blobs.store
     }
@@ -94,6 +94,15 @@ impl<D: crate::store::Store> Handler<D> {
         &self.blobs.rt
     }
 
+    pub(crate) fn from_blobs(blobs: Arc<BlobsInner<D>>) -> Self {
+        Self { blobs }
+    }
+
+    /// Get the blobs ProtocolHandler
+    pub fn blobs(&self) -> Arc<Blobs> {
+        Arc::new(Blobs::from_inner(self.blobs.clone()))
+    }
+
     /// Handle an RPC request
     pub async fn handle_rpc_request<C>(
         self,
@@ -900,14 +909,14 @@ impl<D: crate::store::Store> Handler<D> {
 }
 
 #[derive(Debug)]
-pub(crate) struct RpcHandler {
+pub(crate) struct MemRpcHandler {
     /// Client to hand out
     pub(crate) client: RpcClient<RpcService, MemConnector>,
     /// Handler task
     handler: AbortOnDropHandle<()>,
 }
 
-impl RpcHandler {
+impl MemRpcHandler {
     pub fn new<D: crate::store::Store>(blobs: &Arc<BlobsInner<D>>) -> Self {
         let blobs = blobs.clone();
         let (listener, connector) = quic_rpc::transport::flume::channel(1);