fix(client): isolate query and control transports

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

src/client/src/client.rs

@@ -51,18 +51,34 @@ pub struct Client {
     inner: Arc<Inner>,
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug)]
 struct Inner {
-    channel_manager: ChannelManager,
+    query_channel_manager: ChannelManager,
+    control_channel_manager: ChannelManager,
     peers: Arc<RwLock<Vec<String>>>,
     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<str>,
         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<U, A>(
+        query_channel_manager: ChannelManager,
+        control_channel_manager: ChannelManager,
+        urls: A,
+    ) -> Self
+    where
+        U: AsRef<str>,
+        A: AsRef<[U]>,
+    {
+        let inner = Inner::with_managers(query_channel_manager, control_channel_manager);
         let urls: Vec<String> = 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<FlightClient> {
-        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<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> {
+        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);

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<Peer, Client>,
 }
 
@@ -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);
+    }
+}

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()

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<RecordBatch>>,
-) -> 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<dyn ExecutionPlan>) -> Option<u64> {
     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<Result<RecordBatch>>,
-        table_name: &ResolvedTableReference,
-        table_schema: Arc<Schema>,
-        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<Result<RecordBatch>>,
-        table_name: &ResolvedTableReference,
-        table: &TableRef,
-        query_ctx: QueryContextRef,
-    ) -> Result<usize> {
-        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,

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<api::v1::region::RegionResponse> {
+            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<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(
+                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);