diff --git a/ballista/core/src/client.rs b/ballista/core/src/client.rs index 9f746c90e7..e3c26b158b 100644 --- a/ballista/core/src/client.rs +++ b/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 diff --git a/ballista/core/src/client_pool.rs b/ballista/core/src/client_pool.rs new file mode 100644 index 0000000000..3ae02cef3a --- /dev/null +++ b/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>, + ) -> Result; + + /// 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>, +} + +impl PooledClient { + /// Creates new PooledClient + pub fn new( + client: BallistaClient, + return_fn: Box, + ) -> 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); + } + } +} diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs b/ballista/core/src/execution_plans/shuffle_reader.rs index f7b5f881fd..cc5da2bc21 100644 --- a/ballista/core/src/execution_plans/shuffle_reader.rs +++ b/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, work_dir: Option, + client_pool: Option>, } impl ShuffleReaderExec { @@ -94,12 +96,13 @@ 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(), @@ -107,6 +110,19 @@ impl ShuffleReaderExec { 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) -> 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, config: &SessionConfig, + client_pool: Option>, ) -> 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, prefer_flight: bool, customize_endpoint: Option>, + client_pool: Option>, ) -> result::Result { 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( diff --git a/ballista/core/src/lib.rs b/ballista/core/src/lib.rs index 7211772f53..7e8fad3252 100644 --- a/ballista/core/src/lib.rs +++ b/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. diff --git a/ballista/core/src/utils.rs b/ballista/core/src/utils.rs index 7320926a10..5a6dfb2537 100644 --- a/ballista/core/src/utils.rs +++ b/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, diff --git a/ballista/executor/src/client_pool.rs b/ballista/executor/src/client_pool.rs new file mode 100644 index 0000000000..b3f425870e --- /dev/null +++ b/ballista/executor/src/client_pool.rs @@ -0,0 +1,352 @@ +// 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. +//! +//! `DefaultBallistaClientPool` maintains a `VecDeque` of idle clients per +//! `(host, port, config)` key backed by a `DashMap`. Callers `BallistaClientPool::acquire` +//! a `PooledClient` guard; when the guard is dropped the underlying client is +//! returned to the idle deque automatically. +//! +//! Connections could be discarded calling `PooledClient::discard` which will result +//! of dropping connection rather than returning it to the pool. This could be +//! used for error handling. +//! +//! A optional background tokio task evicts idle connections that have not been used +//! within the configured `idle_timeout`. + +use async_trait::async_trait; +use ballista_core::client::BallistaClient; +use ballista_core::client_pool::{BallistaClientPool, PooledClient}; +use ballista_core::error::Result; +use ballista_core::extension::BallistaConfigGrpcEndpoint; +use ballista_core::utils::GrpcClientConfig; +use dashmap::DashMap; +use std::collections::VecDeque; +use std::fmt::Debug; +use std::sync::{Arc, Weak}; +use std::time::{Duration, Instant}; + +// --------------------------------------------------------------------------- +// DefaultBallistaClientPool +// --------------------------------------------------------------------------- + +struct IdleEntry { + client: BallistaClient, + idle_since: Instant, +} + +type IdleMap = DashMap<(String, u16, GrpcClientConfig), VecDeque>; + +struct Inner { + idle: IdleMap, + idle_timeout: Duration, +} + +/// Default pool implementation. +/// +/// Keeps `BallistaClients` as `VecDeque` per `(host, port, config)`. +/// Idle clients are evicted by a background tokio task that runs at `idle_timeout / 3` +/// intervals (minimum 15 s). The task exits automatically when the pool `Arc` +/// is dropped. +/// +/// The `DefaultBallistaClientPool` uses the (host, port, config) to identify a connection. +/// Therefore changing connection config might leave pooled connections +/// with older config unused until they expire. + +#[derive(Clone)] +pub struct DefaultBallistaClientPool { + inner: Arc, +} + +impl Debug for DefaultBallistaClientPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DefaultBallistaClientPool").finish() + } +} + +impl DefaultBallistaClientPool { + /// Create a pool that evicts connections idle longer + /// than defined `idle_timeout`. + pub fn with_eviction_thread(idle_timeout: Duration) -> Self { + Self::new(idle_timeout, true) + } + + /// Create a pool that evicts connections idle longer than `idle_timeout`, + /// if `enable_eviction_thread` is enabled + pub fn new(idle_timeout: Duration, enable_eviction_thread: bool) -> Self { + let inner = Arc::new(Inner { + idle: DashMap::new(), + idle_timeout, + }); + + let weak: Weak = Arc::downgrade(&inner); + // there is no empirical evidence why 15 is selected. + // we can revisit if interval < 15 is needed + let check_interval = Duration::from_secs((idle_timeout.as_secs() / 3).max(15)); + + if enable_eviction_thread { + tokio::spawn(async move { + log::debug!( + "client connection pool - eviction thread started ... interval: {check_interval:?}" + ); + let mut ticker = tokio::time::interval(check_interval); + loop { + ticker.tick().await; + + match weak.upgrade() { + None => break, + Some(pool) => { + log::trace!("client connection pool - evicting connections"); + evict(&pool.idle, pool.idle_timeout) + } + } + } + log::debug!("client connection pool - eviction thread ... DONE"); + }); + } + + Self { inner } + } + + #[cfg(test)] + /// Total number of idle connections currently held across all endpoints. + pub fn idle_count(&self) -> usize { + self.inner.idle.iter().map(|e| e.value().len()).sum() + } +} + +fn evict(idle: &IdleMap, timeout: Duration) { + let deadline = Instant::now() + .checked_sub(timeout) + .unwrap_or_else(Instant::now); + + // Drain expired entries from the front of each deque (oldest = front). + // This way pool can shrink in case of low utilization. + idle.retain(|_, deque| { + while deque.front().is_some_and(|e| e.idle_since < deadline) { + // evict from front of the queue + deque.pop_front(); + } + !deque.is_empty() + }); +} + +#[async_trait] +impl BallistaClientPool for DefaultBallistaClientPool { + async fn acquire( + &self, + host: &str, + port: u16, + config: &GrpcClientConfig, + customize_endpoint: Option>, + ) -> Result { + let key = (host.to_string(), port, config.clone()); + + // Pop the most-recently-used idle client. The DashMap shard lock is + // held only for the duration of the pop — released before the async + // BallistaClient::try_new call below. + let maybe_idle_client = self + .inner + .idle + .get_mut(&key) + .and_then(|mut deque| deque.pop_back()) // acquire from back of the queue + .map(|e| e.client); + + let client = match maybe_idle_client { + Some(client) => { + log::trace!( + "client connection pool - returning cached connection - host:{host}, port:{port}" + ); + client + } + None => { + log::trace!( + "client connection pool - returning NEW connection - host:{host}, port:{port}" + ); + BallistaClient::try_new( + host, + port, + config.max_message_size, + config.use_tls, + customize_endpoint, + config.io_retries_times, + config.io_retry_wait_time_ms, + ) + .await? + } + }; + + // The return closure captures only an Arc — synchronous, safe for Drop. + let inner_ref = Arc::clone(&self.inner); + let return_key = key; + Ok(PooledClient::new( + client, + Box::new(move |c| { + inner_ref + .idle + .entry(return_key) + .or_default() + .push_back(IdleEntry { + client: c, + idle_since: Instant::now(), + }); + }), + )) + } + + async fn evict_idle(&self) { + evict(&self.inner.idle, self.inner.idle_timeout); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ballista_core::client::BallistaClient; + use std::time::Duration; + + fn make_pool(timeout: Duration) -> DefaultBallistaClientPool { + DefaultBallistaClientPool::new(timeout, false) + } + + /// Inject an `IdleEntry` with a specific `idle_since` directly into the + /// pool's DashMap, bypassing `acquire` so no real server is needed. + fn inject_idle( + pool: &DefaultBallistaClientPool, + host: &str, + port: u16, + age: Duration, + ) { + let client = BallistaClient::new_for_test(host, port); + pool.inner + .idle + .entry((host.to_string(), port, GrpcClientConfig::default())) + .or_default() + .push_back(IdleEntry { + client, + idle_since: Instant::now() - age, + }); + } + + #[tokio::test] + async fn idle_count_starts_at_zero() { + let pool = make_pool(Duration::from_secs(60)); + assert_eq!(pool.idle_count(), 0); + } + + #[tokio::test] + async fn evict_idle_does_not_panic_on_empty_pool() { + let pool = make_pool(Duration::from_secs(60)); + pool.evict_idle().await; + assert_eq!(pool.idle_count(), 0); + } + + /// An entry older than `idle_timeout` must be removed by `evict_idle`. + #[tokio::test] + async fn evict_idle_removes_expired_entries() { + let timeout = Duration::from_millis(100); + let pool = make_pool(timeout); + + inject_idle(&pool, "host-a", 1234, timeout + Duration::from_millis(50)); + assert_eq!(pool.idle_count(), 1); + + pool.evict_idle().await; + assert_eq!(pool.idle_count(), 0); + } + + /// An entry younger than `idle_timeout` must survive eviction. + #[tokio::test] + async fn evict_idle_keeps_fresh_entries() { + let timeout = Duration::from_secs(60); + let pool = make_pool(timeout); + + inject_idle(&pool, "host-b", 2345, Duration::from_millis(10)); + assert_eq!(pool.idle_count(), 1); + + pool.evict_idle().await; + assert_eq!(pool.idle_count(), 1); + } + + /// Dropping a [PooledClient] must return the client to the pool. + #[tokio::test] + async fn pooled_client_returns_on_drop() { + let pool = make_pool(Duration::from_secs(300)); + let client = BallistaClient::new_for_test("host-c", 3456); + let key = ("host-c".to_string(), 3456u16, GrpcClientConfig::default()); + + let inner_ref = Arc::clone(&pool.inner); + let return_key = key.clone(); + let guard = PooledClient::new( + client, + Box::new(move |c| { + inner_ref + .idle + .entry(return_key) + .or_default() + .push_back(IdleEntry { + client: c, + idle_since: Instant::now(), + }); + }), + ); + + assert_eq!(pool.idle_count(), 0); + drop(guard); + assert_eq!(pool.idle_count(), 1); + } + + /// Calling `discard()` must close the connection instead of returning it. + #[tokio::test] + async fn discard_does_not_return_to_pool() { + let pool = make_pool(Duration::from_secs(300)); + let client = BallistaClient::new_for_test("host-d", 4567); + + let inner_ref = Arc::clone(&pool.inner); + let guard = PooledClient::new( + client, + Box::new(move |c| { + inner_ref + .idle + .entry(("host-d".to_string(), 4567u16, GrpcClientConfig::default())) + .or_default() + .push_back(IdleEntry { + client: c, + idle_since: Instant::now(), + }); + }), + ); + + guard.discard(); + assert_eq!(pool.idle_count(), 0); + } + + /// Mixed scenario: one expired and one fresh entry — only the expired one + /// is removed, the other survives. + #[tokio::test] + async fn evict_idle_partial_removal() { + let timeout = Duration::from_millis(100); + let pool = make_pool(timeout); + + inject_idle(&pool, "host-e", 5678, timeout + Duration::from_millis(50)); // stale + inject_idle(&pool, "host-e", 5678, Duration::from_millis(10)); // fresh + assert_eq!(pool.idle_count(), 2); + + pool.evict_idle().await; + assert_eq!(pool.idle_count(), 1); + } +} diff --git a/ballista/executor/src/config.rs b/ballista/executor/src/config.rs index 8606176472..8ddeb4ab3c 100644 --- a/ballista/executor/src/config.rs +++ b/ballista/executor/src/config.rs @@ -171,6 +171,13 @@ pub struct Config { help = "Optional total executor memory budget (e.g. \"8GB\", \"512MiB\"). Each concurrent task receives an equal share." )] pub memory_pool_size: Option, + /// Number of seconds established client connection should be cached if not used (0 means no cache) + #[arg( + long, + default_value_t = 0, + help = "Number of seconds established client connection should be cached if not used (0 means no cache, connection will be disposed)." + )] + pub client_ttl: u64, } impl TryFrom for ExecutorProcessConfig { @@ -208,6 +215,7 @@ impl TryFrom for ExecutorProcessConfig { override_physical_codec: None, override_arrow_flight_service: None, override_create_grpc_client_endpoint: None, + client_ttl: opt.client_ttl, }) } } diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index 93ab176f7d..c57d104227 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -22,6 +22,7 @@ //! for creating query stage executors from physical plans. use async_trait::async_trait; +use ballista_core::client_pool::BallistaClientPool; use ballista_core::execution_plans::sort_shuffle::SortShuffleWriterExec; use ballista_core::execution_plans::{ShuffleReaderExec, ShuffleWriterExec}; use ballista_core::serde::protobuf::ShuffleWritePartition; @@ -83,7 +84,23 @@ pub trait QueryStageExecutor: Sync + Send + Debug + Display { /// /// This implementation expects the input plan to be wrapped in a /// ShuffleWriterExec and creates a DefaultQueryStageExec to execute it. -pub struct DefaultExecutionEngine {} +#[derive(Default)] +pub struct DefaultExecutionEngine { + client_pool: Option>, +} + +impl DefaultExecutionEngine { + /// Creates new Default Execution Engine without client pooling + pub fn new() -> Self { + Self { client_pool: None } + } + /// Creates new Default Execution Engine with client pooling + pub fn with_client_pool(client_pool: Arc) -> Self { + Self { + client_pool: Some(client_pool), + } + } +} impl ExecutionEngine for DefaultExecutionEngine { fn create_query_stage_exec( @@ -98,8 +115,16 @@ impl ExecutionEngine for DefaultExecutionEngine { let plan = plan .transform(|p| { if let Some(reader) = p.as_any().downcast_ref::() { - let reader = Arc::new(reader.change_work_dir(work_dir.to_string())); - Ok(Transformed::yes(reader)) + match &self.client_pool { + Some(client_pool) => Ok(Transformed::yes(Arc::new( + reader + .with_work_dir(work_dir.to_string()) + .with_client_pool(client_pool.clone()), + ))), + None => Ok(Transformed::yes(Arc::new( + reader.with_work_dir(work_dir.to_string()), + ))), + } } else { Ok(Transformed::no(p)) } diff --git a/ballista/executor/src/executor.rs b/ballista/executor/src/executor.rs index f165685a7c..aaff163f7a 100644 --- a/ballista/executor/src/executor.rs +++ b/ballista/executor/src/executor.rs @@ -112,7 +112,7 @@ impl Executor { Arc::new(BallistaFunctionRegistry::default()), Arc::new(LoggingMetricsCollector::default()), concurrent_tasks, - None, + Arc::new(DefaultExecutionEngine::new()), ) } @@ -127,7 +127,30 @@ impl Executor { function_registry: Arc, metrics_collector: Arc, concurrent_tasks: usize, - execution_engine: Option>, + execution_engine: Arc, + ) -> Self { + Self { + metadata, + work_dir: work_dir.to_owned(), + function_registry, + runtime_producer, + config_producer, + metrics_collector, + concurrent_tasks, + abort_handles: Default::default(), + execution_engine, + } + } + /// Creates new Executor with default `ExecutionEngine`. + /// Default `ExecutionEngine` does not cache client connections. + pub fn with_default_execution_engine( + metadata: ExecutorRegistration, + work_dir: &str, + runtime_producer: RuntimeProducer, + config_producer: ConfigProducer, + function_registry: Arc, + metrics_collector: Arc, + concurrent_tasks: usize, ) -> Self { Self { metadata, @@ -138,8 +161,7 @@ impl Executor { metrics_collector, concurrent_tasks, abort_handles: Default::default(), - execution_engine: execution_engine - .unwrap_or_else(|| Arc::new(DefaultExecutionEngine {})), + execution_engine: Arc::new(DefaultExecutionEngine::new()), } } } diff --git a/ballista/executor/src/executor_process.rs b/ballista/executor/src/executor_process.rs index 6447b81063..131a50549f 100644 --- a/ballista/executor/src/executor_process.rs +++ b/ballista/executor/src/executor_process.rs @@ -63,7 +63,8 @@ use ballista_core::utils::{ }; use ballista_core::{BALLISTA_VERSION, ConfigProducer, RuntimeProducer}; -use crate::execution_engine::ExecutionEngine; +use crate::client_pool::DefaultBallistaClientPool; +use crate::execution_engine::{DefaultExecutionEngine, ExecutionEngine}; use crate::executor::{Executor, TasksDrainedFuture}; use crate::executor_server::TERMINATING; use crate::flight_service::BallistaFlightService; @@ -171,6 +172,8 @@ pub struct ExecutorProcessConfig { pub override_arrow_flight_service: Option>, /// Override function for customizing gRPC client endpoints before they are used pub override_create_grpc_client_endpoint: Option, + /// Number of seconds established client connection should be cached (0 means no cache) + pub client_ttl: u64, } impl ExecutorProcessConfig { @@ -219,6 +222,7 @@ impl Default for ExecutorProcessConfig { override_physical_codec: None, override_arrow_flight_service: None, override_create_grpc_client_endpoint: None, + client_ttl: 0, } } } @@ -334,7 +338,17 @@ pub async fn start_executor_process( opt.override_function_registry.clone().unwrap_or_default(), metrics_collector, concurrent_tasks, - opt.override_execution_engine.clone(), + opt.override_execution_engine.clone().unwrap_or_else(|| { + if opt.client_ttl > 0 { + let client_pool = + Arc::new(DefaultBallistaClientPool::with_eviction_thread( + Duration::from_secs(opt.client_ttl), + )); + Arc::new(DefaultExecutionEngine::with_client_pool(client_pool)) + } else { + Arc::new(DefaultExecutionEngine::new()) + } + }), )); let connect_timeout = opt.scheduler_connect_timeout_seconds as u64; diff --git a/ballista/executor/src/lib.rs b/ballista/executor/src/lib.rs index 781b3970e3..3db736e016 100644 --- a/ballista/executor/src/lib.rs +++ b/ballista/executor/src/lib.rs @@ -18,6 +18,8 @@ #![doc = include_str!("../README.md")] #![warn(missing_docs)] +/// Connection pool for `BallistaClient` instances. +mod client_pool; /// Execution plan for collecting distributed query results into a single partition. pub mod collect; /// Command-line configuration for the executor binary. diff --git a/ballista/executor/src/standalone.rs b/ballista/executor/src/standalone.rs index b544b5f5e7..bd72dee4b4 100644 --- a/ballista/executor/src/standalone.rs +++ b/ballista/executor/src/standalone.rs @@ -119,7 +119,7 @@ pub async fn new_standalone_executor_from_builder( info!("work_dir: {work_dir}"); - let executor = Arc::new(Executor::new( + let executor = Arc::new(Executor::with_default_execution_engine( executor_meta, &work_dir, runtime_producer, @@ -127,7 +127,6 @@ pub async fn new_standalone_executor_from_builder( Arc::new(function_registry), Arc::new(LoggingMetricsCollector::default()), concurrent_tasks, - None, )); let service = BallistaFlightService::new(work_dir); diff --git a/examples/examples/mtls-cluster.rs b/examples/examples/mtls-cluster.rs index 00a86d100e..593248183e 100644 --- a/examples/examples/mtls-cluster.rs +++ b/examples/examples/mtls-cluster.rs @@ -381,7 +381,7 @@ async fn run_executor() -> Result<(), Box> { )) }); - let executor = Arc::new(Executor::new( + let executor = Arc::new(Executor::with_default_execution_engine( executor_meta, &work_dir_str, runtime_producer, @@ -389,7 +389,6 @@ async fn run_executor() -> Result<(), Box> { Default::default(), // function_registry Arc::new(LoggingMetricsCollector::default()), // metrics_collector 4, // concurrent_tasks - None, // execution_engine )); // Start Flight service with mTLS for serving shuffle data