atexit-safe close builder draft

Cargo.toml

@@ -120,6 +120,7 @@ libloading = "0.8"
 tracing = "0.1"
 lockfree = "0.5"
 lz4_flex = "0.11"
+nolocal-block-on = "1.0.1"
 nix = { version = "0.29.0", features = ["fs"] }
 num_cpus = "1.16.0"
 num-traits = { version = "0.2.19", default-features = false }

zenoh/Cargo.toml

@@ -79,6 +79,7 @@ git-version = { workspace = true }
 itertools = { workspace = true }
 json5 = { workspace = true }
 lazy_static = { workspace = true }
+nolocal-block-on = { workspace = true }
 tracing = { workspace = true }
 paste = { workspace = true }
 petgraph = { workspace = true }

zenoh/src/api/session.rs

@@ -17,22 +17,23 @@ use std::{
     collections::HashMap,
     convert::TryInto,
     fmt,
+    future::{Future, IntoFuture},
     ops::Deref,
+    pin::Pin,
     sync::{
         atomic::{AtomicU16, Ordering},
         Arc, Mutex, RwLock,
     },
     time::{Duration, SystemTime, UNIX_EPOCH},
 };
-
 use tracing::{error, info, trace, warn};
 use uhlc::Timestamp;
 #[cfg(feature = "internal")]
 use uhlc::HLC;
 use zenoh_buffers::ZBuf;
 use zenoh_collections::SingleOrVec;
 use zenoh_config::{unwrap_or_default, wrappers::ZenohId};
-use zenoh_core::{zconfigurable, zread, Resolve, ResolveClosure, ResolveFuture, Wait};
+use zenoh_core::{zconfigurable, zread, Resolvable, Resolve, ResolveClosure, ResolveFuture, Wait};
 #[cfg(feature = "unstable")]
 use zenoh_protocol::network::{
     declare::{DeclareToken, SubscriberId, TokenId, UndeclareToken},
@@ -615,8 +616,8 @@ impl Session {
     /// subscriber_task.await.unwrap();
     /// # }
     /// ```
-    pub fn close(&self) -> impl Resolve<ZResult<()>> + '_ {
-        self.0.close()
+    pub fn close(&self) -> SessionCloseBuilder {
+        SessionCloseBuilder::new(self.0.clone())
     }
 
     /// Check if the session has been closed.
@@ -1073,38 +1074,39 @@ impl Session {
         })
     }
 }
+
 impl SessionInner {
     pub fn zid(&self) -> ZenohId {
         self.runtime.zid()
     }
 
-    fn close(&self) -> impl Resolve<ZResult<()>> + '_ {
-        ResolveFuture::new(async move {
-            let Some(primitives) = zwrite!(self.state).primitives.take() else {
-                return Ok(());
-            };
-            if self.owns_runtime {
-                info!(zid = %self.zid(), "close session");
-            }
-            self.task_controller.terminate_all(Duration::from_secs(10));
-            if self.owns_runtime {
-                self.runtime.close().await?;
-            } else {
-                primitives.send_close();
-            }
-            let mut state = zwrite!(self.state);
-            state.queryables.clear();
-            state.subscribers.clear();
-            state.liveliness_subscribers.clear();
-            state.local_resources.clear();
-            state.remote_resources.clear();
-            #[cfg(feature = "unstable")]
-            {
-                state.tokens.clear();
-                state.matching_listeners.clear();
-            }
-            Ok(())
-        })
+    async fn close(&self, _timeout: Duration) -> ZResult<()> {
+        let Some(primitives) = zwrite!(self.state).primitives.take() else {
+            return Ok(());
+        };
+        if self.owns_runtime {
+            info!(zid = %self.zid(), "close session");
+        }
+        self.task_controller
+            .terminate_all_async(Duration::from_secs(10))
+            .await;
+        if self.owns_runtime {
+            self.runtime.close().await?
+        } else {
+            primitives.send_close();
+        }
+        let mut state = zwrite!(self.state);
+        state.queryables.clear();
+        state.subscribers.clear();
+        state.liveliness_subscribers.clear();
+        state.local_resources.clear();
+        state.remote_resources.clear();
+        #[cfg(feature = "unstable")]
+        {
+            state.tokens.clear();
+            state.matching_listeners.clear();
+        }
+        Ok(())
     }
 
     pub(crate) fn declare_prefix<'a>(
@@ -2880,3 +2882,56 @@ where
 {
     OpenBuilder::new(config)
 }
+
+/// A builder for closing a [`crate::Session`].
+///
+/// # Examples
+/// ```
+/// # #[tokio::main]
+/// # async fn main() {
+///
+/// let session = zenoh::open(zenoh::Config::default()).await.unwrap();
+/// session.close()
+///     .timeout(std::time::Duration::from_secs(10))
+///     .await
+///     .unwrap();
+/// # }
+/// ```
+
+pub struct SessionCloseBuilder {
+    inner: Arc<SessionInner>,
+    timeout: Duration,
+}
+
+impl SessionCloseBuilder {
+    fn new(inner: Arc<SessionInner>) -> Self {
+        Self {
+            inner,
+            timeout: Duration::MAX,
+        }
+    }
+
+    pub fn timeout(mut self, timeout: Duration) -> Self {
+        self.timeout = timeout;
+        self
+    }
+}
+
+impl Resolvable for SessionCloseBuilder {
+    type To = ZResult<()>;
+}
+
+impl Wait for SessionCloseBuilder {
+    fn wait(self) -> Self::To {
+        nolocal_block_on::block_on(self.into_future())
+    }
+}
+
+impl IntoFuture for SessionCloseBuilder {
+    type Output = <Self as Resolvable>::To;
+    type IntoFuture = Pin<Box<dyn Future<Output = <Self as IntoFuture>::Output> + Send>>;
+
+    fn into_future(self) -> Self::IntoFuture {
+        Box::pin(async move { self.inner.close(self.timeout).await }.into_future())
+    }
+}