Reuse the current runtime if possible.

Signed-off-by: ChenYing Kuo <[email protected]>
eclipse-zenoh · Aug 5, 2024 · c726584 · c726584
1 parent 2333f20
commit c726584
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 13 deletions.
diff --git a/src/files_mgt.rs b/src/files_mgt.rs
@@ -36,7 +36,7 @@ use zenoh::{
     Result as ZResult,
 };
 
-use crate::{data_info_mgt::*, TOKIO_RUNTIME};
+use crate::{blockon_runtime, data_info_mgt::*};
 
 pub const CONFLICT_SUFFIX: &str = ".##z";
 
@@ -377,16 +377,14 @@ impl Drop for FilesMgr {
         match self.on_closure {
             OnClosure::DeleteAll => {
                 // Close data_info_mgr at first
-                tokio::task::block_in_place(|| {
-                    TOKIO_RUNTIME.block_on(async move {
-                        self.data_info_mgr
-                            .close()
-                            .await
-                            .unwrap_or_else(|e| warn!("{}", e));
-                        remove_dir_all(&self.base_dir).unwrap_or_else(|err| {
-                            warn!("Failed to cleanup directory {:?}; {}", self.base_dir, err)
-                        });
-                    })
+                blockon_runtime(async move {
+                    self.data_info_mgr
+                        .close()
+                        .await
+                        .unwrap_or_else(|e| warn!("{}", e));
+                    remove_dir_all(&self.base_dir).unwrap_or_else(|err| {
+                        warn!("Failed to cleanup directory {:?}; {}", self.base_dir, err)
+                    });
                 });
             }
             OnClosure::DoNothing => {

diff --git a/src/lib.rs b/src/lib.rs
@@ -12,7 +12,7 @@
 //   ZettaScale Zenoh Team, <[email protected]>
 //
 
-use std::{collections::HashMap, fs::DirBuilder, io::prelude::*, path::PathBuf};
+use std::{collections::HashMap, fs::DirBuilder, future::Future, io::prelude::*, path::PathBuf};
 
 use async_trait::async_trait;
 use tempfile::tempfile_in;
@@ -38,13 +38,28 @@ use files_mgt::*;
 const WORKER_THREAD_NUM: usize = 2;
 const MAX_BLOCK_THREAD_NUM: usize = 50;
 lazy_static::lazy_static! {
-    pub static ref TOKIO_RUNTIME: tokio::runtime::Runtime = tokio::runtime::Builder::new_multi_thread()
+    // The global runtime is used in the dynamic plugins, which we can't get the current runtime
+    static ref TOKIO_RUNTIME: tokio::runtime::Runtime = tokio::runtime::Builder::new_multi_thread()
                .worker_threads(WORKER_THREAD_NUM)
                .max_blocking_threads(MAX_BLOCK_THREAD_NUM)
                .enable_all()
                .build()
                .expect("Unable to create runtime");
 }
+#[inline(always)]
+fn blockon_runtime<F: Future>(task: F) -> F::Output {
+    // Check whether able to get the current runtime
+    match tokio::runtime::Handle::try_current() {
+        Ok(rt) => {
+            // Able to get the current runtime (standalone binary), spawn on the current runtime
+            tokio::task::block_in_place(|| rt.block_on(task))
+        }
+        Err(_) => {
+            // Unable to get the current runtime (dynamic plugins), spawn on the global runtime
+            tokio::task::block_in_place(|| TOKIO_RUNTIME.block_on(task))
+        }
+    }
+}
 
 /// The environement variable used to configure the root of all storages managed by this FileSystemBackend.
 pub const SCOPE_ENV_VAR: &str = "ZENOH_BACKEND_FS_ROOT";