starred/openobserve
1
0
Fork 0
mirror of https://github.com/openobserve/openobserve.git synced 2026-11-06 00:09:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf: improve user level search group (#3791)

Browse Source
This commit is contained in:
Hengfei Yang
2024-06-23 11:55:41 +08:00
committed by GitHub
parent 2158b0de1a
commit bde4c8a952
2 changed files with 57 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/infra/src/db/nats.rs
+20 -1
View File
@@ -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<Mutex<HashMap<String, Arc<Mutex<bool>>>>> =
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();
src/service/search/cluster/mod.rs
+37 -73
View File
@@ -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))