归还、Drop 与维护任务

pub struct PoolConnection<DB: Database> {
    live: Option<Live<DB>>,
    close_on_drop: bool,
    pub(crate) pool: Arc<PoolInner<DB>>,
}

pub(super) struct Live<DB: Database> {
    pub(super) raw: DB::Connection,
    pub(super) created_at: Instant,
}

pub(super) struct Idle<DB: Database> {
    pub(super) live: Live<DB>,
    pub(super) idle_since: Instant,
}

/// RAII wrapper for connections being handled by functions that may drop them
pub(super) struct Floating<DB: Database, C> {
    pub(super) inner: C,
    pub(super) guard: DecrementSizeGuard<DB>,
}

impl<DB: Database> Drop for PoolConnection<DB> {
    fn drop(&mut self) {
        if self.close_on_drop {
            crate::rt::spawn(self.take_and_close());
            return;
        }

        // We still need to spawn a task to maintain `min_connections`.
        if self.live.is_some() || self.pool.options.min_connections > 0 {
            crate::rt::spawn(self.return_to_pool());
        }
    }
}

    pub fn return_to_pool(&mut self) -> impl Future<Output = ()> + Send + 'static {
        // float the connection in the pool before we move into the task
        // in case the returned `Future` isn't executed, like if it's spawned into a dying runtime
        // https://github.com/launchbadge/sqlx/issues/1396
        // Type hints seem to be broken by `Option` combinators in IntelliJ Rust right now (6/22).
        let floating: Option<Floating<DB, Live<DB>>> =
            self.live.take().map(|live| live.float(self.pool.clone()));

        let pool = self.pool.clone();

        async move {
            let returned_to_pool = if let Some(floating) = floating {
                floating.return_to_pool().await
            } else {
                false
            };

            if !returned_to_pool {
                pool.min_connections_maintenance(None).await;
            }
        }
    }

#[tokio::main(flavor = "current_thread")]
async fn main() -> anyhow::Result<()> {
  // ...
  let pool = PgPoolOptions::new()
    .max_connections(1)
    .connect_with(lab::opts("pool-ct"))
    .await?;

  println!("--- drop, then try_acquire right away ---");
  let conn = pool.acquire().await?;
  drop(conn);
  println!("try_acquire: {:?}", pool.try_acquire().map(|_| ()));
  println!("size={} num_idle={}", pool.size(), pool.num_idle());

  println!("--- drop, then acquire().await ---");
  let started = Instant::now();
  let conn = pool.acquire().await?;
  println!("acquire took {:?}", started.elapsed());
  drop(conn);
  // ...
}

--- drop, then try_acquire right away ---
try_acquire: None
size=1 num_idle=0
--- drop, then acquire().await ---
acquire took 593.625µs
...

    async fn return_to_pool(mut self) -> bool {
        // Immediately close the connection.
        if self.guard.pool.is_closed() {
            self.close().await;
            return false;
        }

        // If the connection is beyond max lifetime, close the connection and
        // immediately create a new connection
        if is_beyond_max_lifetime(&self.inner, &self.guard.pool.options) {
            self.close().await;
            return false;
        }

        if let Some(test) = &self.guard.pool.options.after_release {
            // ... Ok(true) 继续；Ok(false) 正常 close；Err 打 warn 后 close_hard
        }

        // test the connection on-release to ensure it is still viable,
        // and flush anything time-sensitive like transaction rollbacks
        // if an Executor future/stream is dropped during an `.await` call, the connection
        // is likely to be left in an inconsistent state, in which case it should not be
        // returned to the pool; also of course, if it was dropped due to an error
        // this is simply a band-aid as SQLx-next connections should be able
        // to recover from cancellations
        if let Err(error) = self.raw.ping().await {
            tracing::warn!(
                %error,
                "error occurred while testing the connection on-release",
            );

            // Connection is broken, don't try to gracefully close.
            self.close_hard().await;
            false
        } else {
            // if the connection is still viable, release it to the pool
            self.release();
            true
        }
    }

    pub(super) fn release(&self, floating: Floating<DB, Live<DB>>) {
        // `options.after_release` and other checks are in `PoolConnection::return_to_pool()`.

        let Floating { inner: idle, guard } = floating.into_idle();

        if self.idle_conns.push(idle).is_err() {
            panic!("BUG: connection queue overflow in release()");
        }

        // NOTE: we need to make sure we drop the permit *after* we push to the idle queue
        // don't decrease the size
        guard.release_permit();

        self.num_idle.fetch_add(1, Ordering::AcqRel);
    }

