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>
2026-09-06 05:28:57 +00:00 · 2026-09-03 08:57:12 +00:00
parent d664326b1e
commit 945e53e0a3
4 changed files with 323 additions and 27 deletions
@@ -76,7 +76,8 @@ pub struct Client {

 #[derive(Debug)]
 struct Inner {
-    channel_manager: ChannelManager,
+    query_channel_manager: ChannelManager,
+    control_channel_manager: ChannelManager,
    peers: RwLock<Peers>,
    load_balance: Loadbalancer,
    health_check_interval: Duration,
@@ -108,10 +109,20 @@ impl Inner {
        channel_manager: ChannelManager,
        peers: Vec<String>,
        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<String>,
+        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<U, A>(
-        channel_manager: ChannelManager,
+    pub(crate) fn with_managers_and_urls<U, A>(
+        query_channel_manager: ChannelManager,
+        control_channel_manager: ChannelManager,
+        urls: A,
+    ) -> Self
+    where
+        U: AsRef<str>,
+        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<U, A>(
+        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<U, A>(
+        channel_manager: ChannelManager,
+        urls: A,
+        options: ClientOptions,
+    ) -> Self
+    where
+        U: AsRef<str>,
+        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<U, A>(&self, urls: A)
    where
        U: AsRef<str>,
@@ -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<FlightClient> {
-        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<FlightClient> {
+        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<FlightClient> {
+        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<Channel>)> {
        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<Channel>)> {
        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))
@@ -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<Peer, Client>,
 }

@@ -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);
+    }
+}
@@ -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()
@@ -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<Ticket>,
+        ) -> std::result::Result<Response<TonicStream<FlightData>>, tonic::Status> {
+            let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
+                "value",
+                ConcreteDataType::int32_datatype(),
+                false,
+            )]));
+            let stream =
+                futures_util::stream::pending::<common_recordbatch::error::Result<RecordBatch>>();
+            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);