Merge pull request #11 from n0-computer/mem-rpc-client

feat: Add a lazily initialized in mem client for ffi etc.

Author: rklaehn

Cargo.toml

@@ -36,7 +36,7 @@ oneshot = "0.1.8"
 parking_lot = { version = "0.12.1", optional = true }
 portable-atomic = { version = "1", optional = true }
 postcard = { version = "1", default-features = false, features = ["alloc", "use-std", "experimental-derive"] }
-quic-rpc = { version = "0.15.0", optional = true }
+quic-rpc = { version = "0.15.1", optional = true }
 quic-rpc-derive = { version = "0.15.0", optional = true }
 quinn = { package = "iroh-quinn", version = "0.12", features = ["ring"] }
 rand = "0.8"
@@ -130,3 +130,4 @@ iroh-router = { git = "https://github.com/n0-computer/iroh", branch = "main" }
 iroh-net = { git = "https://github.com/n0-computer/iroh", branch = "main" }
 iroh-metrics = { git = "https://github.com/n0-computer/iroh", branch = "main" }
 iroh-base = { git = "https://github.com/n0-computer/iroh", branch = "main" }
+

src/downloader.rs

@@ -645,7 +645,6 @@ impl<G: Getter<Connection = D::Connection>, D: Dialer> Service<G, D> {
     }
 
     /// Handle receiving a [`Message`].
-    ///
     // This is called in the actor loop, and only async because subscribing to an existing transfer
     // sends the initial state.
     async fn handle_message(&mut self, msg: Message) {

src/net_protocol.rs

@@ -3,7 +3,10 @@
 // TODO: reduce API surface and add documentation
 #![allow(missing_docs)]
 
-use std::{collections::BTreeMap, sync::Arc};
+use std::{
+    collections::BTreeMap,
+    sync::{Arc, OnceLock},
+};
 
 use anyhow::{anyhow, Result};
 use futures_lite::future::Boxed as BoxedFuture;
@@ -36,6 +39,8 @@ pub struct Blobs<S> {
     downloader: Downloader,
     batches: tokio::sync::Mutex<BlobBatches>,
     endpoint: Endpoint,
+    #[cfg(feature = "rpc")]
+    pub(crate) rpc_handler: Arc<OnceLock<crate::rpc::RpcHandler>>,
 }
 
 /// Name used for logging when new node addresses are added from gossip.
@@ -107,6 +112,8 @@ impl<S: crate::store::Store> Blobs<S> {
             downloader,
             endpoint,
             batches: Default::default(),
+            #[cfg(feature = "rpc")]
+            rpc_handler: Arc::new(OnceLock::new()),
         }
     }
 

src/protocol.rs

@@ -148,7 +148,8 @@
 //! # use bao_tree::{ChunkNum, ChunkRanges};
 //! # use iroh_blobs::protocol::{GetRequest, RangeSpecSeq};
 //! # let hash: iroh_blobs::Hash = [0; 32].into();
-//! let ranges = &ChunkRanges::from(..ChunkNum(10)) | &ChunkRanges::from(ChunkNum(100)..ChunkNum(110));
+//! let ranges =
+//!     &ChunkRanges::from(..ChunkNum(10)) | &ChunkRanges::from(ChunkNum(100)..ChunkNum(110));
 //! let spec = RangeSpecSeq::from_ranges([ranges]);
 //! let request = GetRequest::new(hash, spec);
 //! ```
@@ -236,8 +237,8 @@
 //! # use iroh_blobs::protocol::{GetRequest, RangeSpecSeq};
 //! # let hash: iroh_blobs::Hash = [0; 32].into();
 //! let spec = RangeSpecSeq::from_ranges_infinite([
-//!   ChunkRanges::all(), // the collection itself
-//!   ChunkRanges::from(..ChunkNum(1)), // the first chunk of each child
+//!     ChunkRanges::all(),               // the collection itself
+//!     ChunkRanges::from(..ChunkNum(1)), // the first chunk of each child
 //! ]);
 //! let request = GetRequest::new(hash, spec);
 //! ```
@@ -252,9 +253,9 @@
 //! # use iroh_blobs::protocol::{GetRequest, RangeSpecSeq};
 //! # let hash: iroh_blobs::Hash = [0; 32].into();
 //! let spec = RangeSpecSeq::from_ranges([
