feat: Cache ballista clients on executor

ballista/core/src/client.rs

@@ -114,6 +114,24 @@ impl BallistaClient {
         })
     }
 
+    /// creates a ballista client to be used for testing
+    /// it connects lazily and which can not really
+    /// be reconfigured.
+    pub fn new_for_test(host: &str, port: u16) -> Self {
+        use tonic::transport::Endpoint;
+        let addr = format!("http://{host}:{port}");
+        let channel = Endpoint::from_shared(addr)
+            .expect("valid address")
+            .connect_lazy();
+        Self {
+            io_retries_times: 3,
+            io_retry_wait_time_ms: 250,
+            host: host.to_string(),
+            port,
+            flight_client: FlightServiceClient::new(channel),
+        }
+    }
+
     /// Retrieves a partition from an executor.
     ///
     /// Depending on the value of the `flight_transport` parameter, this method will utilize either

ballista/core/src/client_pool.rs

@@ -0,0 +1,107 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Connection pool for `BallistaClient` instances.
+
+use crate::client::BallistaClient;
+use crate::error::Result;
+use crate::extension::BallistaConfigGrpcEndpoint;
+use crate::utils::GrpcClientConfig;
+use async_trait::async_trait;
+use std::fmt::Debug;
+use std::ops::{Deref, DerefMut};
+use std::sync::Arc;
+
+// ---------------------------------------------------------------------------
+// Trait
+// ---------------------------------------------------------------------------
+
+/// Manages a pool of reusable [BallistaClient] connections.
+#[async_trait]
+pub trait BallistaClientPool: Send + Sync + Debug {
+    /// Acquire an idle client for `(host, port, config)`, or create a new one if the
+    /// pool is empty for that key. The returned [PooledClient] returns itself
+    /// to the pool on drop.
+    async fn acquire(
+        &self,
+        host: &str,
+        port: u16,
+        config: &GrpcClientConfig,
+        customize_endpoint: Option<Arc<BallistaConfigGrpcEndpoint>>,
+    ) -> Result<PooledClient>;
+
+    /// Remove all idle clients that have been sitting unused longer than the
+    /// configured idle timeout.
+    async fn evict_idle(&self);
+}
+
+// ---------------------------------------------------------------------------
+// PooledClient guard
+// ---------------------------------------------------------------------------
+
+/// A [BallistaClient] checked out from a pool.
+///
+/// Implements [Deref] / [DerefMut] so it can be used exactly like a
+/// [BallistaClient]. On drop, the inner client is returned to the pool
+/// automatically. Call [PooledClient::discard] before dropping if the
+/// connection should **not** be reused (e.g. after a transport error).
+pub struct PooledClient {
+    client: BallistaClient,
+    /// Invoked in `Drop::drop` to push the client back into the idle deque.
+    /// `None` after `discard()` is called.
+    return_fn: Option<Box<dyn FnOnce(BallistaClient) + Send>>,
+}
+
+impl PooledClient {
+    /// Creates new PooledClient
+    pub fn new(
+        client: BallistaClient,
+        return_fn: Box<dyn FnOnce(BallistaClient) + Send>,
+    ) -> Self {
+        Self {
+            client,
+            return_fn: Some(return_fn),
+        }
+    }
+
+    /// Close the connection instead of returning it to the pool.
+    pub fn discard(mut self) {
+        self.return_fn = None;
+    }
+}
+
+impl Deref for PooledClient {
+    type Target = BallistaClient;
+    fn deref(&self) -> &Self::Target {
+        &self.client
+    }
+}
+
+impl DerefMut for PooledClient {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.client
+    }
+}
+
+impl Drop for PooledClient {
+    fn drop(&mut self) {
+        if let Some(f) = self.return_fn.take() {
+            let client = self.client.clone();
+            f(client);
+        }
+    }
+}

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use crate::client::BallistaClient;
+use crate::client_pool::BallistaClientPool;
 use crate::error::BallistaError;
 use crate::execution_plans::sort_shuffle::{
     get_index_path, is_sort_shuffle_output, stream_sort_shuffle_partition,
@@ -72,6 +73,7 @@ pub struct ShuffleReaderExec {
     metrics: ExecutionPlanMetricsSet,
     properties: Arc<PlanProperties>,
     work_dir: Option<String>,
+    client_pool: Option<Arc<dyn BallistaClientPool>>,
 }
 
 impl ShuffleReaderExec {
@@ -94,19 +96,33 @@ impl ShuffleReaderExec {
             partition,
             metrics: ExecutionPlanMetricsSet::new(),
             properties,
-            work_dir: None, // to be updated at the executor side
+            work_dir: None,    // to be updated at the executor side
+            client_pool: None, // to be updated at the executor side
         })
     }
 
     /// changes work dir where shuffle files are located
-    pub fn change_work_dir(&self, work_dir: String) -> Self {
+    pub fn with_work_dir(&self, work_dir: String) -> Self {
         Self {
             stage_id: self.stage_id,
             schema: self.schema.clone(),
             partition: self.partition.clone(),
             metrics: self.metrics.clone(),
             properties: self.properties.clone(),
             work_dir: Some(work_dir),
+            client_pool: self.client_pool.clone(),
+        }
+    }
+    /// creates new shuffle reader with client pool
+    pub fn with_client_pool(&self, client_pool: Arc<dyn BallistaClientPool>) -> Self {
+        Self {
+            stage_id: self.stage_id,
+            schema: self.schema.clone(),
+            partition: self.partition.clone(),
+            metrics: self.metrics.clone(),
+            properties: self.properties.clone(),
+            work_dir: self.work_dir.clone(),
+            client_pool: Some(client_pool),
         }
     }
 }
