fix(cluster): do not to kill the query when the query recv finished.

Author: zhang2014

src/query/service/src/api/rpc/exchange/statistics_receiver.rs

@@ -76,7 +76,7 @@ impl StatisticsReceiver {
                                         recv = Box::pin(flight_exchange.recv());
                                     }
                                     Err(cause) => {
-                                        ctx.get_current_session().force_kill_query();
+                                        ctx.get_current_session().force_kill_query(cause.clone());
                                         return Err(cause);
                                     }
                                 };
@@ -90,14 +90,13 @@ impl StatisticsReceiver {
                             notified = middle;
                             match Self::recv_data(&ctx, res) {
                                 Ok(true) => {
-                                    ctx.get_current_session().force_kill_query();
                                     return Ok(());
                                 }
                                 Ok(false) => {
                                     recv = Box::pin(flight_exchange.recv());
                                 }
                                 Err(cause) => {
-                                    ctx.get_current_session().force_kill_query();
+                                    ctx.get_current_session().force_kill_query(cause.clone());
                                     return Err(cause);
                                 }
                             };
@@ -106,7 +105,7 @@ impl StatisticsReceiver {
                 }
 
                 if let Err(cause) = Self::fetch(&ctx, &flight_exchange, recv).await {
-                    ctx.get_current_session().force_kill_query();
+                    ctx.get_current_session().force_kill_query(cause.clone());
                     return Err(cause);
                 }
 

src/query/service/src/interpreters/interpreter_kill.rs

@@ -48,7 +48,9 @@ impl KillInterpreter {
                 Ok(Box::pin(DataBlockStream::create(schema, None, vec![])))
             }
             Some(kill_session) => {
-                kill_session.force_kill_query();
+                kill_session.force_kill_query(ErrorCode::AbortedQuery(
+                    "Aborted query, because the server is shutting down or the query was killed",
+                ));
                 let schema = Arc::new(DataSchema::empty());
                 Ok(Box::pin(DataBlockStream::create(schema, None, vec![])))
             }

src/query/service/src/servers/http/v1/query/execute_state.rs

@@ -197,7 +197,9 @@ impl Executor {
             Running(r) => {
                 // release session
                 if kill {
-                    r.session.force_kill_query();
+                    r.session.force_kill_query(ErrorCode::AbortedQuery(
+                        "Aborted query, because the server is shutting down or the query was killed",
+                    ));
                 }
                 if let Err(e) = &reason {
                     r.ctx.set_error(e.clone());

src/query/service/src/sessions/query_ctx_shared.rs

@@ -124,11 +124,8 @@ impl QueryContextShared {
         self.query_need_abort.clone()
     }
 
-    pub fn kill(&self) {
-        self.set_error(ErrorCode::AbortedQuery(
-            "Aborted query, because the server is shutting down or the query was killed",
-        ));
-
+    pub fn kill(&self, cause: ErrorCode) {
+        self.set_error(cause);
         self.query_need_abort.store(true, Ordering::Release);
         let mut sources_abort_handle = self.sources_abort_handle.write();
 

src/query/service/src/sessions/session.rs

@@ -106,15 +106,17 @@ impl Session {
     }
 
     pub fn force_kill_session(self: &Arc<Self>) {
-        self.force_kill_query();
+        self.force_kill_query(ErrorCode::AbortedQuery(
+            "Aborted query, because the server is shutting down or the query was killed",
+        ));
         self.kill(/* shutdown io stream */);
     }
 
-    pub fn force_kill_query(self: &Arc<Self>) {
+    pub fn force_kill_query(self: &Arc<Self>, cause: ErrorCode) {
         let session_ctx = self.session_ctx.clone();
 
         if let Some(context_shared) = session_ctx.get_query_context_shared() {
-            context_shared.kill(/* shutdown executing query */);
+            context_shared.kill(cause);
         }
     }