Rocket框架多主从数据库连接池安全复用方案咨询

我来给你梳理几个更优雅的方案，彻底摆脱在路由里写一堆连接Guard的繁琐问题：

Rocket的State机制可以让你把所有数据库连接池集中起来管理，不用在每个路由里逐个注入。核心思路是把每组主从封装成一个集群结构体，再把所有集群放到全局State中，路由里只需要注入一次State，就能按需获取任意集群的主/从连接。

代码示例 #

首先定义集群结构体和全局容器：

use sqlx::{postgres::Postgres, Pool};
use rocket::State;
use rand::Rng;

// 单组主从集群的封装
#[derive(Debug, Clone)]
struct DbCluster {
    master: Pool<Postgres>,
    slaves: Vec<Pool<Postgres>>,
}

impl DbCluster {
    // 获取主库连接（写入用）
    pub async fn get_master(&self) -> Result<sqlx::postgres::PgConnection, sqlx::Error> {
        self.master.acquire().await
    }

    // 随机获取一个从库连接（读取用）
    pub async fn get_slave(&self) -> Result<sqlx::postgres::PgConnection, sqlx::Error> {
        let mut rng = rand::thread_rng();
        let slave_idx = rng.gen_range(0..self.slaves.len());
        self.slaves[slave_idx].acquire().await
    }
}

// 所有数据库集群的全局容器（对应你的5组主从）
#[derive(Debug, Clone)]
struct DbClusters {
    user_db: DbCluster,
    order_db: DbCluster,
    product_db: DbCluster,
    // ... 剩下2组集群
}

然后在启动时初始化所有连接池并挂载到State：

#[rocket::launch]
async fn rocket() -> _ {
    // 初始化用户库的主从连接池
    let user_master_pool = Pool::connect("postgres://user:pass@master-user-db:5432/db").await.unwrap();
    let user_slave_pools = vec![
        Pool::connect("postgres://user:pass@slave-user-1:5432/db").await.unwrap(),
        Pool::connect("postgres://user:pass@slave-user-2:5432/db").await.unwrap(),
        Pool::connect("postgres://user:pass@slave-user-3:5432/db").await.unwrap(),
        Pool::connect("postgres://user:pass@slave-user-4:5432/db").await.unwrap(),
    ];
    let user_db_cluster = DbCluster {
        master: user_master_pool,
        slaves: user_slave_pools,
    };

    // 同理初始化其他4组集群...

    let db_clusters = DbClusters {
        user_db: user_db_cluster,
        order_db: order_db_cluster,
        // ... 其他集群
    };

    rocket::build()
        .manage(db_clusters)
        .mount("/", routes![register_user, get_user])
}

最后在路由里直接用State获取连接：

#[rocket::post("/register")]
async fn register_user(clusters: &State<DbClusters>) -> Result<String, sqlx::Error> {
    // 获取用户库主连接执行写入
    let mut master_conn = clusters.user_db.get_master().await?;
    sqlx::query!("INSERT INTO users (name) VALUES ($1)", "Navin")
        .execute(&mut master_conn)
        .await?;

    Ok("User registered".to_string())
}

#[rocket::get("/user/<id>")]
async fn get_user(clusters: &State<DbClusters>, id: i32) -> Result<String, sqlx::Error> {
    // 获取用户库随机从连接执行读取
    let mut slave_conn = clusters.user_db.get_slave().await?;
    let user = sqlx::query!("SELECT name FROM users WHERE id = $1", id)
        .fetch_one(&mut slave_conn)
        .await?;

    Ok(format!("User name: {}", user.name))
}

如果觉得每次在路由里调用get_master()/get_slave()还是有点麻烦，可以给每个集群的主从分别自定义Guard，让路由参数更简洁。

代码示例 #

比如给用户库的主从写Guard：

// 用户库主连接Guard
#[derive(Debug)]
struct UserDbMaster(sqlx::postgres::PgConnection);

#[rocket::async_trait]
impl rocket::request::FromRequest<'_> for UserDbMaster {
    type Error = sqlx::Error;

    async fn from_request(req: &rocket::Request<'_>) -> rocket::request::Outcome<Self, Self::Error> {
        let clusters = req.guard::<&State<DbClusters>>().await.unwrap();
        match clusters.user_db.get_master().await {
            Ok(conn) => rocket::request::Outcome::Success(UserDbMaster(conn)),
            Err(e) => rocket::request::Outcome::Failure((rocket::http::Status::ServiceUnavailable, e)),
        }
    }
}

// 用户库从连接Guard
#[derive(Debug)]
struct UserDbSlave(sqlx::postgres::PgConnection);

#[rocket::async_trait]
impl rocket::request::FromRequest<'_> for UserDbSlave {
    type Error = sqlx::Error;

    async fn from_request(req: &rocket::Request<'_>) -> rocket::request::Outcome<Self, Self::Error> {
        let clusters = req.guard::<&State<DbClusters>>().await.unwrap();
        match clusters.user_db.get_slave().await {
            Ok(conn) => rocket::request::Outcome::Success(UserDbSlave(conn)),
            Err(e) => rocket::request::Outcome::Failure((rocket::http::Status::ServiceUnavailable, e)),
        }
    }
}

现在路由可以写成这样，非常简洁：

#[rocket::post("/register")]
async fn register_user(master: UserDbMaster) -> Result<String, sqlx::Error> {
    sqlx::query!("INSERT INTO users (name) VALUES ($1)", "Navin")
        .execute(&mut master.0)
        .await?;

    Ok("User registered".to_string())
}

#[rocket::get("/user/<id>")]
async fn get_user(slave: UserDbSlave, id: i32) -> Result<String, sqlx::Error> {
    let user = sqlx::query!("SELECT name FROM users WHERE id = $1", id)
        .fetch_one(&mut slave.0)
        .await?;

    Ok(format!("User name: {}", user.name))
}

1. 彻底解耦：不用再在路由里写一堆重复的连接参数，所有连接池逻辑集中管理
2. 逻辑复用：从库随机选择的逻辑封装在DbCluster里，所有路由共用同一套规则
3. 类型安全：通过结构体明确区分不同业务的数据库集群，避免混淆
4. 可扩展性：后续新增主从集群只需要在DbClusters里加一个字段，不用修改现有路由

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