-//!   ChunkRanges::empty(), // we don't need the collection itself
-//!   ChunkRanges::empty(), // we don't need the first child either
-//!   ChunkRanges::all(), // we need the second child completely
+//!     ChunkRanges::empty(), // we don't need the collection itself
+//!     ChunkRanges::empty(), // we don't need the first child either
+//!     ChunkRanges::all(),   // we need the second child completely
 //! ]);
 //! let request = GetRequest::new(hash, spec);
 //! ```

src/rpc.rs

@@ -7,8 +7,9 @@ use std::{
 
 use anyhow::anyhow;
 use client::{
-    blobs::{BlobInfo, BlobStatus, IncompleteBlobInfo, WrapOption},
+    blobs::{self, BlobInfo, BlobStatus, IncompleteBlobInfo, WrapOption},
     tags::TagInfo,
+    MemConnector,
 };
 use futures_buffered::BufferedStreamExt;
 use futures_lite::StreamExt;
@@ -32,7 +33,11 @@ use proto::{
     },
     Request, RpcError, RpcResult, RpcService,
 };
-use quic_rpc::server::{ChannelTypes, RpcChannel, RpcServerError};
+use quic_rpc::{
+    server::{ChannelTypes, RpcChannel, RpcServerError},
+    RpcClient, RpcServer,
+};
+use tokio_util::task::AbortOnDropHandle;
 
 use crate::{
     export::ExportProgress,
@@ -56,6 +61,16 @@ const RPC_BLOB_GET_CHUNK_SIZE: usize = 1024 * 64;
 const RPC_BLOB_GET_CHANNEL_CAP: usize = 2;
 
 impl<D: crate::store::Store> Blobs<D> {
+    /// Get a client for the blobs protocol
+    pub fn client(self: Arc<Self>) -> blobs::MemClient {
+        let client = self
+            .rpc_handler
+            .get_or_init(|| RpcHandler::new(&self))
+            .client
+            .clone();
+        blobs::Client::new(client)
+    }
+
     /// Handle an RPC request
     pub async fn handle_rpc_request<C>(
         self: Arc<Self>,
@@ -871,3 +886,23 @@ impl<D: crate::store::Store> Blobs<D> {
         Ok(CreateCollectionResponse { hash, tag })
     }
 }
+
+#[derive(Debug)]
+pub(crate) struct RpcHandler {
+    /// Client to hand out
+    client: RpcClient<RpcService, MemConnector>,
+    /// Handler task
+    _handler: AbortOnDropHandle<()>,
+}
+
+impl RpcHandler {
+    fn new<D: crate::store::Store>(blobs: &Arc<Blobs<D>>) -> Self {
+        let blobs = blobs.clone();
+        let (listener, connector) = quic_rpc::transport::flume::channel(1);
+        let listener = RpcServer::new(listener);
+        let client = RpcClient::new(connector);
+        let _handler = listener
+            .spawn_accept_loop(move |req, chan| blobs.clone().handle_rpc_request(req, chan));
+        Self { client, _handler }
+    }
+}

src/rpc/client.rs

@@ -1,10 +1,15 @@
 //! Iroh blobs and tags client
 use anyhow::Result;
 use futures_util::{Stream, StreamExt};
+use quic_rpc::transport::flume::FlumeConnector;
 
 pub mod blobs;
 pub mod tags;
 
+/// Type alias for a memory-backed client.
+pub(crate) type MemConnector =
+    FlumeConnector<crate::rpc::proto::Response, crate::rpc::proto::Request>;
+
 fn flatten<T, E1, E2>(
     s: impl Stream<Item = Result<Result<T, E1>, E2>>,
 ) -> impl Stream<Item = Result<T>>

src/rpc/client/blobs.rs

@@ -111,6 +111,9 @@ pub struct Client<C = BoxedConnector<RpcService>> {
     pub(super) rpc: RpcClient<RpcService, C>,
 }
 
+/// Type alias for a memory-backed client.
+pub type MemClient = Client<crate::rpc::MemConnector>;
+
 impl<C> Client<C>
 where
     C: Connector<RpcService>,
@@ -120,6 +123,11 @@ where
         Self { rpc }
     }
 
+    /// Get a tags client.
+    pub fn tags(&self) -> tags::Client<C> {
+        tags::Client::new(self.rpc.clone())
+    }
+
     /// Check if a blob is completely stored on the node.
     ///
     /// Note that this will return false for blobs that are partially stored on

src/util/fs.rs

@@ -179,7 +179,6 @@ pub struct PathContent {
 }
 
 /// Walks the directory to get the total size and number of files in directory or file
-///
 // TODO: possible combine with `scan_dir`
 pub fn path_content_info(path: impl AsRef<Path>) -> anyhow::Result<PathContent> {
     path_content_info0(path)