feat: record metrics for timed out explain analyze (#8668)

* feat: record metrics for timed out explain analyze Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix: update event recorder test implementations Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * fix: keep timeout metrics payload shape stable Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> --------- Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
2026-08-23 14:38:18 +00:00 · 2026-07-29 14:47:40 +08:00
parent deb688f572
commit 1e9d70fa5c
5 changed files with 315 additions and 22 deletions
@@ -2398,6 +2398,7 @@ dependencies = [
 "humantime",
 "meta-client",
 "serde",
+ "serde_json",
 "snafu 0.8.6",
 "tokio",
 "tonic 0.14.2",
@@ -312,6 +312,28 @@ pub struct SlowQueryTimer {
    sample_ratio: f64,
    record_type: SlowQueriesRecordType,
    recorder: EventRecorderRef,
+    record_state: Arc<RwLock<SlowQueryRecordState>>,
+}
+
+#[derive(Default)]
+struct SlowQueryRecordState {
+    force_record: bool,
+    payload: serde_json::Value,
+}
+
+/// A handle to attach diagnostics to a slow query and force it to be recorded.
+#[derive(Clone)]
+pub struct SlowQueryRecorder {
+    record_state: Arc<RwLock<SlowQueryRecordState>>,
+}
+
+impl SlowQueryRecorder {
+    /// Attaches the payload and marks the query to bypass the threshold and sampling checks.
+    pub fn force_record_with_payload(&self, payload: serde_json::Value) {
+        let mut state = self.record_state.write().unwrap();
+        state.payload = payload;
+        state.force_record = true;
+    }
 }

 impl SlowQueryTimer {
@@ -331,12 +353,20 @@ impl SlowQueryTimer {
            sample_ratio,
            record_type,
            recorder,
+            record_state: Arc::default(),
+        }
+    }
+
+    /// Returns a handle that can attach diagnostics and force this query to be recorded.
+    pub fn recorder(&self) -> SlowQueryRecorder {
+        SlowQueryRecorder {
+            record_state: self.record_state.clone(),
        }
    }
 }

 impl SlowQueryTimer {
-    fn send_slow_query_event(&self, elapsed: Duration) {
+    fn send_slow_query_event(&self, elapsed: Duration, payload: serde_json::Value) {
        let mut slow_query_event = SlowQueryEvent {
            cost: elapsed.as_millis() as u64,
            threshold: self.threshold.as_millis() as u64,
@@ -349,6 +379,7 @@ impl SlowQueryTimer {
            promql_step: None,
            promql_start: None,
            promql_end: None,
+            payload,
        };

        match &self.stmt {
@@ -398,6 +429,7 @@ impl SlowQueryTimer {
                    promql_step = slow_query_event.promql_step,
                    promql_start = slow_query_event.promql_start,
                    promql_end = slow_query_event.promql_end,
+                    payload = slow_query_event.payload.to_string(),
                );
            }
        }
@@ -406,23 +438,98 @@ impl SlowQueryTimer {

 impl Drop for SlowQueryTimer {
    fn drop(&mut self) {
-        // Calculate the elapsed duration since the timer is created.
        let elapsed = self.start.elapsed();
-        if elapsed > self.threshold {
-            // Only capture a portion of slow queries based on sample_ratio.
-            // Generate a random number in [0, 1) and compare it with sample_ratio.
-            if self.sample_ratio >= 1.0 || random::<f64>() <= self.sample_ratio {
-                self.send_slow_query_event(elapsed);
-            }
+        let (force_record, payload) = {
+            let state = self.record_state.read().unwrap();
+            (state.force_record, state.payload.clone())
+        };
+
+        if force_record
+            || (elapsed > self.threshold
+                && (self.sample_ratio >= 1.0 || random::<f64>() <= self.sample_ratio))
+        {
+            self.send_slow_query_event(elapsed, payload);
        }
    }
 }

 #[cfg(test)]
 mod tests {
-    use std::sync::Arc;
+    use std::sync::{Arc, Mutex};
+    use std::time::Duration;

-    use crate::process_manager::ProcessManager;
+    use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef};
+    use common_frontend::slow_query_event::SlowQueryEvent;
+    use common_telemetry::logging::SlowQueriesRecordType;
+    use serde_json::Value;
+
+    use crate::process_manager::{ProcessManager, QueryStatement, SlowQueryTimer};
+
+    #[derive(Debug, Default)]
+    struct RecordingEventRecorder {
+        events: Mutex<Vec<(String, Value)>>,
+    }
+
+    impl EventRecorder for RecordingEventRecorder {
+        fn record(&self, event: Box<dyn Event>) {
+            let event = event
+                .as_any()
+                .downcast_ref::<SlowQueryEvent>()
+                .expect("expected a slow query event");
+            self.events
+                .lock()
+                .unwrap()
+                .push((event.query.clone(), event.payload.clone()));
+        }
+
+        fn event_type_filter(&self) -> EventTypeFilterRef {
+            Arc::new(EventTypeFilter::All)
+        }
+
+        fn close(&self) {}
+    }
+
+    #[test]
+    fn test_forced_slow_query_bypasses_threshold_and_sampling() {
+        let event_recorder = Arc::new(RecordingEventRecorder::default());
+        let timer = SlowQueryTimer::new(
+            QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
+            "public".to_string(),
+            Duration::from_secs(3600),
+            0.0,
+            SlowQueriesRecordType::SystemTable,
+            event_recorder.clone(),
+        );
+        let payload = serde_json::json!({
+            "timed_out": true,
+            "metrics": [{"stage": 0}],
+        });
+        timer.recorder().force_record_with_payload(payload.clone());
+
+        drop(timer);
+
+        let events = event_recorder.events.lock().unwrap();
+        assert_eq!(events.len(), 1);
+        assert_eq!(events[0].0, "EXPLAIN ANALYZE VERBOSE SELECT 1");
+        assert_eq!(events[0].1, payload);
+    }
+
+    #[test]
+    fn test_unforced_fast_query_is_not_recorded() {
+        let event_recorder = Arc::new(RecordingEventRecorder::default());
+        let timer = SlowQueryTimer::new(
+            QueryStatement::Plan("SELECT 1".to_string()),
+            "public".to_string(),
+            Duration::from_secs(3600),
+            0.0,
+            SlowQueriesRecordType::SystemTable,
+            event_recorder.clone(),
+        );
+
+        drop(timer);
+
+        assert!(event_recorder.events.lock().unwrap().is_empty());
+    }

    #[tokio::test]
    async fn test_register_query() {
@@ -17,6 +17,7 @@ greptime-proto.workspace = true
 humantime.workspace = true
 meta-client.workspace = true
 serde.workspace = true
+serde_json.workspace = true
 snafu.workspace = true
 tonic.workspace = true

@@ -45,6 +45,7 @@ pub struct SlowQueryEvent {
    pub promql_step: Option<u64>,
    pub promql_start: Option<i64>,
    pub promql_end: Option<i64>,
+    pub payload: serde_json::Value,
 }

 impl Event for SlowQueryEvent {
@@ -56,6 +57,10 @@ impl Event for SlowQueryEvent {
        SLOW_QUERY_EVENT_TYPE
    }

+    fn json_payload(&self) -> Result<serde_json::Value> {
+        Ok(self.payload.clone())
+    }
+
    fn extra_schema(&self) -> Vec<ColumnSchema> {
        vec![
            ColumnSchema {
@@ -155,6 +160,7 @@ mod tests {
            promql_step: None,
            promql_start: None,
            promql_end: None,
+            payload: serde_json::Value::Null,
        };

        let schema = event.extra_schema();
@@ -177,6 +183,7 @@ mod tests {
            ]
        );
        assert_eq!(schema[8].semantic_type, SemanticType::Field as i32);
+        assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null);

        let rows = event.extra_rows().unwrap();
        assert_eq!(rows.len(), 1);
@@ -185,4 +192,26 @@ mod tests {
            Some(ValueData::StringValue("public".to_string()))
        );
    }
+
+    #[test]
+    fn slow_query_event_includes_timeout_payload() {
+        let payload = serde_json::json!({
+            "timed_out": true,
+            "metrics": [{"stage": 0}],
+        });
+        let event = SlowQueryEvent {
+            cost: 100,
+            threshold: 10,
+            query: "EXPLAIN ANALYZE VERBOSE SELECT 1".to_string(),
+            schema_name: "public".to_string(),
+            is_promql: false,
+            promql_range: None,
+            promql_step: None,
+            promql_start: None,
+            promql_end: None,
+            payload: payload.clone(),
+        };
+
+        assert_eq!(event.json_payload().unwrap(), payload);
+    }
 }
@@ -40,7 +40,7 @@ use auth::{
 };
 use catalog::CatalogManagerRef;
 use catalog::process_manager::{
-    ProcessManagerRef, QueryStatement as CatalogQueryStatement, SlowQueryTimer,
+    ProcessManagerRef, QueryStatement as CatalogQueryStatement, SlowQueryRecorder, SlowQueryTimer,
 };
 use client::OutputData;
 use common_base::Plugins;
@@ -59,6 +59,7 @@ use common_recordbatch::error::StreamTimeoutSnafu;
 use common_telemetry::logging::SlowQueryOptions;
 use common_telemetry::{debug, error, tracing};
 use dashmap::DashMap;
+use datafusion::physical_plan::ExecutionPlan;
 use datafusion_expr::LogicalPlan;
 use futures::{Stream, StreamExt, future};
 use lazy_static::lazy_static;
@@ -214,6 +215,10 @@ fn parse_stmt(sql: &str, dialect: &(dyn Dialect + Send + Sync)) -> Result<Vec<St
    ParserContext::create_with_dialect(sql, dialect, ParseOptions::default()).context(ParseSqlSnafu)
 }

+fn is_explain_analyze_verbose(stmt: &Statement) -> bool {
+    matches!(stmt, Statement::Explain(explain) if explain.analyze && explain.verbose)
+}
+
 fn validate_analyze_stream_statement(stmt: &mut Statement) -> Result<()> {
    let Statement::Explain(explain) = stmt else {
        return InvalidSqlSnafu {
@@ -272,6 +277,9 @@ impl Instance {
            let catalog_name = query_ctx.current_catalog().to_string();
            let schema_name = query_ctx.current_schema();
            let slow_query_timer = self.statement_slow_query_timer(&stmt, schema_name.clone());
+            let timeout_recorder = is_explain_analyze_verbose(&stmt)
+                .then(|| slow_query_timer.as_ref().map(SlowQueryTimer::recorder))
+                .flatten();

            let ticket = self.process_manager.register_query(
                catalog_name,
@@ -282,7 +290,12 @@ impl Instance {
                slow_query_timer,
            );

-            let query_fut = self.exec_statement_with_timeout(stmt, query_ctx, query_interceptor);
+            let query_fut = self.exec_statement_with_timeout(
+                stmt,
+                query_ctx,
+                query_interceptor,
+                timeout_recorder,
+            );

            CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
                .await
@@ -299,7 +312,7 @@ impl Instance {
                    Output { data, meta }
                })
        } else {
-            self.exec_statement_with_timeout(stmt, query_ctx, query_interceptor)
+            self.exec_statement_with_timeout(stmt, query_ctx, query_interceptor, None)
                .await
        }
    }
@@ -309,6 +322,7 @@ impl Instance {
        stmt: Statement,
        query_ctx: QueryContextRef,
        query_interceptor: Option<&SqlQueryInterceptorRef<Error>>,
+        timeout_recorder: Option<SlowQueryRecorder>,
    ) -> Result<Output> {
        let timeout = derive_timeout(&stmt, &query_ctx);
        match timeout {
@@ -323,7 +337,7 @@ impl Instance {
                let output = map_query_output(output)?;
                // compute remaining timeout
                let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default();
-                attach_timeout(output, remaining_timeout)
+                attach_timeout(output, remaining_timeout, timeout_recorder)
            }
            None => self
                .exec_statement(stmt, query_ctx, query_interceptor)
@@ -560,18 +574,44 @@ fn derive_timeout_for_plan(plan: &LogicalPlan, query_ctx: &QueryContextRef) -> O
    }
 }

-fn attach_timeout(output: Output, mut timeout: Duration) -> Result<Output> {
+fn record_explain_analyze_timeout(
+    recorder: Option<&SlowQueryRecorder>,
+    plan: Option<&Arc<dyn ExecutionPlan>>,
+) {
+    let Some(recorder) = recorder else {
+        return;
+    };
+    let metrics = plan
+        .and_then(|plan| query::analyze_plan_metrics_to_json_value(plan, true).ok())
+        .unwrap_or_else(|| serde_json::json!([]));
+    recorder.force_record_with_payload(serde_json::json!({
+        "timed_out": true,
+        "metrics": metrics,
+    }));
+}
+
+fn attach_timeout(
+    output: Output,
+    mut timeout: Duration,
+    timeout_recorder: Option<SlowQueryRecorder>,
+) -> Result<Output> {
    if timeout.is_zero() {
        return StatementTimeoutSnafu.fail();
    }

+    let plan = timeout_recorder
+        .as_ref()
+        .and_then(|_| output.meta.plan.clone());
    let output = match output.data {
        OutputData::AffectedRows(_) | OutputData::RecordBatches(_) => output,
        OutputData::Stream(mut stream) => {
            let schema = stream.schema();
            let s = Box::pin(stream! {
                let mut start = tokio::time::Instant::now();
-                while let Some(item) = tokio::time::timeout(timeout, stream.next()).await.map_err(|_| StreamTimeoutSnafu.build())? {
+                while let Some(item) = tokio::time::timeout(timeout, stream.next()).await.map_err(|_| {
+                    record_explain_analyze_timeout(timeout_recorder.as_ref(), plan.as_ref());
+                    StreamTimeoutSnafu.build()
+                })? {
                    yield item;

                    let now = tokio::time::Instant::now();
@@ -579,6 +619,7 @@ fn attach_timeout(output: Output, mut timeout: Duration) -> Result<Output> {
                    start = now;
                    // tokio::time::timeout may not return an error immediately when timeout is 0.
                    if timeout.is_zero() {
+                        record_explain_analyze_timeout(timeout_recorder.as_ref(), plan.as_ref());
                        StreamTimeoutSnafu.fail()?;
                    }
                }
@@ -677,7 +718,7 @@ impl Instance {
            slow_query_timer,
        );
        let query_fut =
-            self.exec_statement_with_timeout(stmt, query_ctx.clone(), query_interceptor);
+            self.exec_statement_with_timeout(stmt, query_ctx.clone(), query_interceptor, None);
        let output = CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
            .await
            .map_err(|_| error::CancelledSnafu.build())??;
@@ -760,6 +801,7 @@ impl Instance {
        &self,
        plan: LogicalPlan,
        query_ctx: QueryContextRef,
+        timeout_recorder: Option<SlowQueryRecorder>,
    ) -> Result<Output> {
        let timeout = derive_timeout_for_plan(&plan, &query_ctx);
        match timeout {
@@ -770,7 +812,7 @@ impl Instance {
                    .map_err(|_| StatementTimeoutSnafu.build())??;
                let output = map_query_output(output)?;
                let remaining_timeout = timeout.checked_sub(start.elapsed()).unwrap_or_default();
-                attach_timeout(output, remaining_timeout)
+                attach_timeout(output, remaining_timeout, timeout_recorder)
            }
            None => self
                .exec_plan(plan, query_ctx)
@@ -816,6 +858,11 @@ impl Instance {
                None
            };

+            let timeout_recorder = stmt
+                .as_ref()
+                .is_some_and(is_explain_analyze_verbose)
+                .then(|| slow_query_timer.as_ref().map(SlowQueryTimer::recorder))
+                .flatten();
            let ticket = self.process_manager.register_query(
                catalog_name,
                vec![schema_name],
@@ -825,7 +872,7 @@ impl Instance {
                slow_query_timer,
            );

-            let query_fut = self.exec_plan_with_timeout(plan, query_ctx.clone());
+            let query_fut = self.exec_plan_with_timeout(plan, query_ctx.clone(), timeout_recorder);

            CancellableFuture::new(query_fut, ticket.cancellation_handle.clone())
                .await
@@ -842,7 +889,8 @@ impl Instance {
                    Output { data, meta }
                })
        } else {
-            self.exec_plan_with_timeout(plan, query_ctx.clone()).await
+            self.exec_plan_with_timeout(plan, query_ctx.clone(), None)
+                .await
        };

        result.and_then(|output| query_interceptor.post_execute(output, query_ctx))
@@ -1607,10 +1655,12 @@ mod tests {
    use api::prom_store::remote::{LabelMatcher, Query as RemoteQuery, ReadRequest};
    use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse};
    use auth::{PermissionResp, UserInfoRef};
-    use catalog::process_manager::ProcessManager;
+    use catalog::process_manager::{ProcessManager, QueryStatement, SlowQueryTimer};
    use common_base::Plugins;
    use common_error::ext::{BoxedError, PlainError};
    use common_error::status_code::StatusCode;
+    use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef};
+    use common_frontend::slow_query_event::SlowQueryEvent;
    use common_meta::cache::LayeredCacheRegistryBuilder;
    use common_meta::kv_backend::memory::MemoryKvBackend;
    use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor};
@@ -1618,11 +1668,13 @@ mod tests {
    use common_meta::rpc::procedure::{
        MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse,
    };
-    use common_query::Output;
+    use common_query::{Output, OutputMeta};
    use common_recordbatch::{
        OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream,
    };
+    use common_telemetry::logging::SlowQueriesRecordType;
    use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+    use datafusion::physical_plan::empty::EmptyExec;
    use datafusion_expr::dml::InsertOp;
    use datafusion_expr::{LogicalPlanBuilder, LogicalTableSource};
    use datatypes::prelude::ConcreteDataType;
@@ -1654,6 +1706,27 @@ mod tests {
        parse_stmt(sql, &GreptimeDbDialect {}).unwrap()
    }

+    #[derive(Debug, Default)]
+    struct RecordingSlowQueryEventRecorder {
+        payloads: std::sync::Mutex<Vec<serde_json::Value>>,
+    }
+
+    impl EventRecorder for RecordingSlowQueryEventRecorder {
+        fn record(&self, event: Box<dyn Event>) {
+            let event = event
+                .as_any()
+                .downcast_ref::<SlowQueryEvent>()
+                .expect("expected a slow query event");
+            self.payloads.lock().unwrap().push(event.payload.clone());
+        }
+
+        fn event_type_filter(&self) -> EventTypeFilterRef {
+            Arc::new(EventTypeFilter::All)
+        }
+
+        fn close(&self) {}
+    }
+
    #[test]
    fn test_validate_analyze_stream_statement_strictness() {
        for sql in [
@@ -1688,6 +1761,21 @@ mod tests {
            parse_test_sql("explain analyze verbose select 1; select 2").len(),
            2
        );
+
+        assert!(is_explain_analyze_verbose(
+            &parse_test_sql("explain analyze verbose select 1")[0]
+        ));
+        for sql in [
+            "select 1",
+            "explain select 1",
+            "explain analyze select 1",
+            "explain verbose select 1",
+        ] {
+            assert!(
+                !is_explain_analyze_verbose(&parse_test_sql(sql)[0]),
+                "{sql}"
+            );
+        }
    }

    #[derive(Debug, Snafu)]
@@ -1955,6 +2043,73 @@ mod tests {

    impl Unpin for PendingRecordBatchStream {}

+    #[test]
+    fn test_record_explain_analyze_timeout_uses_empty_metrics_without_plan() {
+        let event_recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
+        let timer = SlowQueryTimer::new(
+            QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
+            "public".to_string(),
+            Duration::from_secs(3600),
+            0.0,
+            SlowQueriesRecordType::SystemTable,
+            event_recorder.clone(),
+        );
+        let timeout_recorder = timer.recorder();
+
+        record_explain_analyze_timeout(Some(&timeout_recorder), None);
+        drop(timer);
+
+        let payloads = event_recorder.payloads.lock().unwrap();
+        assert_eq!(payloads.len(), 1);
+        assert_eq!(payloads[0]["timed_out"], true);
+        assert_eq!(payloads[0]["metrics"], serde_json::json!([]));
+    }
+
+    #[tokio::test]
+    async fn test_attach_timeout_records_explain_analyze_metrics() {
+        let event_recorder = Arc::new(RecordingSlowQueryEventRecorder::default());
+        let timer = SlowQueryTimer::new(
+            QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
+            "public".to_string(),
+            Duration::from_secs(3600),
+            0.0,
+            SlowQueriesRecordType::SystemTable,
+            event_recorder.clone(),
+        );
+        let timeout_recorder = timer.recorder();
+        let plan: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
+        let (finish_tx, finish_rx) = oneshot::channel();
+        let stream = PendingRecordBatchStream {
+            schema: Arc::new(GtSchema::new(vec![])),
+            polled_tx: None,
+            _finish_tx: finish_tx,
+            finish_rx: Box::pin(finish_rx),
+        };
+        let output = Output::new(
+            OutputData::Stream(Box::pin(stream)),
+            OutputMeta::new_with_plan(plan),
+        );
+        let output =
+            attach_timeout(output, Duration::from_millis(10), Some(timeout_recorder)).unwrap();
+        let OutputData::Stream(mut stream) = output.data else {
+            unreachable!();
+        };
+
+        let err = stream.next().await.unwrap().unwrap_err();
+        assert_eq!(err.to_string(), "Stream timeout");
+        drop(stream);
+        drop(timer);
+
+        let payloads = event_recorder.payloads.lock().unwrap();
+        assert_eq!(payloads.len(), 1);
+        assert_eq!(payloads[0]["timed_out"], true);
+        assert!(
+            payloads[0]["metrics"]
+                .as_array()
+                .is_some_and(|metrics| !metrics.is_empty())
+        );
+    }
+
    struct PendingDataSource {
        schema: GtSchemaRef,
        polled_tx: std::sync::Mutex<Option<oneshot::Sender<()>>>,