gsadmin/rustfs
1
0
Fork 0
mirror of https://github.com/rustfs/rustfs.git synced 2026-09-02 02:08:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(s3select): report uncompressed input byte metrics (#6865)

Browse Source
This commit is contained in:
GatewayJ
2026-08-30 12:07:20 +08:00
committed by GitHub
parent b428875bed
commit 96239fc034
10 changed files with 1151 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/s3select-query/src/dispatcher/manager.rs
+75 -25
View File
@@ -42,7 +42,7 @@ use rustfs_s3select_api::{
query::{
Query,
ast::ExtStatement,
dispatcher::QueryDispatcher,
dispatcher::{DispatchedQuery, QueryDispatcher},
execution::{Output, QueryStateMachine},
function::FuncMetaManagerRef,
logical_planner::{LogicalPlanner, Plan},
@@ -120,7 +120,7 @@ impl Drop for QueryPhaseGuard<'_> {
#[async_trait]
impl QueryDispatcher for SimpleQueryDispatcher {
async fn execute_query(&self, query: &Query) -> QueryResult<Output> {
self.execute_query_inner(query, None).await
self.execute_query_inner(query, None).await.map(|(_, output)| output)
}
fn try_reserve_query(&self) -> QueryResult<QueryAdmission> {
@@ -133,6 +133,16 @@ impl QueryDispatcher for SimpleQueryDispatcher {
}
async fn execute_query_admitted(&self, query: &Query, admission: QueryAdmission) -> QueryResult<Output> {
self.execute_query_inner(query, Some(admission))
.await
.map(|(_, output)| output)
}
async fn dispatch_query(&self, query: &Query) -> QueryResult<DispatchedQuery> {
self.execute_query_inner(query, None).await
}
async fn dispatch_query_admitted(&self, query: &Query, admission: QueryAdmission) -> QueryResult<DispatchedQuery> {
self.execute_query_inner(query, Some(admission)).await
}
@@ -171,11 +181,12 @@ impl QueryDispatcher for SimpleQueryDispatcher {
};
let logical_plan = self
.statement_to_logical_plan(stmt, &logical_planner, query_state_machine)
.statement_to_logical_plan(stmt, &logical_planner, Arc::clone(&query_state_machine))
.await?;
Ok(logical_plan)
})
.await?;
query_state_machine.query.input_metrics().reset();
if !query_tracker.mark_planned(&self.query_execution_owner) {
drop(logical_plan);
return Err(self.query_tracker_error(&query_tracker));
@@ -212,19 +223,21 @@ impl QueryDispatcher for SimpleQueryDispatcher {
}
async fn build_query_state_machine(&self, query: Query) -> QueryResult<Arc<QueryStateMachine>> {
self.build_query_state_machine_inner(query, None).await
self.build_query_state_machine_inner(query.for_execution(), None).await
}
}
impl SimpleQueryDispatcher {
async fn execute_query_inner(&self, query: &Query, admission: Option<QueryAdmission>) -> QueryResult<Output> {
let query_state_machine = self.build_query_state_machine_inner(query.clone(), admission).await?;
async fn execute_query_inner(&self, query: &Query, admission: Option<QueryAdmission>) -> QueryResult<DispatchedQuery> {
let query_state_machine = self.build_query_state_machine_inner(query.for_execution(), admission).await?;
let execution_query = query_state_machine.query.clone();
let logical_plan = self.build_logical_plan(Arc::clone(&query_state_machine)).await?;
let Some(logical_plan) = logical_plan else {
return Ok(Output::Nil(()));
return Ok((execution_query, Output::Nil(())));
};
self.execute_logical_plan(logical_plan, query_state_machine).await
let output = self.execute_logical_plan(logical_plan, query_state_machine).await?;
Ok((execution_query, output))
}
async fn build_query_state_machine_inner(
@@ -256,29 +269,17 @@ impl SimpleQueryDispatcher {
self.query_timeout.as_secs(),
);
let phase_guard = QueryPhaseGuard::new(&query_tracker, &self.query_execution_owner);
let session = if let Some(snapshot) = query.snapshot().cloned() {
self.run_with_query_deadline(
let session = self
.run_with_query_deadline(
&query_tracker,
self.session_factory
.create_session_ctx_with_snapshot_and_tracker_and_memory_limit(
query.context(),
snapshot,
.create_session_ctx_for_query_with_tracker_and_memory_limit(
&query,
query_tracker.clone(),
self.memory_limit_bytes,
),
)
.await?
} else {
self.run_with_query_deadline(
&query_tracker,
self.session_factory.create_session_ctx_with_tracker_and_memory_limit(
query.context(),
query_tracker.clone(),
self.memory_limit_bytes,
),
)
.await?
};
.await?;
if !query_tracker.mark_admitted(&self.query_execution_owner) {
drop(session);
return Err(self.query_tracker_error(&query_tracker));
@@ -1705,6 +1706,55 @@ mod tests {
assert_eq!(admission.available_permits(), 1);
}
#[tokio::test]
async fn reused_query_gets_execution_local_input_metrics() {
let admission = Arc::new(Semaphore::new(2));
let (dispatcher, input) = test_dispatcher(Arc::clone(&admission), Duration::from_secs(300));
let query = Query::new(QueryContext { input }, "SELECT * FROM S3Object".to_string());
let first = dispatcher
.build_query_state_machine(query.clone())
.await
.expect("first execution state");
let second = dispatcher
.build_query_state_machine(query)
.await
.expect("second execution state");
assert!(!Arc::ptr_eq(first.query.input_metrics(), second.query.input_metrics()));
drop((first, second));
assert_eq!(admission.available_permits(), 2);
}
#[tokio::test]
async fn dispatching_a_reused_query_returns_execution_local_input_metrics() {
let env = snapshot_test_env().await;
let mut input = test_input();
input.bucket = "s3select-reused-query-metrics".to_string();
input.key = "input.csv".to_string();
let input = Arc::new(input);
env.make_bucket(&input.bucket, false).await;
env.put_object_bytes(&input.bucket, &input.key, b"name\nAlice\n".to_vec())
.await;
let snapshot = env.prepare_select_object_snapshot(&input.bucket, &input.key).await;
let dispatcher = production_dispatcher(Arc::clone(&input));
let query = Query::new_with_snapshot(
QueryContext {
input: Arc::clone(&input),
},
input.request.expression.clone(),
snapshot,
);
let (first_query, first_output) = dispatcher.dispatch_query(&query).await.expect("first dispatch should start");
let (second_query, second_output) = dispatcher.dispatch_query(&query).await.expect("second dispatch should start");
assert!(!Arc::ptr_eq(first_query.input_metrics(), second_query.input_metrics()));
assert!(!Arc::ptr_eq(query.input_metrics(), first_query.input_metrics()));
assert!(!Arc::ptr_eq(query.input_metrics(), second_query.input_metrics()));
drop((first_output, second_output));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_planning_claim_has_single_winner() {
let admission = Arc::new(Semaphore::new(1));
crates/s3select-query/src/instance.rs
+93 -6
View File
@@ -69,15 +69,15 @@ where
}
async fn execute(&self, query: &Query) -> QueryResult<QueryHandle> {
let result = self.query_dispatcher.execute_query(query).await?;
let (query, result) = self.query_dispatcher.dispatch_query(query).await?;
Ok(QueryHandle::new(query.clone(), result))
Ok(QueryHandle::new(query, result))
}
async fn execute_admitted(&self, query: &Query, admission: QueryAdmission) -> QueryResult<QueryHandle> {
let result = self.query_dispatcher.execute_query_admitted(query, admission).await?;
let (query, result) = self.query_dispatcher.dispatch_query_admitted(query, admission).await?;
Ok(QueryHandle::new(query.clone(), result))
Ok(QueryHandle::new(query, result))
}
async fn build_query_state_machine(&self, query: Query) -> QueryResult<QueryStateMachineRef> {
@@ -247,8 +247,19 @@ pub async fn make_rustfsms_with_components(
mod tests {
use std::sync::Arc;
use async_trait::async_trait;
use datafusion::{arrow::util::pretty, assert_batches_eq};
use rustfs_s3select_api::query::{Context, Query};
use parking_lot::Mutex;
use rustfs_s3select_api::{
QueryResult, SelectInputMetrics,
query::{
Context, Query,
dispatcher::QueryDispatcher,
execution::{Output, QueryStateMachine},
logical_planner::Plan,
},
server::dbms::DatabaseManagerSystem,
};
use s3s::dto::{
CSVInput, CSVOutput, ExpressionType, FieldDelimiter, FileHeaderInfo, InputSerialization, OutputSerialization,
RecordDelimiter, SelectObjectContentInput, SelectObjectContentRequest,
@@ -257,10 +268,66 @@ mod tests {
use crate::get_global_db;
use super::{
DEFAULT_MAX_CONCURRENT_QUERIES, DEFAULT_MEMORY_LIMIT_BYTES, DEFAULT_QUERY_TIMEOUT_SECS, MAX_QUERY_TIMEOUT_SECS,
DEFAULT_MAX_CONCURRENT_QUERIES, DEFAULT_MEMORY_LIMIT_BYTES, DEFAULT_QUERY_TIMEOUT_SECS, MAX_QUERY_TIMEOUT_SECS, RustFSms,
S3SelectRuntimeConfig, bounded_u64_from_env_value, bounded_usize_from_env_value, target_partitions_from_env_value,
};
#[derive(Default)]
struct FreshMetricsDispatcher {
executed_metrics: Mutex<Vec<Arc<SelectInputMetrics>>>,
}
#[async_trait]
impl QueryDispatcher for FreshMetricsDispatcher {
async fn execute_query(&self, query: &Query) -> QueryResult<Output> {
self.executed_metrics.lock().push(Arc::clone(query.input_metrics()));
Ok(Output::Nil(()))
}
async fn build_logical_plan(&self, _query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Option<Plan>> {
unreachable!("fresh metrics test does not plan queries")
}
async fn execute_logical_plan(
&self,
_logical_plan: Plan,
_query_state_machine: Arc<QueryStateMachine>,
) -> QueryResult<Output> {
unreachable!("fresh metrics test does not execute plans")
}
async fn build_query_state_machine(&self, _query: Query) -> QueryResult<Arc<QueryStateMachine>> {
unreachable!("fresh metrics test does not build state machines")
}
}
fn metrics_test_query() -> Query {
let expression = "SELECT * FROM S3Object";
let input = SelectObjectContentInput {
bucket: "bucket".to_string(),
expected_bucket_owner: None,
key: "input.csv".to_string(),
sse_customer_algorithm: None,
sse_customer_key: None,
sse_customer_key_md5: None,
request: SelectObjectContentRequest {
expression: expression.to_string(),
expression_type: ExpressionType::from_static(ExpressionType::SQL),
input_serialization: InputSerialization {
csv: Some(CSVInput::default()),
..Default::default()
},
output_serialization: OutputSerialization {
csv: Some(CSVOutput::default()),
..Default::default()
},
request_progress: None,
scan_range: None,
},
};
Query::new(Context { input: Arc::new(input) }, expression.to_string())
}
#[test]
fn parses_target_partitions_from_env_value() {
assert_eq!(target_partitions_from_env_value(Some("4")), 4);
@@ -291,6 +358,26 @@ mod tests {
assert_eq!(bounded_u64_from_env_value(None, 300, MAX_QUERY_TIMEOUT_SECS), 300);
}
#[tokio::test]
async fn repeated_execute_returns_the_fresh_dispatched_query_metrics() {
let dispatcher = Arc::new(FreshMetricsDispatcher::default());
let db = RustFSms {
query_dispatcher: Arc::clone(&dispatcher),
};
let query = metrics_test_query();
let first = db.execute(&query).await.expect("first execution should succeed");
let second = db.execute(&query).await.expect("second execution should succeed");
let executed_metrics = dispatcher.executed_metrics.lock();
assert_eq!(executed_metrics.len(), 2);
assert!(Arc::ptr_eq(first.query().input_metrics(), &executed_metrics[0]));
assert!(Arc::ptr_eq(second.query().input_metrics(), &executed_metrics[1]));
assert!(!Arc::ptr_eq(first.query().input_metrics(), second.query().input_metrics()));
assert!(!Arc::ptr_eq(first.query().input_metrics(), query.input_metrics()));
assert!(!Arc::ptr_eq(second.query().input_metrics(), query.input_metrics()));
}
#[tokio::test]
#[ignore = "requires a live RustFS store with a pre-seeded test object (bucket 'dandan')"]
async fn test_simple_sql() {