Nolocal close (#1632)

* Make background close builder nolocal

* make close builders internally-public

* Update lib.rs

* do not use flume channel s it is not nolocal

* fix exports

* switch runtime to Net

* Make close more atexit-safe

* add tests for wait() in close

* Add important note

* Update Cargo.toml and atexit.rs

* Code format

* fix clippy

* Update Cargo.lock

* - review fixes
- get rid of async-channel (use tokio channel instead until it is nolocal)
- improve testing coverage (test background close to be nolocal)

* code format

* fix feature set

* disable session_close_in_atexit test on 1.85 as std has regression rust-lang/rust#138696

* fix tokio executor blockage

* Paic with message describing nolocal thread spawning error

* spelling

* Handle nolocal errors more properly

* refactor: replace `mpsc::channel` with `oneshot::channel` and address the `block_in_place` issue (#26)

---------

Co-authored-by: Yuyuan Yuan <az6980522@gmail.com>

Author: yellowhatter

Cargo.lock

@@ -124,6 +124,7 @@ once_cell = { workspace = true }
 
 [dev-dependencies]
 tokio = { workspace = true }
+libc = { workspace = true }
 
 [build-dependencies]
 rustc_version = { workspace = true }

zenoh/Cargo.toml

@@ -17,4 +17,12 @@ fn main() {
         "cargo:rustc-env=RUSTC_VERSION={}",
         version_meta.short_version_string
     );
+
+    let version = rustc_version::version().unwrap();
+    if version >= rustc_version::Version::parse("1.85.0").unwrap()
+        && version <= rustc_version::Version::parse("1.85.1").unwrap()
+    {
+        // https://github.com/rust-lang/rust/issues/138696
+        println!("cargo:rustc-cfg=nolocal_thread_not_available");
+    }
 }

zenoh/build.rs

@@ -19,6 +19,8 @@ use std::{
 };
 
 use async_trait::async_trait;
+#[cfg(all(feature = "unstable", feature = "internal"))]
+use tokio::sync::oneshot::Receiver;
 use zenoh_core::{Resolvable, Wait};
 use zenoh_result::ZResult;
 use zenoh_runtime::ZRuntime;
@@ -73,7 +75,35 @@ impl<TCloseable: Closeable> Resolvable for CloseBuilder<TCloseable> {
 
 impl<TCloseable: Closeable> Wait for CloseBuilder<TCloseable> {
     fn wait(self) -> Self::To {
-        ZRuntime::Application.block_in_place(self.into_future())
+        let future = self.into_future();
+        match tokio::runtime::Handle::try_current() {
+            Ok(_) => {
+                tracing::trace!("tokio TLS available, closing closeable directly");
+                ZRuntime::Net.block_in_place(future)
+            }
+            Err(e) if e.is_missing_context() => {
+                tracing::trace!("tokio TLS is just missing, closing closeable directly");
+                ZRuntime::Net.block_in_place(future)
+            }
+            Err(_) => {
+                #[cfg(nolocal_thread_not_available)]
+                panic!("Close when thread-local storage is unavailable (typically in atexit()) does not work for this Rust 1.85..1.85.1, see https://github.com/rust-lang/rust/issues/138696");
+
+                #[cfg(not(nolocal_thread_not_available))]
+                {
+                    let evaluate = move || {
+                        // NOTE: tracing logger also panics if used inside atexit() handler!!!
+                        tracing::trace!(
+                            "tokio TLS NOT available, closing closeable in separate thread"
+                        );
+                        ZRuntime::Net.block_in_place(future)
+                    };
+                    std::thread::spawn(evaluate)
+                        .join()
+                        .expect("Error spawning atexit-safe thread")
+                }
+            }
+        }
     }
 }
 
@@ -99,19 +129,13 @@ impl<TCloseable: Closeable> IntoFuture for CloseBuilder<TCloseable> {
 
 #[cfg(all(feature = "unstable", feature = "internal"))]
 /// A builder for close operations running in background
-// NOTE: `Closeable` is only pub(crate) because it is zenoh-internal trait, so we don't
-// care about the `private_bounds` lint in this particular case.
 #[doc(hidden)]
-#[allow(private_bounds)]
 pub struct BackgroundCloseBuilder<TOutput: Send + 'static> {
     inner: Pin<Box<dyn Future<Output = TOutput> + Send>>,
 }
 
 #[cfg(all(feature = "unstable", feature = "internal"))]
 #[doc(hidden)]
-// NOTE: `Closeable` is only pub(crate) because it is zenoh-internal trait, so we don't
-// care about the `private_bounds` lint in this particular case.
-#[allow(private_bounds)]
 impl<TOutput: Send + 'static> BackgroundCloseBuilder<TOutput> {
     fn new(inner: Pin<Box<dyn Future<Output = TOutput> + Send>>) -> Self {
         Self { inner }
@@ -120,13 +144,13 @@ impl<TOutput: Send + 'static> BackgroundCloseBuilder<TOutput> {
 
 #[cfg(all(feature = "unstable", feature = "internal"))]
 impl<TOutput: Send + 'static> Resolvable for BackgroundCloseBuilder<TOutput> {
-    type To = tokio::task::JoinHandle<TOutput>;
+    type To = NolocalJoinHandle<TOutput>;
 }
 
 #[cfg(all(feature = "unstable", feature = "internal"))]
 impl<TOutput: Send + 'static> Wait for BackgroundCloseBuilder<TOutput> {
     fn wait(self) -> Self::To {
-        ZRuntime::Application.block_in_place(self.into_future())
+        ZRuntime::Net.block_in_place(self.into_future())
     }
 }
 
@@ -135,10 +159,61 @@ impl<TOutput: Send + 'static> IntoFuture for BackgroundCloseBuilder<TOutput> {
     type Output = <Self as Resolvable>::To;
     type IntoFuture = Pin<Box<dyn Future<Output = <Self as IntoFuture>::Output> + Send>>;
 
-    // NOTE: yes, we need to return a future that returns JoinHandle
-    #[allow(clippy::async_yields_async)]
     fn into_future(self) -> Self::IntoFuture {
-        Box::pin(async move { ZRuntime::Application.spawn(self.inner) }.into_future())
+        Box::pin(
+            async move {
+                let (tx, rx) = tokio::sync::oneshot::channel();
+                ZRuntime::Net.spawn(async move {
+                    if tx.send(self.inner.await).is_err() {
+                        panic!("BackgroundCloseBuilder: critical error sending the result");
+                    }
+                });
+                NolocalJoinHandle::new(rx)
+            }
+            .into_future(),
+        )
+    }
+}
+
+#[cfg(all(feature = "unstable", feature = "internal"))]
+#[doc(hidden)]
+pub struct NolocalJoinHandle<TOutput: Send + 'static> {
+    rx: Receiver<TOutput>,
+}
+
+#[cfg(all(feature = "unstable", feature = "internal"))]
+impl<TOutput: Send + 'static> NolocalJoinHandle<TOutput> {
+    fn new(rx: Receiver<TOutput>) -> Self {
+        Self { rx }
+    }
+}
+
+#[cfg(all(feature = "unstable", feature = "internal"))]
+impl<TOutput: Send + 'static> Resolvable for NolocalJoinHandle<TOutput> {
+    type To = TOutput;
+}
+
+#[cfg(all(feature = "unstable", feature = "internal"))]
+impl<TOutput: Send + 'static> Wait for NolocalJoinHandle<TOutput> {
+    fn wait(self) -> Self::To {
+        ZRuntime::Net.block_in_place(self.into_future())
+    }
+}
+
+#[cfg(all(feature = "unstable", feature = "internal"))]
+impl<TOutput: Send + 'static> IntoFuture for NolocalJoinHandle<TOutput> {
+    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.rx
+                    .await
+                    .expect("NolocalJoinHandle: critical error receiving the result")
+            }
+            .into_future(),
+        )
     }
 }
 

