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

fix: preserve automatic parallelism for internal Flow requests

Browse Source 
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-10-09 16:35:20 +08:00
parent 17af970ddd
commit fb8a64509f
1 changed files with 107 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/flow/src/batching_mode/frontend_client.rs
+107 -20
View File
@@ -486,10 +486,10 @@ impl FrontendClient {
query, batch_opts, ..
} => {
let query_parallelism = query.parallelism.to_string();
let hints = vec![
(QUERY_PARALLELISM_HINT, query_parallelism.as_str()),
(READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref()),
];
let mut hints = vec![(READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref())];
if query.parallelism > 0 {
hints.push((QUERY_PARALLELISM_HINT, query_parallelism.as_str()));
}
let db = self.get_random_active_frontend(catalog, schema).await?;
*peer_desc = Some(PeerDesc::Dist {
peer: db.peer.clone(),
@@ -510,10 +510,13 @@ impl FrontendClient {
query,
} => {
*peer_desc = Some(PeerDesc::Standalone);
let mut extensions_map = HashMap::from([(
QUERY_PARALLELISM_HINT.to_string(),
query.parallelism.to_string(),
)]);
let mut extensions_map = HashMap::new();
if query.parallelism > 0 {
extensions_map.insert(
QUERY_PARALLELISM_HINT.to_string(),
query.parallelism.to_string(),
);
}
for (key, value) in extensions {
extensions_map.insert((*key).to_string(), (*value).to_string());
}
@@ -587,14 +590,18 @@ impl FrontendClient {
database_client,
query,
} => {
let mut extensions = HashMap::new();
if query.parallelism > 0 {
extensions.insert(
QUERY_PARALLELISM_HINT.to_string(),
query.parallelism.to_string(),
);
}
let ctx = QueryContextBuilder::default()
.current_catalog(catalog.to_string())
.current_schema(schema.to_string())
.channel(Channel::Internal)
.extensions(HashMap::from([(
QUERY_PARALLELISM_HINT.to_string(),
query.parallelism.to_string(),
)]))
.extensions(extensions)
.build();
let ctx = Arc::new(ctx);
let database_client = {
@@ -650,14 +657,16 @@ impl FrontendClient {
peer: db.peer.clone(),
});
let query_parallelism = query.parallelism.to_string();
let mut hints = vec![(READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref())];
if query.parallelism > 0 {
hints.push((QUERY_PARALLELISM_HINT, query_parallelism.as_str()));
}
db.database
.handle_with_retry(
req.clone(),
batch_opts.experimental_grpc_max_retries,
&[
(QUERY_PARALLELISM_HINT, &query.parallelism.to_string()),
(READ_PREFERENCE_HINT, batch_opts.read_preference.as_ref()),
],
&hints,
)
.await
.with_context(|_| InvalidRequestSnafu {
@@ -668,13 +677,17 @@ impl FrontendClient {
database_client,
query,
} => {
let mut extensions = HashMap::new();
if query.parallelism > 0 {
extensions.insert(
QUERY_PARALLELISM_HINT.to_string(),
query.parallelism.to_string(),
);
}
let ctx = QueryContextBuilder::default()
.current_catalog(catalog.to_string())
.current_schema(schema.to_string())
.extensions(HashMap::from([(
QUERY_PARALLELISM_HINT.to_string(),
query.parallelism.to_string(),
)]))
.extensions(extensions)
.build();
let ctx = Arc::new(ctx);
{
@@ -862,6 +875,11 @@ mod tests {
#[derive(Debug)]
struct SnapshotBindingHandler;
#[derive(Debug)]
struct QueryParallelismHandler {
context: Arc<Mutex<Option<QueryContextRef>>>,
}
#[derive(Debug)]
struct RejectUnauthenticatedFlight;
@@ -949,6 +967,19 @@ mod tests {
}
}
#[async_trait::async_trait]
impl GrpcQueryHandlerWithBoxedError for QueryParallelismHandler {
async fn do_query(
&self,
_query: Request,
ctx: QueryContextRef,
) -> std::result::Result<Output, BoxedError> {
ctx.query_option_snapshot().unwrap();
*self.context.lock().unwrap() = Some(ctx);
Ok(Output::new_with_affected_rows(1))
}
}
#[async_trait::async_trait]
impl GrpcQueryHandlerWithBoxedError for SnapshotBindingHandler {
async fn do_query(
@@ -1121,6 +1152,62 @@ mod tests {
assert!(peer_desc.is_none());
}
#[tokio::test]
async fn test_standalone_queries_only_forward_positive_parallelism() {
for (parallelism, expected) in [(0, None), (17, Some("17"))] {
let captured_context = Arc::new(Mutex::new(None));
let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> =
Arc::new(QueryParallelismHandler {
context: captured_context.clone(),
});
let client = FrontendClient::from_grpc_handler(
Arc::downgrade(&handler),
QueryOptions {
parallelism,
..Default::default()
},
);
let query = Request::Query(QueryRequest {
query: Some(Query::Sql("select 1".to_string())),
});
let mut peer_desc = None;
client
.handle(query, "greptime", "public", &mut peer_desc)
.await
.unwrap();
{
let context = captured_context.lock().unwrap().take().unwrap();
assert_eq!(context.extension(QUERY_PARALLELISM_HINT), expected);
context.query_option_snapshot().unwrap();
}
let flow_extension = (query::options::FLOW_RETURN_REGION_SEQ, "true");
let snapshot_seqs = HashMap::from([(1, 10)]);
let query = QueryRequest {
query: Some(Query::Sql("select 1".to_string())),
};
client
.query_with_terminal_metrics(
"greptime",
"public",
query,
&[flow_extension],
&snapshot_seqs,
&mut peer_desc,
)
.await
.unwrap();
{
let context = captured_context.lock().unwrap().take().unwrap();
assert_eq!(context.extension(QUERY_PARALLELISM_HINT), expected);
assert_eq!(context.extension(flow_extension.0), Some(flow_extension.1));
assert_eq!(context.get_snapshot(1), Some(10));
context.query_option_snapshot().unwrap();
}
}
}
#[tokio::test]
async fn test_query_with_terminal_metrics_tracks_watermark_in_standalone_mode() {
let handler: Arc<dyn GrpcQueryHandlerWithBoxedError> = Arc::new(MetricsHandler);