andygrove commented on code in PR #2779:
URL:
https://github.com/apache/datafusion-ballista/pull/2779#discussion_r4493861383
##########
ballista/core/src/client_pool.rs:
##########
@@ -108,11 +113,25 @@ impl SharedClient {
self.client.clone()
}
+ /// Installs the hook the first [discard](Self::discard) runs. Discards
+ /// after the first do not run it again.
+ ///
+ /// Intended for the pool that created this connection. Replaces any
+ /// existing hook, so callers borrowing the connection should not use it.
+ pub fn set_discard_hook(&self, hook: Box<dyn FnOnce() + Send>) {
Review Comment:
Related to the hook comment in the executor pool. This adds a public method
that silently replaces the existing hook, so anyone holding an
`Arc<SharedClient>` can clobber the pool's hook, and the doc comment is the
only guard. Once it's public we're stuck supporting it. If the hook goes away
this goes with it, which I think is the cleanest outcome.
##########
ballista/executor/src/client_pool.rs:
##########
@@ -290,8 +364,32 @@ impl BallistaClientPool for DefaultBallistaClientPool {
.await?;
let shared = Arc::new(SharedClient::new(client));
- *slot.connection.lock().unwrap() = Some(shared.clone());
- Ok(shared)
+ {
+ let inner = Arc::downgrade(&self.inner);
+ let slot_ref = Arc::downgrade(&slot);
+ let me = Arc::downgrade(&shared);
+ let hook_key = key.clone();
+ shared.set_discard_hook(Box::new(move || {
Review Comment:
**Blocking:** I'd like to see the discard hook go. Eviction already drops a
discarded slot on its next pass (`survives_eviction` keeps it only while a
connect is in flight), so the hook only buys freeing it up to 15s earlier.
It also reopens the burst-of-callers race the `connecting` lock was there to
prevent. Once the hook removes the slot, callers already inside `shared()`
reconnect on the old slot while new callers get a fresh one from `or_default()`
and connect again. So a failing peer, which is exactly when a burst of retries
happens, can end up with two connections. Dropping the hook removes that race
and the public API below.
##########
ballista/executor/src/client_pool.rs:
##########
@@ -447,8 +545,8 @@ mod tests {
));
}
- /// Eviction drops the slot of a discarded shared connection, so slots for
- /// executors that have gone away do not pile up, and keeps a live one.
+ /// Discarding a shared connection drops its slot right away; eviction
+ /// keeps the other, live one.
#[tokio::test]
async fn evict_idle_drops_discarded_shared_connections() {
Review Comment:
With the hook in place this test no longer exercises eviction removing a
discarded slot, since the hook removes it first, and nothing else covers that
path. Dropping the hook puts this test back to testing what its name says.
##########
ballista/executor/src/client_pool.rs:
##########
@@ -66,24 +68,59 @@ type IdleMap = DashMap<PoolKey, VecDeque<IdleEntry>>;
#[derive(Default)]
struct SharedSlot {
connecting: tokio::sync::Mutex<()>,
- connection: Mutex<Option<Arc<SharedClient>>>,
+ state: Mutex<SharedSlotState>,
+}
+
+#[derive(Default)]
+struct SharedSlotState {
+ connection: Option<Arc<SharedClient>>,
+ /// When the connection was last handed out by `shared`; eviction uses it
+ /// to tell a connection nobody asks for any more from an active one.
+ last_handout: Option<Instant>,
}
impl SharedSlot {
- /// The open connection, unless there is none or it was discarded.
- fn live(&self) -> Option<Arc<SharedClient>> {
- self.connection
+ /// The open connection for a new caller, unless there is none or it was
+ /// discarded, stamped as requested just now.
+ fn handout(&self) -> Option<Arc<SharedClient>> {
+ let mut state = self.state.lock().unwrap();
+ let client = state
+ .connection
+ .as_ref()
+ .filter(|client| !client.is_discarded())
+ .cloned()?;
+ state.last_handout = Some(Instant::now());
+ Some(client)
+ }
+
+ /// Whether `client` is the connection this slot currently holds. A
+ /// discard from a connection the slot has since replaced must not drop
+ /// the slot.
+ fn holds(&self, client: &Arc<SharedClient>) -> bool {
+ self.state
.lock()
.unwrap()
+ .connection
.as_ref()
- .filter(|client| !client.is_discarded())
- .cloned()
+ .is_some_and(|current| Arc::ptr_eq(current, client))
}
- /// Whether eviction can drop the slot: it holds no live connection and
- /// none is being opened.
- fn is_dead(&self) -> bool {
- self.live().is_none() && self.connecting.try_lock().is_ok()
+ /// Whether the slot survives an eviction pass: without a live
+ /// connection, only while a connect is in flight on it; with one, until
+ /// it has gone unrequested for `idle_timeout`.
+ fn survives_eviction(&self, now: Instant, idle_timeout: Duration) -> bool {
+ let state = self.state.lock().unwrap();
+ let live = state
+ .connection
+ .as_ref()
+ .is_some_and(|client| !client.is_discarded());
+ if !live {
+ return self.connecting.try_lock().is_err();
+ }
+ match state.last_handout {
+ Some(last) => now.duration_since(last) < idle_timeout,
+ None => true,
Review Comment:
Nit: only test-built slots reach this `None => true` arm, and it makes a
live connection immune to idle eviction. Making `last_handout` non-optional and
always set together with `connection` would remove the never-expires case.
##########
ballista/executor/src/client_pool.rs:
##########
@@ -477,6 +575,301 @@ mod tests {
));
}
+ /// A discarded connection's slot leaves the map at once, without waiting
+ /// for an eviction pass, and the next caller reconnects.
+ #[tokio::test]
+ async fn discard_removes_shared_slot_eagerly() {
+ let port = serve_flight().await;
+ let pool = make_pool(Duration::from_secs(60));
+ let config = GrpcClientConfig::default();
+
+ let conn = pool.shared("127.0.0.1", port, &config,
None).await.unwrap();
+ assert_eq!(pool.inner.shared.len(), 1);
+
+ conn.discard();
+ assert_eq!(pool.inner.shared.len(), 0);
+
+ let next = pool.shared("127.0.0.1", port, &config,
None).await.unwrap();
+ assert!(!next.is_discarded());
+ }
+
+ /// A discard that loses the race against a reconnect onto the same slot
+ /// must not drop the slot's fresh connection.
+ #[tokio::test]
+ async fn discard_does_not_drop_slot_that_reconnected_first() {
+ let port = serve_flight().await;
+ let pool = make_pool(Duration::from_secs(60));
+ let config = GrpcClientConfig::default();
+ let key = ("127.0.0.1".to_string(), port, config.clone());
+
+ let old = pool.shared("127.0.0.1", port, &config, None).await.unwrap();
+ let slot = pool.inner.shared.get(&key).unwrap().clone();
+
+ // The reconnect lands first: when the hook runs, the slot already
+ // holds the replacement connection.
+ let replacement =
Arc::new(SharedClient::new(BallistaClient::new_for_test(
+ "127.0.0.1",
+ port,
+ )));
+ slot.state.lock().unwrap().connection = Some(replacement.clone());
+
+ old.discard();
+ assert_eq!(pool.inner.shared.len(), 1);
+ assert!(slot.holds(&replacement));
+ }
+
+ /// A shared connection nobody requests any more is dropped by eviction
+ /// once it has sat unrequested for `idle_timeout`; one with a recent
+ /// handout survives.
+ #[tokio::test]
+ async fn evict_idle_drops_stale_shared_connection() {
+ let port = serve_flight().await;
+ let timeout = Duration::from_millis(100);
Review Comment:
Nit: with a 100ms timeout, the second half of this test expects a
just-handed-out connection to survive `evict_idle`. On a loaded CI runner the
connect plus scheduling could take longer than that and flake. A longer timeout
with only the stale entry back-dated would be safer.
##########
ballista/executor/src/client_pool.rs:
##########
@@ -259,17 +333,17 @@ impl BallistaClientPool for DefaultBallistaClientPool {
customize_endpoint: Option<Arc<BallistaConfigGrpcEndpoint>>,
) -> Result<Arc<SharedClient>> {
let key = (host.to_string(), port, config.clone());
- let slot = self.inner.shared.entry(key).or_default().clone();
+ let slot = self.inner.shared.entry(key.clone()).or_default().clone();
- if let Some(client) = slot.live() {
+ if let Some(client) = slot.handout() {
Review Comment:
Nit: this fast path has the same evicted-slot race the connect path now
handles. If eviction drops the slot between the `entry()` clone above and this
`handout()`, the caller gets a connection the pool no longer tracks, and the
next caller opens a duplicate. Harmless, but it might be worth routing this
through the same reattach check.
##########
ballista/executor/src/client_pool.rs:
##########
@@ -66,24 +68,59 @@ type IdleMap = DashMap<PoolKey, VecDeque<IdleEntry>>;
#[derive(Default)]
struct SharedSlot {
connecting: tokio::sync::Mutex<()>,
- connection: Mutex<Option<Arc<SharedClient>>>,
+ state: Mutex<SharedSlotState>,
+}
+
+#[derive(Default)]
+struct SharedSlotState {
+ connection: Option<Arc<SharedClient>>,
+ /// When the connection was last handed out by `shared`; eviction uses it
+ /// to tell a connection nobody asks for any more from an active one.
+ last_handout: Option<Instant>,
}
impl SharedSlot {
- /// The open connection, unless there is none or it was discarded.
- fn live(&self) -> Option<Arc<SharedClient>> {
- self.connection
+ /// The open connection for a new caller, unless there is none or it was
+ /// discarded, stamped as requested just now.
+ fn handout(&self) -> Option<Arc<SharedClient>> {
+ let mut state = self.state.lock().unwrap();
+ let client = state
+ .connection
+ .as_ref()
+ .filter(|client| !client.is_discarded())
+ .cloned()?;
+ state.last_handout = Some(Instant::now());
+ Some(client)
+ }
+
+ /// Whether `client` is the connection this slot currently holds. A
+ /// discard from a connection the slot has since replaced must not drop
+ /// the slot.
+ fn holds(&self, client: &Arc<SharedClient>) -> bool {
+ self.state
.lock()
.unwrap()
+ .connection
.as_ref()
- .filter(|client| !client.is_discarded())
- .cloned()
+ .is_some_and(|current| Arc::ptr_eq(current, client))
}
- /// Whether eviction can drop the slot: it holds no live connection and
- /// none is being opened.
- fn is_dead(&self) -> bool {
- self.live().is_none() && self.connecting.try_lock().is_ok()
+ /// Whether the slot survives an eviction pass: without a live
+ /// connection, only while a connect is in flight on it; with one, until
+ /// it has gone unrequested for `idle_timeout`.
+ fn survives_eviction(&self, now: Instant, idle_timeout: Duration) -> bool {
Review Comment:
**Question:** retiring on time since last handout means a pipelined stream
that runs longer than `client_ttl` loses its slot, and the next open to that
peer creates a second connection while the first is still busy. Did you
consider having the stream hold a small usage guard so the pool can see the
connection is still in use? That would get around the `strong_count` problem
you described. Not a blocker if the overlap is acceptable, but it'd be good to
call out in the docs.
--
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]
