From 81ec3d8bec0b52ddd348bf9366dbd2a556ce9119 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Wed, 4 Nov 2026 12:19:36 +0000 Subject: [PATCH 1/8] feat: opt-in structured profile for distributed EXPLAIN ANALYZE (#16293) --- src/api/grpc/src/handler/grpc/flight/mod.rs | 87 +-- .../handler/grpc/flight/partition_encoder.rs | 10 +- .../grpc/src/handler/grpc/flight/stream.rs | 183 +++--- .../grpc/src/handler/grpc/flight/visitor.rs | 161 +---- src/config/src/datafusion/request.rs | 9 + src/config/src/meta/search.rs | 4 + src/flight/src/common.rs | 557 ++++++------------ src/flight/src/lib.rs | 1 + src/flight/src/profile.rs | 251 ++++++++ .../src/datafusion/distributed_plan/common.rs | 17 +- .../distributed_plan/decoder_stream.rs | 7 +- .../distribute_analyze_exec.rs | 282 ++++++--- .../src/datafusion/distributed_plan/mod.rs | 1 + .../datafusion/distributed_plan/profile.rs | 380 ++++++++++++ .../distributed_plan/remote_scan_exec.rs | 101 +--- src/search/src/datafusion/optimizer/mod.rs | 17 +- .../physical_optimizer/distribute_analyze.rs | 26 +- src/search/src/datafusion/plan_metrics.rs | 36 +- src/search/src/utils.rs | 6 +- src/search_service/src/cluster/flight.rs | 10 +- src/search_service/src/lib.rs | 1 + .../src/super_cluster/leader.rs | 2 +- 22 files changed, 1239 insertions(+), 910 deletions(-) create mode 100644 src/flight/src/profile.rs create mode 100644 src/search/src/datafusion/distributed_plan/profile.rs diff --git a/src/api/grpc/src/handler/grpc/flight/mod.rs b/src/api/grpc/src/handler/grpc/flight/mod.rs index c839b11435..43be4746d1 100644 --- a/src/api/grpc/src/handler/grpc/flight/mod.rs +++ b/src/api/grpc/src/handler/grpc/flight/mod.rs @@ -15,7 +15,6 @@ use std::{io::Cursor, sync::Arc}; -use arrow::ipc::{CompressionType, writer::IpcWriteOptions}; use arrow_flight::{ Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket, @@ -28,11 +27,14 @@ use datafusion::{ common::{DataFusionError, Result}, physical_plan::execute_stream_partitioned, }; -use flight::common::{MetricsInfo, PreCustomMessage}; +use flight::common::{MetricsInfo, NodeMessages}; use futures::{StreamExt, stream::BoxStream}; use futures_util::pin_mut; use prost::Message; -use search::inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}; +use search::{ + datafusion::distributed_plan::profile, + inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}, +}; use search_service::{grpc::flight as grpcFlight, work_group::DeferredLock}; use tonic::{Request, Response, Status, Streaming}; use tracing::Instrument; @@ -45,11 +47,8 @@ use { use crate::handler::grpc::{ MetadataMap, flight::{ - stream::{FlightEncoderStreamBuilder, MAX_FLIGHT_DATA_SIZE}, - visitor::{ - get_cluster_metrics, get_partial_err, get_peak_memory, get_peak_memory_from_ctx, - get_scan_stats, - }, + stream::FlightEncoderStreamBuilder, + visitor::{get_peak_memory_from_ctx, node_role}, }, }; @@ -164,31 +163,15 @@ impl FlightService for FlightServiceImpl { } let start = std::time::Instant::now(); - let write_options: IpcWriteOptions = IpcWriteOptions::default() - .try_with_compression(Some(CompressionType::ZSTD)) - .map_err(|e| { - log::error!( - "[trace_id {trace_id}] flight->search: do_get create IPC write options error: {e:?}", - ); - Status::internal(e.to_string()) - })?; - - // used for EXPLAIN ANALYZE to collect metrics after stream is done - let metrics = req.search_info.is_analyze.then_some(MetricsInfo { - plan: plan.clone(), - is_super_cluster, - func: Box::new(super_cluster_enabled), - }); - - // Get the peak memory usage from the memory pool - // Note: We get peak memory after stream execution, so we pass it via a shared reference - let peak_memory = get_peak_memory_from_ctx(&ctx); - - // used for super cluster follower leader to get information from follower node - let scan_stats_ref = get_scan_stats(&plan); - let metrics_ref = get_cluster_metrics(&plan); - let peak_memory_ref = get_peak_memory(&plan); - let partial_err_ref = get_partial_err(&plan); + let node_messages = NodeMessages { + scan_stats, + peak_memory: get_peak_memory_from_ctx(&ctx), + profile: req.search_info.is_analyze.then(|| MetricsInfo { + plan: plan.clone(), + build_profile: profile::build_profile, + }), + role: node_role(&plan), + }; // One stream per output partition so they encode in parallel let streams = execute_stream_partitioned(plan, ctx.task_ctx()).map_err(|e| { @@ -198,23 +181,12 @@ impl FlightService for FlightServiceImpl { Status::internal(e.to_string()) })?; - let mut stream = - FlightEncoderStreamBuilder::new(write_options, MAX_FLIGHT_DATA_SIZE, session_guard) - .with_trace_id(trace_id.to_string()) - .with_is_super(is_super_cluster) - .with_defer_lock(lock) - .with_start(start) - .with_custom_message(PreCustomMessage::ScanStats(scan_stats)) - .with_custom_message(PreCustomMessage::ScanStatsRef(scan_stats_ref)) - .with_custom_message(PreCustomMessage::Metrics(metrics)) - .with_custom_message(PreCustomMessage::MetricsRef(metrics_ref)) - .with_custom_message(PreCustomMessage::PeakMemoryRef(Some(peak_memory))) - .with_custom_message(PreCustomMessage::PeakMemoryRef(peak_memory_ref)) - .with_custom_message(PreCustomMessage::PartialErrRefEarly( - partial_err_ref.clone(), - )) - .with_custom_message(PreCustomMessage::PartialErrRef(partial_err_ref)) - .build(streams, span); + let mut stream = FlightEncoderStreamBuilder::new(session_guard, node_messages) + .with_trace_id(trace_id.to_string()) + .with_is_super(is_super_cluster) + .with_defer_lock(lock) + .with_start(start) + .build(streams, span); let stream = async_stream::stream! { let timeout = tokio::time::sleep(tokio::time::Duration::from_secs(timeout)); @@ -367,16 +339,6 @@ fn clear_session_data(trace_id: &str) { log::info!("Cleared session for trace_id: {trace_id}"); } -#[cfg(feature = "enterprise")] -fn super_cluster_enabled() -> bool { - get_o2_config().super_cluster.enabled -} - -#[cfg(not(feature = "enterprise"))] -fn super_cluster_enabled() -> bool { - false -} - #[cfg(test)] mod tests { use super::*; @@ -385,9 +347,4 @@ mod tests { fn test_flight_service_impl_default() { let _impl = FlightServiceImpl; } - - #[test] - fn test_super_cluster_disabled_in_oss() { - assert!(!super_cluster_enabled()); - } } diff --git a/src/api/grpc/src/handler/grpc/flight/partition_encoder.rs b/src/api/grpc/src/handler/grpc/flight/partition_encoder.rs index 3f6d8725a0..756357fa59 100644 --- a/src/api/grpc/src/handler/grpc/flight/partition_encoder.rs +++ b/src/api/grpc/src/handler/grpc/flight/partition_encoder.rs @@ -23,6 +23,8 @@ use datafusion::execution::SendableRecordBatchStream; use flight::encoder::{ChunkCoalescer, FlightDataEncoder}; use futures::{Stream, StreamExt}; +use super::stream::MAX_FLIGHT_DATA_SIZE; + /// Coalesces and ZSTD-encodes one datafusion output partition into FlightData. do_get runs one per /// partition on its own tokio task, so partitions execute and encode in parallel. /// @@ -39,15 +41,11 @@ pub(super) struct PartitionEncoderStream { } impl PartitionEncoderStream { - pub(super) fn new( - inner: SendableRecordBatchStream, - options: IpcWriteOptions, - max_flight_data_size: usize, - ) -> Self { + pub(super) fn new(inner: SendableRecordBatchStream, options: IpcWriteOptions) -> Self { Self { inner, coalescer: ChunkCoalescer::new(), - encoder: FlightDataEncoder::new(options, max_flight_data_size), + encoder: FlightDataEncoder::new(options, MAX_FLIGHT_DATA_SIZE), out: VecDeque::new(), inner_done: false, } diff --git a/src/api/grpc/src/handler/grpc/flight/stream.rs b/src/api/grpc/src/handler/grpc/flight/stream.rs index 519c19ccfe..75ea5368f0 100644 --- a/src/api/grpc/src/handler/grpc/flight/stream.rs +++ b/src/api/grpc/src/handler/grpc/flight/stream.rs @@ -15,11 +15,11 @@ use std::{collections::VecDeque, pin::Pin, task::Poll}; -use arrow::ipc::writer::IpcWriteOptions; +use arrow::ipc::{CompressionType, writer::IpcWriteOptions}; use arrow_flight::{FlightData, error::FlightError}; use config::metrics; use datafusion::execution::SendableRecordBatchStream; -use flight::{common::PreCustomMessage, encoder::FlightDataEncoder}; +use flight::{common::NodeMessages, encoder::FlightDataEncoder}; use futures::{Stream, StreamExt}; use search_service::work_group::DeferredLock; use tokio::{sync::mpsc, task::JoinSet}; @@ -30,12 +30,10 @@ use crate::handler::grpc::flight::{SessionGuard, partition_encoder::PartitionEnc /// Cap on one encoded FlightData message. Mirrors the `ZO_GRPC_MAX_MESSAGE_SIZE` default, /// which is what the receiving channel accepts. -pub(crate) const MAX_FLIGHT_DATA_SIZE: usize = 32 * 1024 * 1024; +pub(super) const MAX_FLIGHT_DATA_SIZE: usize = 32 * 1024 * 1024; pub struct FlightEncoderStreamBuilder { - options: IpcWriteOptions, - max_flight_data_size: usize, - custom_messages: Vec, + node_messages: NodeMessages, // query context trace_id: String, is_super: bool, @@ -45,15 +43,9 @@ pub struct FlightEncoderStreamBuilder { } impl FlightEncoderStreamBuilder { - pub(crate) fn new( - options: IpcWriteOptions, - max_flight_data_size: usize, - session_guard: SessionGuard, - ) -> Self { + pub(crate) fn new(session_guard: SessionGuard, node_messages: NodeMessages) -> Self { Self { - options, - max_flight_data_size, - custom_messages: vec![], + node_messages, trace_id: String::new(), is_super: false, defer_lock: None, @@ -62,11 +54,6 @@ impl FlightEncoderStreamBuilder { } } - pub fn with_custom_message(mut self, message: PreCustomMessage) -> Self { - self.custom_messages.push(message); - self - } - pub fn with_trace_id(mut self, trace_id: String) -> Self { self.trace_id = trace_id; self @@ -98,12 +85,12 @@ impl FlightEncoderStreamBuilder { ) -> FlightEncoderStream { let child_span = info_span!("grpc:search:flight:execute_physical_plan"); let _ = child_span.set_parent(span.context()); + let options = write_options(); let (source, tasks) = if partitions.len() == 1 { // one partition: no parallelism to gain, poll inline (no task, no channel) let inner = partitions.into_iter().next().unwrap(); - let stream = - PartitionEncoderStream::new(inner, self.options.clone(), self.max_flight_data_size); + let stream = PartitionEncoderStream::new(inner, options.clone()); (FlightSource::Inline(Box::new(stream)), JoinSet::new()) } else { // one task per partition, merged over the channel @@ -112,11 +99,7 @@ impl FlightEncoderStreamBuilder { let (tx, rx) = mpsc::channel::, FlightError>>(cap); for inner in partitions { let tx = tx.clone(); - let mut stream = PartitionEncoderStream::new( - inner, - self.options.clone(), - self.max_flight_data_size, - ); + let mut stream = PartitionEncoderStream::new(inner, options.clone()); tasks.spawn(async move { while let Some(item) = stream.next().await { let is_err = item.is_err(); @@ -134,9 +117,9 @@ impl FlightEncoderStreamBuilder { FlightEncoderStream { source, - encoder: FlightDataEncoder::new(self.options, self.max_flight_data_size), + encoder: FlightDataEncoder::new(options, MAX_FLIGHT_DATA_SIZE), queue: VecDeque::new(), - custom_messages: self.custom_messages, + node_messages: self.node_messages, first_data: true, done: false, ctx: StreamContext { @@ -167,7 +150,7 @@ pub struct FlightEncoderStream { /// Only used to encode custom messages (scan stats, peak memory, ...). encoder: FlightDataEncoder, queue: VecDeque, - custom_messages: Vec, + node_messages: NodeMessages, first_data: bool, done: bool, /// query context + resources consumed by the stream's `Drop` teardown. @@ -232,31 +215,20 @@ impl EventLog { } impl FlightEncoderStream { - fn encode_custom_messages(&mut self) -> Result<(), FlightError> { - let custom_messages = std::mem::take(&mut self.custom_messages); - for message in custom_messages.into_iter() { - if let Some(message) = message.get_custom_message() { - let flight_data = self.encoder.encode_custom(&message)?; - self.queue.push_back(flight_data); - } + /// Queues the node's messages due now; a stream with no batch sends its first ones at the end. + fn encode_node_messages(&mut self, at_end: bool) -> Result<(), FlightError> { + let mut messages = Vec::new(); + if self.first_data { + messages.extend(self.node_messages.first()); } - Ok(()) - } - - fn encode_early_emit_messages(&mut self) -> Result<(), FlightError> { - let custom_messages = std::mem::take(&mut self.custom_messages); - let mut remainder_messages = Vec::new(); - for message in custom_messages.into_iter() { - if message.is_early_emit() { - if let Some(message) = message.get_custom_message() { - let flight_data = self.encoder.encode_custom(&message)?; - self.queue.push_back(flight_data); - } - } else { - remainder_messages.push(message); - } + if at_end { + messages.extend(self.node_messages.last()); + } + self.first_data = false; + for message in messages { + let flight_data = self.encoder.encode_custom(&message)?; + self.queue.push_back(flight_data); } - self.custom_messages = remainder_messages; Ok(()) } @@ -318,14 +290,10 @@ impl Stream for FlightEncoderStream { // 4. pull the next encoded FlightData from the partition encoder(s) match this.poll_source(cx) { Poll::Ready(Some(Ok(group))) => { - // before the first data, emit early-emit custom messages - if this.first_data { - this.first_data = false; - if !this.custom_messages.is_empty() - && let Err(e) = this.encode_early_emit_messages() - { - return this.fail(e); - } + if this.first_data + && let Err(e) = this.encode_node_messages(false) + { + return this.fail(e); } // enqueue the whole group so the chunk's dictionary and records stay contiguous this.queue.extend(group); @@ -345,7 +313,7 @@ impl Stream for FlightEncoderStream { if !this.ctx.tasks.is_empty() { return Poll::Pending; } - if let Err(e) = this.encode_custom_messages() { + if let Err(e) = this.encode_node_messages(true) { return this.fail(e); } this.done = true; @@ -393,27 +361,30 @@ impl Drop for FlightEncoderStream { } } +fn write_options() -> IpcWriteOptions { + IpcWriteOptions::default() + .try_with_compression(Some(CompressionType::ZSTD)) + // Cannot fail: compression only needs IPC metadata V5, which is the default. + .expect("default IPC options accept compression") +} + #[cfg(test)] mod tests { use std::sync::Arc; - use arrow::{ - array::{Int32Array, RecordBatch}, - ipc::writer::IpcWriteOptions, - }; + use arrow::array::{Int32Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema}; use config::meta::search::ScanStats; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; - use flight::common::PreCustomMessage; + use flight::common::{CustomMessage, NodeRole}; use futures::TryStreamExt; use super::*; fn make_builder() -> FlightEncoderStreamBuilder { FlightEncoderStreamBuilder::new( - IpcWriteOptions::default(), - 4096, SessionGuard::new("test-trace".to_string()), + follower_messages(), ) } @@ -424,7 +395,6 @@ mod tests { assert!(!b.is_super); assert!(b.defer_lock.is_none()); assert!(b.session_guard.is_some()); - assert!(b.custom_messages.is_empty()); } #[test] @@ -458,11 +428,53 @@ mod tests { let _ = b.start; } - #[test] - fn test_builder_with_custom_message() { - let msg = PreCustomMessage::ScanStats(ScanStats::new()); - let b = make_builder().with_custom_message(msg); - assert_eq!(b.custom_messages.len(), 1); + fn follower_messages() -> NodeMessages { + NodeMessages { + scan_stats: ScanStats::new(), + peak_memory: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + profile: None, + role: NodeRole::Follower, + } + } + + async fn message_kinds(partition: SendableRecordBatchStream) -> Vec<&'static str> { + let fds: Vec = make_builder() + .build(vec![partition], tracing::Span::none()) + .try_collect() + .await + .unwrap(); + fds.iter() + .map( + |fd| match config::utils::json::from_slice(&fd.app_metadata) { + Ok(CustomMessage::ScanStats(_)) => "ScanStats", + Ok(CustomMessage::PeakMemory(_)) => "PeakMemory", + Ok(_) => "other", + Err(_) => "data", + }, + ) + .collect() + } + + #[tokio::test] + async fn test_node_messages_surround_the_batches() { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)])); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1, 2]))]) + .unwrap(); + let with_rows: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + schema.clone(), + futures::stream::iter(vec![Ok(batch)]), + )); + let kinds = message_kinds(with_rows).await; + assert_eq!(kinds.first(), Some(&"ScanStats")); + assert_eq!(kinds.last(), Some(&"PeakMemory")); + assert!(kinds[1..kinds.len() - 1].iter().all(|kind| *kind == "data")); + + let empty: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::empty(), + )); + assert_eq!(message_kinds(empty).await, ["ScanStats", "PeakMemory"]); } #[test] @@ -503,13 +515,9 @@ mod tests { make_partition(total / 2..total), ]; - let options = IpcWriteOptions::default() - .try_with_compression(Some(arrow::ipc::CompressionType::ZSTD)) - .unwrap(); let stream = FlightEncoderStreamBuilder::new( - options, - MAX_FLIGHT_DATA_SIZE, SessionGuard::new("test-trace".to_string()), + follower_messages(), ) .build(partitions, tracing::Span::none()); @@ -518,6 +526,7 @@ mod tests { let dicts = std::collections::HashMap::new(); let mut decoded: Vec = fds .iter() + .filter(|fd| fd.app_metadata.is_empty()) .flat_map(|fd| { let b = arrow_flight::utils::flight_data_to_arrow_batch(fd, schema.clone(), &dicts) .unwrap(); @@ -562,9 +571,8 @@ mod tests { )); let stream = FlightEncoderStreamBuilder::new( - IpcWriteOptions::default(), - MAX_FLIGHT_DATA_SIZE, SessionGuard::new("test-trace".to_string()), + follower_messages(), ) .build(vec![ok_partition, panic_partition], tracing::Span::none()); @@ -676,13 +684,9 @@ mod tests { let id_base_1 = n_per_partition; let partitions = vec![make_partition(0, 0), make_partition(1, id_base_1)]; - let options = IpcWriteOptions::default() - .try_with_compression(Some(arrow::ipc::CompressionType::ZSTD)) - .unwrap(); let stream = FlightEncoderStreamBuilder::new( - options, - MAX_FLIGHT_DATA_SIZE, SessionGuard::new("test-trace".to_string()), + follower_messages(), ) .build(partitions, tracing::Span::none()); let fds: Vec = stream.try_collect().await.unwrap(); @@ -749,13 +753,9 @@ mod tests { futures::stream::iter(batches.into_iter().map(Ok)), )); - let options = IpcWriteOptions::default() - .try_with_compression(Some(arrow::ipc::CompressionType::ZSTD)) - .unwrap(); let stream = FlightEncoderStreamBuilder::new( - options, - MAX_FLIGHT_DATA_SIZE, SessionGuard::new("test-trace".to_string()), + follower_messages(), ) .build(vec![inner], tracing::Span::none()); @@ -767,6 +767,7 @@ mod tests { let dicts = std::collections::HashMap::new(); let decoded: Vec = fds .iter() + .filter(|fd| fd.app_metadata.is_empty()) .flat_map(|fd| { let b = arrow_flight::utils::flight_data_to_arrow_batch(fd, schema.clone(), &dicts) .unwrap(); diff --git a/src/api/grpc/src/handler/grpc/flight/visitor.rs b/src/api/grpc/src/handler/grpc/flight/visitor.rs index 59567c423f..def7730613 100644 --- a/src/api/grpc/src/handler/grpc/flight/visitor.rs +++ b/src/api/grpc/src/handler/grpc/flight/visitor.rs @@ -13,163 +13,40 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::sync::{Arc, atomic::AtomicUsize}; +use std::sync::Arc; -use config::meta::search::ScanStats; use datafusion::{ - common::{ - Result, - tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}, - }, + common::tree_node::{TreeNode, TreeNodeRecursion}, physical_plan::ExecutionPlan, }; -use parking_lot::Mutex; +use flight::common::NodeRole; use search::datafusion::distributed_plan::remote_scan_exec::RemoteScanExec; -pub use search::datafusion::plan_metrics::{get_cluster_metrics, get_peak_memory_from_ctx}; +pub use search::datafusion::plan_metrics::get_peak_memory_from_ctx; -pub fn get_scan_stats(plan: &Arc) -> Option>> { - let mut visitor = RemoteScanVisitor::new(); - let _ = plan.visit(&mut visitor); - visitor.scan_stats -} - -struct RemoteScanVisitor { - scan_stats: Option>>, -} - -impl RemoteScanVisitor { - pub fn new() -> Self { - Self { scan_stats: None } - } -} - -impl<'n> TreeNodeVisitor<'n> for RemoteScanVisitor { - type Node = Arc; - - fn f_up(&mut self, node: &'n Self::Node) -> Result { - let name = node.name(); - if name == "RemoteScanExec" { - let remote_scan_exec = node.downcast_ref::().unwrap(); - self.scan_stats = Some(remote_scan_exec.scan_stats()); - Ok(TreeNodeRecursion::Stop) - } else { - Ok(TreeNodeRecursion::Continue) - } - } -} - -pub fn get_partial_err(plan: &Arc) -> Option>> { - let mut visitor = PartialErrVisitor::new(); - let _ = plan.visit(&mut visitor); - visitor.partial_err -} - -struct PartialErrVisitor { - partial_err: Option>>, -} - -impl PartialErrVisitor { - pub fn new() -> Self { - Self { partial_err: None } - } -} - -impl<'n> TreeNodeVisitor<'n> for PartialErrVisitor { - type Node = Arc; - - fn f_up(&mut self, node: &'n Self::Node) -> Result { - if node.name() == "RemoteScanExec" { - let remote_scan_exec = node.downcast_ref::().unwrap(); - self.partial_err = Some(remote_scan_exec.partial_err()); - Ok(TreeNodeRecursion::Stop) - } else { - Ok(TreeNodeRecursion::Continue) - } - } -} - -pub fn get_peak_memory(plan: &Arc) -> Option> { - let mut visitor = PeakMemoryVisitor::new(); - let _ = plan.visit(&mut visitor); - visitor.peak_memory -} - -struct PeakMemoryVisitor { - peak_memory: Option>, -} - -impl PeakMemoryVisitor { - pub fn new() -> Self { - Self { peak_memory: None } - } -} - -impl<'n> TreeNodeVisitor<'n> for PeakMemoryVisitor { - type Node = Arc; - - fn f_up(&mut self, node: &'n Self::Node) -> Result { - let name = node.name(); - if name == "RemoteScanExec" { - let remote_scan_exec = node.downcast_ref::().unwrap(); - self.peak_memory = Some(remote_scan_exec.peak_memory()); - } - Ok(TreeNodeRecursion::Continue) - } +/// Only a super cluster's cluster leader runs a plan with a remote scan in it. +pub fn node_role(plan: &Arc) -> NodeRole { + let mut role = NodeRole::Follower; + let _ = plan.apply(|node| { + let Some(remote_scan) = node.downcast_ref::() else { + return Ok(TreeNodeRecursion::Continue); + }; + role = NodeRole::ClusterLeader(remote_scan.state.clone()); + Ok(TreeNodeRecursion::Stop) + }); + role } #[cfg(test)] mod tests { - use std::sync::Arc; - use arrow::datatypes::{DataType, Field, Schema}; use datafusion::physical_plan::empty::EmptyExec; use super::*; - fn empty_plan() -> Arc { + #[test] + fn test_plan_without_remote_scan_is_follower() { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - Arc::new(EmptyExec::new(schema)) - } - - #[test] - fn test_remote_scan_visitor_new_has_no_stats() { - let v = RemoteScanVisitor::new(); - assert!(v.scan_stats.is_none()); - } - - #[test] - fn test_partial_err_visitor_new_has_no_err() { - let v = PartialErrVisitor::new(); - assert!(v.partial_err.is_none()); - } - - #[test] - fn test_peak_memory_visitor_new_has_none() { - let v = PeakMemoryVisitor::new(); - assert!(v.peak_memory.is_none()); - } - - #[test] - fn test_get_scan_stats_no_remote_exec_returns_none() { - let plan = empty_plan(); - assert!(get_scan_stats(&plan).is_none()); - } - - #[test] - fn test_get_cluster_metrics_no_remote_exec_returns_empty() { - let plan = empty_plan(); - assert!(get_cluster_metrics(&plan).is_empty()); - } - - #[test] - fn test_get_partial_err_no_remote_exec_returns_none() { - let plan = empty_plan(); - assert!(get_partial_err(&plan).is_none()); - } - - #[test] - fn test_get_peak_memory_no_remote_exec_returns_none() { - let plan = empty_plan(); - assert!(get_peak_memory(&plan).is_none()); + let plan: Arc = Arc::new(EmptyExec::new(schema)); + assert!(matches!(node_role(&plan), NodeRole::Follower)); } } diff --git a/src/config/src/datafusion/request.rs b/src/config/src/datafusion/request.rs index 64f5014722..4ebfd4dbf3 100644 --- a/src/config/src/datafusion/request.rs +++ b/src/config/src/datafusion/request.rs @@ -39,6 +39,8 @@ pub struct Request { /// request because the query is planned again in spawned tasks, where /// that task-local is gone. Never derived from client input. pub chat_history_reader: bool, + /// Adds the structured `profile` column to EXPLAIN ANALYZE output. + pub analyze_profile: bool, } impl Default for Request { @@ -59,6 +61,7 @@ impl Default for Request { overwrite_cache: false, histogram_interval: 0, chat_history_reader: false, + analyze_profile: false, } } } @@ -92,6 +95,7 @@ impl Request { overwrite_cache, histogram_interval, chat_history_reader: false, + analyze_profile: false, } } @@ -112,6 +116,10 @@ impl Request { self.use_cache = use_cache; } + pub fn set_analyze_profile(&mut self, analyze_profile: bool) { + self.analyze_profile = analyze_profile; + } + /// Mark the request as the AI Chat API's history reader when — and only /// when — it is issued from inside that reader's scope. pub fn mark_chat_history_reader(&mut self) { @@ -138,6 +146,7 @@ impl From for Request { histogram_interval: req.search_info.histogram_interval, // Never carried across nodes. chat_history_reader: false, + analyze_profile: false, } } } diff --git a/src/config/src/meta/search.rs b/src/config/src/meta/search.rs index 43ae73516f..f8cdbe2c4d 100644 --- a/src/config/src/meta/search.rs +++ b/src/config/src/meta/search.rs @@ -198,6 +198,9 @@ pub struct Query { /// that don't carry their own 3rd timezone argument. None / "UTC" / "" => UTC. #[serde(default, skip_serializing_if = "Option::is_none")] pub timezone: Option, + /// Adds the structured `profile` column to EXPLAIN ANALYZE output. + #[serde(default)] + pub analyze_profile: bool, } fn default_size() -> i64 { @@ -223,6 +226,7 @@ impl Default for Query { streaming_id: None, histogram_interval: 0, timezone: None, + analyze_profile: false, } } } diff --git a/src/flight/src/common.rs b/src/flight/src/common.rs index 245802aed7..abaed37d97 100644 --- a/src/flight/src/common.rs +++ b/src/flight/src/common.rs @@ -14,7 +14,7 @@ // along with this program. If not, see . use std::{ - fmt, + collections::BTreeMap, sync::{ Arc, atomic::{AtomicUsize, Ordering}, @@ -31,14 +31,15 @@ use config::{ use datafusion::{ self, physical_plan::{ - DisplayFormatType, ExecutionPlan, - display::DisplayableExecutionPlan, + ExecutionPlan, metrics::{self, ExecutionPlanMetricsSet, MetricBuilder, MetricValue}, }, }; use parking_lot::Mutex; use serde::{Deserialize, Serialize}; +use crate::profile::{ProfileNode, RemoteNodeProfile}; + #[derive(Debug, Clone)] pub struct RemoteScanMetrics { /// Time in nanos to execute child operator and fetch batches @@ -115,154 +116,117 @@ pub enum CustomMessage { #[derive(Default, Debug, Serialize, Deserialize, Clone)] pub struct Metrics { + /// Unused; kept, like `metrics`, so that older leaders can still parse the message. pub stage: i32, pub node_address: String, pub node_name: String, + /// Left empty: the leader renders the text plan from `profile`. pub metrics: String, + /// The plan this node ran; older nodes do not send it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub profile: Option, } -/// the CustomMessage is not ready, so we need to use this struct to store -/// the scan stats, and send it to the client when the scan stats is ready -pub enum PreCustomMessage { - ScanStats(ScanStats), - // use for super cluster follower leader - ScanStatsRef(Option>>), - // physical plan, is_super_cluster - Metrics(Option), - // use for super cluster follower leader - MetricsRef(Vec>>>), - // use for storing memory pool reference to extract peak later - PeakMemoryRef(Option>), - // sent before first batch (early signal for fast-failing nodes) - PartialErrRefEarly(Option>>), - // sent after all batches (final accumulated value) - PartialErrRef(Option>>), +/// What a node sends besides its record batches; the shared references are read when sent. +pub struct NodeMessages { + /// This node's own scan stats, known before execution. + pub scan_stats: ScanStats, + /// The peak of this node's memory pool. + pub peak_memory: Arc, + /// The plan to send back as a profile; set only for EXPLAIN ANALYZE. + pub profile: Option, + pub role: NodeRole, +} + +impl NodeMessages { + /// Sent before the first batch, so the caller sees scan stats and fast failures early. + pub fn first(&self) -> Vec { + let mut messages = vec![CustomMessage::ScanStats(self.reported_scan_stats())]; + messages.extend(self.partial_err()); + messages + } + + /// Sent after the last batch. + pub fn last(&self) -> Vec { + let mut messages = vec![CustomMessage::PeakMemory(self.reported_peak_memory())]; + messages.extend(self.partial_err()); + messages.extend( + self.profile + .as_ref() + .map(|info| CustomMessage::Metrics(collect_metrics(info))), + ); + messages + } + + /// Merged here rather than sent twice: the caller's add gives the same result either way. + fn reported_scan_stats(&self) -> ScanStats { + let mut stats = self.scan_stats; + if let NodeRole::ClusterLeader(remote) = &self.role { + stats.add(&remote.scan_stats.lock()); + } + stats + } + + fn reported_peak_memory(&self) -> usize { + let own = self.peak_memory.load(Ordering::Relaxed); + match &self.role { + NodeRole::Follower => own, + NodeRole::ClusterLeader(remote) => own.max(remote.peak_memory.load(Ordering::Relaxed)), + } + } + + /// Only a cluster leader has one: a follower's own failure ends its stream instead. + fn partial_err(&self) -> Option { + let NodeRole::ClusterLeader(remote) = &self.role else { + return None; + }; + let err = remote.partial_err.lock().clone(); + (!err.is_empty()).then_some(CustomMessage::PartialErr(err)) + } +} + +pub enum NodeRole { + /// Scans the data itself. + Follower, + /// A super cluster's cluster leader, which passes the query on to its own followers. + ClusterLeader(RemoteScanState), +} + +/// Per-partition profiles of one remote scan, keyed by partition. +pub type NodeProfiles = Arc>>; + +/// What a remote scan collects from the nodes it calls, shared by every copy of the scan. +#[derive(Debug, Clone, Default)] +pub struct RemoteScanState { + pub scan_stats: Arc>, + pub peak_memory: Arc, + pub partial_err: Arc>, + /// Filled only for EXPLAIN ANALYZE. + pub node_profiles: NodeProfiles, } pub struct MetricsInfo { pub plan: Arc, - pub is_super_cluster: bool, - pub func: Box bool + Send>, + /// Passed in because it lives in the search crate, which depends on this one. + pub build_profile: fn(&dyn ExecutionPlan) -> ProfileNode, } -impl PreCustomMessage { - pub fn is_early_emit(&self) -> bool { - matches!( - self, - PreCustomMessage::ScanStats(_) - | PreCustomMessage::ScanStatsRef(_) - | PreCustomMessage::PartialErrRefEarly(_) - ) - } - - pub fn get_custom_message(&self) -> Option { - match self { - PreCustomMessage::ScanStats(stats) => Some(CustomMessage::ScanStats(*stats)), - PreCustomMessage::ScanStatsRef(stats_ref) => stats_ref - .as_ref() - .map(|stats| CustomMessage::ScanStats(*stats.lock())), - PreCustomMessage::Metrics(metrics_ref) => metrics_ref - .as_ref() - .map(|metrics_info| CustomMessage::Metrics(collect_metrics(metrics_info))), - PreCustomMessage::MetricsRef(metrics) => { - let metrics: Vec = metrics.iter().flat_map(|m| m.lock().clone()).collect(); - (!metrics.is_empty()).then_some(CustomMessage::Metrics(metrics)) - } - PreCustomMessage::PeakMemoryRef(peak_memory_ref) => peak_memory_ref - .as_ref() - .map(|peak| CustomMessage::PeakMemory(peak.load(Ordering::Relaxed))), - PreCustomMessage::PartialErrRefEarly(ref_opt) - | PreCustomMessage::PartialErrRef(ref_opt) => ref_opt.as_ref().and_then(|r| { - let err = r.lock().clone(); - if err.is_empty() { - None - } else { - Some(CustomMessage::PartialErr(err)) - } - }), - } - } -} - -// collect metrics from physical plan fn collect_metrics(metrics_info: &MetricsInfo) -> Vec { - let plan = &metrics_info.plan; - let is_super_cluster = metrics_info.is_super_cluster; - let func = &metrics_info.func; - let plan_with_metrics = if is_super_cluster { - format_leader_plan(plan.as_ref()) - } else { - DisplayableExecutionPlan::with_metrics(plan.as_ref()) - .set_show_statistics(false) - .indent(false) - .to_string() - }; - let stage = if func() { - if is_super_cluster { 1 } else { 2 } - } else { - 1 - }; vec![Metrics { - stage, node_address: LOCAL_NODE.get_grpc_addr().to_string(), node_name: LOCAL_NODE.get_name().to_string(), - metrics: plan_with_metrics, + profile: Some((metrics_info.build_profile)(metrics_info.plan.as_ref())), + ..Default::default() }] } -fn format_leader_plan(plan: &dyn ExecutionPlan) -> String { - let mut output = String::new(); - write_plan(&mut output, plan, 0); - output -} - -struct PlanDisplay<'a>(&'a dyn ExecutionPlan); - -impl fmt::Display for PlanDisplay<'_> { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - self.0.fmt_as(DisplayFormatType::Default, f) - } -} - -fn write_plan(output: &mut String, plan: &dyn ExecutionPlan, indent: usize) { - use fmt::Write; - - write!(output, "{:indent$}", "", indent = indent * 2).expect("write to String"); - write!(output, "{}", PlanDisplay(plan)).expect("write to String"); - - if let Some(metrics) = plan.metrics() { - let metrics = metrics - .aggregate_by_name() - .sorted_for_display() - .timestamps_removed(); - write!(output, ", metrics=[{metrics}]").expect("write to String"); - } else { - write!(output, ", metrics=[]").expect("write to String"); - } - writeln!(output).expect("write to String"); - - if plan.name() == "RemoteScanExec" { - return; - } - - for child in plan.children() { - write_plan(output, child.as_ref(), indent + 1); - } -} - #[cfg(test)] mod tests { use std::sync::Arc; use arrow_schema::{DataType, Field, Schema}; use config::meta::search::ScanStats; - use datafusion::{ - common::Result, - execution::{SendableRecordBatchStream, TaskContext}, - physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, empty::EmptyExec, - }, - }; + use datafusion::physical_plan::{ExecutionPlan, empty::EmptyExec}; use parking_lot::Mutex; use super::*; @@ -285,70 +249,18 @@ mod tests { } } + fn test_profile(plan: &dyn ExecutionPlan) -> ProfileNode { + ProfileNode { + name: plan.name().to_string(), + ..Default::default() + } + } + fn create_test_execution_plan() -> Arc { let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); Arc::new(EmptyExec::new(schema)) } - #[derive(Debug)] - struct NamedExec { - name: &'static str, - inner: Arc, - } - - impl DisplayAs for NamedExec { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!(f, "{}", self.name) - } - } - - impl ExecutionPlan for NamedExec { - fn apply_expressions( - &self, - _f: &mut dyn FnMut( - &Arc, - ) -> datafusion::common::Result< - datafusion::common::tree_node::TreeNodeRecursion, - >, - ) -> datafusion::common::Result { - Ok(datafusion::common::tree_node::TreeNodeRecursion::Continue) - } - - fn name(&self) -> &'static str { - self.name - } - - fn properties(&self) -> &Arc { - self.inner.properties() - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.inner] - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - unimplemented!() - } - - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> Result { - unimplemented!() - } - } - - fn create_named_test_execution_plan( - name: &'static str, - inner: Arc, - ) -> Arc { - Arc::new(NamedExec { name, inner }) - } - #[test] fn test_metrics_default() { let metrics = Metrics::default(); @@ -358,225 +270,97 @@ mod tests { assert_eq!(metrics.metrics, ""); } - #[test] - fn test_pre_custom_message_is_early_emit() { - let scan_stats = create_test_scan_stats(); - - // Test ScanStats variant - let pre_msg = PreCustomMessage::ScanStats(scan_stats); - assert!(pre_msg.is_early_emit()); - - // Test ScanStatsRef variant with Some - let stats_ref = Arc::new(Mutex::new(scan_stats)); - let pre_msg = PreCustomMessage::ScanStatsRef(Some(stats_ref)); - assert!(pre_msg.is_early_emit()); - - // Test ScanStatsRef variant with None - let pre_msg = PreCustomMessage::ScanStatsRef(None); - assert!(pre_msg.is_early_emit()); - - // Test Metrics variant (should be false) - let pre_msg = PreCustomMessage::Metrics(None); - assert!(!pre_msg.is_early_emit()); - - // Test MetricsRef variant (should be false) - let pre_msg = PreCustomMessage::MetricsRef(vec![]); - assert!(!pre_msg.is_early_emit()); + fn describe(messages: Vec) -> Vec { + messages + .into_iter() + .map(|message| match message { + CustomMessage::ScanStats(stats) => format!("ScanStats({})", stats.files), + CustomMessage::PeakMemory(peak) => format!("PeakMemory({peak})"), + CustomMessage::PartialErr(err) => format!("PartialErr({err})"), + CustomMessage::Metrics(metrics) => { + format!("Metrics({})", metrics[0].profile.as_ref().unwrap().name) + } + }) + .collect() } #[test] - fn test_pre_custom_message_get_custom_message_scan_stats() { - let scan_stats = create_test_scan_stats(); - - let pre_msg = PreCustomMessage::ScanStats(scan_stats); - let custom_msg = pre_msg.get_custom_message().unwrap(); - - match custom_msg { - CustomMessage::ScanStats(stats) => { - assert_eq!(stats.files, 10); - assert_eq!(stats.records, 1000); - } - _ => panic!("Expected ScanStats variant"), - } - } - - #[test] - fn test_pre_custom_message_get_custom_message_scan_stats_ref() { - let scan_stats = create_test_scan_stats(); - let stats_ref = Arc::new(Mutex::new(scan_stats)); - - let pre_msg = PreCustomMessage::ScanStatsRef(Some(stats_ref)); - let custom_msg = pre_msg.get_custom_message().unwrap(); - - match custom_msg { - CustomMessage::ScanStats(stats) => { - assert_eq!(stats.files, 10); - assert_eq!(stats.records, 1000); - } - _ => panic!("Expected ScanStats variant"), - } - } - - #[test] - fn test_pre_custom_message_get_custom_message_scan_stats_ref_none() { - let pre_msg = PreCustomMessage::ScanStatsRef(None); - let custom_msg = pre_msg.get_custom_message(); - assert!(custom_msg.is_none()); - } - - #[test] - fn test_pre_custom_message_get_custom_message_metrics() { - let plan = create_test_execution_plan(); - let metrics_info = MetricsInfo { - plan, - is_super_cluster: false, - func: Box::new(|| true), + fn test_follower_sends_only_its_own_values() { + let messages = NodeMessages { + scan_stats: create_test_scan_stats(), + peak_memory: Arc::new(AtomicUsize::new(5)), + profile: None, + role: NodeRole::Follower, }; - - let pre_msg = PreCustomMessage::Metrics(Some(metrics_info)); - let custom_msg = pre_msg.get_custom_message().unwrap(); - - match custom_msg { - CustomMessage::Metrics(metrics) => { - assert_eq!(metrics.len(), 1); - assert_eq!(metrics[0].stage, 2); // func() returns true, is_super_cluster is false - assert!(!metrics[0].node_address.is_empty()); - assert!(!metrics[0].node_name.is_empty()); - assert!(!metrics[0].metrics.is_empty()); - } - _ => panic!("Expected Metrics variant"), - } + assert_eq!(describe(messages.first()), ["ScanStats(10)"]); + assert_eq!(describe(messages.last()), ["PeakMemory(5)"]); } #[test] - fn test_pre_custom_message_get_custom_message_metrics_none() { - let pre_msg = PreCustomMessage::Metrics(None); - let custom_msg = pre_msg.get_custom_message(); - assert!(custom_msg.is_none()); - } - - #[test] - fn test_pre_custom_message_get_custom_message_metrics_ref() { - let metrics1 = vec![Metrics { - stage: 1, - node_address: "node1".to_string(), - node_name: "Node 1".to_string(), - metrics: "metrics1".to_string(), - }]; - let metrics2 = vec![Metrics { - stage: 2, - node_address: "node2".to_string(), - node_name: "Node 2".to_string(), - metrics: "metrics2".to_string(), - }]; - - let metrics_refs = vec![ - Arc::new(Mutex::new(metrics1)), - Arc::new(Mutex::new(metrics2)), - ]; - - let pre_msg = PreCustomMessage::MetricsRef(metrics_refs); - let custom_msg = pre_msg.get_custom_message().unwrap(); - - match custom_msg { - CustomMessage::Metrics(metrics) => { - assert_eq!(metrics.len(), 2); - assert_eq!(metrics[0].node_address, "node1"); - assert_eq!(metrics[1].node_address, "node2"); - } - _ => panic!("Expected Metrics variant"), - } - } - - #[test] - fn test_pre_custom_message_get_custom_message_metrics_ref_empty() { - let metrics_refs = vec![Arc::new(Mutex::new(vec![]))]; - - let pre_msg = PreCustomMessage::MetricsRef(metrics_refs); - let custom_msg = pre_msg.get_custom_message(); - assert!(custom_msg.is_none()); - } - - #[test] - fn test_collect_metrics_super_cluster_true() { - let plan = create_test_execution_plan(); - let metrics_info = MetricsInfo { - plan, - is_super_cluster: true, - func: Box::new(|| true), + fn test_cluster_leader_merges_what_its_remote_scan_collected() { + let remote = RemoteScanState { + scan_stats: Arc::new(Mutex::new(ScanStats { + files: 3, + ..Default::default() + })), + peak_memory: Arc::new(AtomicUsize::new(7)), + ..Default::default() }; + let partial_err = remote.partial_err.clone(); + let messages = NodeMessages { + scan_stats: create_test_scan_stats(), + peak_memory: Arc::new(AtomicUsize::new(5)), + profile: Some(MetricsInfo { + plan: create_test_execution_plan(), + build_profile: test_profile, + }), + role: NodeRole::ClusterLeader(remote), + }; + assert_eq!(describe(messages.first()), ["ScanStats(13)"]); - let metrics = collect_metrics(&metrics_info); + *partial_err.lock() = "node down".to_string(); + assert_eq!( + describe(messages.first()), + ["ScanStats(13)", "PartialErr(node down)"] + ); + assert_eq!( + describe(messages.last()), + [ + "PeakMemory(7)", + "PartialErr(node down)", + "Metrics(EmptyExec)" + ] + ); + } + + #[test] + fn test_profile_carries_this_node() { + let metrics = collect_metrics(&MetricsInfo { + plan: create_test_execution_plan(), + build_profile: test_profile, + }); assert_eq!(metrics.len(), 1); - assert_eq!(metrics[0].stage, 1); // func() returns true, is_super_cluster is true assert!(!metrics[0].node_address.is_empty()); - assert!(!metrics[0].metrics.is_empty()); + assert!(!metrics[0].node_name.is_empty()); } #[test] - fn test_collect_metrics_super_cluster_false() { - let plan = create_test_execution_plan(); - let metrics_info = MetricsInfo { - plan, - is_super_cluster: false, - func: Box::new(|| true), - }; - - let metrics = collect_metrics(&metrics_info); - assert_eq!(metrics.len(), 1); - assert_eq!(metrics[0].stage, 2); // func() returns true, is_super_cluster is false + fn test_metrics_without_profile_deserializes() { + let json = r#"{"stage":1,"node_address":"a","node_name":"b","metrics":"m"}"#; + let metrics: Metrics = serde_json::from_str(json).unwrap(); + assert!(metrics.profile.is_none()); } #[test] - fn test_collect_metrics_func_returns_false() { - let plan = create_test_execution_plan(); - let metrics_info = MetricsInfo { - plan, - is_super_cluster: true, - func: Box::new(|| false), + fn test_metrics_keeps_fields_older_leaders_require() { + let metrics = Metrics { + profile: Some(ProfileNode::default()), + ..Default::default() }; - - let metrics = collect_metrics(&metrics_info); - assert_eq!(metrics.len(), 1); - assert_eq!(metrics[0].stage, 1); // func() returns false, stage defaults to 1 - } - - #[test] - fn test_super_cluster_stops_at_remote_scan() { - let follower_plan = - create_named_test_execution_plan("FollowerOnlyExec", create_test_execution_plan()); - let remote_scan_plan = create_named_test_execution_plan("RemoteScanExec", follower_plan); - let region_leader_plan = - create_named_test_execution_plan("SortPreservingMergeExec", remote_scan_plan); - let metrics_info = MetricsInfo { - plan: region_leader_plan, - is_super_cluster: true, - func: Box::new(|| true), - }; - - let metrics = collect_metrics(&metrics_info); - assert_eq!(metrics.len(), 1); - assert!(metrics[0].metrics.contains("SortPreservingMergeExec")); - assert!(metrics[0].metrics.contains("RemoteScanExec")); - assert!(!metrics[0].metrics.contains("FollowerOnlyExec")); - } - - #[test] - fn test_non_super_cluster_keeps_remote_scan_child_metrics() { - let follower_plan = - create_named_test_execution_plan("FollowerOnlyExec", create_test_execution_plan()); - let remote_scan_plan = create_named_test_execution_plan("RemoteScanExec", follower_plan); - let region_leader_plan = - create_named_test_execution_plan("SortPreservingMergeExec", remote_scan_plan); - let metrics_info = MetricsInfo { - plan: region_leader_plan, - is_super_cluster: false, - func: Box::new(|| true), - }; - - let metrics = collect_metrics(&metrics_info); - assert_eq!(metrics.len(), 1); - assert!(metrics[0].metrics.contains("FollowerOnlyExec, metrics=[]")); + let json = serde_json::to_value(&metrics).unwrap(); + for key in ["stage", "node_address", "node_name", "metrics"] { + assert!(json.get(key).is_some(), "missing {key}"); + } } #[test] @@ -603,6 +387,7 @@ mod tests { node_address: "test_node".to_string(), node_name: "Test Node".to_string(), metrics: "test_metrics".to_string(), + profile: None, }]; let custom_msg = CustomMessage::Metrics(metrics); diff --git a/src/flight/src/lib.rs b/src/flight/src/lib.rs index b0842e58f1..30fa19159d 100644 --- a/src/flight/src/lib.rs +++ b/src/flight/src/lib.rs @@ -16,3 +16,4 @@ pub mod common; pub mod decoder; pub mod encoder; +pub mod profile; diff --git a/src/flight/src/profile.rs b/src/flight/src/profile.rs new file mode 100644 index 0000000000..e6e6a8c9fb --- /dev/null +++ b/src/flight/src/profile.rs @@ -0,0 +1,251 @@ +// 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 +// 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 . + +//! Structured EXPLAIN ANALYZE profile: the executed plan with typed, per-node metrics. +//! +//! The profile is a tree that follows the leader's physical plan. Ordinary operators +//! recurse through `ProfileInputs::Children`. A `RemoteScanExec` stops the recursion +//! with `ProfileInputs::Remote` instead, because the subtree under it on the leader is +//! only the plan template sent to the remote nodes and never ran there. That list +//! holds one `RemoteNodeProfile` per node it called directly (one per partition), +//! including nodes skipped because they had no files. +//! +//! A `RemoteNodeProfile` has two parts: +//! - what the caller observed of that node: start and end time, rows, batches and bytes received +//! over Flight, decode time, and the node's scan stats; +//! - `plan`: the plan the node actually ran, with its own metrics, sent back by the node itself. +//! +//! A caller lists only the nodes it called directly; deeper nodes nest inside +//! `plan`. In a super cluster the cluster leader's plan has its own +//! `RemoteScanExec`, whose `remote` lists that cluster's queriers and ingesters: +//! +//! ```text +//! QueryProfile super leader +//! └─ plan … RemoteScanExec +//! └─ remote: [sc-a, sc-b, sc-c] cluster leaders +//! └─ plan … RemoteScanExec +//! └─ remote: [querier-0 … ingester-1] nodes of that cluster +//! └─ plan … +//! ``` +//! +//! A plan with several remote scans, such as a join of two streams, has one +//! `remote` list per scan, so the same node appears once under each of them. +//! +//! Each plan appears exactly once: a node sends back only its own plan, and its +//! caller stores it under the remote scan that made the call. +//! +//! The leader also renders the text rows of EXPLAIN ANALYZE from this tree, one per +//! node that sent back its plan, which is why `ProfileMetric::display` keeps each +//! value exactly as DataFusion prints it. + +use std::fmt::Write; + +use config::meta::search::ScanStats; +use serde::{Deserialize, Serialize}; + +/// The profile of one EXPLAIN ANALYZE run, as returned to the client. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct QueryProfile { + pub node_name: String, + pub node_address: String, + pub start_us: i64, + pub end_us: i64, + pub output_rows: usize, + pub plan: ProfileNode, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ProfileNode { + pub name: String, + pub detail: String, + pub partitions: usize, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub metrics: Vec, + #[serde(flatten)] + pub inputs: ProfileInputs, +} + +impl ProfileNode { + /// Renders the EXPLAIN ANALYZE text of the plan this node heads, stopping at remote scans. + pub fn to_text(&self) -> String { + let mut out = String::new(); + self.write_text(&mut out, 0); + out + } + + fn write_text(&self, out: &mut String, indent: usize) { + let metrics = self + .metrics + .iter() + .map(|m| format!("{}={}", m.name, m.display)) + .collect::>() + .join(", "); + writeln!( + out, + "{:indent$}{}, metrics=[{metrics}]", + "", + self.detail, + indent = indent * 2 + ) + .expect("write to String"); + if let ProfileInputs::Children(children) = &self.inputs { + for child in children { + child.write_text(out, indent + 1); + } + } + } +} + +/// Where a node's input came from; serialized as a `children` or a `remote` key. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ProfileInputs { + /// Operators that ran on this node; empty for a leaf. + Children(Vec), + /// Only on a remote scan: one entry per node it called directly. + Remote(Vec), +} + +impl Default for ProfileInputs { + fn default() -> Self { + Self::Children(Vec::new()) + } +} + +/// A metric summed across partitions. +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] +pub struct ProfileMetric { + pub name: String, + pub kind: MetricKind, + pub value: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub total: Option, + /// The value as DataFusion prints it, such as `1.20 MB` or `100 total → 20 matched`. + #[serde(default)] + pub display: String, +} + +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum MetricKind { + Count, + /// Nanoseconds. + Time, + Bytes, + Gauge, + Ratio, + Pruning, + #[default] + #[serde(other)] + Other, +} + +/// What the caller saw of one remote node, plus the plan that node ran. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct RemoteNodeProfile { + pub partition: usize, + pub node_name: String, + pub node_address: String, + pub role: String, + pub cluster: String, + pub start_us: i64, + pub end_us: i64, + pub output_rows: u64, + pub output_batches: u64, + pub output_bytes: u64, + pub decode_ns: u64, + pub skipped: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scan_stats: Option, + /// The plan this node ran; its own remote scans nest the nodes below it the same way. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub plan: Option, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_profile_json_round_trip_and_unknown_kind() { + let node = ProfileNode { + name: "FilterExec".to_string(), + metrics: vec![ProfileMetric { + name: "output_rows".to_string(), + kind: MetricKind::Count, + value: 1, + total: None, + display: "1".to_string(), + }], + ..Default::default() + }; + let json = serde_json::to_string(&node).unwrap(); + assert!(json.contains(r#""children":[]"#) && !json.contains("remote")); + let back: ProfileNode = serde_json::from_str(&json).unwrap(); + assert_eq!(back.metrics, node.metrics); + assert!(matches!(back.inputs, ProfileInputs::Children(c) if c.is_empty())); + let m: ProfileMetric = + serde_json::from_str(r#"{"name":"x","kind":"future_kind","value":1}"#).unwrap(); + assert_eq!(m.kind, MetricKind::Other); + } + + #[test] + fn test_remote_inputs_serialize_under_remote_key() { + let node = ProfileNode { + name: "RemoteScanExec".to_string(), + inputs: ProfileInputs::Remote(vec![RemoteNodeProfile { + node_name: "q-0".to_string(), + ..Default::default() + }]), + ..Default::default() + }; + let json = serde_json::to_string(&node).unwrap(); + assert!(json.contains(r#""remote":[{"#) && !json.contains("children")); + let back: ProfileNode = serde_json::from_str(&json).unwrap(); + let ProfileInputs::Remote(remote) = back.inputs else { + panic!("expected remote inputs"); + }; + assert_eq!(remote[0].node_name, "q-0"); + } + + #[test] + fn test_to_text_indents_children_and_stops_at_remote() { + let metric = |name: &str, display: &str| ProfileMetric { + name: name.to_string(), + display: display.to_string(), + ..Default::default() + }; + let node = ProfileNode { + detail: "SortExec: expr=[a@0 ASC]".to_string(), + metrics: vec![metric("output_rows", "3"), metric("spilled_bytes", "0.0 B")], + inputs: ProfileInputs::Children(vec![ProfileNode { + detail: "RemoteScanExec: partitions=2".to_string(), + inputs: ProfileInputs::Remote(vec![RemoteNodeProfile { + plan: Some(ProfileNode { + detail: "FilterExec".to_string(), + ..Default::default() + }), + ..Default::default() + }]), + ..Default::default() + }]), + ..Default::default() + }; + assert_eq!( + node.to_text(), + "SortExec: expr=[a@0 ASC], metrics=[output_rows=3, spilled_bytes=0.0 B]\n RemoteScanExec: partitions=2, metrics=[]\n" + ); + } +} diff --git a/src/search/src/datafusion/distributed_plan/common.rs b/src/search/src/datafusion/distributed_plan/common.rs index f3f4470207..d021b36f2b 100644 --- a/src/search/src/datafusion/distributed_plan/common.rs +++ b/src/search/src/datafusion/distributed_plan/common.rs @@ -22,9 +22,10 @@ use datafusion::{ common::Result, execution::SendableRecordBatchStream, physical_plan::stream::RecordBatchStreamAdapter, }; -use flight::common::Metrics; use parking_lot::Mutex; +use super::profile::NodeProfileRecorder; + #[derive(Debug)] pub struct QueryContext { pub trace_id: String, @@ -33,13 +34,13 @@ pub struct QueryContext { pub is_querier: bool, pub scan_stats: Arc>, pub partial_err: Arc>, - pub cluster_metrics: Arc>>, pub peak_memory: Arc, pub start: std::time::Instant, pub num_rows: usize, pub req_id: u64, pub req_last_time: std::time::Instant, pub print_key_event: bool, + pub node_profile: NodeProfileRecorder, } impl QueryContext { @@ -50,7 +51,6 @@ impl QueryContext { is_querier: false, scan_stats: Arc::new(Mutex::new(ScanStats::new())), partial_err: Arc::new(Mutex::new(String::new())), - cluster_metrics: Arc::new(Mutex::new(Vec::new())), peak_memory: Arc::new(AtomicUsize::new(0)), trace_id: String::new(), start: std::time::Instant::now(), @@ -58,6 +58,7 @@ impl QueryContext { req_id: 0, req_last_time: std::time::Instant::now(), print_key_event: config::get_config().common.print_key_event, + node_profile: NodeProfileRecorder::default(), } } @@ -86,11 +87,6 @@ impl QueryContext { self } - pub fn with_cluster_metrics(mut self, metrics: Arc>>) -> Self { - self.cluster_metrics = metrics; - self - } - pub fn with_start_time(mut self, start: std::time::Instant) -> Self { self.start = start; self @@ -100,6 +96,11 @@ impl QueryContext { self.peak_memory = peak_memory; self } + + pub fn with_node_profile(mut self, node_profile: NodeProfileRecorder) -> Self { + self.node_profile = node_profile; + self + } } /// EmptyStream is a stream that returns an empty record batch diff --git a/src/search/src/datafusion/distributed_plan/decoder_stream.rs b/src/search/src/datafusion/distributed_plan/decoder_stream.rs index 1aa382da1f..000759e828 100644 --- a/src/search/src/datafusion/distributed_plan/decoder_stream.rs +++ b/src/search/src/datafusion/distributed_plan/decoder_stream.rs @@ -67,7 +67,7 @@ impl FlightDecoderStream { self.query_context.scan_stats.lock().add(&stats); } CustomMessage::Metrics(metrics) => { - self.query_context.cluster_metrics.lock().extend(metrics); + self.query_context.node_profile.record_plan(metrics); } CustomMessage::PeakMemory(peak_memory) => { self.query_context @@ -191,6 +191,10 @@ impl Stream for FlightDecoderStream { impl Drop for FlightDecoderStream { fn drop(&mut self) { + let scan_stats = self.scan_stats; + self.query_context + .node_profile + .finish(&self.metrics, scan_stats); let QueryContext { trace_id, node, @@ -206,7 +210,6 @@ impl Drop for FlightDecoderStream { } else { "follower".to_string() }; - let scan_stats = self.scan_stats; let batches = self.metrics.batches_received.value(); let bytes = self.metrics.bytes_received.value(); let avg_ms = self.metrics.avg_batch_time_ms(); diff --git a/src/search/src/datafusion/distributed_plan/distribute_analyze_exec.rs b/src/search/src/datafusion/distributed_plan/distribute_analyze_exec.rs index 5903d86def..1f7c28904e 100644 --- a/src/search/src/datafusion/distributed_plan/distribute_analyze_exec.rs +++ b/src/search/src/datafusion/distributed_plan/distribute_analyze_exec.rs @@ -16,12 +16,12 @@ use std::sync::Arc; use arrow::{ - array::{Int32Builder, StringBuilder}, + array::{ArrayBuilder, ArrayRef, Int32Builder, StringArray, StringBuilder}, datatypes::SchemaRef, record_batch::RecordBatch, }; use arrow_schema::{DataType, Field, Schema}; -use config::{cluster::LOCAL_NODE, meta::cluster::NodeInfo}; +use config::{cluster::LOCAL_NODE, meta::cluster::NodeInfo, utils::time::now_micros}; use datafusion::{ common::{DataFusionError, Result, internal_err, tree_node::TreeNodeRecursion}, execution::TaskContext, @@ -32,24 +32,22 @@ use datafusion::{ stream::RecordBatchStreamAdapter, }, }; -use flight::common::Metrics; +use flight::profile::{ProfileInputs, ProfileNode, QueryProfile, RemoteNodeProfile}; use futures::StreamExt; -use crate::datafusion::{ - distributed_plan::display::DisplayableExecutionPlan, plan_metrics::get_cluster_metrics, +use crate::datafusion::distributed_plan::{ + display::DisplayableExecutionPlan, profile::build_profile, }; -/// query EXPLAIN ANALYZE in distributed mode -/// the output recordbatch's schema is -/// phase: current we have three phase -/// node: the ip or partition id -/// plan: plan with metrics -#[derive(Debug)] +/// EXPLAIN ANALYZE across nodes: one text row per node that ran a plan, phase = its depth. +#[derive(Debug, Clone)] pub struct DistributeAnalyzeExec { /// Control how much extra to print verbose: bool, /// If statistics should be displayed show_statistics: bool, + /// Adds the `profile` column, the whole profile as JSON on the first row. + profile: bool, /// The input plan (the plan being analyzed) pub input: Arc, /// The output schema for RecordBatches of this exec node @@ -58,25 +56,90 @@ pub struct DistributeAnalyzeExec { } impl DistributeAnalyzeExec { - /// Create a new AnalyzeExec - pub fn new(verbose: bool, show_statistics: bool, input: Arc) -> Self { - let fields = vec![ + pub fn new( + verbose: bool, + show_statistics: bool, + profile: bool, + input: Arc, + ) -> Self { + let mut fields = vec![ Field::new("phase", DataType::Int32, false), Field::new("node_address", DataType::Utf8, false), Field::new("node_name", DataType::Utf8, false), Field::new("plan", DataType::Utf8, false), ]; + if profile { + fields.push(Field::new("profile", DataType::Utf8, true)); + } let schema = Arc::new(Schema::new(fields)); let cache = Self::compute_properties(&input, Arc::clone(&schema)); DistributeAnalyzeExec { verbose, show_statistics, + profile, input, schema, cache, } } + fn output_batch( + &self, + profile: &QueryProfile, + duration: std::time::Duration, + ) -> Result { + let mut phase_builder = Int32Builder::new(); + let mut address_builder = StringBuilder::with_capacity(1, 1024); + let mut name_builder = StringBuilder::with_capacity(1, 1024); + let mut plan_builder = StringBuilder::with_capacity(1, 1024); + + phase_builder.append_value(0); + address_builder.append_value(&profile.node_address); + name_builder.append_value(&profile.node_name); + let annotated_plan = DisplayableExecutionPlan::new(self.input.as_ref()) + .indent(self.verbose) + .to_string(); + plan_builder.append_value(annotated_plan); + + let mut remote_rows = Vec::new(); + collect_remote_rows(&profile.plan, 1, &mut remote_rows); + for (stage, node, plan) in remote_rows { + phase_builder.append_value(stage); + address_builder.append_value(&node.node_address); + name_builder.append_value(&node.node_name); + plan_builder.append_value(plan.to_text()); + } + + if self.verbose { + phase_builder.append_value(0); + address_builder.append_value("Output Rows"); + name_builder.append_value(""); + plan_builder.append_value(profile.output_rows.to_string()); + + phase_builder.append_value(0); + address_builder.append_value("Duration"); + name_builder.append_value(""); + plan_builder.append_value(format!("{duration:?}")); + } + + let rows = phase_builder.len(); + let mut columns: Vec = vec![ + Arc::new(phase_builder.finish()), + Arc::new(address_builder.finish()), + Arc::new(name_builder.finish()), + Arc::new(plan_builder.finish()), + ]; + if self.profile { + let json = + serde_json::to_string(profile).map_err(|e| DataFusionError::External(e.into()))?; + columns.push(Arc::new(StringArray::from_iter( + std::iter::once(Some(json)).chain(std::iter::repeat_n(None, rows - 1)), + ))); + } + + RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(DataFusionError::from) + } + /// This function creates the cache object that stores the plan properties such as schema, /// equivalence properties, ordering, partitioning, etc. fn compute_properties( @@ -130,6 +193,7 @@ impl ExecutionPlan for DistributeAnalyzeExec { Ok(Arc::new(Self::new( self.verbose, self.show_statistics, + self.profile, children.pop().unwrap(), ))) } @@ -147,10 +211,8 @@ impl ExecutionPlan for DistributeAnalyzeExec { // Create future that computes the final output let start = std::time::Instant::now(); - let captured_input = Arc::clone(&self.input); - let captured_schema = Arc::clone(&self.schema); - let verbose = self.verbose; - let show_statistics = self.show_statistics; + let start_us = now_micros(); + let exec = self.clone(); let mut stream = execute_stream(self.input.clone(), context.clone())?; let output = async move { @@ -158,16 +220,19 @@ impl ExecutionPlan for DistributeAnalyzeExec { while let Some(batch) = stream.next().await.transpose()? { total_rows += batch.num_rows(); } + // remote nodes publish their profiles when their streams are dropped + drop(stream); let duration = std::time::Instant::now() - start; - create_output_batch( - verbose, - show_statistics, - total_rows, - duration, - captured_input, - captured_schema, - ) + let profile = QueryProfile { + node_name: LOCAL_NODE.get_name(), + node_address: LOCAL_NODE.get_grpc_addr(), + start_us, + end_us: now_micros(), + output_rows: total_rows, + plan: build_profile(exec.input.as_ref()), + }; + exec.output_batch(&profile, duration) }; Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -177,75 +242,35 @@ impl ExecutionPlan for DistributeAnalyzeExec { } } -fn create_output_batch( - verbose: bool, - _show_statistics: bool, - total_rows: usize, - duration: std::time::Duration, - input: Arc, - schema: SchemaRef, -) -> Result { - let mut phase_builder = Int32Builder::new(); - let mut address_builder = StringBuilder::with_capacity(1, 1024); - let mut name_builder = StringBuilder::with_capacity(1, 1024); - let mut plan_builder = StringBuilder::with_capacity(1, 1024); - - phase_builder.append_value(0); // Phase 0 for the main plan - address_builder.append_value(LOCAL_NODE.get_grpc_addr()); - name_builder.append_value(LOCAL_NODE.get_name()); - - let annotated_plan = DisplayableExecutionPlan::new(input.as_ref()) - .indent(verbose) - .to_string(); - plan_builder.append_value(annotated_plan); - - // add metrics from remote scan - let metrics = collect_metrics(input); - - for m in metrics { - phase_builder.append_value(m.stage); - address_builder.append_value(m.node_address); - name_builder.append_value(m.node_name); - plan_builder.append_value(m.metrics); +/// Lists, depth first, each remote node that sent back its plan, with its stage below the leader. +fn collect_remote_rows<'a>( + node: &'a ProfileNode, + stage: i32, + rows: &mut Vec<(i32, &'a RemoteNodeProfile, &'a ProfileNode)>, +) { + match &node.inputs { + ProfileInputs::Children(children) => { + for child in children { + collect_remote_rows(child, stage, rows); + } + } + ProfileInputs::Remote(remotes) => { + for (remote, plan) in remotes + .iter() + .filter_map(|r| r.plan.as_ref().map(|plan| (r, plan))) + { + rows.push((stage, remote, plan)); + collect_remote_rows(plan, stage + 1, rows); + } + } } - - if verbose { - phase_builder.append_value(0); // Phase 0 for the main plan - address_builder.append_value("Output Rows"); - name_builder.append_value(""); - plan_builder.append_value(total_rows.to_string()); - - phase_builder.append_value(0); // Phase 0 for the main plan - address_builder.append_value("Duration"); - name_builder.append_value(""); - plan_builder.append_value(format!("{duration:?}")); - } - - RecordBatch::try_new( - schema, - vec![ - Arc::new(phase_builder.finish()), - Arc::new(address_builder.finish()), - Arc::new(name_builder.finish()), - Arc::new(plan_builder.finish()), - ], - ) - .map_err(DataFusionError::from) -} - -fn collect_metrics(plan: Arc) -> Vec { - let mut metrics = Vec::new(); - for m in get_cluster_metrics(&plan) { - let m = m.lock().clone(); - metrics.extend(m); - } - metrics } #[cfg(test)] mod tests { use std::sync::Arc; + use arrow::array::Array; use arrow_schema::{DataType, Field, Schema}; use datafusion::physical_plan::empty::EmptyExec; @@ -259,7 +284,7 @@ mod tests { #[test] fn test_distribute_analyze_exec_new_non_verbose() { let input = empty_input(); - let exec = DistributeAnalyzeExec::new(false, false, input); + let exec = DistributeAnalyzeExec::new(false, false, false, input); assert_eq!(exec.name(), "DistributeAnalyzeExec"); assert_eq!(exec.children().len(), 1); } @@ -267,16 +292,91 @@ mod tests { #[test] fn test_distribute_analyze_exec_new_verbose() { let input = empty_input(); - let exec = DistributeAnalyzeExec::new(true, true, input); + let exec = DistributeAnalyzeExec::new(true, true, false, input); assert_eq!(exec.name(), "DistributeAnalyzeExec"); } #[test] fn test_distribute_analyze_exec_schema_has_four_fields() { let input = empty_input(); - let exec = DistributeAnalyzeExec::new(false, false, input); + let exec = DistributeAnalyzeExec::new(false, false, false, input); assert_eq!(exec.schema().fields().len(), 4); assert_eq!(exec.schema().field(0).name(), "phase"); assert_eq!(exec.schema().field(1).name(), "node_address"); } + + #[test] + fn test_distribute_analyze_exec_schema_adds_profile_when_requested() { + let exec = DistributeAnalyzeExec::new(false, false, true, empty_input()); + assert_eq!(exec.schema().fields().len(), 5); + assert_eq!(exec.schema().field(4).name(), "profile"); + assert!(exec.schema().field(4).is_nullable()); + } + + #[tokio::test] + async fn test_distribute_analyze_exec_without_profile_emits_four_columns() { + let exec = Arc::new(DistributeAnalyzeExec::new( + true, + false, + false, + empty_input(), + )); + let batches = datafusion::physical_plan::collect(exec, Arc::new(TaskContext::default())) + .await + .unwrap(); + assert_eq!(batches[0].num_columns(), 4); + } + + #[tokio::test] + async fn test_distribute_analyze_exec_emits_profile_on_first_row() { + let exec = Arc::new(DistributeAnalyzeExec::new(true, false, true, empty_input())); + let batches = datafusion::physical_plan::collect(exec, Arc::new(TaskContext::default())) + .await + .unwrap(); + let profile = batches[0] + .column(4) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(profile.is_valid(0)); + assert!((1..profile.len()).all(|i| profile.is_null(i))); + let parsed: QueryProfile = serde_json::from_str(profile.value(0)).unwrap(); + assert_eq!(parsed.plan.name, "EmptyExec"); + assert!(parsed.end_us >= parsed.start_us); + } + + #[test] + fn test_collect_remote_rows_nests_stages_and_skips_nodes_without_plan() { + let node = |name: &str, plan: Option| RemoteNodeProfile { + node_name: name.to_string(), + plan, + ..Default::default() + }; + let scan = |remotes| ProfileNode { + name: "RemoteScanExec".to_string(), + inputs: ProfileInputs::Remote(remotes), + ..Default::default() + }; + let cluster_leader = node( + "b-querier-0", + Some(scan(vec![node( + "b-querier-1", + Some(ProfileNode::default()), + )])), + ); + let root = ProfileNode { + inputs: ProfileInputs::Children(vec![scan(vec![ + cluster_leader, + node("skipped", None), + ])]), + ..Default::default() + }; + let mut rows = Vec::new(); + collect_remote_rows(&root, 1, &mut rows); + let rows: Vec<_> = rows + .iter() + .map(|(stage, node, _)| (*stage, node.node_name.as_str())) + .collect(); + assert_eq!(rows, vec![(1, "b-querier-0"), (2, "b-querier-1")]); + } } diff --git a/src/search/src/datafusion/distributed_plan/mod.rs b/src/search/src/datafusion/distributed_plan/mod.rs index 8b455ee154..3a6323a177 100644 --- a/src/search/src/datafusion/distributed_plan/mod.rs +++ b/src/search/src/datafusion/distributed_plan/mod.rs @@ -39,6 +39,7 @@ pub mod enrichment_exec; pub mod metadata_count_exec; pub mod node; mod once_async; +pub mod profile; pub mod remote_scan_exec; pub mod rewrite; pub mod streaming_aggs_exec; diff --git a/src/search/src/datafusion/distributed_plan/profile.rs b/src/search/src/datafusion/distributed_plan/profile.rs new file mode 100644 index 0000000000..2849c21b2b --- /dev/null +++ b/src/search/src/datafusion/distributed_plan/profile.rs @@ -0,0 +1,380 @@ +// 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 +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::fmt; + +use config::{ + meta::{cluster::NodeInfo, search::ScanStats}, + utils::time::now_micros, +}; +use datafusion::physical_plan::{ + DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + metrics::{MetricValue, MetricsSet}, +}; +use flight::{ + common::{Metrics, NodeProfiles, RemoteScanMetrics}, + profile::{MetricKind, ProfileInputs, ProfileMetric, ProfileNode, RemoteNodeProfile}, +}; + +use super::remote_scan_exec::RemoteScanExec; + +/// One remote node as the leader sees it, published to its scan on drop; the default records +/// nothing. +#[derive(Debug, Default)] +pub struct NodeProfileRecorder { + profiles: Option, + profile: RemoteNodeProfile, +} + +impl NodeProfileRecorder { + pub fn start(profiles: NodeProfiles, partition: usize, node: &dyn NodeInfo) -> Self { + let role = match (node.is_querier(), node.is_ingester()) { + (true, true) => "querier,ingester", + (true, false) => "querier", + (false, true) => "ingester", + (false, false) => "", + }; + Self { + profiles: Some(profiles), + profile: RemoteNodeProfile { + partition, + node_name: node.get_name(), + node_address: node.get_grpc_addr(), + role: role.to_string(), + cluster: node.get_cluster(), + start_us: now_micros(), + ..Default::default() + }, + } + } + + /// Keeps the plan the node sent back about itself. + pub fn record_plan(&mut self, metrics: Vec) { + if let Some(plan) = metrics.into_iter().find_map(|m| m.profile) { + self.profile.plan = Some(plan); + } + } + + pub fn mark_skipped(&mut self) { + self.profile.skipped = true; + } + + pub fn finish(&mut self, metrics: &RemoteScanMetrics, scan_stats: ScanStats) { + let p = &mut self.profile; + p.output_rows = metrics.output_rows.value() as u64; + p.output_batches = metrics.batches_received.value() as u64; + p.output_bytes = metrics.bytes_received.value() as u64; + p.decode_ns = metrics.decode_time.value() as u64; + p.scan_stats = Some(scan_stats); + } +} + +impl Drop for NodeProfileRecorder { + fn drop(&mut self) { + if let Some(profiles) = self.profiles.take() { + self.profile.end_us = now_micros(); + let profile = std::mem::take(&mut self.profile); + profiles.lock().insert(profile.partition, profile); + } + } +} + +/// A plan node's own line, exactly as EXPLAIN prints it. +struct PlanDetail<'a>(&'a dyn ExecutionPlan); + +impl fmt::Display for PlanDetail<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt_as(DisplayFormatType::Default, f) + } +} + +/// Profile of an executed plan; a remote scan lists the nodes it called, as its children never ran. +pub fn build_profile(plan: &dyn ExecutionPlan) -> ProfileNode { + let inputs = match plan.downcast_ref::() { + Some(remote) => ProfileInputs::Remote( + remote + .state + .node_profiles + .lock() + .values() + .cloned() + .collect(), + ), + None => ProfileInputs::Children( + plan.children() + .into_iter() + .map(|child| build_profile(child.as_ref())) + .collect(), + ), + }; + ProfileNode { + name: plan.name().to_string(), + detail: PlanDetail(plan).to_string(), + partitions: plan.output_partitioning().partition_count(), + metrics: plan + .metrics() + .map(|m| convert_metrics(&m)) + .unwrap_or_default(), + inputs, + } +} + +/// Aggregates and orders the metrics the same way the text plan does, so `display` matches it. +fn convert_metrics(metrics: &MetricsSet) -> Vec { + metrics + .aggregate_by_name() + .sorted_for_display() + .timestamps_removed() + .iter() + .map(|metric| { + let value = metric.value(); + let (kind, number, total) = metric_value(value); + ProfileMetric { + name: value.name().to_string(), + kind, + value: number, + total, + display: value.to_string(), + } + }) + .collect() +} + +fn metric_value(value: &MetricValue) -> (MetricKind, u64, Option) { + let as_u64 = |v: usize| v as u64; + match value { + MetricValue::OutputRows(c) + | MetricValue::SpillCount(c) + | MetricValue::OutputBatches(c) + | MetricValue::SpilledRows(c) + | MetricValue::Count { count: c, .. } => (MetricKind::Count, as_u64(c.value()), None), + MetricValue::SpilledBytes(c) | MetricValue::OutputBytes(c) => { + (MetricKind::Bytes, as_u64(c.value()), None) + } + MetricValue::CurrentMemoryUsage(g) | MetricValue::PeakMemoryUsage { gauge: g, .. } => { + (MetricKind::Bytes, as_u64(g.value()), None) + } + MetricValue::ElapsedCompute(t) | MetricValue::Time { time: t, .. } => { + (MetricKind::Time, as_u64(t.value()), None) + } + MetricValue::Gauge { gauge: g, .. } => (MetricKind::Gauge, as_u64(g.value()), None), + MetricValue::PruningMetrics { + pruning_metrics, .. + } => { + let pruned = as_u64(pruning_metrics.pruned()); + let total = pruned + as_u64(pruning_metrics.matched()); + (MetricKind::Pruning, pruned, Some(total)) + } + MetricValue::Ratio { ratio_metrics, .. } => ( + MetricKind::Ratio, + as_u64(ratio_metrics.part()), + Some(as_u64(ratio_metrics.total())), + ), + // Timestamps never get here: convert_metrics removes them first. + MetricValue::Custom { .. } + | MetricValue::StartTimestamp(_) + | MetricValue::EndTimestamp(_) => (MetricKind::Other, as_u64(value.as_usize()), None), + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{Int32Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use datafusion::{ + physical_plan::{ + collect, + display::DisplayableExecutionPlan, + empty::EmptyExec, + metrics::{ExecutionPlanMetricsSet, MetricBuilder}, + union::UnionExec, + }, + prelude::SessionContext, + }; + + use super::*; + use crate::datafusion::distributed_plan::node::RemoteScanNode; + + #[derive(Debug)] + struct TestNode; + + impl NodeInfo for TestNode { + fn is_querier(&self) -> bool { + true + } + fn get_grpc_addr(&self) -> String { + "10.0.0.1:5081".to_string() + } + fn get_auth_token(&self) -> String { + String::new() + } + fn get_name(&self) -> String { + "q-0".to_string() + } + fn is_local(&self) -> bool { + false + } + } + + fn empty() -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + Arc::new(EmptyExec::new(schema)) + } + + fn batch(rows: i32) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from_iter_values(0..rows))], + ) + .unwrap() + } + + #[test] + fn test_recorder_publishes_one_node_on_drop() { + let profiles = NodeProfiles::default(); + let mut recorder = NodeProfileRecorder::start(profiles.clone(), 2, &TestNode); + let metrics = RemoteScanMetrics::new(2, &ExecutionPlanMetricsSet::new()); + metrics.record_output(&batch(3)); + metrics.record_output(&batch(4)); + recorder.record_plan(vec![Metrics { + profile: Some(ProfileNode { + name: "FilterExec".to_string(), + ..Default::default() + }), + ..Default::default() + }]); + recorder.finish( + &metrics, + ScanStats { + files: 5, + ..Default::default() + }, + ); + assert!(profiles.lock().is_empty()); + drop(recorder); + + let p = profiles.lock().get(&2).cloned().unwrap(); + assert_eq!(p.node_name, "q-0"); + assert_eq!(p.role, "querier"); + assert_eq!(p.output_rows, 7); + assert_eq!(p.output_batches, 2); + assert!(p.end_us >= p.start_us); + assert_eq!(p.scan_stats.unwrap().files, 5); + assert_eq!(p.plan.unwrap().name, "FilterExec"); + } + + #[test] + fn test_recorder_publishes_skipped_and_unfinished_nodes() { + let profiles = NodeProfiles::default(); + NodeProfileRecorder::start(profiles.clone(), 0, &TestNode).mark_skipped(); + drop(NodeProfileRecorder::start(profiles.clone(), 1, &TestNode)); + NodeProfileRecorder::default().mark_skipped(); + let profiles = profiles.lock(); + assert_eq!(profiles.len(), 2); + assert!(profiles[&0].skipped); + assert!(!profiles[&1].skipped && profiles[&1].scan_stats.is_none()); + } + + #[test] + fn test_convert_metrics_sums_partitions() { + let set = ExecutionPlanMetricsSet::new(); + MetricBuilder::new(&set).output_rows(0).add(10); + MetricBuilder::new(&set).output_rows(1).add(30); + MetricBuilder::new(&set) + .elapsed_compute(0) + .add_duration(std::time::Duration::from_nanos(5)); + let metrics = convert_metrics(&set.clone_inner()); + let rows = metrics.iter().find(|m| m.name == "output_rows").unwrap(); + assert_eq!(rows.kind, MetricKind::Count); + assert_eq!(rows.value, 40); + let time = metrics + .iter() + .find(|m| m.name == "elapsed_compute") + .unwrap(); + assert_eq!(time.kind, MetricKind::Time); + assert_eq!(time.value, 5); + assert_eq!(time.display, "5ns"); + } + + #[tokio::test] + async fn test_profile_text_matches_datafusion() { + let ctx = SessionContext::new(); + let sql = "SELECT a % 7 AS k, count(*) FROM (SELECT unnest(range(1, 50)) AS a) \ + WHERE a > 10 GROUP BY k ORDER BY k"; + let plan = ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + collect(plan.clone(), ctx.task_ctx()).await.unwrap(); + let expected = DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .indent(false) + .to_string(); + let profile = build_profile(plan.as_ref()); + assert_eq!(profile.to_text(), expected); + assert!(expected.contains("selectivity="), "{expected}"); + } + + #[test] + fn test_memory_metrics_are_bytes() { + let set = ExecutionPlanMetricsSet::new(); + MetricBuilder::new(&set).mem_used(0).set(2048); + let metrics = convert_metrics(&set.clone_inner()); + assert_eq!(metrics[0].kind, MetricKind::Bytes); + assert_eq!(metrics[0].value, 2048); + assert_eq!(metrics[0].display, "2.0 KB"); + } + + #[test] + fn test_build_profile_walks_children() { + let plan = UnionExec::try_new(vec![empty(), empty()]).unwrap(); + let profile = build_profile(plan.as_ref()); + assert_eq!(profile.name, "UnionExec"); + let ProfileInputs::Children(children) = profile.inputs else { + panic!("expected children"); + }; + assert_eq!(children.len(), 2); + assert_eq!(children[0].name, "EmptyExec"); + } + + #[test] + fn test_build_profile_stops_at_remote() { + let node = RemoteScanNode { + nodes: vec![Arc::new(TestNode)], + opentelemetry_context: opentelemetry::Context::current(), + query_identifier: Default::default(), + search_infos: Default::default(), + index_info: Default::default(), + super_cluster_info: Default::default(), + }; + let remote = RemoteScanExec::new(empty(), node).unwrap(); + drop(NodeProfileRecorder::start( + remote.state.node_profiles.clone(), + 0, + &TestNode, + )); + let profile = build_profile(&remote); + assert_eq!(profile.name, "RemoteScanExec"); + let ProfileInputs::Remote(remote) = profile.inputs else { + panic!("expected remote inputs"); + }; + assert_eq!(remote[0].node_name, "q-0"); + } +} diff --git a/src/search/src/datafusion/distributed_plan/remote_scan_exec.rs b/src/search/src/datafusion/distributed_plan/remote_scan_exec.rs index 1ca5a00db9..6a014af2a9 100644 --- a/src/search/src/datafusion/distributed_plan/remote_scan_exec.rs +++ b/src/search/src/datafusion/distributed_plan/remote_scan_exec.rs @@ -13,14 +13,11 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::sync::{Arc, atomic::AtomicUsize}; +use std::sync::Arc; use arrow_schema::SchemaRef; use config::{ - meta::{ - inverted_index::IndexOptimizeMode, - search::{ScanStats, SearchEventType}, - }, + meta::{inverted_index::IndexOptimizeMode, search::SearchEventType}, utils::rand::generate_random_string, }; use datafusion::{ @@ -35,14 +32,13 @@ use datafusion::{ }, }; use datafusion_proto::bytes::physical_plan_to_bytes_with_extension_codec; -use flight::common::{Metrics, RemoteScanMetrics}; +use flight::common::{RemoteScanMetrics, RemoteScanState}; use futures::{StreamExt, TryStreamExt}; use futures_util::pin_mut; use opentelemetry::context::FutureExt as _; -use parking_lot::Mutex; use rand::prelude::SliceRandom; -use super::node::RemoteScanNode; +use super::{node::RemoteScanNode, profile::NodeProfileRecorder}; use crate::datafusion::{ distributed_plan::{ codec::get_physical_extension_codec, @@ -60,10 +56,7 @@ pub struct RemoteScanExec { remote_scan_node: RemoteScanNode, partitions: usize, cache: Arc, - pub scan_stats: Arc>, - pub peak_memory: Arc, - pub partial_err: Arc>, - pub cluster_metrics: Arc>>, + pub state: RemoteScanState, pub enrich_mode_node_idx: usize, pub metrics: ExecutionPlanMetricsSet, } @@ -101,53 +94,19 @@ impl RemoteScanExec { remote_scan_node, partitions: output_partitions, cache, - scan_stats: Arc::new(Mutex::new(ScanStats::default())), - partial_err: Arc::new(Mutex::new(String::new())), - cluster_metrics: Arc::new(Mutex::new(Vec::new())), - peak_memory: Arc::new(AtomicUsize::new(0)), + state: RemoteScanState::default(), enrich_mode_node_idx, metrics: ExecutionPlanMetricsSet::new(), }) } - pub fn scan_stats(&self) -> Arc> { - self.scan_stats.clone() - } - - pub fn partial_err(&self) -> Arc> { - self.partial_err.clone() - } - - pub fn cluster_metrics(&self) -> Arc>> { - self.cluster_metrics.clone() - } - - pub fn peak_memory(&self) -> Arc { - self.peak_memory.clone() - } - /// The input plan bytes exactly as every follower receives them. pub fn encoded_plan(&self) -> &[u8] { &self.remote_scan_node.search_infos.plan } - pub fn with_scan_stats(mut self, scan_stats: Arc>) -> Self { - self.scan_stats = scan_stats; - self - } - - pub fn with_partial_err(mut self, partial_err: Arc>) -> Self { - self.partial_err = partial_err; - self - } - - pub fn with_cluster_metrics(mut self, cluster_metrics: Arc>>) -> Self { - self.cluster_metrics = cluster_metrics; - self - } - - pub fn with_peak_memory(mut self, peak_memory: Arc) -> Self { - self.peak_memory = peak_memory; + pub fn with_state(mut self, state: RemoteScanState) -> Self { + self.state = state; self } @@ -254,10 +213,7 @@ impl ExecutionPlan for RemoteScanExec { return Ok(self); } let remote_scan = Self::new(children[0].clone(), self.remote_scan_node.clone())? - .with_scan_stats(self.scan_stats.clone()) - .with_partial_err(self.partial_err.clone()) - .with_cluster_metrics(self.cluster_metrics.clone()) - .with_peak_memory(self.peak_memory.clone()) + .with_state(self.state.clone()) .with_metrics(self.metrics.clone()); Ok(Arc::new(remote_scan)) } @@ -273,10 +229,7 @@ impl ExecutionPlan for RemoteScanExec { partition, partition == self.enrich_mode_node_idx, self.input.schema().clone(), - self.scan_stats(), - self.partial_err(), - self.cluster_metrics(), - self.peak_memory(), + self.state.clone(), baseline_metrics, ); let stream = futures::stream::once(fut).try_flatten(); @@ -297,12 +250,15 @@ async fn get_remote_batch( partition: usize, enrich_mode: bool, schema: SchemaRef, - scan_stats: Arc>, - partial_err: Arc>, - cluster_metrics: Arc>>, - peak_memory: Arc, + state: RemoteScanState, metrics: RemoteScanMetrics, ) -> Result { + let RemoteScanState { + scan_stats, + peak_memory, + partial_err, + node_profiles, + } = state; let start = std::time::Instant::now(); let cfg = config::get_config(); let trace_id = remote_scan_node.query_identifier.trace_id.clone(); @@ -335,6 +291,12 @@ async fn get_remote_batch( } } + let mut recorder = if remote_scan_node.search_infos.is_analyze { + NodeProfileRecorder::start(node_profiles, partition, node.as_ref()) + } else { + NodeProfileRecorder::default() + }; + let empty_stream = EmptyStream::new( &trace_id, schema.clone(), @@ -351,6 +313,7 @@ async fn get_remote_batch( && !enrich_mode && remote_scan_node.is_file_list_empty(partition) { + recorder.mark_skipped(); return Ok(get_empty_stream(empty_stream)); } @@ -412,9 +375,9 @@ async fn get_remote_batch( .with_is_querier(is_querier) .with_scan_stats(scan_stats) .with_partial_err(partial_err.clone()) - .with_cluster_metrics(cluster_metrics) .with_peak_memory(peak_memory) - .with_start_time(start); + .with_start_time(start) + .with_node_profile(recorder); let mut stream = FlightDecoderStream::new(stream, schema.clone(), metrics, query_context); let stream = async_stream::stream! { @@ -529,7 +492,7 @@ mod tests { #[test] fn test_remote_scan_exec_partial_err_empty() { let exec = make_remote_exec(); - assert!(exec.partial_err().lock().is_empty()); + assert!(exec.state.partial_err.lock().is_empty()); } #[test] @@ -541,11 +504,11 @@ mod tests { } #[test] - fn test_remote_scan_exec_with_partial_err() { - let exec = make_remote_exec(); - let new_err = Arc::new(parking_lot::Mutex::new("pre-set".to_string())); - let exec = exec.with_partial_err(new_err.clone()); - assert_eq!(*exec.partial_err().lock(), "pre-set"); + fn test_remote_scan_exec_with_state() { + let state = RemoteScanState::default(); + *state.partial_err.lock() = "pre-set".to_string(); + let exec = make_remote_exec().with_state(state); + assert_eq!(*exec.state.partial_err.lock(), "pre-set"); } #[test] diff --git a/src/search/src/datafusion/optimizer/mod.rs b/src/search/src/datafusion/optimizer/mod.rs index d1190adba5..10dd8109e9 100644 --- a/src/search/src/datafusion/optimizer/mod.rs +++ b/src/search/src/datafusion/optimizer/mod.rs @@ -297,6 +297,7 @@ pub fn generate_physical_optimizer_rules( pub async fn create_physical_plan( ctx: &SessionContext, sql: &str, + analyze_profile: bool, ) -> Result> { let plan = ctx.state().create_logical_plan(sql).await?; // Defense in depth: reject DDL/DML/COPY plans so search SQL cannot write files on the node. @@ -306,7 +307,7 @@ pub async fn create_physical_plan( .with_allow_statements(false) .verify_plan(&plan)?; let physical_plan = ctx.state().create_physical_plan(&plan).await?; - optimize_distribute_analyze(physical_plan) + optimize_distribute_analyze(physical_plan, analyze_profile) } #[cfg(test)] @@ -542,12 +543,24 @@ mod tests { "INSERT INTO t2 VALUES (1)", ] { assert!( - create_physical_plan(&ctx, sql).await.is_err(), + create_physical_plan(&ctx, sql, false).await.is_err(), "should reject: {sql}" ); } } + #[tokio::test] + async fn test_create_physical_plan_adds_profile_only_when_asked() { + let sql = "EXPLAIN ANALYZE SELECT a FROM t"; + let ctx = SessionContext::new(); + register_tables(&ctx); + let plan = create_physical_plan(&ctx, sql, false).await.unwrap(); + assert_eq!(plan.schema().fields().len(), 4); + + let plan = create_physical_plan(&ctx, sql, true).await.unwrap(); + assert_eq!(plan.schema().field(4).name(), "profile"); + } + #[tokio::test] async fn test_struct_field_access_matches_datafusion() { let mut sql = test_sql(); diff --git a/src/search/src/datafusion/optimizer/physical_optimizer/distribute_analyze.rs b/src/search/src/datafusion/optimizer/physical_optimizer/distribute_analyze.rs index 4f316888e9..4db2a1b3e7 100644 --- a/src/search/src/datafusion/optimizer/physical_optimizer/distribute_analyze.rs +++ b/src/search/src/datafusion/optimizer/physical_optimizer/distribute_analyze.rs @@ -28,11 +28,15 @@ use crate::datafusion::distributed_plan::{ }; // replace the AnalyzeExec to DistributeAnalyzeExec -pub fn optimize_distribute_analyze(plan: Arc) -> Result> { +pub fn optimize_distribute_analyze( + plan: Arc, + profile: bool, +) -> Result> { if let Some(analyze) = plan.downcast_ref::() { let distribute_analyze = Arc::new(DistributeAnalyzeExec::new( analyze.verbose(), analyze.show_statistics(), + profile, analyze.input().clone(), )) as Arc; let mut rewriter = AnalyzeRewrite::new(); @@ -78,9 +82,8 @@ mod tests { ])) } - #[test] - fn test_optimize_distribute_analyze() { - let plan = Arc::new( + fn analyze_plan() -> Arc { + Arc::new( AnalyzeExec::builder( false, false, @@ -89,9 +92,20 @@ mod tests { ) .with_metric_types(vec![]) .build(), - ); - let optimized_plan = optimize_distribute_analyze(plan).unwrap(); + ) + } + + #[test] + fn test_optimize_distribute_analyze() { + let optimized_plan = optimize_distribute_analyze(analyze_plan(), false).unwrap(); assert_eq!(optimized_plan.name(), "DistributeAnalyzeExec"); + assert_eq!(optimized_plan.schema().fields().len(), 4); + } + + #[test] + fn test_optimize_distribute_analyze_with_profile() { + let optimized_plan = optimize_distribute_analyze(analyze_plan(), true).unwrap(); + assert_eq!(optimized_plan.schema().fields().len(), 5); } #[test] diff --git a/src/search/src/datafusion/plan_metrics.rs b/src/search/src/datafusion/plan_metrics.rs index a81c74daad..f6600a4c4e 100644 --- a/src/search/src/datafusion/plan_metrics.rs +++ b/src/search/src/datafusion/plan_metrics.rs @@ -15,41 +15,9 @@ use std::sync::{Arc, atomic::AtomicUsize}; -use datafusion::{ - common::{ - Result, - tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}, - }, - physical_plan::ExecutionPlan, - prelude::SessionContext, -}; -use flight::common::Metrics; -use parking_lot::Mutex; +use datafusion::prelude::SessionContext; -use super::{distributed_plan::remote_scan_exec::RemoteScanExec, peak_memory_pool::PeakMemoryPool}; - -pub fn get_cluster_metrics(plan: &Arc) -> Vec>>> { - let mut visitor = MetricsVisitor::default(); - let _ = plan.visit(&mut visitor); - visitor.metrics -} - -#[derive(Default)] -struct MetricsVisitor { - metrics: Vec>>>, -} - -impl<'n> TreeNodeVisitor<'n> for MetricsVisitor { - type Node = Arc; - - fn f_up(&mut self, node: &'n Self::Node) -> Result { - if node.name() == "RemoteScanExec" { - let remote_scan_exec = node.downcast_ref::().unwrap(); - self.metrics.push(remote_scan_exec.cluster_metrics()); - } - Ok(TreeNodeRecursion::Continue) - } -} +use super::peak_memory_pool::PeakMemoryPool; /// Get the peak memory from the SessionContext's memory pool. /// This should be called after the query execution to get the peak memory usage. diff --git a/src/search/src/utils.rs b/src/search/src/utils.rs index 367957170e..c5556658e3 100644 --- a/src/search/src/utils.rs +++ b/src/search/src/utils.rs @@ -87,17 +87,17 @@ impl ExecutionPlanVisitor for ScanStatsVisitor { let mayby_remote_scan_exec = plan.downcast_ref::(); if let Some(remote_scan_exec) = mayby_remote_scan_exec { { - let guard = remote_scan_exec.scan_stats.lock(); + let guard = remote_scan_exec.state.scan_stats.lock(); let stats = *guard; self.scan_stats.add(&stats); } { - let guard = remote_scan_exec.partial_err.lock(); + let guard = remote_scan_exec.state.partial_err.lock(); let err = (*guard).clone(); self.partial_err.push_str(&err); } { - let peak_memory = remote_scan_exec.peak_memory().load(Ordering::Relaxed); + let peak_memory = remote_scan_exec.state.peak_memory.load(Ordering::Relaxed); self.peak_memory = peak_memory.max(self.peak_memory); } } diff --git a/src/search_service/src/cluster/flight.rs b/src/search_service/src/cluster/flight.rs index d285b6d0c4..ba7c447513 100644 --- a/src/search_service/src/cluster/flight.rs +++ b/src/search_service/src/cluster/flight.rs @@ -397,10 +397,12 @@ pub async fn run_datafusion( register_table(&ctx, &sql).await?; // create physical plan - let physical_plan = create_physical_plan(&ctx, &sql.sql).await.map_err(|e| { - log::error!("[trace_id {trace_id}] flight->search: create physical plan error: {e}"); - e - })?; + let physical_plan = create_physical_plan(&ctx, &sql.sql, req.analyze_profile) + .await + .map_err(|e| { + log::error!("[trace_id {trace_id}] flight->search: create physical plan error: {e}"); + e + })?; if cfg.common.print_key_sql { log::info!("[trace_id {trace_id}] leader physical plan"); diff --git a/src/search_service/src/lib.rs b/src/search_service/src/lib.rs index d6fd7926a8..e1bb5da309 100644 --- a/src/search_service/src/lib.rs +++ b/src/search_service/src/lib.rs @@ -163,6 +163,7 @@ pub async fn search( request.set_local_mode(Some(v)); } request.set_use_cache(in_req.use_cache); + request.set_analyze_profile(in_req.query.analyze_profile); // Before the query is planned again in spawned tasks (cluster::http). request.mark_chat_history_reader(); let meta = Sql::new_from_req(&request, &query).await?; diff --git a/src/search_service/src/super_cluster/leader.rs b/src/search_service/src/super_cluster/leader.rs index f8d126ffe3..8e57a0985c 100644 --- a/src/search_service/src/super_cluster/leader.rs +++ b/src/search_service/src/super_cluster/leader.rs @@ -243,7 +243,7 @@ async fn run_datafusion( register_table(&ctx, &sql).await?; // create physical plan - let physical_plan = create_physical_plan(&ctx, &sql.sql).await?; + let physical_plan = create_physical_plan(&ctx, &sql.sql, req.analyze_profile).await?; if cfg.common.print_key_sql { log::info!("[trace_id {trace_id}] super cluster leader: physical plan"); From af4bad4167e5953b787a359ffbaa5e70fbd69fb9 Mon Sep 17 00:00:00 2001 From: Bhargav Date: Wed, 4 Nov 2026 12:21:32 +0000 Subject: [PATCH 2/8] feat(overview): org-scoped overview endpoint and Home Overview v2 page (#15931) Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- scripts/translations/.translation_state.json | 222 ++- src/api/http/src/handler/http/router/mod.rs | 1 + .../http/src/handler/http/router/openapi.rs | 1 + .../src/request/organization/org.rs | 104 ++ src/config/src/config.rs | 6 + src/config/src/router.rs | 51 +- src/core/src/lib.rs | 1 + src/core/src/overview.rs | 1610 +++++++++++++++++ src/infra/src/schema/mod.rs | 63 + src/infra/src/table/alerts/mod.rs | 111 +- src/stream/src/lib.rs | 133 +- .../Onboarding/home-first-data.spec.js | 124 +- web/src/lib/core/Button/OButton.spec.ts | 56 +- web/src/lib/core/Button/OButton.types.ts | 6 +- web/src/lib/core/Button/OButton.vue | 27 +- .../lib/core/EmptyState/OEmptyState.spec.ts | 23 + web/src/lib/core/EmptyState/OEmptyState.vue | 32 +- .../core/RefreshButton/ORefreshButton.spec.ts | 11 + .../RefreshButton/ORefreshButton.types.ts | 4 + .../lib/core/RefreshButton/ORefreshButton.vue | 2 + .../data/SegmentedBar/OSegmentedBar.spec.ts | 23 +- .../data/SegmentedBar/OSegmentedBar.types.ts | 19 +- .../lib/data/SegmentedBar/OSegmentedBar.vue | 17 +- web/src/locales/languages/ar-SA.json | 218 ++- web/src/locales/languages/de-DE.json | 218 ++- web/src/locales/languages/en-US.json | 187 +- web/src/locales/languages/es-ES.json | 216 ++- web/src/locales/languages/fr-FR.json | 218 ++- web/src/locales/languages/it-IT.json | 216 ++- web/src/locales/languages/ja-JP.json | 220 ++- web/src/locales/languages/ko-KR.json | 216 ++- web/src/locales/languages/nl-NL.json | 220 ++- web/src/locales/languages/pl-PL.json | 216 ++- web/src/locales/languages/pt-PT.json | 216 ++- web/src/locales/languages/ru-RU.json | 218 ++- web/src/locales/languages/th-TH.json | 224 ++- web/src/locales/languages/tr-TR.json | 224 ++- web/src/locales/languages/vi-VN.json | 216 ++- web/src/locales/languages/zh-CN.json | 220 ++- web/src/locales/languages/zh-TW.json | 222 ++- web/src/services/overview.queries.spec.ts | 57 + web/src/services/overview.queries.ts | 32 + web/src/services/overview.querykeys.ts | 20 + web/src/services/overview.spec.ts | 38 + web/src/services/overview.ts | 31 + web/src/ts/interfaces/overview.ts | 94 + web/src/views/HomeView.spec.ts | 7 + web/src/views/HomeView.vue | 13 +- web/src/views/OverviewSectionEmpty.spec.ts | 50 + web/src/views/OverviewSectionEmpty.vue | 66 + web/src/views/OverviewSectionError.vue | 9 +- web/src/views/OverviewSkeleton.vue | 22 +- web/src/views/OverviewTab.spec.ts | 1197 +++++++++++- web/src/views/OverviewTab.vue | 1339 +++++++++----- web/src/views/overview/OverviewAlertsCard.vue | 295 +++ web/src/views/overview/OverviewCards.spec.ts | 1032 +++++++++++ .../views/overview/OverviewDataHealthCard.vue | 227 +++ .../views/overview/OverviewDataHealthRow.vue | 109 ++ .../overview/OverviewHeadlineStrip.spec.ts | 373 ++++ .../views/overview/OverviewHeadlineStrip.vue | 321 ++++ web/src/views/overview/OverviewLimitCard.vue | 260 +++ web/src/views/overview/OverviewRowButton.vue | 26 + .../views/overview/OverviewSectionCard.vue | 84 + .../views/overview/OverviewSectionState.vue | 88 + .../views/overview/OverviewSignalCards.vue | 124 ++ web/src/views/overview/OverviewStatusLine.vue | 75 + .../views/overview/OverviewTopStreamsCard.vue | 118 ++ web/src/views/overview/overviewFixtures.ts | 111 ++ web/src/views/overview/overviewFormat.spec.ts | 459 +++++ web/src/views/overview/overviewFormat.ts | 343 ++++ .../views/overview/useOverviewTime.spec.ts | 51 + web/src/views/overview/useOverviewTime.ts | 34 + 72 files changed, 11921 insertions(+), 1516 deletions(-) create mode 100644 src/core/src/overview.rs create mode 100644 web/src/services/overview.queries.spec.ts create mode 100644 web/src/services/overview.queries.ts create mode 100644 web/src/services/overview.querykeys.ts create mode 100644 web/src/services/overview.spec.ts create mode 100644 web/src/services/overview.ts create mode 100644 web/src/ts/interfaces/overview.ts create mode 100644 web/src/views/OverviewSectionEmpty.spec.ts create mode 100644 web/src/views/OverviewSectionEmpty.vue create mode 100644 web/src/views/overview/OverviewAlertsCard.vue create mode 100644 web/src/views/overview/OverviewCards.spec.ts create mode 100644 web/src/views/overview/OverviewDataHealthCard.vue create mode 100644 web/src/views/overview/OverviewDataHealthRow.vue create mode 100644 web/src/views/overview/OverviewHeadlineStrip.spec.ts create mode 100644 web/src/views/overview/OverviewHeadlineStrip.vue create mode 100644 web/src/views/overview/OverviewLimitCard.vue create mode 100644 web/src/views/overview/OverviewRowButton.vue create mode 100644 web/src/views/overview/OverviewSectionCard.vue create mode 100644 web/src/views/overview/OverviewSectionState.vue create mode 100644 web/src/views/overview/OverviewSignalCards.vue create mode 100644 web/src/views/overview/OverviewStatusLine.vue create mode 100644 web/src/views/overview/OverviewTopStreamsCard.vue create mode 100644 web/src/views/overview/overviewFixtures.ts create mode 100644 web/src/views/overview/overviewFormat.spec.ts create mode 100644 web/src/views/overview/overviewFormat.ts create mode 100644 web/src/views/overview/useOverviewTime.spec.ts create mode 100644 web/src/views/overview/useOverviewTime.ts diff --git a/scripts/translations/.translation_state.json b/scripts/translations/.translation_state.json index ca7c686ff2..96bd10958a 100644 --- a/scripts/translations/.translation_state.json +++ b/scripts/translations/.translation_state.json @@ -252,22 +252,6 @@ "feature_traces": "998741a58ed2b411", "feature_vrl_functions": "4484dd94e9355480", "feature_workload_management": "8004f04c30d8bddc", - "features": "5697d03daef4de9c", - "features_all_unlocked": "0791a3840eb6061a", - "features_filter_all": "a52ace420f2175d0", - "features_filter_free": "f411a1fb62758b4c", - "features_filter_new": "18fdd549b2ed367a", - "features_free_heading": "5892004a2a606ee5", - "features_group_heading": "98bb11ae86a0c557", - "features_missing": "694943b3db420a2e", - "features_new": "d9fe8aa562864fe0", - "features_no_match": "c7655586749bb9b3", - "features_none": "6b474f2c0cba33e7", - "features_paid_ratio": "3257b4364ee32149", - "features_request_access": "22d6edf4b21d5e25", - "features_search": "55e327604647ce88", - "features_summary": "5d2957d8f8c9d8e7", - "features_unlocked": "531c7b28abf27600", "foreverNoLimits": "c4e4f35814b1022b", "fullyManaged": "49964f86a467ca7b", "gbPerDay": "dd9d1b63afeb0945", @@ -13076,19 +13060,6 @@ "tabLabel": "cb702378f31507ef", "title": "2a932f90172e99ea" }, - "featureLock": { - "bannerMessage": "2ad54d7dcf939937", - "bannerMessageGeneric": "9e4736ff1b60331d", - "contactSales": "604abea32b499a4c", - "description": "052f7504abd2c39a", - "headerLocked": "ce0f1697c8a3a6d5", - "routeTitle": "694943b3db420a2e", - "thisFeature": "2598604c06f0d223", - "title": "84027eb92891f1fd", - "toastMessage": "b17241a153486b9c", - "toastMessageGeneric": "4cf1f345d3caeb96", - "toastTitle": "feff10da75dad3ed" - }, "fleet": { "addCollector": "6d30a69cb01ed55a", "audit": { @@ -13229,7 +13200,6 @@ "restart": "6b983a81e5e8099e", "restartManyTitle": "b48d3498ba6feded", "restartMessage": "22ee1481ad969aff", - "restartOffline": "35842abc83dff75b", "restartQueued": "43da4197b31c94b2", "restartUnsupported": "322b4fd3d40b0ff7", "search": "773656193742d0b1", @@ -13565,7 +13535,7 @@ "confirmFixForward": "6e986df0401cda4a", "confirmResume": "323c118fa4613787", "confirmRollback": "ad34d76a7f53f6b5", - "confirmRollbackNoStable": "5b00aa2620aa662d", + "confirmRollbackNoStable": "76cba5181ffc44a2", "confirmSkip": "fbf2627372d3e0cc", "confirmed": "fe00b67b6dd1143f", "emptyFleet": "40e87c8e8e200c76", @@ -13626,7 +13596,6 @@ "phaseStopped": "1a4f630ac1b69fd0", "phaseTimedOut": "9718cc761cfd5810", "phaseTimeoutNote": "f11dc4d3138e81ac", - "phaseTimeoutNoteSeconds": "fafc2b0d6a049798", "phaseVerifying": "a7d4a9f301edaa61", "phases": "b4b0b75a774761b3", "pickConfig": "90ea303d63b6e1b9", @@ -13693,8 +13662,8 @@ "noAccess": "cb92165578c84ffb", "offlineHeartbeats": "b82c0b4fc3ceead6", "offlineHeartbeatsHelp": "91687a5a34666ebf", - "offlineRatio": "98affbf40cf01c0c", - "offlineRatioHelp": "fcf7765ecca25ef0", + "offlineRatio": "ab13785c48fe3a97", + "offlineRatioHelp": "8fa2bcd13b66f2f3", "outOfRange": "f9a9c12c180ebfe6", "phaseTimeout": "47d919e06d87345e", "phaseTimeoutHelp": "019f8feb0a24eaa2", @@ -20357,32 +20326,109 @@ "updateOrganization": "f927b01e6f33a249" }, "overview": { - "activeAnomalies": "e96dd052f0d944fa", - "activeIncidents": "f5239d5e7c01c5a5", + "activeAnomalies": "981ade4159cacdec", + "activeIncidents": "0d5b14c485598692", + "alerts": { + "criticalLevel": "78282475e12bce5a", + "emptyAction": "0fd0bf350a56af51", + "emptyDesc": "8627367fc7f62316", + "emptyTitle": "567533d895c27065", + "firingFor": "11cb9d91c8532103", + "firingNow": "c9061aa892c72868", + "noneFiring": "081fc480527c69f5", + "rulesActive": "72fdbccd8fb3aed6", + "rulesEvaluated": "9fedbef3944b2a0e", + "rulesFailing": "3212b6c9a5d4a814", + "rulesFailingToEvaluate": "2b3301a6487190ce", + "showMore": "1557a0aa0041cf6e", + "subtitle": "c2a5c180dd02641a", + "viewAll": "1553d761825d0509", + "warningLevel": "d7322f65981a285b" + }, "alertsSuffix": "deb3e366c077c228", - "allClear": "22d352f617758292", - "allClearDesc": "bfe3044f55b21533", - "allClearDescScoped": "7c66a0cc083a98ba", - "allClearPartial": "f5a637211cc64316", - "allClearPartialDesc": "e5ed97f9af61f860", "anomalyDetectedDescription": "67a554a9305c968d", "anomalyDetectedTitle": "42b3e5f78e8b2967", - "areaAlerts": "f62e9267d196c761", - "areaAnomalies": "a0a660cc9daa76b9", - "areaIncidents": "bfe868931580fb3b", - "areaServices": "604dce445e326d95", + "cards": { + "criticalUnit": "da88d8560092eff8", + "deliveryFailures": "3295bdbfe78e3c01", + "deliveryNote": "31f287d5888f7ae8", + "failingNone": "d6a57f4f5d585d7f", + "failingNote": "797d69b6351b9d99", + "failingRules": "b1829503b811b86f", + "firingAlerts": "d9e293a211a1aeab", + "firingNone": "b3c57801a954c8df", + "incidentsNone": "1b303835cec87ac7", + "noAlertRules": "7364c4d5c6836fda", + "noEnabledRules": "3113508de7297ef5", + "noPipelines": "29cf5a22b101cdea", + "noStreams": "06e3aaafc23cc49c", + "oldestFiring": "4642274888e72100", + "openIncidents": "0021b3e95793e696", + "pipelineErrors": "ee50a619caee4c07", + "pipelinesOf": "6f1293ff044ce6fb", + "quietNote": "8e9fd378a75b7afc", + "quietStreams": "254025432ea334ec" + }, "colErrorRate": "33a224722fc7b1e4", "colLatency": "e0e7d293043c8cf6", "colReqs": "61985debc6417c53", + "dataAsOf": "559ef63277c32024", + "dataHealth": { + "compaction": "ea4d27f706d82e1a", + "compactionNote": "781dd2c3af67a05b", + "delivery": "3295bdbfe78e3c01", + "deliveryNote": "31f287d5888f7ae8", + "failing": "276f4b0bff63da18", + "failingNote": "4476ec65febe1e7a", + "noPipelines": "72183d31fea778ed", + "pipelinesActive": "ebc50e31e444d623", + "pipelinesErrored": "c8b743fe3a5980ab", + "pipelinesNote": "6cb0fae496a12e5e", + "quiet": "0a6365d180913968", + "quietNote": "e9aa7568a24f29e1", + "schema": "03a2d5849d4086cb", + "schemaNote": "9b199e7827a03769", + "status": { + "healthy": "7f1e323b6272c620", + "info": "170322a32f3c35b2", + "needsAction": "e382d466c25613b6", + "notSetUp": "94b768af74236a67", + "warning": "e981ddae45d8f4ca" + }, + "subtitle": "f7b738721f218463", + "viewPipelines": "da5dae86bafae8a8" + }, "daysAgo": "a41234397f3c63b2", - "emptyActionAlerts": "6ac085a35d9963e7", - "emptyActionAlertsDesc": "e6bc20e23b0b8156", + "empty": { + "anomalies": "bea29d1c1f129a3e", + "anomalyNoChecks": "b33b808d6c169339", + "incidents": "4d3495d2550aef11", + "recentEvents": "b3d8551f4c8f3f2d", + "services": "da1d75df527ee5a7" + }, "emptyActionLogs": "3f8105f14af710a4", "emptyActionLogsDesc": "635d266e4ad84bde", - "emptyActionTraces": "5a292d4771830cc9", - "emptyActionTracesDesc": "94a3bd51c2b7b54e", "hoursAgo": "520c737ea62430c4", "investigate": "289ddb3c038edede", + "level": { + "critical": "427dd2969bd140be", + "okFiring": "1d99d0eaede7d441", + "warning": "e981ddae45d8f4ca" + }, + "limit": { + "approachingLimit": "9f01b44bd1611c17", + "blockNote": "db801da4ce61853e", + "daysOver": "ac9d0939417258ff", + "daysOverOf": "8557fdf7bd6b3d31", + "expiredLine": "5de72173d91c820a", + "ofLimit": "6303f8c21765ae2a", + "ofUnlimited": "05cdf654143fbde8", + "overLimit": "66a81e7aafc36cc6", + "resets": "f11e2b11d0662d61", + "subtitle": "66ed51c1f474925a", + "unavailable": "c29fa9075eeecae4", + "used": "ffc22437b1c278a7" + }, "minutesAgo": "0a03167c208c5cd7", "notMonitoredAlert": "0fd0bf350a56af51", "notMonitoredAlertDesc": "0e7bd079c07ed0b8", @@ -20390,14 +20436,88 @@ "notMonitoredAnomalyDesc": "523aa0a36890ff4c", "notMonitoredDesc": "da8cc880b82d0c93", "notMonitoredTitle": "5ed64bcad471cff7", - "recentEvents": "9acf8ff6c2901930", + "recentEvents": "8ecce94dc301d033", "scrollLeft": "73befbd7d4f0f62a", "scrollRight": "6d42701317bf1cc0", "secondsAgo": "244073ecb2be8fe8", + "sections": { + "alerts": "f62e9267d196c761", + "alertsAndIncidents": "3421fbf2a4841e49", + "dataHealth": "6e5625b8c13013a0", + "limitAndStorage": "5fe708f05c7e2cc3", + "storage": "0db6a0a3bbfd6e3b", + "topStreams": "bb3ca650f2b0aa42" + }, "services": "604dce445e326d95", "showingOf": "f3d25c9265aac7c1", + "signals": { + "events": "8d14f6e72de8f18a", + "logs": "ea2100dc89ae9fe2", + "metrics": "a58da793c7250f1b", + "open": "ed077f3d8125d60d", + "openAria": { + "logs": "591889ac42e04819", + "metrics": "4975865997142959", + "traces": "ffde39c05cfe85ff" + }, + "samples": "bc228b269ce0d87f", + "spans": "b3393bfb2a55adff", + "stored": "0db6a0a3bbfd6e3b", + "streams": "bfd297afbf9e7430", + "traces": "998741a58ed2b411" + }, + "stale": { + "message": "ae19c182cdc829fe", + "retry": "942087cc2d41e013" + }, + "state": { + "sectionForbidden": "f03bf794d964b129", + "sectionLocked": "2fc5a6a63df6538b", + "tileError": "81bc72e54792e9fc", + "tileForbidden": "9b1f0823459d06f7", + "tileLocked": "a424e33d90931d1e", + "tileUnavailable": "67a926f7008e2b0e" + }, + "status": { + "attention": "c1ebc7817870e5be", + "checking": "ec963ffc911b8401", + "clear": "21ad0f972eb181fc", + "items": { + "failing": "36a2f30f9248ca02", + "firing": "5591cea01bf2d35e", + "incidents": "c61409d4f4cddc66", + "pipelines": "eb13ac2dc3a4b5d7", + "quiet": "1cac853c1d0dd6d3" + }, + "partial": "279219da8364b5fa", + "unmonitored": "f0a10061d1008a98" + }, + "storage": { + "bySignal": "7fe6b7b7ffaad7f4", + "emptyTitle": "e2e18a767a2b6db5", + "splitNote": "4a533a47757998fd", + "used": "0db6a0a3bbfd6e3b" + }, + "streams": { + "allStreams": "7ba01d16e615fc96", + "emptyDesc": "f100c4b2cb3c6942", + "emptyTitle": "eddceedac16ae5a7", + "shareOfTotal": "5892fce97211d1ec", + "shareUnderOne": "b540ca1185a12840", + "startIngesting": "ccd00cde6f176b41", + "subtitle": "0473127f6a49ead0", + "typeShare": "a1d4e3f6c7380584" + }, + "time": { + "daysHours": "be43ecf3668b6b23", + "hoursMinutes": "425dc6c72de6fdc5", + "minutes": "f7047179845887e8" + }, + "timeRange": { + "heading": "b7fdbb97bb7418c0" + }, "timesSymbol": "8db71ed28b0f2f14", - "untitledIncident": "4bea619d75d7fa3f", + "untitledIncident": "a3e4994bcd734dd5", "viewAll": "30a64216eaeac5d3", "viewingLabel": "b28f71ee763aebbe" }, diff --git a/src/api/http/src/handler/http/router/mod.rs b/src/api/http/src/handler/http/router/mod.rs index cf99e180e3..df3633e1c7 100644 --- a/src/api/http/src/handler/http/router/mod.rs +++ b/src/api/http/src/handler/http/router/mod.rs @@ -1344,6 +1344,7 @@ pub fn service_routes() -> Router { // Org info .route("/{org_id}/summary", get(organization::org::org_summary)) + .route("/{org_id}/overview", get(organization::org::org_overview)) .route("/{org_id}/ingest/recent_rejections", get(ingest_rejections::recent_rejections)) .route("/{org_id}/passcode", get(organization::org::get_user_passcode).put(organization::org::update_user_passcode)) .route("/{org_id}/rumtoken", get(organization::org::get_user_rumtoken).post(organization::org::create_user_rumtoken).put(organization::org::update_user_rumtoken)) diff --git a/src/api/http/src/handler/http/router/openapi.rs b/src/api/http/src/handler/http/router/openapi.rs index 86c21bb9f1..d0d53ecb7e 100644 --- a/src/api/http/src/handler/http/router/openapi.rs +++ b/src/api/http/src/handler/http/router/openapi.rs @@ -47,6 +47,7 @@ use crate::{ openobserve_api_management::request::organization::org::rename_org, openobserve_api_management::request::organization::assume_service_account::assume_service_account, openobserve_api_management::request::organization::org::org_summary, + openobserve_api_management::request::organization::org::org_overview, openobserve_api_management::request::organization::org::get_user_passcode, openobserve_api_management::request::organization::org::update_user_passcode, openobserve_api_management::request::organization::org::get_user_rumtoken, diff --git a/src/api/management/src/request/organization/org.rs b/src/api/management/src/request/organization/org.rs index cb9d93ccd0..20cf01280c 100644 --- a/src/api/management/src/request/organization/org.rs +++ b/src/api/management/src/request/organization/org.rs @@ -28,6 +28,7 @@ use openobserve_api_common::extractors::Headers; use openobserve_core::{ auth::{UserEmail, check_permissions}, organization::{self, get_passcode, get_rum_token, update_passcode, update_rum_token}, + overview, }; #[cfg(feature = "cloud")] use { @@ -384,6 +385,65 @@ pub async fn org_summary(Path(org_id): Path) -> impl IntoResponse { Json(org_summary) } +/// GetOrganizationOverview + +#[utoipa::path( + get, + path = "/{org_id}/overview", + context_path = "/api", + tag = "Organizations", + operation_id = "GetOrganizationOverview", + summary = "Get organization overview", + description = "Returns the streams, alerts, pipelines and license limits sections the Overview page needs in one call. Each section carries its own status (ok, error, locked or forbidden), so one failing section never fails the response, and results only include objects the caller may see. Responses are cached briefly per caller.", + security( + ("Authorization"= []) + ), + params( + ("org_id" = String, Path, description = "Organization name"), + ("top" = Option, Query, description = "Rows per list (top streams per type, firing alerts). Default 5, clamped to 1..20."), + ("refresh" = Option, Query, description = "Set true (or 1) to skip the short-lived cache and rebuild now."), + ), + responses( + (status = 200, description = "Success", content_type = "application/json", body = Object), + (status = 400, description = "Failure", content_type = "application/json", body = ()), + (status = 403, description = "Forbidden", content_type = "application/json", body = ()), + ), + extensions( + ("x-o2-ratelimit" = json!({"module": "Summary", "operation": "get"})), + ("x-o2-mcp" = json!({"enabled": false})) + ) +)] +pub async fn org_overview( + Path(org_id): Path, + Headers(user_email): Headers, + Query(query): Query>, +) -> Response { + let top = match parse_top(&query) { + Ok(top) => top, + Err(message) => return MetaHttpResponse::bad_request(message), + }; + let top_n = overview::normalize_top_n(top); + let refresh = parse_refresh(&query); + if !refresh && let Some(hit) = overview::cached_overview(&org_id, &user_email.user_id, top_n) { + return Json(hit).into_response(); + } + let scope = match overview::resolve_scope(&org_id, &user_email.user_id).await { + Ok(scope) => scope, + Err(e) => { + log::error!("[overview] failed to resolve permissions for {org_id}: {e}"); + return match overview::scope_unavailable(&org_id, &user_email.user_id, top_n).await { + Some(degraded) => Json(degraded).into_response(), + None => MetaHttpResponse::forbidden("Forbidden"), + }; + } + }; + if scope.denies_everything() { + return MetaHttpResponse::forbidden("Forbidden"); + } + let result = overview::get_overview(&org_id, &scope, top_n).await; + Json(result).into_response() +} + /// GetIngestToken #[utoipa::path( @@ -2007,14 +2067,58 @@ fn unknown_quota_pool_message(pool: &str) -> String { format!("unknown quota pool '{pool}' (expected one of: {expected})") } +fn parse_top(query: &HashMap) -> Result, &'static str> { + match query.get("top").map(|v| v.parse::()) { + Some(Err(_)) => Err("'top' must be a non-negative integer; values are clamped to 1..20"), + Some(Ok(n)) => Ok(Some(n)), + None => Ok(None), + } +} + +fn parse_refresh(query: &HashMap) -> bool { + query + .get("refresh") + .is_some_and(|v| v.eq_ignore_ascii_case("true") || v == "1") +} + #[cfg(test)] mod tests { + use std::collections::HashMap; + + use openobserve_core::overview; #[cfg(feature = "cloud")] use openobserve_core::trial_quota::TrialQuotaPool; #[cfg(feature = "cloud")] use super::*; + fn q(key: &str, value: &str) -> HashMap { + HashMap::from([(key.to_string(), value.to_string())]) + } + + #[test] + fn parse_top_handles_absent_invalid_and_numeric_values() { + assert_eq!(super::parse_top(&HashMap::new()), Ok(None)); + assert!(super::parse_top(&q("top", "abc")).is_err()); + assert!(super::parse_top(&q("top", "-1")).is_err()); + let top_n = |v: &str| overview::normalize_top_n(super::parse_top(&q("top", v)).unwrap()); + assert_eq!(top_n("0"), 1); + assert_eq!(top_n("7"), 7); + assert_eq!(top_n("500"), 20); + assert_eq!(overview::normalize_top_n(None), 5); + } + + #[test] + fn parse_refresh_accepts_only_true_or_one() { + for v in ["true", "1", "TRUE"] { + assert!(super::parse_refresh(&q("refresh", v)), "{v}"); + } + for v in ["false", "0", "yes", ""] { + assert!(!super::parse_refresh(&q("refresh", v)), "{v}"); + } + assert!(!super::parse_refresh(&HashMap::new())); + } + /// The route accepts every key in `ALL_POOLS`, so a hand-written list leaves an admin who /// typos the pool they want reading a 400 that never names it. #[cfg(feature = "cloud")] diff --git a/src/config/src/config.rs b/src/config/src/config.rs index 41dd45242b..bcee22b7c9 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -2863,6 +2863,12 @@ pub struct Limit { pub req_cols_per_record_limit: usize, #[env_config(name = "ZO_NODE_HEARTBEAT_TTL", default = 30)] // seconds pub node_heartbeat_ttl: i64, + #[env_config( + name = "ZO_OVERVIEW_CACHE_TTL", + default = 30, + help = "Seconds the org overview response is cached per caller; 0 disables the cache." + )] + pub overview_cache_ttl: u64, // How long an o2-ai session->owner claim survives. Must exceed the longest // expected conversation; mirrors o2-ai's O2_AI_SESSION_OWNER_TTL default. #[env_config(name = "ZO_AI_SESSION_OWNER_TTL", default = 86400)] // seconds diff --git a/src/config/src/router.rs b/src/config/src/router.rs index d2c9ad5e9c..b87f2e42ea 100644 --- a/src/config/src/router.rs +++ b/src/config/src/router.rs @@ -14,10 +14,11 @@ // along with this program. If not, see . /// usize indicates the number of parts to skip based on their actual paths. -const QUERIER_ROUTES: [(&str, usize); 47] = [ +const QUERIER_ROUTES: [(&str, usize); 48] = [ ("config", 0), // /config (unauthenticated bootstrap) ("config", 2), // /api/{org_id}/config (authenticated full) ("summary", 2), // /api/{org_id}/summary + ("overview", 2), // /api/{org_id}/overview ("organizations", 1), // /api/organizations ("settings", 2), // /api/{org_id}/settings/... ("schema", 3), // /api/{org_id}/streams/{stream_name}/schema @@ -67,6 +68,8 @@ const QUERIER_ROUTES: [(&str, usize); 47] = [ ("llm/models/built-in", 2), // /api/{org_id}/llm/models/built-in ("llm/models/refresh-built-in", 2), // /api/{org_id}/llm/models/refresh-built-in ]; +/// Entries that match their whole remaining path, so `overview_x/...` stays unrouted. +const EXACT_QUERIER_ROUTES: [&str; 1] = ["overview"]; const QUERIER_ROUTES_BY_BODY: [&str; 9] = [ "/_search", "/_search_partition", @@ -117,7 +120,12 @@ pub fn is_querier_route(path: &str) -> bool { } // Join remaining segments without collecting into a Vec let route_part = segments.collect::>().join("/"); - route_part.starts_with(route) + let matched = if EXACT_QUERIER_ROUTES.contains(route) { + extract_path_without_query(&route_part) == *route + } else { + route_part.starts_with(route) + }; + matched && INGESTER_ROUTES .iter() .all(|ingest_route| !route_part.ends_with(ingest_route)) @@ -159,7 +167,16 @@ pub fn is_public_dashboard_slug(slug: &str) -> bool { #[inline] pub fn is_fixed_querier_route(path: &str) -> bool { let path = remove_base_uri(path); - FIXED_QUERIER_ROUTES.iter().any(|x| path.contains(x)) + FIXED_QUERIER_ROUTES.iter().any(|x| path.contains(x)) || is_org_overview_route(path) +} + +/// Only `/api/{org}/overview` itself; a substring match would pin streams named `overview_*`. +fn is_org_overview_route(path: &str) -> bool { + let segments: Vec<&str> = extract_path_without_query(path) + .split('/') + .filter(|s| !s.is_empty()) + .collect(); + matches!(segments.as_slice(), ["api", _, "overview"]) } /// The public routes, and the status calls that answer from the querier owning a link's files. @@ -205,6 +222,14 @@ mod tests { assert!(is_querier_route("/api/org1/summary")); assert!(!is_querier_route("/summary")); // Should fail as it needs org_id + assert!(is_querier_route("/api/org1/overview")); + assert!(!is_querier_route("/overview")); + assert!(is_querier_route("/api/org1/overview?top=5")); + assert!(is_querier_route("/api/org1/overview/")); + assert!(!is_querier_route("/api/default/overview_x/_json")); + assert!(!is_querier_route("/api/default/overview_events/_search")); + assert!(!is_querier_route("/api/default/overview/extra")); + // Test streams route assert!(is_querier_route("/api/org1/streams")); assert!(is_querier_route("/api/org1/streams/mystream")); @@ -283,6 +308,7 @@ mod tests { #[test] fn test_is_fixed_querier_route() { assert!(is_fixed_querier_route("/summary")); + assert!(is_fixed_querier_route("/api/org1/overview")); assert!(is_fixed_querier_route("/schema")); assert!(is_fixed_querier_route("/streams")); @@ -290,6 +316,25 @@ mod tests { assert!(is_fixed_querier_route("/summary_other")); } + #[test] + fn test_overview_pin_matches_only_the_org_route() { + assert!(is_fixed_querier_route("/api/default/overview")); + assert!(is_fixed_querier_route("/api/default/overview/")); + assert!(is_fixed_querier_route("/api/default/overview?top=5")); + assert!(!is_fixed_querier_route( + "/api/default/overview_events/_search" + )); + for path in [ + "/api/default/streams/overview/schema", + "/api/default/streams/overview_x/_around", + "/api/default/overview_events/_search", + "/api/default/overview/extra", + "/overview", + ] { + assert!(!is_org_overview_route(path), "{path}"); + } + } + #[test] fn test_is_querier_route_ingester_exclusion() { // streams/_json is a querier-prefix route but ends with an ingester suffix diff --git a/src/core/src/lib.rs b/src/core/src/lib.rs index 02ada8c213..063da748e5 100644 --- a/src/core/src/lib.rs +++ b/src/core/src/lib.rs @@ -57,6 +57,7 @@ pub mod org_cleanup; #[cfg(feature = "cloud")] pub mod org_usage; pub mod organization; +pub mod overview; pub mod pipeline; pub mod profiles; pub mod prompts; diff --git a/src/core/src/overview.rs b/src/core/src/overview.rs new file mode 100644 index 0000000000..5b356d3c9f --- /dev/null +++ b/src/core/src/overview.rs @@ -0,0 +1,1610 @@ +// 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 +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::{ + collections::{BTreeMap, HashMap, HashSet}, + sync::LazyLock, +}; + +use common::meta::stream::Stream; +use config::{ + get_config, + meta::{ + alerts::state::AlertState, pipeline::Pipeline, + self_reporting::ai_chat::is_protected_ai_chat_stream, stream::StreamType, + }, + utils::time::now_micros, +}; +use dashmap::DashMap; +use futures::{StreamExt, TryStreamExt}; +use infra::{db::get_orm_client_ro, table::alerts::AlertSummaryRow}; +use serde::Serialize; +use stream::try_get_streams; + +const DEFAULT_TOP_N: usize = 5; +const MAX_TOP_N: usize = 20; +/// User-facing stream types; derived and internal streams are left out. +const OVERVIEW_STREAM_TYPES: [StreamType; 4] = [ + StreamType::Logs, + StreamType::Metrics, + StreamType::Traces, + StreamType::EnrichmentTables, +]; + +/// Stats refresh lags ingestion, so a shorter gap would flag streams that are still writing. +const QUIET_AFTER_MICROS: i64 = 3_600 * 1_000_000; +/// Streams silent longer than this are treated as abandoned, not newly quiet. +const QUIET_FORGOTTEN_AFTER_MICROS: i64 = 7 * 24 * 3_600 * 1_000_000; +const DAY_MICROS: i64 = 24 * 3_600 * 1_000_000; + +const ALERT_ID_CHUNK: usize = 500; +const ALERT_STATE_CONCURRENCY: usize = 4; +const CACHE_MAX_ENTRIES: usize = 1024; +const ALERT_STATE_KEYS: [&str; 9] = [ + "firing", + "pending", + "normal", + "succeeded", + "error", + "skipped", + "notify_failed", + "suppressed", + "never_evaluated", +]; +const ALERT_LEVEL_KEYS: [&str; 5] = ["ok", "warning", "critical", "no_data", "unclassified"]; + +/// Entries are keyed per caller so one user's visible set is never served to another. +static CACHE: LazyLock> = LazyLock::new(DashMap::new); + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +struct CacheKey { + org_id: String, + user_id: String, + top_n: usize, +} + +/// What the caller may see; a stream type or pipeline list that is absent is unrestricted. +#[derive(Clone, Debug, Default)] +pub struct OverviewScope { + pub user_id: String, + pub stream_grants: HashMap>, + pub pipeline_grants: Option>, + pub alerts_forbidden: bool, +} + +impl OverviewScope { + /// True when the caller may see none of the streams, alerts or pipelines. + pub fn denies_everything(&self) -> bool { + streams_forbidden(self) + && self.alerts_forbidden + && self.pipeline_grants.as_ref().is_some_and(Vec::is_empty) + } +} + +/// Module-level LIST decisions, matching the auth middleware on the sibling list routes. +struct ModuleAccess { + streams: Vec, + alerts: bool, + pipelines: bool, +} + +impl ModuleAccess { + #[cfg(any(not(feature = "enterprise"), test))] + fn all() -> Self { + Self { + streams: OVERVIEW_STREAM_TYPES.to_vec(), + alerts: true, + pipelines: true, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum SectionStatus { + Ok, + Error, + Locked, + Forbidden, +} + +#[derive(Clone, Debug, Serialize)] +pub struct Section { + pub status: SectionStatus, + #[serde(skip_serializing_if = "Option::is_none")] + pub data: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, + /// Why a section is locked, so the UI can pick the right message. + #[serde(skip_serializing_if = "Option::is_none")] + pub reason: Option, +} + +impl Section { + fn ok(data: T) -> Self { + Self { + status: SectionStatus::Ok, + data: Some(data), + error: None, + reason: None, + } + } + + fn error(message: &str) -> Self { + Self { + status: SectionStatus::Error, + data: None, + error: Some(message.to_string()), + reason: None, + } + } + + fn forbidden() -> Self { + Self { + status: SectionStatus::Forbidden, + data: None, + error: None, + reason: None, + } + } + + #[cfg(any(not(feature = "enterprise"), feature = "cloud", test))] + fn locked(reason: &str) -> Self { + Self { + status: SectionStatus::Locked, + data: None, + error: None, + reason: Some(reason.to_string()), + } + } + + fn from_result(result: Result, what: &str) -> Self { + match result { + Ok(data) => Self::ok(data), + Err(e) => { + log::error!("[overview] failed to load {what}: {e}"); + Self::error(&format!("failed to load {what}")) + } + } + } +} + +#[derive(Clone, Debug, Serialize)] +pub struct Overview { + pub streams: Section, + pub alerts: Section, + pub pipelines: Section, + pub limits: Section, + pub generated_at: i64, + pub cached: bool, +} + +impl Overview { + fn has_error(&self) -> bool { + self.streams.status == SectionStatus::Error + || self.alerts.status == SectionStatus::Error + || self.pipelines.status == SectionStatus::Error + || self.limits.status == SectionStatus::Error + } +} + +#[derive(Clone, Debug, Serialize)] +pub struct StreamsData { + pub total_streams: u64, + pub total_storage_size_mb: f64, + pub by_type: Vec, + pub quiet_count: usize, + /// Schema versions created in the last 24 hours across the visible user streams. + pub schema_changes_24h: u64, + /// `None` when the lookup fails. + pub compaction_pending: Option, +} + +#[derive(Clone, Debug, Serialize)] +pub struct StreamTypeSummary { + pub stream_type: String, + pub count: u64, + pub doc_num: i64, + pub storage_size_mb: f64, + pub compressed_size_mb: f64, + pub top: Vec, +} + +#[derive(Clone, Debug, Serialize)] +pub struct TopStream { + pub name: String, + pub doc_num: i64, + pub storage_size_mb: f64, +} + +#[derive(Clone, Debug, Serialize)] +pub struct AlertsData { + pub total: u64, + pub enabled: u64, + pub disabled: u64, + /// Enabled rules by the outcome of their last evaluation. + pub by_state: BTreeMap, + /// Enabled rules by the severity level of their last classification. + pub by_level: BTreeMap, + /// Same count as `by_state.error`; do not add the two together. + pub failed_evaluation: u64, + /// Oldest-firing first, capped at `top`. + pub firing: Vec, + /// Same keys as `by_level`, counting only firing rules, uncapped. + pub firing_by_level: BTreeMap, +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct FiringAlert { + pub name: String, + pub stream: String, + pub level: String, + pub since: i64, +} + +#[derive(Clone, Debug, Serialize)] +pub struct PipelinesData { + pub total: u64, + pub enabled: u64, + pub disabled: u64, + /// Enabled pipelines that still have an error row the retention cleanup has not removed. + pub errored: u64, +} + +/// Deployment-wide license ingestion figures, not per org. +#[derive(Clone, Debug, Serialize)] +pub struct LimitsData { + pub scope: &'static str, + pub ingestion_used_percent: f64, + pub expired: bool, + pub days_over_limit: u8, + pub used_today_gb: Option, + pub limit_today_gb: Option, + pub allowed_days: Option, + /// Next UTC midnight, when the daily ingestion counter resets. + pub resets_at: i64, +} + +pub fn normalize_top_n(requested: Option) -> usize { + requested.unwrap_or(DEFAULT_TOP_N).clamp(1, MAX_TOP_N) +} + +/// Microsecond timestamp of the next UTC midnight after `now`. +pub fn next_utc_midnight_micros(now: i64) -> i64 { + (now.div_euclid(DAY_MICROS) + 1) * DAY_MICROS +} + +/// Resolves which streams and pipelines the caller may see; errors when the grant lookup fails. +pub async fn resolve_scope(org_id: &str, user_id: &str) -> Result { + let stream_lookups = futures::future::try_join_all(OVERVIEW_STREAM_TYPES.map(|stream_type| { + let resource = stream_type.to_string(); + async move { list_grants(org_id, user_id, &resource).await } + })); + let (stream_results, pipeline_grants) = + tokio::try_join!(stream_lookups, list_grants(org_id, user_id, "pipelines"))?; + let stream_grants = OVERVIEW_STREAM_TYPES + .into_iter() + .zip(stream_results) + .filter_map(|(stream_type, grants)| grants.map(|g| (stream_type, g))) + .collect(); + let mut scope = OverviewScope { + user_id: user_id.to_string(), + stream_grants, + pipeline_grants, + alerts_forbidden: false, + }; + apply_module_access(&mut scope, &module_access(org_id, user_id).await); + Ok(scope) +} + +/// Overview when grants cannot be resolved; `None` when the module LIST check denies everything. +pub async fn scope_unavailable(org_id: &str, user_id: &str, top_n: usize) -> Option { + let mut scope = OverviewScope::default(); + apply_module_access(&mut scope, &module_access(org_id, user_id).await); + if scope.denies_everything() { + return None; + } + let (alerts, limits) = tokio::join!( + async { + if scope.alerts_forbidden { + return Section::forbidden(); + } + Section::from_result(load_alerts(org_id, user_id, top_n).await, "alerts") + }, + limits_section(), + ); + Some(Overview { + streams: unresolved_section(streams_forbidden(&scope)), + alerts, + pipelines: unresolved_section(scope.pipeline_grants.as_ref().is_some_and(Vec::is_empty)), + limits, + generated_at: now_micros(), + cached: false, + }) +} + +/// Returns a fresh cached overview for the caller, so a hit can skip the permission lookups. +pub fn cached_overview(org_id: &str, user_id: &str, top_n: usize) -> Option { + let ttl_micros = cache_ttl_micros(); + if ttl_micros <= 0 { + return None; + } + let key = CacheKey { + org_id: org_id.to_string(), + user_id: user_id.to_string(), + top_n, + }; + let hit = CACHE.get(&key)?; + if now_micros() - hit.generated_at >= ttl_micros { + return None; + } + let mut overview = hit.clone(); + overview.cached = true; + Some(overview) +} + +/// Builds the overview and caches it per caller when every section loaded. +pub async fn get_overview(org_id: &str, scope: &OverviewScope, top_n: usize) -> Overview { + let ttl_micros = cache_ttl_micros(); + let overview = build_overview(org_id, scope, top_n).await; + if ttl_micros > 0 && !overview.has_error() { + let key = CacheKey { + org_id: org_id.to_string(), + user_id: scope.user_id.clone(), + top_n, + }; + cache_put(key, overview.clone(), ttl_micros); + } + overview +} + +/// A section the module check denied stays forbidden; an allowed one failed to resolve. +fn unresolved_section(denied: bool) -> Section { + if denied { + Section::forbidden() + } else { + Section::error("failed to resolve permissions") + } +} + +/// The module LIST check gates a section even when per-object grants exist, as on list routes. +fn apply_module_access(scope: &mut OverviewScope, access: &ModuleAccess) { + for stream_type in OVERVIEW_STREAM_TYPES { + if !access.streams.contains(&stream_type) { + scope.stream_grants.insert(stream_type, Vec::new()); + } + } + if !access.pipelines { + scope.pipeline_grants = Some(Vec::new()); + } + scope.alerts_forbidden = !access.alerts; +} + +#[cfg(feature = "enterprise")] +async fn module_access(org_id: &str, user_id: &str) -> ModuleAccess { + let role = db::user::get(Some(org_id), user_id) + .await + .ok() + .flatten() + .map_or_else(common::meta::user::get_default_user_role, |user| user.role); + let role = &role; + let stream_checks = + futures::future::join_all(OVERVIEW_STREAM_TYPES.map(|stream_type| async move { + let path = [org_id, "streams"]; + ( + stream_type, + can_list(org_id, user_id, role, &path, Some(stream_type)).await, + ) + })); + let alerts_path = ["v2", org_id, "alerts"]; + let pipelines_path = [org_id, "pipelines"]; + let (stream_checks, alerts, pipelines) = tokio::join!( + stream_checks, + can_list(org_id, user_id, role, &alerts_path, None), + can_list(org_id, user_id, role, &pipelines_path, None), + ); + let streams = stream_checks + .into_iter() + .filter_map(|(stream_type, allowed)| allowed.then_some(stream_type)) + .collect(); + ModuleAccess { + streams, + alerts, + pipelines, + } +} + +#[cfg(not(feature = "enterprise"))] +async fn module_access(_org_id: &str, _user_id: &str) -> ModuleAccess { + ModuleAccess::all() +} + +/// Mirrors the middleware list check; alerts need LIST on all folders, stricter than one folder. +#[cfg(feature = "enterprise")] +async fn can_list( + org_id: &str, + user_id: &str, + role: &config::meta::user::UserRole, + path: &[&str], + stream_type: Option, +) -> bool { + use o2_openfga::meta::route_permissions::{self as rp, StreamType as RpStreamType}; + + let rp_type = stream_type.map(|st| match st { + StreamType::Metrics => RpStreamType::Metrics, + StreamType::Traces => RpStreamType::Traces, + StreamType::Profiles => RpStreamType::Profiles, + StreamType::EnrichmentTables => RpStreamType::EnrichmentTables, + _ => RpStreamType::Logs, + }); + let Some(route) = rp::resolve_permission(path, "GET", org_id, "", rp_type) else { + return false; + }; + if route.bypass_check { + return true; + } + let auth_info = crate::auth::AuthExtractor { + auth: String::new(), + method: route.method, + o2_type: route.o2_type, + org_id: route.org_id, + bypass_check: false, + parent_id: route.parent_id, + use_all_org: route.use_all_org, + use_self_context: route.use_self_context, + use_self_parent: route.use_self_parent, + feature: route.feature, + }; + crate::authz::check_permissions(user_id, auth_info, role.clone(), true).await +} + +async fn list_grants( + org_id: &str, + user_id: &str, + resource: &str, +) -> Result>, anyhow::Error> { + crate::authz::list_objects_for_user(org_id, user_id, "GET", &ofga_object_type(resource)).await +} + +async fn build_overview(org_id: &str, scope: &OverviewScope, top_n: usize) -> Overview { + let (streams, alerts, pipelines, limits) = tokio::join!( + async { + if streams_forbidden(scope) { + return Section::forbidden(); + } + Section::from_result(load_streams(org_id, scope, top_n).await, "streams") + }, + async { + if scope.alerts_forbidden { + return Section::forbidden(); + } + Section::from_result(load_alerts(org_id, &scope.user_id, top_n).await, "alerts") + }, + async { + if scope.pipeline_grants.as_ref().is_some_and(Vec::is_empty) { + return Section::forbidden(); + } + let loaded = load_pipelines(org_id, scope.pipeline_grants.clone()).await; + Section::from_result(loaded, "pipelines") + }, + limits_section(), + ); + Overview { + streams, + alerts, + pipelines, + limits, + generated_at: now_micros(), + cached: false, + } +} + +fn streams_forbidden(scope: &OverviewScope) -> bool { + OVERVIEW_STREAM_TYPES + .iter() + .all(|t| scope.stream_grants.get(t).is_some_and(Vec::is_empty)) +} + +fn cache_ttl_micros() -> i64 { + ttl_secs_to_micros(get_config().limit.overview_cache_ttl) +} + +fn ttl_secs_to_micros(secs: u64) -> i64 { + i64::try_from(secs) + .unwrap_or(i64::MAX) + .saturating_mul(1_000_000) +} + +fn cache_put(key: CacheKey, overview: Overview, ttl_micros: i64) { + make_room(&CACHE, now_micros(), ttl_micros, CACHE_MAX_ENTRIES); + CACHE.insert(key, overview); +} + +/// Drops expired entries, then the oldest ones, so a full cache never wipes fresh entries +/// wholesale. +fn make_room(cache: &DashMap, now: i64, ttl_micros: i64, max: usize) { + if cache.len() < max { + return; + } + cache.retain(|_, v| now - v.generated_at < ttl_micros); + while cache.len() >= max { + let oldest = cache + .iter() + .min_by_key(|e| e.value().generated_at) + .map(|e| e.key().clone()); + match oldest { + Some(key) => cache.remove(&key), + None => break, + }; + } +} + +#[cfg(feature = "enterprise")] +fn ofga_object_type(name: &str) -> String { + use o2_openfga::meta::mapping::OFGA_MODELS; + + OFGA_MODELS + .get(name) + .map_or(name, |model| model.key) + .to_string() +} + +#[cfg(not(feature = "enterprise"))] +fn ofga_object_type(name: &str) -> String { + name.to_string() +} + +async fn load_streams( + org_id: &str, + scope: &OverviewScope, + top_n: usize, +) -> Result { + let streams = fetch_streams(org_id, scope).await?; + let visible_types: HashSet = + streams.iter().map(|s| s.stream_type.to_string()).collect(); + let now = now_micros(); + let schema_changes_24h = count_schema_changes(org_id, &streams, now - DAY_MICROS).await; + let mut data = build_streams_data(streams, top_n, now); + data.schema_changes_24h = schema_changes_24h; + data.compaction_pending = pending_compaction(org_id, &visible_types).await; + Ok(data) +} + +/// Per-type lookups run concurrently; per-type grants stay separate and type order is kept. +async fn fetch_streams(org_id: &str, scope: &OverviewScope) -> Result, anyhow::Error> { + let per_type = futures::future::try_join_all(OVERVIEW_STREAM_TYPES.map(|stream_type| { + let grants = scope.stream_grants.get(&stream_type).cloned(); + async move { try_get_streams(org_id, Some(stream_type), false, grants).await } + })) + .await?; + Ok(per_type.into_iter().flatten().collect()) +} + +async fn count_schema_changes(org_id: &str, streams: &[Stream], since: i64) -> u64 { + let counted: Vec<(StreamType, &str)> = streams + .iter() + .filter(|s| is_counted_user_stream(s)) + .map(|s| (s.stream_type, s.name.as_str())) + .collect(); + infra::schema::count_cached_versions_since(org_id, &counted, since).await as u64 +} + +async fn pending_compaction(org_id: &str, visible_types: &HashSet) -> Option { + match infra::file_list::get_pending_jobs_count().await { + Ok(counts) => Some(sum_pending_for_org(&counts, org_id, visible_types)), + Err(e) => { + log::warn!("[overview] failed to load pending compaction jobs: {e}"); + None + } + } +} + +fn sum_pending_for_org( + counts: &HashMap>, + org_id: &str, + visible_types: &HashSet, +) -> i64 { + counts.get(org_id).map_or(0, |by_type| { + by_type + .iter() + .filter(|(t, _)| visible_types.contains(*t)) + .map(|(_, n)| n) + .sum() + }) +} + +/// User data streams only; enrichment tables, AI chat and internal streams are left out. +fn is_counted_user_stream(stream: &Stream) -> bool { + infra::schema::is_user_data_stream(stream.stream_type, &stream.name) + && !is_protected_ai_chat_stream(&stream.name) +} + +fn quiet_count(streams: &[Stream], now: i64) -> usize { + streams + .iter() + .filter(|s| is_counted_user_stream(s)) + .filter(|s| { + let last = s.stats.doc_time_max; + last < now - QUIET_AFTER_MICROS && last > now - QUIET_FORGOTTEN_AFTER_MICROS + }) + .count() +} + +fn build_streams_data(streams: Vec, top_n: usize, now: i64) -> StreamsData { + let mut by_type = Vec::with_capacity(OVERVIEW_STREAM_TYPES.len()); + for stream_type in OVERVIEW_STREAM_TYPES { + let mut of_type: Vec<&Stream> = streams + .iter() + .filter(|s| s.stream_type == stream_type && !is_protected_ai_chat_stream(&s.name)) + .collect(); + of_type.sort_by(|a, b| { + b.stats + .storage_size + .total_cmp(&a.stats.storage_size) + .then_with(|| a.name.cmp(&b.name)) + }); + by_type.push(StreamTypeSummary { + stream_type: stream_type.to_string(), + count: of_type.len() as u64, + doc_num: of_type.iter().map(|s| s.stats.doc_num).sum(), + storage_size_mb: of_type + .iter() + .fold(0.0, |acc, s| acc + s.stats.storage_size), + compressed_size_mb: of_type + .iter() + .fold(0.0, |acc, s| acc + s.stats.compressed_size), + top: of_type + .iter() + .take(top_n) + .map(|s| TopStream { + name: s.name.clone(), + doc_num: s.stats.doc_num, + storage_size_mb: s.stats.storage_size, + }) + .collect(), + }); + } + let quiet_count = quiet_count(&streams, now); + StreamsData { + total_streams: by_type.iter().map(|t| t.count).sum(), + total_storage_size_mb: by_type.iter().fold(0.0, |acc, t| acc + t.storage_size_mb), + by_type, + quiet_count, + schema_changes_24h: 0, + compaction_pending: None, + } +} + +async fn load_alerts( + org_id: &str, + user_id: &str, + top_n: usize, +) -> Result { + let conn = get_orm_client_ro().await; + let permissions = crate::alerts::alert::permitted_alerts(org_id, Some(user_id), None).await?; + let alerts = filter_permitted_alerts( + infra::table::alerts::list_summary_rows(conn, org_id).await?, + permissions, + org_id, + ); + let ids: Vec = alerts + .iter() + .filter(|a| a.enabled) + .map(|a| a.id.clone()) + .collect(); + let chunk_states: Vec> = futures::stream::iter( + ids.chunks(ALERT_ID_CHUNK) + .map(infra::table::alert_states::get_rollups), + ) + .buffered(ALERT_STATE_CONCURRENCY) + .try_collect() + .await?; + let states: Vec = chunk_states.into_iter().flatten().collect(); + Ok(build_alerts_data(&alerts, &states, top_n)) +} + +/// Mirrors `alerts::alert::list_v2`: `None` permissions means no per-object filtering. +fn filter_permitted_alerts( + alerts: Vec, + permissions: Option>, + org_id: &str, +) -> Vec { + let Some(permissions) = permissions else { + return alerts; + }; + if permissions.contains(&format!("alert:_all_{org_id}")) { + return alerts; + } + alerts + .into_iter() + .filter(|a| { + permissions.contains(&format!("alert:{}", a.name)) + || permissions.contains(&format!("alert:{}/{}", a.folder_id, a.id)) + || permissions.contains(&format!("alert:{}", a.id)) + }) + .collect() +} + +fn build_alerts_data( + alerts: &[AlertSummaryRow], + states: &[AlertState], + top_n: usize, +) -> AlertsData { + let by_id: HashMap = + states.iter().map(|s| (s.alert_id.clone(), s)).collect(); + let mut by_state: BTreeMap = ALERT_STATE_KEYS + .iter() + .map(|k| (k.to_string(), 0)) + .collect(); + let mut by_level = zeroed_levels(); + let mut firing_by_level = zeroed_levels(); + let mut firing = Vec::new(); + let mut enabled = 0; + let mut failed_evaluation = 0; + for alert in alerts.iter().filter(|a| a.enabled) { + enabled += 1; + let state = by_id.get(&alert.id).copied(); + let outcome = state.and_then(|s| s.last_outcome.as_ref()); + let state_key = outcome.map_or("never_evaluated".to_string(), |o| o.to_string()); + *by_state.entry(state_key).or_insert(0) += 1; + let level_key = state + .and_then(|s| s.level) + .map_or("unclassified".to_string(), |l| l.to_string()); + if let Some(state) = + state.filter(|s| s.last_outcome.as_ref().is_some_and(|o| o.is_firing())) + { + *firing_by_level.entry(level_key.clone()).or_insert(0) += 1; + firing.push(FiringAlert { + name: alert.name.clone(), + stream: alert.stream_name.clone(), + level: level_key.clone(), + since: state.since.or(state.last_outcome_at).unwrap_or(0), + }); + } + *by_level.entry(level_key).or_insert(0) += 1; + if outcome.is_some_and(|o| o.as_str() == "error") { + failed_evaluation += 1; + } + } + // since 0 means unknown, so those sort after every known start + firing.sort_by_key(|f| (f.since <= 0, f.since, f.name.clone())); + firing.truncate(top_n); + AlertsData { + total: alerts.len() as u64, + enabled, + disabled: alerts.len() as u64 - enabled, + by_state, + by_level, + failed_evaluation, + firing, + firing_by_level, + } +} + +fn zeroed_levels() -> BTreeMap { + ALERT_LEVEL_KEYS + .iter() + .map(|k| (k.to_string(), 0)) + .collect() +} + +async fn load_pipelines( + org_id: &str, + grants: Option>, +) -> Result { + let pipelines = crate::pipeline::list_user_pipelines(org_id, grants).await?; + let errored: HashSet = db::pipeline_errors::list_by_org(org_id) + .await? + .into_iter() + .map(|e| e.pipeline_id) + .collect(); + Ok(build_pipelines_data(&pipelines, &errored)) +} + +fn build_pipelines_data(pipelines: &[Pipeline], errored: &HashSet) -> PipelinesData { + let enabled = pipelines.iter().filter(|p| p.enabled).count() as u64; + PipelinesData { + total: pipelines.len() as u64, + enabled, + disabled: pipelines.len() as u64 - enabled, + errored: pipelines + .iter() + .filter(|p| p.enabled && errored.contains(&p.id)) + .count() as u64, + } +} + +#[cfg(all(feature = "enterprise", not(feature = "cloud")))] +async fn limits_section() -> Section { + use o2_enterprise::enterprise::license::{ + ingestion_limit_exceeded_count, ingestion_used, license_expired, license_status, + }; + + let decision = license_status().await.decision; + Section::ok(LimitsData { + scope: "deployment", + ingestion_used_percent: ingestion_used() * 100.0, + expired: license_expired().await, + days_over_limit: ingestion_limit_exceeded_count(), + used_today_gb: Some(decision.used_today_gb), + limit_today_gb: decision.limit_today_gb, + allowed_days: Some(decision.allowed_days), + resets_at: next_utc_midnight_micros(now_micros()), + }) +} + +#[cfg(any(not(feature = "enterprise"), feature = "cloud"))] +async fn limits_section() -> Section { + Section::locked("not available in this edition") +} + +#[cfg(test)] +mod tests { + use arrow_schema::{DataType, Field, Schema}; + use config::meta::{ + alerts::level::AlertLevel, pipeline::Pipeline, self_reporting::usage::RunOutcome, + stream::StreamStats, + }; + use infra::schema::{STREAM_SCHEMAS_LATEST, SchemaCache, insert_latest}; + + use super::*; + + const MIN: i64 = 60 * 1_000_000; + + fn stream(name: &str, stream_type: StreamType, size: f64, docs: i64) -> Stream { + Stream { + name: name.to_string(), + storage_type: "disk".to_string(), + stream_type, + stats: StreamStats { + storage_size: size, + compressed_size: size / 2.0, + doc_num: docs, + ..Default::default() + }, + schema: vec![], + uds_schema: vec![], + settings: Default::default(), + metrics_meta: None, + total_fields: 0, + pattern_associations: vec![], + is_derived: None, + is_internal: false, + } + } + + fn alert(id: &str, enabled: bool) -> AlertSummaryRow { + named_alert(id, "", enabled) + } + + fn named_alert(id: &str, name: &str, enabled: bool) -> AlertSummaryRow { + AlertSummaryRow { + id: id.to_string(), + name: name.to_string(), + folder_id: "default".to_string(), + stream_name: "s".to_string(), + enabled, + } + } + + fn state( + alert: &AlertSummaryRow, + outcome: Option, + level: Option, + ) -> AlertState { + let mut s = AlertState::empty(&alert.id, ""); + s.last_outcome = outcome; + s.level = level; + s + } + + fn pipeline(id: &str, enabled: bool) -> Pipeline { + Pipeline { + id: id.to_string(), + version: 0, + enabled, + org: String::new(), + name: id.to_string(), + description: String::new(), + source: Default::default(), + kind: Default::default(), + nodes: vec![], + edges: vec![], + } + } + + #[test] + fn test_normalize_top_n() { + assert_eq!(normalize_top_n(None), DEFAULT_TOP_N); + assert_eq!(normalize_top_n(Some(0)), 1); + assert_eq!(normalize_top_n(Some(7)), 7); + assert_eq!(normalize_top_n(Some(500)), MAX_TOP_N); + } + + #[test] + fn test_streams_section_totals_and_top_n_per_type() { + let streams = vec![ + stream("small", StreamType::Logs, 1.0, 10), + stream("big", StreamType::Logs, 9.0, 90), + stream("mid", StreamType::Logs, 5.0, 50), + stream("m1", StreamType::Metrics, 2.0, 20), + stream("internal", StreamType::Index, 100.0, 1), + ]; + let data = build_streams_data(streams, 2, now_micros()); + assert_eq!(data.total_streams, 4); + assert_eq!(data.total_storage_size_mb, 17.0); + let logs = &data.by_type[0]; + assert_eq!(logs.stream_type, "logs"); + assert_eq!(logs.count, 3); + assert_eq!(logs.doc_num, 150); + let top: Vec<&str> = logs.top.iter().map(|t| t.name.as_str()).collect(); + assert_eq!(top, vec!["big", "mid"]); + assert_eq!(data.by_type[2].count, 0); + assert!(data.by_type[2].storage_size_mb.is_sign_positive()); + assert!(data.by_type[2].compressed_size_mb.is_sign_positive()); + } + + #[test] + fn test_alerts_section_counts_by_state_level_and_failures() { + let a1 = alert("2Z8rVTbWbyUjb3hIibV02e0sqHW", true); + let a2 = alert("2Z8rVTbWbyUjb3hIibV02e0sqHX", true); + let a3 = alert("2Z8rVTbWbyUjb3hIibV02e0sqHY", true); + let a4 = alert("2Z8rVTbWbyUjb3hIibV02e0sqHZ", false); + let states = vec![ + state(&a1, Some(RunOutcome::Firing), Some(AlertLevel::Critical)), + state(&a2, Some(RunOutcome::Error), None), + state(&a4, Some(RunOutcome::Firing), Some(AlertLevel::Warning)), + ]; + let data = build_alerts_data(&[a1, a2, a3, a4], &states, 5); + assert_eq!((data.total, data.enabled, data.disabled), (4, 3, 1)); + assert_eq!(data.by_state["firing"], 1); + assert_eq!(data.by_state["error"], 1); + assert_eq!(data.by_state["never_evaluated"], 1); + assert_eq!(data.by_level["critical"], 1); + assert_eq!(data.by_level["unclassified"], 2); + assert_eq!(data.failed_evaluation, 1); + } + + #[test] + fn test_pipelines_section_counts_errored_among_given_pipelines() { + let pipelines = vec![ + pipeline("p1", true), + pipeline("p2", false), + pipeline("p3", true), + ]; + let errored: HashSet = + ["p1".to_string(), "p2".to_string(), "hidden".to_string()].into(); + let data = build_pipelines_data(&pipelines, &errored); + assert_eq!( + (data.total, data.enabled, data.disabled, data.errored), + (3, 2, 1, 1) + ); + } + + #[test] + fn test_section_serialization_carries_status_and_omits_empty_fields() { + let locked = serde_json::to_value(Section::::locked("no license")).unwrap(); + assert_eq!( + locked, + serde_json::json!({"status": "locked", "reason": "no license"}) + ); + let err = serde_json::to_value(Section::::error("boom")).unwrap(); + assert_eq!(err, serde_json::json!({"status": "error", "error": "boom"})); + } + + #[cfg(any(not(feature = "enterprise"), feature = "cloud"))] + #[tokio::test] + async fn test_limits_locked_without_enterprise() { + let limits = limits_section().await; + assert_eq!(limits.status, SectionStatus::Locked); + assert_eq!( + limits.reason.as_deref(), + Some("not available in this edition") + ); + } + + #[tokio::test] + async fn test_permitted_list_is_honoured_without_a_stream_type() { + let org = "overview_no_type_org"; + seed_stream(org, StreamType::Logs, "visible", 1.0).await; + seed_stream(org, StreamType::Logs, "hidden", 2.0).await; + seed_stream(org, StreamType::Metrics, "hidden_metric", 3.0).await; + let permitted = Some(vec!["logs:visible".to_string()]); + let streams = try_get_streams(org, None, false, permitted).await.unwrap(); + let names: Vec<&str> = streams.iter().map(|s| s.name.as_str()).collect(); + assert_eq!(names, vec!["visible"]); + } + + #[test] + fn test_streams_forbidden_only_when_every_type_has_empty_grants() { + let none: HashMap> = + OVERVIEW_STREAM_TYPES.iter().map(|t| (*t, vec![])).collect(); + let mut scope = OverviewScope { + user_id: "u".to_string(), + stream_grants: none, + pipeline_grants: None, + alerts_forbidden: false, + }; + assert!(streams_forbidden(&scope)); + scope + .stream_grants + .insert(StreamType::Logs, vec!["logs:app".to_string()]); + assert!(!streams_forbidden(&scope)); + scope.stream_grants.remove(&StreamType::Logs); + assert!(!streams_forbidden(&scope)); + } + + async fn seed_stream(org: &str, stream_type: StreamType, name: &str, size: f64) { + let schema = Schema::new(vec![Field::new("f", DataType::Utf8, true)]); + { + let mut w = STREAM_SCHEMAS_LATEST.write().await; + insert_latest( + &mut w, + format!("{org}/{stream_type}/{name}"), + SchemaCache::new(schema), + ); + } + infra::cache::stats::set_stream_stats( + org, + name, + stream_type, + StreamStats { + storage_size: size * config::SIZE_IN_MB, + compressed_size: size * config::SIZE_IN_MB, + doc_num: 1, + ..Default::default() + }, + ); + } + + #[tokio::test] + async fn test_restricted_caller_never_sees_hidden_streams() { + let org = "overview_rbac_org"; + seed_stream(org, StreamType::Logs, "visible", 1.0).await; + seed_stream(org, StreamType::Logs, "secret", 1000.0).await; + seed_stream(org, StreamType::Metrics, "secret_metric", 500.0).await; + + let restricted = OverviewScope { + user_id: "limited@example.com".to_string(), + stream_grants: HashMap::from([ + (StreamType::Logs, vec!["logs:visible".to_string()]), + (StreamType::Metrics, vec![]), + ]), + pipeline_grants: None, + alerts_forbidden: false, + }; + let data = load_streams(org, &restricted, 5).await.unwrap(); + assert_eq!(data.total_streams, 1); + assert_eq!(data.total_storage_size_mb, 1.0); + let logs = &data.by_type[0]; + assert_eq!((logs.count, logs.storage_size_mb), (1, 1.0)); + let names: Vec<&str> = logs.top.iter().map(|t| t.name.as_str()).collect(); + assert_eq!(names, vec!["visible"]); + assert_eq!(data.by_type[1].count, 0); + assert!(data.by_type[1].top.is_empty()); + + let unrestricted = OverviewScope { + user_id: "admin@example.com".to_string(), + ..Default::default() + }; + let data = load_streams(org, &unrestricted, 5).await.unwrap(); + assert_eq!(data.total_streams, 3); + assert_eq!(data.total_storage_size_mb, 1501.0); + } + + fn blank_scope() -> OverviewScope { + OverviewScope { + user_id: "u".to_string(), + ..Default::default() + } + } + + #[test] + fn test_summary_only_caller_gets_streams_alerts_pipelines_forbidden() { + let mut scope = blank_scope(); + let none = ModuleAccess { + streams: vec![], + alerts: false, + pipelines: false, + }; + apply_module_access(&mut scope, &none); + assert!(streams_forbidden(&scope)); + assert!(scope.alerts_forbidden); + assert_eq!(scope.pipeline_grants, Some(vec![])); + assert!(scope.denies_everything()); + } + + #[test] + fn test_unresolved_section_is_forbidden_only_when_module_denied() { + let mut scope = blank_scope(); + let none = ModuleAccess { + streams: vec![], + alerts: false, + pipelines: false, + }; + apply_module_access(&mut scope, &none); + let streams: Section = unresolved_section(streams_forbidden(&scope)); + assert_eq!(streams.status, SectionStatus::Forbidden); + let allowed: Section = unresolved_section(false); + assert_eq!(allowed.status, SectionStatus::Error); + } + + #[test] + fn test_caller_with_all_list_grants_stays_unrestricted() { + let mut scope = blank_scope(); + apply_module_access(&mut scope, &ModuleAccess::all()); + assert!(scope.stream_grants.is_empty()); + assert!(!scope.alerts_forbidden); + assert!(scope.pipeline_grants.is_none()); + assert!(!scope.denies_everything()); + } + + #[test] + fn test_module_access_only_closes_the_sections_it_denies() { + let mut scope = blank_scope(); + let access = ModuleAccess { + streams: vec![StreamType::Logs], + alerts: true, + pipelines: false, + }; + apply_module_access(&mut scope, &access); + assert!(!streams_forbidden(&scope)); + assert!(!scope.stream_grants.contains_key(&StreamType::Logs)); + assert_eq!( + scope.stream_grants[&StreamType::Metrics], + Vec::::new() + ); + assert!(!scope.alerts_forbidden); + assert_eq!(scope.pipeline_grants, Some(vec![])); + assert!(!scope.denies_everything()); + } + + fn object_only_scope() -> OverviewScope { + let mut scope = blank_scope(); + for stream_type in OVERVIEW_STREAM_TYPES { + scope + .stream_grants + .insert(stream_type, vec![format!("{stream_type}:app")]); + } + scope.pipeline_grants = Some(vec!["pipelines:p".to_string()]); + scope + } + + #[test] + fn test_object_grants_do_not_bypass_a_missing_module_list_permission() { + let mut scope = object_only_scope(); + let none = ModuleAccess { + streams: vec![], + alerts: false, + pipelines: false, + }; + apply_module_access(&mut scope, &none); + assert!(scope.denies_everything()); + } + + #[test] + fn test_object_grants_are_kept_when_module_list_is_allowed() { + let mut scope = object_only_scope(); + apply_module_access(&mut scope, &ModuleAccess::all()); + assert_eq!(scope.stream_grants[&StreamType::Logs], vec!["logs:app"]); + assert_eq!(scope.pipeline_grants, Some(vec!["pipelines:p".to_string()])); + assert!(!scope.alerts_forbidden); + } + + #[test] + fn test_list_on_streams_only_leaves_alerts_and_pipelines_forbidden() { + for mut scope in [blank_scope(), object_only_scope()] { + let access = ModuleAccess { + streams: OVERVIEW_STREAM_TYPES.to_vec(), + alerts: false, + pipelines: false, + }; + apply_module_access(&mut scope, &access); + assert!(!streams_forbidden(&scope)); + assert!(scope.alerts_forbidden); + assert_eq!(scope.pipeline_grants, Some(vec![])); + } + } + + #[test] + fn test_summary_only_caller_is_denied_with_and_without_object_filtering() { + let none = ModuleAccess { + streams: vec![], + alerts: false, + pipelines: false, + }; + for mut scope in [blank_scope(), object_only_scope()] { + apply_module_access(&mut scope, &none); + assert!(streams_forbidden(&scope)); + assert!(scope.alerts_forbidden); + assert!(scope.denies_everything()); + } + } + + #[test] + fn test_internal_and_ai_chat_streams_are_not_counted_as_user_streams() { + let now = now_micros(); + let old = now - 2 * QUIET_AFTER_MICROS; + assert!(is_counted_user_stream(&stream_at( + "app", + StreamType::Logs, + old + ))); + assert!(!is_counted_user_stream(&stream_at( + "_o2_internal", + StreamType::Logs, + old + ))); + let streams = vec![ + stream_at("app", StreamType::Logs, old), + stream_at("_o2_internal", StreamType::Logs, old), + ]; + assert_eq!(quiet_count(&streams, now), 1); + } + + #[test] + fn test_filter_permitted_alerts_by_name_id_folder_id_or_all() { + let alerts = vec![ + named_alert("id1", "by_name", true), + named_alert("id2", "by_folder_id", true), + named_alert("id3", "by_id", true), + named_alert("id4", "hidden", true), + ]; + let permissions = Some(vec![ + "alert:by_name".to_string(), + "alert:default/id2".to_string(), + "alert:id3".to_string(), + ]); + let names = |rows: Vec| -> Vec { + rows.into_iter().map(|r| r.name).collect() + }; + assert_eq!( + names(filter_permitted_alerts(alerts.clone(), permissions, "o")), + vec!["by_name", "by_folder_id", "by_id"] + ); + assert_eq!(filter_permitted_alerts(alerts.clone(), None, "o").len(), 4); + let all = Some(vec!["alert:_all_o".to_string()]); + assert_eq!(filter_permitted_alerts(alerts, all, "o").len(), 4); + } + + #[tokio::test] + async fn test_concurrent_stream_fetch_matches_sequential_per_type_order() { + let org = "overview_fetch_org"; + seed_stream(org, StreamType::Metrics, "m_a", 1.0).await; + seed_stream(org, StreamType::Logs, "l_a", 1.0).await; + seed_stream(org, StreamType::Logs, "l_hidden", 1.0).await; + seed_stream(org, StreamType::Traces, "t_a", 1.0).await; + let scope = OverviewScope { + user_id: "u".to_string(), + stream_grants: HashMap::from([ + (StreamType::Logs, vec!["logs:l_a".to_string()]), + (StreamType::Traces, vec![]), + ]), + ..Default::default() + }; + let mut sequential = Vec::new(); + for stream_type in OVERVIEW_STREAM_TYPES { + let grants = scope.stream_grants.get(&stream_type).cloned(); + sequential.extend( + try_get_streams(org, Some(stream_type), false, grants) + .await + .unwrap(), + ); + } + let concurrent = fetch_streams(org, &scope).await.unwrap(); + let key = |streams: &[Stream]| -> Vec<(String, String)> { + streams + .iter() + .map(|s| (s.stream_type.to_string(), s.name.clone())) + .collect() + }; + assert_eq!(key(&concurrent), key(&sequential)); + assert_eq!( + key(&concurrent), + vec![ + ("logs".to_string(), "l_a".to_string()), + ("metrics".to_string(), "m_a".to_string()) + ] + ); + } + + #[tokio::test] + async fn test_cache_is_keyed_per_caller_and_marks_hits() { + let org = "overview_cache_org"; + let (admin, limited) = ("admin@example.com", "limited@example.com"); + let key = CacheKey { + org_id: org.to_string(), + user_id: admin.to_string(), + top_n: 5, + }; + let seeded = Overview { + streams: Section::ok(build_streams_data( + vec![stream("only_for_admin", StreamType::Logs, 7.0, 1)], + 5, + now_micros(), + )), + alerts: Section::locked("test"), + pipelines: Section::locked("test"), + limits: Section::locked("test"), + generated_at: now_micros(), + cached: false, + }; + cache_put(key, seeded, 30 * 1_000_000); + + let hit = cached_overview(org, admin, 5).unwrap(); + assert!(hit.cached); + assert_eq!(hit.streams.data.unwrap().total_streams, 1); + assert!(cached_overview(org, limited, 5).is_none()); + + let limited_key = CacheKey { + org_id: org.to_string(), + user_id: limited.to_string(), + top_n: 5, + }; + assert!(CACHE.get(&limited_key).is_none()); + } + + fn stream_at(name: &str, stream_type: StreamType, doc_time_max: i64) -> Stream { + let mut s = stream(name, stream_type, 1.0, 1); + s.stats.doc_time_max = doc_time_max; + s + } + + #[test] + fn test_quiet_streams_classified_at_boundaries() { + let now = 100 * DAY_MICROS; + let streams = vec![ + stream_at("m59", StreamType::Logs, now - 59 * MIN), + stream_at("m61", StreamType::Logs, now - 61 * MIN), + stream_at("d6", StreamType::Metrics, now - 6 * DAY_MICROS), + stream_at("d8", StreamType::Traces, now - 8 * DAY_MICROS), + stream_at("unknown", StreamType::Logs, 0), + stream_at("lookup", StreamType::EnrichmentTables, now - 2 * 60 * MIN), + stream_at("index", StreamType::Index, now - 2 * 60 * MIN), + ]; + assert_eq!(quiet_count(&streams, now), 2); + } + + #[test] + fn test_streams_data_serializes_added_fields_when_empty() { + let data = build_streams_data(vec![], 5, now_micros()); + let v = serde_json::to_value(&data).unwrap(); + assert_eq!(v["quiet_count"], 0); + assert_eq!(v["schema_changes_24h"], 0); + assert!(v["compaction_pending"].is_null()); + } + + #[test] + fn test_sum_pending_for_org_adds_stream_types() { + let counts = HashMap::from([ + ( + "o1".to_string(), + HashMap::from([("logs".to_string(), 3), ("metrics".to_string(), 4)]), + ), + ("o2".to_string(), HashMap::from([("logs".to_string(), 9)])), + ]); + let all: HashSet = ["logs", "metrics"].map(String::from).into(); + let logs_only: HashSet = ["logs".to_string()].into(); + assert_eq!(sum_pending_for_org(&counts, "o1", &all), 7); + assert_eq!(sum_pending_for_org(&counts, "o1", &logs_only), 3); + assert_eq!(sum_pending_for_org(&counts, "o1", &HashSet::new()), 0); + assert_eq!(sum_pending_for_org(&counts, "none", &all), 0); + } + + #[tokio::test] + async fn test_count_schema_changes_counts_recent_versions_of_user_streams() { + let org = "overview_schema_changes_org"; + let schema = Schema::new(vec![Field::new("f", DataType::Utf8, true)]); + let now = now_micros(); + { + let mut w = infra::schema::STREAM_SCHEMAS.write().await; + w.insert( + format!("{org}/logs/app"), + vec![ + (now - 2 * DAY_MICROS, schema.clone()), + (now - 1_000_000, schema.clone()), + (now - 2_000_000, schema.clone()), + ], + ); + w.insert( + format!("{org}/logs/fresh"), + vec![(now - 2 * 3_600 * 1_000_000, schema.clone())], + ); + w.insert( + format!("{org}/logs/_o2_internal"), + vec![ + (now - 2 * DAY_MICROS, schema.clone()), + (now - 1_000_000, schema.clone()), + ], + ); + w.insert( + format!("{org}/enrichment_tables/lookup"), + vec![(now - 1_000_000, schema)], + ); + } + let streams = vec![ + stream("app", StreamType::Logs, 1.0, 1), + stream("lookup", StreamType::EnrichmentTables, 1.0, 1), + stream("absent", StreamType::Logs, 1.0, 1), + stream("fresh", StreamType::Logs, 1.0, 1), + stream("_o2_internal", StreamType::Logs, 1.0, 1), + ]; + assert_eq!( + count_schema_changes(org, &streams, now - DAY_MICROS).await, + 2 + ); + } + + #[test] + fn test_ttl_secs_to_micros_clamps_instead_of_wrapping() { + assert_eq!(ttl_secs_to_micros(30), 30_000_000); + assert_eq!(ttl_secs_to_micros(0), 0); + assert_eq!(ttl_secs_to_micros(u64::MAX), i64::MAX); + assert_eq!(ttl_secs_to_micros(i64::MAX as u64 + 1), i64::MAX); + } + + #[test] + fn test_quiet_boundaries_exactly_one_hour_and_seven_days() { + let now = 100 * DAY_MICROS; + let streams = vec![ + stream_at("h1", StreamType::Logs, now - QUIET_AFTER_MICROS), + stream_at("d7", StreamType::Logs, now - QUIET_FORGOTTEN_AFTER_MICROS), + stream_at("h1_more", StreamType::Logs, now - QUIET_AFTER_MICROS - 1), + ]; + // exactly 1h is still active and exactly 7d is already forgotten + assert_eq!(quiet_count(&streams, now), 1); + } + + #[test] + fn test_make_room_evicts_oldest_entries_and_keeps_fresh_ones() { + let cache: DashMap = DashMap::new(); + let key = |n: usize| CacheKey { + org_id: "o".to_string(), + user_id: format!("u{n}"), + top_n: 5, + }; + let entry = |generated_at: i64| Overview { + streams: Section::locked("test"), + alerts: Section::locked("test"), + pipelines: Section::locked("test"), + limits: Section::locked("test"), + generated_at, + cached: false, + }; + let now = 1_000_000_000; + let ttl = 30 * 1_000_000; + for n in 0..4 { + cache.insert(key(n), entry(now - 4 + n as i64)); + } + make_room(&cache, now, ttl, 4); + assert_eq!(cache.len(), 3); + assert!(!cache.contains_key(&key(0))); + assert!((1..4).all(|n| cache.contains_key(&key(n)))); + + cache.insert(key(9), entry(now - 10 * ttl)); + make_room(&cache, now, ttl, 4); + assert!(!cache.contains_key(&key(9))); + assert_eq!(cache.len(), 3); + } + + #[test] + fn test_cache_key_includes_top_n() { + let org = "overview_top_key_org"; + let key = |top_n| CacheKey { + org_id: org.to_string(), + user_id: "u".to_string(), + top_n, + }; + let seeded = Overview { + streams: Section::locked("test"), + alerts: Section::locked("test"), + pipelines: Section::locked("test"), + limits: Section::locked("test"), + generated_at: now_micros(), + cached: false, + }; + cache_put(key(3), seeded, 30 * 1_000_000); + assert!(cached_overview(org, "u", 3).is_some()); + assert!(cached_overview(org, "u", 4).is_none()); + } + + #[test] + fn test_firing_alert_with_unknown_since_sorts_last_and_serializes_zero() { + let ids = [ + "2Z8rVTbWbyUjb3hIibV02e0sqHA", + "2Z8rVTbWbyUjb3hIibV02e0sqHB", + "2Z8rVTbWbyUjb3hIibV02e0sqHC", + ]; + let alerts: Vec = ids + .iter() + .enumerate() + .map(|(i, id)| named_alert(id, &format!("a{i}"), true)) + .collect(); + let mut states: Vec = alerts + .iter() + .map(|a| state(a, Some(RunOutcome::Firing), None)) + .collect(); + states[0].since = None; + states[1].since = Some(500); + states[2].since = Some(900); + let data = build_alerts_data(&alerts, &states, 3); + let names: Vec<&str> = data.firing.iter().map(|f| f.name.as_str()).collect(); + assert_eq!(names, vec!["a1", "a2", "a0"]); + assert_eq!(serde_json::to_value(&data.firing[2]).unwrap()["since"], 0); + let capped = build_alerts_data(&alerts, &states, 2); + let names: Vec<&str> = capped.firing.iter().map(|f| f.name.as_str()).collect(); + assert_eq!(names, vec!["a1", "a2"]); + } + + #[test] + fn test_next_utc_midnight_micros() { + let midnight = 20_000 * DAY_MICROS; + assert_eq!(next_utc_midnight_micros(midnight), midnight + DAY_MICROS); + assert_eq!( + next_utc_midnight_micros(midnight + 5), + midnight + DAY_MICROS + ); + assert_eq!( + next_utc_midnight_micros(midnight + DAY_MICROS - 1), + midnight + DAY_MICROS + ); + } + + #[test] + fn test_firing_alerts_sorted_oldest_first_capped_with_level_counts() { + let ids = [ + "2Z8rVTbWbyUjb3hIibV02e0sqHA", + "2Z8rVTbWbyUjb3hIibV02e0sqHB", + "2Z8rVTbWbyUjb3hIibV02e0sqHC", + "2Z8rVTbWbyUjb3hIibV02e0sqHD", + ]; + let alerts: Vec = ids + .iter() + .enumerate() + .map(|(i, id)| named_alert(id, &format!("a{i}"), true)) + .collect(); + let mut states = vec![ + state( + &alerts[0], + Some(RunOutcome::Firing), + Some(AlertLevel::Critical), + ), + state(&alerts[1], Some(RunOutcome::NotifyFailed), None), + state( + &alerts[2], + Some(RunOutcome::Firing), + Some(AlertLevel::Critical), + ), + state(&alerts[3], Some(RunOutcome::Normal), Some(AlertLevel::Ok)), + ]; + states[0].since = Some(300); + states[1].since = Some(100); + states[2].since = Some(200); + let data = build_alerts_data(&alerts, &states, 2); + let names: Vec<&str> = data.firing.iter().map(|f| f.name.as_str()).collect(); + assert_eq!(names, vec!["a1", "a2"]); + assert_eq!(data.firing[0].level, "unclassified"); + assert_eq!(data.firing[1].level, "critical"); + assert_eq!(data.firing[0].since, 100); + assert_eq!(data.firing_by_level["critical"], 2); + assert_eq!(data.firing_by_level["unclassified"], 1); + assert_eq!(data.firing_by_level["ok"], 0); + assert_eq!(data.by_level["ok"], 1); + let empty = build_alerts_data(&[], &[], 5); + let v = serde_json::to_value(&empty).unwrap(); + assert_eq!(v["firing"], serde_json::json!([])); + assert_eq!(v["firing_by_level"]["warning"], 0); + } +} diff --git a/src/infra/src/schema/mod.rs b/src/infra/src/schema/mod.rs index 81c670b5d9..5f06d7ee87 100644 --- a/src/infra/src/schema/mod.rs +++ b/src/infra/src/schema/mod.rs @@ -241,6 +241,27 @@ pub async fn get_from_db( }) } +/// Sums the versions created at or after `since` across `streams`, excluding each stream's oldest +/// (its creation), under one cache read lock. +pub async fn count_cached_versions_since( + org_id: &str, + streams: &[(StreamType, &str)], + since: i64, +) -> usize { + let r = STREAM_SCHEMAS.read().await; + streams + .iter() + .map(|(stream_type, stream_name)| { + let cache_key = format!("{org_id}/{stream_type}/{stream_name}"); + r.get(&cache_key).map_or(0, |versions| { + let oldest = versions.iter().map(|(t, _)| *t).min(); + let in_window = versions.iter().filter(|(t, _)| *t >= since).count(); + in_window - usize::from(oldest.is_some_and(|t| t >= since)) + }) + }) + .sum() +} + #[tracing::instrument(name = "infra:schema:get_versions", skip_all)] pub async fn get_versions( org_id: &str, @@ -1150,6 +1171,48 @@ fn user_data_org(key: &str) -> Option<&str> { mod tests { use super::*; + #[tokio::test] + async fn test_count_cached_versions_since_skips_creation_and_includes_boundary() { + let schema = Schema::empty(); + let versions = vec![ + (100, schema.clone()), + (200, schema.clone()), + (300, schema.clone()), + ]; + { + let mut w = STREAM_SCHEMAS.write().await; + w.insert("count_versions_org/logs/three".to_string(), versions); + w.insert( + "count_versions_org/logs/one".to_string(), + vec![(300, schema)], + ); + } + let count = |name: &'static str, since: i64| async move { + count_cached_versions_since("count_versions_org", &[(StreamType::Logs, name)], since) + .await + }; + assert_eq!(count("three", 200).await, 2); + assert_eq!(count("three", 201).await, 1); + assert_eq!(count("three", 0).await, 2); + assert_eq!(count("one", 0).await, 0); + assert_eq!(count("one", 300).await, 0); + assert_eq!(count("missing", 0).await, 0); + + let all = [ + (StreamType::Logs, "three"), + (StreamType::Logs, "one"), + (StreamType::Logs, "missing"), + ]; + assert_eq!( + count_cached_versions_since("count_versions_org", &all, 0).await, + 2 + ); + assert_eq!( + count_cached_versions_since("count_versions_org", &[], 0).await, + 0 + ); + } + #[tokio::test] async fn test_stream_settings_mutation_api() { // unique org prefix so the shared global caches can host parallel tests diff --git a/src/infra/src/table/alerts/mod.rs b/src/infra/src/table/alerts/mod.rs index 1a5e2f36a5..d2a410af8f 100644 --- a/src/infra/src/table/alerts/mod.rs +++ b/src/infra/src/table/alerts/mod.rs @@ -30,9 +30,10 @@ use config::meta::{ use hashbrown::HashMap; use itertools::Itertools; use sea_orm::{ - ActiveModelTrait, ColumnTrait, Condition, ConnectionTrait, DeriveIden, EntityTrait, ModelTrait, - Order, PaginatorTrait, QueryFilter, QueryOrder, Set, TransactionTrait, TryIntoModel, - prelude::Expr, sea_query::Func, + ActiveModelTrait, ColumnTrait, Condition, ConnectionTrait, DeriveIden, EntityTrait, + FromQueryResult, JoinType, ModelTrait, Order, PaginatorTrait, QueryFilter, QueryOrder, + QuerySelect, RelationTrait, Set, TransactionTrait, TryIntoModel, prelude::Expr, + sea_query::Func, }; use svix_ksuid::{Ksuid, KsuidLike}; @@ -47,6 +48,17 @@ use crate::{ pub mod intermediate; +/// The alert columns a summary needs, read without parsing the stored alert body. +#[derive(Clone, Debug, FromQueryResult)] +pub struct AlertSummaryRow { + pub id: String, + pub name: String, + /// The folder's public id, not its surrogate key. + pub folder_id: String, + pub stream_name: String, + pub enabled: bool, +} + impl TryFrom for MetaAlert { type Error = errors::Error; @@ -467,6 +479,28 @@ pub async fn list( Ok(alerts) } +/// Lists every alert in the org as [`AlertSummaryRow`], selecting only the columns it needs. +pub async fn list_summary_rows( + conn: &C, + org_id: &str, +) -> Result, errors::Error> { + Ok(alerts::Entity::find() + .select_only() + .columns([ + alerts::Column::Id, + alerts::Column::Name, + alerts::Column::StreamName, + alerts::Column::Enabled, + ]) + .column_as(folders::Column::FolderId, "folder_id") + .join(JoinType::InnerJoin, alerts::Relation::Folders.def()) + .filter(folders::Column::Type.eq::(folder_type_into_i16(FolderType::Alerts))) + .filter(folders::Column::Org.eq(org_id)) + .into_model::() + .all(conn) + .await?) +} + /// The `(long, short)` burn-window pairs of every alert pointing at this SLO. /// /// The reverse lookup D60 put the indexed `slo_id` column there for: it runs @@ -1869,6 +1903,77 @@ pub(super) mod tests { assert_eq!(ids_of(&db, only_enabled).await, vec![on]); } + #[tokio::test] + async fn list_summary_rows_returns_public_folder_id_and_enabled_flag() { + let db = db_with_folder().await; + seed_folder(&db, FOLDER_B_PK, ORG, "ops").await; + let (on, off) = (alert_id(), alert_id()); + insert_alert(&db, AlertSpec::plain(&on, "on")).await; + insert_alert( + &db, + AlertSpec::plain(&off, "off") + .disabled() + .in_folder(ORG, FOLDER_B_PK), + ) + .await; + + let mut rows = super::super::list_summary_rows(&db, ORG).await.unwrap(); + rows.sort_by(|a, b| a.name.cmp(&b.name)); + let got: Vec<(&str, &str, &str, bool)> = rows + .iter() + .map(|r| { + ( + r.id.as_str(), + r.name.as_str(), + r.folder_id.as_str(), + r.enabled, + ) + }) + .collect(); + assert_eq!( + got, + vec![ + (off.as_str(), "alert-off", "ops", false), + (on.as_str(), "alert-on", "default", true), + ] + ); + assert_eq!(rows[0].stream_name, "default"); + } + + #[tokio::test] + async fn list_summary_rows_excludes_other_orgs_and_non_alert_folders() { + let db = db_with_folder().await; + seed_folder(&db, FOLDER_OTHER_ORG_PK, OTHER_ORG, "default").await; + folders::Entity::insert(folders::ActiveModel { + id: Set("folder_pk_00000000000000009".to_string()), + org: Set(ORG.to_string()), + folder_id: Set("dash".to_string()), + name: Set("folder-dash".to_string()), + description: Set(None), + r#type: Set(folder_type_into_i16(FolderType::Dashboards)), + icon: Set(None), + }) + .exec(&db) + .await + .unwrap(); + let (mine, theirs, dash) = (alert_id(), alert_id(), alert_id()); + insert_alert(&db, AlertSpec::plain(&mine, "mine")).await; + insert_alert( + &db, + AlertSpec::plain(&theirs, "theirs").in_folder(OTHER_ORG, FOLDER_OTHER_ORG_PK), + ) + .await; + insert_alert( + &db, + AlertSpec::plain(&dash, "dash").in_folder(ORG, "folder_pk_00000000000000009"), + ) + .await; + + let rows = super::super::list_summary_rows(&db, ORG).await.unwrap(); + let ids: Vec<&str> = rows.iter().map(|r| r.id.as_str()).collect(); + assert_eq!(ids, vec![mine.as_str()]); + } + /// No filter = historical behaviour, including non-SLO alerts. #[tokio::test] async fn list_models_without_slo_filter_returns_everything() { diff --git a/src/stream/src/lib.rs b/src/stream/src/lib.rs index 92a7cd41a9..26d4c8c44f 100644 --- a/src/stream/src/lib.rs +++ b/src/stream/src/lib.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 . -use std::{future::Future, io::Error}; +use std::{collections::HashSet, future::Future, io::Error}; use arrow_schema::{DataType, Field, Schema}; use axum::{http, response::Response as HttpResponse}; @@ -93,33 +93,31 @@ pub async fn get_streams( fetch_schema: bool, permitted_streams: Option>, ) -> Vec { - let indices = db::schema::list(org_id, stream_type, fetch_schema) + try_get_streams(org_id, stream_type, fetch_schema, permitted_streams) .await - .unwrap_or_default(); + .unwrap_or_default() +} - let filtered_indices = if let Some(s_type) = stream_type { - let s_type = match s_type { - StreamType::EnrichmentTables => "enrichment_table", - _ => s_type.as_str(), - }; - match permitted_streams { - Some(permitted_streams) => { - if permitted_streams.contains(&format!("{s_type}:_all_{org_id}")) { - indices - } else { - indices - .into_iter() - .filter(|stream_loc| { - permitted_streams - .contains(&format!("{}:{}", s_type, stream_loc.stream_name)) - }) - .collect::>() - } - } - None => indices, +/// Like [`get_streams`] but reports a failed schema listing instead of returning an empty list. +pub async fn try_get_streams( + org_id: &str, + stream_type: Option, + fetch_schema: bool, + permitted_streams: Option>, +) -> Result, anyhow::Error> { + let indices = db::schema::list(org_id, stream_type, fetch_schema).await?; + + let filtered_indices = match permitted_streams { + Some(permitted) => { + let grants: HashSet<&str> = permitted.iter().map(String::as_str).collect(); + indices + .into_iter() + .filter(|loc| { + is_stream_permitted(&grants, org_id, loc.stream_type, &loc.stream_name) + }) + .collect::>() } - } else { - indices + None => indices, }; let mut indices_res = Vec::with_capacity(filtered_indices.len()); for stream_loc in filtered_indices { @@ -155,7 +153,7 @@ pub async fn get_streams( )); } } - indices_res + Ok(indices_res) } pub fn stream_res( @@ -373,6 +371,24 @@ pub async fn save_stream_settings( Ok(MetaHttpResponse::ok("")) } +/// Checks a stream against the caller's permitted object list, using the stream's own type. +fn is_stream_permitted( + permitted: &HashSet<&str>, + org_id: &str, + stream_type: StreamType, + stream_name: &str, +) -> bool { + let s_type = match stream_type { + StreamType::EnrichmentTables => "enrichment_table", + _ => stream_type.as_str(), + }; + if permitted.contains(format!("{s_type}:_all_{org_id}").as_str()) { + return true; + } + let grant_name = config::utils::str::into_ofga_supported_format(stream_name); + permitted.contains(format!("{s_type}:{grant_name}").as_str()) +} + async fn sync_associated_metadata_stream_retention( org_id: &str, stream_name: &str, @@ -1208,6 +1224,71 @@ mod tests { // unique org so the process-wide schema caches can host these tests in parallel const SETTINGS_TEST_ORG: &str = "settings_cache_test_org"; + #[test] + fn test_is_stream_permitted_uses_each_streams_own_type() { + let permitted = HashSet::from(["logs:app", "metrics:cpu"]); + assert!(is_stream_permitted( + &permitted, + "o", + StreamType::Logs, + "app" + )); + assert!(is_stream_permitted( + &permitted, + "o", + StreamType::Metrics, + "cpu" + )); + assert!(!is_stream_permitted( + &permitted, + "o", + StreamType::Logs, + "cpu" + )); + assert!(!is_stream_permitted( + &permitted, + "o", + StreamType::Traces, + "app" + )); + } + + #[test] + fn test_is_stream_permitted_matches_sanitized_grant_names() { + let permitted = HashSet::from(["logs:app.v2"]); + assert!(is_stream_permitted( + &permitted, + "o", + StreamType::Logs, + "app:v2" + )); + assert!(!is_stream_permitted( + &permitted, + "o", + StreamType::Logs, + "app_v2" + )); + } + + #[test] + fn test_is_stream_permitted_all_grant_and_enrichment_key() { + let all = HashSet::from(["logs:_all_o"]); + assert!(is_stream_permitted(&all, "o", StreamType::Logs, "anything")); + assert!(!is_stream_permitted( + &all, + "o", + StreamType::Metrics, + "anything" + )); + let enrich = HashSet::from(["enrichment_table:geo"]); + assert!(is_stream_permitted( + &enrich, + "o", + StreamType::EnrichmentTables, + "geo" + )); + } + #[test] fn test_associated_metadata_streams() { assert_eq!( diff --git a/tests/ui-testing/playwright-tests/Onboarding/home-first-data.spec.js b/tests/ui-testing/playwright-tests/Onboarding/home-first-data.spec.js index 8ac2805b2a..07cd84594c 100644 --- a/tests/ui-testing/playwright-tests/Onboarding/home-first-data.spec.js +++ b/tests/ui-testing/playwright-tests/Onboarding/home-first-data.spec.js @@ -14,6 +14,33 @@ const OVERVIEW_SECTIONS = [ /\/api\/v2\/[^/]+\/alerts\/incidents/, /\/api\/[^/]+\/traces\/service_graph\/topology\/current/, ]; +const OVERVIEW_ENDPOINT = /\/api\/[^/]+\/overview(\?.*)?$/; +const FIRING_RULE = 'checkout_latency_p95'; + +// one firing rule, pipelines refused and streams failed, so every section state shows on one page +function overviewWithGaps() { + const nowUs = Date.now() * 1000; + return { + streams: { status: 'error', error: 'stubbed failure' }, + alerts: { + status: 'ok', + data: { + total: 4, + enabled: 4, + disabled: 0, + by_state: { firing: 1, normal: 3, notify_failed: 0 }, + by_level: { ok: 0, warning: 0, critical: 1, no_data: 0, unclassified: 0 }, + failed_evaluation: 0, + firing: [{ name: FIRING_RULE, stream: 'checkout_traces', level: 'critical', since: nowUs - 25 * 60 * 1_000_000 }], + firing_by_level: { ok: 0, warning: 0, critical: 1, no_data: 0, unclassified: 0 }, + }, + }, + pipelines: { status: 'forbidden' }, + limits: { status: 'locked' }, + generated_at: nowUs, + cached: false, + }; +} // Cloud lets only root create orgs, so the empty screens use the stack's one empty org and never write to it. async function emptyOrgId(page, name) { @@ -81,6 +108,26 @@ async function failSections(page, status, which = OVERVIEW_SECTIONS) { } } +async function answerSectionsEmpty(page) { + await page.route(OVERVIEW_SECTIONS[0], (route) => route.fulfill({ json: { total: 0, hits: [] } })); + await page.route(OVERVIEW_SECTIONS[1], (route) => route.fulfill({ json: [] })); + await page.route(OVERVIEW_SECTIONS[2], (route) => route.fulfill({ json: { incidents: [], total: 0 } })); + await page.route(OVERVIEW_SECTIONS[3], (route) => route.fulfill({ json: { nodes: [], edges: [] } })); +} + +async function stubOverview(page, reply) { + await page.route(OVERVIEW_ENDPOINT, (route) => route.fulfill(reply)); +} + +// an org whose summary counts no rules takes the not-monitored hero, which would hide the cards under test +async function pinSummaryRules(page) { + await page.route(/\/api\/[^/]+\/summary(\?.*)?$/, async (route) => { + const response = await route.fetch(); + const body = await response.json().catch(() => ({})); + await route.fulfill({ response, json: { ...body, alerts: { ...(body.alerts || {}), num_realtime: 4, num_scheduled: 0 } } }); + }); +} + async function openHome(page, orgId, tab) { await page.evaluate((t) => localStorage.setItem('o2_home_active_tab', t), tab); await page.goto(`${base()}/web/?org_identifier=${orgId}`); @@ -91,7 +138,6 @@ const tid = (page, id) => page.locator(`[data-test="${id}"]`); async function expectNoOverviewToolbar(page) { await expect(tid(page, 'date-time-btn')).toHaveCount(0); - await expect(page.locator('[data-test="home-page"] [data-test="overview-all-clear-empty-state"]')).toHaveCount(0); } const SCREENS = [ @@ -169,7 +215,6 @@ const SCREENS = [ await expect(tid(page, 'first-data-panel-arrived-summary')).toHaveText('1,204 records in stream default'); await expect(tid(page, 'overview-not-monitored-empty-state')).toBeVisible({ timeout: 20000 }); await expect(tid(page, 'overview-not-monitored-alert-card')).toBeVisible(); - await expect(tid(page, 'overview-all-clear-empty-state')).toHaveCount(0); expect(loads).toEqual([]); }, }, @@ -180,24 +225,19 @@ const SCREENS = [ await openHome(page, orgId, 'overview'); await expect(tid(page, 'overview-load-error-empty-state')).toBeVisible({ timeout: 30000 }); await expect(tid(page, 'overview-load-error-retry-card')).toBeVisible(); - await expect(tid(page, 'overview-all-clear-empty-state')).toHaveCount(0); await expect(tid(page, 'first-data-panel')).toHaveCount(0); }, }, { id: 'h-home-overview-summary-error', async init({ page }) { - await page.route(OVERVIEW_SECTIONS[0], (route) => route.fulfill({ json: { total: 0, hits: [] } })); - await page.route(OVERVIEW_SECTIONS[1], (route) => route.fulfill({ json: [] })); - await page.route(OVERVIEW_SECTIONS[2], (route) => route.fulfill({ json: { incidents: [], total: 0 } })); - await page.route(OVERVIEW_SECTIONS[3], (route) => route.fulfill({ json: { nodes: [], edges: [] } })); + await answerSectionsEmpty(page); await page.route(/\/api\/[^/]+\/summary(\?.*)?$/, (route) => route.fulfill({ status: 500, json: { code: 500, message: 'stubbed failure' } })); }, async run({ page, orgId }) { await openHome(page, orgId, 'overview'); await expect(tid(page, 'overview-load-error-empty-state')).toBeVisible({ timeout: 30000 }); await expect(tid(page, 'overview-load-error-retry-card')).toBeVisible(); - await expect(tid(page, 'overview-all-clear-empty-state')).toHaveCount(0); await expect(tid(page, 'overview-not-monitored-empty-state')).toHaveCount(0); }, }, @@ -283,7 +323,6 @@ const SCREENS = [ async run({ page, orgId }) { await openHome(page, orgId, 'overview'); await expect(tid(page, 'overview-load-error-empty-state')).toBeVisible({ timeout: 30000 }); - await expect(tid(page, 'overview-all-clear-empty-state')).toHaveCount(0); await expect(tid(page, 'first-data-panel')).toHaveCount(0); await expect(tid(page, 'first-event-status')).toHaveCount(0); }, @@ -307,7 +346,6 @@ const SCREENS = [ await expect(tid(page, 'date-time-btn')).toBeVisible(); await expect(page.locator('[data-test="home-page"] [data-test="refresh-button"]')).toBeVisible(); await expect(page.locator('[data-test="home-page"]').getByText('just now')).toBeVisible(); - await expect(tid(page, 'overview-all-clear-empty-state')).toHaveCount(0); await expect(tid(page, 'first-data-panel')).toHaveCount(0); await expect(tid(page, 'first-event-status')).toHaveCount(0); // the load's one-shot first-data notice probes are all; no watcher polls behind the flag-off hero @@ -334,7 +372,6 @@ const SCREENS = [ await openHome(page, orgId, 'overview'); await expect(tid(page, 'overview-load-error-empty-state')).toBeVisible({ timeout: 30000 }); await expect(tid(page, 'overview-load-error-retry-card')).toBeVisible(); - await expect(tid(page, 'overview-all-clear-empty-state')).toHaveCount(0); await expect(tid(page, 'home-first-data-hero')).toHaveCount(0); await expect(tid(page, 'first-data-panel')).toHaveCount(0); await expect(page.locator('[data-test="home-page"]').getByText('just now')).toHaveCount(0); @@ -342,7 +379,6 @@ const SCREENS = [ await tid(page, 'overview-load-error-retry-card').click(); await expect.poll(() => streamReads.length, { timeout: 15000 }).toBeGreaterThan(beforeRetry); await expect(tid(page, 'overview-load-error-empty-state')).toBeVisible({ timeout: 30000 }); - await expect(tid(page, 'overview-all-clear-empty-state')).toHaveCount(0); await page.mouse.move(760, 880); }, }, @@ -379,7 +415,69 @@ const SCREENS = [ await expect(tid(page, 'overview-recent-events-section')).toBeVisible(); await expect(tid(page, 'overview-services-section')).toBeVisible(); await expect(tid(page, 'overview-load-error-empty-state')).toHaveCount(0); - await expect(tid(page, 'overview-all-clear-empty-state')).toHaveCount(0); + }, + }, + { + id: 'h-home-overview-section-states', + async init({ page }) { + await answerSectionsEmpty(page); + await pinSummaryRules(page); + await stubOverview(page, { json: overviewWithGaps() }); + }, + async run({ page, orgId }) { + await openHome(page, orgId, 'overview'); + const tileValue = (key) => page.locator(`[data-test="overview-headline-${key}"] [data-test="overview-tile-value"]`); + const tileNote = (key) => page.locator(`[data-test="overview-headline-${key}"] [data-test="overview-tile-note"]`); + + await expect(tid(page, 'overview-status-danger')).toBeVisible({ timeout: 30000 }); + await expect(tid(page, 'overview-status-text')).toHaveText('Needs attention'); + await expect(page.locator('[data-test^="overview-status-item-"]')).toHaveText(['1 alert firing']); + await expect(tid(page, 'overview-status-clear')).toHaveCount(0); + + await expect(tileValue('firing')).toHaveText('1'); + await expect(tileValue('failing')).toHaveText('0'); + await expect(tileValue('delivery')).toHaveText('0'); + await expect(tileValue('pipelines')).toHaveText('—'); + await expect(tileNote('pipelines')).toHaveText('No access'); + await expect(tileValue('quiet')).toHaveText('—'); + await expect(tileNote('quiet')).toHaveText('Could not load'); + + await expect(tid(page, `overview-alerts-firing-${FIRING_RULE}`)).toBeVisible(); + await expect(tid(page, 'overview-alerts-error')).toHaveCount(0); + await expect(tid(page, 'overview-alerts-forbidden')).toHaveCount(0); + + await expect(tid(page, 'overview-data-health-status-pipelines')).toHaveText('No access'); + await expect(tid(page, 'overview-data-health-value-pipelines')).toHaveText('—'); + await expect(tid(page, 'overview-data-health-status-quiet')).toHaveText('Could not load'); + await expect(tid(page, 'overview-data-health-value-quiet')).toHaveText('—'); + await expect(tid(page, 'overview-data-health-retry-btn')).toHaveCount(0); + + await expect(tid(page, 'overview-streams-error')).toBeVisible(); + await expect(tid(page, 'overview-streams-retry-btn')).toBeVisible(); + await expect(tid(page, 'overview-streams-top')).toHaveCount(0); + + await expect(tid(page, 'overview-load-error-empty-state')).toHaveCount(0); + await expect(tid(page, 'overview-no-access-empty-state')).toHaveCount(0); + }, + }, + { + id: 'h-home-overview-forbidden', + async init({ page }) { + await failSections(page, 403); + await stubOverview(page, { status: 403, json: { code: 403, message: 'Unauthorized Access' } }); + }, + async run({ page, orgId }) { + await openHome(page, orgId, 'overview'); + const noAccess = tid(page, 'overview-no-access-empty-state'); + await expect(noAccess).toBeVisible({ timeout: 30000 }); + await expect(noAccess).toContainText("You don't have access"); + await expect(tid(page, 'overview-load-error-empty-state')).toHaveCount(0); + await expect(tid(page, 'overview-load-error-retry-card')).toHaveCount(0); + await expect(page.locator('[data-test$="-retry-btn"]')).toHaveCount(0); + await expect(tid(page, 'overview-headline-strip')).toHaveCount(0); + await expect(page.locator('[data-test="home-page"] [data-test="refresh-button"]')).toHaveCount(0); + await expectNoOverviewToolbar(page); + await expect(page.locator('[data-test^="o-toast-"]')).toHaveCount(0); }, }, ]; diff --git a/web/src/lib/core/Button/OButton.spec.ts b/web/src/lib/core/Button/OButton.spec.ts index 7f6c2df5f1..ea0beec37b 100644 --- a/web/src/lib/core/Button/OButton.spec.ts +++ b/web/src/lib/core/Button/OButton.spec.ts @@ -43,6 +43,58 @@ describe("OButton", () => { expect(classes).toContain("opacity-60"); }); + it("row variant lays out start name and end value on an auto-height compact row", () => { + const classes = mount(OButton, { props: { variant: "row", block: true } }).classes(); + for (const c of ["h-auto!", "justify-between!", "gap-3!", "px-0!", "py-1.5!", "text-start"]) { + expect(classes).toContain(c); + } + expect(classes).toContain("focus-visible:ring-3"); + }); + + it("locked-body variant starts and wraps its message on an outlined body", () => { + const classes = mount(OButton, { props: { variant: "locked-body", block: true } }).classes(); + for (const c of [ + "h-auto!", + "justify-start!", + "whitespace-normal!", + "px-3!", + "py-3!", + "border", + ]) { + expect(classes).toContain(c); + } + }); + + it("row and locked-body dim when unavailable but focusable", () => { + const row = mount(OButton, { + props: { variant: "row", disabled: true, focusableUnavailable: true }, + }).classes(); + expect(row).toContain("text-text-disabled"); + const locked = mount(OButton, { + props: { variant: "locked-body", disabled: true, focusableUnavailable: true }, + }).classes(); + expect(locked).toContain("opacity-50"); + }); + + it.each(["row", "locked-body"] satisfies ButtonVariant[])( + "%s keeps start alignment, focus and aria-disabled when unavailable", + async (variant) => { + const wrapper = mount(OButton, { + props: { variant, disabled: true, focusableUnavailable: true }, + attachTo: document.body, + }); + expect(wrapper.classes()).toContain("text-start"); + expect(wrapper.classes()).toContain("focus-visible:ring-3"); + expect(wrapper.attributes("disabled")).toBeUndefined(); + expect(wrapper.attributes("aria-disabled")).toBe("true"); + (wrapper.element as HTMLElement).focus(); + expect(document.activeElement).toBe(wrapper.element); + await wrapper.trigger("click"); + expect(wrapper.emitted("click")).toBeUndefined(); + wrapper.unmount(); + }, + ); + it('defaults to type="button"', () => { const wrapper = mount(OButton); expect(wrapper.attributes("type")).toBe("button"); @@ -460,6 +512,8 @@ describe("OButton", () => { "tint-purple", "tint-blue", "tint-orange", + "row", + "locked-body", ] satisfies ButtonVariant[])( "preserves native disabled styling for %s when unavailable", (variant) => { @@ -480,7 +534,7 @@ describe("OButton", () => { .filter( (value) => groups.includes(value.split("-")[0]) && - !/^text-(xs|sm|base|lg|compact|inherit)$/.test(value), + !/^text-(xs|sm|base|lg|compact|inherit|start)$/.test(value), ) .sort(), ).toEqual(disabledClasses.sort()); diff --git a/web/src/lib/core/Button/OButton.types.ts b/web/src/lib/core/Button/OButton.types.ts index c642473db7..d1611e64cc 100644 --- a/web/src/lib/core/Button/OButton.types.ts +++ b/web/src/lib/core/Button/OButton.types.ts @@ -66,7 +66,11 @@ export type ButtonVariant = | "tint-blue" | "tint-orange" // Table-cell content that is itself the action: no chrome, full cell width, stacked lines. - | "cell"; + | "cell" + // Borderless list row whose whole content is the action. + | "row" + // Outlined prompt whose message is the action, e.g. a locked section. + | "locked-body"; /** Size controls height, padding, font-size, and border-radius */ export type ButtonSize = diff --git a/web/src/lib/core/Button/OButton.vue b/web/src/lib/core/Button/OButton.vue index 362fa6d17a..8d4bd27093 100644 --- a/web/src/lib/core/Button/OButton.vue +++ b/web/src/lib/core/Button/OButton.vue @@ -297,6 +297,24 @@ const variantClasses: Record, string> = { "focus-visible:ring-3 focus-visible:ring-button-ghost-focus-ring", "disabled:opacity-60", ].join(" "), + // row: list row whose content is itself the action. + row: [ + "bg-transparent text-button-ghost-text border-0 text-start rounded-default", + "h-auto! justify-between! gap-3! px-0! py-1.5!", + "enabled:hover:bg-button-ghost-hover-bg", + "enabled:active:bg-button-ghost-active-bg", + "focus-visible:ring-3 focus-visible:ring-button-ghost-focus-ring", + "disabled:text-text-disabled", + ].join(" "), + // locked-body: outlined prompt whose message is the action. + "locked-body": [ + "bg-transparent text-button-outline-text border border-button-outline-border text-start", + "h-auto! justify-start! px-3! py-3! whitespace-normal!", + "enabled:hover:bg-button-outline-hover-bg enabled:hover:border-button-outline-hover-border", + "enabled:active:bg-button-outline-active-bg", + "focus-visible:ring-3 focus-visible:ring-button-outline-hover-border", + "disabled:opacity-50", + ].join(" "), // pricing-chip: Pill-shaped toggle chip for model pricing quick-setup template selection "pricing-chip": [ "bg-transparent text-inherit border border-border-default", @@ -371,6 +389,8 @@ const disabledVariantClasses: Record, string "tint-blue": "opacity-60", "tint-orange": "opacity-60", cell: "opacity-60", + row: "text-text-disabled", + "locked-body": "opacity-50", }; const unavailableClasses = computed(() => { @@ -381,9 +401,10 @@ const unavailableClasses = computed(() => { .split(" ") .filter( (value) => - !value.startsWith("enabled:") && - !value.startsWith("disabled:") && - !overriddenGroups.includes(value.split("-")[0]), + value === "text-start" || + (!value.startsWith("enabled:") && + !value.startsWith("disabled:") && + !overriddenGroups.includes(value.split("-")[0])), ), ...disabledClasses, "cursor-not-allowed!", diff --git a/web/src/lib/core/EmptyState/OEmptyState.spec.ts b/web/src/lib/core/EmptyState/OEmptyState.spec.ts index eff1cff916..2359493a2c 100644 --- a/web/src/lib/core/EmptyState/OEmptyState.spec.ts +++ b/web/src/lib/core/EmptyState/OEmptyState.spec.ts @@ -221,6 +221,29 @@ describe("OEmptyState", () => { expect(wrapper.find("p").exists()).toBe(true); }); + it("renders a text-only p title for slim, with no icon or illustration", () => { + wrapper = buildWrapper({ size: "slim" }); + + expect(wrapper.find("h2").exists()).toBe(false); + expect(wrapper.find("p").exists()).toBe(true); + expect(wrapper.find('[data-test="o-icon-stub"]').exists()).toBe(false); + expect(wrapper.find('[data-test="illustration-stub"]').exists()).toBe(false); + }); + + it("draws no dot-grid backdrop for slim or inline unless asked for", () => { + expect(buildWrapper({ size: "slim" }).find('[aria-hidden="true"].absolute').exists()).toBe( + false, + ); + expect(buildWrapper({ size: "inline" }).find('[aria-hidden="true"].absolute').exists()).toBe( + false, + ); + expect( + buildWrapper({ size: "slim", backdrop: true }) + .find('[aria-hidden="true"].absolute') + .exists(), + ).toBe(true); + }); + it("does not render the illustration when size is inline", () => { // Arrange & Act — inline never renders the full illustration wrapper = buildWrapper({ size: "inline" }); diff --git a/web/src/lib/core/EmptyState/OEmptyState.vue b/web/src/lib/core/EmptyState/OEmptyState.vue index c56fe677f3..8d7f881432 100644 --- a/web/src/lib/core/EmptyState/OEmptyState.vue +++ b/web/src/lib/core/EmptyState/OEmptyState.vue @@ -21,6 +21,7 @@ along with this program. If not, see . • hero full page / section (dashboards list, alerts page) — big illustration, rich copy, fills its container, dot-grid backdrop • block inside a card or panel — medium illustration, tighter copy + • slim one text line inside a card section — no icon, no illustration • inline inside a table body / dropdown — tiny: just icon + one line Driven by a PRESET (the common path) or by props/slots (the escape hatch): @@ -43,7 +44,7 @@ along with this program. If not, see . ]" data-test="o2-empty-state" > - +