Merge pull request #8944 from yufan022/main

fix(mysql): fix mysql conns leak
databendlabs · Nov 24, 2022 · 5f78113 · 5f78113 · vercel · Nov 24, 2022
2 parents bb2c9a9 + 265f21e
commit 5f78113
Showing 1 changed file with 15 additions and 6 deletions.
diff --git a/src/query/service/src/sessions/session_mgr.rs b/src/query/service/src/sessions/session_mgr.rs
@@ -95,7 +95,7 @@ impl SessionManager {
         let config = self.get_conf();
         {
             let sessions = self.active_sessions.read();
-            if sessions.len() == self.max_sessions {
+            if sessions.len() >= self.max_sessions {
                 return Err(ErrorCode::TooManyUserConnections(
                     "The current accept connection has exceeded max_active_sessions config",
                 ));
@@ -106,11 +106,9 @@ impl SessionManager {
         let mut mysql_conn_id = None;
         match session_typ {
             SessionType::MySQL => {
-                let mut conn_id_session_id = self.mysql_conn_map.write();
+                let conn_id_session_id = self.mysql_conn_map.read();
                 mysql_conn_id = Some(self.mysql_basic_conn_id.fetch_add(1, Ordering::Relaxed));
-                if conn_id_session_id.len() < self.max_sessions {
-                    conn_id_session_id.insert(mysql_conn_id, id.clone());
-                } else {
+                if conn_id_session_id.len() >= self.max_sessions {
                     return Err(ErrorCode::TooManyUserConnections(
                         "The current accept connection has exceeded max_active_sessions config",
                     ));
@@ -128,7 +126,7 @@ impl SessionManager {
         let user_api = UserApiProvider::instance();
         let session_settings = Settings::try_create(&config, user_api, tenant).await?;
         let session_ctx = SessionContext::try_create(config.clone(), session_settings)?;
-        let session = Session::try_create(id, typ.clone(), session_ctx, mysql_conn_id)?;
+        let session = Session::try_create(id.clone(), typ.clone(), session_ctx, mysql_conn_id)?;
 
         let mut sessions = self.active_sessions.write();
         if sessions.len() < self.max_sessions {
@@ -142,6 +140,17 @@ impl SessionManager {
                 sessions.insert(session.get_id(), Arc::downgrade(&session));
             }
 
+            if let SessionType::MySQL = session_typ {
+                let mut conn_id_session_id = self.mysql_conn_map.write();
+                if conn_id_session_id.len() < self.max_sessions {
+                    conn_id_session_id.insert(mysql_conn_id, id);
+                } else {
+                    return Err(ErrorCode::TooManyUserConnections(
+                        "The current accept connection has exceeded max_active_sessions config",
+                    ));
+                }
+            }
+
             Ok(session)
         } else {
             Err(ErrorCode::TooManyUserConnections(