diff --git a/Cargo.lock b/Cargo.lock index b632005409..53aacb5ad5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2398,6 +2398,7 @@ dependencies = [ "humantime", "meta-client", "serde", + "serde_json", "snafu 0.8.6", "tokio", "tonic 0.14.2", diff --git a/src/catalog/src/process_manager.rs b/src/catalog/src/process_manager.rs index 8730821c3e..dfe3f7f283 100644 --- a/src/catalog/src/process_manager.rs +++ b/src/catalog/src/process_manager.rs @@ -312,6 +312,28 @@ pub struct SlowQueryTimer { sample_ratio: f64, record_type: SlowQueriesRecordType, recorder: EventRecorderRef, + record_state: Arc>, +} + +#[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>, +} + +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::() <= 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::() <= 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>, + } + + impl EventRecorder for RecordingEventRecorder { + fn record(&self, event: Box) { + let event = event + .as_any() + .downcast_ref::() + .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() { diff --git a/src/common/frontend/Cargo.toml b/src/common/frontend/Cargo.toml index c775944286..32a5146b56 100644 --- a/src/common/frontend/Cargo.toml +++ b/src/common/frontend/Cargo.toml @@ -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 diff --git a/src/common/frontend/src/slow_query_event.rs b/src/common/frontend/src/slow_query_event.rs index d7ec630c5d..920f4c3363 100644 --- a/src/common/frontend/src/slow_query_event.rs +++ b/src/common/frontend/src/slow_query_event.rs @@ -45,6 +45,7 @@ pub struct SlowQueryEvent { pub promql_step: Option, pub promql_start: Option, pub promql_end: Option, + 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 { + Ok(self.payload.clone()) + } + fn extra_schema(&self) -> Vec { 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); + } } diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 07fce9feb6..5be2d617ab 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -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 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>, + timeout_recorder: Option, ) -> Result { 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 { +fn record_explain_analyze_timeout( + recorder: Option<&SlowQueryRecorder>, + plan: Option<&Arc>, +) { + 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, +) -> Result { 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 { 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, ) -> Result { 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>, + } + + impl EventRecorder for RecordingSlowQueryEventRecorder { + fn record(&self, event: Box) { + let event = event + .as_any() + .downcast_ref::() + .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 = 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>>,