From 945e53e0a385f155f8b8799b1e07ba1f133a17d2 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 3 Sep 2026 08:57:12 +0000 Subject: [PATCH] fix(client): isolate query and control transports (#8990) * refactor(client): isolate query and control transports Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(client): cover retained Flight transport isolation Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * style(client): satisfy retained Flight test lint Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(client): clarify transport lane routing Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/client/src/client.rs | 149 +++++++++++++++++++++++---- src/client/src/client_manager.rs | 73 +++++++++++-- src/client/src/database.rs | 2 +- tests-integration/src/grpc/flight.rs | 126 +++++++++++++++++++++- 4 files changed, 323 insertions(+), 27 deletions(-) diff --git a/src/client/src/client.rs b/src/client/src/client.rs index 8b6671cb5c..fc40c98a72 100644 --- a/src/client/src/client.rs +++ b/src/client/src/client.rs @@ -76,7 +76,8 @@ pub struct Client { #[derive(Debug)] struct Inner { - channel_manager: ChannelManager, + query_channel_manager: ChannelManager, + control_channel_manager: ChannelManager, peers: RwLock, load_balance: Loadbalancer, health_check_interval: Duration, @@ -108,10 +109,20 @@ impl Inner { channel_manager: ChannelManager, peers: Vec, options: ClientOptions, + ) -> Self { + Self::with_managers_and_peers(channel_manager.clone(), channel_manager, peers, options) + } + + fn with_managers_and_peers( + query_channel_manager: ChannelManager, + control_channel_manager: ChannelManager, + peers: Vec, + options: ClientOptions, ) -> Self { let peer_count = peers.len(); Self { - channel_manager, + query_channel_manager, + control_channel_manager, peers: RwLock::new(Peers { addresses: peers, states: PeerStates { @@ -188,7 +199,7 @@ impl Inner { } async fn check_peer_health(&self, addr: &str) -> bool { - let Ok(channel) = self.channel_manager.get(addr) else { + let Ok(channel) = self.control_channel_manager.get(addr) else { return false; }; let mut client = HealthCheckClient::new(channel); @@ -269,9 +280,26 @@ impl Client { Self::with_manager_and_urls_and_options(channel_manager, urls, ClientOptions::default()) } - /// Creates a client with a channel manager, URLs, and custom options. - pub fn with_manager_and_urls_and_options( - channel_manager: ChannelManager, + pub(crate) fn with_managers_and_urls( + query_channel_manager: ChannelManager, + control_channel_manager: ChannelManager, + urls: A, + ) -> Self + where + U: AsRef, + A: AsRef<[U]>, + { + Self::with_managers_and_urls_and_options( + query_channel_manager, + control_channel_manager, + urls, + ClientOptions::default(), + ) + } + + fn with_managers_and_urls_and_options( + query_channel_manager: ChannelManager, + control_channel_manager: ChannelManager, urls: A, options: ClientOptions, ) -> Self @@ -285,14 +313,34 @@ impl Client { .map(|peer| peer.as_ref().to_string()) .collect(); Self { - inner: Arc::new(Inner::with_manager_and_peers( - channel_manager, + inner: Arc::new(Inner::with_managers_and_peers( + query_channel_manager, + control_channel_manager, urls, options, )), } } + /// Creates a client with a channel manager, URLs, and custom options. + pub fn with_manager_and_urls_and_options( + channel_manager: ChannelManager, + urls: A, + options: ClientOptions, + ) -> Self + where + U: AsRef, + A: AsRef<[U]>, + { + let channel_manager_for_query = channel_manager.clone(); + Self::with_managers_and_urls_and_options( + channel_manager_for_query, + channel_manager, + urls, + options, + ) + } + pub fn start(&self, urls: A) where U: AsRef, @@ -350,7 +398,7 @@ impl Client { let channel = self .inner - .channel_manager + .control_channel_manager .get(&addr) .context(error::CreateChannelSnafu { addr: &addr })?; Ok((addr, channel)) @@ -358,7 +406,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 @@ -366,22 +414,63 @@ 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. + /// + /// This public name is retained for compatibility. 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, + ) + } + + 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 { + self.trigger_health_check(); + 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); @@ -396,16 +485,40 @@ impl Client { pub(crate) fn raw_region_client(&self) -> Result<(String, PbRegionClient)> { let (addr, channel) = self.find_channel()?; let client = PbRegionClient::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( + self.inner + .control_channel_manager + .config() + .max_recv_message_size + .as_bytes() as usize, + ) + .max_encoding_message_size( + self.inner + .control_channel_manager + .config() + .max_send_message_size + .as_bytes() as usize, + ); Ok((addr, client)) } pub(crate) fn raw_flow_client(&self) -> Result<(String, PbFlowClient)> { let (addr, channel) = self.find_channel()?; let client = PbFlowClient::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( + self.inner + .control_channel_manager + .config() + .max_recv_message_size + .as_bytes() as usize, + ) + .max_encoding_message_size( + self.inner + .control_channel_manager + .config() + .max_send_message_size + .as_bytes() as usize, + ) .accept_compressed(CompressionEncoding::Zstd) .send_compressed(CompressionEncoding::Zstd); Ok((addr, client)) diff --git a/src/client/src/client_manager.rs b/src/client/src/client_manager.rs index edac45a9fe..008ada00d2 100644 --- a/src/client/src/client_manager.rs +++ b/src/client/src/client_manager.rs @@ -26,7 +26,11 @@ use crate::flow::FlowRequester; use crate::region::RegionRequester; pub struct NodeClients { - channel_manager: ChannelManager, + // Keep the channel managers independent by lane: + // query = Flight DoGet/distributed reads; control/mutation = health, unary RPCs, + // gateway, and Flight DoPut. + query_channel_manager: ChannelManager, + control_channel_manager: ChannelManager, clients: Cache, } @@ -39,7 +43,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 +58,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 +79,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 +91,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 +105,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_use_isolated_reused_channel_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; + + client.make_flight_client(false, false).unwrap(); + client.make_flight_client(false, false).unwrap(); + assert_pool_has_one_address(&node_clients.query_channel_manager); + let mut control_count = 0; + node_clients.control_channel_manager.retain_channel(|_, _| { + control_count += 1; + true + }); + assert_eq!(0, control_count); + + 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); + + 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 7c33794f47..881730cfaa 100644 --- a/src/client/src/database.rs +++ b/src/client/src/database.rs @@ -847,7 +847,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/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index 1a2d841600..f37d5ee638 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -18,24 +18,30 @@ mod test { use std::net::SocketAddr; use std::pin::Pin; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::{Duration, Instant}; use api::v1::auth_header::AuthScheme; use api::v1::greptime_request::Request as GreptimeQueryRequest; + use api::v1::health_check_server::HealthCheckServer; use api::v1::query_request::Query; use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, QueryRequest, SemanticType}; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::{FlightData, FlightDescriptor, Ticket}; use auth::user_provider_from_option; + use client::client_manager::NodeClients; use client::region::RegionRequester; 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, DoPutResponse}; - use common_grpc::flight::{FlightEncoder, FlightMessage}; + use common_grpc::flight::{FlightDecoder, FlightEncoder, FlightMessage}; + use common_meta::peer::Peer; use common_query::{Output, OutputData}; use common_recordbatch::adapter::RegionWatermarkEntry; - use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; + use common_recordbatch::{ + RecordBatch, RecordBatchStreamWrapper, RecordBatches, SendableRecordBatchStream, + }; use common_telemetry::tracing_context::TracingContext; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; @@ -53,6 +59,8 @@ mod test { use servers::query_handler::grpc::GrpcQueryHandler; use servers::server::Server; use session::context::QueryContextRef; + use tokio::net::TcpListener; + use tokio_stream::wrappers::TcpListenerStream; use tonic::transport::Server as TonicServer; use tonic::{Response, Status}; use tower::service_fn; @@ -64,6 +72,8 @@ mod test { struct SlowFlightCraft; + struct RetainedFlightCraft; + struct ErrorFlightCraft; struct SlowRemoteQueryHandler { @@ -109,6 +119,32 @@ mod test { } } + #[async_trait::async_trait] + impl 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( + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + Box::pin(recordbatches), + )), + TracingContext::default(), + FlightCompression::default(), + session::context::QueryContext::arc(), + ); + Ok(Response::new(Box::pin(stream))) + } + } + #[async_trait::async_trait] impl FlightCraft for ErrorFlightCraft { async fn do_get( @@ -201,6 +237,92 @@ mod test { Client::with_manager_and_urls(channel_manager, [addr]) } + #[tokio::test(flavor = "multi_thread")] + async fn test_retained_flight_stream_uses_separate_control_connection() { + // This models a cursor-like transport condition: a DoGet response whose first Flight + // message is available while the rest remains active. The control lane must stay live + // without relying on an unconditional production deadlock to reproduce the transport risk. + 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.inspect(|_| { + server_accepted_connections.fetch_add(1, Ordering::SeqCst); + }) + }); + let mut server = tokio::spawn(async move { + TonicServer::builder() + .add_service(FlightServiceServer::new(FlightCraftWrapper( + RetainedFlightCraft, + ))) + .add_service(HealthCheckServer::new(servers::grpc::HealthCheckHandler)) + .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 client = node_clients.get_client(&peer).await; + + let mut flight_client = client.make_flight_client(false, false).unwrap(); + let mut retained_stream = flight_client + .mut_inner() + .do_get(tonic::Request::new(Ticket::default())) + .await + .unwrap() + .into_inner(); + let first_data = retained_stream.message().await.unwrap().unwrap(); + let mut decoder = FlightDecoder::default(); + assert!(matches!( + decoder.try_decode(&first_data).unwrap(), + Some(FlightMessage::Schema(_)) + )); + + tokio::time::timeout(Duration::from_secs(2), client.health_check()) + .await + .expect("control RPC should not wait for retained DoGet") + .unwrap(); + assert_eq!(2, accepted_connections.load(Ordering::SeqCst)); + + // A second DoGet and control RPC must reuse their respective physical connections. + let mut second_flight_client = client.make_flight_client(false, false).unwrap(); + let mut second_stream = second_flight_client + .mut_inner() + .do_get(tonic::Request::new(Ticket::default())) + .await + .unwrap() + .into_inner(); + let second_data = second_stream.message().await.unwrap().unwrap(); + let mut second_decoder = FlightDecoder::default(); + assert!(matches!( + second_decoder.try_decode(&second_data).unwrap(), + Some(FlightMessage::Schema(_)) + )); + client.health_check().await.unwrap(); + assert_eq!(2, accepted_connections.load(Ordering::SeqCst)); + + // Release the retained and secondary responses before gracefully stopping the real server. + drop(retained_stream); + drop(second_stream); + drop(flight_client); + drop(second_flight_client); + drop(client); + drop(node_clients); + shutdown_tx.send(()).unwrap(); + let server_result = tokio::time::timeout(Duration::from_secs(2), &mut server).await; + if server_result.is_err() { + server.abort(); + let _ = server.await; + panic!("Flight test server did not stop after stream release"); + } + server_result.unwrap().unwrap(); + } + #[tokio::test(flavor = "multi_thread")] async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { let client = client_for_flight_craft("slow-flight", SlowFlightCraft);