feat: /v1/status merge status of HttpQueryManager. (#18778)

- merge count and last_query_finished_at from http query manager
- update last_query_finished_at each time a query finished （not need count=0）, same with last_query_started_at

Author: youngsofun

scripts/ci/deploy/config/databend-query-node-1.toml

@@ -28,6 +28,7 @@ http_handler_host = "0.0.0.0"
 http_handler_port = 8000
 # mainly for test/debug
 # http_session_timeout_secs = 90
+http_handler_result_timeout_secs = 60
 
 # Databend Query FlightSQL Handler.
 flight_sql_handler_host = "0.0.0.0"

src/query/service/src/servers/admin/v1/instance_status.rs

@@ -19,6 +19,7 @@ use poem::IntoResponse;
 use serde::Deserialize;
 use serde::Serialize;
 
+use crate::servers::HttpQueryManager;
 use crate::sessions::QueriesQueueManager;
 use crate::sessions::SessionManager;
 
@@ -40,8 +41,6 @@ pub struct InstanceStatus {
     pub instance_started_at: u64,
     // the local timestamp, may be useful to avoid the clock drift issues
     pub instance_timestamp: u64,
-    // the timestamp on last valid request to /v1/query
-    pub last_query_request_at: Option<u64>,
 }
 
 // lightweight way to get status
@@ -51,17 +50,27 @@ pub struct InstanceStatus {
 pub async fn instance_status_handler() -> poem::Result<impl IntoResponse> {
     let session_manager = SessionManager::instance();
     let queue_manager = QueriesQueueManager::instance();
+    let (http_query_count, last_query_finished_at_http) = HttpQueryManager::instance().status();
     let status = session_manager.get_current_session_status();
+
+    let last_query_finished_at = [
+        status.last_query_started_at.map(unix_timestamp_secs),
+        last_query_finished_at_http,
+    ]
+    .iter()
+    .flatten()
+    .max()
+    .copied();
+
     let status = InstanceStatus {
-        running_queries_count: status.running_queries_count,
+        running_queries_count: status.running_queries_count.max(http_query_count),
         active_sessions_count: status.active_sessions_count,
         queuing_queries_count: queue_manager.length() as u64,
         last_query_started_at: status.last_query_started_at.map(unix_timestamp_secs),
-        last_query_finished_at: status.last_query_finished_at.map(unix_timestamp_secs),
+        last_query_finished_at,
         max_running_query_executed_secs: status.max_running_query_executed_secs,
         instance_started_at: unix_timestamp_secs(status.instance_started_at),
         instance_timestamp: unix_timestamp_secs(SystemTime::now()),
-        last_query_request_at: status.last_query_request_at,
     };
     Ok(Json(status))
 }

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

@@ -60,7 +60,7 @@ use super::query::ExecuteStateKind;
 use super::query::HttpQuery;
 use super::query::HttpQueryRequest;
 use super::query::HttpQueryResponseInternal;
-use super::query::RemoveReason;
+use super::query::StopReason;
 use crate::clusters::ClusterDiscovery;
 use crate::servers::http::error::HttpErrorCode;
 use crate::servers::http::error::QueryError;
@@ -93,7 +93,6 @@ use crate::servers::http::v1::HttpQueryManager;
 use crate::servers::http::v1::HttpSessionConf;
 use crate::servers::HttpHandlerKind;
 use crate::sessions::QueryAffect;
-use crate::sessions::SessionManager;
 
 pub fn make_page_uri(query_id: &str, page_no: usize) -> String {
     format!("/v1/query/{}/page/{}", query_id, page_no)
@@ -167,10 +166,10 @@ pub struct QueryResponse {
 }
 
 impl QueryResponse {
-    pub(crate) fn removed(query_id: &str, remove_reason: RemoveReason) -> impl IntoResponse {
+    pub(crate) fn removed(query_id: &str, remove_reason: StopReason) -> impl IntoResponse {
         let id = query_id.to_string();
         let state = match remove_reason {
-            RemoveReason::Finished => ExecuteStateKind::Succeeded,
+            StopReason::Finished => ExecuteStateKind::Succeeded,
             _ => ExecuteStateKind::Failed,
         };
         Json(QueryResponse {
@@ -306,10 +305,10 @@ async fn query_final_handler(
         );
         let http_query_manager = HttpQueryManager::instance();
         match http_query_manager
-            .remove_query(
+            .stop_query(
                 &query_id,
                 &ctx.client_session_id,
-                RemoveReason::Finished,
+                StopReason::Finished,
                 ErrorCode::ClosedQuery("Query closed by client"),
             )
             .await?
@@ -347,10 +346,10 @@ async fn query_cancel_handler(
         );
         let http_query_manager = HttpQueryManager::instance();
         match http_query_manager
-            .remove_query(
+            .stop_query(
                 &query_id,
                 &ctx.client_session_id,
-                RemoveReason::Canceled,
+                StopReason::Canceled,
                 ErrorCode::AbortedQuery("canceled by client"),
             )
             .await?
@@ -449,9 +448,8 @@ async fn query_page_handler(
                 query.update_expire_time(false).await;
                 let (resp, next_is_final) = QueryResponse::from_internal(query_id, resp, false);
                 if next_is_final {
-                    query.wait_for_final()
+                    query.set_data_drained()
                 }
-                SessionManager::instance().new_query_request();
                 Ok(resp)
             }
         }
@@ -501,7 +499,6 @@ pub(crate) async fn query_handler(
             client_session_id_info,
             mask_connection_info(&format!("{:?}", req))
         );
-        SessionManager::instance().new_query_request();
         let sql = req.sql.clone();
 
         match HttpQuery::try_create(ctx, req.clone()).await {
@@ -522,7 +519,6 @@ pub(crate) async fn query_handler(
                 let http_query_manager = HttpQueryManager::instance();
                 let query = http_query_manager.add_query(query).await;
 
-                query.update_expire_time(true).await;
                 // tmp workaround to tolerant old clients
                 let resp = query
                     .get_response_page(0)
@@ -550,7 +546,7 @@ pub(crate) async fn query_handler(
                 let (resp, next_is_final) =
                     QueryResponse::from_internal(query.id.to_string(), resp, false);
                 if next_is_final {
-                    query.wait_for_final()
+                    query.set_data_drained()
                 }
                 Ok(resp.into_response())
             }
@@ -842,7 +838,7 @@ pub fn query_route() -> Route {
     route
 }
 
-fn query_id_removed(query_id: &str, remove_reason: RemoveReason) -> PoemError {
+fn query_id_removed(query_id: &str, remove_reason: StopReason) -> PoemError {
     PoemError::from_string(
         format!("[HTTP-QUERY] Query ID {query_id} {}", remove_reason),
         StatusCode::BAD_REQUEST,

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

@@ -52,7 +52,7 @@ use serde::Serializer;
 
 use super::execute_state::ExecutionError;
 use super::HttpQueryContext;
-use super::RemoveReason;
+use super::StopReason;
 use crate::servers::http::error::QueryError;
 use crate::servers::http::v1::http_query_handlers::QueryResponseField;
 use crate::servers::http::v1::query::blocks_serializer::BlocksSerializer;
@@ -405,15 +405,15 @@ impl HttpSessionConf {
                 session.set_all_variables(state.get_variables()?)
             }
             if let Some(id) = state.last_query_ids.first() {
-                if let Some(last_query) = http_query_manager.queries.get(id) {
+                if let Some(last_query) = http_query_manager.get_query(id) {
                     let state = *last_query.state.lock();
                     if !matches!(
                         state,
-                        HttpQueryState::Removed(RemoveReason::Finished | RemoveReason::Canceled)
+                        HttpQueryState::Stopped(StopReason::Finished | StopReason::Canceled)
                     ) {
                         warn!(
-                            "[HTTP-QUERY] Last query id not finished yet, id = {}, state = {:?}",
-                            id, state
+                            "[HTTP-QUERY] Last query id not finished yet, id = {}, state = {:?}, is_data_drained={}",
+                            id, state, last_query.is_data_drained.load(Ordering::Relaxed)
                         );
                     }
                 }
@@ -511,17 +511,15 @@ pub struct HttpQueryResponseInternal {
 
 #[derive(Debug, Clone, Copy)]
 pub enum HttpQueryState {
-    Working,
-    WaitForFinal,
     ExpireAt(Instant),
-    Removed(RemoveReason),
+    Stopped(StopReason),
 }
 
 #[derive(Debug)]
 pub enum ExpireResult {
     Expired,
     Sleep(Duration),
-    Removed,
+    Stopped,
 }
 
 pub struct HttpQuery {
@@ -533,6 +531,7 @@ pub struct HttpQuery {
     executor: Arc<Mutex<Executor>>,
     page_manager: Arc<TokioMutex<PageManager>>,
     state: Arc<Mutex<HttpQueryState>>,
+    is_data_drained: AtomicBool,
     /// The timeout for the query result polling. In the normal case, the client driver
     /// should fetch the paginated result in a timely manner, and the interval should not
     /// exceed this result_timeout_secs.
@@ -624,6 +623,8 @@ impl HttpQuery {
 
         let settings = session.get_settings();
         let result_timeout_secs = settings.get_http_handler_result_timeout_secs()?;
+        let deadline = Instant::now()
+            + Duration::from_secs(result_timeout_secs + req.pagination.wait_time_secs as u64);
 
         Ok(HttpQuery {
             id: query_id,
@@ -635,11 +636,12 @@ impl HttpQuery {
             page_manager: Arc::new(TokioMutex::new(page_manager)),
             result_timeout_secs,
 
-            state: Arc::new(Mutex::new(HttpQueryState::Working)),
+            state: Arc::new(Mutex::new(HttpQueryState::ExpireAt(deadline))),
             temp_tbl_mgr: session.temp_tbl_mgr().clone(),
             query_mem_stat: ctx.get_query_memory_tracking(),
             is_txn_mgr_saved: Default::default(),
             last_session_conf: Default::default(),
+            is_data_drained: AtomicBool::new(false),
         })
     }
 
@@ -847,35 +849,32 @@ impl HttpQuery {
 
     #[async_backtrace::framed]
     pub async fn update_expire_time(&self, before_wait: bool) {
-        let duration = Duration::from_secs(self.result_timeout_secs)
-            + if before_wait {
-                Duration::from_secs(self.request.pagination.wait_time_secs as u64)
-            } else {
-                Duration::new(0, 0)
-            };
-        let deadline = Instant::now() + duration;
+        let mut to = Duration::from_secs(self.result_timeout_secs);
+        if before_wait {
+            to += Duration::from_secs(self.request.pagination.wait_time_secs as u64)
+        };
+        let deadline = Instant::now() + to;
         let mut t = self.state.lock();
         *t = HttpQueryState::ExpireAt(deadline);
     }
 
-    pub fn mark_removed(&self, remove_reason: RemoveReason) -> bool {
+    pub fn mark_stopped(&self, remove_reason: StopReason) -> bool {
         let mut t = self.state.lock();
-        if !matches!(*t, HttpQueryState::Removed(_)) {
-            *t = HttpQueryState::Removed(remove_reason);
+        if !matches!(*t, HttpQueryState::Stopped(_)) {
+            *t = HttpQueryState::Stopped(remove_reason);
             true
         } else {
             false
         }
     }
 
-    pub fn wait_for_final(&self) {
-        let mut t = self.state.lock();
-        *t = HttpQueryState::WaitForFinal;
+    pub fn set_data_drained(&self) {
+        self.is_data_drained.store(true, Ordering::Relaxed);
     }
 
-    pub fn check_removed(&self) -> Option<RemoveReason> {
+    pub fn check_removed(&self) -> Option<StopReason> {
         let t = self.state.lock();
-        if let HttpQueryState::Removed(r) = *t {
+        if let HttpQueryState::Stopped(r) = *t {
             Some(r)
         } else {
             None
@@ -895,8 +894,7 @@ impl HttpQuery {
                     ExpireResult::Sleep(expire_at - now)
                 }
             }
-            HttpQueryState::Removed(_) => ExpireResult::Removed,
-            _ => ExpireResult::Sleep(Duration::from_secs(self.result_timeout_secs)),
+            HttpQueryState::Stopped(_) => ExpireResult::Stopped,
         }
     }
 
@@ -911,8 +909,7 @@ impl HttpQuery {
                 *expire_state = HttpQueryState::ExpireAt(deadline);
                 true
             }
-            HttpQueryState::Removed(_) => false,
-            HttpQueryState::Working | HttpQueryState::WaitForFinal => true,
+            HttpQueryState::Stopped(_) => false,
         }
     }