martin-g commented on code in PR #1578:
URL:
https://github.com/apache/datafusion-ballista/pull/1578#discussion_r3161253476
##########
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 {
Review Comment:
There is no `#[cfg(test)]`.
Does it need to be here ? It could be a helper method in `mod tests` too
##########
ballista/executor/src/client_pool.rs:
##########
@@ -0,0 +1,351 @@
+// 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
Review Comment:
```suggestion
//! `DefaultBallistaClientPool` maintains a `VecDeque` of idle clients per
```
##########
ballista/executor/src/client_pool.rs:
##########
@@ -0,0 +1,351 @@
+// 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<IdleEntry>>;
+
+struct Inner {
+ idle: IdleMap,
+ idle_timeout: Duration,
+}
+
+/// Default pool implementation.
+///
+/// Keeps a `VecDeque<BallistaClient>` 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<Inner>,
+}
+
+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<Inner> = Arc::downgrade(&inner);
+ // TODO: do we limit minimum interval here?
+ 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<Arc<BallistaConfigGrpcEndpoint>>,
+ ) -> Result<PooledClient> {
+ 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
Review Comment:
does it need to check whether it is still connected ?
if the remote party has been restarted in the meantime then the client will
probably fail without reconnecting.
##########
ballista/core/src/execution_plans/shuffle_reader.rs:
##########
@@ -94,19 +96,49 @@ 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),
+ }
+ }
+ /// creates new shuffle reader with client pool and work dir
+ pub fn with_client_pool_and_work_dir(
Review Comment:
Is this method really needed ?
The user could could it in two steps:
`with_client_pool(client_pool).with_work_dir(work_dir)`
At the moment `fn with_client_pool()` is not used anywhere.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
