fix: Fix the task leak with the lazy in-mem rpc client while still keeping it lazy (#31)

* Make blobs more cheaply cloneable by by giving it an Inner

* Remove the lazy part.

The lazy handler kept a reference to Blobs alive. This caused both the
task and the blobs to never be dropped. To solve this you can just split
the inner part in 2 parts, one that has the handle and one that has the
logic. But that is not nice. I think it is best for the mem rpc handler
to exist completely separately, especially given that rpc is a non-default
feature.

* spawn_rpc should make it sufficiently clear that this is a thing you need to put away somewhere.

Or maybe spawn_client?

* back to the lazy client

* add comment about the purpose of the handler

Author: rklaehn

examples/custom-protocol.rs

@@ -122,7 +122,7 @@ async fn main() -> Result<()> {
 
             // Print out our query results.
             for hash in hashes {
-                read_and_print(&blobs_client, hash).await?;
+                read_and_print(blobs_client, hash).await?;
             }
         }
     }

src/net_protocol.rs

@@ -5,28 +5,22 @@
 
 use std::{collections::BTreeSet, fmt::Debug, ops::DerefMut, sync::Arc};
 
-use anyhow::{anyhow, bail, Result};
+use anyhow::{bail, Result};
 use futures_lite::future::Boxed as BoxedFuture;
 use futures_util::future::BoxFuture;
 use iroh::{endpoint::Connecting, protocol::ProtocolHandler, Endpoint, NodeAddr};
 use iroh_base::hash::{BlobFormat, Hash};
 use serde::{Deserialize, Serialize};
-use tracing::{debug, warn};
+use tracing::debug;
 
 use crate::{
-    downloader::{DownloadRequest, Downloader},
-    get::{
-        db::{DownloadProgress, GetState},
-        Stats,
-    },
+    downloader::Downloader,
     provider::EventSender,
     store::GcConfig,
     util::{
         local_pool::{self, LocalPoolHandle},
-        progress::{AsyncChannelProgressSender, ProgressSender},
         SetTagOption,
     },
-    HashAndFormat,
 };
 
 /// A callback that blobs can ask about a set of hashes that should not be garbage collected.
@@ -47,16 +41,21 @@ impl Default for GcState {
     }
 }
 