println!("--- close_on_drop ---");
p.quiet.store(true, Ordering::SeqCst);
let pl = pool(p.port, "pool-cod").await?;
let mut conn = pl.acquire().await?;
p.quiet.store(false, Ordering::SeqCst);
conn.close_on_drop();
drop(conn);
println!("dropped: size={}", pl.size());
settle().await;
println!("later:   size={} backends={}", pl.size(), lab::backends(&mut obs, "pool-cod").await.len());
p.quiet.store(true, Ordering::SeqCst);
pl.close().await;

println!("--- detach ---");
let pl = pool(p.port, "pool-detach").await?;
let conn = pl.acquire().await?;
let mut raw = conn.detach();
println!("detached: size={} num_idle={}", pl.size(), pl.num_idle());
let other = pl.acquire().await?;
println!("acquired another: size={} backends={}", pl.size(), lab::backends(&mut obs, "pool-detach").await.len());
let one: i32 = sqlx::query_scalar("select 1").fetch_one(&mut raw).await?;
println!("detached connection still works: {one}");
drop(other);
raw.close().await?;
pl.close().await;

println!("--- leak ---");
let pl = pool(p.port, "pool-leak").await?;
let conn = pl.acquire().await?;
let raw = conn.leak();
settle().await;
println!("leaked: size={} num_idle={}", pl.size(), pl.num_idle());
println!("acquire: {:?}", pl.acquire().await.err());
let closed = tokio::time::timeout(Duration::from_secs(1), pl.close()).await;
println!("pool.close() within 1s: {}", if closed.is_ok() { "done" } else { "still waiting" });
raw.close().await?;

--- close_on_drop ---
dropped: size=1
  [conn 1] C->S Terminate
later:   size=0 backends=0
--- detach ---
detached: size=0 num_idle=0
acquired another: size=1 backends=2
detached connection still works: 1
--- leak ---
leaked: size=1 num_idle=0
acquire: Some(PoolTimedOut)
pool.close() within 1s: still waiting

fn spawn_maintenance_tasks<DB: Database>(pool: &Arc<PoolInner<DB>>) {
    // NOTE: use `pool_weak` for the maintenance tasks
    // so they don't keep `PoolInner` from being dropped.
    let pool_weak = Arc::downgrade(pool);

    let period = match (pool.options.max_lifetime, pool.options.idle_timeout) {
        (Some(it), None) | (None, Some(it)) => it,

        (Some(a), Some(b)) => cmp::min(a, b),

        (None, None) => {
            if pool.options.min_connections > 0 {
                crate::rt::spawn(async move {
                    if let Some(pool) = pool_weak.upgrade() {
                        pool.min_connections_maintenance(None).await;
                    }
                });
            }

            return;
        }
    };

    // Immediately cancel this task if the pool is closed.
    let mut close_event = pool.close_event();

    crate::rt::spawn(async move {
        let _ = close_event
            .do_until(async {
                // If the last handle to the pool was dropped while we were sleeping
                while let Some(pool) = pool_weak.upgrade() {
                    if pool.is_closed() {
                        return;
                    }

                    let next_run = Instant::now() + period;

                    // Go over all idle connections, check for idleness and lifetime,
                    // and if we have fewer than min_connections after reaping a connection,
                    // open a new one immediately. Note that other connections may be popped from
                    // the queue in the meantime - that's fine, there is no harm in checking more
                    for _ in 0..pool.num_idle() {
                        if let Some(conn) = pool.try_acquire() {
                            if is_beyond_idle_timeout(&conn, &pool.options)
                                || is_beyond_max_lifetime(&conn, &pool.options)
                            {
                                let _ = conn.close().await;
                                pool.min_connections_maintenance(Some(next_run)).await;
                            } else {
                                pool.release(conn.into_live());
                            }
                        }
                    }

                    // Don't hold a reference to the pool while sleeping.
                    drop(pool);

                    // ... sleep 到 next_run；已经过了就 yield_now
                }
            })
            .await;
    });
}

let pool = PgPoolOptions::new()
  .min_connections(min)
  .idle_timeout(Duration::from_secs(1))
  .max_lifetime(None)
  .connect_with(lab::opts(app))
  .await?;
