Pool checkout tweak (#615)

Author: levkk

cli.sh

@@ -22,11 +22,18 @@ function bench_init() {
     PGPASSWORD=pgdog pgbench -h 127.0.0.1 -p 6432 -U pgdog pgdog -i
 }
 
+function psql_cmd() {
+    PGPASSWORD=pgdog psql -h 127.0.0.1 -p 6432 -U pgdog pgdog
+}
+
 # Parse command
 case "$1" in
     admin)
         admin
         ;;
+    psql)
+        psql_cmd
+        ;;
     binit)
         bench_init
         ;;

pgdog/src/backend/pool/healthcheck.rs

@@ -49,16 +49,16 @@ impl<'a> Healtcheck<'a> {
 
     /// Perform the healtcheck if it's required.
     pub async fn healthcheck(&mut self) -> Result<(), Error> {
-        let healtcheck_age = self.conn.healthcheck_age(self.now);
+        let health_check_age = self.conn.healthcheck_age(self.now);
 
-        if healtcheck_age < self.healthcheck_interval {
+        if health_check_age < self.healthcheck_interval {
             return Ok(());
         }
 
         match timeout(self.healthcheck_timeout, self.conn.healthcheck(";")).await {
             Ok(Ok(())) => Ok(()),
             Ok(Err(err)) => {
-                error!("healthcheck server error: {} [{}]", err, self.pool.addr());
+                error!("health check server error: {} [{}]", err, self.pool.addr());
                 Err(Error::HealthcheckError)
             }
             Err(_) => Err(Error::HealthcheckError),

pgdog/src/backend/pool/pool_impl.rs

@@ -6,7 +6,7 @@ use std::time::Duration;
 
 use once_cell::sync::{Lazy, OnceCell};
 use parking_lot::{lock_api::MutexGuard, Mutex, RawMutex};
-use tokio::time::Instant;
+use tokio::time::{timeout, Instant};
 use tracing::error;
 
 use crate::backend::{Server, ServerOptions};
@@ -96,60 +96,78 @@ impl Pool {
     }
 
     pub async fn get(&self, request: &Request) -> Result<Guard, Error> {
-        self.get_internal(request).await
+        match timeout(self.config().checkout_timeout, self.get_internal(request)).await {
+            Ok(Ok(conn)) => Ok(conn),
+            Err(_) => {
+                self.inner.health.toggle(false);
+                Err(Error::CheckoutTimeout)
+            }
+            Ok(Err(err)) => {
+                self.inner.health.toggle(false);
+                Err(err.into())
+            }
+        }
     }
 
     /// Get a connection from the pool.
     async fn get_internal(&self, request: &Request) -> Result<Guard, Error> {
-        let pool = self.clone();
-
-        // Fast path, idle connection probably available.
-        let (server, granted_at, paused) = {
-            // Ask for time before we acquire the lock
-            // and only if we actually waited for a connection.
-            let granted_at = request.created_at;
-            let elapsed = granted_at.saturating_duration_since(request.created_at);
-            let mut guard = self.lock();
+        loop {
+            let pool = self.clone();
 
-            if !guard.online {
-                return Err(Error::Offline);
-            }
+            // Fast path, idle connection probably available.
+            let (server, granted_at, paused) = {
+                // Ask for time before we acquire the lock
+                // and only if we actually waited for a connection.
+                let granted_at = request.created_at;
+                let elapsed = granted_at.saturating_duration_since(request.created_at);
+                let mut guard = self.lock();
 
-            let conn = guard.take(request);
+                if !guard.online {
+                    return Err(Error::Offline);
+                }
 
-            if conn.is_some() {
-                guard.stats.counts.wait_time += elapsed;
-                guard.stats.counts.server_assignment_count += 1;
-            }
+                let conn = guard.take(request);
 
-            (conn, granted_at, guard.paused)
-        };
+                if conn.is_some() {
+                    guard.stats.counts.wait_time += elapsed;
+                    guard.stats.counts.server_assignment_count += 1;
+                }
 
-        if paused {
-            self.comms().ready.notified().await;
-        }
+                (conn, granted_at, guard.paused)
+            };
 
-        let (server, granted_at) = if let Some(mut server) = server {
-            server
-                .prepared_statements_mut()
-                .set_capacity(self.inner.config.prepared_statements_limit);
-            server.set_pooler_mode(self.inner.config.pooler_mode);
-            (Guard::new(pool, server, granted_at), granted_at)
-        } else {
-            // Slow path, pool is empty, will create new connection
-            // or wait for one to be returned if the pool is maxed out.
-            let mut waiting = Waiting::new(pool, request)?;
-            waiting.wait().await?
-        };
+            if paused {
+                self.comms().ready.notified().await;
+            }
 
-        return self
-            .maybe_healthcheck(
-                server,
-                self.inner.config.healthcheck_timeout,
-                self.inner.config.healthcheck_interval,
-                granted_at,
-            )
-            .await;
+            let (server, granted_at) = if let Some(mut server) = server {
+                server
+                    .prepared_statements_mut()
+                    .set_capacity(self.inner.config.prepared_statements_limit);
+                server.set_pooler_mode(self.inner.config.pooler_mode);
+                (Guard::new(pool, server, granted_at), granted_at)
+            } else {
+                // Slow path, pool is empty, will create new connection
+                // or wait for one to be returned if the pool is maxed out.
+                let mut waiting = Waiting::new(pool, request)?;
+                waiting.wait().await?
+            };
+
+            match self
+                .maybe_healthcheck(
+                    server,
+                    self.inner.config.healthcheck_timeout,
+                    self.inner.config.healthcheck_interval,
+                    granted_at,
+                )
+                .await
+            {
+                Ok(conn) => return Ok(conn),
+                // Try another connection.
+                Err(Error::HealthcheckError) => continue,
+                Err(err) => return Err(err),
+            }
+        }
     }
 
     /// Get server parameters, fetch them if necessary.

pgdog/src/backend/pool/test/mod.rs

@@ -513,6 +513,33 @@ async fn test_idle_healthcheck_loop() {
     );
 }
 
+#[tokio::test]
+async fn test_checkout_timeout() {
+    crate::logger();
+
+    let config = Config {
+        max: 1,
+        min: 1,
+        checkout_timeout: Duration::from_millis(100),
+        ..Default::default()
+    };
+
+    let pool = Pool::new(&PoolConfig {
+        address: Address::new_test(),
+        config,
+    });
+    pool.launch();
+
+    // Hold the only connection
+    let _conn = pool.get(&Request::default()).await.unwrap();
+
+    // Try to get another connection - should timeout
+    let result = pool.get(&Request::default()).await;
+    assert!(result.is_err());
+    assert_eq!(result.unwrap_err(), Error::CheckoutTimeout);
+    assert!(pool.lock().waiting.is_empty());
+}
+
 #[tokio::test]
 async fn test_move_conns_to() {
     crate::logger();

pgdog/src/backend/pool/waiting.rs

@@ -1,12 +1,7 @@
-use std::ops::Deref;
-
 use crate::backend::Server;
 
 use super::{Error, Guard, Pool, Request};
-use tokio::{
-    sync::oneshot::*,
-    time::{timeout, Instant},
-};
+use tokio::{sync::oneshot::*, time::Instant};
 
 pub(super) struct Waiting {
     pool: Pool,
@@ -24,6 +19,10 @@ impl Drop for Waiting {
 }
 
 impl Waiting {
+    /// Create new waiter.
+    ///
+    /// N.B. You must call and await `Waiting::wait`, otherwise you'll leak waiters.
+    ///
     pub(super) fn new(pool: Pool, request: &Request) -> Result<Self, Error> {
         let request = *request;
         let (tx, rx) = channel();
@@ -49,68 +48,27 @@ impl Waiting {
 
     /// Wait for connection from the pool.
     pub(super) async fn wait(&mut self) -> Result<(Guard, Instant), Error> {
-        let checkout_timeout = self.pool.inner().config.checkout_timeout;
         let rx = self.rx.take().expect("waiter rx taken");
 
-        // Make this cancellation-safe.
-        let mut wait_guard = WaitGuard::new(self);
-        let server = timeout(checkout_timeout, rx).await;
-        wait_guard.disarm();
-        drop(wait_guard);
+        // Can be cancelled. Drop will remove the waiter from the queue.
+        let server = rx.await;
+
+        // Disarm the guard. We can't be cancelled beyond this point.
+        self.waiting = false;
 
         let now = Instant::now();
         match server {
-            Ok(Ok(server)) => {
+            Ok(server) => {
                 let server = server?;
                 Ok((Guard::new(self.pool.clone(), server, now), now))
             }
 
-            Err(_err) => {
-                let mut guard = self.pool.lock();
-                guard.remove_waiter(&self.request.id);
-                self.pool.inner().health.toggle(false);
+            Err(_) => {
+                // Should not be possible.
+                // This means someone else removed my waiter from the wait queue,
+                // indicating a bug in the pool.
                 Err(Error::CheckoutTimeout)
             }
-
-            // Should not be possible.
-            // This means someone else removed my waiter from the wait queue,
-            // indicating a bug in the pool.
-            Ok(Err(_)) => Err(Error::CheckoutTimeout),
-        }
-    }
-}
-
-struct WaitGuard<'a> {
-    waiting: &'a Waiting,
-    armed: bool,
-}
-
-impl<'a> Deref for WaitGuard<'a> {
-    type Target = &'a Waiting;
-
-    fn deref(&self) -> &Self::Target {
-        &self.waiting
-    }
-}
-
-impl<'a> WaitGuard<'a> {
-    fn new(waiting: &'a Waiting) -> Self {
-        Self {
-            waiting,
-            armed: true,
-        }
-    }
-
-    fn disarm(&mut self) {
-        self.armed = false;
-    }
-}
-
-impl Drop for WaitGuard<'_> {
-    fn drop(&mut self) {
-        if self.armed {
-            let id = self.waiting.request.id;
-            self.waiting.pool.lock().remove_waiter(&id);
         }
     }
 }
@@ -126,7 +84,7 @@ mod tests {
     use super::*;
     use crate::backend::pool::Pool;
     use crate::net::messages::BackendKeyData;
-    use tokio::time::{sleep, Duration};
+    use tokio::time::{sleep, timeout, Duration};
 
     #[tokio::test]
     async fn test_cancellation_safety() {
@@ -200,12 +158,14 @@ mod tests {
         let _conn = pool.get(&Request::default()).await.unwrap();
 
         let request = Request::new(BackendKeyData::new());
-        let mut waiting = Waiting::new(pool.clone(), &request).unwrap();
-
-        let result = waiting.wait().await;
+        let waiter_pool = pool.clone();
+        let get_conn = async move {
+            let mut waiting = Waiting::new(waiter_pool.clone(), &request).unwrap();
+            waiting.wait().await
+        };
+        let result = timeout(Duration::from_millis(100), get_conn).await;
 
         assert!(result.is_err());
-        assert!(matches!(result.unwrap_err(), Error::CheckoutTimeout));
 
         let pool_guard = pool.lock();
         assert!(

pgdog/src/backend/server.rs

@@ -653,6 +653,7 @@ impl Server {
         debug!("running healthcheck \"{}\" [{}]", query, self.addr);
 
         self.execute(query).await?;
+
         self.stats.healthcheck();
 
         Ok(())