feat(iroh): pass a runtime to Doc client to spawn close task on drop (#1758)

We spawn a tokio task when dropping a client `Doc` handle to notify the
node that the document was closed. This does not work in the FFI
currently because the FFI does not run in a global tokio context.
Instead we now always pass a `runtime::Handle` into the client which can
be used to spawn tasks independently of running in a tokio context.

---------

Co-authored-by: Franz Heinzmann (Frando) <frando@unbiskant.org>

Author: dignifiedquire

iroh/src/client.rs

@@ -17,6 +17,7 @@ use futures::stream::BoxStream;
 use futures::{SinkExt, Stream, StreamExt, TryStreamExt};
 use iroh_bytes::provider::AddProgress;
 use iroh_bytes::store::ValidateProgress;
+use iroh_bytes::util::runtime;
 use iroh_bytes::Hash;
 use iroh_bytes::{BlobFormat, Tag};
 use iroh_net::{key::PublicKey, magic_endpoint::ConnectionInfo, PeerAddr};
@@ -67,11 +68,14 @@ where
     C: ServiceConnection<ProviderService>,
 {
     /// Create a new high-level client to a Iroh node from the low-level RPC client.
-    pub fn new(rpc: RpcClient<ProviderService, C>) -> Self {
+    pub fn new(rpc: RpcClient<ProviderService, C>, rt: runtime::Handle) -> Self {
         Self {
             node: NodeClient { rpc: rpc.clone() },
             blobs: BlobsClient { rpc: rpc.clone() },
-            docs: DocsClient { rpc: rpc.clone() },
+            docs: DocsClient {
+                rpc: rpc.clone(),
+                rt,
+            },
             authors: AuthorsClient { rpc: rpc.clone() },
             tags: TagsClient { rpc },
         }
@@ -129,6 +133,7 @@ where
 #[derive(Debug, Clone)]
 pub struct DocsClient<C> {
     rpc: RpcClient<ProviderService, C>,
+    rt: runtime::Handle,
 }
 
 impl<C> DocsClient<C>
@@ -138,7 +143,7 @@ where
     /// Create a new document.
     pub async fn create(&self) -> Result<Doc<C>> {
         let res = self.rpc.rpc(DocCreateRequest {}).await??;
-        let doc = Doc::new(self.rpc.clone(), res.id);
+        let doc = Doc::new(self.rt.clone(), self.rpc.clone(), res.id);
         Ok(doc)
     }
 
@@ -155,7 +160,7 @@ where
     /// Import a document from a ticket and join all peers in the ticket.
     pub async fn import(&self, ticket: DocTicket) -> Result<Doc<C>> {
         let res = self.rpc.rpc(DocImportRequest(ticket)).await??;
-        let doc = Doc::new(self.rpc.clone(), res.doc_id);
+        let doc = Doc::new(self.rt.clone(), self.rpc.clone(), res.doc_id);
         Ok(doc)
     }
 
@@ -168,7 +173,7 @@ where
     /// Get a [`Doc`] client for a single document. Return None if the document cannot be found.
     pub async fn open(&self, id: NamespaceId) -> Result<Option<Doc<C>>> {
         self.rpc.rpc(DocOpenRequest { doc_id: id }).await??;
-        let doc = Doc::new(self.rpc.clone(), id);
+        let doc = Doc::new(self.rt.clone(), self.rpc.clone(), id);
         Ok(Some(doc))
     }
 }
@@ -523,6 +528,7 @@ struct DocInner<C: ServiceConnection<ProviderService>> {
     id: NamespaceId,
     rpc: RpcClient<ProviderService, C>,
     closed: AtomicBool,
+    rt: runtime::Handle,
 }
 
 impl<C> Drop for DocInner<C>
@@ -532,7 +538,7 @@ where
     fn drop(&mut self) {
         let doc_id = self.id;
         let rpc = self.rpc.clone();
-        tokio::task::spawn(async move {
+        self.rt.main().spawn(async move {
             rpc.rpc(DocCloseRequest { doc_id }).await.ok();
         });
     }
@@ -542,11 +548,12 @@ impl<C> Doc<C>
 where
     C: ServiceConnection<ProviderService>,
 {
-    fn new(rpc: RpcClient<ProviderService, C>, id: NamespaceId) -> Self {
+    fn new(rt: runtime::Handle, rpc: RpcClient<ProviderService, C>, id: NamespaceId) -> Self {
         Self(Arc::new(DocInner {
             rpc,
             id,
             closed: AtomicBool::new(false),
+            rt,
         }))
     }
 
@@ -741,3 +748,35 @@ where
         Err(err) => Err(err.into()),
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use iroh_bytes::util::runtime;
+
+    #[tokio::test]
+    async fn test_drop_doc_client_sync() -> Result<()> {
+        let db = iroh_bytes::store::readonly_mem::Store::default();
+        let doc_store = iroh_sync::store::memory::Store::default();
+        let rt = runtime::Handle::from_current(1)?;
+        let node = crate::node::Node::builder(db, doc_store)
+            .runtime(&rt)
+            .spawn()
+            .await?;
+
+        let client = node.client();
+        let doc = client.docs.create().await?;
+
+        let res = std::thread::spawn(move || {
+            drop(doc);
+            drop(client);
+            drop(node);
+        });
+
+        tokio::task::spawn_blocking(move || res.join().map_err(|e| anyhow::anyhow!("{:?}", e)))
+            .await??;
+
+        Ok(())
+    }
+}

iroh/src/client/quic.rs

@@ -6,6 +6,7 @@ use std::{
     time::Duration,
 };
 
+use iroh_bytes::util::runtime;
 use quic_rpc::transport::quinn::QuinnConnection;
 
 use crate::rpc_protocol::{NodeStatusRequest, ProviderRequest, ProviderResponse, ProviderService};
@@ -26,9 +27,13 @@ pub type Iroh = super::Iroh<QuinnConnection<ProviderResponse, ProviderRequest>>;
 pub type Doc = super::Doc<QuinnConnection<ProviderResponse, ProviderRequest>>;
 
 /// Connect to an iroh node running on the same computer, but in a different process.
-pub async fn connect(rpc_port: u16) -> anyhow::Result<Iroh> {
+pub async fn connect(rpc_port: u16, rt: Option<runtime::Handle>) -> anyhow::Result<Iroh> {
+    let rt = match rt {
+        Some(rt) => rt,
+        None => runtime::Handle::from_current(1)?,
+    };
     let client = connect_raw(rpc_port).await?;
-    Ok(Iroh::new(client))
+    Ok(Iroh::new(client, rt))
 }
 
 /// Create a raw RPC client to an iroh node running on the same computer, but in a different

iroh/src/commands.rs

@@ -84,15 +84,15 @@ pub struct FullArgs {
 }
 
 impl Cli {
-    pub async fn run(self, rt: &runtime::Handle) -> Result<()> {
+    pub async fn run(self, rt: runtime::Handle) -> Result<()> {
         match self.command {
             Commands::Console => {
-                let iroh = iroh::client::quic::connect(self.rpc_args.rpc_port).await?;
+                let iroh = iroh::client::quic::connect(self.rpc_args.rpc_port, Some(rt)).await?;
                 let env = ConsoleEnv::for_console()?;
                 repl::run(&iroh, &env).await
             }
             Commands::Rpc(command) => {
-                let iroh = iroh::client::quic::connect(self.rpc_args.rpc_port).await?;
+                let iroh = iroh::client::quic::connect(self.rpc_args.rpc_port, Some(rt)).await?;
                 let env = ConsoleEnv::for_cli()?;
                 command.run(&iroh, &env).await
             }
@@ -106,9 +106,9 @@ impl Cli {
                 let config = NodeConfig::from_env(cfg.as_deref())?;
 
                 #[cfg(feature = "metrics")]
-                let metrics_fut = start_metrics_server(metrics_addr, rt);
+                let metrics_fut = start_metrics_server(metrics_addr, &rt);
 
-                let res = command.run(rt, &config, keylog).await;
+                let res = command.run(&rt, &config, keylog).await;
 
                 #[cfg(feature = "metrics")]
                 if let Some(metrics_fut) = metrics_fut {

iroh/src/main.rs

@@ -33,5 +33,5 @@ async fn main_impl() -> Result<()> {
         .init();
 
     let cli = Cli::parse();
-    cli.run(&rt).await
+    cli.run(rt).await
 }

iroh/src/node.rs

@@ -743,7 +743,7 @@ impl<D: ReadableStore> Node<D> {
 
     /// Return a client to control this node over an in-memory channel.
     pub fn client(&self) -> crate::client::mem::Iroh {
-        crate::client::Iroh::new(self.controller())
+        crate::client::Iroh::new(self.controller(), self.inner.rt.clone())
     }
 
     /// Return a single token containing everything needed to get a hash.