rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-08-18 12:08:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: scope live analyze metrics to streaming requests (#8644)

Browse Source 
Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
2026-07-28 17:19:23 +08:00
committed by GitHub
parent 775a9af3b8
commit 255572020a
6 changed files with 66 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/datanode/src/region_server.rs
+3 -2
View File
@@ -1805,7 +1805,8 @@ impl RegionServerInner {
request: QueryRequest,
query_ctx: QueryContextRef,
) -> Result<SendableRecordBatchStream> {
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)) => {
src/query/src/dist_plan/merge_scan.rs
+3 -2
View File
@@ -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()),
src/servers/src/grpc/flight/stream.rs
+24 -19
View File
@@ -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::<TonicResult<FlightMessage>>(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<QueryContext> {
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<QueryContext> {
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"
);
}
src/servers/src/http/handler.rs
+1
View File
@@ -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 {
src/session/src/context.rs
+29 -1
View File
@@ -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<String, String> {
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());
}
}
src/session/src/hints.rs
+6 -1
View File
@@ -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));
}
}