diff --git a/src/config/src/utils/mod.rs b/src/config/src/utils/mod.rs index 1340708f4e..555d66339d 100644 --- a/src/config/src/utils/mod.rs +++ b/src/config/src/utils/mod.rs @@ -39,4 +39,5 @@ pub mod str; pub mod sysinfo; pub mod tantivy; pub mod time; +pub mod took_watcher; pub mod util; diff --git a/src/config/src/utils/took_watcher.rs b/src/config/src/utils/took_watcher.rs new file mode 100644 index 0000000000..e7709cfb83 --- /dev/null +++ b/src/config/src/utils/took_watcher.rs @@ -0,0 +1,129 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::fmt; + +use tokio::time::{Duration, Instant}; + +#[derive(Clone, Debug)] +pub struct Split { + pub name: String, + pub duration: Duration, +} + +#[derive(Clone)] +pub struct TookWatcher { + start_time: Instant, + last_split: Option, + splits: Vec, +} + +impl TookWatcher { + /// Creates and starts the internal timer + pub fn new() -> Self { + TookWatcher { + start_time: Instant::now(), + last_split: None, + splits: Vec::new(), + } + } + + /// Get total elapsed time since stopwatch started + pub fn total_elapsed(&self) -> Duration { + self.start_time.elapsed() + } + + /// Convenience method to get total time in milliseconds + pub fn total_millis(&self) -> u64 { + self.total_elapsed().as_millis() as u64 + } + + /// Convenience method to get total time as f64 seconds + pub fn total_secs_f64(&self) -> f64 { + self.total_elapsed().as_secs_f64() + } + + /// Get elapsed time since last split WITHOUT recording a new split. + /// This is safe to call in loops for timeout checking. + pub fn split_elapsed(&self) -> Duration { + match self.last_split { + Some(ref instant) => instant.elapsed(), + None => self.start_time.elapsed(), + } + } + + /// Convenience method to get elapsed time in milliseconds without mutating + pub fn split_elapsed_millis(&self) -> u64 { + self.split_elapsed().as_millis() as u64 + } + + /// Convenience method to get elapsed time as f64 seconds without mutating + pub fn split_elapsed_secs_f64(&self) -> f64 { + self.split_elapsed().as_secs_f64() + } + + /// Record a named split and return the duration since last split + pub fn record_split(&mut self, name: impl Into) -> Duration { + let duration = self.split_elapsed(); + + self.splits.push(Split { + name: name.into(), + duration, + }); + + self.last_split = Some(Instant::now()); + duration + } + + /// Get all recorded splits + pub fn get_splits(&self) -> &[Split] { + &self.splits + } + + /// Generate a formatted summary of all splits + pub fn get_summary(&self) -> String { + format!( + "total: {} ms, breakdown: {}", + self.total_millis(), + self.splits + .iter() + .map(|s| format!("{}={}ms", s.name, s.duration.as_millis())) + .collect::>() + .join(", ") + ) + } +} + +impl Default for TookWatcher { + fn default() -> Self { + Self::new() + } +} + +impl fmt::Debug for TookWatcher { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("StopWatch") + .field("elapsed_ms", &self.split_elapsed_millis()) + .field("total_ms", &self.total_millis()) + .field("splits", &self.splits.len()) + .finish() + } +} + +impl fmt::Display for TookWatcher { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}ms", self.total_millis()) + } +} diff --git a/src/handler/grpc/flight/mod.rs b/src/handler/grpc/flight/mod.rs index e102b9ae39..c026e7a608 100644 --- a/src/handler/grpc/flight/mod.rs +++ b/src/handler/grpc/flight/mod.rs @@ -53,7 +53,7 @@ use crate::{ service::search::{ grpc::flight as grpcFlight, inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}, - utils::AsyncDefer, + work_group::DeferredLock, }, }; @@ -148,7 +148,7 @@ impl FlightService for FlightServiceImpl { ); // prepare dataufion context - let (ctx, physical_plan, defer, scan_stats) = match result { + let (ctx, physical_plan, lock, scan_stats) = match result { Ok(v) => v, Err(e) => { // clear session data @@ -215,7 +215,7 @@ impl FlightService for FlightServiceImpl { let mut stream = FlightEncoderStreamBuilder::new(write_options, 33554432) .with_trace_id(trace_id.to_string()) .with_is_super(is_super_cluster) - .with_defer(defer) + .with_defer_lock(lock) .with_start(start) .with_custom_message(PreCustomMessage::ScanStats(scan_stats)) .with_custom_message(PreCustomMessage::ScanStatsRef(scan_stats_ref)) @@ -313,7 +313,7 @@ impl FlightService for FlightServiceImpl { type PlanResult = ( datafusion::prelude::SessionContext, Arc, - Option, + Option, ScanStats, ); @@ -324,9 +324,9 @@ async fn get_ctx_and_physical_plan( req: &FlightSearchRequest, ) -> Result { if req.super_cluster_info.is_super_cluster { - let (ctx, physical_plan, defer, scan_stats) = + let (ctx, physical_plan, lock, scan_stats) = crate::service::search::super_cluster::follower::search(trace_id, req).await?; - Ok((ctx, physical_plan, Some(defer), scan_stats)) + Ok((ctx, physical_plan, Some(lock), scan_stats)) } else { let (ctx, physical_plan, scan_stats) = grpcFlight::search(trace_id, req).await?; Ok((ctx, physical_plan, None, scan_stats)) diff --git a/src/handler/grpc/flight/stream.rs b/src/handler/grpc/flight/stream.rs index 4196d3416e..4e43983923 100644 --- a/src/handler/grpc/flight/stream.rs +++ b/src/handler/grpc/flight/stream.rs @@ -25,7 +25,7 @@ use futures_core::ready; use tracing::info_span; use tracing_opentelemetry::OpenTelemetrySpanExt; -use crate::{handler::grpc::flight::clear_session_data, service::search::utils::AsyncDefer}; +use crate::{handler::grpc::flight::clear_session_data, service::search::work_group::DeferredLock}; pub struct FlightEncoderStreamBuilder { encoder: FlightDataEncoder, @@ -34,7 +34,7 @@ pub struct FlightEncoderStreamBuilder { // query context trace_id: String, is_super: bool, - defer: Option, + defer_lock: Option, start: std::time::Instant, } @@ -46,7 +46,7 @@ impl FlightEncoderStreamBuilder { custom_messages: vec![], trace_id: String::new(), is_super: false, - defer: None, + defer_lock: None, start: std::time::Instant::now(), } } @@ -66,8 +66,8 @@ impl FlightEncoderStreamBuilder { self } - pub fn with_defer(mut self, defer: Option) -> Self { - self.defer = defer; + pub fn with_defer_lock(mut self, lock: Option) -> Self { + self.defer_lock = lock; self } @@ -91,7 +91,7 @@ impl FlightEncoderStreamBuilder { custom_messages: self.custom_messages, trace_id: self.trace_id, is_super: self.is_super, - defer: self.defer, + defer_lock: self.defer_lock, start: self.start, first_batch: true, span, @@ -113,7 +113,7 @@ pub struct FlightEncoderStream { // query context trace_id: String, is_super: bool, - defer: Option, + defer_lock: Option, start: std::time::Instant, span: tracing::Span, child_span: tracing::Span, @@ -269,7 +269,7 @@ impl Drop for FlightEncoderStream { .inc(); // defer is only set for super cluster follower leader - if let Some(defer) = self.defer.take() { + if let Some(defer) = self.defer_lock.take() { drop(defer); } else { log::info!( diff --git a/src/service/promql/search/mod.rs b/src/service/promql/search/mod.rs index e2d8c443c9..6ae6095d27 100644 --- a/src/service/promql/search/mod.rs +++ b/src/service/promql/search/mod.rs @@ -27,7 +27,10 @@ use config::{ self_reporting::usage::{RequestStats, UsageType}, stream::StreamType, }, - utils::time::{now_micros, second_micros}, + utils::{ + time::{now_micros, second_micros}, + took_watcher::TookWatcher, + }, }; use futures::future::try_join_all; use hashbrown::HashMap; @@ -35,6 +38,8 @@ use infra::{ client::grpc::make_grpc_metrics_client, errors::{Error, ErrorCodes, Result}, }; +#[cfg(feature = "enterprise")] +use o2_enterprise::enterprise::search::WorkGroup; use proto::cluster_rpc; use tracing::{Instrument, info_span}; @@ -85,9 +90,14 @@ async fn search_in_cluster( req: cluster_rpc::MetricsQueryRequest, user_email: &str, ) -> Result { - let op_start = std::time::Instant::now(); + let mut stop_watch = TookWatcher::new(); let started_at = now_micros(); let cfg = get_config(); + let timeout = if req.timeout > 0 { + req.timeout as u64 + } else { + cfg.limit.query_timeout + }; let &cluster_rpc::MetricsQueryStmt { ref query, @@ -114,6 +124,32 @@ async fn search_in_cluster( query, ); + // Check work group (OSS uses dist_lock, Enterprise uses WorkGroup::Short) + #[cfg(not(feature = "enterprise"))] + let _lock = crate::service::search::work_group::check_work_group( + trace_id, + &req.org_id, + timeout, + &mut stop_watch, + "metrics", + ) + .await + .map_err(|e| Error::ErrorCode(ErrorCodes::ServerInternalError(e.to_string())))?; + + // Enterprise: Always use Short workgroup for metrics queries + #[cfg(feature = "enterprise")] + let _lock = crate::service::search::work_group::check_work_group( + trace_id, + &req.org_id, + Some(user_email), + timeout, + WorkGroup::Short, + &mut stop_watch, + "metrics", + ) + .await + .map_err(|e| Error::ErrorCode(ErrorCodes::ServerInternalError(e.to_string())))?; + // get querier nodes from cluster let mut nodes = cluster::get_cached_online_querier_nodes(Some(RoleGroup::Interactive)) .await @@ -170,6 +206,7 @@ async fn search_in_cluster( } else { merge_matrix_query(&cached_values) }; + return Ok(values); } @@ -279,7 +316,12 @@ async fn search_in_cluster( } }; for res in task_results { - results.push(res?); + match res { + Ok(response) => results.push(response), + Err(err) => { + return Err(err); + } + } } // merge multiple instances data @@ -319,17 +361,18 @@ async fn search_in_cluster( } else { return Err(server_internal_error("invalid result type")); }; + let result_time = stop_watch.record_split("result").as_secs_f64(); log::info!( "[trace_id {trace_id}] promql->search->result: files: {}, scan_size: {} mb, took: {} ms", scan_stats.files, scan_stats.original_size, - op_start.elapsed().as_millis(), + (result_time * 1000.0) as u64, ); let req_stats = RequestStats { records: scan_stats.records, size: scan_stats.original_size as f64, - response_time: op_start.elapsed().as_secs_f64(), + response_time: result_time, request_body: Some(query.to_string()), user_email: Some(user_email.to_string()), min_ts: Some(start), @@ -367,6 +410,11 @@ async fn search_in_cluster( log::error!("[trace_id {trace_id}] promql->search->cache: set cache err: {err:?}"); } + log::info!( + "[trace_id {trace_id}] promql->search search finished {}", + stop_watch.get_summary() + ); + Ok(values) } diff --git a/src/service/search/cluster/flight.rs b/src/service/search/cluster/flight.rs index 4edd93eff5..0112e36108 100644 --- a/src/service/search/cluster/flight.rs +++ b/src/service/search/cluster/flight.rs @@ -28,14 +28,13 @@ use config::{ stream::{QueryPartitionStrategy, StreamType}, }, metrics, - utils::{json, time::now_micros}, + utils::{json, time::now_micros, took_watcher::TookWatcher}, }; use datafusion::{ common::TableReference, physical_plan::visit_execution_plan, prelude::SessionContext, }; use hashbrown::{HashMap, HashSet}; use infra::{ - dist_lock, errors::{Error, ErrorCodes, Result}, file_list::FileId, }; @@ -65,7 +64,7 @@ use crate::{ }, inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}, sql::Sql, - utils::{AsyncDefer, ScanStatsVisitor, check_query_default_limit_exceeded}, + utils::{ScanStatsVisitor, check_query_default_limit_exceeded}, }, }, }; @@ -77,7 +76,7 @@ use crate::{ fields(org_id = req.org_id) )] pub async fn search(trace_id: &str, sql: Arc, mut req: Request) -> Result { - let start = std::time::Instant::now(); + let mut stop_watch = TookWatcher::new(); let cfg = get_config(); log::info!("[trace_id {trace_id}] flight->search: start {sql}"); @@ -108,7 +107,7 @@ pub async fn search(trace_id: &str, sql: Arc, mut req: Request) -> Result>(); let file_id_list_num = file_id_list_vec.len(); let file_id_list_records = file_id_list_vec.iter().map(|v| v.records).sum::(); - let file_id_list_took = start.elapsed().as_millis() as usize; + let file_id_list_took = stop_watch.record_split("get_file_list").as_millis() as usize; log::info!( "{}", search_inspector_fields( @@ -137,7 +136,6 @@ pub async fn search(trace_id: &str, sql: Arc, mut req: Request) -> Result, mut req: Request) -> Result, mut req: Request) -> Result, mut req: Request) -> Result, mut req: Request) -> Result, ScanStats, String) = match task { @@ -373,7 +322,10 @@ pub async fn search(trace_id: &str, sql: Arc, mut req: Request) -> Result Err(err), }?; - log::info!("[trace_id {trace_id}] flight->search: search finished"); + log::info!( + "[trace_id {trace_id}] flight->search: search finished, {}", + stop_watch.get_summary() + ); scan_stats.format_to_mb(); scan_stats.file_list_took += file_id_list_took as i64; @@ -508,149 +460,6 @@ pub async fn get_online_querier_nodes( Ok(nodes) } -#[cfg(not(feature = "enterprise"))] -#[tracing::instrument(name = "service:search:cluster:flight:check_work_group", skip_all)] -pub async fn check_work_group( - req: &Request, - trace_id: &str, - start: std::time::Instant, - file_list_took: usize, // the time took to get file list -) -> Result<(usize, String, Option)> { - let cfg = get_config(); - let work_group_str = "global".to_string(); - - let locker_key = format!("/search/cluster_queue/{work_group_str}"); - let locker = if cfg.common.local_mode || !cfg.common.feature_query_queue_enabled { - None - } else { - dist_lock::lock_with_trace_id(trace_id, &locker_key, req.timeout as u64) - .await - .map_err(|e| { - metrics::QUERY_PENDING_NUMS - .with_label_values(&[&req.org_id]) - .dec(); - Error::Message(e.to_string()) - })? - }; - - // done in the queue - let took_wait = start.elapsed().as_millis() as usize - file_list_took; - log::info!( - "[trace_id {trace_id}] search: wait in queue took: {} ms", - took_wait, - ); - Ok((took_wait, work_group_str, locker)) -} - -#[cfg(feature = "enterprise")] -#[tracing::instrument(name = "service:search:cluster:flight:check_work_group", skip_all)] -pub async fn check_work_group( - req: &Request, - trace_id: &str, - nodes: &[Node], - file_id_list_vec: &[&FileId], - start: std::time::Instant, - file_list_took: usize, // the time took to get file list - search_role: String, -) -> Result<( - usize, - String, - Option, -)> { - let cfg = get_config(); - let user_id = req.user_id.clone(); - let user_id = user_id.as_deref(); - - // 1. get work group - // Check if this is a background task (alerts, reports, derived streams) - let is_background_task = req - .search_event_type - .as_ref() - .and_then(|st| config::meta::search::SearchEventType::try_from(st.as_str()).ok()) - .map(|st| st.is_background()) - .unwrap_or(false); - - let work_group: Option = - Some(o2_enterprise::enterprise::search::work_group::predict( - nodes, - file_id_list_vec, - is_background_task, - )); - - SEARCH_SERVER - .add_work_group(trace_id, work_group.clone()) - .await; - - let work_group_str = work_group.as_ref().unwrap().to_string(); - - let locker_key = format!("/search/cluster_queue/{work_group_str}"); - // 2. get a cluster search queue lock - let locker = if cfg.common.local_mode || !cfg.common.feature_query_queue_enabled { - None - } else { - dist_lock::lock_with_trace_id(trace_id, &locker_key, req.timeout as u64) - .await - .map_err(|e| { - metrics::QUERY_PENDING_NUMS - .with_label_values(&[&req.org_id]) - .dec(); - Error::Message(e.to_string()) - })? - }; - - // 3. check global concurrency - super::work_group_checking(trace_id, start, req, &work_group, &locker, None).await?; - - // 4. check user concurrency - if user_id.is_some() && !is_background_task { - super::work_group_checking(trace_id, start, req, &work_group, &locker, user_id).await?; - } - - // 5. process the search in the work group - if let Err(e) = work_group - .as_ref() - .unwrap() - .process(trace_id, user_id) - .await - { - metrics::QUERY_PENDING_NUMS - .with_label_values(&[&req.org_id]) - .dec(); - dist_lock::unlock_with_trace_id(trace_id, &locker).await?; - return Err(Error::Message(e.to_string())); - } - - // 6. unlock the queue in no enterprise version, - if let Err(e) = dist_lock::unlock_with_trace_id(trace_id, &locker).await { - metrics::QUERY_PENDING_NUMS - .with_label_values(&[&req.org_id]) - .dec(); - work_group - .as_ref() - .unwrap() - .done(trace_id, user_id) - .await - .map_err(|e| Error::Message(e.to_string()))?; - return Err(e); - } - - // done in the queue - let took_wait = start.elapsed().as_millis() as usize - file_list_took; - log::info!( - "{}", - search_inspector_fields( - format!("[trace_id {trace_id}] search: wait in queue took: {took_wait} ms"), - SearchInspectorFieldsBuilder::new() - .node_name(LOCAL_NODE.name.clone()) - .component("flight:check_work_group".to_string()) - .search_role(search_role) - .duration(took_wait) - .build() - ) - ); - Ok((took_wait, work_group_str, work_group)) -} - #[tracing::instrument(name = "service:search:cluster:flight:partition_file_lists", skip_all)] pub async fn partition_file_lists( file_id_lists: HashMap>, diff --git a/src/service/search/mod.rs b/src/service/search/mod.rs index a55a5b789d..287e256aad 100644 --- a/src/service/search/mod.rs +++ b/src/service/search/mod.rs @@ -115,6 +115,7 @@ pub(crate) mod streaming; #[cfg(feature = "enterprise")] pub(crate) mod super_cluster; pub(crate) mod utils; +pub(crate) mod work_group; /// The result of search in cluster /// data, scan_stats, wait_in_queue, is_partial, partial_err diff --git a/src/service/search/super_cluster/follower.rs b/src/service/search/super_cluster/follower.rs index c00ab2c757..ba167108f8 100644 --- a/src/service/search/super_cluster/follower.rs +++ b/src/service/search/super_cluster/follower.rs @@ -41,7 +41,7 @@ use crate::service::{ db::enrichment_table, search::{ SEARCH_SERVER, - cluster::flight::{check_work_group, get_online_querier_nodes, partition_file_list}, + cluster::flight::{get_online_querier_nodes, partition_file_list}, datafusion::{ distributed_plan::{ NewEmptyExecVisitor, @@ -52,7 +52,7 @@ use crate::service::{ exec::{DataFusionContextBuilder, register_udf}, }, inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}, - utils::AsyncDefer, + work_group::DeferredLock, }, }; @@ -70,10 +70,11 @@ pub async fn search( ) -> Result<( SessionContext, Arc, - AsyncDefer, + DeferredLock, ScanStats, )> { - let start = std::time::Instant::now(); + let mut stop_watch = config::utils::took_watcher::TookWatcher::new(); + let cfg = config::get_config(); let mut req: Request = (*flight_request).clone().into(); let trace_id = trace_id.to_string(); @@ -116,7 +117,8 @@ pub async fn search( get_file_id_lists(&trace_id, &req.org_id, stream_type, &stream, req.time_range).await?; let file_id_list_vec = file_id_list.iter().collect::>(); let file_id_list_num = file_id_list_vec.len(); - let file_id_list_took = start.elapsed().as_millis() as usize; + + let file_id_list_took = stop_watch.record_split("get_file_list").as_millis() as usize; log::info!( "{}", search_inspector_fields( @@ -196,42 +198,25 @@ pub async fn search( ); // check work group - let (took_wait, work_group_str, work_group) = check_work_group( - &req, + let _lock = crate::service::search::work_group::acquire_work_group_lock( &trace_id, + &req, + &mut stop_watch, + "super_cluster_follower", &nodes, &file_id_list_vec, - start, - file_id_list_took, - "leader".to_string(), ) .await?; + + let took_wait = _lock.took_wait; + let work_group_str = _lock.work_group_str.clone(); + // add work_group req.add_work_group(Some(work_group_str.clone())); log::info!( "[trace_id {trace_id}] flight->follower_leader: add work_group: {work_group_str}, took: {took_wait} ms" ); - // release work_group in flight follow search - let user_id = req.user_id.clone(); - let trace_id_move = trace_id.to_string(); - let defer = AsyncDefer::new({ - async move { - let _ = work_group - .as_ref() - .unwrap() - .done(&trace_id_move, user_id.as_deref()) - .await - .map_err(|e| { - log::error!( - "[trace_id {trace_id_move}] release work_group in flight follow search error: {e}", - ); - e.to_string(); - }); - log::info!("[trace_id {trace_id_move}] release work_group in flight follow search"); - } - }); - // partition file list let partition_file_lists = partition_file_list(file_id_list, &nodes, role_group).await?; log::info!( @@ -303,7 +288,7 @@ pub async fn search( ..Default::default() }; - Ok((ctx, physical_plan, defer, scan_stats)) + Ok((ctx, physical_plan, _lock, scan_stats)) } #[tracing::instrument( diff --git a/src/service/search/work_group.rs b/src/service/search/work_group.rs new file mode 100644 index 0000000000..9b1fdef33d --- /dev/null +++ b/src/service/search/work_group.rs @@ -0,0 +1,343 @@ +// Copyright 2025 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . +#[cfg(feature = "enterprise")] +use config::meta::search::SearchEventType; +use config::{ + datafusion::request::Request, get_config, meta::cluster::Node, metrics, + utils::took_watcher::TookWatcher, +}; +use infra::{ + errors::{Error, Result}, + file_list::FileId, +}; +#[cfg(feature = "enterprise")] +use o2_enterprise::enterprise::search::WorkGroup; + +use super::utils::AsyncDefer; +#[cfg(feature = "enterprise")] +use crate::service::search::SEARCH_SERVER; + +/// Guard that automatically releases work group lock when dropped +pub struct DeferredLock { + pub took_wait: usize, + pub work_group_str: String, + #[cfg(feature = "enterprise")] + #[allow(dead_code)] // Available for external use/debugging + pub work_group: Option, + _guard: AsyncDefer, +} + +/// OSS version: Uses distributed lock for concurrency control +#[cfg(not(feature = "enterprise"))] +#[tracing::instrument( + name = "service:search:work_group:check", + skip_all, + fields(caller = caller, trace_id = trace_id) +)] +pub async fn check_work_group( + trace_id: &str, + org_id: &str, + timeout: u64, + stop_watch: &mut TookWatcher, + caller: &str, +) -> Result { + let cfg = get_config(); + let work_group_str = "global".to_string(); + + let locker_key = format!("/search/cluster_queue/{work_group_str}"); + let locker = if cfg.common.local_mode || !cfg.common.feature_query_queue_enabled { + None + } else { + infra::dist_lock::lock_with_trace_id(trace_id, &locker_key, timeout) + .await + .map_err(|e| { + metrics::QUERY_PENDING_NUMS + .with_label_values(&[org_id]) + .dec(); + Error::Message(e.to_string()) + })? + }; + + let took_wait = stop_watch.record_split("queue_wait").as_millis() as usize; + log::info!("[trace_id {trace_id}] {caller}->search: wait in queue took: {took_wait} ms"); + + // Create cleanup guard + let trace_id_owned = trace_id.to_string(); + let caller_owned = caller.to_string(); + let guard = AsyncDefer::new(async move { + if let Err(e) = infra::dist_lock::unlock_with_trace_id(&trace_id_owned, &locker).await { + log::error!( + "[trace_id {trace_id_owned}] {caller_owned}->search: failed to unlock: {e:?}" + ); + } + }); + + Ok(DeferredLock { + took_wait, + work_group_str, + _guard: guard, + }) +} + +/// Enterprise version: Uses workgroup concurrency control with waiting +#[cfg(feature = "enterprise")] +#[tracing::instrument( + name = "service:search:work_group:check", + skip_all, + fields(caller = caller, trace_id = trace_id, work_group = ?work_group) +)] +pub async fn check_work_group( + trace_id: &str, + org_id: &str, + user_id: Option<&str>, + timeout: u64, + work_group: WorkGroup, + stop_watch: &mut TookWatcher, + caller: &str, +) -> Result { + let cfg = get_config(); + let work_group_str = work_group.to_string(); + + // Get distributed lock temporarily (for queue coordination) + let locker_key = format!("/search/cluster_queue/{work_group_str}"); + let locker = if cfg.common.local_mode || !cfg.common.feature_query_queue_enabled { + None + } else { + infra::dist_lock::lock_with_trace_id(trace_id, &locker_key, timeout) + .await + .map_err(|e| { + metrics::QUERY_PENDING_NUMS + .with_label_values(&[org_id]) + .dec(); + Error::Message(e.to_string()) + })? + }; + + // Check global concurrency + work_group_need_wait( + trace_id, + stop_watch, + org_id, + timeout, + &work_group, + None, + caller, + ) + .await?; + + // Check user concurrency (skip for background tasks) + let is_background = matches!(work_group, WorkGroup::Background); + if user_id.is_some() && !is_background { + work_group_need_wait( + trace_id, + stop_watch, + org_id, + timeout, + &work_group, + user_id, + caller, + ) + .await?; + } + + // Add to workgroup queue + if let Err(e) = work_group.process(trace_id, user_id).await { + metrics::QUERY_PENDING_NUMS + .with_label_values(&[org_id]) + .dec(); + if let Err(unlock_err) = infra::dist_lock::unlock_with_trace_id(trace_id, &locker).await { + log::error!( + "[trace_id {trace_id}] {caller}->search: failed to unlock after process error: {unlock_err:?}" + ); + } + return Err(Error::Message(e.to_string())); + } + + // Release distributed lock (workgroup queue now owns the concurrency control) + if let Err(e) = infra::dist_lock::unlock_with_trace_id(trace_id, &locker).await { + metrics::QUERY_PENDING_NUMS + .with_label_values(&[org_id]) + .dec(); + if let Err(done_err) = work_group.done(trace_id, user_id).await { + log::error!( + "[trace_id {trace_id}] {caller}->search: failed to mark work group as done after unlock error: {done_err:?}" + ); + } + return Err(e); + } + + let took_wait = stop_watch.record_split("queue_wait").as_millis() as usize; + log::info!("[trace_id {trace_id}] {caller}->search: wait in queue took: {took_wait} ms"); + + // Create cleanup guard + let trace_id_owned = trace_id.to_string(); + let user_id_owned = user_id.map(|s| s.to_string()); + let org_id_owned = org_id.to_string(); + let work_group_clone = work_group.clone(); + let caller_owned = caller.to_string(); + let guard = AsyncDefer::new(async move { + if let Err(e) = work_group_clone + .done(&trace_id_owned, user_id_owned.as_deref()) + .await + { + log::error!( + "[trace_id {trace_id_owned}] {caller_owned}->search: failed to mark work group as done: {e:?}" + ); + } + + metrics::QUERY_RUNNING_NUMS + .with_label_values(&[&org_id_owned]) + .dec(); + log::info!("[trace_id {trace_id_owned}] search completed, metrics decremented"); + }); + + Ok(DeferredLock { + took_wait, + work_group_str, + work_group: Some(work_group), + _guard: guard, + }) +} + +/// Enterprise: Wait for workgroup slot to become available +#[cfg(feature = "enterprise")] +#[tracing::instrument( + name = "service:search:work_group:need_wait", + skip_all, + fields(caller = caller, user_id = user_id) +)] +async fn work_group_need_wait( + trace_id: &str, + stop_watch: &mut TookWatcher, + org_id: &str, + timeout: u64, + work_group: &WorkGroup, + user_id: Option<&str>, + caller: &str, +) -> Result<()> { + let mut log_wait = false; + loop { + // Check timeout - use total_millis() to check total time from start + if stop_watch.total_millis() >= timeout * 1000 { + metrics::QUERY_TIMEOUT_NUMS + .with_label_values(&[org_id]) + .inc(); + return Err(Error::Message(format!( + "[trace_id {trace_id}] {caller}->search: request timeout in queue" + ))); + } + + match work_group.need_wait(user_id).await { + Ok((true, cur, max)) => { + // Need to wait - concurrency limit reached + if !log_wait { + log::info!( + "[trace_id {trace_id}] user: {user_id:?} is waiting in work_group {work_group:?}[{cur}/{max}]" + ); + log_wait = true; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + Ok((false, cur, max)) => { + // Got approval - slot available + if log_wait { + log::info!( + "[trace_id {trace_id}] user: {user_id:?} get approved in work_group {work_group:?}[{cur}/{max}]" + ); + } + return Ok(()); + } + Err(e) => { + return Err(Error::Message(e.to_string())); + } + } + } +} + +/// Acquire work group lock with automatic work group prediction (enterprise) +/// or simple distributed lock (OSS) +/// +/// This is a high-level helper that encapsulates: +/// - Determining if request is a background task +/// - Predicting appropriate work group (enterprise only) +/// - Adding work group to search server (enterprise only) +/// - Acquiring the work group lock +#[cfg(not(feature = "enterprise"))] +pub async fn acquire_work_group_lock( + trace_id: &str, + req: &Request, + stop_watch: &mut TookWatcher, + caller: &str, + _nodes: &[Node], + _file_id_list_vec: &[&FileId], +) -> Result { + check_work_group( + trace_id, + &req.org_id, + req.timeout as u64, + stop_watch, + caller, + ) + .await +} + +/// Acquire work group lock with automatic work group prediction (enterprise) +/// or simple distributed lock (OSS) +/// +/// This is a high-level helper that encapsulates: +/// - Determining if request is a background task +/// - Predicting appropriate work group (enterprise only) +/// - Adding work group to search server (enterprise only) +/// - Acquiring the work group lock +#[cfg(feature = "enterprise")] +pub async fn acquire_work_group_lock( + trace_id: &str, + req: &Request, + stop_watch: &mut TookWatcher, + caller: &str, + nodes: &[Node], + file_id_list_vec: &[&FileId], +) -> Result { + // Predict workgroup first + let is_background_task = req + .search_event_type + .as_ref() + .and_then(|st| SearchEventType::try_from(st.as_str()).ok()) + .map(|st| st.is_background()) + .unwrap_or(false); + + let work_group = o2_enterprise::enterprise::search::work_group::predict( + nodes, + file_id_list_vec, + is_background_task, + ); + + SEARCH_SERVER + .add_work_group(trace_id, Some(work_group.clone())) + .await; + + let user_id = req.user_id.as_deref(); + + check_work_group( + trace_id, + &req.org_id, + user_id, + req.timeout as u64, + work_group, + stop_watch, + caller, + ) + .await +}