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

feat: add mem rpc client

Author: rklaehn

Cargo.toml

@@ -59,7 +59,7 @@ tracing = "0.1"
 
 # rpc
 nested_enum_utils = { version = "0.1.0", optional = true }
-quic-rpc = { version = "0.15", optional = true }
+quic-rpc = { version = "0.15.1", optional = true }
 quic-rpc-derive = { version = "0.15", optional = true }
 serde-error = { version = "0.1.3", optional = true }
 portable-atomic = { version = "1.9.0", optional = true }

src/engine.rs

@@ -56,6 +56,8 @@ pub struct Engine<D> {
     content_status_cb: ContentStatusCallback,
     local_pool_handle: LocalPoolHandle,
     blob_store: D,
+    #[cfg(feature = "rpc")]
+    pub(crate) rpc_handler: Arc<std::sync::OnceLock<crate::rpc::RpcHandler>>,
 }
 
 impl<D: iroh_blobs::store::Store> Engine<D> {
@@ -118,6 +120,8 @@ impl<D: iroh_blobs::store::Store> Engine<D> {
             default_author: Arc::new(default_author),
             local_pool_handle,
             blob_store: bao_store,
+            #[cfg(feature = "rpc")]
+            rpc_handler: Default::default(),
         })
     }
 

src/rpc.rs

@@ -1,7 +1,11 @@
 //! Quic RPC implementation for docs.
 
-use proto::RpcService;
-use quic_rpc::server::{ChannelTypes, RpcChannel};
+use proto::{Request, RpcService};
+use quic_rpc::{
+    server::{ChannelTypes, RpcChannel},
+    RpcClient, RpcServer,
+};
+use tokio_util::task::AbortOnDropHandle;
 
 use crate::engine::Engine;
 
@@ -14,15 +18,22 @@ type RpcError = serde_error::Error;
 type RpcResult<T> = std::result::Result<T, RpcError>;
 
 impl<D: iroh_blobs::store::Store> Engine<D> {
+    /// Get an in memory client to interact with the docs engine.
+    pub fn client(&self) -> &client::docs::MemClient {
+        &self
+            .rpc_handler
+            .get_or_init(|| RpcHandler::new(self))
+            .client
+    }
+
     /// Handle a docs request from the RPC server.
     pub async fn handle_rpc_request<C: ChannelTypes<RpcService>>(
-        &self,
-        msg: crate::rpc::proto::Request,
+        self,
+        msg: Request,
         chan: RpcChannel<RpcService, C>,
     ) -> Result<(), quic_rpc::server::RpcServerError<C>> {
-        use crate::rpc::proto::Request::*;
-
-        let this = self.clone();
+        use Request::*;
+        let this = self;
         match msg {
             Open(msg) => chan.rpc(msg, this, Self::doc_open).await,
             Close(msg) => chan.rpc(msg, this, Self::doc_close).await,
@@ -65,3 +76,23 @@ impl<D: iroh_blobs::store::Store> Engine<D> {
         }
     }
 }
+
+#[derive(Debug)]
+pub(crate) struct RpcHandler {
+    /// Client to hand out
+    client: client::docs::MemClient,
+    /// Handler task
+    _handler: AbortOnDropHandle<()>,
+}
+
+impl RpcHandler {
+    fn new<D: iroh_blobs::store::Store>(engine: &Engine<D>) -> Self {
+        let engine = engine.clone();
+        let (listener, connector) = quic_rpc::transport::flume::channel(1);
+        let listener = RpcServer::new(listener);
+        let client = client::docs::MemClient::new(RpcClient::new(connector));
+        let _handler = listener
+            .spawn_accept_loop(move |req, chan| engine.clone().handle_rpc_request(req, chan));
+        Self { client, _handler }
+    }
+}

src/rpc/client/authors.rs

@@ -19,6 +19,7 @@ use crate::{
 
 /// Iroh docs client.
 #[derive(Debug, Clone)]
+#[repr(transparent)]
 pub struct Client<C = BoxedConnector<RpcService>> {
     pub(super) rpc: quic_rpc::RpcClient<RpcService, C>,
 }

src/rpc/client/docs.rs

@@ -16,10 +16,12 @@ use iroh_base::node_addr::AddrInfoOptions;
 use iroh_blobs::{export::ExportProgress, store::ExportMode, Hash};
 use iroh_net::NodeAddr;
 use portable_atomic::{AtomicBool, Ordering};
-use quic_rpc::{client::BoxedConnector, message::RpcMsg, Connector};
+use quic_rpc::{
+    client::BoxedConnector, message::RpcMsg, transport::flume::FlumeConnector, Connector,
+};
 use serde::{Deserialize, Serialize};
 
-use super::flatten;
+use super::{authors, flatten};
 use crate::{
     actor::OpenState,
     rpc::proto::{
@@ -38,8 +40,13 @@ pub use crate::{
     Entry,
 };
 
+/// Type alias for a memory-backed client.
+pub type MemClient =
+    Client<FlumeConnector<crate::rpc::proto::Response, crate::rpc::proto::Request>>;
+
 /// Iroh docs client.
 #[derive(Debug, Clone)]
+#[repr(transparent)]
 pub struct Client<C = BoxedConnector<RpcService>> {
     pub(super) rpc: quic_rpc::RpcClient<RpcService, C>,
 }
@@ -50,6 +57,11 @@ impl<C: Connector<RpcService>> Client<C> {
         Self { rpc }
     }
 
+    /// Returns an authors client.
+    pub fn authors(&self) -> authors::Client<C> {
+        authors::Client::new(self.rpc.clone())
+    }
+
     /// Creates a client.
     pub async fn create(&self) -> Result<Doc<C>> {
         let res = self.rpc.rpc(CreateRequest {}).await??;