GetConn: move queue_id up one level, remove duplicated code

Author: blackbeam

src/conn/pool/futures/get_conn.rs

@@ -26,7 +26,6 @@ use crate::{
 /// States of the GetConn future.
 pub(crate) enum GetConnInner {
     New,
-    Queued(QueueId),
     Done,
     // TODO: one day this should be an existential
     Connecting(crate::BoxFuture<'static, Conn>),
@@ -38,10 +37,6 @@ impl fmt::Debug for GetConnInner {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             GetConnInner::New => f.debug_tuple("GetConnInner::New").finish(),
-            GetConnInner::Queued(queue_id) => f
-                .debug_tuple("GetConnInner::Queued")
-                .field(queue_id)
-                .finish(),
             GetConnInner::Done => f.debug_tuple("GetConnInner::Done").finish(),
             GetConnInner::Connecting(_) => f
                 .debug_tuple("GetConnInner::Connecting")
@@ -66,13 +61,15 @@ impl GetConnInner {
 #[derive(Debug)]
 #[must_use = "futures do nothing unless you `.await` or poll them"]
 pub struct GetConn {
+    pub(crate) queue_id: Option<QueueId>,
     pub(crate) pool: Option<Pool>,
     pub(crate) inner: GetConnInner,
 }
 
 impl GetConn {
     pub(crate) fn new(pool: &Pool) -> GetConn {
         GetConn {
+            queue_id: None,
             pool: Some(pool.clone()),
             inner: GetConnInner::New,
         }
@@ -100,40 +97,11 @@ impl Future for GetConn {
         loop {
             match self.inner {
                 GetConnInner::New => {
-                    let queue_id = QueueId::next();
-                    match Pin::new(self.pool_mut()).poll_new_conn(cx, false, queue_id) {
-                        Poll::Pending => {
-                            self.inner = GetConnInner::Queued(queue_id);
-                            return Poll::Pending;
-                        }
-                        Poll::Ready(res) => match res?.inner.take() {
-                            GetConnInner::Connecting(conn_fut) => {
-                                self.inner = GetConnInner::Connecting(conn_fut);
-                            }
-                            GetConnInner::Checking(conn_fut) => {
-                                self.inner = GetConnInner::Checking(conn_fut);
-                            }
-                            GetConnInner::Done => unreachable!(
-                                "Pool::poll_new_conn never gives out already-consumed GetConns"
-                            ),
-                            GetConnInner::New => {
-                                unreachable!(
-                                    "Pool::poll_new_conn never gives out GetConnInner::New"
-                                )
-                            }
-                            GetConnInner::Queued(_) => {
-                                unreachable!(
-                                    "Pool::poll_new_conn never gives out GetConnInner::Queued"
-                                )
-                            }
-                        },
-                    }
-                }
-                GetConnInner::Queued(queue_id) => {
-                    match ready!(Pin::new(self.pool_mut()).poll_new_conn(cx, true, queue_id))?
-                        .inner
-                        .take()
-                    {
+                    let queued = self.queue_id.is_some();
+                    let queue_id = *self.queue_id.get_or_insert_with(|| QueueId::next());
+                    let next =
+                        ready!(Pin::new(self.pool_mut()).poll_new_conn(cx, queued, queue_id))?;
+                    match next {
                         GetConnInner::Connecting(conn_fut) => {
                             self.inner = GetConnInner::Connecting(conn_fut);
                         }
@@ -146,9 +114,6 @@ impl Future for GetConn {
                         GetConnInner::New => {
                             unreachable!("Pool::poll_new_conn never gives out GetConnInner::New")
                         }
-                        GetConnInner::Queued(_) => {
-                            unreachable!("Pool::poll_new_conn never gives out GetConnInner::Queued")
-                        }
                     }
                 }
                 GetConnInner::Done => {

src/conn/pool/mod.rs

@@ -273,7 +273,7 @@ impl Pool {
         cx: &mut Context<'_>,
         queued: bool,
         queue_id: QueueId,
-    ) -> Poll<Result<GetConn>> {
+    ) -> Poll<Result<GetConnInner>> {
         self.poll_new_conn_inner(cx, queued, queue_id)
     }
 
@@ -282,7 +282,7 @@ impl Pool {
         cx: &mut Context<'_>,
         queued: bool,
         queue_id: QueueId,
-    ) -> Poll<Result<GetConn>> {
+    ) -> Poll<Result<GetConnInner>> {
         let mut exchange = self.inner.exchange.lock().unwrap();
 
         // NOTE: this load must happen while we hold the lock,
@@ -304,16 +304,13 @@ impl Pool {
 
         while let Some(IdlingConn { mut conn, .. }) = exchange.available.pop_back() {
             if !conn.expired() {
-                return Poll::Ready(Ok(GetConn {
-                    pool: Some(self.clone()),
-                    inner: GetConnInner::Checking(
-                        async move {
-                            conn.stream_mut()?.check().await?;
-                            Ok(conn)
-                        }
-                        .boxed(),
-                    ),
-                }));
+                return Poll::Ready(Ok(GetConnInner::Checking(
+                    async move {
+                        conn.stream_mut()?.check().await?;
+                        Ok(conn)
+                    }
+                    .boxed(),
+                )));
             } else {
                 self.send_to_recycler(conn);
             }
@@ -325,10 +322,9 @@ impl Pool {
             // we are allowed to make a new connection, so we will!
             exchange.exist += 1;
 
-            return Poll::Ready(Ok(GetConn {
-                pool: Some(self.clone()),
-                inner: GetConnInner::Connecting(Conn::new(self.opts.clone()).boxed()),
-            }));
+            return Poll::Ready(Ok(GetConnInner::Connecting(
+                Conn::new(self.opts.clone()).boxed(),
+            )));
         }
 
         // Polled, but no conn available? Back into the queue.