feat: add metrics workgroup (#9400)

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;

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 <http://www.gnu.org/licenses/>.
+
+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<Instant>,
+    splits: Vec<Split>,
+}
+
+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<String>) -> 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::<Vec<_>>()
+                .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())
+    }
+}

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<dyn datafusion::physical_plan::ExecutionPlan>,
-    Option<AsyncDefer>,
+    Option<DeferredLock>,
     ScanStats,
 );
 
@@ -324,9 +324,9 @@ async fn get_ctx_and_physical_plan(
     req: &FlightSearchRequest,
 ) -> Result<PlanResult, infra::errors::Error> {
     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))

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<AsyncDefer>,
+    defer_lock: Option<DeferredLock>,
     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<AsyncDefer>) -> Self {
-        self.defer = defer;
+    pub fn with_defer_lock(mut self, lock: Option<DeferredLock>) -> 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<AsyncDefer>,
+    defer_lock: Option<DeferredLock>,
     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!(

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<Value> {
-    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)
 }
 

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<Sql>, mut req: Request) -> Result<SearchResult> {
-    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<Sql>, mut req: Request) -> Result<S
     let file_id_list_vec = file_id_list.values().flatten().collect::<Vec<_>>();
     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::<i64>();
-    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<Sql>, mut req: Request) -> Result<S
     };
 
     // 3. get nodes
-    let get_node_start = std::time::Instant::now();
     let is_local_mode = req.local_mode.unwrap_or_default();
     let role_group = if is_local_mode {
         None
@@ -184,7 +182,7 @@ pub async fn search(trace_id: &str, sql: Arc<Sql>, mut req: Request) -> Result<S
                 .node_name(LOCAL_NODE.name.clone())
                 .component("flight:leader get nodes".to_string())
                 .search_role("leader".to_string())
-                .duration(get_node_start.elapsed().as_millis() as usize)
+                .duration(stop_watch.record_split("get_nodes").as_millis() as usize)
                 .desc(format!(
                     "get nodes num: {}, querier num: {}",
                     nodes.len(),
@@ -199,22 +197,19 @@ pub async fn search(trace_id: &str, sql: Arc<Sql>, mut req: Request) -> Result<S
         .with_label_values(&[&req.org_id])
         .inc();
 
-    // 4. check work group
-    let file_list_took = start.elapsed().as_millis() as usize;
-    #[cfg(not(feature = "enterprise"))]
-    let (took_wait, work_group_str, locker) =
-        check_work_group(&req, trace_id, start, file_list_took).await?;
-    #[cfg(feature = "enterprise")]
-    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,
+        "logs",
         &nodes,
         &file_id_list_vec,
-        start,
-        file_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));
 
@@ -225,52 +220,6 @@ pub async fn search(trace_id: &str, sql: Arc<Sql>, mut req: Request) -> Result<S
         .with_label_values(&[&sql.org_id])
         .inc();
 
-    // release lock when search done or get error
-    let trace_id_move = trace_id.to_string();
-    let org_id_move = sql.org_id.clone();
-    #[cfg(not(feature = "enterprise"))]
-    let _defer = AsyncDefer::new({
-        async move {
-            metrics::QUERY_RUNNING_NUMS
-                .with_label_values(&[&org_id_move])
-                .dec();
-            // search done, release lock
-            let _ = dist_lock::unlock_with_trace_id(&trace_id_move, &locker)
-                .await
-                .map_err(|e| {
-                    log::error!(
-                        "[trace_id {trace_id_move}] release lock in flight search error: {e}"
-                    );
-                    Error::Message(e.to_string())
-                });
-            log::info!("[trace_id {trace_id_move}] release lock in flight search");
-        }
-    });
-
-    #[cfg(feature = "enterprise")]
-    let user_id = req.user_id.clone();
-    #[cfg(feature = "enterprise")]
-    let _defer = AsyncDefer::new({
-        async move {
-            metrics::QUERY_RUNNING_NUMS
-                .with_label_values(&[&org_id_move])
-                .dec();
-            // search done, release lock
-            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 search error: {e}",
-                    );
-                    e.to_string();
-                });
-            log::info!("[trace_id {trace_id_move}] release work group in flight search");
-        }
-    });
-
     // 5. partition file list
     let partitioned_file_lists = partition_file_lists(file_id_list, &nodes, role_group).await?;
     let mut need_ingesters = 0;
@@ -364,7 +313,7 @@ pub async fn search(trace_id: &str, sql: Arc<Sql>, mut req: Request) -> Result<S
     };
 
     // release source
-    drop(_defer);
+    drop(_lock);
 
     // 9. get data from datafusion
     let (data, mut scan_stats, partial_err): (Vec<RecordBatch>, ScanStats, String) = match task {
@@ -373,7 +322,10 @@ pub async fn search(trace_id: &str, sql: Arc<Sql>, mut req: Request) -> Result<S
         Err(err) => 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<infra::dist_lock::Locker>)> {
-    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<o2_enterprise::enterprise::search::WorkGroup>,
-)> {
-    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<o2_enterprise::enterprise::search::WorkGroup> =
-        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<TableReference, Vec<FileId>>,

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

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<dyn ExecutionPlan>,
-    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::<Vec<_>>();
     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(

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 <http://www.gnu.org/licenses/>.
+#[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<WorkGroup>,
+    _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<DeferredLock> {
+    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<DeferredLock> {
+    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<DeferredLock> {
+    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<DeferredLock> {
+    // 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
+}