From bde4c8a9520e2f95e0ffa29978b659f9ee488e59 Mon Sep 17 00:00:00 2001 From: Hengfei Yang Date: Sun, 23 Jun 2024 11:55:41 +0800 Subject: [PATCH] perf: improve user level search group (#3791) --- src/infra/src/db/nats.rs | 21 +++++- src/service/search/cluster/mod.rs | 110 ++++++++++-------------------- 2 files changed, 57 insertions(+), 74 deletions(-) diff --git a/src/infra/src/db/nats.rs b/src/infra/src/db/nats.rs index 64ed9505d6..f19c562c90 100644 --- a/src/infra/src/db/nats.rs +++ b/src/infra/src/db/nats.rs @@ -27,8 +27,9 @@ use bytes::Bytes; use config::{cluster, get_config, ider, utils::base64}; use futures::{StreamExt, TryStreamExt}; use hashbrown::HashMap; +use once_cell::sync::Lazy; use tokio::{ - sync::{mpsc, OnceCell}, + sync::{mpsc, Mutex, OnceCell}, task::JoinHandle, }; @@ -661,6 +662,10 @@ pub async fn connect() -> async_nats::Client { } } +/// global locker for nats +static LOCAL_LOCKER: Lazy>>>> = + Lazy::new(|| Mutex::new(HashMap::new())); + pub(crate) struct Locker { key: String, lock_id: String, @@ -689,6 +694,20 @@ impl Locker { chrono::Utc::now().timestamp_micros() + Duration::from_secs(timeout).as_micros() as i64; let value = Bytes::from(format!("{}:{}", self.lock_id, expiration)); let key = key_encode(new_key); + + // check local global locker + let mut local_mutex = LOCAL_LOCKER.lock().await; + let locker = match local_mutex.get(&key) { + Some(v) => v.clone(), + None => { + let locker = Arc::new(Mutex::new(false)); + local_mutex.insert(key.clone(), locker.clone()); + locker + } + }; + drop(local_mutex); + let _lock_guard = locker.lock().await; + // check if the locker already expired, clean it if let Ok(Some(ret)) = bucket.get(&key).await { let ret = String::from_utf8_lossy(&ret).to_string(); diff --git a/src/service/search/cluster/mod.rs b/src/service/search/cluster/mod.rs index 833242fa9c..4cd2edbb5e 100644 --- a/src/service/search/cluster/mod.rs +++ b/src/service/search/cluster/mod.rs @@ -281,42 +281,31 @@ pub async fn search( } else { dist_lock::lock(&locker_key, req.timeout as u64).await? }; + // check global concurrency #[cfg(feature = "enterprise")] - { - // check user concurrency - if user_id.is_some() { - loop { - if start.elapsed().as_millis() as u64 >= req.timeout as u64 * 1000 { - dist_lock::unlock(&locker).await?; - return Err(Error::Message(format!( - "[trace_id {trace_id}] search: timeout in queue" - ))); - } - match work_group.as_ref().unwrap().need_wait(user_id).await { - Ok(true) => { - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - } - Ok(false) => { - break; - } - Err(e) => { - dist_lock::unlock(&locker).await?; - return Err(Error::Message(e.to_string())); - } - } + loop { + if start.elapsed().as_millis() as u64 >= req.timeout as u64 * 1000 { + dist_lock::unlock(&locker).await?; + return Err(Error::Message(format!( + "[trace_id {trace_id}] search: timeout in queue" + ))); + } + match work_group.as_ref().unwrap().need_wait(None).await { + Ok(true) => { + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; } - // process the search in the work group - if let Err(e) = work_group - .as_ref() - .unwrap() - .process(trace_id, user_id) - .await - { + Ok(false) => { + break; + } + Err(e) => { dist_lock::unlock(&locker).await?; return Err(Error::Message(e.to_string())); } } - // check global concurrency + } + // check user concurrency + #[cfg(feature = "enterprise")] + if user_id.is_some() { loop { if start.elapsed().as_millis() as u64 >= req.timeout as u64 * 1000 { dist_lock::unlock(&locker).await?; @@ -324,7 +313,7 @@ pub async fn search( "[trace_id {trace_id}] search: timeout in queue" ))); } - match work_group.as_ref().unwrap().need_wait(None).await { + match work_group.as_ref().unwrap().need_wait(user_id).await { Ok(true) => { tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; } @@ -337,14 +326,21 @@ pub async fn search( } } } - // 3. process the search in the work group - if let Err(e) = work_group.as_ref().unwrap().process(trace_id, None).await { - dist_lock::unlock(&locker).await?; - return Err(Error::Message(e.to_string())); - } - dist_lock::unlock(&locker).await?; } + // 3. process the search in the work group + #[cfg(feature = "enterprise")] + if let Err(e) = work_group + .as_ref() + .unwrap() + .process(trace_id, user_id) + .await + { + dist_lock::unlock(&locker).await?; + return Err(Error::Message(e.to_string())); + } + #[cfg(feature = "enterprise")] + dist_lock::unlock(&locker).await?; // done in the queue let took_wait = start.elapsed().as_millis() as usize - file_list_took; log::info!( @@ -481,17 +477,9 @@ pub async fn search( work_group .as_ref() .unwrap() - .done(&trace_id, None) + .done(&trace_id, user_id) .await .map_err(|e| Error::Message(e.to_string()))?; - if user_id.is_some() { - work_group - .as_ref() - .unwrap() - .done(&trace_id, user_id) - .await - .map_err(|e| Error::Message(e.to_string()))?; - } return Err(Error::ErrorCode(ErrorCodes::SearchCancelQuery(format!( "[trace_id {trace_id}] search->grpc: search canceled before call search->grpc" )))); @@ -620,17 +608,9 @@ pub async fn search( work_group .as_ref() .unwrap() - .done(trace_id, None) + .done(trace_id, user_id) .await .map_err(|e| Error::Message(e.to_string()))?; - if user_id.is_some() { - work_group - .as_ref() - .unwrap() - .done(trace_id, user_id) - .await - .map_err(|e| Error::Message(e.to_string()))?; - } } return Err(Error::ErrorCode(ErrorCodes::ServerInternalError( e.to_string(), @@ -656,17 +636,9 @@ pub async fn search( work_group .as_ref() .unwrap() - .done(trace_id, None) + .done(trace_id, user_id) .await .map_err(|e| Error::Message(e.to_string()))?; - if user_id.is_some() { - work_group - .as_ref() - .unwrap() - .done(trace_id, user_id) - .await - .map_err(|e| Error::Message(e.to_string()))?; - } } return Err(e); } @@ -681,17 +653,9 @@ pub async fn search( work_group .as_ref() .unwrap() - .done(trace_id, None) + .done(trace_id, user_id) .await .map_err(|e| Error::Message(e.to_string()))?; - if user_id.is_some() { - work_group - .as_ref() - .unwrap() - .done(trace_id, user_id) - .await - .map_err(|e| Error::Message(e.to_string()))?; - } } Ok((merge_batches, scan_stats, took_wait, is_partial))