Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 21 additions & 2 deletions src/conn/pool/futures/disconnect_pool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,13 @@ use std::{
task::{Context, Poll},
};

use futures_core::ready;
use tokio::sync::mpsc::UnboundedSender;

use crate::{
conn::pool::{Inner, Pool},
error::Error,
Conn,
};

use std::sync::{atomic, Arc};
Expand All @@ -27,20 +31,23 @@ use std::sync::{atomic, Arc};
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct DisconnectPool {
pool_inner: Arc<Inner>,
drop: Option<UnboundedSender<Option<Conn>>>,
}

impl DisconnectPool {
pub(crate) fn new(pool: Pool) -> Self {
Self {
pool_inner: pool.inner,
drop: Some(pool.drop),
}
}
}

impl Future for DisconnectPool {
type Output = Result<(), Error>;

fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.pool_inner.close.store(true, atomic::Ordering::Release);
let mut exchange = self.pool_inner.exchange.lock().unwrap();
exchange.spawn_futures_if_needed(&self.pool_inner);
exchange.waiting.push_back(cx.waker().clone());
Expand All @@ -49,7 +56,19 @@ impl Future for DisconnectPool {
if self.pool_inner.closed.load(atomic::Ordering::Acquire) {
Poll::Ready(Ok(()))
} else {
Poll::Pending
match self.drop.take() {
Some(drop) => match drop.send(None) {
Ok(_) => {
// Recycler is alive. Waiting for it to finish.
Poll::Ready(Ok(ready!(Box::pin(drop.closed()).as_mut().poll(cx))))
}
Err(_) => {
// Recycler seem dead. No one will wake us.
Poll::Ready(Ok(()))
}
},
None => Poll::Pending,
}
}
}
}
9 changes: 0 additions & 9 deletions src/conn/pool/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -162,15 +162,6 @@ impl Pool {
/// **Note:** This Future won't resolve until all active connections, taken from it,
/// are dropped or disonnected. Also all pending and new `GetConn`'s will resolve to error.
pub fn disconnect(self) -> DisconnectPool {
let was_closed = self.inner.close.swap(true, atomic::Ordering::AcqRel);
if !was_closed {
// make sure we wake up the Recycler.
//
// note the lack of an .expect() here, because the Recycler may decide that there are
// no connections to wait for and exit quickly!
let _ = self.drop.send(None).is_ok();
}

DisconnectPool::new(self)
}

Expand Down