基于Tokio与Axum实现异步数据库连接池的延迟初始化

需求 #

希望一次性延迟初始化包含异步数据库连接池的结构体，并能在全局各处使用。

现有基础代码 #

最初定义的Database结构体及异步初始化方法如下：

#[derive(Debug)]
pub struct Database {
    pub pool: PgPool,
}

impl Database {
    pub async fn init() -> Self {
        let pool = Database::create_db_pool().await;

        Database { pool }
    }

    pub async fn create_db_pool() -> PgPool {
        PgPoolOptions::new()
            .max_connections(20)
            .acquire_timeout(Duration::from_secs(3))
            .connect(API_ENV.db_url)
            .await
            .expect("Can't connect to database")
    }
}

尝试的方案及问题 #

方案1：使用OnceCell #

通过OnceCell实现全局实例，但每次调用都需要写DATABASE.get().unwrap()，不够简洁：

pub static DATABASE: OnceCell<Database> = OnceCell::new();

async fn main() {
    let db = Database::init().await;
    DATABASE.set(db).unwrap();

    let db = DATABASE.get().unwrap();
}

方案2：使用Lazy配合block_on #

尝试用Lazy结合block_on直接初始化全局实例，却触发运行时错误：

pub static DATABASE: Lazy<Database> = Lazy::new(|| {
    let pool: Database = Runtime::new()
        .unwrap()
        .block_on(async { Database::init().await });

    pool
});

async fn main() {
    let db = DATABASE.pool;
}

错误信息：

Cannot start a runtime from within a runtime. This happens because a function (like block_on) attempted to block the current thread while the thread is being used to drive asynchronous tasks.

最终实现方案 #

参考@Chayim Friedman的思路，拆分连接池与结构体的初始化逻辑，实现简洁的全局调用：

修改后的结构体及全局变量定义 #

#[derive(Debug)]
pub struct Database {
    pub pool: PgPool,
}

impl Database {
    pub fn init() -> Self {
        let pool = POOL.get().unwrap();

        Database { pool: pool.clone() }
    }

    pub async fn create_db_pool() -> PgPool {
        PgPoolOptions::new()
            .max_connections(20)
            .acquire_timeout(Duration::from_secs(3))
            .connect(API_ENV.db_url)
            .await
            .expect("Can't connect to database")
    }
}

pub static DATABASE: Lazy<Database> = Lazy::new(Database::init);
pub static POOL: OnceCell<PgPool> = OnceCell::new();

main函数初始化逻辑 #

在异步的main中完成连接池的异步初始化：

async fn main() {
    let pool = Database::create_db_pool().await;
    POOL.set(pool).unwrap();
}

全局使用方式 #

在任意位置直接访问全局DATABASE即可，无需重复unwrap：

fn anywhere() {
    let pool = &DATABASE.pool;
}

内容的提问来源于stack exchange，提问作者ZiiMakc