refactor(pool): simplify idle task using async/await (#812)

Author: 0x676e67

src/core/client/mod.rs

@@ -649,7 +649,7 @@ where
                                     // Errors are sent via err_tx to ensure they can be checked if the sender (tx) fails.
                                     executor.execute(
                                         conn.with_upgrades()
-                                                                        .map_err(|e| {
+                                                .map_err(|e| {
                                                 // Log the connection error at debug level for diagnostic purposes.
                                                 debug!("client connection error: {:?}", e);
                                                 // Log that the error is being sent to the error channel.

src/core/client/pool.rs

@@ -23,7 +23,7 @@ use crate::{
             exec::{self, Exec},
             timer::Timer,
         },
-        rt::{Sleep, Timer as _},
+        rt::Timer as _,
     },
     sync::Mutex,
 };
@@ -415,13 +415,11 @@ impl<T: Poolable, K: Key> PoolInner<T, K> {
         let interval = IdleTask {
             timer: timer.clone(),
             duration: dur,
-            deadline: Instant::now(),
-            fut: timer.sleep_until(Instant::now()), // ready at first tick
             pool: WeakOpt::downgrade(pool_ref),
             pool_drop_notifier: rx,
         };
 
-        self.exec.execute(interval);
+        self.exec.execute(interval.run());
     }
 }
 
@@ -744,55 +742,43 @@ impl Expiration {
     }
 }
 
-pin_project_lite::pin_project! {
-    struct IdleTask<T, K: Key> {
-        timer: Timer,
-        duration: Duration,
-        deadline: Instant,
-        fut: Pin<Box<dyn Sleep>>,
-        pool: WeakOpt<Mutex<PoolInner<T, K>>>,
-        // This allows the IdleTask to be notified as soon as the entire
-        // Pool is fully dropped, and shutdown. This channel is never sent on,
-        // but Err(Canceled) will be received when the Pool is dropped.
-        #[pin]
-        pool_drop_notifier: oneshot::Receiver<Infallible>,
-    }
+struct IdleTask<T, K: Key> {
+    timer: Timer,
+    duration: Duration,
+    pool: WeakOpt<Mutex<PoolInner<T, K>>>,
+    // This allows the IdleTask to be notified as soon as the entire
+    // Pool is fully dropped, and shutdown. This channel is never sent on,
+    // but Err(Canceled) will be received when the Pool is dropped.
+    pool_drop_notifier: oneshot::Receiver<Infallible>,
 }
 
-impl<T: Poolable + 'static, K: Key> Future for IdleTask<T, K> {
-    type Output = ();
+impl<T: Poolable + 'static, K: Key> IdleTask<T, K> {
+    async fn run(self) {
+        use futures_util::future;
 
-    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
-        let mut this = self.project();
+        let mut sleep = self.timer.sleep_until(Instant::now() + self.duration);
+        let mut on_pool_drop = self.pool_drop_notifier;
         loop {
-            match this.pool_drop_notifier.as_mut().poll(cx) {
-                Poll::Ready(Ok(n)) => match n {},
-                Poll::Pending => (),
-                Poll::Ready(Err(_canceled)) => {
-                    trace!("pool closed, canceling idle interval");
-                    return Poll::Ready(());
+            match future::select(&mut on_pool_drop, &mut sleep).await {
+                future::Either::Left(_) => {
+                    // pool dropped, bah-bye
+                    break;
                 }
-            }
-
-            ready!(Pin::new(&mut this.fut).poll(cx));
-            // Set this task to run after the next deadline
-            // If the poll missed the deadline by a lot, set the deadline
-            // from the current time instead
-            *this.deadline += *this.duration;
-            if *this.deadline < Instant::now() - Duration::from_millis(5) {
-                *this.deadline = Instant::now() + *this.duration;
-            }
-            *this.fut = this.timer.sleep_until(*this.deadline);
+                future::Either::Right(((), _)) => {
+                    if let Some(inner) = self.pool.upgrade() {
+                        let mut inner = inner.lock();
+                        trace!("idle interval checking for expired");
+                        inner.clear_expired();
+                        drop(inner);
+                    }
 
-            if let Some(inner) = this.pool.upgrade() {
-                let mut inner = inner.lock();
-                inner.clear_expired();
-                drop(inner);
-                trace!("idle interval checking for expired");
-                continue;
+                    let deadline = Instant::now() + self.duration;
+                    self.timer.reset(&mut sleep, deadline);
+                }
             }
-            return Poll::Ready(());
         }
+
+        trace!("pool closed, canceling idle interval");
     }
 }