PoolTimedOut error with OnceCell due to OnceCell dropped with the test which initialized it

[kingwingfly]

Bug Description

I put PgPool in a ModelManager, so that I can manage it conveniently. Then I init it using tokio::sync::OnceCell, so that all the test func use the same ModelManager.

And I have 2 tests, one acquire a PgPool connection, and another print something through a method on ModelManager.

Run test again and again, and about 1/5 tests will fail with error PoolTimedOut.

Minimal Reproduction

Find the the cause and fix solution here:
https://github.com/kingwingfly/share_runtime_example

For convenient:

#![allow(dead_code)]

use sqlx::postgres::{PgPool, PgPoolOptions};
use tokio::sync::OnceCell;

const PG_URL: &str = "postgres://postgres:postgres@localhost:5432/postgres";
static INIT: OnceCell<ModelManager> = OnceCell::const_new();

/// Init ModelManager for test, so that all the test func use the same ModelManager
async fn init_test() -> ModelManager {
    let mm = INIT
        .get_or_init(|| async {
            println!("INIT ModelManager");
            ModelManager::new().await
        })
        .await;
    mm.clone()
}

#[derive(Clone)]
struct ModelManager {
    // postgres connection pool
    pg: PgPool,
}

impl ModelManager {
    async fn new() -> Self {
        let pg = PgPoolOptions::new()
            .acquire_timeout(std::time::Duration::from_secs(2))
            .max_connections(4)
            .connect(PG_URL)
            .await
            .unwrap();
        Self { pg }
    }

    async fn pg(&self) {
        self.pg.acquire().await.unwrap();
    }
    async fn foo(&self) {
        println!("foo");
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test1() {
        let mm = init_test().await;
        mm.pg().await;
    }

    #[tokio::test]
    async fn test2() {
        let _mm = init_test().await;
    }
}

Info

• SQLx version: 0.7.2; 0.6.3
• SQLx features enabled: ["runtime-tokio-rustls", "postgres"]
• Database server and version: Postgres 16.1 (docker); MySql:latest (docker)
• Operating system: OSX
• rustc --version: rustc 1.74.0 (79e9716c9 2023-11-13) and rustc 1.76.0-nightly (a57770440 2023-11-16)

[kingwingfly]

Enable this will lower the possibility the bug happens, but cannot fix it.

PgPoolOptions::new()
    .test_before_acquire(false)
    .connect(&url);

[kingwingfly]

I also added this

PgPoolOptions::new()
    .before_acquire(|_, _| {
        Box::pin(async {
            println!("1");
            Ok(true)
        })
    })

If the bug happen, nothing printed, if not, "1" printed.

Does it means acquire is even not called, then I got a TimedOut error?

[kingwingfly]

I'm sorry, but it's due to I'm using tokio wrong way.
The reason is tokio-rs/tokio#6066 (comment)

Some issue concerned:
tokio-rs/tokio#6066
tokio-rs/tokio#2374

[LennDG]

For anyone else coming here: the simplest way to solve this is by setting the max_connections to 1 during test, i.e.:

async fn new() -> Self {
        let max_connections = if cfg!(test) { 1 } else {4 };

        let pg = PgPoolOptions::new()
            .acquire_timeout(std::time::Duration::from_secs(2))
            .max_connections(max_connections )
            .connect(PG_URL)
            .await
            .unwrap();
        Self { pg }
    }

I am not entirely sure what causes this to fix it but I found it solved by this commit in a tutorial.

[kingwingfly]

I have put the reason and fixing method here