// check out three at once so the pool has to open three
let (a, b, c) = (pool.acquire().await?, pool.acquire().await?, pool.acquire().await?);
drop((a, b, c));
tokio::time::sleep(Duration::from_millis(50)).await;
for _ in 0..(if min == 0 { 10 } else { 18 }) {
  println!(
    "{:>5}ms size={} idle={} pids={:?}",
    clock.ms(),
    pool.size(),
    pool.num_idle(),
    lab::backends(&mut obs, app).await
  );
  tokio::time::sleep(Duration::from_millis(250)).await;
}

--- idle_timeout(1s), max_lifetime(None), min_connections(0) ---
  136ms size=3 idle=3 pids=[328, 329, 330]
  398ms size=3 idle=3 pids=[328, 329, 330]
  653ms size=3 idle=3 pids=[328, 329, 330]
  908ms size=3 idle=3 pids=[328, 329, 330]
 1164ms size=3 idle=3 pids=[328, 329, 330]
 1420ms size=3 idle=3 pids=[328, 329, 330]
 1674ms size=3 idle=3 pids=[328, 329, 330]
 1929ms size=3 idle=3 pids=[328, 329, 330]
 2183ms size=0 idle=0 pids=[]
 2437ms size=0 idle=0 pids=[]
--- idle_timeout(1s), max_lifetime(None), min_connections(1) ---
  146ms size=3 idle=3 pids=[334, 335, 336]
  406ms size=3 idle=3 pids=[334, 335, 336]
  659ms size=3 idle=3 pids=[334, 335, 336]
  913ms size=3 idle=3 pids=[334, 335, 336]
 1167ms size=3 idle=3 pids=[334, 335, 336]
 1420ms size=3 idle=3 pids=[334, 335, 336]
 1674ms size=3 idle=3 pids=[334, 335, 336]
 1930ms size=3 idle=3 pids=[334, 335, 336]
 2182ms size=1 idle=1 pids=[337]
 2437ms size=1 idle=1 pids=[337]
 2691ms size=1 idle=1 pids=[337]
 2945ms size=1 idle=1 pids=[337]
 3200ms size=1 idle=1 pids=[337]
 3455ms size=1 idle=1 pids=[337]
 3710ms size=1 idle=1 pids=[337]
 3965ms size=1 idle=1 pids=[337]
 4220ms size=1 idle=1 pids=[338]
 4475ms size=1 idle=1 pids=[338]

let pool = PgPoolOptions::new()
  .max_lifetime(Duration::from_secs(2))
  .idle_timeout(None)
  .connect_with(lab::opts(APP))
  .await?;
println!("{:>5}ms built: size={} pids={:?}", clock.ms(), pool.size(), lab::backends(&mut obs, APP).await);

tokio::time::sleep(Duration::from_millis(2600)).await;
let mut conn = pool.acquire().await?;
let pid: i32 = sqlx::query_scalar("select pg_backend_pid()").fetch_one(&mut *conn).await?;
println!("{:>5}ms acquired, pid={pid}", clock.ms());
drop(conn);
tokio::time::sleep(Duration::from_millis(50)).await;
println!("{:>5}ms after drop: size={} idle={} pids={:?}", clock.ms(), pool.size(), pool.num_idle(), lab::backends(&mut obs, APP).await);

// a fresh connection that just sits in the idle queue
drop(pool.acquire().await?);
tokio::time::sleep(Duration::from_millis(50)).await;
for _ in 0..8 {
  println!("{:>5}ms size={} idle={} pids={:?}", clock.ms(), pool.size(), pool.num_idle(), lab::backends(&mut obs, APP).await);
  tokio::time::sleep(Duration::from_millis(500)).await;
}

   27ms built: size=1 pids=[353]
 2639ms acquired, pid=353
 2692ms after drop: size=0 idle=0 pids=[]
 2773ms size=1 idle=1 pids=[356]
 3278ms size=1 idle=1 pids=[356]
 3782ms size=1 idle=1 pids=[356]
 4287ms size=1 idle=1 pids=[356]
 4793ms size=1 idle=1 pids=[356]
 5296ms size=1 idle=1 pids=[356]
 5800ms size=1 idle=1 pids=[356]
 6304ms size=0 idle=0 pids=[]

let pool = PgPoolOptions::new()
  .min_connections(3)
  .connect_with(lab::opts("pool-min-a"))
  .await?;
watch("connect               ", &pool, "pool-min-a").await;
pool.close().await;

let pool = PgPoolOptions::new()
  .min_connections(3)
  .connect_lazy_with(lab::opts("pool-min-b"));
