From 037bbbb2ebb44f11e4667e4f0f4dcc959b709ab8 Mon Sep 17 00:00:00 2001
From: Hengfei Yang <hengfei.yang@gmail.com>
Date: Tue, 18 Apr 2023 14:58:13 +0800
Subject: [PATCH] test: search dead lock (#588)

* test: search dead lock

* feat: improve search took response

* feat: improve search response
---
 src/handler/http/request/search/mod.rs        | 18 +++++++++++--
 src/handler/http/router/openapi.rs            |  1 +
 src/infra/config.rs                           |  4 +++
 src/meta/search.rs                            | 27 +++++++++++++++++--
 .../search/datafusion/storage/memory.rs       |  1 +
 src/service/search/mod.rs                     | 21 +++++++--------
 6 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/src/handler/http/request/search/mod.rs b/src/handler/http/request/search/mod.rs
index e0d437ab3b..5a815eb8be 100644
--- a/src/handler/http/request/search/mod.rs
+++ b/src/handler/http/request/search/mod.rs
@@ -21,7 +21,7 @@ use std::time::Instant;
 
 use crate::common::http::get_stream_type_from_request;
 use crate::common::json;
-use crate::infra::config::CONFIG;
+use crate::infra::config::{CONFIG, SEARCH_LOCKER};
 use crate::infra::{errors, metrics};
 use crate::meta::http::HttpResponse as MetaHttpResponse;
 use crate::meta::{self, StreamType};
@@ -115,9 +115,14 @@ pub async fn search(
         return Ok(bad_request(e));
     }
 
+    // get a local search queue lock
+    let locker = SEARCH_LOCKER.clone();
+    let _locker = locker.lock().await;
+    let took_wait = start.elapsed().as_millis() as usize;
+
     // do search
     match SearchService::search(&org_id, stream_type, &req).await {
-        Ok(res) => {
+        Ok(mut res) => {
             let time = start.elapsed().as_secs_f64();
             metrics::HTTP_RESPONSE_TIME
                 .with_label_values(&[
@@ -137,6 +142,7 @@ pub async fn search(
                     stream_type.to_string().as_str(),
                 ])
                 .inc();
+            res.set_local_took(start.elapsed().as_millis() as usize, took_wait);
             Ok(HttpResponse::Ok().json(res))
         }
         Err(err) => {
@@ -236,6 +242,10 @@ pub async fn around(
         .get("size")
         .map_or(10, |v| v.parse::<usize>().unwrap_or(0));
 
+    // get a local search queue lock
+    let locker = SEARCH_LOCKER.clone();
+    let _locker = locker.lock().await;
+
     // search forward
     let req = meta::search::Request {
         query: meta::search::Query {
@@ -444,6 +454,10 @@ pub async fn values(
         end_time = chrono::Utc::now().timestamp_micros();
     }
 
+    // get a local search queue lock
+    let locker = SEARCH_LOCKER.clone();
+    let _locker = locker.lock().await;
+
     // search
     let mut req = meta::search::Request {
         query: meta::search::Query {
diff --git a/src/handler/http/router/openapi.rs b/src/handler/http/router/openapi.rs
index fc8f5e82a1..500202a48b 100644
--- a/src/handler/http/router/openapi.rs
+++ b/src/handler/http/router/openapi.rs
@@ -99,6 +99,7 @@ use crate::meta;
             meta::search::Request,
             meta::search::RequestEncoding,
             meta::search::Response,
+            meta::search::ResponseTook,
             meta::alert::Alert,
             meta::alert::AlertList,
             meta::alert::Condition,
diff --git a/src/infra/config.rs b/src/infra/config.rs
index 24681a4155..2790836805 100644
--- a/src/infra/config.rs
+++ b/src/infra/config.rs
@@ -19,8 +19,10 @@ use dotenvy::dotenv;
 use once_cell::sync::Lazy;
 use reqwest::Client;
 use std::sync::atomic::AtomicU8;
+use std::sync::Arc;
 use std::time::Duration;
 use sys_info::hostname;
+use tokio::sync::Mutex;
 
 use crate::common::file::get_file_meta;
 use crate::meta::alert::{AlertDestination, AlertList, DestinationTemplate, Trigger, TriggerTimer};
@@ -39,6 +41,8 @@ pub static SEARCHING_IN_CACHE: AtomicU8 = AtomicU8::new(0);
 
 pub static CONFIG: Lazy<Config> = Lazy::new(init);
 pub static INSTANCE_ID: Lazy<DashMap<String, String>> = Lazy::new(DashMap::new);
+
+pub static SEARCH_LOCKER: Lazy<Arc<Mutex<bool>>> = Lazy::new(|| Arc::new(Mutex::const_new(false)));
 pub static TELEMETRY_CLIENT: Lazy<segment::HttpClient> = Lazy::new(|| {
     segment::HttpClient::new(
         Client::builder()
diff --git a/src/meta/search.rs b/src/meta/search.rs
index e0c895a4ec..0862f05357 100644
--- a/src/meta/search.rs
+++ b/src/meta/search.rs
@@ -136,6 +136,8 @@ impl Request {
 #[schema(as = SearchResponse)]
 pub struct Response {
     pub took: usize,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub took_detail: Option<ResponseTook>,
     #[schema(value_type = Vec<Object>)]
     pub hits: Vec<json::Value>,
     #[serde(skip_serializing_if = "HashMap::is_empty")]
@@ -150,11 +152,19 @@ pub struct Response {
     #[serde(skip_serializing_if = "String::is_empty")]
     pub response_type: String,
 }
+#[derive(Clone, Debug, Serialize, Deserialize, Default, ToSchema)]
+pub struct ResponseTook {
+    pub total: usize,
+    pub wait_queue: usize,
+    pub cluster_total: usize,
+    pub cluster_wait_queue: usize,
+}
 
 impl Response {
     pub fn new(from: usize, size: usize) -> Self {
         Response {
             took: 0,
+            took_detail: None,
             total: 0,
             from,
             size,
@@ -176,8 +186,21 @@ impl Response {
         val.push(hit.to_owned());
     }
 
-    pub fn set_took(&mut self, val: usize) {
-        self.took = val;
+    pub fn set_cluster_took(&mut self, val: usize, wait: usize) {
+        self.took = val - wait;
+        self.took_detail = Some(ResponseTook {
+            total: 0,
+            wait_queue: 0,
+            cluster_total: val,
+            cluster_wait_queue: wait,
+        });
+    }
+
+    pub fn set_local_took(&mut self, val: usize, wait: usize) {
+        if self.took_detail.is_some() {
+            self.took_detail.as_mut().unwrap().total = val;
+            self.took_detail.as_mut().unwrap().wait_queue = wait;
+        }
     }
 
     pub fn set_total(&mut self, val: usize) {
diff --git a/src/service/search/datafusion/storage/memory.rs b/src/service/search/datafusion/storage/memory.rs
index 43e0b62f5b..cbc2915a2c 100644
--- a/src/service/search/datafusion/storage/memory.rs
+++ b/src/service/search/datafusion/storage/memory.rs
@@ -173,6 +173,7 @@ impl InMemory {
     async fn get_bytes(&self, location: &Path) -> Result<Bytes> {
         let file = location.to_string();
         let data = file_data::get(&file);
+        tokio::task::yield_now().await;
         if data.is_ok() {
             return Ok(data.unwrap());
         }
diff --git a/src/service/search/mod.rs b/src/service/search/mod.rs
index 102be4d932..db247f8dc0 100644
--- a/src/service/search/mod.rs
+++ b/src/service/search/mod.rs
@@ -50,14 +50,12 @@ pub async fn search(
     stream_type: StreamType,
     req: &meta::search::Request,
 ) -> Result<Response, Error> {
-    let root_span = info_span!("service:search:enter");
+    let root_span = info_span!("srv:search:enter");
 
     let mut req: cluster_rpc::SearchRequest = req.to_owned().into();
-
     req.org_id = org_id.to_string();
     req.stype = cluster_rpc::SearchType::User as i32;
     req.stream_type = stream_type.to_string();
-
     search_in_cluster(req).instrument(root_span).await
 }
 
@@ -123,7 +121,7 @@ async fn partition_by_file_list(sql: &sql::Sql, stream_type: StreamType) -> Vec<
     }
 }
 
-#[instrument(name = "service:search:cluster", skip(req))]
+#[instrument(name = "srv:search:cluster", skip(req))]
 async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response, Error> {
     let start = std::time::Instant::now();
 
@@ -133,9 +131,10 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
     } else {
         Some(get_queue_lock().await?)
     };
+    let took_wait = start.elapsed().as_millis() as usize;
 
     //XXX span1.exit(); // drop span1
-    //XXX let span2 = info_span!("service:search:cluster:prepare_base").entered();
+    //XXX let span2 = info_span!("srv:search:cluster:prepare_base").entered();
 
     // get nodes from cluster
     let mut nodes = cluster::get_cached_online_query_nodes().unwrap();
@@ -186,7 +185,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
     };
 
     //XXX span3.exit(); // drop span3
-    //XXX let span4 = info_span!("service:search:cluster:do_search").entered();
+    //XXX let span4 = info_span!("srv:search:cluster:do_search").entered();
 
     // make grpc auth token
     let user = ROOT_USER.get("root").unwrap();
@@ -226,7 +225,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
             req_file_list_end,
         );
 
-        let grpc_span = info_span!("service:search:cluster:grpc_search");
+        let grpc_span = info_span!("srv:search:cluster:grpc_search");
 
         let node_addr = node.grpc_addr.clone();
         let credentials_str = credentials.clone();
@@ -325,7 +324,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
     }
 
     //XXX span4.exit(); // drop span4
-    //XXX let span6 = info_span!("service:search:cluster:release_queue_lock").entered();
+    //XXX let span6 = info_span!("srv:search:cluster:release_queue_lock").entered();
 
     // search done, release lock
     if locker.is_some() {
@@ -335,7 +334,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
     }
 
     //XXX span6.exit(); // drop span6
-    //XXX let span7 = info_span!("service:search:cluster:merge_result").entered();
+    //XXX let span7 = info_span!("srv:search:cluster:merge_result").entered();
 
     // merge multiple instances data
     let mut file_count = 0;
@@ -406,7 +405,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
     }
 
     //XXX span7.exit(); // drop span7
-    //XXX let _span8 = info_span!("service:search:cluster:response").entered();
+    //XXX let _span8 = info_span!("srv:search:cluster:response").entered();
 
     // final result
     let mut result = Response::new(sql.meta.offset, sql.meta.limit);
@@ -471,7 +470,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result<Response,
     result.aggs.remove("_count");
 
     result.set_total(total);
-    result.set_took(start.elapsed().as_millis() as usize);
+    result.set_cluster_took(start.elapsed().as_millis() as usize, took_wait);
     result.set_file_count(file_count as usize);
     result.set_scan_size(scan_size as usize);