fix(cluster): fix the possibility of connection leak when the cluster state is broken. (#16842)

* fix(cluster): fix connect leak if server is broken status

* fix(cluster): fix connect leak if server is broken status

Author: zhang2014

src/query/service/src/schedulers/scheduler.rs

@@ -114,13 +114,20 @@ pub async fn build_distributed_pipeline(
 
     let exchange_manager = ctx.get_exchange_manager();
 
-    let mut build_res = exchange_manager
+    match exchange_manager
         .commit_actions(ctx.clone(), fragments_actions)
-        .await?;
-
-    let settings = ctx.get_settings();
-    build_res.set_max_threads(settings.get_max_threads()? as usize);
-    Ok(build_res)
+        .await
+    {
+        Ok(mut build_res) => {
+            let settings = ctx.get_settings();
+            build_res.set_max_threads(settings.get_max_threads()? as usize);
+            Ok(build_res)
+        }
+        Err(error) => {
+            exchange_manager.on_finished_query(&ctx.get_id(), Some(error.clone()));
+            Err(error)
+        }
+    }
 }
 
 pub struct ServiceQueryExecutor {

src/query/service/src/servers/flight/v1/actions/init_query_env.rs

@@ -38,7 +38,7 @@ pub async fn init_query_env(env: QueryEnv) -> Result<()> {
             .init_query_env(&env, ctx)
             .await
         {
-            DataExchangeManager::instance().on_finished_query(&env.query_id);
+            DataExchangeManager::instance().on_finished_query(&env.query_id, Some(e.clone()));
             return Err(e);
         }
 

src/query/service/src/servers/flight/v1/actions/init_query_fragments.rs

@@ -37,7 +37,7 @@ pub async fn init_query_fragments(fragments: QueryFragments) -> Result<()> {
     }));
 
     if let Err(cause) = join_handler.await.flatten() {
-        DataExchangeManager::instance().on_finished_query(&query_id);
+        DataExchangeManager::instance().on_finished_query(&query_id, Some(cause.clone()));
         return Err(cause);
     }
 

src/query/service/src/servers/flight/v1/actions/start_prepared_query.rs

@@ -27,7 +27,7 @@ pub async fn start_prepared_query(id: String) -> Result<()> {
 
     debug!("start prepared query {}", id);
     if let Err(cause) = DataExchangeManager::instance().execute_partial_query(&id) {
-        DataExchangeManager::instance().on_finished_query(&id);
+        DataExchangeManager::instance().on_finished_query(&id, Some(cause.clone()));
         return Err(cause);
     }
     Ok(())

src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs

@@ -242,7 +242,13 @@ impl DataExchangeManager {
                 "Query {} cannot start command while in 180 seconds",
                 query_id
             );
-            self.on_finished_query(&query_id);
+            self.on_finished_query(
+                &query_id,
+                Some(ErrorCode::Internal(format!(
+                    "Query {} cannot start command while in 180 seconds",
+                    query_id
+                ))),
+            );
         }
     }
 
@@ -379,25 +385,25 @@ impl DataExchangeManager {
         }
     }
 
-    pub fn shutdown_query(&self, query_id: &str) {
+    pub fn shutdown_query(&self, query_id: &str, cause: Option<ErrorCode>) {
         let queries_coordinator_guard = self.queries_coordinator.lock();
         let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() };
 
         if let Some(query_coordinator) = queries_coordinator.get_mut(query_id) {
-            query_coordinator.shutdown_query();
+            query_coordinator.shutdown_query(cause);
         }
     }
 
     #[fastrace::trace]
-    pub fn on_finished_query(&self, query_id: &str) {
+    pub fn on_finished_query(&self, query_id: &str, cause: Option<ErrorCode>) {
         let queries_coordinator_guard = self.queries_coordinator.lock();
         let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() };
 
         if let Some(mut query_coordinator) = queries_coordinator.remove(query_id) {
             // Drop mutex guard to avoid deadlock during shutdown,
             drop(queries_coordinator_guard);
 
-            query_coordinator.shutdown_query();
+            query_coordinator.shutdown_query(cause);
             query_coordinator.on_finished();
         }
     }
@@ -476,7 +482,8 @@ impl DataExchangeManager {
                         let mut statistics_receiver = statistics_receiver.lock();
 
                         statistics_receiver.shutdown(info.res.is_err());
-                        ctx.get_exchange_manager().on_finished_query(&query_id);
+                        ctx.get_exchange_manager()
+                            .on_finished_query(&query_id, info.res.clone().err());
                         statistics_receiver.wait_shutdown()
                     },
                 ));
@@ -768,10 +775,10 @@ impl QueryCoordinator {
         Err(ErrorCode::Unimplemented("ExchangeSource is unimplemented"))
     }
 
-    pub fn shutdown_query(&mut self) {
+    pub fn shutdown_query(&mut self, cause: Option<ErrorCode>) {
         if let Some(query_info) = &mut self.info {
             if let Some(query_executor) = &query_info.query_executor {
-                query_executor.finish(None);
+                query_executor.finish(cause);
             }
 
             if let Some(worker) = query_info.remove_leak_query_worker.take() {
@@ -871,10 +878,11 @@ impl QueryCoordinator {
 
         Thread::named_spawn(Some(String::from("Distributed-Executor")), move || {
             let _g = span.set_local_parent();
-            statistics_sender.shutdown(executor.execute().err());
+            let error = executor.execute().err();
+            statistics_sender.shutdown(error.clone());
             query_ctx
                 .get_exchange_manager()
-                .on_finished_query(&query_id);
+                .on_finished_query(&query_id, error);
         });
 
         Ok(())

src/query/service/src/servers/flight/v1/exchange/statistics_sender.rs

@@ -81,8 +81,9 @@ impl StatisticsSender {
                             notified = right;
                             sleep_future = Box::pin(sleep(Duration::from_millis(100)));
 
-                            if let Err(_cause) = Self::send_progress(&ctx, &tx).await {
-                                ctx.get_exchange_manager().shutdown_query(&query_id);
+                            if let Err(cause) = Self::send_progress(&ctx, &tx).await {
+                                ctx.get_exchange_manager()
+                                    .shutdown_query(&query_id, Some(cause));
                                 return;
                             }