zenoh/src/api/builders/close.rs

@@ -477,6 +477,12 @@ compile_error!(
 
 #[zenoh_macros::internal]
 pub mod internal {
+    #[zenoh_macros::unstable]
+    pub mod builders {
+        pub mod close {
+            pub use crate::api::builders::close::{BackgroundCloseBuilder, NolocalJoinHandle};
+        }
+    }
     pub mod traits {
         pub use crate::api::builders::sample::{
             EncodingBuilderTrait, QoSBuilderTrait, SampleBuilderTrait, TimestampBuilderTrait,

zenoh/src/lib.rs

@@ -0,0 +1,165 @@
+//
+// Copyright (c) 2025 ZettaScale Technology
+//
+// This program and the accompanying materials are made available under the
+// terms of the Eclipse Public License 2.0 which is available at
+// http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0
+// which is available at https://www.apache.org/licenses/LICENSE-2.0.
+//
+// SPDX-License-Identifier: EPL-2.0 OR Apache-2.0
+//
+// Contributors:
+//   ZettaScale Zenoh Team, <zenoh@zettascale.tech>
+//
+
+#![cfg(feature = "internal_config")]
+
+fn run_in_separate_process(main_name: &str, must_panic: bool) {
+    let output = std::process::Command::new(std::env::current_exe().unwrap())
+        .arg(main_name)
+        .arg("--nocapture")
+        .arg("--show-output")
+        .arg("--exact")
+        .arg("--include-ignored")
+        .env(
+            {
+                match must_panic {
+                    true => "PROBE_PANIC_IN_TEST",
+                    false => "PROBE_NO_PANIC_IN_TEST",
+                }
+            },
+            "true",
+        )
+        .output()
+        .expect("Failed to run test in separate process");
+
+    match must_panic {
+        true => {
+            assert!(
+                !output.status.success(),
+                "Inner test should have been failed"
+            );
+        }
+        false => {
+            assert!(output.status.success(), "Inner test failed");
+        }
+    };
+}
+
+#[ignore]
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn panic_in_separate_process() {
+    if std::env::var("PROBE_PANIC_IN_TEST").is_ok() {
+        panic!("Panic in test because PROBE_PANIC_IN_TEST is set!");
+    }
+}
+
+#[ignore]
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn atexit_panic_in_separate_process() {
+    unsafe { libc::atexit(panic_in_atexit) };
+}
+
+extern "C" fn panic_in_atexit() {
+    if std::env::var("PROBE_PANIC_IN_TEST").is_ok() {
+        panic!("Panic in atexit because PROBE_PANIC_IN_TEST is set!");
+    }
+}
+
+#[test]
+fn panic_is_seen_in_separate_process() {
+    run_in_separate_process("panic_in_separate_process", true);
+}
+
+#[test]
+fn panic_is_seen_in_separate_process_atexit() {
+    run_in_separate_process("atexit_panic_in_separate_process", true);
+}
+
+use std::sync::OnceLock;
+
+use zenoh::{Session, Wait};
+static SESSION: OnceLock<Session> = OnceLock::new();
+
+#[ignore]
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn session_close_in_atexit_main() {
+    if std::env::var("PROBE_PANIC_IN_TEST").is_ok() {
+        panic!("Panic in test because PROBE_PANIC_IN_TEST is set!");
+    }
+
+    // Open the sessions
+    let mut config = zenoh::Config::default();
+    config
+        .listen
+        .endpoints
+        .set(vec!["tcp/127.0.0.1:19446".parse().unwrap()])
+        .unwrap();
+    config.scouting.multicast.set_enabled(Some(false)).unwrap();
+
+    let s = zenoh::open(config).await.unwrap();
+    let _session = SESSION.get_or_init(move || s);
+
+    unsafe { libc::atexit(close_session_in_atexit) };
+}
+
+extern "C" fn close_session_in_atexit() {
+    let session = SESSION.get().unwrap();
+    session.close().wait().unwrap();
+}
+
+#[cfg(not(nolocal_thread_not_available))]
+#[test]
+fn session_close_in_atexit() {
+    run_in_separate_process("session_close_in_atexit_main", true);
+    run_in_separate_process("session_close_in_atexit_main", false);
+}
+
+#[cfg(all(feature = "unstable", feature = "internal"))]
+use std::sync::Mutex;
+
+#[cfg(all(feature = "unstable", feature = "internal"))]
+use zenoh::{internal::builders::close::NolocalJoinHandle, Result};
+
+#[cfg(all(feature = "unstable", feature = "internal"))]
+static BACKGROUND_SESSION: OnceLock<Session> = OnceLock::new();
+#[cfg(all(feature = "unstable", feature = "internal"))]
+static CLOSE: OnceLock<Mutex<Option<NolocalJoinHandle<Result<()>>>>> = OnceLock::new();
+
+#[cfg(all(feature = "unstable", feature = "internal"))]
+#[ignore]
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn background_session_close_in_atexit_main() {
+    if std::env::var("PROBE_PANIC_IN_TEST").is_ok() {
+        panic!("Panic in test because PROBE_PANIC_IN_TEST is set!");
+    }
+
+    // Open the sessions
+    let mut config = zenoh::Config::default();
+    config
+        .listen
+        .endpoints
+        .set(vec!["tcp/127.0.0.1:19445".parse().unwrap()])
+        .unwrap();
+    config.scouting.multicast.set_enabled(Some(false)).unwrap();
+
+    let s = zenoh::open(config).await.unwrap();
+    let session = BACKGROUND_SESSION.get_or_init(move || s);
+
+    let _close = CLOSE.get_or_init(|| Mutex::new(Some(session.close().in_background().wait())));
+
+    unsafe { libc::atexit(background_close_session_in_atexit) };
+}
+
+#[cfg(all(feature = "unstable", feature = "internal"))]
+extern "C" fn background_close_session_in_atexit() {
+    let mut close = CLOSE.get().unwrap().lock().unwrap();
+    close.take().unwrap().wait().unwrap();
+}
+
+#[cfg(all(feature = "unstable", feature = "internal"))]
+#[test]
+fn background_session_close_in_atexit() {
+    run_in_separate_process("background_session_close_in_atexit_main", true);
+    run_in_separate_process("background_session_close_in_atexit_main", false);
+}