From a1d2dbb0504c1acd3a15f1b7ef8214b7fac217d1 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:54:07 +0800 Subject: [PATCH] fix(client): isolate query and control transports Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/client/src/client.rs | 89 +++++++++++++++--- src/client/src/client_manager.rs | 72 +++++++++++++-- src/client/src/database.rs | 2 +- src/query/src/datafusion.rs | 126 +++++--------------------- tests-integration/src/grpc/flight.rs | 131 ++++++++++++++++++++++++++- 5 files changed, 300 insertions(+), 120 deletions(-) diff --git a/src/client/src/client.rs b/src/client/src/client.rs index 39cb5c30aa..96c7047bb9 100644 --- a/src/client/src/client.rs +++ b/src/client/src/client.rs @@ -51,18 +51,34 @@ pub struct Client { inner: Arc, } -#[derive(Debug, Default)] +#[derive(Debug)] struct Inner { - channel_manager: ChannelManager, + query_channel_manager: ChannelManager, + control_channel_manager: ChannelManager, peers: Arc>>, load_balance: Loadbalancer, } +impl Default for Inner { + fn default() -> Self { + Self::with_manager(ChannelManager::new()) + } +} + impl Inner { fn with_manager(channel_manager: ChannelManager) -> Self { + Self::with_managers(channel_manager.clone(), channel_manager) + } + + fn with_managers( + query_channel_manager: ChannelManager, + control_channel_manager: ChannelManager, + ) -> Self { Self { - channel_manager, - ..Default::default() + query_channel_manager, + control_channel_manager, + peers: Arc::new(RwLock::new(Vec::new())), + load_balance: Loadbalancer::default(), } } @@ -107,7 +123,19 @@ impl Client { U: AsRef, A: AsRef<[U]>, { - let inner = Inner::with_manager(channel_manager); + Self::with_managers_and_urls(channel_manager.clone(), channel_manager, urls) + } + + pub(crate) fn with_managers_and_urls( + query_channel_manager: ChannelManager, + control_channel_manager: ChannelManager, + urls: A, + ) -> Self + where + U: AsRef, + A: AsRef<[U]>, + { + let inner = Inner::with_managers(query_channel_manager, control_channel_manager); let urls: Vec = urls .as_ref() .iter() @@ -143,7 +171,7 @@ impl Client { let channel = self .inner - .channel_manager + .control_channel_manager .get(&addr) .context(error::CreateChannelSnafu { addr: &addr })?; Ok((addr, channel)) @@ -151,7 +179,7 @@ impl Client { pub fn max_grpc_recv_message_size(&self) -> usize { self.inner - .channel_manager + .control_channel_manager .config() .max_recv_message_size .as_bytes() as usize @@ -159,22 +187,61 @@ impl Client { pub fn max_grpc_send_message_size(&self) -> usize { self.inner - .channel_manager + .control_channel_manager .config() .max_send_message_size .as_bytes() as usize } + /// Creates a Flight client on the query lane for DoGet/distributed reads. pub fn make_flight_client( &self, send_compression: bool, accept_compression: bool, ) -> Result { - let (addr, channel) = self.find_channel()?; + self.make_flight_client_with_manager( + &self.inner.query_channel_manager, + send_compression, + accept_compression, + ) + } + + /// Creates a Flight client on the control lane for DoPut/mutations. + pub(crate) fn make_control_flight_client( + &self, + send_compression: bool, + accept_compression: bool, + ) -> Result { + self.make_flight_client_with_manager( + &self.inner.control_channel_manager, + send_compression, + accept_compression, + ) + } + + fn make_flight_client_with_manager( + &self, + channel_manager: &ChannelManager, + send_compression: bool, + accept_compression: bool, + ) -> Result { + let addr = self + .inner + .get_peer() + .context(error::IllegalGrpcClientStateSnafu { + err_msg: "No available peer found", + })?; + let channel = channel_manager + .get(&addr) + .context(error::CreateChannelSnafu { addr: &addr })?; let mut client = FlightServiceClient::new(channel) - .max_decoding_message_size(self.max_grpc_recv_message_size()) - .max_encoding_message_size(self.max_grpc_send_message_size()); + .max_decoding_message_size( + channel_manager.config().max_recv_message_size.as_bytes() as usize + ) + .max_encoding_message_size( + channel_manager.config().max_send_message_size.as_bytes() as usize + ); // todo(hl): support compression methods. if send_compression { client = client.send_compressed(CompressionEncoding::Zstd); diff --git a/src/client/src/client_manager.rs b/src/client/src/client_manager.rs index edac45a9fe..6f048a9ef2 100644 --- a/src/client/src/client_manager.rs +++ b/src/client/src/client_manager.rs @@ -26,7 +26,10 @@ use crate::flow::FlowRequester; use crate::region::RegionRequester; pub struct NodeClients { - channel_manager: ChannelManager, + // Keep query and control traffic on independent pools. Query DoGet streams can + // remain active without consuming the control lane used by mutations. + query_channel_manager: ChannelManager, + control_channel_manager: ChannelManager, clients: Cache, } @@ -39,7 +42,8 @@ impl Default for NodeClients { impl Debug for NodeClients { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("NodeClients") - .field("channel_manager", &self.channel_manager) + .field("query_channel_manager", &self.query_channel_manager) + .field("control_channel_manager", &self.control_channel_manager) .finish() } } @@ -53,7 +57,7 @@ impl DatanodeManager for NodeClients { send_compression, accept_compression, .. - } = self.channel_manager.config(); + } = self.control_channel_manager.config(); Arc::new(RegionRequester::new( client, *send_compression, @@ -74,7 +78,8 @@ impl FlownodeManager for NodeClients { impl NodeClients { pub fn new(config: ChannelConfig) -> Self { Self { - channel_manager: ChannelManager::with_config(config, None), + query_channel_manager: ChannelManager::with_config(config.clone(), None), + control_channel_manager: ChannelManager::with_config(config, None), clients: CacheBuilder::new(1024) .time_to_live(Duration::from_secs(30 * 60)) .time_to_idle(Duration::from_secs(5 * 60)) @@ -85,8 +90,9 @@ impl NodeClients { pub async fn get_client(&self, datanode: &Peer) -> Client { self.clients .get_with_by_ref(datanode, async move { - Client::with_manager_and_urls( - self.channel_manager.clone(), + Client::with_managers_and_urls( + self.query_channel_manager.clone(), + self.control_channel_manager.clone(), vec![datanode.addr.clone()], ) }) @@ -98,3 +104,57 @@ impl NodeClients { self.clients.insert(datanode, client).await } } + +#[cfg(test)] +mod tests { + use common_grpc::channel_manager::ChannelManager; + use common_meta::peer::Peer; + + use super::{ChannelConfig, NodeClients}; + use crate::Client; + + const PEER_ADDR: &str = "127.0.0.1:3001"; + + fn assert_pool_has_one_address(manager: &ChannelManager) { + let mut count = 0; + let mut addresses = Vec::new(); + manager.retain_channel(|addr, _| { + count += 1; + addresses.push(addr.clone()); + true + }); + assert_eq!(1, count); + assert_eq!([PEER_ADDR], addresses.as_slice()); + } + + #[tokio::test] + async fn test_node_clients_route_flight_lanes_to_independent_pools() { + let node_clients = NodeClients::new(ChannelConfig::default()); + let peer = Peer { + id: 1, + addr: PEER_ADDR.to_string(), + }; + let client = node_clients.get_client(&peer).await; + + // DoGet/query uses only the query pool. + client.make_flight_client(false, false).unwrap(); + client.make_flight_client(false, false).unwrap(); + assert_pool_has_one_address(&node_clients.query_channel_manager); + node_clients + .control_channel_manager + .retain_channel(|_, _| false); + + // DoPut/mutation uses only the independently pooled control lane. + client.make_control_flight_client(false, false).unwrap(); + client.make_control_flight_client(false, false).unwrap(); + assert_pool_has_one_address(&node_clients.query_channel_manager); + assert_pool_has_one_address(&node_clients.control_channel_manager); + + // Direct constructors retain their legacy shared-manager behavior. + let manager = ChannelManager::new(); + let legacy = Client::with_manager_and_urls(manager.clone(), [PEER_ADDR]); + legacy.make_flight_client(false, false).unwrap(); + legacy.make_control_flight_client(false, false).unwrap(); + assert_pool_has_one_address(&manager); + } +} diff --git a/src/client/src/database.rs b/src/client/src/database.rs index 4fff3636b6..35e586ed4d 100644 --- a/src/client/src/database.rs +++ b/src/client/src/database.rs @@ -836,7 +836,7 @@ impl Database { MetadataValue::from_str(db_to_put).context(InvalidTonicMetadataValueSnafu)?, ); - let mut client = self.client.make_flight_client(false, false)?; + let mut client = self.client.make_control_flight_client(false, false)?; let response = client.mut_inner().do_put(request).await?; let response = response .into_inner() diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index e6e879a584..2d4f237dd6 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -30,7 +30,7 @@ use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; use common_query::{Output, OutputData, OutputMeta}; use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters}; -use common_recordbatch::{EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream}; +use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream}; use common_telemetry::tracing; use datafusion::catalog::TableFunction; use datafusion::dataframe::DataFrame; @@ -44,7 +44,6 @@ use datafusion_expr::{ use datatypes::prelude::VectorRef; use datatypes::schema::Schema; use futures_util::StreamExt; -use futures_util::future::try_join; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use sqlparser::ast::AnalyzeFormat; @@ -81,29 +80,6 @@ pub const QUERY_PARALLELISM_HINT: &str = "query_parallelism"; /// Whether to fallback to the original plan when failed to push down. pub const QUERY_FALLBACK_HINT: &str = "query_fallback"; -// An unbounded queue keeps draining source RPCs while mutation RPCs on a shared -// HTTP/2 connection are pending, trading bounded memory for request liveness. -async fn forward_record_batches( - mut stream: SendableRecordBatchStream, - batch_tx: tokio::sync::mpsc::UnboundedSender>, -) -> Result<()> { - while let Some(batch) = stream.next().await { - match batch.context(CreateRecordBatchSnafu) { - Ok(batch) => { - if batch_tx.send(Ok(batch)).is_err() { - break; - } - tokio::task::yield_now().await; - } - Err(error) => { - let _ = batch_tx.send(Err(error)); - break; - } - } - } - Ok(()) -} - fn query_load_region_id(plan: &Arc) -> Option { let mut region_id = None; let mut stack = vec![plan.clone()]; @@ -229,7 +205,7 @@ impl DatafusionQueryEngine { let Output { data, meta } = self .exec_query_plan((*dml.input).clone(), query_ctx.clone()) .await?; - let stream = match data { + let mut stream = match data { OutputData::RecordBatches(batches) => batches.as_stream(), OutputData::Stream(stream) => stream, _ => unreachable!(), @@ -238,35 +214,30 @@ impl DatafusionQueryEngine { let mut affected_rows = 0; let mut insert_cost = 0; - match dml.op { - WriteOp::Insert(_) => { - let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); - let producer = forward_record_batches(stream, batch_tx); - let consumer = self.consume_insert_record_batches( - batch_rx, - &table_name, - table.schema(), - query_ctx.clone(), - ); - let ((), (rows, cost)) = try_join(producer, consumer).await?; - affected_rows += rows; - insert_cost += cost; + while let Some(batch) = stream.next().await { + let batch = batch.context(CreateRecordBatchSnafu)?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table.schema()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + + match dml.op { + WriteOp::Insert(_) => { + // We ignore the insert op. + let output = self + .insert(&table_name, column_vectors, query_ctx.clone()) + .await?; + let (rows, cost) = output.extract_rows_and_cost(); + affected_rows += rows; + insert_cost += cost; + } + WriteOp::Delete => { + affected_rows += self + .delete(&table_name, &table, column_vectors, query_ctx.clone()) + .await?; + } + _ => unreachable!("guarded by the 'ensure!' at the beginning"), } - WriteOp::Delete => { - // Keep DELETE on the same producer/consumer schedule as INSERT so the source - // stream can continue draining while mutation RPCs are pending. - let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); - let producer = forward_record_batches(stream, batch_tx); - let consumer = self.consume_delete_record_batches( - batch_rx, - &table_name, - &table, - query_ctx.clone(), - ); - let (rows, ()) = try_join(consumer, producer).await?; - affected_rows += rows; - } - _ => unreachable!("guarded by the 'ensure!' at the beginning"), } Ok(Output::new( OutputData::AffectedRows(affected_rows), @@ -274,53 +245,6 @@ impl DatafusionQueryEngine { )) } - async fn consume_insert_record_batches( - &self, - mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, - table_name: &ResolvedTableReference, - table_schema: Arc, - query_ctx: QueryContextRef, - ) -> Result<(usize, usize)> { - let mut affected_rows = 0; - let mut insert_cost = 0; - while let Some(batch) = batch_rx.recv().await { - let batch = batch?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table_schema.clone()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - // We ignore the insert op. - let output = self - .insert(table_name, column_vectors, query_ctx.clone()) - .await?; - let (rows, cost) = output.extract_rows_and_cost(); - affected_rows += rows; - insert_cost += cost; - } - Ok((affected_rows, insert_cost)) - } - - async fn consume_delete_record_batches( - &self, - mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, - table_name: &ResolvedTableReference, - table: &TableRef, - query_ctx: QueryContextRef, - ) -> Result { - let mut affected_rows = 0; - while let Some(batch) = batch_rx.recv().await { - let batch = batch?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table.schema()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - affected_rows += self - .delete(table_name, table, column_vectors, query_ctx.clone()) - .await?; - } - Ok(affected_rows) - } - #[tracing::instrument(skip_all)] async fn delete( &self, diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index 3dee75e066..9dd8b952ed 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -17,22 +17,33 @@ mod test { use std::collections::HashMap; use std::net::SocketAddr; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::{Duration, Instant}; use api::v1::auth_header::AuthScheme; use api::v1::query_request::Query; + use api::v1::region::{FlushRequest, RegionRequest, RegionRequestHeader, region_request}; use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, QueryRequest, SemanticType}; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::{FlightData, FlightDescriptor, Ticket}; + use async_trait::async_trait; use auth::user_provider_from_option; + use client::client_manager::NodeClients; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_grpc::flight::do_put::DoPutMetadata; use common_grpc::flight::{FlightEncoder, FlightMessage}; + use common_meta::node_manager::DatanodeManager; + use common_meta::peer::Peer; use common_query::OutputData; + use common_query::request::QueryRequest as RegionQueryRequest; use common_recordbatch::adapter::RegionWatermarkEntry; - use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; + use common_recordbatch::{ + RecordBatch, RecordBatchStreamWrapper, RecordBatches, SendableRecordBatchStream, + }; + use common_runtime::Builder as RuntimeBuilder; + use common_runtime::runtime::BuilderBuild; use common_telemetry::tracing_context::TracingContext; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; @@ -46,8 +57,11 @@ mod test { TonicStream, }; use servers::grpc::greptime_handler::GreptimeRequestHandler; + use servers::grpc::region_server::{RegionServerHandler, RegionServerRequestHandler}; use servers::grpc::{FlightCompression, GrpcServerConfig}; use servers::server::Server; + use tokio::net::TcpListener; + use tokio_stream::wrappers::TcpListenerStream; use tonic::Response; use tonic::transport::Server as TonicServer; use tower::service_fn; @@ -59,6 +73,52 @@ mod test { struct SlowFlightCraft; + struct RetainedFlightCraft; + + struct RetainedRegionHandler; + + #[async_trait] + impl RegionServerHandler for RetainedRegionHandler { + async fn handle( + &self, + _request: region_request::Body, + ) -> servers::error::Result { + Ok(api::v1::region::RegionResponse { + header: Some(api::v1::ResponseHeader { + status: Some(api::v1::Status { + status_code: common_error::status_code::StatusCode::Success as _, + ..Default::default() + }), + }), + ..Default::default() + }) + } + } + + #[async_trait] + impl servers::grpc::flight::FlightCraft for RetainedFlightCraft { + async fn do_get( + &self, + _: tonic::Request, + ) -> std::result::Result>, tonic::Status> { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "value", + ConcreteDataType::int32_datatype(), + false, + )])); + let stream = + futures_util::stream::pending::>(); + let recordbatches = RecordBatchStreamWrapper::new(schema, stream); + let stream = FlightRecordBatchStream::new( + FlightRecordBatchSource::RecordBatches(Box::pin(recordbatches)), + TracingContext::default(), + FlightCompression::default(), + session::context::QueryContext::arc(), + ); + Ok(Response::new(Box::pin(stream))) + } + } + fn slow_recordbatch_stream() -> SendableRecordBatchStream { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "value", @@ -96,6 +156,75 @@ mod test { } } + #[tokio::test(flavor = "multi_thread")] + async fn test_node_clients_route_retained_region_query_and_control_to_separate_connections() { + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server_accepted_connections = accepted_connections.clone(); + let incoming = TcpListenerStream::new(listener).map(move |result| { + result.map(|stream| { + server_accepted_connections.fetch_add(1, Ordering::SeqCst); + stream + }) + }); + let runtime = RuntimeBuilder::default() + .worker_threads(2) + .thread_name("retained-region-test") + .build() + .unwrap(); + let region_handler = + RegionServerRequestHandler::new(Arc::new(RetainedRegionHandler), runtime); + let mut server = tokio::spawn(async move { + TonicServer::builder() + .add_service(FlightServiceServer::new(FlightCraftWrapper( + RetainedFlightCraft, + ))) + .add_service(api::v1::region::region_server::RegionServer::new( + region_handler, + )) + .serve_with_incoming_shutdown(incoming, async { + let _ = shutdown_rx.await; + }) + .await + .unwrap(); + }); + + let peer = Peer::new(1, addr.to_string()); + let node_clients = NodeClients::new(ChannelConfig::new().timeout(None)); + let datanode = node_clients.datanode(&peer).await; + let query_request = RegionQueryRequest { + header: None, + region_id: store_api::storage::RegionId::new(1, 0), + plan: datafusion_expr::LogicalPlanBuilder::empty(false) + .build() + .unwrap(), + }; + let retained_query = datanode.handle_query(query_request.clone()).await.unwrap(); + + let control_request = RegionRequest { + header: Some(RegionRequestHeader::default()), + body: Some(region_request::Body::Flush(FlushRequest::default())), + }; + datanode.handle(control_request.clone()).await.unwrap(); + assert_eq!(2, accepted_connections.load(Ordering::SeqCst)); + + let second_query = datanode.handle_query(query_request).await.unwrap(); + datanode.handle(control_request).await.unwrap(); + assert_eq!(2, accepted_connections.load(Ordering::SeqCst)); + + drop(retained_query); + drop(second_query); + drop(datanode); + drop(node_clients); + shutdown_tx.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(2), &mut server) + .await + .unwrap() + .unwrap(); + } + #[tokio::test(flavor = "multi_thread")] async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { let (client_io, server_io) = tokio::io::duplex(1024);