mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-09-05 21:18:57 +00:00
feat!: stabilize streaming analyze metrics (#8966)
* feat: stabilize streaming analyze metrics Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * feat: expose analyze memory usage Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor: simplify analyze stream handling Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix: preserve analyze stream sequence on panic Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * chore: log analyze stream worker panic Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
@@ -37,7 +37,6 @@
|
||||
| `http.body_limit` | String | `64MB` | HTTP request body limit.<br/>The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.<br/>Set to 0 to disable limit. |
|
||||
| `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default<br/>This allows browser to access http APIs without CORS restrictions |
|
||||
| `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. |
|
||||
| `http.experimental_enable_explain_analyze_stream` | Bool | `true` | Experimental: enable POST /v1/sql/analyze/stream for streaming EXPLAIN ANALYZE VERBOSE metrics. |
|
||||
| `http.enable_api_server` | Bool | `false` | Whether to start the dedicated public HTTP **API** server. This server serves<br/>only the `v1` interfaces plus the dashboard, and shares every other `[http]`<br/>option with the main server. It is disabled by default; set to `true` to enable. |
|
||||
| `http.api_server_addr` | String | `127.0.0.1:4006` | The address to bind the dedicated HTTP API server, in the same form as `addr`.<br/>Defaults to `127.0.0.1:4006`. |
|
||||
| `grpc` | -- | -- | The gRPC server options. |
|
||||
@@ -274,7 +273,6 @@
|
||||
| `http.body_limit` | String | `64MB` | HTTP request body limit.<br/>The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.<br/>Set to 0 to disable limit. |
|
||||
| `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default<br/>This allows browser to access http APIs without CORS restrictions |
|
||||
| `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. |
|
||||
| `http.experimental_enable_explain_analyze_stream` | Bool | `true` | Experimental: enable POST /v1/sql/analyze/stream for streaming EXPLAIN ANALYZE VERBOSE metrics. |
|
||||
| `http.enable_api_server` | Bool | `false` | Whether to start the dedicated public HTTP **API** server. This server serves<br/>only the `v1` interfaces plus the dashboard, and shares every other `[http]`<br/>option with the main server. It is disabled by default; set to `true` to enable. |
|
||||
| `http.api_server_addr` | String | `127.0.0.1:4006` | The address to bind the dedicated HTTP API server, in the same form as `addr`.<br/>Defaults to `127.0.0.1:4006`. |
|
||||
| `grpc` | -- | -- | The gRPC server options. |
|
||||
|
||||
@@ -65,8 +65,6 @@ enable_cors = true
|
||||
## Customize allowed origins for HTTP CORS.
|
||||
## @toml2docs:none-default
|
||||
cors_allowed_origins = ["https://example.com"]
|
||||
## Experimental: enable POST /v1/sql/analyze/stream for streaming EXPLAIN ANALYZE VERBOSE metrics.
|
||||
experimental_enable_explain_analyze_stream = true
|
||||
|
||||
## Whether to start the dedicated public HTTP **API** server. This server serves
|
||||
## only the `v1` interfaces plus the dashboard, and shares every other `[http]`
|
||||
|
||||
@@ -91,8 +91,6 @@ enable_cors = true
|
||||
## @toml2docs:none-default
|
||||
cors_allowed_origins = ["https://example.com"]
|
||||
|
||||
## Experimental: enable POST /v1/sql/analyze/stream for streaming EXPLAIN ANALYZE VERBOSE metrics.
|
||||
experimental_enable_explain_analyze_stream = true
|
||||
|
||||
## Whether to start the dedicated public HTTP **API** server. This server serves
|
||||
## only the `v1` interfaces plus the dashboard, and shares every other `[http]`
|
||||
|
||||
@@ -607,7 +607,6 @@ mod test {
|
||||
vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let mut recordbatches = recordbatches_from_flight_message_stream(
|
||||
"test-peer".to_string(),
|
||||
stream::iter(vec![
|
||||
@@ -679,6 +678,35 @@ mod test {
|
||||
assert!(recordbatches.next().await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_record_batch_stream_captures_final_metrics_after_record_batch() {
|
||||
let schema = test_schema();
|
||||
let batch = RecordBatch::new(
|
||||
schema.clone(),
|
||||
vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
|
||||
)
|
||||
.unwrap();
|
||||
let final_metrics = serde_json::to_string(&RecordBatchMetrics {
|
||||
elapsed_compute: 99,
|
||||
..Default::default()
|
||||
})
|
||||
.unwrap();
|
||||
let mut recordbatches = recordbatches_from_flight_message_stream(
|
||||
"test-peer".to_string(),
|
||||
stream::iter(vec![
|
||||
Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
|
||||
Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())),
|
||||
Ok(FlightMessage::Metrics(final_metrics)),
|
||||
]),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(recordbatches.next().await.unwrap().unwrap().num_rows(), 1);
|
||||
assert!(recordbatches.next().await.is_none());
|
||||
assert_eq!(recordbatches.metrics().unwrap().elapsed_compute, 99);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_record_batch_stream_exposes_error_after_pre_batch_metrics() {
|
||||
let schema = test_schema();
|
||||
|
||||
@@ -327,7 +327,9 @@ impl RecordBatchStreamAdapter {
|
||||
query_load_region_id: Option<u64>,
|
||||
) -> RecordBatchMetrics {
|
||||
if explain_verbose {
|
||||
collect_full_metrics(df_plan, explain_verbose, query_load_region_id)
|
||||
// Verbose in-progress snapshots have the same complete topology as
|
||||
// the final snapshot, but use compact (Default) plan formatting.
|
||||
collect_full_metrics(df_plan, false, query_load_region_id)
|
||||
} else {
|
||||
collect_lightweight_query_load_metrics(df_plan, query_load_region_id)
|
||||
}
|
||||
@@ -479,13 +481,17 @@ impl RecordBatchStream for RecordBatchStreamAdapter {
|
||||
match &self.metrics_2 {
|
||||
Metrics::Unresolved(df_plan) => {
|
||||
if self.explain_verbose {
|
||||
Some(self.collect_plan_metrics(df_plan))
|
||||
Some(Self::collect_partial_metrics(
|
||||
df_plan.as_ref(),
|
||||
true,
|
||||
self.query_load_region_id,
|
||||
))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
Metrics::PartialResolved(df_plan, metrics) => Some(if self.explain_verbose {
|
||||
self.collect_plan_metrics(df_plan)
|
||||
Self::collect_partial_metrics(df_plan.as_ref(), true, self.query_load_region_id)
|
||||
} else {
|
||||
metrics.clone()
|
||||
}),
|
||||
@@ -954,6 +960,7 @@ mod test {
|
||||
struct TestMetricsExec {
|
||||
properties: Arc<PlanProperties>,
|
||||
metrics: ExecutionPlanMetricsSet,
|
||||
format_plan: bool,
|
||||
}
|
||||
|
||||
impl TestMetricsExec {
|
||||
@@ -961,7 +968,19 @@ mod test {
|
||||
Self::with_output_bytes(schema, &[24])
|
||||
}
|
||||
|
||||
fn new_without_plan_formatting(schema: DfSchemaRef) -> Self {
|
||||
Self::with_output_bytes_and_formatting(schema, &[24], false)
|
||||
}
|
||||
|
||||
fn with_output_bytes(schema: DfSchemaRef, output_bytes_by_partition: &[usize]) -> Self {
|
||||
Self::with_output_bytes_and_formatting(schema, output_bytes_by_partition, true)
|
||||
}
|
||||
|
||||
fn with_output_bytes_and_formatting(
|
||||
schema: DfSchemaRef,
|
||||
output_bytes_by_partition: &[usize],
|
||||
format_plan: bool,
|
||||
) -> Self {
|
||||
let metrics = ExecutionPlanMetricsSet::new();
|
||||
let elapsed_compute = MetricBuilder::new(&metrics).elapsed_compute(0);
|
||||
elapsed_compute.add_duration(Duration::from_nanos(42));
|
||||
@@ -978,13 +997,22 @@ mod test {
|
||||
Boundedness::Bounded,
|
||||
)),
|
||||
metrics,
|
||||
format_plan,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl DisplayAs for TestMetricsExec {
|
||||
fn fmt_as(&self, _t: DisplayFormatType, _f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
||||
panic!("non-verbose lightweight partial metrics must not format the plan")
|
||||
fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
||||
assert!(
|
||||
self.format_plan,
|
||||
"non-verbose lightweight partial metrics must not format the plan"
|
||||
);
|
||||
write!(f, "RegionScanExec")?;
|
||||
if matches!(t, DisplayFormatType::Verbose) {
|
||||
write!(f, ": files=[file-1.parquet]")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1070,7 +1098,9 @@ mod test {
|
||||
futures::stream::iter(vec![Ok(batch1), Ok(batch2)]),
|
||||
),
|
||||
);
|
||||
let plan = Arc::new(TestMetricsExec::new(schema.arrow_schema().clone()));
|
||||
let plan = Arc::new(TestMetricsExec::new_without_plan_formatting(
|
||||
schema.arrow_schema().clone(),
|
||||
));
|
||||
|
||||
let mut adapter = RecordBatchStreamAdapter::try_new(df_stream).unwrap();
|
||||
adapter.set_metrics2(plan);
|
||||
@@ -1087,12 +1117,49 @@ mod test {
|
||||
assert_eq!(metrics.plan_metrics.len(), 1);
|
||||
assert_eq!(metrics.plan_metrics[0].plan, REGION_SCAN_EXEC_NAME);
|
||||
assert_eq!(metrics.plan_metrics[0].plan_name, REGION_SCAN_EXEC_NAME);
|
||||
assert_eq!(
|
||||
metrics.plan_metrics[0].metrics,
|
||||
vec![("output_bytes".to_string(), 24)]
|
||||
assert!(
|
||||
metrics.plan_metrics[0]
|
||||
.metrics
|
||||
.iter()
|
||||
.any(|(name, value)| name == "output_bytes" && *value == 24)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_record_batch_stream_adapter_uses_compact_partial_and_verbose_final_metrics() {
|
||||
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
|
||||
"a",
|
||||
ConcreteDataType::int32_datatype(),
|
||||
false,
|
||||
)]));
|
||||
let batch = RecordBatch::new(
|
||||
schema.clone(),
|
||||
vec![Arc::new(Int32Vector::from_slice([1])) as _],
|
||||
)
|
||||
.unwrap()
|
||||
.into_df_record_batch();
|
||||
let df_stream = Box::pin(
|
||||
datafusion::physical_plan::stream::RecordBatchStreamAdapter::new(
|
||||
schema.arrow_schema().clone(),
|
||||
futures::stream::iter(vec![Ok(batch)]),
|
||||
),
|
||||
);
|
||||
let plan = Arc::new(TestMetricsExec::new(schema.arrow_schema().clone()));
|
||||
let mut adapter = RecordBatchStreamAdapter::try_new(df_stream).unwrap();
|
||||
adapter.set_metrics2(plan);
|
||||
adapter.set_explain_verbose(true);
|
||||
|
||||
adapter.next().await.unwrap().unwrap();
|
||||
let partial = adapter.metrics().unwrap();
|
||||
assert_eq!(partial.plan_metrics.len(), 1);
|
||||
assert_eq!(partial.plan_metrics[0].plan.trim_end(), "RegionScanExec");
|
||||
assert!(!partial.plan_metrics[0].plan.contains("files"));
|
||||
|
||||
assert!(adapter.next().await.is_none());
|
||||
let final_metrics = adapter.metrics().unwrap();
|
||||
assert!(final_metrics.plan_metrics[0].plan.contains("files"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_record_batch_stream_adapter_reuses_partial_query_stats_on_drop() {
|
||||
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
|
||||
|
||||
@@ -231,31 +231,37 @@ fn parse_stmt(sql: &str, dialect: &(dyn Dialect + Send + Sync)) -> Result<Vec<St
|
||||
|
||||
fn is_explain_analyze_verbose(stmt: &Statement) -> bool {
|
||||
matches!(stmt, Statement::Explain(explain) if explain.analyze && explain.verbose)
|
||||
|| matches!(stmt, Statement::Tql(Tql::Analyze(analyze)) if analyze.is_verbose)
|
||||
}
|
||||
|
||||
fn validate_analyze_stream_statement(stmt: &mut Statement) -> Result<()> {
|
||||
let Statement::Explain(explain) = stmt else {
|
||||
return InvalidSqlSnafu {
|
||||
err_msg: "only EXPLAIN ANALYZE VERBOSE statement is supported",
|
||||
let (is_verbose, format) = match stmt {
|
||||
Statement::Explain(explain) => (explain.analyze && explain.verbose, &mut explain.format),
|
||||
Statement::Tql(Tql::Analyze(analyze)) => (analyze.is_verbose, &mut analyze.format),
|
||||
_ => {
|
||||
return InvalidSqlSnafu {
|
||||
err_msg: "only EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE statement is supported",
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
.fail();
|
||||
};
|
||||
|
||||
ensure!(
|
||||
explain.analyze && explain.verbose,
|
||||
is_verbose,
|
||||
InvalidSqlSnafu {
|
||||
err_msg: "statement must be EXPLAIN ANALYZE VERBOSE"
|
||||
err_msg: "statement must be EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE"
|
||||
}
|
||||
);
|
||||
match explain.format {
|
||||
match format {
|
||||
None | Some(AnalyzeFormat::JSON) => {
|
||||
// Keep explicit FORMAT JSON accepted, but pass JSON through
|
||||
// QueryContext.explain_format instead of the statement to avoid the
|
||||
// planner's current `EXPLAIN VERBOSE with FORMAT` limitation.
|
||||
explain.format = None;
|
||||
// QueryContext.explain_format instead of the statement to avoid
|
||||
// the planner's current `EXPLAIN VERBOSE with FORMAT` limitation.
|
||||
*format = None;
|
||||
Ok(())
|
||||
}
|
||||
Some(_) => InvalidSqlSnafu {
|
||||
err_msg: "only FORMAT JSON is supported for analyze stream",
|
||||
err_msg: "only FORMAT JSON is supported for EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE",
|
||||
}
|
||||
.fail(),
|
||||
}
|
||||
@@ -711,7 +717,7 @@ impl Instance {
|
||||
ensure!(
|
||||
stmts.len() == 1,
|
||||
InvalidSqlSnafu {
|
||||
err_msg: "only single EXPLAIN ANALYZE VERBOSE statement is supported"
|
||||
err_msg: "only a single EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE statement is supported"
|
||||
}
|
||||
);
|
||||
let mut stmt = stmts.remove(0);
|
||||
@@ -1877,6 +1883,9 @@ mod tests {
|
||||
"explain analyze select 1",
|
||||
"explain analyze verbose format text select 1",
|
||||
"explain analyze verbose format graphviz select 1",
|
||||
"TQL ANALYZE (0, 10, '5s') physical_metric",
|
||||
"TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
|
||||
"TQL ANALYZE VERBOSE FORMAT TEXT (0, 10, '5s') physical_metric",
|
||||
] {
|
||||
let mut stmts = parse_test_sql(sql);
|
||||
assert!(
|
||||
@@ -1888,16 +1897,19 @@ mod tests {
|
||||
for sql in [
|
||||
"explain analyze verbose select 1",
|
||||
"explain analyze verbose format json select 1",
|
||||
"TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric",
|
||||
"TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') physical_metric",
|
||||
] {
|
||||
let mut stmts = parse_test_sql(sql);
|
||||
assert!(
|
||||
validate_analyze_stream_statement(&mut stmts[0]).is_ok(),
|
||||
"{sql}"
|
||||
);
|
||||
let Statement::Explain(explain) = &stmts[0] else {
|
||||
unreachable!();
|
||||
};
|
||||
assert!(explain.format.is_none());
|
||||
match &stmts[0] {
|
||||
Statement::Explain(explain) => assert!(explain.format.is_none()),
|
||||
Statement::Tql(Tql::Analyze(analyze)) => assert!(analyze.format.is_none()),
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
@@ -1908,11 +1920,16 @@ mod tests {
|
||||
assert!(is_explain_analyze_verbose(
|
||||
&parse_test_sql("explain analyze verbose select 1")[0]
|
||||
));
|
||||
assert!(is_explain_analyze_verbose(
|
||||
&parse_test_sql("TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric")[0]
|
||||
));
|
||||
for sql in [
|
||||
"select 1",
|
||||
"explain select 1",
|
||||
"explain analyze select 1",
|
||||
"explain verbose select 1",
|
||||
"TQL ANALYZE (0, 10, '5s') physical_metric",
|
||||
"TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
|
||||
] {
|
||||
assert!(
|
||||
!is_explain_analyze_verbose(&parse_test_sql(sql)[0]),
|
||||
|
||||
@@ -332,6 +332,8 @@ struct JsonMetrics {
|
||||
|
||||
// other metrics
|
||||
metrics: HashMap<String, usize>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
memory_usage: Option<usize>,
|
||||
children: Vec<JsonMetrics>,
|
||||
}
|
||||
|
||||
@@ -339,12 +341,14 @@ impl JsonMetrics {
|
||||
fn from_record_batch_metrics(record_batch_metrics: RecordBatchMetrics) -> Self {
|
||||
let mut layers: HashMap<usize, Vec<Self>> = HashMap::default();
|
||||
|
||||
let memory_usage = record_batch_metrics.memory_usage;
|
||||
for plan_metrics in record_batch_metrics.plan_metrics.into_iter().rev() {
|
||||
let (level, mut metrics) = Self::from_plan_metrics(plan_metrics);
|
||||
if let Some(next_layer) = layers.remove(&(level + 1)) {
|
||||
metrics.children = next_layer;
|
||||
}
|
||||
if level == 0 {
|
||||
metrics.memory_usage = Some(memory_usage);
|
||||
return metrics;
|
||||
}
|
||||
layers.entry(level).or_default().push(metrics);
|
||||
@@ -382,6 +386,7 @@ impl JsonMetrics {
|
||||
output_rows,
|
||||
elapsed_compute,
|
||||
metrics: other_metrics,
|
||||
memory_usage: None,
|
||||
children: vec![],
|
||||
},
|
||||
)
|
||||
@@ -498,4 +503,32 @@ mod tests {
|
||||
assert_eq!(metrics.name, "FilterExec");
|
||||
assert_eq!(metrics.param, "predicate");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qbs_analyze_json_includes_memory_usage_only_at_root() {
|
||||
let metrics = JsonMetrics::from_record_batch_metrics(RecordBatchMetrics {
|
||||
memory_usage: 42,
|
||||
plan_metrics: vec![
|
||||
PlanMetrics {
|
||||
plan: "RootExec".to_string(),
|
||||
plan_name: "RootExec".to_string(),
|
||||
level: 0,
|
||||
metrics: vec![("mem_used".to_string(), 24)],
|
||||
},
|
||||
PlanMetrics {
|
||||
plan: "ChildExec".to_string(),
|
||||
plan_name: "ChildExec".to_string(),
|
||||
level: 1,
|
||||
metrics: vec![("mem_used".to_string(), 18)],
|
||||
},
|
||||
],
|
||||
..Default::default()
|
||||
});
|
||||
let value = serde_json::to_value(metrics).unwrap();
|
||||
|
||||
assert_eq!(value["memory_usage"], 42);
|
||||
assert!(value["children"][0].get("memory_usage").is_none());
|
||||
assert_eq!(value["metrics"]["mem_used"], 24);
|
||||
assert_eq!(value["children"][0]["metrics"]["mem_used"], 18);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -754,6 +754,19 @@ impl MergeScanExec {
|
||||
stream.next().instrument(region_span.clone()).await
|
||||
};
|
||||
let Some(batch) = batch else {
|
||||
// The remote Flight stream publishes its terminal metrics
|
||||
// immediately after EOF. Capture them before leaving the
|
||||
// loop so the final verbose snapshot is not lost.
|
||||
if let Some(metrics) = stream.metrics() {
|
||||
let load = region_scan_load(&metrics);
|
||||
let (c, s) = parse_catalog_and_schema_from_db_string(&dbname);
|
||||
let value = read_meter!(c, s, load, current_channel as u8);
|
||||
metric.record_greptime_exec_cost(value as usize);
|
||||
sub_stage_metrics_moved
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(region_id, metrics);
|
||||
}
|
||||
break;
|
||||
};
|
||||
let poll_elapsed = poll_timer.elapsed();
|
||||
@@ -823,18 +836,6 @@ impl MergeScanExec {
|
||||
);
|
||||
}
|
||||
|
||||
// process metrics after all data is drained.
|
||||
if let Some(metrics) = stream.metrics() {
|
||||
let load = region_scan_load(&metrics);
|
||||
let (c, s) = parse_catalog_and_schema_from_db_string(&dbname);
|
||||
let value = read_meter!(c, s, load, current_channel as u8);
|
||||
metric.record_greptime_exec_cost(value as usize);
|
||||
|
||||
// record metrics from sub sgates
|
||||
let mut sub_stage_metrics = sub_stage_metrics_moved.lock().unwrap();
|
||||
sub_stage_metrics.insert(region_id, metrics);
|
||||
}
|
||||
|
||||
MERGE_SCAN_POLL_ELAPSED.observe(poll_duration.as_secs_f64());
|
||||
}
|
||||
|
||||
|
||||
@@ -700,6 +700,17 @@ mod test {
|
||||
}
|
||||
other => panic!("expected record batch after pending metrics, got {other:?}"),
|
||||
}
|
||||
|
||||
drop(tx);
|
||||
let final_metrics_data = tokio::time::timeout(Duration::from_secs(2), stream.next())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
decoder.try_decode(&final_metrics_data).unwrap().unwrap(),
|
||||
FlightMessage::Metrics(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
+36
-44
@@ -253,8 +253,6 @@ pub struct HttpOptions {
|
||||
|
||||
pub enable_cors: bool,
|
||||
|
||||
pub experimental_enable_explain_analyze_stream: bool,
|
||||
|
||||
/// Whether to start the dedicated public HTTP **API** server, which serves
|
||||
/// only the `v1` interfaces plus the dashboard. It shares every other
|
||||
/// `[http]` option with the main server and only differs by its bound
|
||||
@@ -274,7 +272,6 @@ impl Default for HttpOptions {
|
||||
body_limit: DEFAULT_BODY_LIMIT,
|
||||
cors_allowed_origins: Vec::new(),
|
||||
enable_cors: true,
|
||||
experimental_enable_explain_analyze_stream: true,
|
||||
enable_api_server: false,
|
||||
api_server_addr: format!("127.0.0.1:{}", DEFAULT_HTTP_API_ADDR_PORT),
|
||||
}
|
||||
@@ -601,7 +598,6 @@ impl From<NullResponse> for HttpResponse {
|
||||
#[derive(Clone)]
|
||||
pub struct ApiState {
|
||||
pub sql_handler: ServerSqlQueryHandlerRef,
|
||||
pub experimental_enable_explain_analyze_stream: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -638,12 +634,7 @@ impl HttpServerBuilder {
|
||||
}
|
||||
|
||||
pub fn with_sql_handler(self, sql_handler: ServerSqlQueryHandlerRef) -> Self {
|
||||
let sql_router = HttpServer::route_sql(ApiState {
|
||||
sql_handler,
|
||||
experimental_enable_explain_analyze_stream: self
|
||||
.options
|
||||
.experimental_enable_explain_analyze_stream,
|
||||
});
|
||||
let sql_router = HttpServer::route_sql(ApiState { sql_handler });
|
||||
|
||||
Self {
|
||||
router: self
|
||||
@@ -1315,7 +1306,7 @@ impl HttpServer {
|
||||
}
|
||||
|
||||
fn route_sql<S>(api_state: ApiState) -> Router<S> {
|
||||
let mut router = Router::new()
|
||||
Router::new()
|
||||
.route("/sql", routing::get(handler::sql).post(handler::sql))
|
||||
.route(
|
||||
"/sql/parse",
|
||||
@@ -1328,16 +1319,12 @@ impl HttpServer {
|
||||
.route(
|
||||
"/promql",
|
||||
routing::get(handler::promql).post(handler::promql),
|
||||
);
|
||||
|
||||
if api_state.experimental_enable_explain_analyze_stream {
|
||||
router = router.route(
|
||||
)
|
||||
.route(
|
||||
"/sql/analyze/stream",
|
||||
routing::post(handler::sql_analyze_stream),
|
||||
);
|
||||
}
|
||||
|
||||
router.with_state(api_state)
|
||||
)
|
||||
.with_state(api_state)
|
||||
}
|
||||
|
||||
fn route_logs<S>(log_handler: LogQueryHandlerRef) -> Router<S> {
|
||||
@@ -1652,31 +1639,6 @@ mod test {
|
||||
server.build(app).unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
pub async fn test_analyze_stream_route_config_gate() {
|
||||
let (tx, _rx) = mpsc::channel(100);
|
||||
let options = HttpOptions {
|
||||
experimental_enable_explain_analyze_stream: false,
|
||||
..Default::default()
|
||||
};
|
||||
let app = make_test_app_custom(tx, options);
|
||||
let client = TestClient::new(app).await;
|
||||
let res = client
|
||||
.post("/v1/sql/analyze/stream?sql=EXPLAIN%20ANALYZE%20VERBOSE%20SELECT%201")
|
||||
.send()
|
||||
.await;
|
||||
assert_eq!(res.status(), StatusCode::NOT_FOUND);
|
||||
|
||||
let (tx, _rx) = mpsc::channel(100);
|
||||
let app = make_test_app_custom(tx, HttpOptions::default());
|
||||
let client = TestClient::new(app).await;
|
||||
let res = client
|
||||
.post("/v1/sql/analyze/stream?sql=EXPLAIN%20ANALYZE%20VERBOSE%20SELECT%201")
|
||||
.send()
|
||||
.await;
|
||||
assert_ne!(res.status(), StatusCode::NOT_FOUND);
|
||||
}
|
||||
|
||||
fn make_split_builder() -> HttpServerBuilder {
|
||||
let (tx, _rx) = mpsc::channel(100);
|
||||
let instance = Arc::new(DummyInstance { _tx: tx });
|
||||
@@ -1987,6 +1949,36 @@ mod test {
|
||||
assert_eq!(Duration::from_secs(0), default.timeout)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_http_options_legacy_analyze_stream_config_is_ignored() {
|
||||
let options: HttpOptions = serde_json::from_value(serde_json::json!({
|
||||
"addr": "127.0.0.1:4000",
|
||||
"timeout": "0s",
|
||||
"body_limit": "64MiB",
|
||||
"cors_allowed_origins": [],
|
||||
"enable_cors": true,
|
||||
"experimental_enable_explain_analyze_stream": false,
|
||||
"enable_api_server": false,
|
||||
"api_server_addr": "127.0.0.1:4006"
|
||||
}))
|
||||
.unwrap();
|
||||
let serialized = serde_json::to_string(&options).unwrap();
|
||||
assert!(!serialized.contains("experimental_enable_explain_analyze_stream"));
|
||||
|
||||
let (tx, _rx) = mpsc::channel(100);
|
||||
let app = make_test_app_custom(tx, options);
|
||||
let client = TestClient::new(app).await;
|
||||
let response = client
|
||||
.post("/v1/sql/analyze/stream")
|
||||
.form(&handler::SqlQuery {
|
||||
sql: Some("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
|
||||
..Default::default()
|
||||
})
|
||||
.send()
|
||||
.await;
|
||||
assert_ne!(response.status(), StatusCode::NOT_FOUND);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_http_server_request_timeout() {
|
||||
common_telemetry::init_default_ut_logging();
|
||||
|
||||
+199
-84
@@ -13,7 +13,9 @@
|
||||
// limitations under the License.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::panic::AssertUnwindSafe;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use axum::extract::rejection::FormRejection;
|
||||
@@ -26,11 +28,10 @@ use common_error::ext::ErrorExt;
|
||||
use common_error::status_code::StatusCode;
|
||||
use common_plugins::GREPTIME_EXEC_WRITE_COST;
|
||||
use common_query::{Output, OutputData};
|
||||
use common_recordbatch::{RecordBatch, SendableRecordBatchStream, util};
|
||||
use common_recordbatch::util;
|
||||
use common_telemetry::tracing;
|
||||
use datafusion::physical_plan::ExecutionPlan;
|
||||
use datatypes::schema::SchemaRef;
|
||||
use futures::StreamExt;
|
||||
use futures::{FutureExt, StreamExt};
|
||||
use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
@@ -39,6 +40,7 @@ use snafu::ResultExt;
|
||||
use sql::dialect::GreptimeDbDialect;
|
||||
use sql::parser::{ParseOptions, ParserContext};
|
||||
use sql::statements::statement::Statement;
|
||||
use tokio::sync::{Notify, watch};
|
||||
|
||||
use crate::error::{FailedToParseQuerySnafu, InvalidQuerySnafu, Result};
|
||||
use crate::http::header::collect_plan_metrics;
|
||||
@@ -100,15 +102,32 @@ struct AnalyzeStreamPayload {
|
||||
code: Option<u32>,
|
||||
}
|
||||
|
||||
struct AnalyzeStreamState {
|
||||
stream: SendableRecordBatchStream,
|
||||
schema: SchemaRef,
|
||||
plan: Option<Arc<dyn ExecutionPlan>>,
|
||||
batches: Vec<RecordBatch>,
|
||||
seq: u64,
|
||||
start: Instant,
|
||||
requested_interval_ms: u64,
|
||||
current_interval_ms: u64,
|
||||
#[derive(Clone, Debug)]
|
||||
#[doc(hidden)]
|
||||
pub struct AnalyzeStreamMessage {
|
||||
pub event_name: &'static str,
|
||||
pub payload: String,
|
||||
}
|
||||
|
||||
struct AnalyzeStreamWorkerGuard {
|
||||
cancel: watch::Sender<bool>,
|
||||
handle: common_runtime::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl Drop for AnalyzeStreamWorkerGuard {
|
||||
fn drop(&mut self) {
|
||||
let _ = self.cancel.send(true);
|
||||
self.handle.abort();
|
||||
}
|
||||
}
|
||||
|
||||
struct AnalyzeStreamBodyState {
|
||||
latest_metrics: watch::Receiver<Option<String>>,
|
||||
terminal: watch::Receiver<Option<AnalyzeStreamMessage>>,
|
||||
notify: Arc<Notify>,
|
||||
// Keeping the guard in the body state makes dropping the response cancel and
|
||||
// abort the worker instead of leaving an owned query stream detached.
|
||||
_worker: AnalyzeStreamWorkerGuard,
|
||||
done: bool,
|
||||
}
|
||||
|
||||
@@ -195,7 +214,7 @@ pub async fn sql(
|
||||
|
||||
/// Handler to stream partial `EXPLAIN ANALYZE VERBOSE` metrics as SSE.
|
||||
///
|
||||
/// This experimental endpoint is POST-only SSE, so browser `EventSource` does
|
||||
/// This endpoint is POST-only SSE, so browser `EventSource` does
|
||||
/// not apply. Each `metrics` event carries a complete snapshot (not a delta);
|
||||
/// large snapshots are throttled but never truncated. `final`, `canceled`, and
|
||||
/// `error` are terminal events. If the client disconnects it won't receive a
|
||||
@@ -284,88 +303,182 @@ pub async fn sql_analyze_stream(
|
||||
};
|
||||
let schema = stream.schema();
|
||||
|
||||
let sse_stream = futures::stream::unfold(
|
||||
AnalyzeStreamState {
|
||||
stream,
|
||||
schema,
|
||||
plan,
|
||||
batches: Vec::new(),
|
||||
seq: 0,
|
||||
start,
|
||||
requested_interval_ms: interval_ms,
|
||||
current_interval_ms: interval_ms,
|
||||
let (metrics_tx, metrics_rx) = watch::channel::<Option<String>>(None);
|
||||
let (terminal_tx, terminal_rx) = watch::channel::<Option<AnalyzeStreamMessage>>(None);
|
||||
let notify = Arc::new(Notify::new());
|
||||
let (cancel_tx, mut cancel_rx) = watch::channel(false);
|
||||
let worker_notify = notify.clone();
|
||||
let panic_notify = notify.clone();
|
||||
let sequence = Arc::new(AtomicU64::new(0));
|
||||
let worker_sequence = sequence.clone();
|
||||
let worker_terminal_tx = terminal_tx.clone();
|
||||
let worker = common_runtime::spawn_global(async move {
|
||||
let worker_result = AssertUnwindSafe(async move {
|
||||
let mut stream = stream;
|
||||
let mut batches = Vec::new();
|
||||
let mut current_interval_ms = interval_ms;
|
||||
let tick = tokio::time::sleep(Duration::from_millis(current_interval_ms));
|
||||
tokio::pin!(tick);
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = cancel_rx.changed() => return,
|
||||
item = stream.next() => {
|
||||
match item {
|
||||
Some(Ok(next_batch)) => batches.push(next_batch),
|
||||
Some(Err(err)) => {
|
||||
let status = err.status_code();
|
||||
let event_name = if status == StatusCode::Cancelled { "canceled" } else { "error" };
|
||||
let (payload, _) = make_analyze_payload(AnalyzePayloadArgs {
|
||||
seq: worker_sequence.load(Ordering::Relaxed),
|
||||
state: event_name,
|
||||
partial: false,
|
||||
start,
|
||||
plan: plan.as_ref(),
|
||||
output: None,
|
||||
reason: Some(err.output_msg()),
|
||||
code: Some(status as u32),
|
||||
});
|
||||
send_analyze_terminal(&terminal_tx, &worker_notify, event_name, payload);
|
||||
return;
|
||||
}
|
||||
None => {
|
||||
let output = HttpRecordsOutput::try_new(schema.clone(), batches)
|
||||
.map(GreptimeQueryOutput::Records);
|
||||
let (event_name, payload) = make_final_analyze_event(
|
||||
output.map_err(|err| (err.output_msg(), err.status_code() as u32)),
|
||||
worker_sequence.load(Ordering::Relaxed),
|
||||
start,
|
||||
plan.as_ref(),
|
||||
);
|
||||
send_analyze_terminal(&terminal_tx, &worker_notify, event_name, payload);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = &mut tick, if plan.is_some() => {
|
||||
let (payload, payload_bytes) = make_analyze_payload(AnalyzePayloadArgs {
|
||||
seq: worker_sequence.load(Ordering::Relaxed),
|
||||
state: "metrics",
|
||||
partial: true,
|
||||
start,
|
||||
plan: plan.as_ref(),
|
||||
output: None,
|
||||
reason: None,
|
||||
code: None,
|
||||
});
|
||||
current_interval_ms = adaptive_interval_ms(payload_bytes, interval_ms);
|
||||
worker_sequence.fetch_add(1, Ordering::Relaxed);
|
||||
if metrics_tx.send(Some(payload)).is_err() {
|
||||
return;
|
||||
}
|
||||
worker_notify.notify_one();
|
||||
tick.as_mut().reset(tokio::time::Instant::now() + Duration::from_millis(current_interval_ms));
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
.catch_unwind()
|
||||
.await;
|
||||
|
||||
if worker_result.is_err() {
|
||||
tracing::debug!("analyze stream worker panicked");
|
||||
let (payload, _) = make_analyze_payload(AnalyzePayloadArgs {
|
||||
seq: sequence.load(Ordering::Relaxed),
|
||||
state: "error",
|
||||
partial: false,
|
||||
start,
|
||||
plan: None,
|
||||
output: None,
|
||||
reason: Some("analyze stream worker panicked".to_string()),
|
||||
code: Some(StatusCode::Internal as u32),
|
||||
});
|
||||
send_analyze_terminal(&worker_terminal_tx, &panic_notify, "error", payload);
|
||||
}
|
||||
});
|
||||
|
||||
let sse_stream = analyze_stream_body(metrics_rx, terminal_rx, notify, cancel_tx, worker);
|
||||
|
||||
Sse::new(sse_stream)
|
||||
.keep_alive(KeepAlive::new().interval(Duration::from_secs(15)))
|
||||
.into_response()
|
||||
}
|
||||
|
||||
#[doc(hidden)]
|
||||
pub fn analyze_stream_body(
|
||||
metrics: watch::Receiver<Option<String>>,
|
||||
terminal: watch::Receiver<Option<AnalyzeStreamMessage>>,
|
||||
notify: Arc<Notify>,
|
||||
cancel: watch::Sender<bool>,
|
||||
worker: common_runtime::JoinHandle<()>,
|
||||
) -> impl futures::Stream<Item = std::result::Result<Event, std::convert::Infallible>> {
|
||||
futures::stream::unfold(
|
||||
AnalyzeStreamBodyState {
|
||||
latest_metrics: metrics,
|
||||
terminal,
|
||||
notify,
|
||||
_worker: AnalyzeStreamWorkerGuard {
|
||||
cancel,
|
||||
handle: worker,
|
||||
},
|
||||
done: false,
|
||||
},
|
||||
|mut state| async move {
|
||||
if state.done {
|
||||
return None;
|
||||
}
|
||||
let tick = tokio::time::sleep(Duration::from_millis(state.current_interval_ms));
|
||||
tokio::pin!(tick);
|
||||
loop {
|
||||
tokio::select! {
|
||||
item = state.stream.next() => {
|
||||
match item {
|
||||
Some(Ok(batch)) => state.batches.push(batch),
|
||||
Some(Err(err)) => {
|
||||
let status = err.status_code();
|
||||
let event_name = if status == StatusCode::Cancelled { "canceled" } else { "error" };
|
||||
let (payload, _) = make_analyze_payload(AnalyzePayloadArgs {
|
||||
seq: state.seq,
|
||||
state: event_name,
|
||||
partial: false,
|
||||
start: state.start,
|
||||
plan: state.plan.as_ref(),
|
||||
output: None,
|
||||
reason: Some(err.output_msg()),
|
||||
code: Some(status as u32),
|
||||
});
|
||||
state.seq += 1;
|
||||
state.done = true;
|
||||
return Some((Ok::<Event, std::convert::Infallible>(Event::default().event(event_name).data(payload)), state));
|
||||
}
|
||||
None => {
|
||||
let batches = std::mem::take(&mut state.batches);
|
||||
let output = HttpRecordsOutput::try_new(state.schema.clone(), batches)
|
||||
.map(GreptimeQueryOutput::Records);
|
||||
let (event_name, payload) = make_final_analyze_event(
|
||||
output.map_err(|err| (err.output_msg(), err.status_code() as u32)),
|
||||
state.seq,
|
||||
state.start,
|
||||
state.plan.as_ref(),
|
||||
);
|
||||
state.seq += 1;
|
||||
state.done = true;
|
||||
return Some((Ok::<Event, std::convert::Infallible>(Event::default().event(event_name).data(payload)), state));
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = &mut tick => {
|
||||
if state.plan.is_some() {
|
||||
let (payload, payload_bytes) = make_analyze_payload(AnalyzePayloadArgs {
|
||||
seq: state.seq,
|
||||
state: "metrics",
|
||||
partial: true,
|
||||
start: state.start,
|
||||
plan: state.plan.as_ref(),
|
||||
output: None,
|
||||
reason: None,
|
||||
code: None,
|
||||
});
|
||||
state.current_interval_ms = adaptive_interval_ms(payload_bytes, state.requested_interval_ms);
|
||||
state.seq += 1;
|
||||
return Some((Ok::<Event, std::convert::Infallible>(Event::default().event("metrics").data(payload)), state));
|
||||
}
|
||||
tick.as_mut().reset(tokio::time::Instant::now() + Duration::from_millis(state.current_interval_ms));
|
||||
let notify = Arc::clone(&state.notify);
|
||||
let notified = notify.notified();
|
||||
tokio::pin!(notified);
|
||||
// Register before checking the slots to avoid a lost wakeup.
|
||||
notified.as_mut().enable();
|
||||
let latest = {
|
||||
let latest = state.latest_metrics.borrow_and_update();
|
||||
latest.has_changed().then(|| latest.clone())
|
||||
};
|
||||
if let Some(Some(payload)) = latest {
|
||||
return Some((Ok(Event::default().event("metrics").data(payload)), state));
|
||||
}
|
||||
// Inspect the terminal slot directly because a closed watch channel
|
||||
// can otherwise hide a value published while a worker was unwinding.
|
||||
if state.terminal.borrow().is_some() {
|
||||
let terminal = { state.terminal.borrow_and_update().clone() };
|
||||
if let Some(AnalyzeStreamMessage {
|
||||
event_name,
|
||||
payload,
|
||||
}) = terminal
|
||||
{
|
||||
state.done = true;
|
||||
return Some((Ok(Event::default().event(event_name).data(payload)), state));
|
||||
}
|
||||
}
|
||||
notified.await;
|
||||
}
|
||||
},
|
||||
);
|
||||
)
|
||||
}
|
||||
|
||||
Sse::new(sse_stream)
|
||||
.keep_alive(KeepAlive::new().interval(Duration::from_secs(15)))
|
||||
.into_response()
|
||||
#[doc(hidden)]
|
||||
pub fn send_analyze_terminal(
|
||||
terminal_tx: &watch::Sender<Option<AnalyzeStreamMessage>>,
|
||||
notify: &Notify,
|
||||
event_name: &'static str,
|
||||
payload: String,
|
||||
) {
|
||||
if terminal_tx.send_if_modified(|terminal| {
|
||||
if terminal.is_none() {
|
||||
*terminal = Some(AnalyzeStreamMessage {
|
||||
event_name,
|
||||
payload,
|
||||
});
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}) {
|
||||
notify.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
fn adaptive_interval_ms(payload_bytes: usize, requested_ms: u64) -> u64 {
|
||||
@@ -438,7 +551,9 @@ fn make_analyze_payload(args: AnalyzePayloadArgs<'_>) -> (String, usize) {
|
||||
reason,
|
||||
code,
|
||||
} = args;
|
||||
let metrics = plan.and_then(|plan| query::analyze_plan_metrics_to_json_value(plan, true).ok());
|
||||
// Periodic snapshots are compact; terminal snapshots retain verbose plan details.
|
||||
let metrics =
|
||||
plan.and_then(|plan| query::analyze_plan_metrics_to_json_value(plan, !partial).ok());
|
||||
let payload = AnalyzeStreamPayload {
|
||||
seq,
|
||||
state,
|
||||
|
||||
@@ -30,15 +30,15 @@ pub type ServerSqlQueryHandlerRef = Arc<dyn SqlQueryHandler + Send + Sync>;
|
||||
pub trait SqlQueryHandler {
|
||||
async fn do_query(&self, query: &str, query_ctx: QueryContextRef) -> Vec<Result<Output>>;
|
||||
|
||||
/// Executes the experimental HTTP analyze-stream query path.
|
||||
/// Executes an HTTP analyze-stream query.
|
||||
///
|
||||
/// Implementations must validate that `query` is exactly one explicit
|
||||
/// `EXPLAIN ANALYZE VERBOSE` statement and must return a streaming output.
|
||||
/// `OutputMeta.plan` is used by the HTTP layer to emit metrics snapshots;
|
||||
/// when it is absent, partial metrics may not be available. The returned
|
||||
/// stream should support cancel-on-drop semantics (as the production
|
||||
/// frontend implementation does) so client disconnect can best-effort cancel
|
||||
/// the underlying query.
|
||||
/// Implementations must validate that `query` is exactly one supported
|
||||
/// SQL `EXPLAIN ANALYZE VERBOSE` or TQL `TQL ANALYZE VERBOSE` statement and
|
||||
/// must return a streaming output. `OutputMeta.plan` is used by the HTTP
|
||||
/// layer to emit metrics snapshots; when it is absent, partial metrics may
|
||||
/// not be available. The returned stream should support cancel-on-drop
|
||||
/// semantics (as the production frontend implementation does) so client
|
||||
/// disconnect can best-effort cancel the underlying query.
|
||||
async fn do_analyze_stream_query(
|
||||
&self,
|
||||
query: &str,
|
||||
|
||||
@@ -12,9 +12,12 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::any::Any;
|
||||
use std::collections::HashMap;
|
||||
use std::fmt;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -27,9 +30,15 @@ use bytes::Bytes;
|
||||
use common_query::{Output, OutputData};
|
||||
use common_recordbatch::adapter::RecordBatchMetrics;
|
||||
use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream};
|
||||
use datafusion::execution::TaskContext;
|
||||
use datafusion::physical_expr::{EquivalenceProperties, Partitioning};
|
||||
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
|
||||
use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
|
||||
use datafusion_common::Result as DfResult;
|
||||
use datafusion_expr::LogicalPlan;
|
||||
use datatypes::schema::SchemaRef;
|
||||
use futures::Stream;
|
||||
use futures_util::StreamExt;
|
||||
use headers::HeaderValue;
|
||||
use mime_guess::mime;
|
||||
use query::parser::PromQuery;
|
||||
@@ -47,6 +56,7 @@ use servers::query_handler::sql::{ServerSqlQueryHandlerRef, SqlQueryHandler};
|
||||
use session::context::{QueryContext, QueryContextRef};
|
||||
use sql::statements::statement::Statement;
|
||||
use table::test_util::MemTable;
|
||||
use tokio::sync::{Notify, watch};
|
||||
|
||||
use crate::create_testing_sql_query_handler;
|
||||
|
||||
@@ -98,9 +108,79 @@ impl RecordBatchStream for DelayedRecordBatchStream {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct PanickingMetricsExec {
|
||||
properties: Arc<PlanProperties>,
|
||||
metrics_calls: Arc<AtomicUsize>,
|
||||
panic_after: usize,
|
||||
}
|
||||
|
||||
impl PanickingMetricsExec {
|
||||
fn with_panic_after(schema: SchemaRef, panic_after: usize) -> Self {
|
||||
Self {
|
||||
properties: Arc::new(PlanProperties::new(
|
||||
EquivalenceProperties::new(schema.arrow_schema().clone()),
|
||||
Partitioning::UnknownPartitioning(1),
|
||||
EmissionType::Incremental,
|
||||
Boundedness::Bounded,
|
||||
)),
|
||||
metrics_calls: Arc::new(AtomicUsize::new(0)),
|
||||
panic_after,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl DisplayAs for PanickingMetricsExec {
|
||||
fn fmt_as(&self, _t: DisplayFormatType, _f: &mut fmt::Formatter) -> fmt::Result {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl ExecutionPlan for PanickingMetricsExec {
|
||||
fn name(&self) -> &str {
|
||||
"PanickingMetricsExec"
|
||||
}
|
||||
|
||||
fn as_any(&self) -> &dyn Any {
|
||||
self
|
||||
}
|
||||
|
||||
fn properties(&self) -> &Arc<PlanProperties> {
|
||||
&self.properties
|
||||
}
|
||||
|
||||
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
|
||||
vec![]
|
||||
}
|
||||
|
||||
fn with_new_children(
|
||||
self: Arc<Self>,
|
||||
_children: Vec<Arc<dyn ExecutionPlan>>,
|
||||
) -> DfResult<Arc<dyn ExecutionPlan>> {
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
fn execute(
|
||||
&self,
|
||||
_partition: usize,
|
||||
_context: Arc<TaskContext>,
|
||||
) -> DfResult<common_recordbatch::DfSendableRecordBatchStream> {
|
||||
unimplemented!("test plan is never executed")
|
||||
}
|
||||
|
||||
fn metrics(&self) -> Option<datafusion::physical_plan::metrics::MetricsSet> {
|
||||
if self.metrics_calls.fetch_add(1, Ordering::Relaxed) >= self.panic_after {
|
||||
panic!("metrics collection panicked")
|
||||
}
|
||||
Some(datafusion::physical_plan::metrics::MetricsSet::new())
|
||||
}
|
||||
}
|
||||
|
||||
struct SlowAnalyzeStreamHandler {
|
||||
inner: ServerSqlQueryHandlerRef,
|
||||
delay: Duration,
|
||||
panic_metrics: bool,
|
||||
panic_after: usize,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -115,7 +195,17 @@ impl SqlQueryHandler for SlowAnalyzeStreamHandler {
|
||||
query_ctx: QueryContextRef,
|
||||
) -> Result<Output> {
|
||||
let output = self.inner.do_analyze_stream_query(query, query_ctx).await?;
|
||||
let Output { data, meta } = output;
|
||||
let Output { data, mut meta } = output;
|
||||
if self.panic_metrics {
|
||||
let schema = match &data {
|
||||
OutputData::Stream(stream) => stream.schema(),
|
||||
_ => unreachable!(),
|
||||
};
|
||||
meta.plan = Some(Arc::new(PanickingMetricsExec::with_panic_after(
|
||||
schema,
|
||||
self.panic_after,
|
||||
)));
|
||||
}
|
||||
let data = match data {
|
||||
OutputData::Stream(stream) => {
|
||||
OutputData::Stream(Box::pin(DelayedRecordBatchStream::new(stream, self.delay)))
|
||||
@@ -160,10 +250,7 @@ async fn test_sql_not_provided() {
|
||||
let sql_handler = create_testing_sql_query_handler(MemTable::default_numbers_table());
|
||||
let ctx = QueryContext::with_db_name(None);
|
||||
ctx.set_current_user(auth::userinfo_by_name(None));
|
||||
let api_state = ApiState {
|
||||
sql_handler,
|
||||
experimental_enable_explain_analyze_stream: false,
|
||||
};
|
||||
let api_state = ApiState { sql_handler };
|
||||
|
||||
for format in ["greptimedb_v1", "influxdb_v1", "csv", "table"] {
|
||||
let query = http_handler::SqlQuery {
|
||||
@@ -194,10 +281,7 @@ async fn test_sql_output_rows() {
|
||||
|
||||
let ctx = QueryContext::with_db_name(None);
|
||||
ctx.set_current_user(auth::userinfo_by_name(None));
|
||||
let api_state = ApiState {
|
||||
sql_handler,
|
||||
experimental_enable_explain_analyze_stream: false,
|
||||
};
|
||||
let api_state = ApiState { sql_handler };
|
||||
|
||||
let query_sql = "select sum(uint32s) from numbers limit 20";
|
||||
for format in ["greptimedb_v1", "influxdb_v1", "csv", "table"] {
|
||||
@@ -302,10 +386,7 @@ async fn test_dashboard_sql_limit() {
|
||||
let sql_handler = create_testing_sql_query_handler(MemTable::specified_numbers_table(2000));
|
||||
let ctx = QueryContext::with_db_name(None);
|
||||
ctx.set_current_user(auth::userinfo_by_name(None));
|
||||
let api_state = ApiState {
|
||||
sql_handler,
|
||||
experimental_enable_explain_analyze_stream: false,
|
||||
};
|
||||
let api_state = ApiState { sql_handler };
|
||||
for format in ["greptimedb_v1", "csv", "table"] {
|
||||
let query = create_query(format, "select * from numbers", Some(1000));
|
||||
let sql_response = http_handler::sql(
|
||||
@@ -348,10 +429,7 @@ async fn test_sql_form() {
|
||||
|
||||
let ctx = QueryContext::with_db_name(None);
|
||||
ctx.set_current_user(auth::userinfo_by_name(None));
|
||||
let api_state = ApiState {
|
||||
sql_handler,
|
||||
experimental_enable_explain_analyze_stream: false,
|
||||
};
|
||||
let api_state = ApiState { sql_handler };
|
||||
|
||||
for format in ["greptimedb_v1", "influxdb_v1", "csv", "table", "null"] {
|
||||
let form = create_form(format);
|
||||
@@ -526,6 +604,9 @@ async fn test_analyze_stream_route_rejects_invalid_sql() {
|
||||
"EXPLAIN ANALYZE SELECT 1",
|
||||
"EXPLAIN ANALYZE VERBOSE FORMAT TEXT SELECT 1",
|
||||
"EXPLAIN ANALYZE VERBOSE SELECT 1; SELECT 2",
|
||||
"TQL ANALYZE (0, 10, '5s') physical_metric",
|
||||
"TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric",
|
||||
"TQL ANALYZE VERBOSE FORMAT TEXT (0, 10, '5s') physical_metric",
|
||||
] {
|
||||
let response = client
|
||||
.post("/v1/sql/analyze/stream")
|
||||
@@ -568,6 +649,173 @@ async fn test_analyze_stream_route_accepts_explicit_format_json() {
|
||||
assert!(sse_event_payload(&body, "final").is_some(), "{body}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_analyze_stream_body_orders_latest_metrics_before_terminal() {
|
||||
let (metrics_tx, metrics_rx) = watch::channel::<Option<String>>(None);
|
||||
let (terminal_tx, terminal_rx) =
|
||||
watch::channel::<Option<http_handler::AnalyzeStreamMessage>>(None);
|
||||
let notify = Arc::new(Notify::new());
|
||||
metrics_tx.send(Some("latest".to_string())).unwrap();
|
||||
terminal_tx
|
||||
.send(Some(http_handler::AnalyzeStreamMessage {
|
||||
event_name: "final",
|
||||
payload: "terminal".to_string(),
|
||||
}))
|
||||
.unwrap();
|
||||
drop(metrics_tx);
|
||||
let (cancel_tx, mut cancel_rx) = watch::channel(false);
|
||||
let worker = tokio::spawn(async move {
|
||||
let _ = cancel_rx.changed().await;
|
||||
});
|
||||
let body =
|
||||
http_handler::analyze_stream_body(metrics_rx, terminal_rx, notify, cancel_tx, worker);
|
||||
futures::pin_mut!(body);
|
||||
let first = tokio::time::timeout(Duration::from_secs(1), body.next())
|
||||
.await
|
||||
.expect("timed out waiting for metrics")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let second = tokio::time::timeout(Duration::from_secs(1), body.next())
|
||||
.await
|
||||
.expect("timed out waiting for terminal")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(format!("{first:?}").contains("latest"));
|
||||
assert!(format!("{second:?}").contains("terminal"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_analyze_stream_body_drop_cancels_worker() {
|
||||
let (_, metrics_rx) = watch::channel::<Option<String>>(None);
|
||||
let (_, terminal_rx) = watch::channel::<Option<http_handler::AnalyzeStreamMessage>>(None);
|
||||
let notify = Arc::new(Notify::new());
|
||||
let (cancel_tx, mut canceled_rx) = watch::channel(false);
|
||||
let worker = tokio::spawn(async { futures::future::pending::<()>().await });
|
||||
let mut body = Box::pin(http_handler::analyze_stream_body(
|
||||
metrics_rx,
|
||||
terminal_rx,
|
||||
notify,
|
||||
cancel_tx,
|
||||
worker,
|
||||
));
|
||||
assert!(
|
||||
tokio::time::timeout(Duration::from_millis(1), body.next())
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
drop(body);
|
||||
tokio::time::timeout(Duration::from_secs(1), async {
|
||||
while !*canceled_rx.borrow() {
|
||||
canceled_rx.changed().await.unwrap();
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("timed out waiting for worker cancellation");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_analyze_stream_terminal_is_sent_once() {
|
||||
let (terminal_tx, mut terminal_rx) =
|
||||
watch::channel::<Option<http_handler::AnalyzeStreamMessage>>(None);
|
||||
let notify = Notify::new();
|
||||
http_handler::send_analyze_terminal(&terminal_tx, ¬ify, "error", "one".to_string());
|
||||
http_handler::send_analyze_terminal(&terminal_tx, ¬ify, "error", "two".to_string());
|
||||
terminal_rx.changed().await.unwrap();
|
||||
let terminal = terminal_rx.borrow_and_update().clone().unwrap();
|
||||
assert_eq!(terminal.payload, "one");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_analyze_stream_worker_panic_emits_one_error_terminal() {
|
||||
common_telemetry::init_default_ut_logging();
|
||||
|
||||
let inner = create_testing_sql_query_handler(MemTable::default_numbers_table());
|
||||
let sql_handler = Arc::new(SlowAnalyzeStreamHandler {
|
||||
inner,
|
||||
delay: Duration::ZERO,
|
||||
panic_metrics: true,
|
||||
panic_after: 0,
|
||||
});
|
||||
let ctx = QueryContext::with_db_name(None);
|
||||
ctx.set_current_user(auth::userinfo_by_name(None));
|
||||
let response = http_handler::sql_analyze_stream(
|
||||
State(ApiState { sql_handler }),
|
||||
Query(http_handler::SqlQuery {
|
||||
sql: Some("EXPLAIN ANALYZE VERBOSE SELECT sum(uint32s) FROM numbers".to_string()),
|
||||
snapshot_interval_ms: Some(1000),
|
||||
..Default::default()
|
||||
}),
|
||||
axum::Extension(ctx),
|
||||
Ok(Form(http_handler::SqlQuery::default())),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = tokio::time::timeout(
|
||||
Duration::from_secs(3),
|
||||
axum::body::to_bytes(response.into_body(), usize::MAX),
|
||||
)
|
||||
.await
|
||||
.expect("timed out waiting for panic terminal")
|
||||
.unwrap();
|
||||
let body = String::from_utf8(body.to_vec()).unwrap();
|
||||
assert_eq!(sse_event_payloads(&body, "error").len(), 1, "{body}");
|
||||
assert!(sse_event_payload(&body, "final").is_none(), "{body}");
|
||||
let payload: Value = serde_json::from_str(&sse_event_payload(&body, "error").unwrap()).unwrap();
|
||||
assert_eq!(payload["state"], "error");
|
||||
assert_eq!(payload["reason"], "analyze stream worker panicked");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_analyze_stream_worker_panic_after_metrics_uses_current_sequence() {
|
||||
common_telemetry::init_default_ut_logging();
|
||||
|
||||
let inner = create_testing_sql_query_handler(MemTable::default_numbers_table());
|
||||
let sql_handler = Arc::new(SlowAnalyzeStreamHandler {
|
||||
inner,
|
||||
delay: Duration::from_millis(2500),
|
||||
panic_metrics: true,
|
||||
panic_after: 1,
|
||||
});
|
||||
let ctx = QueryContext::with_db_name(None);
|
||||
ctx.set_current_user(auth::userinfo_by_name(None));
|
||||
let response = http_handler::sql_analyze_stream(
|
||||
State(ApiState { sql_handler }),
|
||||
Query(http_handler::SqlQuery {
|
||||
sql: Some("EXPLAIN ANALYZE VERBOSE SELECT sum(uint32s) FROM numbers".to_string()),
|
||||
snapshot_interval_ms: Some(1000),
|
||||
..Default::default()
|
||||
}),
|
||||
axum::Extension(ctx),
|
||||
Ok(Form(http_handler::SqlQuery::default())),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = tokio::time::timeout(
|
||||
Duration::from_secs(5),
|
||||
axum::body::to_bytes(response.into_body(), usize::MAX),
|
||||
)
|
||||
.await
|
||||
.expect("timed out waiting for panic terminal after metrics")
|
||||
.unwrap();
|
||||
let body = String::from_utf8(body.to_vec()).unwrap();
|
||||
let metrics_payload: Value =
|
||||
serde_json::from_str(&sse_event_payload(&body, "metrics").expect(&body)).unwrap();
|
||||
let error_payloads = sse_event_payloads(&body, "error");
|
||||
assert_eq!(error_payloads.len(), 1, "{body}");
|
||||
let error_payload: Value = serde_json::from_str(&error_payloads[0]).unwrap();
|
||||
assert_eq!(metrics_payload["state"], "metrics");
|
||||
assert_eq!(error_payload["state"], "error");
|
||||
assert_eq!(error_payload["reason"], "analyze stream worker panicked");
|
||||
let metrics_seq = metrics_payload["seq"].as_u64().unwrap();
|
||||
let error_seq = error_payload["seq"].as_u64().unwrap();
|
||||
assert!(
|
||||
error_seq > metrics_seq,
|
||||
"terminal sequence should follow metrics sequence: {body}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_analyze_stream_emits_metrics_before_final_when_stream_is_pending() {
|
||||
common_telemetry::init_default_ut_logging();
|
||||
@@ -576,12 +824,10 @@ async fn test_analyze_stream_emits_metrics_before_final_when_stream_is_pending()
|
||||
let sql_handler = Arc::new(SlowAnalyzeStreamHandler {
|
||||
inner,
|
||||
delay: Duration::from_millis(1500),
|
||||
panic_metrics: false,
|
||||
panic_after: 0,
|
||||
});
|
||||
let options = HttpOptions {
|
||||
experimental_enable_explain_analyze_stream: true,
|
||||
..Default::default()
|
||||
};
|
||||
let server = HttpServerBuilder::new(options)
|
||||
let server = HttpServerBuilder::new(HttpOptions::default())
|
||||
.with_sql_handler(sql_handler)
|
||||
.build();
|
||||
let app = server.build(server.make_app()).unwrap();
|
||||
@@ -617,27 +863,29 @@ async fn test_analyze_stream_emits_metrics_before_final_when_stream_is_pending()
|
||||
);
|
||||
}
|
||||
|
||||
fn sse_event_payload(body: &str, event_name: &str) -> Option<String> {
|
||||
body.split("\n\n").find_map(|event| {
|
||||
let mut found = false;
|
||||
let mut data = Vec::new();
|
||||
for line in event.lines() {
|
||||
if line.strip_prefix("event: ") == Some(event_name) {
|
||||
found = true;
|
||||
} else if let Some(value) = line.strip_prefix("data: ") {
|
||||
data.push(value);
|
||||
fn sse_event_payloads(body: &str, event_name: &str) -> Vec<String> {
|
||||
body.split("\n\n")
|
||||
.filter_map(|event| {
|
||||
let mut found = false;
|
||||
let mut data = Vec::new();
|
||||
for line in event.lines() {
|
||||
if line.strip_prefix("event: ") == Some(event_name) {
|
||||
found = true;
|
||||
} else if let Some(value) = line.strip_prefix("data: ") {
|
||||
data.push(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
found.then(|| data.join("\n"))
|
||||
})
|
||||
found.then(|| data.join("\n"))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn sse_event_payload(body: &str, event_name: &str) -> Option<String> {
|
||||
sse_event_payloads(body, event_name).into_iter().next()
|
||||
}
|
||||
|
||||
async fn analyze_stream_test_client(sql_handler: ServerSqlQueryHandlerRef) -> TestClient {
|
||||
let options = HttpOptions {
|
||||
experimental_enable_explain_analyze_stream: true,
|
||||
..Default::default()
|
||||
};
|
||||
let server = HttpServerBuilder::new(options)
|
||||
let server = HttpServerBuilder::new(HttpOptions::default())
|
||||
.with_sql_handler(sql_handler)
|
||||
.build();
|
||||
let app = server.build(server.make_app()).unwrap();
|
||||
|
||||
@@ -119,6 +119,7 @@ macro_rules! http_tests {
|
||||
|
||||
test_http_auth,
|
||||
test_sql_api,
|
||||
test_http_analyze_stream_tql,
|
||||
test_http_sql_slow_query,
|
||||
test_prometheus_promql_api,
|
||||
test_prometheus_label_replace_response,
|
||||
@@ -807,6 +808,61 @@ async fn test_sql_format_api() {
|
||||
guard.remove_all().await;
|
||||
}
|
||||
|
||||
pub async fn test_http_analyze_stream_tql(store_type: StorageType) {
|
||||
common_telemetry::init_default_ut_logging();
|
||||
let (app, mut guard) =
|
||||
setup_test_prom_app_with_frontend(store_type, "analyze_stream_tql").await;
|
||||
let client = TestClient::new(app).await;
|
||||
|
||||
let res = client
|
||||
.post("/v1/sql/analyze/stream")
|
||||
.header("Accept", "text/event-stream")
|
||||
.form(&BTreeMap::from([(
|
||||
"sql".to_string(),
|
||||
"TQL ANALYZE VERBOSE (0, 10, '5s') demo".to_string(),
|
||||
)]))
|
||||
.send()
|
||||
.await;
|
||||
assert_eq!(res.status(), StatusCode::OK);
|
||||
assert!(
|
||||
res.headers()
|
||||
.get("content-type")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.is_some_and(|value| value.starts_with("text/event-stream"))
|
||||
);
|
||||
let body = res.text().await;
|
||||
let final_event = body
|
||||
.split("\n\n")
|
||||
.find(|event| event.lines().any(|line| line == "event: final"))
|
||||
.expect(&body);
|
||||
let payload = final_event
|
||||
.lines()
|
||||
.find_map(|line| line.strip_prefix("data: "))
|
||||
.map(|data| serde_json::from_str::<Value>(data).unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(payload["state"], "final");
|
||||
assert!(payload["metrics"].as_array().is_some_and(|v| !v.is_empty()));
|
||||
assert!(
|
||||
payload["output"]["records"]["rows"]
|
||||
.as_array()
|
||||
.is_some_and(|v| !v.is_empty())
|
||||
);
|
||||
|
||||
let res = client
|
||||
.post("/v1/sql/analyze/stream")
|
||||
.header("Accept", "text/event-stream")
|
||||
.form(&BTreeMap::from([(
|
||||
"sql".to_string(),
|
||||
"TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') demo".to_string(),
|
||||
)]))
|
||||
.send()
|
||||
.await;
|
||||
assert_eq!(res.status(), StatusCode::OK);
|
||||
assert!(res.text().await.contains("event: final"));
|
||||
|
||||
guard.remove_all().await;
|
||||
}
|
||||
|
||||
pub async fn test_http_sql_slow_query(store_type: StorageType) {
|
||||
let (app, mut guard) = setup_test_http_app_with_frontend_and_slow_query_threshold(
|
||||
store_type,
|
||||
@@ -2147,7 +2203,6 @@ timeout = "0s"
|
||||
body_limit = "64MiB"
|
||||
cors_allowed_origins = []
|
||||
enable_cors = true
|
||||
experimental_enable_explain_analyze_stream = true
|
||||
enable_api_server = false
|
||||
api_server_addr = "127.0.0.1:4006"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user