@@ -212,8 +228,12 @@ impl ExecutionPlan for ShuffleReaderExec {
                 "ShuffleReader work dir should have been set by executor".to_owned(),
             ))?;
 
-        let response_receiver =
-            send_fetch_partitions(work_dir, partition_locations, config);
+        let response_receiver = send_fetch_partitions(
+            work_dir,
+            partition_locations,
+            config,
+            self.client_pool.clone(),
+        );
 
         let input_stream = Box::pin(RecordBatchStreamAdapter::new(
             self.schema.clone(),
@@ -415,6 +435,7 @@ fn send_fetch_partitions(
     work_dir: &str,
     partition_locations: Vec<PartitionLocation>,
     config: &SessionConfig,
+    client_pool: Option<Arc<dyn BallistaClientPool>>,
 ) -> AbortableReceiverStream {
     let max_request_num = config.ballista_shuffle_reader_maximum_concurrent_requests();
     let sort_shuffle_enabled = config.ballista_sort_shuffle_enabled();
@@ -467,6 +488,7 @@ fn send_fetch_partitions(
         spawned_tasks.push(SpawnedTask::spawn({
             let customize_endpoint = customize_endpoint.clone();
             let grpc_config = grpc_config.clone();
+            let client_pool = client_pool.clone();
             async move {
                 // Block if exceeds max request number.
                 let permit = semaphore.acquire_owned().await.unwrap();
@@ -475,6 +497,7 @@ fn send_fetch_partitions(
                     grpc_config,
                     prefer_flight,
                     customize_endpoint,
+                    client_pool,
                 )
                 .await;
                 // Block if the channel buffer is full.
@@ -518,6 +541,7 @@ async fn fetch_partition_remote(
     config: Arc<GrpcClientConfig>,
     prefer_flight: bool,
     customize_endpoint: Option<Arc<BallistaConfigGrpcEndpoint>>,
+    client_pool: Option<Arc<dyn BallistaClientPool>>,
 ) -> result::Result<SendableRecordBatchStream, BallistaError> {
     let metadata = &location.executor_meta;
     let partition_id = &location.partition_id;
@@ -526,13 +550,11 @@ async fn fetch_partition_remote(
     let host = metadata.host.as_str();
     let port = metadata.port;
 
-    // TODO for shuffle client connections, we should avoid creating new connections again and again.
-    // And we should also avoid to keep alive too many connections for long time.
-    let mut ballista_client =
-        new_ballista_client(host, port, &config, customize_endpoint)
+    if let Some(pool) = client_pool {
+        let mut pooled = pool
+            .acquire(host, port, &config, customize_endpoint)
             .await
             .map_err(|error| match error {
-                // map grpc connection error to partition fetch error.
                 BallistaError::GrpcConnectionError(msg) => BallistaError::FetchFailed(
                     metadata.id.clone(),
                     partition_id.stage_id,
@@ -542,15 +564,47 @@ async fn fetch_partition_remote(
                 other => other,
             })?;
 
-    ballista_client
-        .fetch_partition(
-            &metadata.id,
-            partition_id,
-            file_id,
-            is_sort_shuffle,
-            prefer_flight,
-        )
-        .await
+        let result = pooled
+            .fetch_partition(
+                &metadata.id,
+                partition_id,
+                file_id,
+                is_sort_shuffle,
+                prefer_flight,
+            )
+            .await;
+        if result.is_err() {
+            pooled.discard();
+        }
+        result
+    } else {
+        // TODO for shuffle client connections, we should avoid creating new connections again and again.
+        // And we should also avoid to keep alive too many connections for long time.
+        let mut ballista_client =
+            new_ballista_client(host, port, &config, customize_endpoint)
+                .await
+                .map_err(|error| match error {
+                    BallistaError::GrpcConnectionError(msg) => {
+                        BallistaError::FetchFailed(
+                            metadata.id.clone(),
+                            partition_id.stage_id,
+                            partition_id.partition_id,
+                            msg,
+                        )
+                    }
+                    other => other,
+                })?;
+
+        ballista_client
+            .fetch_partition(
+                &metadata.id,
+                partition_id,
+                file_id,
+                is_sort_shuffle,
+                prefer_flight,
+            )
+            .await
+    }
 }
 
 fn fetch_partition_local(
@@ -999,7 +1053,7 @@ mod tests {
             Arc::new(schema),
             Partitioning::UnknownPartitioning(4),
         )?
-        .change_work_dir(work_dir);
+        .with_work_dir(work_dir);
         let mut stream = shuffle_reader_exec.execute(0, task_ctx)?;
         let batches = utils::collect_stream(&mut stream).await;
 
@@ -1143,6 +1197,7 @@ mod tests {
             &work_dir.to_string_lossy(),
             partition_locations,
             &config,
+            None,
         );
 
         let stream = RecordBatchStreamAdapter::new(

ballista/core/src/lib.rs

@@ -33,6 +33,8 @@ pub fn print_version() {
 
 /// Client utilities for connecting to Ballista schedulers.
 pub mod client;
+/// Connection pool for reusing `BallistaClient` instances across requests.
+pub mod client_pool;
 /// Configuration options and settings for Ballista components.
 pub mod config;
 /// Utilities for generating execution plan diagrams.

ballista/core/src/utils.rs

@@ -54,7 +54,7 @@ use tonic::transport::{Channel, Endpoint, Error, Server};
 /// let ballista_config = BallistaConfig::default();
 /// let grpc_config = GrpcClientConfig::from(&ballista_config);
 /// ```
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct GrpcClientConfig {
     /// Connection timeout in seconds
     pub connect_timeout_seconds: u64,