diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 367127df16..55dc6a670a 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -1805,7 +1805,8 @@ impl RegionServerInner { request: QueryRequest, query_ctx: QueryContextRef, ) -> Result { - let explain_verbose = query_ctx.explain_verbose(); + let live_analyze_metrics = + query_ctx.explain_verbose() && query_ctx.live_analyze_metrics_enabled(); let inner = self.clone(); let mut stream = common_runtime::spawn_query(async move { inner.handle_read_inner(request, query_ctx).await @@ -1822,7 +1823,7 @@ impl RegionServerInner { let producer_metrics = metrics.clone(); let producer_handle = common_runtime::spawn_query(async move { - if explain_verbose { + if live_analyze_metrics { loop { match time::timeout(FLIGHT_METRICS_HEARTBEAT_INTERVAL, stream.next()).await { Ok(Some(batch)) => { diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 562c33757b..07dfe3de7d 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -358,6 +358,7 @@ impl MergeScanExec { let current_channel = self.query_ctx.channel(); let read_preference = self.query_ctx.read_preference(); let explain_verbose = self.query_ctx.explain_verbose(); + let live_analyze_metrics = explain_verbose && self.query_ctx.live_analyze_metrics_enabled(); let remote_dyn_filter_registry_lease = acquire_remote_dyn_filter_registry_lease( context.as_ref(), &query_ctx, @@ -398,7 +399,7 @@ impl MergeScanExec { remote_dyn_filter_registry_lease.as_ref(), &captured_remote_dyn_filters, ); - if explain_verbose { + if live_analyze_metrics { let remote_query_id = region_query_ctx.remote_query_id().map(str::to_string); if let Some(remote_query_id) = remote_query_id { region_query_ctx.set_extension( @@ -452,7 +453,7 @@ impl MergeScanExec { let mut poll_duration = Duration::ZERO; let mut poll_timer = Instant::now(); loop { - let batch = if explain_verbose { + let batch = if live_analyze_metrics { match time::timeout( FLIGHT_METRICS_HEARTBEAT_INTERVAL, stream.next().instrument(region_span.clone()), diff --git a/src/servers/src/grpc/flight/stream.rs b/src/servers/src/grpc/flight/stream.rs index 22dcd7baba..c222c72e1c 100644 --- a/src/servers/src/grpc/flight/stream.rs +++ b/src/servers/src/grpc/flight/stream.rs @@ -142,10 +142,12 @@ impl FlightRecordBatchStream { query_ctx: QueryContextRef, ) -> Self { let should_send_partial_metrics = query_ctx.explain_verbose(); - let can_send_metrics_before_batch = query_ctx - .remote_query_id() - .zip(query_ctx.extension(SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY)) - .is_some_and(|(remote_query_id, capability)| capability == remote_query_id); + let can_send_metrics_before_batch = query_ctx.explain_verbose() + && query_ctx.live_analyze_metrics_enabled() + && query_ctx + .remote_query_id() + .zip(query_ctx.extension(SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY)) + .is_some_and(|(remote_query_id, capability)| capability == remote_query_id); let (tx, rx) = mpsc::channel::>(1); let join_handle = common_runtime::spawn_global(async move { Self::flight_data_stream( @@ -335,7 +337,10 @@ mod test { use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::vectors::Int32Vector; use futures::StreamExt; - use session::context::{QueryContext, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY}; + use session::context::{ + LIVE_ANALYZE_METRICS_EXTENSION_KEY, QueryContext, + SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, + }; use super::*; @@ -364,12 +369,9 @@ mod test { Arc::new(query_ctx) } - fn query_context_with_capability(capability: &str) -> Arc { - let mut query_ctx = (*QueryContext::arc()).clone(); - query_ctx.set_extension( - SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, - capability, - ); + fn query_context_with_live_metrics_and_matching_capability() -> Arc { + let mut query_ctx = (*query_context_with_matching_capability()).clone(); + query_ctx.enable_live_analyze_metrics(); Arc::new(query_ctx) } @@ -481,7 +483,7 @@ mod test { schema: schema.clone(), metrics, }); - let query_ctx = query_context_with_matching_capability(); + let query_ctx = query_context_with_live_metrics_and_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( recordbatches, @@ -536,7 +538,7 @@ mod test { metrics, rx, }); - let query_ctx = query_context_with_matching_capability(); + let query_ctx = query_context_with_live_metrics_and_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( recordbatches, @@ -577,7 +579,7 @@ mod test { } #[tokio::test] - async fn test_flight_record_batch_stream_requires_capability_for_pre_batch_metrics() { + async fn test_flight_record_batch_stream_requires_live_metrics_for_pre_batch_metrics() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "a", ConcreteDataType::int32_datatype(), @@ -590,7 +592,7 @@ mod test { ..Default::default() }, }); - let query_ctx = QueryContext::arc(); + let query_ctx = query_context_with_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( recordbatches, @@ -612,12 +614,12 @@ mod test { ) .await .is_err(), - "pre-batch Metrics must be gated by client capability" + "pre-batch Metrics must be gated by live analyze metrics" ); } #[tokio::test] - async fn test_flight_record_batch_stream_rejects_spoofed_capability_for_pre_batch_metrics() { + async fn test_flight_record_batch_stream_rejects_spoofed_live_metrics_for_pre_batch_metrics() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "a", ConcreteDataType::int32_datatype(), @@ -630,7 +632,10 @@ mod test { ..Default::default() }, }); - let query_ctx = query_context_with_capability("true"); + let query_ctx = query_context_with_live_metrics_and_matching_capability(); + let mut query_ctx = (*query_ctx).clone(); + query_ctx.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, "true"); + let query_ctx = Arc::new(query_ctx); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( recordbatches, @@ -652,7 +657,7 @@ mod test { ) .await .is_err(), - "pre-batch Metrics must reject a spoofed capability" + "pre-batch Metrics must reject spoofed live analyze metrics" ); } diff --git a/src/servers/src/http/handler.rs b/src/servers/src/http/handler.rs index a6421d4e55..f92f2413c7 100644 --- a/src/servers/src/http/handler.rs +++ b/src/servers/src/http/handler.rs @@ -234,6 +234,7 @@ pub async fn sql_analyze_stream( query_ctx.set_current_schema(&schema); } query_ctx.set_channel(Channel::HttpSql); + query_ctx.enable_live_analyze_metrics(); let query_ctx = Arc::new(query_ctx); let Some(sql) = query_params.sql.or(form_params.sql) else { diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 827950123a..c554c0955f 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -33,7 +33,8 @@ use derive_builder::Builder; use sql::dialect::{Dialect, GenericDialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect}; pub use crate::hints::{ - REMOTE_QUERY_ID_EXTENSION_KEY, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, + LIVE_ANALYZE_METRICS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY, + SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, }; use crate::protocol_ctx::ProtocolCtx; use crate::query_id::QueryId; @@ -379,6 +380,18 @@ impl QueryContext { .and_then(|query_id| query_id.parse().ok()) } + pub fn enable_live_analyze_metrics(&mut self) { + if let Some(remote_query_id) = self.remote_query_id().map(str::to_string) { + self.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, remote_query_id); + } + } + + pub fn live_analyze_metrics_enabled(&self) -> bool { + self.remote_query_id() + .zip(self.extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY)) + .is_some_and(|(remote_query_id, value)| value == remote_query_id) + } + pub fn extensions(&self) -> HashMap { self.extensions.clone() } @@ -855,4 +868,19 @@ mod test { query_id ); } + + #[test] + fn test_live_analyze_metrics_requires_matching_remote_query_id() { + let mut ctx = QueryContext::arc().as_ref().clone(); + assert!(!ctx.live_analyze_metrics_enabled()); + + ctx.enable_live_analyze_metrics(); + assert!(ctx.live_analyze_metrics_enabled()); + + ctx.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, "true"); + assert!(!ctx.live_analyze_metrics_enabled()); + + ctx.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, "another-query-id"); + assert!(!ctx.live_analyze_metrics_enabled()); + } } diff --git a/src/session/src/hints.rs b/src/session/src/hints.rs index 35f5101e45..e656b34c82 100644 --- a/src/session/src/hints.rs +++ b/src/session/src/hints.rs @@ -21,12 +21,14 @@ pub const INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY: &str = "initial_remote_dyn_filter_registrations"; pub const SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY: &str = "query.support_flight_metrics_before_batch"; +pub const LIVE_ANALYZE_METRICS_EXTENSION_KEY: &str = "query.live_analyze_metrics"; pub const READ_PREFERENCE_HINT: &str = "read_preference"; -pub const RESERVED_EXTENSION_KEYS: [&str; 3] = [ +pub const RESERVED_EXTENSION_KEYS: [&str; 4] = [ REMOTE_QUERY_ID_EXTENSION_KEY, INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, + LIVE_ANALYZE_METRICS_EXTENSION_KEY, ]; /// Deprecated, use `HINTS_KEY` instead. @@ -57,6 +59,9 @@ mod tests { assert!(is_reserved_extension_key( SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY )); + assert!(is_reserved_extension_key( + LIVE_ANALYZE_METRICS_EXTENSION_KEY + )); assert!(!is_reserved_extension_key(READ_PREFERENCE_HINT)); } }