andygrove commented on code in PR #2779:
URL:
https://github.com/apache/datafusion-ballista/pull/2779#discussion_r4486562618
##########
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,
Review Comment:
One thing I'd love a test for: once eviction retires a live slot, nothing
holds a `Channel` for that connection any more, since the reader drops its
`Arc` after `fetch_pipelined_partition` returns and `FlightDataStream` only
owns the tonic `Streaming`. So this relies on hyper keeping the h2 connection
up until its open streams finish. Could you add a test that opens a pipelined
stream, ages the handout and runs `evict_idle`, then keeps reading to the end?
That would pin down the behavior the PR description relies on.
##########
ballista/core/src/client_pool.rs:
##########
@@ -108,11 +113,21 @@ impl SharedClient {
self.client.clone()
}
+ /// Installs the hook the first [discard](Self::discard) runs. Discards
+ /// after the first do not run it again.
+ pub fn set_discard_hook(&self, hook: Box<dyn FnOnce() + Send + Sync>) {
Review Comment:
Since this is public and replaces any existing hook, maybe say in the doc
that it's meant for the pool that created the connection, so a caller doesn't
clobber the pool's hook by accident.
##########
ballista/core/src/client_pool.rs:
##########
@@ -108,11 +113,21 @@ impl SharedClient {
self.client.clone()
}
+ /// Installs the hook the first [discard](Self::discard) runs. Discards
+ /// after the first do not run it again.
+ pub fn set_discard_hook(&self, hook: Box<dyn FnOnce() + Send + Sync>) {
+ *self.on_discard.lock().unwrap() = Some(hook);
+ }
+
/// Marks the connection as gone, after a failure that suggests it is, so
/// that the pool opens a new one for the next caller. Callers already on
/// it keep it until they are done.
pub fn discard(&self) {
- self.discarded.store(true, Ordering::Release);
+ if !self.discarded.swap(true, Ordering::AcqRel)
+ && let Some(hook) = self.on_discard.lock().unwrap().take()
+ {
+ hook();
Review Comment:
Small one: the `on_discard` guard is still held while `hook()` runs, since
it lives through the `if let` body. Nothing re-enters today, but taking the
hook out first (`let hook = self.on_discard.lock().unwrap().take();` then call
it) would make that impossible.
##########
ballista/core/src/client_pool.rs:
##########
@@ -92,6 +92,10 @@ pub trait BallistaClientPool: Send + Sync + Debug {
pub struct SharedClient {
client: BallistaClient,
discarded: AtomicBool,
+ /// Runs once, on the first [discard](Self::discard). A pool uses it to
+ /// drop the connection's slot eagerly instead of waiting for the next
+ /// eviction pass.
+ on_discard: Mutex<Option<Box<dyn FnOnce() + Send + Sync>>>,
Review Comment:
The `+ Sync` on the hook type isn't needed since it's behind a `Mutex`, so
`Box<dyn FnOnce() + Send>` would be enough (same for `set_discard_hook`).
##########
ballista/executor/src/client_pool.rs:
##########
@@ -183,7 +219,40 @@ fn evict(inner: &Inner) {
}
!deque.is_empty()
});
- inner.shared.retain(|_, slot| !slot.is_dead());
+ inner
+ .shared
+ .retain(|_, slot| slot.survives_eviction(now, inner.idle_timeout));
+}
+
+/// Re-attach `slot`, carrying the connection `ours` just opened, to the map.
+/// Eviction can drop the slot while its `shared` caller was still
+/// connecting; without this the fresh connection would be invisible to the
+/// next caller, which would open another one. When a different slot has
+/// already reconnected, defer to its connection and let `ours` close on
+/// drop instead.
+fn reattach_shared_slot(
+ inner: &Inner,
+ key: &PoolKey,
+ slot: Arc<SharedSlot>,
+ ours: Arc<SharedClient>,
+) -> Arc<SharedClient> {
+ match inner.shared.entry(key.clone()) {
+ Entry::Occupied(mut e) => {
+ if Arc::ptr_eq(e.get(), &slot) {
+ ours
+ } else if let Some(other) = e.get().handout() {
Review Comment:
In the defer branch the doc comment says `ours` closes on drop, but
`_connecting` was already released, so callers waiting on the detached slot get
`ours` from `handout()` anyway. Harmless (one extra short-lived connection),
but maybe tweak the comment, or reattach before releasing the lock and store
`other` in the slot so they converge too.
--
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]