-#[derive(Debug, Clone)]
-pub struct Blobs<S> {
-    rt: LocalPoolHandle,
+#[derive(Debug)]
+pub(crate) struct BlobsInner<S> {
+    pub(crate) rt: LocalPoolHandle,
     pub(crate) store: S,
     events: EventSender,
-    downloader: Downloader,
+    pub(crate) downloader: Downloader,
+    pub(crate) endpoint: Endpoint,
+    gc_state: std::sync::Mutex<GcState>,
     #[cfg(feature = "rpc")]
-    batches: Arc<tokio::sync::Mutex<BlobBatches>>,
-    endpoint: Endpoint,
-    gc_state: Arc<std::sync::Mutex<GcState>>,
+    pub(crate) batches: tokio::sync::Mutex<BlobBatches>,
+}
+
+#[derive(Debug, Clone)]
+pub struct Blobs<S> {
+    pub(crate) inner: Arc<BlobsInner<S>>,
     #[cfg(feature = "rpc")]
     pub(crate) rpc_handler: Arc<std::sync::OnceLock<crate::rpc::RpcHandler>>,
 }
@@ -76,7 +75,7 @@ pub(crate) struct BlobBatches {
 #[derive(Debug, Default)]
 struct BlobBatch {
     /// The tags in this batch.
-    tags: std::collections::BTreeMap<HashAndFormat, Vec<crate::TempTag>>,
+    tags: std::collections::BTreeMap<iroh::hash::HashAndFormat, Vec<crate::TempTag>>,
 }
 
 #[cfg(feature = "rpc")]
@@ -95,7 +94,11 @@ impl BlobBatches {
     }
 
     /// Remove a tag from a batch.
-    pub fn remove_one(&mut self, batch: BatchId, content: &HashAndFormat) -> Result<()> {
+    pub fn remove_one(
+        &mut self,
+        batch: BatchId,
+        content: &iroh::hash::HashAndFormat,
+    ) -> Result<()> {
         if let Some(batch) = self.batches.get_mut(&batch) {
             if let Some(tags) = batch.tags.get_mut(content) {
                 tags.pop();
@@ -178,40 +181,46 @@ impl<S: crate::store::Store> Blobs<S> {
         endpoint: Endpoint,
     ) -> Self {
         Self {
-            rt,
-            store,
-            events,
-            downloader,
-            endpoint,
-            #[cfg(feature = "rpc")]
-            batches: Default::default(),
-            gc_state: Default::default(),
+            inner: Arc::new(BlobsInner {
+                rt,
+                store,
+                events,
+                downloader,
+                endpoint,
+                #[cfg(feature = "rpc")]
+                batches: Default::default(),
+                gc_state: Default::default(),
+            }),
             #[cfg(feature = "rpc")]
             rpc_handler: Default::default(),
         }
     }
 
     pub fn store(&self) -> &S {
-        &self.store
+        &self.inner.store
+    }
+
+    pub fn events(&self) -> &EventSender {
+        &self.inner.events
     }
 
     pub fn rt(&self) -> &LocalPoolHandle {
-        &self.rt
+        &self.inner.rt
     }
 
     pub fn downloader(&self) -> &Downloader {
-        &self.downloader
+        &self.inner.downloader
     }
 
     pub fn endpoint(&self) -> &Endpoint {
-        &self.endpoint
+        &self.inner.endpoint
     }
 
     /// Add a callback that will be called before the garbage collector runs.
     ///
     /// This can only be called before the garbage collector has started, otherwise it will return an error.
     pub fn add_protected(&self, cb: ProtectCb) -> Result<()> {
-        let mut state = self.gc_state.lock().unwrap();
+        let mut state = self.inner.gc_state.lock().unwrap();
         match &mut *state {
             GcState::Initial(cbs) => {
                 cbs.push(cb);
@@ -225,7 +234,7 @@ impl<S: crate::store::Store> Blobs<S> {
 
     /// Start garbage collection with the given settings.
     pub fn start_gc(&self, config: GcConfig) -> Result<()> {
-        let mut state = self.gc_state.lock().unwrap();
+        let mut state = self.inner.gc_state.lock().unwrap();
         let protected = match state.deref_mut() {
             GcState::Initial(items) => std::mem::take(items),
             GcState::Started(_) => bail!("gc already started"),
@@ -241,161 +250,20 @@ impl<S: crate::store::Store> Blobs<S> {
                 set
             }
         };
-        let store = self.store.clone();
+        let store = self.store().clone();
         let run = self
-            .rt
+            .rt()
             .spawn(move || async move { store.gc_run(config, protected_cb).await });
         *state = GcState::Started(Some(run));
         Ok(())
     }
-
-    #[cfg(feature = "rpc")]
-    pub(crate) async fn batches(&self) -> tokio::sync::MutexGuard<'_, BlobBatches> {
-        self.batches.lock().await
-    }
-
-    pub(crate) async fn download(
-        &self,
-        endpoint: Endpoint,
-        req: BlobDownloadRequest,
-        progress: AsyncChannelProgressSender<DownloadProgress>,
-    ) -> Result<()> {
-        let BlobDownloadRequest {
-            hash,
-            format,
-            nodes,
-            tag,
-            mode,
-        } = req;
-        let hash_and_format = HashAndFormat { hash, format };
-        let temp_tag = self.store.temp_tag(hash_and_format);
-        let stats = match mode {
-            DownloadMode::Queued => {
-                self.download_queued(endpoint, hash_and_format, nodes, progress.clone())
-                    .await?
-            }
-            DownloadMode::Direct => {
-                self.download_direct_from_nodes(endpoint, hash_and_format, nodes, progress.clone())
-                    .await?
-            }
-        };
-
-        progress.send(DownloadProgress::AllDone(stats)).await.ok();
-        match tag {
-            SetTagOption::Named(tag) => {
-                self.store.set_tag(tag, Some(hash_and_format)).await?;
-            }
-            SetTagOption::Auto => {
-                self.store.create_tag(hash_and_format).await?;
-            }
-        }
-        drop(temp_tag);
-
-        Ok(())
-    }
-
-    async fn download_queued(
-        &self,
-        endpoint: Endpoint,
-        hash_and_format: HashAndFormat,
-        nodes: Vec<NodeAddr>,
-        progress: AsyncChannelProgressSender<DownloadProgress>,
-    ) -> Result<Stats> {
-        /// Name used for logging when new node addresses are added from gossip.
-        const BLOB_DOWNLOAD_SOURCE_NAME: &str = "blob_download";
-
-        let mut node_ids = Vec::with_capacity(nodes.len());
-        let mut any_added = false;
-        for node in nodes {
-            node_ids.push(node.node_id);
-            if !node.info.is_empty() {
-                endpoint.add_node_addr_with_source(node, BLOB_DOWNLOAD_SOURCE_NAME)?;
-                any_added = true;
-            }
-        }
-        let can_download = !node_ids.is_empty() && (any_added || endpoint.discovery().is_some());
-        anyhow::ensure!(can_download, "no way to reach a node for download");
-        let req = DownloadRequest::new(hash_and_format, node_ids).progress_sender(progress);
-        let handle = self.downloader.queue(req).await;
-        let stats = handle.await?;
-        Ok(stats)
-    }
-
-    #[tracing::instrument("download_direct", skip_all, fields(hash=%hash_and_format.hash.fmt_short()))]
-    async fn download_direct_from_nodes(
-        &self,
-        endpoint: Endpoint,
-        hash_and_format: HashAndFormat,
-        nodes: Vec<NodeAddr>,
-        progress: AsyncChannelProgressSender<DownloadProgress>,
-    ) -> Result<Stats> {
-        let mut last_err = None;
-        let mut remaining_nodes = nodes.len();
-        let mut nodes_iter = nodes.into_iter();
-        'outer: loop {
-            match crate::get::db::get_to_db_in_steps(
-                self.store.clone(),
-                hash_and_format,
-                progress.clone(),
-            )
-            .await?
-            {
-                GetState::Complete(stats) => return Ok(stats),
-                GetState::NeedsConn(needs_conn) => {
-                    let (conn, node_id) = 'inner: loop {
-                        match nodes_iter.next() {
-                            None => break 'outer,
-                            Some(node) => {
-                                remaining_nodes -= 1;
-                                let node_id = node.node_id;
-                                if node_id == endpoint.node_id() {
-                                    debug!(
-                                        ?remaining_nodes,
-                                        "skip node {} (it is the node id of ourselves)",
-                                        node_id.fmt_short()
-                                    );
-                                    continue 'inner;
-                                }
-                                match endpoint.connect(node, crate::protocol::ALPN).await {
-                                    Ok(conn) => break 'inner (conn, node_id),
-                                    Err(err) => {
-                                        debug!(
-                                            ?remaining_nodes,
-                                            "failed to connect to {}: {err}",
-                                            node_id.fmt_short()
-                                        );
-                                        continue 'inner;
-                                    }
-                                }
-                            }
-                        }
-                    };
-                    match needs_conn.proceed(conn).await {
-                        Ok(stats) => return Ok(stats),
-                        Err(err) => {
-                            warn!(
-                                ?remaining_nodes,
-                                "failed to download from {}: {err}",
-                                node_id.fmt_short()
-                            );
-                            last_err = Some(err);
-                        }
-                    }
-                }
-            }
-        }
-        match last_err {
-            Some(err) => Err(err.into()),
-            None => Err(anyhow!("No nodes to download from provided")),
-        }
-    }
 }
 
 impl<S: crate::store::Store> ProtocolHandler for Blobs<S> {
     fn accept(&self, conn: Connecting) -> BoxedFuture<Result<()>> {
-        let db = self.store.clone();
-        let events = self.events.clone();
-        let rt = self.rt.clone();
+        let db = self.store().clone();
+        let events = self.events().clone();
+        let rt = self.rt().clone();
 
         Box::pin(async move {
             crate::provider::handle_connection(conn.await?, db, events, rt).await;
@@ -404,7 +272,7 @@ impl<S: crate::store::Store> ProtocolHandler for Blobs<S> {
     }
 
     fn shutdown(&self) -> BoxedFuture<()> {
-        let store = self.store.clone();
+        let store = self.store().clone();
         Box::pin(async move {
             store.shutdown().await;
         })

src/rpc.rs

@@ -2,19 +2,21 @@
 
 use std::{
     io,
+    ops::Deref,
     sync::{Arc, Mutex},
 };
 
 use anyhow::anyhow;
 use client::{
-    blobs::{self, BlobInfo, BlobStatus, IncompleteBlobInfo, WrapOption},
+    blobs::{self, BlobInfo, BlobStatus, DownloadMode, IncompleteBlobInfo, MemClient, WrapOption},
     tags::TagInfo,
     MemConnector,
 };
 use futures_buffered::BufferedStreamExt;
 use futures_lite::StreamExt;
 use futures_util::{FutureExt, Stream};
 use genawaiter::sync::{Co, Gen};
+use iroh::{Endpoint, NodeAddr};
 use iroh_base::hash::{BlobFormat, HashAndFormat};
 use iroh_io::AsyncSliceReader;
 use proto::{
@@ -38,15 +40,21 @@ use quic_rpc::{
     RpcClient, RpcServer,
 };
 use tokio_util::task::AbortOnDropHandle;
+use tracing::{debug, warn};
 
 use crate::{
+    downloader::{DownloadRequest, Downloader},
     export::ExportProgress,
     format::collection::Collection,
-    get::db::DownloadProgress,
-    net_protocol::{BlobDownloadRequest, Blobs},
+    get::{
+        db::{DownloadProgress, GetState},
+        Stats,
+    },
+    net_protocol::{BlobDownloadRequest, Blobs, BlobsInner},
     provider::{AddProgress, BatchAddPathProgress},
     store::{ConsistencyCheckProgress, ImportProgress, MapEntry, ValidateProgress},
     util::{
+        local_pool::LocalPoolHandle,
         progress::{AsyncChannelProgressSender, ProgressSender},
         SetTagOption,
     },
@@ -62,13 +70,63 @@ 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) -> blobs::MemClient {
-        let client = self
+    pub fn client(&self) -> &blobs::MemClient {
+        &self
             .rpc_handler
-            .get_or_init(|| RpcHandler::new(self))
+            .get_or_init(|| RpcHandler::new(&self.inner))
             .client
-            .clone();
-        blobs::Client::new(client)
+    }
+
+    /// Handle an RPC request
+    pub async fn handle_rpc_request<C>(
+        self,
+        msg: Request,
+        chan: RpcChannel<RpcService, C>,
+    ) -> std::result::Result<(), RpcServerError<C>>
+    where
+        C: ChannelTypes<RpcService>,
+    {
+        Handler(self.inner.clone())
+            .handle_rpc_request(msg, chan)
+            .await
+    }
+}
+
+/// This is just an internal helper so I don't have to
+/// define all the rpc methods on `self: Arc<BlobsInner<S>>`
+#[derive(Clone)]
+struct Handler<S>(Arc<BlobsInner<S>>);
+
+impl<S> Deref for Handler<S> {
+    type Target = BlobsInner<S>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<D: crate::store::Store> Handler<D> {
+    fn store(&self) -> &D {
+        &self.0.store
+    }
+
+    fn rt(&self) -> &LocalPoolHandle {
+        &self.0.rt
+    }
+
+    fn endpoint(&self) -> &Endpoint {
+        &self.0.endpoint
+    }
+
+    fn downloader(&self) -> &Downloader {
+        &self.0.downloader
+    }
+
+    #[cfg(feature = "rpc")]
+    pub(crate) async fn batches(
+        &self,
+    ) -> tokio::sync::MutexGuard<'_, crate::net_protocol::BlobBatches> {
+        self.0.batches.lock().await
     }
 
     /// Handle an RPC request
@@ -872,24 +930,178 @@ impl<D: crate::store::Store> Blobs<D> {
 
         Ok(CreateCollectionResponse { hash, tag })
     }
+
+    pub(crate) async fn download(
+        &self,
+        endpoint: Endpoint,
+        req: BlobDownloadRequest,
+        progress: AsyncChannelProgressSender<DownloadProgress>,
+    ) -> anyhow::Result<()> {
+        let BlobDownloadRequest {
+            hash,
+            format,
+            nodes,
+            tag,
+            mode,
+        } = req;
+        let hash_and_format = HashAndFormat { hash, format };
+        let temp_tag = self.store().temp_tag(hash_and_format);
+        let stats = match mode {
+            DownloadMode::Queued => {
+                self.download_queued(endpoint, hash_and_format, nodes, progress.clone())
+                    .await?
+            }
+            DownloadMode::Direct => {
+                self.download_direct_from_nodes(endpoint, hash_and_format, nodes, progress.clone())
+                    .await?
+            }
+        };
+
+        progress.send(DownloadProgress::AllDone(stats)).await.ok();
+        match tag {
+            SetTagOption::Named(tag) => {
+                self.store().set_tag(tag, Some(hash_and_format)).await?;
+            }
+            SetTagOption::Auto => {
+                self.store().create_tag(hash_and_format).await?;
+            }
+        }
+        drop(temp_tag);
+
+        Ok(())
+    }
+
+    async fn download_queued(
+        &self,
+        endpoint: Endpoint,
+        hash_and_format: HashAndFormat,
+        nodes: Vec<NodeAddr>,
+        progress: AsyncChannelProgressSender<DownloadProgress>,
+    ) -> anyhow::Result<Stats> {
+        /// Name used for logging when new node addresses are added from gossip.
+        const BLOB_DOWNLOAD_SOURCE_NAME: &str = "blob_download";
+
+        let mut node_ids = Vec::with_capacity(nodes.len());
+        let mut any_added = false;
+        for node in nodes {
+            node_ids.push(node.node_id);
+            if !node.info.is_empty() {
+                endpoint.add_node_addr_with_source(node, BLOB_DOWNLOAD_SOURCE_NAME)?;
+                any_added = true;
+            }
+        }
+        let can_download = !node_ids.is_empty() && (any_added || endpoint.discovery().is_some());
+        anyhow::ensure!(can_download, "no way to reach a node for download");
+        let req = DownloadRequest::new(hash_and_format, node_ids).progress_sender(progress);
+        let handle = self.downloader().queue(req).await;
+        let stats = handle.await?;
+        Ok(stats)
+    }
+
+    #[tracing::instrument("download_direct", skip_all, fields(hash=%hash_and_format.hash.fmt_short()))]
+    async fn download_direct_from_nodes(
+        &self,
+        endpoint: Endpoint,
+        hash_and_format: HashAndFormat,
+        nodes: Vec<NodeAddr>,
+        progress: AsyncChannelProgressSender<DownloadProgress>,
+    ) -> anyhow::Result<Stats> {
+        let mut last_err = None;
+        let mut remaining_nodes = nodes.len();
+        let mut nodes_iter = nodes.into_iter();
+        'outer: loop {
+            match crate::get::db::get_to_db_in_steps(
+                self.store().clone(),
+                hash_and_format,
+                progress.clone(),
+            )
+            .await?
+            {
+                GetState::Complete(stats) => return Ok(stats),
+                GetState::NeedsConn(needs_conn) => {
+                    let (conn, node_id) = 'inner: loop {
+                        match nodes_iter.next() {
+                            None => break 'outer,
+                            Some(node) => {
+                                remaining_nodes -= 1;
+                                let node_id = node.node_id;
+                                if node_id == endpoint.node_id() {
+                                    debug!(
+                                        ?remaining_nodes,
+                                        "skip node {} (it is the node id of ourselves)",
+                                        node_id.fmt_short()
+                                    );
+                                    continue 'inner;
+                                }
+                                match endpoint.connect(node, crate::protocol::ALPN).await {
+                                    Ok(conn) => break 'inner (conn, node_id),
+                                    Err(err) => {
+                                        debug!(
+                                            ?remaining_nodes,
+                                            "failed to connect to {}: {err}",
+                                            node_id.fmt_short()
+                                        );
+                                        continue 'inner;
+                                    }
+                                }
+                            }
+                        }
+                    };
+                    match needs_conn.proceed(conn).await {
+                        Ok(stats) => return Ok(stats),
+                        Err(err) => {
+                            warn!(
+                                ?remaining_nodes,
+                                "failed to download from {}: {err}",
+                                node_id.fmt_short()
+                            );
+                            last_err = Some(err);
+                        }
+                    }
+                }
+            }
+        }
+        match last_err {
+            Some(err) => Err(err.into()),
+            None => Err(anyhow!("No nodes to download from provided")),
+        }
+    }
 }
 
+/// An in memory rpc handler for the blobs rpc protocol
+///
+/// This struct contains both a task that handles rpc requests and a client
+/// that can be used to send rpc requests.
+///
+/// Dropping it will stop the handler task, so you need to put it somewhere
+/// where it will be kept alive. This struct will capture a copy of
+/// [`crate::net_protocol::Blobs`] and keep it alive.
 #[derive(Debug)]
 pub(crate) struct RpcHandler {
     /// Client to hand out
-    client: RpcClient<RpcService, MemConnector>,
+    client: MemClient,
     /// Handler task
     _handler: AbortOnDropHandle<()>,
 }
 
+impl Deref for RpcHandler {
+    type Target = MemClient;
+
+    fn deref(&self) -> &Self::Target {
+        &self.client
+    }
+}
+
 impl RpcHandler {
-    fn new<D: crate::store::Store>(blobs: &Blobs<D>) -> Self {
+    fn new<D: crate::store::Store>(blobs: &Arc<BlobsInner<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));
+        let client = MemClient::new(client);
+        let _handler = listener.spawn_accept_loop(move |req, chan| {
+            Handler(blobs.clone()).handle_rpc_request(req, chan)
+        });
         Self { client, _handler }
     }
 }

tests/blobs.rs

@@ -13,7 +13,7 @@ async fn blobs_gc_smoke() -> TestResult<()> {
     let pool = LocalPool::default();
     let endpoint = Endpoint::builder().bind().await?;
     let blobs = Blobs::memory().build(pool.handle(), &endpoint);
-    let client = blobs.clone().client();
+    let client = blobs.client();
     blobs.start_gc(GcConfig {
         period: Duration::from_millis(1),
         done_callback: None,
@@ -32,12 +32,7 @@ async fn blobs_gc_protected() -> TestResult<()> {
     let pool = LocalPool::default();
     let endpoint = Endpoint::builder().bind().await?;
     let blobs = Blobs::memory().build(pool.handle(), &endpoint);
-    let client: iroh_blobs::rpc::client::blobs::Client<
-        quic_rpc::transport::flume::FlumeConnector<
-            iroh_blobs::rpc::proto::Response,
-            iroh_blobs::rpc::proto::Request,
-        >,
-    > = blobs.clone().client();
+    let client = blobs.client();
     let h1 = client.add_bytes(b"test".to_vec()).await?;
     let protected = Arc::new(Mutex::new(Vec::new()));
     blobs.add_protected(Box::new({

tests/gc.rs

@@ -66,8 +66,8 @@ impl<S: Store> Node<S> {
     }
 
     /// Returns an in-memory blobs client
-    pub fn blobs(&self) -> blobs::MemClient {
-        self.blobs.clone().client()
+    pub fn blobs(&self) -> &blobs::MemClient {
+        self.blobs.client()
     }
 
     /// Returns an in-memory tags client