Merge pull request #9242 from Xuanwo/fix-testing-singletion-drop

fix: Address unit test hang on test_insert

Author: Xuanwo

src/common/base/src/base/singleton_instance.rs

@@ -17,37 +17,45 @@ use std::fmt::Debug;
 use once_cell::sync::OnceCell;
 use state::Container;
 
+/// SINGLETON_TYPE is the global singleton type.
+static SINGLETON_TYPE: OnceCell<SingletonType> = OnceCell::new();
+
+/// GLOBAL is a static type that holding all global data.
+static GLOBAL: OnceCell<Container![Send + Sync]> = OnceCell::new();
+
+#[cfg(debug_assertions)]
+/// LOCAL is a static type that holding all global data only for local tests.
+static LOCAL: OnceCell<
+    parking_lot::RwLock<std::collections::HashMap<String, Container![Send + Sync]>>,
+> = OnceCell::new();
+
 /// Singleton is a wrapper enum for `Container![Send + Sync]`.
 ///
 /// - `Production` is used in our production code.
 /// - `Testing` is served for test only and gated under `debug_assertions`.
-pub enum Singleton {
-    Production(Container![Send + Sync]),
+pub enum SingletonType {
+    Production,
 
     #[cfg(debug_assertions)]
-    Testing(parking_lot::Mutex<std::collections::HashMap<String, Container![Send + Sync]>>),
+    Testing,
 }
 
-unsafe impl Send for Singleton {}
-unsafe impl Sync for Singleton {}
-
-impl Singleton {
+impl SingletonType {
     fn get<T: Clone + 'static>(&self) -> T {
         match self {
-            Singleton::Production(c) => {
-                let v: &T = c.get();
+            SingletonType::Production => {
+                let v: &T = GLOBAL.wait().get();
                 v.clone()
             }
             #[cfg(debug_assertions)]
-            Singleton::Testing(c) => {
-                let thread = std::thread::current();
-                let thread_name = match thread.name() {
-                    Some(name) => name,
-                    None => panic!("thread doesn't have name"),
-                };
-                let guard = c.lock();
+            SingletonType::Testing => {
+                let thread_name = std::thread::current()
+                    .name()
+                    .expect("thread doesn't have name")
+                    .to_string();
+                let guard = LOCAL.wait().read();
                 let v: &T = guard
-                    .get(thread_name)
+                    .get(&thread_name)
                     .unwrap_or_else(|| panic!("thread {thread_name} is not initiated"))
                     .get();
                 v.clone()
@@ -57,37 +65,35 @@ impl Singleton {
 
     fn set<T: Send + Sync + 'static>(&self, value: T) -> bool {
         match self {
-            Singleton::Production(c) => c.set(value),
+            SingletonType::Production => GLOBAL.wait().set(value),
             #[cfg(debug_assertions)]
-            Singleton::Testing(c) => {
-                let thread = std::thread::current();
-                let thread_name = match thread.name() {
-                    Some(name) => name,
-                    None => panic!("thread doesn't have name"),
-                };
-                let mut guard = c.lock();
-                let c = guard.entry(thread_name.to_string()).or_default();
-                c.set(value)
+            SingletonType::Testing => {
+                let thread_name = std::thread::current()
+                    .name()
+                    .expect("thread doesn't have name")
+                    .to_string();
+                let guard = LOCAL.wait().read();
+                guard
+                    .get(&thread_name)
+                    .unwrap_or_else(|| panic!("thread {thread_name} is not initiated"))
+                    .set(value)
             }
         }
     }
 }
 
-impl Debug for Singleton {
+impl Debug for SingletonType {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("Singleton")
             .field("type", &match self {
-                Self::Production(_) => "Production",
+                Self::Production => "Production",
                 #[cfg(debug_assertions)]
-                Self::Testing(_) => "Testing",
+                Self::Testing => "Testing",
             })
             .finish()
     }
 }
 
-/// GLOBAL is a static type that holding all global data.
-static GLOBAL: OnceCell<Singleton> = OnceCell::new();
-
 /// Global is an empty struct that only used to carry associated functions.
 pub struct GlobalInstance;
 
@@ -96,44 +102,48 @@ impl GlobalInstance {
     ///
     /// Should only be initiated once.
     pub fn init_production() {
-        let _ = GLOBAL.set(Singleton::Production(<Container![Send + Sync]>::new()));
+        let _ = SINGLETON_TYPE.set(SingletonType::Production);
+        let _ = GLOBAL.set(<Container![Send + Sync]>::new());
     }
 
     /// init testing global data registry.
     ///
     /// Should only be initiated once and only used in testing.
     #[cfg(debug_assertions)]
-    pub fn init_testing() {
-        let _ = GLOBAL.set(Singleton::Testing(parking_lot::Mutex::default()));
+    pub fn init_testing(thread_name: &str) {
+        let _ = SINGLETON_TYPE.set(SingletonType::Testing);
+        let _ = LOCAL.set(parking_lot::RwLock::default());
+
+        let v = LOCAL
+            .wait()
+            .write()
+            .insert(thread_name.to_string(), <Container![Send + Sync]>::new());
+        assert!(
+            v.is_none(),
+            "thread {thread_name} has been initiated before"
+        )
     }
 
     /// drop testing global data by thread name.
     ///
     /// Should only be used in testing code.
     #[cfg(debug_assertions)]
     pub fn drop_testing(thread_name: &str) {
-        match GLOBAL.wait() {
-            Singleton::Production(_) => {
-                unreachable!("drop_testing should never be called on production global")
-            }
-            Singleton::Testing(c) => {
-                let mut guard = c.lock();
-                let _ = guard.remove(thread_name);
-            }
-        }
+        // Make sure the write lock is released before alling drop.
+        let _ = { LOCAL.wait().write().remove(thread_name) };
     }
 
     /// Get data from global data registry.
     pub fn get<T: Clone + 'static>() -> T {
-        GLOBAL
+        SINGLETON_TYPE
             .get()
             .expect("global data registry must be initiated")
             .get()
     }
 
     /// Set data into global data registry.
     pub fn set<T: Send + Sync + 'static>(value: T) {
-        let set = GLOBAL
+        let set = SINGLETON_TYPE
             .get()
             .expect("global data registry must be initiated")
             .set(value);

src/query/service/src/global_services.rs

@@ -36,23 +36,11 @@ pub struct GlobalServices;
 
 impl GlobalServices {
     pub async fn init(config: Config) -> Result<()> {
-        GlobalServices::init_with(config, true).await
+        GlobalInstance::init_production();
+        GlobalServices::init_with(config).await
     }
 
-    pub async fn init_with(config: Config, prodcution: bool) -> Result<()> {
-        #[cfg(debug_assertions)]
-        if prodcution {
-            GlobalInstance::init_production();
-        } else {
-            GlobalInstance::init_testing();
-        }
-        #[cfg(not(debug_assertions))]
-        if prodcution {
-            GlobalInstance::init_production();
-        } else {
-            unreachable!("release build must init with production global")
-        }
-
+    pub async fn init_with(config: Config) -> Result<()> {
         // The order of initialization is very important
         GlobalConfig::init(config.clone())?;