watch("connect_lazy (default)", &pool, "pool-min-b").await;
pool.close().await;

let pool = PgPoolOptions::new()
  .min_connections(3)
  .idle_timeout(None)
  .max_lifetime(None)
  .connect_lazy_with(lab::opts("pool-min-c"));
watch("connect_lazy (no reap)", &pool, "pool-min-c").await;
pool.close().await;

// lazy pool with default timeouts: does using it once fill the rest?
let pool = PgPoolOptions::new()
  .min_connections(3)
  .connect_lazy_with(lab::opts("pool-min-d"));
drop(pool.acquire().await?);
watch("lazy + one acquire    ", &pool, "pool-min-d").await;
pool.close().await;

    pub(super) fn close(self: &Arc<Self>) -> impl Future<Output = ()> + '_ {
        self.mark_closed();

        async move {
            // For child pools, we need to acquire permits we actually have rather than
            // max_connections
            let permits_to_acquire = if self.options.parent_pool.is_some() {
                // Child pools start with 0 permits, so we acquire based on current size
                self.size()
            } else {
                // Parent pools can acquire all max_connections permits
                self.options.max_connections
            };

            let _permits = self.semaphore.acquire(permits_to_acquire).await;

            while let Some(idle) = self.idle_conns.pop() {
                let _ = idle.live.raw.close().await;
            }

            self.num_idle.store(0, Ordering::Release);
            self.size.store(0, Ordering::Release);
        }
    }

let pool = PgPoolOptions::new()
  .max_connections(2)
  .connect_with(lab::via_proxy("pool-close", p.port))
  .await?;
let held = pool.acquire().await?; // conn 1
drop(pool.acquire().await?); // conn 2, back in the idle queue
tokio::time::sleep(Duration::from_millis(100)).await;
let clock = Arc::new(lab::Clock::new());
println!("{:>4}ms size={} idle={}", clock.ms(), pool.size(), pool.num_idle());
p.quiet.store(false, Ordering::SeqCst);

let c = clock.clone();
let holder = tokio::spawn(async move {
  tokio::time::sleep(Duration::from_millis(300)).await;
  println!("{:>4}ms holder drops its connection", c.ms());
  drop(held);
});
let (pool2, c) = (pool.clone(), clock.clone());
let late = tokio::spawn(async move {
  tokio::time::sleep(Duration::from_millis(100)).await;
  println!("{:>4}ms acquire after close: {:?}", c.ms(), pool2.acquire().await.err());
});

println!("{:>4}ms calling close()", clock.ms());
let closing = pool.close();
println!("{:>4}ms is_closed={} (close future not polled yet)", clock.ms(), pool.is_closed());
closing.await;
println!("{:>4}ms close() returned: size={} idle={}", clock.ms(), pool.size(), pool.num_idle());

   0ms size=2 idle=1
   0ms calling close()
   0ms is_closed=true (close future not polled yet)
 103ms acquire after close: Some(PoolClosed)
 302ms holder drops its connection
 303ms close() returned: size=0 idle=0
  [conn 1] C->S Terminate
  [conn 2] C->S Terminate

child.close().await;
println!("child closed; parent permits still lent to the child: parent.close() within 1s: {}",
  if tokio::time::timeout(std::time::Duration::from_secs(1), parent.close()).await.is_ok() { "done" } else { "still waiting" });
drop(child);
parent.close().await;
println!("after drop(child): parent closed");

...
child closed; parent permits still lent to the child: parent.close() within 1s: still waiting
after drop(child): parent closed

connect                t=   0ms: size=3 num_idle=3 backends=3
connect                t= 500ms: size=3 num_idle=3 backends=3
connect                t=2000ms: size=3 num_idle=3 backends=3
connect_lazy (default) t=   0ms: size=0 num_idle=0 backends=0
connect_lazy (default) t= 500ms: size=0 num_idle=0 backends=0
connect_lazy (default) t=2000ms: size=0 num_idle=0 backends=0
connect_lazy (no reap) t=   0ms: size=1 num_idle=0 backends=1
connect_lazy (no reap) t= 500ms: size=3 num_idle=3 backends=3
connect_lazy (no reap) t=2000ms: size=3 num_idle=3 backends=3
lazy + one acquire     t=   0ms: size=1 num_idle=1 backends=1
lazy + one acquire     t= 500ms: size=1 num_idle=1 backends=1
lazy + one acquire     t=2000ms: size=1 num_idle=1 backends=1