diff --git a/.zed/debug.json b/.zed/debug.json deleted file mode 100644 index 1e3c0929c9..0000000000 --- a/.zed/debug.json +++ /dev/null @@ -1,29 +0,0 @@ -// Project-local debug tasks -// -// For more documentation on how to configure debug tasks, -// see: https://zed.dev/docs/debugger -[ - { - "label": "DBG OpenObserve", - "build": { - "command": "cargo", - "args": ["build"] - }, - "program": "$ZED_WORKTREE_ROOT/target/debug/openobserve", - "env": { - "RUST_LOG": "debug", - "ZO_ROOT_USER_EMAIL": "root@example.com", - "ZO_ROOT_USER_PASSWORD": "Complexpass#123", - "ZO_QUERY_RECOMMENDATION_DURATION": "3600000000", - "ZO_QUERY_RECOMMENDATION_INTERVAL": "5", - "ZO_QUERY_RECOMMENDATION_TOP_K": "128", - "VITE_OPENOBSERVE_ENTERPRISE": "true", - "ZO_USAGE_BATCH_SIZE": "200", - "ZO_USAGE_PUBLISH_INTERVAL": "2", - "ZO_USAGE_REPORTING_ENABLED": "true" - }, - "request": "launch", - "sourceLanguages": ["rust"], - "adapter": "CodeLLDB" - } -] diff --git a/benchmarks/pyroscope/README.md b/benchmarks/pyroscope/README.md deleted file mode 100644 index 559e1375b4..0000000000 --- a/benchmarks/pyroscope/README.md +++ /dev/null @@ -1,11 +0,0 @@ -## Pyroscope for profiling - -Pyroscope can be used for basic profiling of `OpenObserve`. -The following steps can be followed: -- Setup a local `pyroscope` server from `benchmarks/pyroscope` directory: - - `docker-compose up docker-compose.yml` -- In another terminal, compile and run the `openobserve` code with the debug profile. -- Note that we are setting `RUST_LOG=warn` to restrict logging only to `warn`. -- `cargo build --profile release-profiling && RUST_LOG=warn ZO_ROOT_USER_EMAIL=root@example.com ZO_ROOT_USER_PASSWORD=Complexpass#123 ./target/release-profiling/openobserve` - -- Launch http://localhost:4040 diff --git a/benchmarks/pyroscope/docker-compose.yml b/benchmarks/pyroscope/docker-compose.yml deleted file mode 100644 index a6618bed39..0000000000 --- a/benchmarks/pyroscope/docker-compose.yml +++ /dev/null @@ -1,10 +0,0 @@ -version: '3' -services: - pyroscope: - image: pyroscope/pyroscope:latest - environment: - - PYROSCOPE_LOG_LEVEL=debug - ports: - - '4040:4040' - command: - - 'server' diff --git a/src/common/infra/cluster/mod.rs b/src/common/infra/cluster/mod.rs index 6e22dbf660..46ea3aef10 100644 --- a/src/common/infra/cluster/mod.rs +++ b/src/common/infra/cluster/mod.rs @@ -216,77 +216,5 @@ mod tests { assert!(get_cached_online_query_nodes(None).await.is_some()); assert!(get_cached_online_ingester_nodes().await.is_some()); assert!(get_cached_online_querier_nodes(None).await.is_some()); - - // Reset the global state. - reset_consistent_hash().await; - - // Test consistent hash logic. - let node = LOCAL_NODE.clone(); - for i in 0..10 { - let node_q = Node { - name: format!("node-q-{i}").to_string(), - role: [Role::Querier].to_vec(), - ..node.clone() - }; - let node_c = Node { - name: format!("node-c-{i}").to_string(), - role: [Role::Compactor].to_vec(), - ..node.clone() - }; - add_node_to_consistent_hash(&node_q, &Role::Querier, Some(RoleGroup::Interactive)) - .await; - add_node_to_consistent_hash(&node_q, &Role::Querier, Some(RoleGroup::Background)).await; - add_node_to_consistent_hash(&node_q, &Role::Querier, None).await; - add_node_to_consistent_hash(&node_c, &Role::Compactor, None).await; - add_node_to_consistent_hash(&node_c, &Role::FlattenCompactor, None).await; - } - - for key in ["test", "test1", "test2", "test3", "test4", "test5", "test6"] { - println!( - "{key}-q: {}", - get_node_from_consistent_hash(key, &Role::Querier, None) - .await - .unwrap() - ); - println!( - "{key}-c: {}", - get_node_from_consistent_hash(key, &Role::Compactor, None) - .await - .unwrap() - ); - } - - // gxhash hash - let data = [ - ["test", "node-q-2", "node-c-7"], - ["test1", "node-q-3", "node-c-0"], - ["test2", "node-q-6", "node-c-5"], - ["test3", "node-q-9", "node-c-9"], - ["test4", "node-q-2", "node-c-0"], - ["test5", "node-q-5", "node-c-8"], - ["test6", "node-q-5", "node-c-8"], - ]; - - remove_node_from_consistent_hash(&node, &Role::Querier, Some(RoleGroup::Interactive)).await; - remove_node_from_consistent_hash(&node, &Role::Querier, Some(RoleGroup::Background)).await; - remove_node_from_consistent_hash(&node, &Role::Compactor, None).await; - remove_node_from_consistent_hash(&node, &Role::FlattenCompactor, None).await; - for key in data { - assert_eq!( - get_node_from_consistent_hash(key.first().unwrap(), &Role::Querier, None).await, - Some(key.get(1).unwrap().to_string()) - ); - assert_eq!( - get_node_from_consistent_hash(key.first().unwrap(), &Role::Compactor, None).await, - Some(key.get(2).unwrap().to_string()) - ); - } - - let ret = print_consistent_hash().await; - assert_eq!(ret.len(), 4); - assert_eq!(ret["querier_interactive"].len(), 10); - assert_eq!(ret["querier_background"].len(), 10); - assert_eq!(ret["compactor"].len(), 10); - assert_eq!(ret["flatten_compactor"].len(), 10); } } diff --git a/src/config/src/config.rs b/src/config/src/config.rs index 04bd08d415..aa9c74e95f 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -2630,6 +2630,8 @@ fn check_limit_config(cfg: &mut Config) -> Result<(), anyhow::Error> { cfg.limit.batch_size = 8192; } cfg.limit.batch_size = cfg.limit.batch_size.clamp(1024, 8192); + // clamp datafusion_min_partition_num to 1 + cfg.limit.datafusion_min_partition_num = cfg.limit.datafusion_min_partition_num.max(1); Ok(()) } @@ -3377,6 +3379,13 @@ fn check_inverted_index_config(cfg: &mut Config) -> Result<(), anyhow::Error> { Ok(()) } +pub fn ensure_not_empty(s: &str, name: &str) -> Result<(), anyhow::Error> { + if s.trim().is_empty() { + return Err(anyhow::anyhow!("{} is empty", name)); + } + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -3505,4 +3514,24 @@ mod tests { std::env::remove_var("ZO_USAGE_REPORT_TO_OWN_ORG"); } } + + #[test] + fn test_ensure_not_empty_valid() { + assert!(ensure_not_empty("valid", "TEST").is_ok()); + } + + #[test] + fn test_ensure_not_empty_invalid() { + assert!(ensure_not_empty("", "TEST").is_err()); + } + + #[test] + fn test_ensure_not_empty_with_whitespace() { + assert!(ensure_not_empty(" value ", "TEST").is_ok()); + } + + #[test] + fn test_ensure_not_empty_single_char() { + assert!(ensure_not_empty("a", "TEST").is_ok()); + } } diff --git a/src/config/src/meta/mod.rs b/src/config/src/meta/mod.rs index 3f111917fd..1765a90cba 100644 --- a/src/config/src/meta/mod.rs +++ b/src/config/src/meta/mod.rs @@ -1,4 +1,4 @@ -// Copyright 2025 OpenObserve Inc. +// Copyright 2026 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 diff --git a/src/config/src/utils/async_file.rs b/src/config/src/utils/async_file.rs index 71a80fe850..50eaeaa99a 100644 --- a/src/config/src/utils/async_file.rs +++ b/src/config/src/utils/async_file.rs @@ -193,8 +193,8 @@ pub fn create_wal_dir_datetime_filter( let month = match components.next().map(|c| c.parse::()) { Some(Ok(m @ 1..=12)) => m, - Some(_) => return false, // Parsed, but invalid month number - None => start_time.month(), // Not present or failed to parse + Some(_) => return false, // Parsed, but invalid month number + None => return true, // Not present or failed to parse, could be a skippable path }; let month_days = [31u32, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]; @@ -209,13 +209,25 @@ pub fn create_wal_dir_datetime_filter( } } Some(_) => return false, - None => start_time.day(), + None => { + if month == start_time.month() { + start_time.day() + } else { + end_time.day() + } + } }; let hour = match components.next().map(|c| c.parse::()) { Some(Ok(hour @ 0..24)) => hour, Some(_) => return false, - None => start_time.hour(), + None => { + if month == start_time.month() && day == start_time.day() { + start_time.hour() + } else { + end_time.hour() + } + } }; let date_range_check = diff --git a/src/config/src/utils/hash/mod.rs b/src/config/src/utils/hash/mod.rs index 0e056dc046..5a7679839c 100644 --- a/src/config/src/utils/hash/mod.rs +++ b/src/config/src/utils/hash/mod.rs @@ -1,4 +1,4 @@ -// Copyright 2025 OpenObserve Inc. +// Copyright 2026 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 @@ -24,6 +24,10 @@ pub trait Sum64 { fn sum64(&mut self, key: &str) -> u64; } +pub fn sum64(key: &str) -> u64 { + gxhash::new().sum64(key) +} + pub fn get_passcode_hash(pass: &str, salt: &str) -> String { let t_cost = 4; let m_cost = 2048; diff --git a/src/handler/grpc/flight/mod.rs b/src/handler/grpc/flight/mod.rs index 52f71d323c..c5ff9c9e2f 100644 --- a/src/handler/grpc/flight/mod.rs +++ b/src/handler/grpc/flight/mod.rs @@ -154,11 +154,14 @@ impl FlightService for FlightServiceImpl { let (ctx, physical_plan, lock, scan_stats) = match result { Ok(v) => v, Err(e) => { - // clear session data clear_session_data(&trace_id); - log::error!( - "[trace_id {trace_id}] flight->search: do_get physical plan generate error: {e:?}", - ); + #[cfg(feature = "enterprise")] + if get_o2_config().work_group.max_nodes_per_query > 0 { + o2_enterprise::enterprise::search::admission::ledger::release(&trace_id); + log::error!( + "[trace_id {trace_id}] flight->search: do_get physical plan generate error: {e:?}", + ); + } return Err(Status::internal(e.to_string())); } }; @@ -211,11 +214,14 @@ impl FlightService for FlightServiceImpl { let peak_memory_ref = get_peak_memory(&physical_plan); let stream = execute_stream(physical_plan, ctx.task_ctx().clone()).map_err(|e| { - // clear session data clear_session_data(&trace_id); - log::error!( - "[trace_id {trace_id}] flight->search: do_get physical plan execution error: {e:?}", - ); + #[cfg(feature = "enterprise")] + if get_o2_config().work_group.max_nodes_per_query > 0 { + o2_enterprise::enterprise::search::admission::ledger::release(&trace_id); + log::error!( + "[trace_id {trace_id}] flight->search: do_get physical plan execution error: {e:?}", + ); + } Status::internal(e.to_string()) })?; diff --git a/src/handler/grpc/flight/stream.rs b/src/handler/grpc/flight/stream.rs index 4e43983923..b1c5e91e62 100644 --- a/src/handler/grpc/flight/stream.rs +++ b/src/handler/grpc/flight/stream.rs @@ -1,4 +1,4 @@ -// Copyright 2025 OpenObserve Inc. +// Copyright 2026 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 @@ -22,6 +22,8 @@ use datafusion::execution::SendableRecordBatchStream; use flight::{common::PreCustomMessage, encoder::FlightDataEncoder}; use futures::{Stream, StreamExt}; use futures_core::ready; +#[cfg(feature = "enterprise")] +use o2_enterprise::enterprise::common::config::get_config as get_o2_config; use tracing::info_span; use tracing_opentelemetry::OpenTelemetrySpanExt; @@ -268,15 +270,22 @@ impl Drop for FlightEncoderStream { .with_label_values(&["/search/flight/do_get", "200", "", "", "", ""]) .inc(); + // Release the node-level slot reservation for this Follow node. + #[cfg(feature = "enterprise")] + if get_o2_config().work_group.max_nodes_per_query > 0 { + o2_enterprise::enterprise::search::admission::ledger::release(trace_id); + log::info!("[trace_id {trace_id}] flight->search: releasing slot"); + } + // defer is only set for super cluster follower leader if let Some(defer) = self.defer_lock.take() { drop(defer); } else { + // clear session data + clear_session_data(&self.trace_id); log::info!( "[trace_id {trace_id}] flight->search: drop FlightEncoderStream, is_super: {is_super}", ); - // clear session data - clear_session_data(&self.trace_id); } } } diff --git a/src/handler/grpc/request/search/mod.rs b/src/handler/grpc/request/search/mod.rs index d8b7ea11a6..5d9707deb9 100644 --- a/src/handler/grpc/request/search/mod.rs +++ b/src/handler/grpc/request/search/mod.rs @@ -1,4 +1,4 @@ -// Copyright 2025 OpenObserve Inc. +// Copyright 2026 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 @@ -28,9 +28,15 @@ use tonic::{Request, Response, Status}; use tracing_opentelemetry::OpenTelemetrySpanExt; #[cfg(feature = "enterprise")] use { - config::meta::search::ScanStats, - config::metrics, - o2_enterprise::enterprise::search::{QueryManager, TaskStatus, WorkGroup}, + config::{meta::search::ScanStats, metrics}, + o2_enterprise::enterprise::search::{ + QueryManager, TaskStatus, WorkGroup, admission::ledger::NODE_LEDGER, + }, + proto::cluster_rpc::{ + ReleaseQueryRequest, ReleaseQueryResponse, StartQueryRequest, StartQueryResponse, + TryAcquireRequest, TryAcquireResponse, + }, + std::str::FromStr, }; use crate::{handler::grpc::MetadataMap, service::search as SearchService}; @@ -405,4 +411,77 @@ impl Search for Searcher { file_data: res.to_vec(), })) } + + // --- Slot-based admission RPCs --- + + #[cfg(feature = "enterprise")] + async fn try_acquire( + &self, + req: Request, + ) -> Result, Status> { + let r = req.into_inner(); + let wg = WorkGroup::from_str(&r.work_group).unwrap_or(WorkGroup::Short); + match NODE_LEDGER.try_acquire(&r.trace_id, &wg, r.ttl_ms) { + Ok(()) => Ok(Response::new(TryAcquireResponse { + success: true, + reason: String::new(), + })), + Err(reason) => Ok(Response::new(TryAcquireResponse { + success: false, + reason, + })), + } + } + + #[cfg(not(feature = "enterprise"))] + async fn try_acquire( + &self, + _req: Request, + ) -> Result, Status> { + Err(Status::unimplemented("Not Supported")) + } + + #[cfg(feature = "enterprise")] + async fn start_query( + &self, + req: Request, + ) -> Result, Status> { + let trace_id = req.into_inner().trace_id; + match NODE_LEDGER.start(&trace_id) { + Ok(()) => Ok(Response::new(StartQueryResponse { + success: true, + reason: String::new(), + })), + Err(reason) => Ok(Response::new(StartQueryResponse { + success: false, + reason, + })), + } + } + + #[cfg(not(feature = "enterprise"))] + async fn start_query( + &self, + _req: Request, + ) -> Result, Status> { + Err(Status::unimplemented("Not Supported")) + } + + #[cfg(feature = "enterprise")] + async fn release_query( + &self, + req: Request, + ) -> Result, Status> { + let trace_id = req.into_inner().trace_id; + NODE_LEDGER.release(&trace_id); + Ok(Response::new(ReleaseQueryResponse { success: true })) + } + + #[cfg(not(feature = "enterprise"))] + async fn release_query( + &self, + _req: Request, + ) -> Result, Status> { + Err(Status::unimplemented("Not Supported")) + } } diff --git a/src/handler/http/request/ai/chat.rs b/src/handler/http/request/ai/chat.rs index ddcfe64d1e..e292c117ff 100644 --- a/src/handler/http/request/ai/chat.rs +++ b/src/handler/http/request/ai/chat.rs @@ -177,14 +177,14 @@ pub async fn chat(Path(org_id): Path, in_req: axum::extract::Request) -> let trace_id = auth_data.get_trace_id(); let user_id = &auth_data.user_id; let org_id_str = org_id.as_str(); - let config = get_o2_config(); + let o2_cfg = get_o2_config(); // Check if AI/agent is enabled - if !config.ai.enabled { + if !o2_cfg.ai.enabled { return MetaHttpResponse::bad_request("AI is not enabled"); } - if config.ai.agent_url.is_empty() { + if o2_cfg.ai.agent_url.is_empty() { return MetaHttpResponse::bad_request("AI agent URL is not set"); } @@ -244,7 +244,7 @@ pub async fn chat(Path(org_id): Path, in_req: axum::extract::Request) -> // Create agent client let zo_config = get_config(); let client = match RcaAgentClient::new( - &config.ai.agent_url, + &o2_cfg.ai.agent_url, &zo_config.auth.root_user_email, &zo_config.auth.root_user_password, ) { @@ -473,8 +473,7 @@ pub async fn chat_stream(Path(org_id): Path, in_req: axum::extract::Requ // don't get properly exported when held inside async_stream::stream!). #[cfg(feature = "enterprise")] let otel_chat_span = { - let config = get_o2_config(); - if config.ai.tracing_enabled { + if get_o2_config().ai.tracing_enabled { use opentelemetry::trace::{SpanKind, TraceContextExt, Tracer}; // Extract parent context from incoming traceparent header @@ -560,7 +559,6 @@ pub async fn chat_stream(Path(org_id): Path, in_req: axum::extract::Requ // Extract and merge passthrough headers from config #[cfg(feature = "enterprise")] { - let _config = get_o2_config(); // Note: passthrough_headers config field needs to be added to o2_enterprise Ai config // For now, use empty string (no passthrough) until config is updated let passthrough_config = ""; // TODO: Replace with _config.ai.passthrough_headers once field is added @@ -602,14 +600,13 @@ pub async fn chat_stream(Path(org_id): Path, in_req: axum::extract::Requ use futures::StreamExt; use o2_enterprise::enterprise::alerts::rca_agent::{QueryRequest, get_agent_client}; - let config = get_o2_config(); let org_id_str = org_id.clone(); let mut code = StatusCode::OK.as_u16(); let body_bytes_str = serde_json::to_string(&prompt_body).unwrap_or_default(); // Check if AI/agent is enabled - if !config.ai.enabled { + if !get_o2_config().ai.enabled { let error_message = Some("AI is not enabled".to_string()); code = StatusCode::BAD_REQUEST.as_u16(); report_to_audit( @@ -929,9 +926,7 @@ pub async fn feedback(Path(org_id): Path, in_req: axum::extract::Request { use o2_enterprise::enterprise::alerts::rca_agent::get_agent_client; - let config = get_o2_config(); - - if !config.ai.enabled { + if !get_o2_config().ai.enabled { return MetaHttpResponse::bad_request("AI is not enabled"); } @@ -1024,9 +1019,8 @@ pub async fn confirm_action( { use o2_enterprise::enterprise::alerts::rca_agent::get_agent_client; - let config = get_o2_config(); - - if !config.ai.enabled { + let o2_cfg = get_o2_config(); + if !o2_cfg.ai.enabled { return MetaHttpResponse::bad_request("AI is not enabled"); } @@ -1061,7 +1055,7 @@ pub async fn confirm_action( // Forward the confirmation to the agent's /confirm endpoint let confirm_url = format!( "{}/confirm/{}", - config.ai.agent_url.trim_end_matches('/'), + o2_cfg.ai.agent_url.trim_end_matches('/'), session_id ); diff --git a/src/handler/http/request/alerts/dedup_stats.rs b/src/handler/http/request/alerts/dedup_stats.rs index 22c5323c04..3dbe9638ad 100644 --- a/src/handler/http/request/alerts/dedup_stats.rs +++ b/src/handler/http/request/alerts/dedup_stats.rs @@ -70,7 +70,7 @@ pub async fn get_dedup_summary(Path(org_id): Path) -> Response { .filter(|alert| alert.deduplication.as_ref().is_some_and(|d| d.enabled)) .count(); - // Get current pending batches count from in-memory state (enterprise only) + // Get current pending batches count from in-memory state #[cfg(feature = "enterprise")] let pending_batches = crate::service::alerts::grouping::get_pending_batch_count(&org_id); diff --git a/src/handler/http/request/alerts/incidents.rs b/src/handler/http/request/alerts/incidents.rs index f8b41eb131..4103ec10c1 100644 --- a/src/handler/http/request/alerts/incidents.rs +++ b/src/handler/http/request/alerts/incidents.rs @@ -13,7 +13,7 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -//! HTTP handlers for alert incident management (Enterprise only) +//! HTTP handlers for alert incident management use axum::{ Json, @@ -400,14 +400,14 @@ pub async fn trigger_incident_rca( common::config::get_config as get_o2_config, }; - let o2_config = get_o2_config(); + let o2_cfg = get_o2_config(); // Check if RCA is enabled - if !o2_config.incidents.enabled || !o2_config.incidents.rca_enabled { + if !o2_cfg.incidents.enabled || !o2_cfg.incidents.rca_enabled { return MetaHttpResponse::bad_request("RCA is not enabled"); } - if o2_config.ai.agent_url.is_empty() { + if o2_cfg.ai.agent_url.is_empty() { return axum::response::Response::builder() .status(axum::http::StatusCode::SERVICE_UNAVAILABLE) .header(axum::http::header::CONTENT_TYPE, "application/json") @@ -419,7 +419,7 @@ pub async fn trigger_incident_rca( // In-flight guard { - let cooldown = o2_config.incidents.reanalysis_cooldown_minutes; + let cooldown = o2_cfg.incidents.reanalysis_cooldown_minutes; let events = infra::table::incident_events::get(&org_id, &incident_id) .await .unwrap_or_default(); @@ -481,7 +481,7 @@ pub async fn trigger_incident_rca( // Create RCA agent client let zo_config = config::get_config(); let client = match RcaAgentClient::new( - &o2_config.ai.agent_url, + &o2_cfg.ai.agent_url, &zo_config.auth.root_user_email, &zo_config.auth.root_user_password, ) { diff --git a/src/handler/http/request/patterns/mod.rs b/src/handler/http/request/patterns/mod.rs index f9b30f4bc2..1619d90ce1 100644 --- a/src/handler/http/request/patterns/mod.rs +++ b/src/handler/http/request/patterns/mod.rs @@ -179,9 +179,9 @@ pub async fn extract_patterns( // This determines the maximum number of logs to analyze for pattern detection // Default is 10K which aligns with industry standards for pattern quality if req.query.size == 0 || req.query.size == -1 { - let o2_config = get_o2_config(); - req.query.size = if o2_config.log_patterns.max_logs_for_extraction > 0 { - o2_config.log_patterns.max_logs_for_extraction as i64 + let o2_cfg = get_o2_config(); + req.query.size = if o2_cfg.log_patterns.max_logs_for_extraction > 0 { + o2_cfg.log_patterns.max_logs_for_extraction as i64 } else { config::get_config().limit.query_default_limit }; @@ -190,14 +190,14 @@ pub async fn extract_patterns( // Set default sampling ratio for pattern extraction if not explicitly provided // Pattern extraction benefits from sampling to improve performance on large datasets if req.query.sampling_ratio.is_none() { - let o2_config = get_o2_config(); - let default_ratio = o2_config.search_group.query_default_sampling_ratio / 100.0; + let o2_cfg = get_o2_config(); + let default_ratio = o2_cfg.common.query_default_sampling_ratio / 100.0; if default_ratio > 0.0 && default_ratio <= 1.0 { req.query.sampling_ratio = Some(default_ratio); log::info!( "[PATTERNS trace_id {}] Applied default sampling_ratio for pattern extraction: {}%", trace_id, - o2_config.search_group.query_default_sampling_ratio + o2_cfg.common.query_default_sampling_ratio ); } } diff --git a/src/handler/http/request/pipelines/backfill.rs b/src/handler/http/request/pipelines/backfill.rs index 09ce11816d..859ed3e961 100644 --- a/src/handler/http/request/pipelines/backfill.rs +++ b/src/handler/http/request/pipelines/backfill.rs @@ -13,8 +13,7 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -//! HTTP handlers for backfill job management (Enterprise only) - +//! HTTP handlers for backfill job management use axum::{ Json, extract::{Path, Query}, diff --git a/src/handler/http/request/traces/user.rs b/src/handler/http/request/traces/user.rs index 834af3a8ce..86d9219494 100644 --- a/src/handler/http/request/traces/user.rs +++ b/src/handler/http/request/traces/user.rs @@ -200,7 +200,7 @@ pub async fn get_latest_users( let stream_type = StreamType::Traces; let user_id_opt = Some(user_id.to_string()); - // Phase 1: Get paginated user list with trace_ids per user + // Get paginated user list with trace_ids per user let user_filter = if filter.is_empty() { format!("{user_id_col} IS NOT NULL AND {user_id_col} != ''") } else { diff --git a/src/infra/src/cluster/mod.rs b/src/infra/src/cluster/mod.rs index 66f0ba0126..af4d8ec9fb 100644 --- a/src/infra/src/cluster/mod.rs +++ b/src/infra/src/cluster/mod.rs @@ -1,4 +1,4 @@ -// Copyright 2025 OpenObserve Inc. +// Copyright 2026 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 @@ -13,7 +13,7 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::{collections::HashMap, ops::Bound, sync::Arc, time::Duration}; +use std::{ops::Bound, sync::Arc, time::Duration}; use config::{ RwAHashMap, RwBTreeMap, @@ -26,6 +26,7 @@ use config::{ sysinfo::{NodeMetrics, get_node_metrics}, }, }; +use hashbrown::{HashMap, HashSet}; use once_cell::sync::Lazy; use crate::{ @@ -121,41 +122,97 @@ pub async fn get_node_from_consistent_hash( None } -pub async fn print_consistent_hash() -> HashMap>> { - let mut map = HashMap::new(); - let r = QUERIER_INTERACTIVE_CONSISTENT_HASH.read().await; - let mut node_map = HashMap::new(); - for (k, v) in r.iter() { - let entry = node_map.entry(v.clone()).or_insert(Vec::new()); - entry.push(*k); +/// Like [`get_node_from_consistent_hash`] but restricts the ring walk to nodes +/// whose names are in `allowed`. +/// +/// Use this when a query runs on a subset of all queriers (e.g. +/// `max_nodes_per_query=3` with `strategy=org/stream`). Walking the global +/// ring but skipping nodes outside the allowed set preserves all consistent- +/// hashing properties (minimal remapping on membership change, even +/// distribution via vnodes) while confining dispatch to the selected subset. +pub async fn get_node_from_consistent_hash_within( + key: &str, + role: &Role, + group: Option, + allowed: &HashSet, +) -> Option { + let hash = config::utils::hash::gxhash::new().sum64(key); + let nodes = match role { + Role::Querier => match group { + Some(RoleGroup::Interactive) => QUERIER_INTERACTIVE_CONSISTENT_HASH.read().await, + Some(RoleGroup::Background) => QUERIER_BACKGROUND_CONSISTENT_HASH.read().await, + _ => QUERIER_INTERACTIVE_CONSISTENT_HASH.read().await, + }, + Role::Compactor => COMPACTOR_CONSISTENT_HASH.read().await, + Role::FlattenCompactor => FLATTEN_COMPACTOR_CONSISTENT_HASH.read().await, + _ => return None, + }; + if nodes.is_empty() { + return None; } - drop(r); - map.insert("querier_interactive".to_string(), node_map); - let r = QUERIER_BACKGROUND_CONSISTENT_HASH.read().await; - let mut node_map = HashMap::new(); - for (k, v) in r.iter() { - let entry = node_map.entry(v.clone()).or_insert(Vec::new()); - entry.push(*k); + + // Walk clockwise from the hash position, stop at the first node in `allowed`. + let mut iter = nodes.lower_bound(Bound::Included(&hash)); + while let Some((_, name)) = iter.next() { + if allowed.contains(name) { + return Some(name.clone()); + } } - drop(r); - map.insert("querier_background".to_string(), node_map); - let r = COMPACTOR_CONSISTENT_HASH.read().await; - let mut node_map = HashMap::new(); - for (k, v) in r.iter() { - let entry = node_map.entry(v.clone()).or_insert(Vec::new()); - entry.push(*k); + // Wrap around to the beginning of the ring. + for (_, name) in nodes.iter() { + if allowed.contains(name) { + return Some(name.clone()); + } } - drop(r); - map.insert("compactor".to_string(), node_map); - let r = FLATTEN_COMPACTOR_CONSISTENT_HASH.read().await; - let mut node_map = HashMap::new(); - for (k, v) in r.iter() { - let entry = node_map.entry(v.clone()).or_insert(Vec::new()); - entry.push(*k); + None +} + +/// Returns up to `n` distinct node names starting from the position of `key` +/// on the consistent hash ring, wrapping around if necessary. +/// +/// Useful for selecting a stable subset of nodes for a given key (e.g. org_id +/// or stream_key) without pinning to a single node. +pub async fn get_nodes_from_consistent_hash( + key: &str, + role: &Role, + group: Option, + n: usize, +) -> HashSet { + if n == 0 { + return HashSet::new(); } - drop(r); - map.insert("flatten_compactor".to_string(), node_map); - map + let hash = config::utils::hash::gxhash::new().sum64(key); + let nodes = match role { + Role::Querier => match group { + Some(RoleGroup::Interactive) => QUERIER_INTERACTIVE_CONSISTENT_HASH.read().await, + Some(RoleGroup::Background) => QUERIER_BACKGROUND_CONSISTENT_HASH.read().await, + _ => QUERIER_INTERACTIVE_CONSISTENT_HASH.read().await, + }, + Role::Compactor => COMPACTOR_CONSISTENT_HASH.read().await, + Role::FlattenCompactor => FLATTEN_COMPACTOR_CONSISTENT_HASH.read().await, + _ => return HashSet::new(), + }; + if nodes.is_empty() { + return HashSet::new(); + } + + // Start from the hash position and walk clockwise until we have `n` distinct nodes. + let mut result = HashSet::with_capacity(n); + let mut iter = nodes.lower_bound(Bound::Included(&hash)); + while let Some((_, name)) = iter.next() { + result.insert(name.to_string()); + if result.len() >= n { + return result; + } + } + // Wrap around to the beginning of the ring. + for (_, name) in nodes.iter() { + result.insert(name.to_string()); + if result.len() >= n { + break; + } + } + result } pub async fn count_consistent_hash() -> HashMap { @@ -644,4 +701,107 @@ mod tests { async fn test_list_nodes() { assert!(list_nodes().await.unwrap().is_empty()); } + + #[tokio::test] + async fn test_cluster_consistent_hash() { + // Reset the global state. + reset_consistent_hash().await; + + // Test consistent hash logic. + let node = LOCAL_NODE.clone(); + for i in 0..10 { + let node_q = Node { + name: format!("node-q-{i}").to_string(), + role: [Role::Querier].to_vec(), + ..node.clone() + }; + let node_c = Node { + name: format!("node-c-{i}").to_string(), + role: [Role::Compactor].to_vec(), + ..node.clone() + }; + add_node_to_consistent_hash(&node_q, &Role::Querier, Some(RoleGroup::Interactive)) + .await; + add_node_to_consistent_hash(&node_q, &Role::Querier, Some(RoleGroup::Background)).await; + add_node_to_consistent_hash(&node_q, &Role::Querier, None).await; + add_node_to_consistent_hash(&node_c, &Role::Compactor, None).await; + add_node_to_consistent_hash(&node_c, &Role::FlattenCompactor, None).await; + } + + for key in ["test", "test1", "test2", "test3", "test4", "test5", "test6"] { + println!( + "{key}-q: {}", + get_node_from_consistent_hash(key, &Role::Querier, None) + .await + .unwrap() + ); + println!( + "{key}-c: {}", + get_node_from_consistent_hash(key, &Role::Compactor, None) + .await + .unwrap() + ); + } + + // gxhash hash + let data = [ + ["test", "node-q-2", "node-c-7"], + ["test1", "node-q-3", "node-c-0"], + ["test2", "node-q-6", "node-c-5"], + ["test3", "node-q-9", "node-c-9"], + ["test4", "node-q-2", "node-c-0"], + ["test5", "node-q-5", "node-c-8"], + ["test6", "node-q-5", "node-c-8"], + ]; + + remove_node_from_consistent_hash(&node, &Role::Querier, Some(RoleGroup::Interactive)).await; + remove_node_from_consistent_hash(&node, &Role::Querier, Some(RoleGroup::Background)).await; + remove_node_from_consistent_hash(&node, &Role::Compactor, None).await; + remove_node_from_consistent_hash(&node, &Role::FlattenCompactor, None).await; + for key in data { + assert_eq!( + get_node_from_consistent_hash(key.first().unwrap(), &Role::Querier, None).await, + Some(key.get(1).unwrap().to_string()) + ); + assert_eq!( + get_node_from_consistent_hash(key.first().unwrap(), &Role::Compactor, None).await, + Some(key.get(2).unwrap().to_string()) + ); + } + } + + #[tokio::test] + async fn test_get_nodes_from_consistent_hash() { + // Reset the global state. + reset_consistent_hash().await; + + // Test consistent hash logic. + let node = LOCAL_NODE.clone(); + for i in 0..10 { + let node_q = Node { + name: format!("node-q-{i}").to_string(), + role: [Role::Querier].to_vec(), + ..node.clone() + }; + let node_c = Node { + name: format!("node-c-{i}").to_string(), + role: [Role::Compactor].to_vec(), + ..node.clone() + }; + add_node_to_consistent_hash(&node_q, &Role::Querier, Some(RoleGroup::Interactive)) + .await; + add_node_to_consistent_hash(&node_q, &Role::Querier, Some(RoleGroup::Background)).await; + add_node_to_consistent_hash(&node_q, &Role::Querier, None).await; + add_node_to_consistent_hash(&node_c, &Role::Compactor, None).await; + add_node_to_consistent_hash(&node_c, &Role::FlattenCompactor, None).await; + } + + let nodes = get_nodes_from_consistent_hash("test", &Role::Querier, None, 3).await; + + // sort the nodes by name + let mut nodes: Vec = nodes.into_iter().collect(); + nodes.sort(); + assert_eq!(nodes.len(), 3); + assert_eq!(nodes, vec!["node-q-2", "node-q-5", "node-q-8"]); + } } diff --git a/src/infra/src/file_list/postgres.rs b/src/infra/src/file_list/postgres.rs index a9bd72bd66..3d87b9e9ce 100644 --- a/src/infra/src/file_list/postgres.rs +++ b/src/infra/src/file_list/postgres.rs @@ -2809,7 +2809,7 @@ CREATE TABLE IF NOT EXISTS stream_stats ) .await?; - // Phase 1: Autovacuum tuning + column width compatibility + // Autovacuum tuning + column width compatibility if cfg.common.meta_partition_mode == "auto" { apply_autovacuum_tuning(&pool).await?; apply_column_width_compat(&pool).await?; diff --git a/src/infra/src/schema/mod.rs b/src/infra/src/schema/mod.rs index 1c7debcecf..6f974d3008 100644 --- a/src/infra/src/schema/mod.rs +++ b/src/infra/src/schema/mod.rs @@ -88,6 +88,13 @@ pub fn mk_key(org_id: &str, stream_type: StreamType, stream_name: &str) -> Strin format!("{SCHEMA_KEY}{org_id}/{stream_type}/{stream_name}") } +pub async fn exists(org_id: &str, stream_type: StreamType, stream_name: &str) -> bool { + let Ok(schema) = get_cache(org_id, stream_name, stream_type).await else { + return false; + }; + !schema.is_empty() +} + pub async fn get(org_id: &str, stream_name: &str, stream_type: StreamType) -> Result { let schema = get_cache(org_id, stream_name, stream_type).await?; Ok(schema.schema().as_ref().clone()) @@ -796,6 +803,10 @@ impl SchemaCache { size += std::mem::size_of::() + self.hash_key.len(); size } + + pub fn is_empty(&self) -> bool { + self.fields_map.is_empty() + } } impl MemorySize for SchemaCache { diff --git a/src/job/mod.rs b/src/job/mod.rs index c358ba4e0a..319d66a9c8 100644 --- a/src/job/mod.rs +++ b/src/job/mod.rs @@ -14,11 +14,12 @@ // along with this program. If not, see . use config::{cluster::LOCAL_NODE, spawn_pausable_job}; -#[cfg(feature = "enterprise")] -use o2_enterprise::enterprise::common::config::get_config as get_enterprise_config; -#[cfg(feature = "enterprise")] -use o2_openfga::config::get_config as get_openfga_config; use regex::Regex; +#[cfg(feature = "enterprise")] +use { + o2_enterprise::enterprise::{common::config::get_config as get_o2_config, search::admission}, + o2_openfga::config::get_config as get_openfga_config, +}; use crate::{ common::meta::{ @@ -143,6 +144,7 @@ pub async fn init() -> Result<(), anyhow::Error> { .expect("Email regex is valid"); let cfg = config::get_config(); + // init root user if !db::user::root_user_exists().await { if cfg.auth.root_user_email.is_empty() @@ -379,14 +381,26 @@ pub async fn init() -> Result<(), anyhow::Error> { } config_watcher::run(); + #[cfg(feature = "enterprise")] - if LOCAL_NODE.is_querier() && get_enterprise_config().ai.enabled { + if LOCAL_NODE.is_querier() && get_o2_config().ai.enabled { tokio::task::spawn(async move { o2_enterprise::enterprise::ai::agent::prompt::prompts::load_system_prompt() .await .expect("load system prompt failed"); }); } + + // Initialize slot-based admission ledger on querier nodes + #[cfg(feature = "enterprise")] + if LOCAL_NODE.is_querier() && get_o2_config().work_group.max_nodes_per_query > 0 { + admission::init_slot_ledger(cfg.limit.real_cpu_num as f64, cfg.limit.mem_total as f64); + // Run the TTL sweep at ≤ 1/4 of the reservation TTL so that expired + // reservations are reaped promptly without spinning. + let sweep_ms = (get_o2_config().work_group.slot_reserved_ttl_ms / 4).max(100); + admission::ledger::spawn_ttl_cleanup_task(sweep_ms); + } + tokio::task::spawn(files::run()); tokio::task::spawn(stats::run()); tokio::task::spawn(compactor::run()); @@ -532,11 +546,11 @@ pub async fn init() -> Result<(), anyhow::Error> { #[cfg(feature = "enterprise")] spawn_pausable_job!( "service_streams_batch_processor", - get_enterprise_config().service_streams.batch_flush_interval_secs, + get_o2_config().service_streams.batch_flush_interval_secs, { o2_enterprise::enterprise::service_streams::batch_processor::run_once().await; }, - pause_if: !get_enterprise_config().service_streams.enabled + pause_if: !get_o2_config().service_streams.enabled ); #[cfg(feature = "enterprise")] spawn_pausable_job!( @@ -603,7 +617,7 @@ pub async fn init() -> Result<(), anyhow::Error> { } } }, - pause_if: !get_enterprise_config().service_streams.enabled + pause_if: !get_o2_config().service_streams.enabled ); #[cfg(feature = "enterprise")] tokio::task::spawn(pipeline::run()); @@ -680,12 +694,15 @@ pub async fn init() -> Result<(), anyhow::Error> { pub async fn init_deferred() -> Result<(), anyhow::Error> { #[cfg(feature = "enterprise")] { - o2_enterprise::enterprise::license::start_license_check( - crate::service::self_reporting::search::get_usage, - LOCAL_NODE.is_router() && LOCAL_NODE.is_single_role(), - ) - .await; - tokio::task::spawn(db::license::watch()); + // only querier nodes watch license keys + if LOCAL_NODE.is_querier() { + o2_enterprise::enterprise::license::start_license_check( + crate::service::self_reporting::search::get_usage, + LOCAL_NODE.is_single_role(), + ) + .await; + tokio::task::spawn(db::license::watch()); + } } if !LOCAL_NODE.is_ingester() && !LOCAL_NODE.is_querier() && !LOCAL_NODE.is_alert_manager() { diff --git a/src/proto/proto/cluster/search.proto b/src/proto/proto/cluster/search.proto index d3275f73d7..0658b0bbea 100644 --- a/src/proto/proto/cluster/search.proto +++ b/src/proto/proto/cluster/search.proto @@ -19,6 +19,10 @@ service Search { rpc DeleteResult(DeleteResultRequest) returns (DeleteResultResponse) {} rpc GetTable(GetTableRequest) returns (GetTableResponse) {} rpc GetSourcemapFile(GetSourcemapFileRequest) returns (GetSourcemapFileResponse) {} + // Slot-based admission RPCs + rpc TryAcquire(TryAcquireRequest) returns (TryAcquireResponse) {} + rpc StartQuery(StartQueryRequest) returns (StartQueryResponse) {} + rpc ReleaseQuery(ReleaseQueryRequest) returns (ReleaseQueryResponse) {} } message QueryStatusRequest {} @@ -128,14 +132,10 @@ message QueryStatus { optional string stream_type = 7; optional Query query = 8; optional ScanStats scan_stats = 9; - optional WorkGroup work_group = 10; + reserved 10; // was WorkGroup enum (varint); now reserved to avoid rolling-upgrade wire-type clash optional string search_type = 11; optional SearchEventContext search_event_context = 12; -} - -enum WorkGroup { - SHORT = 0; - LONG = 1; + optional string work_group = 13; } message Query { @@ -160,4 +160,33 @@ message GetSourcemapFileRequest { message GetSourcemapFileResponse { bytes file_data = 1; -} \ No newline at end of file +} + +// Slot-based admission messages +message TryAcquireRequest { + string trace_id = 1; + string work_group = 2; // "short" | "long" | "background" + uint64 ttl_ms = 3; +} + +message TryAcquireResponse { + bool success = 1; + string reason = 2; +} + +message StartQueryRequest { + string trace_id = 1; +} + +message StartQueryResponse { + bool success = 1; + string reason = 2; +} + +message ReleaseQueryRequest { + string trace_id = 1; +} + +message ReleaseQueryResponse { + bool success = 1; +} diff --git a/src/proto/src/generated/cluster.rs b/src/proto/src/generated/cluster.rs index 36fb8c9c51..2a29187e67 100644 --- a/src/proto/src/generated/cluster.rs +++ b/src/proto/src/generated/cluster.rs @@ -1125,12 +1125,12 @@ pub struct QueryStatus { pub query: ::core::option::Option, #[prost(message, optional, tag = "9")] pub scan_stats: ::core::option::Option, - #[prost(enumeration = "WorkGroup", optional, tag = "10")] - pub work_group: ::core::option::Option, #[prost(string, optional, tag = "11")] pub search_type: ::core::option::Option<::prost::alloc::string::String>, #[prost(message, optional, tag = "12")] pub search_event_context: ::core::option::Option, + #[prost(string, optional, tag = "13")] + pub work_group: ::core::option::Option<::prost::alloc::string::String>, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct Query { @@ -1165,31 +1165,45 @@ pub struct GetSourcemapFileResponse { #[prost(bytes = "vec", tag = "1")] pub file_data: ::prost::alloc::vec::Vec, } -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] -#[repr(i32)] -pub enum WorkGroup { - Short = 0, - Long = 1, +/// Slot-based admission messages +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct TryAcquireRequest { + #[prost(string, tag = "1")] + pub trace_id: ::prost::alloc::string::String, + /// "short" | "long" | "background" + #[prost(string, tag = "2")] + pub work_group: ::prost::alloc::string::String, + #[prost(uint64, tag = "3")] + pub ttl_ms: u64, } -impl WorkGroup { - /// String value of the enum field names used in the ProtoBuf definition. - /// - /// The values are not transformed in any way and thus are considered stable - /// (if the ProtoBuf definition does not change) and safe for programmatic use. - pub fn as_str_name(&self) -> &'static str { - match self { - Self::Short => "SHORT", - Self::Long => "LONG", - } - } - /// Creates an enum from field names used in the ProtoBuf definition. - pub fn from_str_name(value: &str) -> ::core::option::Option { - match value { - "SHORT" => Some(Self::Short), - "LONG" => Some(Self::Long), - _ => None, - } - } +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct TryAcquireResponse { + #[prost(bool, tag = "1")] + pub success: bool, + #[prost(string, tag = "2")] + pub reason: ::prost::alloc::string::String, +} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct StartQueryRequest { + #[prost(string, tag = "1")] + pub trace_id: ::prost::alloc::string::String, +} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct StartQueryResponse { + #[prost(bool, tag = "1")] + pub success: bool, + #[prost(string, tag = "2")] + pub reason: ::prost::alloc::string::String, +} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ReleaseQueryRequest { + #[prost(string, tag = "1")] + pub trace_id: ::prost::alloc::string::String, +} +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ReleaseQueryResponse { + #[prost(bool, tag = "1")] + pub success: bool, } /// Generated client implementations. pub mod search_client { @@ -1507,6 +1521,77 @@ pub mod search_client { .insert(GrpcMethod::new("cluster.Search", "GetSourcemapFile")); self.inner.unary(req, path, codec).await } + /// Slot-based admission RPCs + pub async fn try_acquire( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::unknown( + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic_prost::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/cluster.Search/TryAcquire", + ); + let mut req = request.into_request(); + req.extensions_mut().insert(GrpcMethod::new("cluster.Search", "TryAcquire")); + self.inner.unary(req, path, codec).await + } + pub async fn start_query( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::unknown( + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic_prost::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/cluster.Search/StartQuery", + ); + let mut req = request.into_request(); + req.extensions_mut().insert(GrpcMethod::new("cluster.Search", "StartQuery")); + self.inner.unary(req, path, codec).await + } + pub async fn release_query( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::unknown( + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic_prost::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/cluster.Search/ReleaseQuery", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert(GrpcMethod::new("cluster.Search", "ReleaseQuery")); + self.inner.unary(req, path, codec).await + } } } /// Generated server implementations. @@ -1586,6 +1671,28 @@ pub mod search_server { tonic::Response, tonic::Status, >; + /// Slot-based admission RPCs + async fn try_acquire( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn start_query( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; + async fn release_query( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; } #[derive(Debug)] pub struct SearchServer { @@ -2105,6 +2212,137 @@ pub mod search_server { }; Box::pin(fut) } + "/cluster.Search/TryAcquire" => { + #[allow(non_camel_case_types)] + struct TryAcquireSvc(pub Arc); + impl tonic::server::UnaryService + for TryAcquireSvc { + type Response = super::TryAcquireResponse; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::try_acquire(&inner, request).await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let method = TryAcquireSvc(inner); + let codec = tonic_prost::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.unary(method, req).await; + Ok(res) + }; + Box::pin(fut) + } + "/cluster.Search/StartQuery" => { + #[allow(non_camel_case_types)] + struct StartQuerySvc(pub Arc); + impl tonic::server::UnaryService + for StartQuerySvc { + type Response = super::StartQueryResponse; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::start_query(&inner, request).await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let method = StartQuerySvc(inner); + let codec = tonic_prost::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.unary(method, req).await; + Ok(res) + }; + Box::pin(fut) + } + "/cluster.Search/ReleaseQuery" => { + #[allow(non_camel_case_types)] + struct ReleaseQuerySvc(pub Arc); + impl< + T: Search, + > tonic::server::UnaryService + for ReleaseQuerySvc { + type Response = super::ReleaseQueryResponse; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::release_query(&inner, request).await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let method = ReleaseQuerySvc(inner); + let codec = tonic_prost::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.unary(method, req).await; + Ok(res) + }; + Box::pin(fut) + } _ => { Box::pin(async move { let mut response = http::Response::new( diff --git a/src/service/alerts/deduplication.rs b/src/service/alerts/deduplication.rs index 41ed642b5e..9b390f8cef 100644 --- a/src/service/alerts/deduplication.rs +++ b/src/service/alerts/deduplication.rs @@ -1,4 +1,4 @@ -// Copyright 2025 OpenObserve Inc. +// Copyright 2026 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 @@ -13,7 +13,7 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -//! Alert deduplication logic (Enterprise only) +//! Alert deduplication logic //! //! This module provides the deduplication interface that delegates to the enterprise //! implementation. It is only compiled when the enterprise feature is enabled. diff --git a/src/service/alerts/incidents.rs b/src/service/alerts/incidents.rs index de023bf828..c7a145b397 100644 --- a/src/service/alerts/incidents.rs +++ b/src/service/alerts/incidents.rs @@ -1,4 +1,4 @@ -// Copyright 2025 OpenObserve Inc. +// Copyright 2026 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 @@ -13,7 +13,7 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -//! Alert Incident Correlation Service (Enterprise only) +//! Alert Incident Correlation Service //! //! Correlates fired alerts into unified incidents to reduce alert fatigue. diff --git a/src/service/db/file_list/broadcast.rs b/src/service/db/file_list/broadcast.rs index b714c4a233..040aeee703 100644 --- a/src/service/db/file_list/broadcast.rs +++ b/src/service/db/file_list/broadcast.rs @@ -1,4 +1,4 @@ -// Copyright 2025 OpenObserve Inc. +// Copyright 2026 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 @@ -50,6 +50,20 @@ pub async fn send(items: &[FileKey]) -> Result<(), anyhow::Error> { }) .await .unwrap_or_default(); + + // Enterprise: when max_nodes_per_query > 0 the query side uses a filtered + // ring walk (get_node_from_consistent_hash_within) to dispatch files only + // to the selected subset of nodes. Pre-compute per-item target nodes here + // using the same logic so broadcast and dispatch always agree on which node + // owns each file. When max_nodes_per_query == 0 (strategy=all) this + // returns None and the code falls through to the existing global-ring path. + #[cfg(feature = "enterprise")] + let item_slot_targets = + o2_enterprise::enterprise::search::admission::node_selection::broadcast_item_targets( + items, &nodes, + ) + .await; + for node in nodes { if node.uuid.eq(&LOCAL_NODE.uuid) { continue; @@ -61,6 +75,25 @@ pub async fn send(items: &[FileKey]) -> Result<(), anyhow::Error> { node_items.push(item.clone()); continue; } + + // Enterprise: use pre-computed selection-aware targets so that + // broadcast routes each file to exactly the node that the filtered + // ring walk would pick at query time. + #[cfg(feature = "enterprise")] + { + if let Some(ref targets) = item_slot_targets + && let Some((interactive, background)) = targets.get(&item.id) + { + if interactive.as_deref() == Some(node.name.as_str()) + || background.as_deref() == Some(node.name.as_str()) + { + node_items.push(item.clone()); + } + continue; + } + } + + // OSS / strategy=all: use the global ring unchanged. // check if the item is for interactive node if let Some(node_name) = cluster::get_node_from_consistent_hash( &item.id.to_string(), diff --git a/src/service/promql/search/mod.rs b/src/service/promql/search/mod.rs index 64a8a4c9bc..cdaed64368 100644 --- a/src/service/promql/search/mod.rs +++ b/src/service/promql/search/mod.rs @@ -21,7 +21,7 @@ use std::{ use config::{ get_config, meta::{ - cluster::RoleGroup, + cluster::{Node, RoleGroup}, promql::value::*, search::ScanStats, self_reporting::usage::{RequestStats, UsageType}, @@ -36,7 +36,6 @@ use futures::future::try_join_all; use hashbrown::HashMap; use infra::{ client::grpc::make_grpc_metrics_client, - cluster::get_cached_online_querier_nodes, errors::{Error, ErrorCodes, Result}, runtime::DATAFUSION_RUNTIME, }; @@ -44,8 +43,9 @@ use proto::cluster_rpc; use tracing::{Instrument, info_span}; #[cfg(feature = "enterprise")] use { - crate::service::search::SEARCH_SERVER, o2_enterprise::enterprise::search::TaskStatus, - o2_enterprise::enterprise::search::WorkGroup, + crate::service::search::SEARCH_SERVER, + o2_enterprise::enterprise::common::config::get_config as get_o2_config, + o2_enterprise::enterprise::search::{TaskStatus, WorkGroup}, }; use crate::service::{ @@ -82,7 +82,7 @@ pub async fn search( ) -> Result { #[cfg(feature = "enterprise")] { - let sql = Some(req.query.clone()); + let query = Some(req.query.clone()); let start_time = Some(req.start); let end_time = Some(req.end); @@ -95,7 +95,7 @@ pub async fn search( Some(user_email.to_string()), Some(org_id.to_string()), Some(StreamType::Metrics.to_string()), - sql, + query, start_time, end_time, None, @@ -118,6 +118,31 @@ pub async fn search( let mut stop_watch = TookWatcher::new(); + // get querier nodes from cluster. + // Use "metrics/" (stream-type sentinel) rather than a bare "metrics" string + // so that select_nodes with strategy=stream receives a non-empty, + // deterministic key. A single PromQL request may cover multiple metric + // names, so we cannot provide a more specific stream name here. + let nodes = crate::service::search::cluster::flight::get_online_querier_nodes( + trace_id, + &req.org_id, + "metrics/", + Some(RoleGroup::Interactive), + ) + .await?; + let nodes: Vec<_> = nodes.into_iter().filter(|node| node.is_querier()).collect(); + if nodes.is_empty() { + log::error!("no querier node online"); + return Err(Error::ErrorCode(ErrorCodes::ServerInternalError( + "no querier node found".to_string(), + ))); + } + + log::info!( + "[trace_id {trace_id}] promql->search: get nodes num: {}", + nodes.len(), + ); + // 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( @@ -144,6 +169,41 @@ pub async fn search( .await .map_err(|e| Error::ErrorCode(ErrorCodes::ServerInternalError(e.to_string())))?; + let mut _took_wait = _lock.took_wait; + + // Node-level slot admission for PromQL (enterprise) + // Placed here alongside the work_group check so both admission stages are co-located. + // The guard is held for the duration of search_in_cluster via the select! below. + // PromQL always runs as "short". Nodes are cached so the extra call is cheap. + #[cfg(feature = "enterprise")] + let _node_slot_guard = { + let wg = &get_o2_config().work_group; + if wg.max_nodes_per_query > 0 { + let start_time = std::time::Instant::now(); + let elapsed_ms = stop_watch.total_millis(); + let remaining_ms = (timeout * 1000).saturating_sub(elapsed_ms); + let guard = Some( + o2_enterprise::enterprise::search::admission::acquire_node_slots( + trace_id, + &nodes, + &WorkGroup::Short, + remaining_ms, + ) + .await + .map_err(|e| Error::ErrorCode(ErrorCodes::ServerInternalError(e.to_string())))?, + ); + let took_wait = start_time.elapsed().as_millis() as usize; + _took_wait += took_wait; + log::info!( + "[trace_id {trace_id}] promql->search: wait in node slot queue took: {took_wait} ms" + ); + guard + } else { + None + } + }; + log::info!("[trace_id {trace_id}] promql->search: total wait in queue took: {_took_wait} ms"); + #[cfg(feature = "enterprise")] let (abort_sender, abort_receiver) = tokio::sync::oneshot::channel(); #[cfg(feature = "enterprise")] @@ -162,8 +222,9 @@ pub async fn search( let trace_id_move = trace_id.to_string(); let user_email_move = user_email.to_string(); - let query_task = DATAFUSION_RUNTIME - .spawn(async move { search_in_cluster(&trace_id_move, req, &user_email_move).await }); + let query_task = DATAFUSION_RUNTIME.spawn(async move { + search_in_cluster(&trace_id_move, req, &user_email_move, &nodes).await + }); tokio::pin!(query_task); // run search in cluster @@ -213,6 +274,7 @@ async fn search_in_cluster( trace_id: &str, req: cluster_rpc::MetricsQueryRequest, user_email: &str, + nodes: &[Node], ) -> Result { let start_ins = std::time::Instant::now(); let started_at = now_micros(); @@ -228,6 +290,7 @@ async fn search_in_cluster( query_data: _, label_selector: _, } = req.query.as_ref().unwrap(); + let nr_queriers = nodes.len() as i64; // cache enabled if result cache is enabled and use_cache is true and start != end let use_cache = cfg.common.result_cache_enabled && req.use_cache && start != end; @@ -244,22 +307,6 @@ async fn search_in_cluster( query, ); - // get querier nodes from cluster - let mut nodes = get_cached_online_querier_nodes(Some(RoleGroup::Interactive)) - .await - .unwrap(); - // sort nodes by node_id this will improve hit cache ratio - nodes.sort_by(|a, b| a.grpc_addr.cmp(&b.grpc_addr)); - nodes.dedup_by(|a, b| a.grpc_addr == b.grpc_addr); - nodes.sort_by_key(|x| x.id); - let nodes = nodes; - if nodes.is_empty() { - return Err(Error::ErrorCode(ErrorCodes::ServerInternalError( - "no querier node found".to_string(), - ))); - } - let nr_queriers = nodes.len() as i64; - // The number of resolution steps; see the diagram at // https://promlabs.com/blog/2020/06/18/the-anatomy-of-a-promql-query/#range-queries let partition_step = max(micros(DEFAULT_LOOKBACK) * 2, step); diff --git a/src/service/search/cache/mod.rs b/src/service/search/cache/mod.rs index 9d04fab6a8..4f6b81c56b 100644 --- a/src/service/search/cache/mod.rs +++ b/src/service/search/cache/mod.rs @@ -1205,9 +1205,6 @@ pub async fn apply_regex_to_response( ctx: &str, ) -> Result<(), infra::errors::Error> { if res.hits.is_empty() { - log::info!( - "[trace_id {trace_id}] SDR patterns application : context {ctx} , took 0 ms, processed 0 hits", - ); return Ok(()); } @@ -1231,6 +1228,9 @@ pub async fn apply_regex_to_response( sql, ) .await?; + if projections.is_empty() { + return Ok(()); + } let ret = match pattern_manager.process_at_search( org_id, @@ -1240,13 +1240,15 @@ pub async fn apply_regex_to_response( ) { Ok(_) => Ok(()), Err(e) => { - log::error!("error in processing records for patterns for stream {all_streams} : {e}"); + log::error!( + "[trace_id {trace_id}] SDR patterns application: error in processing records for stream: {all_streams}: {e}" + ); Err(infra::errors::Error::Message(e.to_string())) } }; let took = start.elapsed().as_millis(); log::info!( - "[trace_id {trace_id}] SDR patterns application : context : {ctx} , took {took} ms, processed {} hits", + "[trace_id {trace_id}] SDR patterns application: context: {ctx}, took: {took} ms, processed: {} hits", res.hits.len() ); ret diff --git a/src/service/search/cluster/flight.rs b/src/service/search/cluster/flight.rs index a63dba8d25..84a5c8958a 100644 --- a/src/service/search/cluster/flight.rs +++ b/src/service/search/cluster/flight.rs @@ -44,9 +44,13 @@ use itertools::Itertools; use parking_lot::Mutex; use tracing::{Instrument, info_span}; use tracing_opentelemetry::OpenTelemetrySpanExt; - #[cfg(feature = "enterprise")] -use crate::service::search::SEARCH_SERVER; +use { + crate::service::search::SEARCH_SERVER, + o2_enterprise::enterprise::common::config::get_config as get_o2_config, + o2_enterprise::enterprise::search::{WorkGroup, admission}, +}; + use crate::{ handler::grpc::flight::visitor::get_peak_memory_from_ctx, service::{ @@ -138,6 +142,10 @@ 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 0 { + let start_time = std::time::Instant::now(); + let elapsed_ms = took_watch.total_millis(); + let remaining_ms = (timeout * 1000).saturating_sub(elapsed_ms); + let guard = Some( + admission::acquire_node_slots(trace_id, &nodes, &work_group, remaining_ms).await?, + ); + let took_wait = start_time.elapsed().as_millis() as usize; + _took_wait += took_wait; + log::info!( + "[trace_id {trace_id}] flight->search: wait in node slot queue took: {took_wait} ms" + ); + guard + } else { + None + } + }; + log::info!("[trace_id {trace_id}] flight->search: total wait in queue took: {_took_wait} ms"); + + // Move metric updates to here, after all admission stages complete. + // A request blocked on slot budget should not be counted as "running". metrics::QUERY_PENDING_NUMS .with_label_values(&[&sql.org_id]) .dec(); @@ -336,7 +372,7 @@ pub async fn search(trace_id: &str, sql: Arc, mut req: Request) -> Result, ) -> Result> { // get nodes from cluster @@ -456,14 +494,27 @@ pub async fn get_online_querier_nodes( let querier_num = nodes.iter().filter(|node| node.is_querier()).count(); if querier_num == 0 { - log::error!("no querier node online"); + log::error!("[trace_id {trace_id}] flight->search: no querier node online"); return Err(Error::Message("no querier node online".to_string())); } // use enterprise scheduler to filter nodes #[cfg(feature = "enterprise")] { + // filter nodes by cpu nodes = o2_enterprise::enterprise::search::scheduler::filter_nodes_by_cpu(nodes); + // filter nodes by slot admission + nodes = o2_enterprise::enterprise::search::admission::node_selection::select_nodes( + _org_id, + _stream_key, + nodes, + role_group, + ) + .await; + log::debug!( + "[trace_id {trace_id}] flight->search: slot admission select_nodes for org: {_org_id}, nodes: {:?}", + nodes.iter().map(|node| node.name.clone()).collect_vec() + ); } Ok(nodes) @@ -581,9 +632,25 @@ pub(crate) async fn partition_file_by_hash( idx += 1; } let mut partitions = vec![Vec::new(); idx]; + if idx == 0 { + return partitions; + } + + // Build the allowed set from the selected querier nodes. When all online + // queriers are selected this degrades to the normal global-ring behaviour. + // When only a subset is selected (e.g. max_nodes_per_query=3, strategy=org) + // the ring walk skips nodes outside the subset, so every file maps to one + // of the selected nodes without any fallback-to-0 imbalance. + let allowed: HashSet = node_idx.keys().map(|s| s.to_string()).collect(); + for fk in file_id_list { - let node_name = - cluster::get_node_from_consistent_hash(&fk.id.to_string(), &Role::Querier, group).await; + let node_name = cluster::get_node_from_consistent_hash_within( + &fk.id.to_string(), + &Role::Querier, + group, + &allowed, + ) + .await; let idx = match node_name { Some(node_name) => match node_idx.get(&node_name) { Some(idx) => *idx, diff --git a/src/service/search/datafusion/exec.rs b/src/service/search/datafusion/exec.rs index 1f79493d6a..7c54079f91 100644 --- a/src/service/search/datafusion/exec.rs +++ b/src/service/search/datafusion/exec.rs @@ -48,7 +48,7 @@ use datafusion::{ prelude::{SessionContext, col}, }; #[cfg(feature = "enterprise")] -use o2_enterprise::enterprise::{common::config::get_config as get_o2_config, search::WorkGroup}; +use o2_enterprise::enterprise::search::WorkGroup; #[cfg(all(feature = "enterprise", feature = "vortex"))] use { vortex::{VortexSessionDefault, io::session::RuntimeSessionExt, session::VortexSession}, @@ -65,7 +65,6 @@ use crate::service::search::{ }; pub(crate) const DATAFUSION_MIN_MEM: usize = 1024 * 1024 * 256; // 256MB -pub(crate) const DATAFUSION_MIN_PARTITION: usize = 2; // CPU cores pub fn create_session_config( sorted_by_time: bool, @@ -77,10 +76,7 @@ pub fn create_session_config( } else { target_partitions }; - let target_partitions = max( - DATAFUSION_MIN_PARTITION, - max(cfg.limit.datafusion_min_partition_num, target_partitions), - ); + let target_partitions = max(cfg.limit.datafusion_min_partition_num, target_partitions); let mut config = SessionConfig::from_env()? .with_batch_size(get_batch_size()) .with_target_partitions(target_partitions) @@ -397,10 +393,7 @@ impl TableBuilder { } else { session.target_partitions }; - let target_partitions = max( - DATAFUSION_MIN_PARTITION, - max(cfg.limit.datafusion_min_partition_num, target_partitions), - ); + let target_partitions = max(cfg.limit.datafusion_min_partition_num, target_partitions); #[cfg(feature = "enterprise")] let (target_partitions, _) = get_cpu_and_mem_limit( @@ -569,29 +562,24 @@ impl TableBuilder { async fn get_cpu_and_mem_limit( trace_id: &str, work_group: Option, - mut target_partitions: usize, - mut memory_size: usize, + target_partitions: usize, + memory_size: usize, ) -> Result<(usize, usize)> { - if let Some(wg) = work_group.as_ref() + let (target_partitions, memory_size) = if let Some(wg) = work_group.as_ref() && let Ok(wg) = WorkGroup::from_str(wg) { - let (cpu, mem) = wg.get_dynamic_resource().await.map_err(|e| { - DataFusionError::Execution(format!("Failed to get dynamic resource: {e}")) - })?; - if get_o2_config().search_group.cpu_limit_enabled { - target_partitions = std::cmp::max( - get_config().limit.datafusion_min_partition_num, - target_partitions * cpu as usize / 100, - ); - } - memory_size = memory_size * mem as usize / 100; - } - - let target_partitions = if target_partitions == 0 { - DATAFUSION_MIN_PARTITION + wg.get_resource(trace_id, target_partitions, memory_size) + .await + .map_err(|e| { + DataFusionError::Execution(format!("Failed to get dynamic resource: {e}")) + })? } else { - target_partitions + (target_partitions, memory_size) }; + let target_partitions = std::cmp::max( + get_config().limit.datafusion_min_partition_num, + target_partitions, + ); log::info!( "[trace_id: {trace_id}] work_group: {work_group:?}, target_partitions: {target_partitions}, memory_size: {memory_size}" @@ -627,7 +615,6 @@ mod tests { get_config() .limit .cpu_num - .max(DATAFUSION_MIN_PARTITION) .max(get_config().limit.datafusion_min_partition_num) ); assert_eq!(config.options().execution.batch_size, get_batch_size()); @@ -644,11 +631,8 @@ mod tests { let config = create_session_config(true, target_partitions)?; let expected_partitions = std::cmp::max( - DATAFUSION_MIN_PARTITION, - std::cmp::max( - get_config().limit.datafusion_min_partition_num, - target_partitions, - ), + get_config().limit.datafusion_min_partition_num, + target_partitions, ); assert_eq!( @@ -910,7 +894,7 @@ mod tests { let config = create_session_config(false, 1)?; // Very small number let actual_partitions = config.options().execution.target_partitions; - assert!(actual_partitions >= DATAFUSION_MIN_PARTITION); + assert!(actual_partitions >= get_config().limit.datafusion_min_partition_num); Ok(()) } diff --git a/src/service/search/datafusion/merge/downsampling.rs b/src/service/search/datafusion/merge/downsampling.rs index 7162cbf456..610a9514fb 100644 --- a/src/service/search/datafusion/merge/downsampling.rs +++ b/src/service/search/datafusion/merge/downsampling.rs @@ -44,7 +44,7 @@ use vortex::{ }; use crate::service::search::datafusion::{ - exec::{DATAFUSION_MIN_PARTITION, DataFusionContextBuilder}, + exec::DataFusionContextBuilder, merge::{MergeParquetResult, append_metadata}, table_provider::uniontable::NewUnionTable, }; @@ -73,7 +73,7 @@ pub async fn merge_parquet_files_with_downsampling( // create datafusion context let ctx = DataFusionContextBuilder::new() .sorted_by_time(true) - .build(DATAFUSION_MIN_PARTITION) + .build(get_config().limit.datafusion_min_partition_num) .await?; // register union table let union_table = Arc::new(NewUnionTable::new(schema.clone(), tables)); diff --git a/src/service/search/datafusion/merge/mod.rs b/src/service/search/datafusion/merge/mod.rs index 9e65218f4c..47128e01f9 100644 --- a/src/service/search/datafusion/merge/mod.rs +++ b/src/service/search/datafusion/merge/mod.rs @@ -31,9 +31,7 @@ use futures::TryStreamExt; use parquet::{arrow::AsyncArrowWriter, file::metadata::KeyValue}; use super::table_provider::uniontable::NewUnionTable; -use crate::service::search::datafusion::exec::{ - DATAFUSION_MIN_PARTITION, DataFusionContextBuilder, -}; +use crate::service::search::datafusion::exec::DataFusionContextBuilder; #[cfg(feature = "enterprise")] pub mod downsampling; @@ -107,7 +105,7 @@ pub async fn merge_parquet_files( let ctx = DataFusionContextBuilder::new() .sorted_by_time(true) - .build(DATAFUSION_MIN_PARTITION) + .build(get_config().limit.datafusion_min_partition_num) .await?; // register union table let union_table = Arc::new(NewUnionTable::new(schema.clone(), tables)); diff --git a/src/service/search/mod.rs b/src/service/search/mod.rs index 8b7d2d0a49..2c5c2d3663 100644 --- a/src/service/search/mod.rs +++ b/src/service/search/mod.rs @@ -69,7 +69,7 @@ use { o2_enterprise::enterprise::{ common::config::get_config as get_o2_config, search::{ - TaskStatus, WorkGroup, + TaskStatus, cache::streaming_agg::{ create_aggregation_cache_file_path, discover_cache_for_query, generate_optimal_partitions, get_aggregation_cache_key_from_request, @@ -259,13 +259,9 @@ pub async fn search( { if let Some(status) = SEARCH_SERVER.remove(&trace_id, false).await && let Some((_, stat)) = status.first() + && let Some(wg) = stat.work_group.as_ref() { - match stat.work_group.as_ref() { - Some(WorkGroup::Short) => _work_group = Some("short".to_string()), - Some(WorkGroup::Long) => _work_group = Some("long".to_string()), - Some(WorkGroup::Background) => _work_group = Some("background".to_string()), - None => _work_group = None, - } + _work_group = Some(wg.to_string()); }; } @@ -832,13 +828,37 @@ pub async fn search_partition( return Ok(response); }; - let nodes = get_cached_online_querier_nodes(Some(RoleGroup::Interactive)) + let role_group = if is_http_req { + Some(RoleGroup::Interactive) + } else { + Some(RoleGroup::Background) + }; + let nodes = get_cached_online_querier_nodes(role_group) .await .unwrap_or_default(); if nodes.is_empty() { log::error!("[trace_id {trace_id}] search_partition: no querier node online"); return Err(Error::Message("no querier node online".to_string())); } + + // Use the configured node selection strategy so cpu_cores reflects the + // actual nodes this query will fan out to — not the entire cluster. + // With "org" or "stream" strategies only a subset of nodes are used; + // summing all nodes' CPU would under-estimate total_secs and create too + // few partitions. + #[cfg(feature = "enterprise")] + let cpu_cores = { + let stream_key = sql.get_first_stream_key(); + let selected = o2_enterprise::enterprise::search::admission::node_selection::select_nodes( + org_id, + &stream_key, + nodes, + role_group, + ) + .await; + selected.iter().map(|n| n.cpu_num).sum::() as usize + }; + #[cfg(not(feature = "enterprise"))] let cpu_cores = nodes.iter().map(|n| n.cpu_num).sum::() as usize; let (records, original_size) = files.iter().fold((0, 0), |(records, original_size), f| { @@ -1136,7 +1156,7 @@ pub async fn search_partition( resp.streaming_id = streaming_id.clone(); } - // Get cache strategy for streaming aggregates (enterprise only) + // Get cache strategy for streaming aggregates #[cfg(feature = "enterprise")] let stremaing_aggs_cache_strategy = if streaming_aggs && let Some(streaming_id_ref) = streaming_id.as_deref() @@ -1286,13 +1306,7 @@ pub async fn query_status() -> Result { } else { "processing" }; - let work_group = if result.work_group.is_none() { - "Unknown" - } else if *result.work_group.as_ref().unwrap() == 0 { - "Short" - } else { - "Long" - }; + let work_group = result.work_group.clone().unwrap_or_default(); let search_type: Option = result.search_type.map(|s_event_type| { search::SearchEventType::try_from(s_event_type.as_str()) .unwrap_or(search::SearchEventType::UI) @@ -1307,7 +1321,7 @@ pub async fn query_status() -> Result { stream_type: result.stream_type, query, scan_stats, - work_group: work_group.to_string(), + work_group, search_type, search_event_context: result.search_event_context.map(Into::into), }); diff --git a/src/service/search/partition.rs b/src/service/search/partition.rs index bd57fea9fc..46933b64c4 100644 --- a/src/service/search/partition.rs +++ b/src/service/search/partition.rs @@ -1,4 +1,4 @@ -// Copyright 2025 OpenObserve Inc. +// Copyright 2026 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 @@ -48,7 +48,7 @@ impl PartitionGenerator { /// * `order_by` - Sort order for partitions /// * `is_aggregate` - Whether this is an aggregate query /// * `add_mini_partition` - Whether to add a mini partition for faster initial results - /// * `cache_strategy` - Optional cache-aware partition strategy (enterprise only) + /// * `cache_strategy` - Optional cache-aware partition strategy /// /// # Returns /// Vector of [start, end] time ranges in microseconds @@ -65,7 +65,7 @@ impl PartitionGenerator { StreamingAggsPartitionStrategy, >, ) -> Vec<[i64; 2]> { - // If cache-aware strategy is provided, use it (enterprise only) + // If cache-aware strategy is provided, use it #[cfg(feature = "enterprise")] if let Some(strategy) = streaming_aggs_cache_strategy { return strategy.to_time_partitions(order_by); diff --git a/src/service/search/sql/mod.rs b/src/service/search/sql/mod.rs index 792f0e8eed..65ead518d7 100644 --- a/src/service/search/sql/mod.rs +++ b/src/service/search/sql/mod.rs @@ -106,6 +106,23 @@ impl Sql { Self::new_with_options(query, org_id, stream_type, search_event_type, false).await } + pub fn get_first_stream_key(&self) -> String { + self.stream_names + .first() + .map(|s| { + format!( + "{}/{}", + s.get_stream_type(self.stream_type), + s.stream_name() + ) + }) + // For multi-stream / cross-index queries there is no single stream + // name. Fall back to the stream-type prefix so select_nodes always + // receives a non-empty, deterministic key (rather than "" which + // would silently select all nodes and defeat org/stream affinity). + .unwrap_or_else(|| format!("{}/", self.stream_type)) + } + pub async fn new_with_options( query: &SearchQuery, org_id: &str, diff --git a/src/service/search/super_cluster/follower.rs b/src/service/search/super_cluster/follower.rs index eb7a35ee62..d10972de3c 100644 --- a/src/service/search/super_cluster/follower.rs +++ b/src/service/search/super_cluster/follower.rs @@ -111,6 +111,7 @@ pub async fn search( // get stream name let stream = TableReference::from(empty_exec.name()); let stream_type = stream.get_stream_type(req.stream_type); + let stream_key = format!("{}/{}", stream_type, stream.stream_name()); // 1. get file id list let file_id_list = get_file_id_lists( @@ -161,7 +162,8 @@ pub async fn search( .map(RoleGroup::from) }) .unwrap_or(Some(RoleGroup::Interactive)); - let mut nodes = get_online_querier_nodes(&trace_id, role_group).await?; + let mut nodes = + get_online_querier_nodes(&trace_id, &req.org_id, &stream_key, role_group).await?; // local mode, only use local node as querier node if req.local_mode.unwrap_or_default() { diff --git a/src/service/search/work_group.rs b/src/service/search/work_group.rs index 73c1fa9bd9..e43fe5e11d 100644 --- a/src/service/search/work_group.rs +++ b/src/service/search/work_group.rs @@ -32,10 +32,9 @@ use super::utils::AsyncDefer; /// 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, + pub work_group_str: String, _guard: AsyncDefer, } @@ -109,7 +108,6 @@ pub async fn check_work_group( ) -> 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 { @@ -194,8 +192,8 @@ pub async fn check_work_group( Ok(DeferredLock { took_wait, - work_group_str, work_group: Some(work_group), + work_group_str, _guard: guard, }) } @@ -287,9 +285,7 @@ async fn work_group_need_wait( // 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:?}[global:{}/{}, org:{}/{}, user:{}/{}]", - status.global_current, - status.global_limit, + "[trace_id {trace_id}] user: {user_id:?} is waiting in work_group {work_group:?}[org:{}/{}, user:{}/{}]", status.org_current, status.org_limit, status.user_current, @@ -303,9 +299,7 @@ async fn work_group_need_wait( // Got approval - slot available if log_wait { log::info!( - "[trace_id {trace_id}] user: {user_id:?} get approved in work_group {work_group:?}[global:{}/{}, org:{}/{}, user:{}/{}]", - status.global_current, - status.global_limit, + "[trace_id {trace_id}] user: {user_id:?} get approved in work_group {work_group:?}[org:{}/{}, user:{}/{}]", status.org_current, status.org_limit, status.user_current, @@ -326,8 +320,8 @@ async fn work_group_need_wait( /// /// 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) +/// - Predicting appropriate work group +/// - Adding work group to search server /// - Acquiring the work group lock #[cfg(not(feature = "enterprise"))] pub async fn acquire_work_group_lock( @@ -348,14 +342,6 @@ pub async fn acquire_work_group_lock( .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, @@ -365,7 +351,6 @@ pub async fn acquire_work_group_lock( nodes: &[Node], file_id_list_vec: &[&FileId], ) -> Result { - // Predict workgroup first let is_background_task = req .search_event_type .as_ref()