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

refactor: split flight's common module by responsibility

Browse Source 
common.rs held the wire messages, what a node sends and when, and what
a remote scan's caller tracks. Each now has its own module:

- message: CustomMessage and Metrics, the protocol between nodes.
- node_messages: NodeMessages, NodeRole and MetricsInfo, the sending side.
- remote_scan: RemoteScanMetrics, RemoteScanState and NodeProfiles, the
  receiving side.

FlightMessage, the decoder's output, moves into decoder. Code moves
without changes; callers import from the new modules.
This commit is contained in:
haohuaijin
2026-11-04 17:07:39 +08:00
parent 3678d69d3a
commit 33e173e5f5
13 changed files with 485 additions and 427 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/api/grpc/src/handler/grpc/flight/mod.rs
+1 -1
View File
@@ -28,7 +28,7 @@ use datafusion::{
common::{DataFusionError, Result},
physical_plan::execute_stream_partitioned,
};
use flight::common::{MetricsInfo, NodeMessages};
use flight::node_messages::{MetricsInfo, NodeMessages};
use futures::{StreamExt, stream::BoxStream};
use futures_util::pin_mut;
use prost::Message;
src/api/grpc/src/handler/grpc/flight/stream.rs
+2 -2
View File
@@ -19,7 +19,7 @@ use arrow::ipc::writer::IpcWriteOptions;
use arrow_flight::{FlightData, error::FlightError};
use config::metrics;
use datafusion::execution::SendableRecordBatchStream;
use flight::{common::NodeMessages, encoder::FlightDataEncoder};
use flight::{encoder::FlightDataEncoder, node_messages::NodeMessages};
use futures::{Stream, StreamExt};
use search_service::work_group::DeferredLock;
use tokio::{sync::mpsc, task::JoinSet};
@@ -392,7 +392,7 @@ mod tests {
use arrow_schema::{DataType, Field, Schema};
use config::meta::search::ScanStats;
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use flight::common::{CustomMessage, NodeRole};
use flight::{message::CustomMessage, node_messages::NodeRole};
use futures::TryStreamExt;
use super::*;
src/api/grpc/src/handler/grpc/flight/visitor.rs
+1 -1
View File
@@ -19,7 +19,7 @@ use datafusion::{
common::tree_node::{TreeNode, TreeNodeRecursion},
physical_plan::ExecutionPlan,
};
use flight::common::NodeRole;
use flight::node_messages::NodeRole;
use search::datafusion::distributed_plan::remote_scan_exec::RemoteScanExec;
pub use search::datafusion::plan_metrics::get_peak_memory_from_ctx;
src/flight/src/common.rs
-413
View File
@@ -1,413 +0,0 @@
// 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 <http://www.gnu.org/licenses/>.
use std::{
collections::BTreeMap,
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
};
use arrow::array::RecordBatch;
use arrow_schema::SchemaRef;
use config::{
cluster::LOCAL_NODE,
meta::{cluster::NodeInfo, search::ScanStats},
utils::record_batch_ext::RecordBatchExt,
};
use datafusion::{
self,
physical_plan::{
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
pub fetch_time: metrics::Time,
/// Time in nanos to perform decoding
pub decode_time: metrics::Time,
/// output rows
pub output_rows: metrics::Count,
/// number of RecordBatches received from the remote node
pub batches_received: metrics::Count,
/// total in-memory bytes of RecordBatches received; registered as
/// `MetricValue::OutputBytes` so DataFusion formats it with `human_readable_size`
/// (e.g. "1.5 GB") rather than `human_readable_count` ("1.57 B" = billion)
pub bytes_received: metrics::Count,
}
impl RemoteScanMetrics {
pub fn new(input_partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self {
let fetch_time = MetricBuilder::new(metrics).subset_time("fetch_time", input_partition);
let decode_time = MetricBuilder::new(metrics).subset_time("decode_time", input_partition);
let output_rows = MetricBuilder::new(metrics).output_rows(input_partition);
let batches_received =
MetricBuilder::new(metrics).counter("batches_received", input_partition);
// Register as OutputBytes so the display uses human_readable_size (KB/MB/GB)
// instead of human_readable_count (K/M/B where B means billion).
let bytes_received = metrics::Count::new();
MetricBuilder::new(metrics)
.with_partition(input_partition)
.build(MetricValue::OutputBytes(bytes_received.clone()));
Self {
fetch_time,
decode_time,
output_rows,
batches_received,
bytes_received,
}
}
pub fn record_output(&self, batch: &RecordBatch) {
self.output_rows.add(batch.num_rows());
self.batches_received.add(1);
self.bytes_received.add(batch.size());
}
/// Average fetch time per batch in milliseconds, or 0 if no batches received.
pub fn avg_batch_time_ms(&self) -> u64 {
let batches = self.batches_received.value();
if batches == 0 {
return 0;
}
// fetch_time.value() is total nanoseconds (usize)
let total_ns = self.fetch_time.value() as u64;
total_ns / batches as u64 / 1_000_000
}
}
#[derive(Debug)]
pub enum FlightMessage {
Schema(SchemaRef),
RecordBatch(RecordBatch),
CustomMessage(CustomMessage),
}
/// Custom message for ser/deserialize
#[derive(Debug, Serialize, Deserialize, Clone)]
pub enum CustomMessage {
ScanStats(ScanStats),
Metrics(Vec<Metrics>),
PeakMemory(usize),
PartialErr(String),
}
#[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<ProfileNode>,
}
/// 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<AtomicUsize>,
/// The plan to send back as a profile; set only for EXPLAIN ANALYZE.
pub profile: Option<MetricsInfo>,
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<CustomMessage> {
let mut messages = vec![CustomMessage::ScanStats(self.scan_stats)];
if let NodeRole::ClusterLeader(remote) = &self.role {
messages.push(CustomMessage::ScanStats(*remote.scan_stats.lock()));
messages.extend(remote.partial_err());
}
messages
}
/// Sent after the last batch.
pub fn last(&self) -> Vec<CustomMessage> {
let mut messages = vec![CustomMessage::PeakMemory(
self.peak_memory.load(Ordering::Relaxed),
)];
if let NodeRole::ClusterLeader(remote) = &self.role {
messages.push(CustomMessage::PeakMemory(
remote.peak_memory.load(Ordering::Relaxed),
));
messages.extend(remote.partial_err());
}
if let Some(info) = &self.profile {
messages.push(CustomMessage::Metrics(collect_metrics(info)));
}
messages
}
}
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<Mutex<BTreeMap<usize, RemoteNodeProfile>>>;
/// 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<Mutex<ScanStats>>,
pub peak_memory: Arc<AtomicUsize>,
pub partial_err: Arc<Mutex<String>>,
/// Filled only for EXPLAIN ANALYZE.
pub node_profiles: NodeProfiles,
}
impl RemoteScanState {
fn partial_err(&self) -> Option<CustomMessage> {
let err = self.partial_err.lock().clone();
(!err.is_empty()).then_some(CustomMessage::PartialErr(err))
}
}
pub struct MetricsInfo {
pub plan: Arc<dyn ExecutionPlan>,
/// Passed in because it lives in the search crate, which depends on this one.
pub build_profile: fn(&dyn ExecutionPlan) -> ProfileNode,
}
fn collect_metrics(metrics_info: &MetricsInfo) -> Vec<Metrics> {
vec![Metrics {
node_address: LOCAL_NODE.get_grpc_addr().to_string(),
node_name: LOCAL_NODE.get_name().to_string(),
profile: Some((metrics_info.build_profile)(metrics_info.plan.as_ref())),
..Default::default()
}]
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use arrow_schema::{DataType, Field, Schema};
use config::meta::search::ScanStats;
use datafusion::physical_plan::{ExecutionPlan, empty::EmptyExec};
use parking_lot::Mutex;
use super::*;
fn create_test_scan_stats() -> ScanStats {
ScanStats {
files: 10,
records: 1000,
original_size: 2048,
compressed_size: 1024,
querier_files: 5,
querier_memory_cached_files: 2,
querier_disk_cached_files: 3,
idx_scan_size: 512,
idx_took: 100,
file_list_took: 50,
aggs_cache_ratio: 80,
peak_memory_usage: 1024000,
wait_in_queue: 0,
}
}
fn test_profile(plan: &dyn ExecutionPlan) -> ProfileNode {
ProfileNode {
name: plan.name().to_string(),
..Default::default()
}
}
fn create_test_execution_plan() -> Arc<dyn ExecutionPlan> {
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
Arc::new(EmptyExec::new(schema))
}
#[test]
fn test_metrics_default() {
let metrics = Metrics::default();
assert_eq!(metrics.stage, 0);
assert_eq!(metrics.node_address, "");
assert_eq!(metrics.node_name, "");
assert_eq!(metrics.metrics, "");
}
fn describe(messages: Vec<CustomMessage>) -> Vec<String> {
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_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,
};
assert_eq!(describe(messages.first()), ["ScanStats(10)"]);
assert_eq!(describe(messages.last()), ["PeakMemory(5)"]);
}
#[test]
fn test_cluster_leader_also_sends_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(10)", "ScanStats(3)"]
);
*partial_err.lock() = "node down".to_string();
assert_eq!(
describe(messages.last()),
[
"PeakMemory(5)",
"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!(!metrics[0].node_address.is_empty());
assert!(!metrics[0].node_name.is_empty());
}
#[test]
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_metrics_keeps_fields_older_leaders_require() {
let metrics = Metrics {
profile: Some(ProfileNode::default()),
..Default::default()
};
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]
fn test_custom_message_serialization() {
let scan_stats = create_test_scan_stats();
let custom_msg = CustomMessage::ScanStats(scan_stats);
let serialized = serde_json::to_string(&custom_msg).unwrap();
let deserialized: CustomMessage = serde_json::from_str(&serialized).unwrap();
match deserialized {
CustomMessage::ScanStats(stats) => {
assert_eq!(stats.files, 10);
assert_eq!(stats.records, 1000);
}
_ => panic!("Expected ScanStats variant"),
}
}
#[test]
fn test_metrics_serialization() {
let metrics = vec![Metrics {
stage: 1,
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);
let serialized = serde_json::to_string(&custom_msg).unwrap();
let deserialized: CustomMessage = serde_json::from_str(&serialized).unwrap();
match deserialized {
CustomMessage::Metrics(metrics) => {
assert_eq!(metrics.len(), 1);
assert_eq!(metrics[0].stage, 1);
assert_eq!(metrics[0].node_address, "test_node");
assert_eq!(metrics[0].node_name, "Test Node");
assert_eq!(metrics[0].metrics, "test_metrics");
}
_ => panic!("Expected Metrics variant"),
}
}
#[test]
fn test_custom_message_peak_memory_serialization() {
let peak_memory = 1024 * 1024 * 100; // 100 MB
let custom_msg = CustomMessage::PeakMemory(peak_memory);
let serialized = serde_json::to_string(&custom_msg).unwrap();
let deserialized: CustomMessage = serde_json::from_str(&serialized).unwrap();
match deserialized {
CustomMessage::PeakMemory(mem) => {
assert_eq!(mem, peak_memory);
}
_ => panic!("Expected PeakMemory variant"),
}
}
}
src/flight/src/decoder.rs
+14 -3
View File
@@ -15,7 +15,11 @@
use std::{collections::HashMap, fmt::Debug, pin::Pin, sync::Arc, task::Poll};
use arrow::{array::ArrayRef, buffer::Buffer, ipc::MessageHeader};
use arrow::{
array::{ArrayRef, RecordBatch},
buffer::Buffer,
ipc::MessageHeader,
};
use arrow_flight::{
FlightData,
error::{FlightError, Result},
@@ -26,7 +30,14 @@ use datafusion::parquet::data_type::AsBytes;
use futures::{Stream, StreamExt, ready};
use tonic::Streaming;
use crate::common::{CustomMessage, FlightMessage, RemoteScanMetrics};
use crate::{message::CustomMessage, remote_scan::RemoteScanMetrics};
#[derive(Debug)]
pub enum FlightMessage {
Schema(SchemaRef),
RecordBatch(RecordBatch),
CustomMessage(CustomMessage),
}
pub struct FlightDataDecoder {
response: Streaming<FlightData>,
@@ -198,7 +209,7 @@ mod tests {
use flatbuffers::FlatBufferBuilder;
use super::*;
use crate::common::CustomMessage;
use crate::message::CustomMessage;
fn create_test_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
src/flight/src/encoder.rs
+2 -2
View File
@@ -26,7 +26,7 @@ use arrow_flight::{FlightData, error::Result};
use bytes::Bytes;
use flatbuffers::FlatBufferBuilder;
use crate::common::CustomMessage;
use crate::message::CustomMessage;
/// Encodes CustomMessages and RecordBatch chunks into FlightData (coalescing is separate — see
/// [`ChunkCoalescer`]). Owns the zstd compressor, so it's not `Clone`: use one per thread.
@@ -232,7 +232,7 @@ mod tests {
use config::meta::search::ScanStats;
use super::*;
use crate::common::CustomMessage;
use crate::message::CustomMessage;
fn create_test_schema() -> Schema {
Schema::new(vec![
src/flight/src/lib.rs
+3 -1
View File
@@ -13,7 +13,9 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
pub mod common;
pub mod decoder;
pub mod encoder;
pub mod message;
pub mod node_messages;
pub mod profile;
pub mod remote_scan;
src/flight/src/message.rs
+153
View File
@@ -0,0 +1,153 @@
// 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 <http://www.gnu.org/licenses/>.
//! What nodes send each other over Flight besides record batches; changing it changes the protocol.
use config::meta::search::ScanStats;
use serde::{Deserialize, Serialize};
use crate::profile::ProfileNode;
/// Custom message for ser/deserialize
#[derive(Debug, Serialize, Deserialize, Clone)]
pub enum CustomMessage {
ScanStats(ScanStats),
Metrics(Vec<Metrics>),
PeakMemory(usize),
PartialErr(String),
}
#[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<ProfileNode>,
}
#[cfg(test)]
mod tests {
use super::*;
fn create_test_scan_stats() -> ScanStats {
ScanStats {
files: 10,
records: 1000,
original_size: 2048,
compressed_size: 1024,
querier_files: 5,
querier_memory_cached_files: 2,
querier_disk_cached_files: 3,
idx_scan_size: 512,
idx_took: 100,
file_list_took: 50,
aggs_cache_ratio: 80,
peak_memory_usage: 1024000,
wait_in_queue: 0,
}
}
#[test]
fn test_metrics_default() {
let metrics = Metrics::default();
assert_eq!(metrics.stage, 0);
assert_eq!(metrics.node_address, "");
assert_eq!(metrics.node_name, "");
assert_eq!(metrics.metrics, "");
}
#[test]
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_metrics_keeps_fields_older_leaders_require() {
let metrics = Metrics {
profile: Some(ProfileNode::default()),
..Default::default()
};
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]
fn test_custom_message_serialization() {
let scan_stats = create_test_scan_stats();
let custom_msg = CustomMessage::ScanStats(scan_stats);
let serialized = serde_json::to_string(&custom_msg).unwrap();
let deserialized: CustomMessage = serde_json::from_str(&serialized).unwrap();
match deserialized {
CustomMessage::ScanStats(stats) => {
assert_eq!(stats.files, 10);
assert_eq!(stats.records, 1000);
}
_ => panic!("Expected ScanStats variant"),
}
}
#[test]
fn test_metrics_serialization() {
let metrics = vec![Metrics {
stage: 1,
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);
let serialized = serde_json::to_string(&custom_msg).unwrap();
let deserialized: CustomMessage = serde_json::from_str(&serialized).unwrap();
match deserialized {
CustomMessage::Metrics(metrics) => {
assert_eq!(metrics.len(), 1);
assert_eq!(metrics[0].stage, 1);
assert_eq!(metrics[0].node_address, "test_node");
assert_eq!(metrics[0].node_name, "Test Node");
assert_eq!(metrics[0].metrics, "test_metrics");
}
_ => panic!("Expected Metrics variant"),
}
}
#[test]
fn test_custom_message_peak_memory_serialization() {
let peak_memory = 1024 * 1024 * 100; // 100 MB
let custom_msg = CustomMessage::PeakMemory(peak_memory);
let serialized = serde_json::to_string(&custom_msg).unwrap();
let deserialized: CustomMessage = serde_json::from_str(&serialized).unwrap();
match deserialized {
CustomMessage::PeakMemory(mem) => {
assert_eq!(mem, peak_memory);
}
_ => panic!("Expected PeakMemory variant"),
}
}
}
src/flight/src/node_messages.rs
+202
View File
@@ -0,0 +1,202 @@
// 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 <http://www.gnu.org/licenses/>.
//! What a node sends back to its caller besides record batches, and when.
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
};
use config::{
cluster::LOCAL_NODE,
meta::{cluster::NodeInfo, search::ScanStats},
};
use datafusion::physical_plan::ExecutionPlan;
use crate::{
message::{CustomMessage, Metrics},
profile::ProfileNode,
remote_scan::RemoteScanState,
};
/// 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<AtomicUsize>,
/// The plan to send back as a profile; set only for EXPLAIN ANALYZE.
pub profile: Option<MetricsInfo>,
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<CustomMessage> {
let mut messages = vec![CustomMessage::ScanStats(self.scan_stats)];
if let NodeRole::ClusterLeader(remote) = &self.role {
messages.push(CustomMessage::ScanStats(*remote.scan_stats.lock()));
messages.extend(partial_err(remote));
}
messages
}
/// Sent after the last batch.
pub fn last(&self) -> Vec<CustomMessage> {
let mut messages = vec![CustomMessage::PeakMemory(
self.peak_memory.load(Ordering::Relaxed),
)];
if let NodeRole::ClusterLeader(remote) = &self.role {
messages.push(CustomMessage::PeakMemory(
remote.peak_memory.load(Ordering::Relaxed),
));
messages.extend(partial_err(remote));
}
if let Some(info) = &self.profile {
messages.push(CustomMessage::Metrics(collect_metrics(info)));
}
messages
}
}
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),
}
pub struct MetricsInfo {
pub plan: Arc<dyn ExecutionPlan>,
/// Passed in because it lives in the search crate, which depends on this one.
pub build_profile: fn(&dyn ExecutionPlan) -> ProfileNode,
}
fn partial_err(remote: &RemoteScanState) -> Option<CustomMessage> {
let err = remote.partial_err.lock().clone();
(!err.is_empty()).then_some(CustomMessage::PartialErr(err))
}
fn collect_metrics(metrics_info: &MetricsInfo) -> Vec<Metrics> {
vec![Metrics {
node_address: LOCAL_NODE.get_grpc_addr().to_string(),
node_name: LOCAL_NODE.get_name().to_string(),
profile: Some((metrics_info.build_profile)(metrics_info.plan.as_ref())),
..Default::default()
}]
}
#[cfg(test)]
mod tests {
use arrow_schema::{DataType, Field, Schema};
use datafusion::physical_plan::empty::EmptyExec;
use parking_lot::Mutex;
use super::*;
fn create_test_scan_stats() -> ScanStats {
ScanStats {
files: 10,
..Default::default()
}
}
fn test_profile(plan: &dyn ExecutionPlan) -> ProfileNode {
ProfileNode {
name: plan.name().to_string(),
..Default::default()
}
}
fn create_test_execution_plan() -> Arc<dyn ExecutionPlan> {
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
Arc::new(EmptyExec::new(schema))
}
fn describe(messages: Vec<CustomMessage>) -> Vec<String> {
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_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,
};
assert_eq!(describe(messages.first()), ["ScanStats(10)"]);
assert_eq!(describe(messages.last()), ["PeakMemory(5)"]);
}
#[test]
fn test_cluster_leader_also_sends_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(10)", "ScanStats(3)"]
);
*partial_err.lock() = "node down".to_string();
assert_eq!(
describe(messages.last()),
[
"PeakMemory(5)",
"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!(!metrics[0].node_address.is_empty());
assert!(!metrics[0].node_name.is_empty());
}
}
src/flight/src/remote_scan.rs
+101
View File
@@ -0,0 +1,101 @@
// 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 <http://www.gnu.org/licenses/>.
//! What the caller of a remote scan tracks about the nodes it calls.
use std::{
collections::BTreeMap,
sync::{Arc, atomic::AtomicUsize},
};
use arrow::array::RecordBatch;
use config::{meta::search::ScanStats, utils::record_batch_ext::RecordBatchExt};
use datafusion::physical_plan::metrics::{
self, ExecutionPlanMetricsSet, MetricBuilder, MetricValue,
};
use parking_lot::Mutex;
use crate::profile::RemoteNodeProfile;
/// Per-partition profiles of one remote scan, keyed by partition.
pub type NodeProfiles = Arc<Mutex<BTreeMap<usize, RemoteNodeProfile>>>;
#[derive(Debug, Clone)]
pub struct RemoteScanMetrics {
/// Time in nanos to execute child operator and fetch batches
pub fetch_time: metrics::Time,
/// Time in nanos to perform decoding
pub decode_time: metrics::Time,
/// output rows
pub output_rows: metrics::Count,
/// number of RecordBatches received from the remote node
pub batches_received: metrics::Count,
/// total in-memory bytes of RecordBatches received; registered as
/// `MetricValue::OutputBytes` so DataFusion formats it with `human_readable_size`
/// (e.g. "1.5 GB") rather than `human_readable_count` ("1.57 B" = billion)
pub bytes_received: metrics::Count,
}
impl RemoteScanMetrics {
pub fn new(input_partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self {
let fetch_time = MetricBuilder::new(metrics).subset_time("fetch_time", input_partition);
let decode_time = MetricBuilder::new(metrics).subset_time("decode_time", input_partition);
let output_rows = MetricBuilder::new(metrics).output_rows(input_partition);
let batches_received =
MetricBuilder::new(metrics).counter("batches_received", input_partition);
// Register as OutputBytes so the display uses human_readable_size (KB/MB/GB)
// instead of human_readable_count (K/M/B where B means billion).
let bytes_received = metrics::Count::new();
MetricBuilder::new(metrics)
.with_partition(input_partition)
.build(MetricValue::OutputBytes(bytes_received.clone()));
Self {
fetch_time,
decode_time,
output_rows,
batches_received,
bytes_received,
}
}
pub fn record_output(&self, batch: &RecordBatch) {
self.output_rows.add(batch.num_rows());
self.batches_received.add(1);
self.bytes_received.add(batch.size());
}
/// Average fetch time per batch in milliseconds, or 0 if no batches received.
pub fn avg_batch_time_ms(&self) -> u64 {
let batches = self.batches_received.value();
if batches == 0 {
return 0;
}
// fetch_time.value() is total nanoseconds (usize)
let total_ns = self.fetch_time.value() as u64;
total_ns / batches as u64 / 1_000_000
}
}
/// 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<Mutex<ScanStats>>,
pub peak_memory: Arc<AtomicUsize>,
pub partial_err: Arc<Mutex<String>>,
/// Filled only for EXPLAIN ANALYZE.
pub node_profiles: NodeProfiles,
}
src/search/src/datafusion/distributed_plan/decoder_stream.rs
+3 -2
View File
@@ -23,8 +23,9 @@ use config::{
utils::{record_batch_ext::RecordBatchExt, size::bytes_to_human_readable},
};
use flight::{
common::{CustomMessage, FlightMessage, RemoteScanMetrics},
decoder::FlightDataDecoder,
decoder::{FlightDataDecoder, FlightMessage},
message::CustomMessage,
remote_scan::RemoteScanMetrics,
};
use futures::{Stream, StreamExt, ready};
use tonic::Streaming;
src/search/src/datafusion/distributed_plan/profile.rs
+2 -1
View File
@@ -24,8 +24,9 @@ use datafusion::physical_plan::{
metrics::{MetricValue, MetricsSet},
};
use flight::{
common::{Metrics, NodeProfiles, RemoteScanMetrics},
message::Metrics,
profile::{MetricKind, ProfileInputs, ProfileMetric, ProfileNode, RemoteNodeProfile},
remote_scan::{NodeProfiles, RemoteScanMetrics},
};
use super::remote_scan_exec::RemoteScanExec;
src/search/src/datafusion/distributed_plan/remote_scan_exec.rs
+1 -1
View File
@@ -32,7 +32,7 @@ use datafusion::{
},
};
use datafusion_proto::bytes::physical_plan_to_bytes_with_extension_codec;
use flight::common::{RemoteScanMetrics, RemoteScanState};
use flight::remote_scan::{RemoteScanMetrics, RemoteScanState};
use futures::{StreamExt, TryStreamExt};
use futures_util::pin_mut;
use opentelemetry::context::FutureExt as _;