fix: add some debug logs for query slow (#9072)

src/config/src/config.rs

@@ -2765,8 +2765,6 @@ fn check_disk_cache_config(cfg: &mut Config) -> Result<(), anyhow::Error> {
         cfg.common.result_cache_enabled = false;
         cfg.common.metrics_cache_enabled = false;
         cfg.common.feature_query_streaming_aggs = false;
-        cfg.cache_latest_files.enabled = false;
-        cfg.cache_latest_files.delete_merge_files = false;
     }
 
     let disks = sysinfo::disk::get_disk_usage();

src/config/src/utils/record_batch_ext.rs

@@ -25,7 +25,7 @@ use arrow::{
     record_batch::RecordBatch,
 };
 use arrow_schema::{ArrowError, DataType, Schema};
-use datafusion::error::DataFusionError;
+use datafusion::{error::DataFusionError, physical_plan::spill::get_record_batch_memory_size};
 use hashbrown::{HashMap, HashSet};
 
 use super::{
@@ -43,7 +43,8 @@ pub trait RecordBatchExt {
 
 impl RecordBatchExt for RecordBatch {
     fn size(&self) -> usize {
-        self.schema().size() + self.get_array_memory_size() + USIZE_SIZE
+        // per RecordBatch size = schema size + data size + usize size for num_rows
+        self.schema().size() + get_record_batch_memory_size(self) + USIZE_SIZE
     }
 }
 

src/config/src/utils/util.rs

@@ -20,6 +20,10 @@ where
     if v == Default::default() { def } else { v }
 }
 
+pub fn is_power_of_two(n: u64) -> bool {
+    n == 0 || (n & (n - 1)) == 0
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -34,4 +38,17 @@ mod tests {
         assert_eq!(zero_or("", "v"), "v");
         assert_eq!(zero_or("vv", "v"), "vv");
     }
+
+    #[test]
+    fn test_is_power_of_two() {
+        assert!(is_power_of_two(0));
+        assert!(is_power_of_two(1));
+        assert!(is_power_of_two(2));
+        assert!(!is_power_of_two(3));
+        assert!(is_power_of_two(4));
+        assert!(!is_power_of_two(5));
+        assert!(!is_power_of_two(6));
+        assert!(!is_power_of_two(7));
+        assert!(is_power_of_two(8));
+    }
 }

src/handler/grpc/flight/mod.rs

@@ -214,6 +214,7 @@ impl FlightService for FlightServiceImpl {
 
         let mut stream = FlightEncoderStreamBuilder::new(write_options, 33554432)
             .with_trace_id(trace_id.to_string())
+            .with_is_super(is_super_cluster)
             .with_defer(defer)
             .with_start(start)
             .with_custom_message(PreCustomMessage::ScanStats(scan_stats))

src/handler/grpc/flight/stream.rs

@@ -33,6 +33,7 @@ pub struct FlightEncoderStreamBuilder {
     custom_messages: Vec<PreCustomMessage>,
     // query context
     trace_id: String,
+    is_super: bool,
     defer: Option<AsyncDefer>,
     start: std::time::Instant,
 }
@@ -44,6 +45,7 @@ impl FlightEncoderStreamBuilder {
             queue: VecDeque::new(),
             custom_messages: vec![],
             trace_id: String::new(),
+            is_super: false,
             defer: None,
             start: std::time::Instant::now(),
         }
@@ -59,6 +61,11 @@ impl FlightEncoderStreamBuilder {
         self
     }
 
+    pub fn with_is_super(mut self, is_super: bool) -> Self {
+        self.is_super = is_super;
+        self
+    }
+
     pub fn with_defer(mut self, defer: Option<AsyncDefer>) -> Self {
         self.defer = defer;
         self
@@ -83,11 +90,15 @@ impl FlightEncoderStreamBuilder {
             done: false,
             custom_messages: self.custom_messages,
             trace_id: self.trace_id,
+            is_super: self.is_super,
             defer: self.defer,
             start: self.start,
             first_batch: true,
             span,
             child_span,
+            req_id: 0,
+            req_last_time: std::time::Instant::now(),
+            print_key_event: config::get_config().common.print_key_event,
         }
     }
 }
@@ -101,10 +112,14 @@ pub struct FlightEncoderStream {
     custom_messages: Vec<PreCustomMessage>,
     // query context
     trace_id: String,
+    is_super: bool,
     defer: Option<AsyncDefer>,
     start: std::time::Instant,
     span: tracing::Span,
     child_span: tracing::Span,
+    req_id: u64,
+    req_last_time: std::time::Instant,
+    print_key_event: bool,
 }
 
 impl FlightEncoderStream {
@@ -155,10 +170,31 @@ impl Stream for FlightEncoderStream {
     ) -> Poll<Option<Self::Item>> {
         loop {
             if self.done && self.queue.is_empty() {
+                if self.print_key_event {
+                    log::info!(
+                        "[trace_id {}] flight->search: stream Poll::Ready(None) is_super: {}, took: {} ms",
+                        self.trace_id,
+                        self.is_super,
+                        self.req_last_time.elapsed().as_millis(),
+                    );
+                }
                 return Poll::Ready(None);
             }
 
             if let Some(data) = self.queue.pop_front() {
+                self.req_id += 1;
+                let took = self.req_last_time.elapsed().as_millis();
+                self.req_last_time = std::time::Instant::now();
+                if self.print_key_event
+                    && (took > 100 || config::utils::util::is_power_of_two(self.req_id))
+                {
+                    log::info!(
+                        "[trace_id {}] flight->search: stream Poll::Ready(#{}) message, is_super: {}, took: {took} ms",
+                        self.trace_id,
+                        self.req_id,
+                        self.is_super,
+                    );
+                }
                 return Poll::Ready(Some(Ok(data)));
             }
 
@@ -181,8 +217,9 @@ impl Stream for FlightEncoderStream {
                     self.done = true;
                     self.queue.clear();
                     log::error!(
-                        "[trace_id {}] flight->search: stream error: {e:?}, took: {} ms",
+                        "[trace_id {}] flight->search: stream error: {e:?}, is_super: {}, took: {} ms",
                         self.trace_id,
+                        self.is_super,
                         self.start.elapsed().as_millis()
                     );
                     return Poll::Ready(Some(Err(tonic::Status::internal(e.to_string()))));
@@ -211,8 +248,11 @@ impl Stream for FlightEncoderStream {
 impl Drop for FlightEncoderStream {
     fn drop(&mut self) {
         let trace_id = &self.trace_id;
+        let is_super = self.is_super;
         let took = self.start.elapsed().as_millis();
-        log::info!("[trace_id {trace_id}] flight->search: stream end, took: {took} ms",);
+        log::info!(
+            "[trace_id {trace_id}] flight->search: stream end, is_super: {is_super}, took: {took} ms"
+        );
 
         let _child_enter = self.child_span.enter();
         let _enter = self.span.enter();
@@ -232,7 +272,9 @@ impl Drop for FlightEncoderStream {
         if let Some(defer) = self.defer.take() {
             drop(defer);
         } else {
-            log::info!("[trace_id {trace_id}] flight->search: drop FlightEncoderStream",);
+            log::info!(
+                "[trace_id {trace_id}] flight->search: drop FlightEncoderStream, is_super: {is_super}",
+            );
             // clear session data
             clear_session_data(&self.trace_id);
         }

src/service/search/datafusion/distributed_plan/common.rs

@@ -36,6 +36,9 @@ pub struct QueryContext {
     pub cluster_metrics: Arc<Mutex<Vec<Metrics>>>,
     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,
 }
 
 impl QueryContext {
@@ -50,6 +53,9 @@ impl QueryContext {
             trace_id: String::new(),
             start: std::time::Instant::now(),
             num_rows: 0,
+            req_id: 0,
+            req_last_time: std::time::Instant::now(),
+            print_key_event: config::get_config().common.print_key_event,
         }
     }
 

src/service/search/datafusion/distributed_plan/decoder_stream.rs

@@ -18,7 +18,10 @@ use std::{fmt::Debug, pin::Pin, task::Poll};
 use arrow::array::RecordBatch;
 use arrow_flight::{FlightData, error::Result};
 use arrow_schema::SchemaRef;
-use config::{meta::search::ScanStats, utils::size::bytes_to_human_readable};
+use config::{
+    meta::search::ScanStats,
+    utils::{record_batch_ext::RecordBatchExt, size::bytes_to_human_readable},
+};
 use flight::{
     common::{CustomMessage, FlightMessage, RemoteScanMetrics},
     decoder::FlightDataDecoder,
@@ -85,6 +88,16 @@ impl Stream for FlightDecoderStream {
             match res {
                 // Inner exhausted
                 None => {
+                    if self.query_context.print_key_event {
+                        log::info!(
+                            "[trace_id {}] flight->search: stream receive Poll::Ready(None) from node: {}, name: {}, is_super: {}, took: {} ms",
+                            self.query_context.trace_id,
+                            self.query_context.node.get_grpc_addr(),
+                            self.query_context.node.get_name(),
+                            self.query_context.is_super,
+                            self.query_context.req_last_time.elapsed().as_millis()
+                        );
+                    }
                     poll = Poll::Ready(None);
                     break;
                 }
@@ -106,10 +119,43 @@ impl Stream for FlightDecoderStream {
                 Some(Ok(data)) => match data {
                     FlightMessage::RecordBatch(batch) => {
                         self.query_context.num_rows += batch.num_rows();
+                        self.query_context.req_id += 1;
+                        let took = self.query_context.req_last_time.elapsed().as_millis();
+                        self.query_context.req_last_time = std::time::Instant::now();
+                        if self.query_context.print_key_event
+                            && (took > 100
+                                || config::utils::util::is_power_of_two(self.query_context.req_id))
+                        {
+                            let num_bytes = batch.size();
+                            log::info!(
+                                "[trace_id {}] flight->search: stream receive RecordBatch #{} from node: {}, name: {}, is_super: {}, bytes: {num_bytes}, took: {took} ms",
+                                self.query_context.trace_id,
+                                self.query_context.req_id,
+                                self.query_context.node.get_grpc_addr(),
+                                self.query_context.node.get_name(),
+                                self.query_context.is_super,
+                            );
+                        }
                         poll = Poll::Ready(Some(Ok(batch)));
                         break;
                     }
                     FlightMessage::CustomMessage(message) => {
+                        self.query_context.req_id += 1;
+                        let took = self.query_context.req_last_time.elapsed().as_millis();
+                        self.query_context.req_last_time = std::time::Instant::now();
+                        if self.query_context.print_key_event
+                            && (took > 100
+                                || config::utils::util::is_power_of_two(self.query_context.req_id))
+                        {
+                            log::info!(
+                                "[trace_id {}] flight->search: stream receive CustomMessage #{} from node: {}, name: {}, is_super: {}, took: {took} ms",
+                                self.query_context.trace_id,
+                                self.query_context.req_id,
+                                self.query_context.node.get_grpc_addr(),
+                                self.query_context.node.get_name(),
+                                self.query_context.is_super,
+                            );
+                        }
                         self.process_custom_message(message);
                     }
                     FlightMessage::Schema(_) => {

src/service/search/grpc/storage.rs

@@ -784,7 +784,7 @@ pub async fn tantivy_search(
                 .search_role("follower".to_string())
                 .duration(search_start.elapsed().as_millis() as usize)
                 .desc(format!(
-                    "found {} , is_add_filter_back: {}, file_num: {}",
+                    "found {}, is_add_filter_back: {}, file_num: {}",
                     tantivy_result,
                     is_add_filter_back,
                     file_list_map.len(),

src/service/search/mod.rs

@@ -768,7 +768,7 @@ pub async fn search_partition(
             let records = (stats.doc_num as i64 / data_retention) * query_duration;
             let original_size = (stats.storage_size as i64 / data_retention) * query_duration;
             log::info!(
-                "[trace_id {trace_id}] using approximation: stream: {stream_name}, records: {records}, original_size: {original_size} , data_retention in seconds: {data_retention}",
+                "[trace_id {trace_id}] using approximation: stream: {stream_name}, records: {records}, original_size: {original_size}, data_retention in seconds: {data_retention}",
             );
             files.push(infra::file_list::FileId {
                 id: Utc::now().timestamp_micros(),