fix: Address unit tests hang on multi partition

Signed-off-by: Xuanwo <github@xuanwo.io>

Author: Xuanwo

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

@@ -56,7 +56,7 @@ impl SingletonType {
                 let guard = LOCAL.wait().read();
                 let v: &T = guard
                     .get(&thread_name)
-                    .unwrap_or_else(|| panic!("thread {thread_name} is not initiated"))
+                    .unwrap_or_else(|| panic!("thread {thread_name} is not initiated, don't worry if we are in dropping"))
                     .get();
                 v.clone()
             }
@@ -75,7 +75,7 @@ impl SingletonType {
                 let guard = LOCAL.wait().read();
                 guard
                     .get(&thread_name)
-                    .unwrap_or_else(|| panic!("thread {thread_name} is not initiated"))
+                    .unwrap_or_else(|| panic!("thread {thread_name} is not initiated, don't worry if we are in dropping"))
                     .set(value)
             }
         }

src/query/service/src/sessions/session_mgr.rs

@@ -158,8 +158,10 @@ impl SessionManager {
 
         // stop tracking session
         {
-            let mut sessions = self.active_sessions.write();
-            sessions.remove(session_id);
+            // Make sure this write lock has been released before dropping.
+            // Becuase droping session could re-enter `destroy_session`.
+            let weak_session = { self.active_sessions.write().remove(session_id) };
+            drop(weak_session);
         }
 
         // also need remove mysql_conn_map

src/query/service/tests/it/servers/http/http_query_handlers.rs

@@ -1383,7 +1383,7 @@ async fn test_func_object_keys() -> Result<()> {
     Ok(())
 }
 
-#[tokio::test(flavor = "current_thread")]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn test_multi_partition() -> Result<()> {
     let _guard = TestGlobalServices::setup(ConfigBuilder::create().build()).await?;
 

src/query/service/tests/it/tests/sessions.rs

@@ -17,9 +17,7 @@ use common_config::Config;
 use common_exception::Result;
 use common_tracing::set_panic_hook;
 use databend_query::clusters::ClusterDiscovery;
-use databend_query::sessions::SessionManager;
 use databend_query::GlobalServices;
-use tracing::debug;
 use tracing::info;
 
 pub struct TestGlobalServices;
@@ -64,19 +62,6 @@ pub struct TestGuard {
 
 impl Drop for TestGuard {
     fn drop(&mut self) {
-        // Check if session manager sill have active sessions.
-        {
-            let session_mgr = SessionManager::instance();
-            // Destory all sessions.
-            for process in session_mgr.processes_info() {
-                session_mgr.destroy_session(&process.id);
-            }
-            // Double check again.
-            for process in session_mgr.processes_info() {
-                debug!("process {process:?} still running after drop, something must be wrong");
-            }
-        }
-
         GlobalInstance::drop_testing(&self.thread_name);
     }
 }