fix: http handler kill not working. (#15558)

* fix: http handler kill not working.

* ci: add logic test for http kill.

Author: youngsofun

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

@@ -234,13 +234,16 @@ async fn query_final_handler(
             query_id, query_id
         );
         let http_query_manager = HttpQueryManager::instance();
-        match http_query_manager.remove_query(&query_id, RemoveReason::Finished) {
+        match http_query_manager
+            .remove_query(
+                &query_id,
+                RemoveReason::Finished,
+                ErrorCode::ClosedQuery("closed by client"),
+            )
+            .await
+        {
             Some(query) => {
                 let mut response = query.get_response_state_only().await;
-                if query.check_removed().is_none() && !response.state.state.is_stopped() {
-                    query.kill(ErrorCode::ClosedQuery("closed by client")).await;
-                    response = query.get_response_state_only().await;
-                }
                 // it is safe to set these 2 fields to None, because client now check for null/None first.
                 response.session = None;
                 response.state.affect = None;
@@ -268,15 +271,15 @@ async fn query_cancel_handler(
             query_id, query_id
         );
         let http_query_manager = HttpQueryManager::instance();
-        match http_query_manager.remove_query(&query_id, RemoveReason::Canceled) {
-            Some(query) => {
-                if query.check_removed().is_none() {
-                    query
-                        .kill(ErrorCode::AbortedQuery("canceled by client"))
-                        .await;
-                }
-                Ok(StatusCode::OK)
-            }
+        match http_query_manager
+            .remove_query(
+                &query_id,
+                RemoveReason::Canceled,
+                ErrorCode::AbortedQuery("canceled by client"),
+            )
+            .await
+        {
+            Some(_) => Ok(StatusCode::OK),
             None => Err(query_id_not_found(&query_id, &ctx.node_id)),
         }
     }

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

@@ -338,7 +338,13 @@ impl HttpQuery {
                             "last query on the session not finished",
                         ));
                     }
-                    let _ = http_query_manager.remove_query(&query_id, RemoveReason::Canceled);
+                    let _ = http_query_manager
+                        .remove_query(
+                            &query_id,
+                            RemoveReason::Canceled,
+                            ErrorCode::ClosedQuery("closed by next query"),
+                        )
+                        .await;
                 }
                 // wait for Arc<QueryContextShared> to drop and detach itself from session
                 // should not take too long

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

@@ -29,7 +29,6 @@ use databend_common_config::InnerConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_storages_common_txn::TxnManagerRef;
-use log::warn;
 use parking_lot::Mutex;
 use tokio::task;
 
@@ -150,19 +149,13 @@ impl HttpQueryManager {
                             "http query {} timeout after {} s",
                             &query_id_clone, query_result_timeout_secs
                         );
-                        match self_clone.remove_query(&query_id_clone, RemoveReason::Timeout) {
-                            Some(_) => {
-                                warn!("{msg}");
-                                if let Some(query) = http_query_weak.upgrade() {
-                                    if query.check_removed().is_none() {
-                                        query.kill(ErrorCode::AbortedQuery(&msg)).await;
-                                    }
-                                }
-                            }
-                            None => {
-                                warn!("{msg}, but already evict, too many queries?");
-                            }
-                        };
+                        _ = self_clone
+                            .remove_query(
+                                &query_id_clone,
+                                RemoveReason::Timeout,
+                                ErrorCode::AbortedQuery(&msg),
+                            )
+                            .await;
                         break;
                     }
                     ExpireResult::Sleep(t) => {
@@ -177,26 +170,24 @@ impl HttpQueryManager {
     }
 
     #[async_backtrace::framed]
-    pub(crate) fn remove_query(
+    pub(crate) async fn remove_query(
         self: &Arc<Self>,
         query_id: &str,
         reason: RemoveReason,
+        error: ErrorCode,
     ) -> Option<Arc<HttpQuery>> {
         // deref at once to avoid holding DashMap shard guard for too long.
         let query = self.queries.get(query_id).map(|q| q.clone());
-        query.map(|q| {
-            let not_removed_yet = !q.mark_removed(reason);
-            let to_evict = not_removed_yet
-                .then(|| {
-                    let mut queue = self.removed_queries.lock();
-                    queue.push(q.id.to_string())
-                })
-                .flatten();
-            if let Some(qid) = to_evict {
-                self.queries.remove(&qid);
+        if let Some(q) = &query {
+            if q.mark_removed(reason) {
+                q.kill(error).await;
+                let mut queue = self.removed_queries.lock();
+                if let Some(to_evict) = queue.push(q.id.to_string()) {
+                    self.queries.remove(&to_evict);
+                };
             }
-            q.clone()
-        })
+        }
+        query
     }
 
     #[async_backtrace::framed]

tests/suites/1_stateful/09_http_handler/09_0005_kill.result

@@ -0,0 +1,13 @@
+## query
+"Running"
+## kill
+200
+## page
+{"error":{"code":"400","message":"query id QID canceled"}}
+400
+## final
+{
+  "code": 1043,
+  "message": "canceled by client",
+  "detail": ""
+}

tests/suites/1_stateful/09_http_handler/09_0005_kill.sh

@@ -0,0 +1,22 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. "$CURDIR"/../../../shell_env.sh
+
+QID="my_query_for_kill_${RANDOM}"
+echo "## query"
+curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H "x-databend-query-id:${QID}"  -H 'Content-Type: application/json' -d '{"sql": "select sleep(0.5) from numbers(15000000);",  "pagination": { "wait_time_secs": 6}}' | jq ".state"
+echo "## kill"
+curl -s -u root: -XGET -w "%{http_code}\n"  "http://localhost:8000/v1/query/${QID}/kill"
+echo "## page"
+curl -s -u root: -XGET -w "\n%{http_code}\n" "http://localhost:8000/v1/query/${QID}/page/0" | sed "s/${QID}/QID/g"
+echo "## final"
+curl -s -u root: -XGET -w "\n" "http://localhost:8000/v1/query/${QID}/final" | jq ".error"
+
+## todo: this is flaky on ci, may lost the second row, can not reproduce locally for now
+#echo "## query_log"
+#echo "select exception_code, exception_text, log_type from system.query_log where query_id='${QID}' order by log_type" | $BENDSQL_CLIENT_CONNECT | sed "s/${QID}/QID/g"
+#----
+### query_log
+#0		1
+#1043	AbortedQuery. Code: 1043, Text = canceled by client.	4