文盤Rust -- r2d2 實現redis連接池

作者：賈世聞

我們在開發應用後端系統的時候經常要和各種資料庫、緩存等資源打交道。這一期，我們聊聊如何訪問redis 並將資源池化。

在一個應用後端程式訪問redis主要要做的工作有兩個，單例和池化。

在後端應用集成redis，我們主要用到以下幾個crate:​ ​once_cell​​​、​ ​redis-rs​​​、​ ​r2d2​​​.once_cell 實現單例；redis-rs 是 redis的 rust 驅動；r2d2 是一個池化連接的工具包。本期代碼均出現在​ ​fullstack-rs​​​項目中。​ ​fullstack-rs​​​是我新開的一個實驗性項目，目標是做一個類似​ ​gin-vue-admin​​的集成開發框架。

redis資源的定義主要是在​ ​https://github.com/jiashiwen/fullstack-rs/blob/main/backend/src/resources/redis_resource.rs​​ 中實現的。

一、redis-rs 封裝

在實際開發中，我們面對的redis資源可能是單實例也有可能是集群，在這裡我們對redis-rs進行了簡單封裝，便於適應這兩種情況。

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Clone)]
#[serde(rename_all = "lowercase")]
pub struct RedisInstance {
    #[serde(default = "RedisInstance::urls_default")]
    pub urls: Vec<String>,
    #[serde(default = "RedisInstance::password_default")]
    pub password: String,
    #[serde(default = "RedisInstance::instance_type_default")]
    pub instance_type: InstanceType,
}

RedisInstance,定義redis資源的描述，與配置文件相對應。詳細的配置描述可以參考 ​ ​https://github.com/jiashiwen/fullstack-rs/blob/main/backend/src/configure/config_global.rs​​ 文件中 RedisConfig 和 RedisPool 兩個 struct 描述。

#[derive(Clone)]
pub enum RedisClient {
    Single(redis::Client),
    Cluster(redis::cluster::ClusterClient),
}


impl RedisClient {
    pub fn get_redis_connection(&self) -> RedisResult<RedisConnection> {
        return match self {
            RedisClient::Single(s) => {
                let conn = s.get_connection()?;
                Ok(RedisConnection::Single(Box::new(conn)))
            }
            RedisClient::Cluster(c) => {
                let conn = c.get_connection()?;
                Ok(RedisConnection::Cluster(Box::new(conn)))
            }
        };
    }
}


pub enum RedisConnection {
    Single(Box<redis::Connection>),
    Cluster(Box<redis::cluster::ClusterConnection>),
}


impl RedisConnection {
    pub fn is_open(&self) -> bool {
        return match self {
            RedisConnection::Single(sc) => sc.is_open(),
            RedisConnection::Cluster(cc) => cc.is_open(),
        };
    }


    pub fn query<T: FromRedisValue>(&mut self, cmd: &redis::Cmd) -> RedisResult<T> {
        return match self {
            RedisConnection::Single(sc) => match sc.as_mut().req_command(cmd) {
                Ok(val) => from_redis_value(&val),
                Err(e) => Err(e),
            },
            RedisConnection::Cluster(cc) => match cc.req_command(cmd) {
                Ok(val) => from_redis_value(&val),
                Err(e) => Err(e),
            },
        };
    }
}

RedisClient 和 RedisConnection 對redis 的鏈接進行了封裝，用來實現統一的調用介面。

二、基於 r2d2 實現 redis 連接池

以上，基本完成的reids資源的準備工作，下麵來實現一個redis鏈接池。

#[derive(Clone)]
pub struct RedisConnectionManager {
    pub redis_client: RedisClient,
}


impl r2d2::ManageConnection for RedisConnectionManager {
    type Connection = RedisConnection;
    type Error = RedisError;


    fn connect(&self) -> Result<RedisConnection, Self::Error> {
        let conn = self.redis_client.get_redis_connection()?;
        Ok(conn)
    }


    fn is_valid(&self, conn: &mut RedisConnection) -> Result<(), Self::Error> {
        match conn {
            RedisConnection::Single(sc) => {
                redis::cmd("PING").query(sc)?;
            }
            RedisConnection::Cluster(cc) => {
                redis::cmd("PING").query(cc)?;
            }
        }
        Ok(())
    }


    fn has_broken(&self, conn: &mut RedisConnection) -> bool {
        !conn.is_open()
    }
}

利用 r2d2 來實現連接池需要實現 r2d2::ManageConnection trait。connect 函數獲取連接；is_valid 函數校驗連通性；has_broken 判斷連接是否崩潰不可用。

pub fn gen_redis_conn_pool() -> Result<Pool<RedisConnectionManager>> {
    let config = get_config()?;
    let redis_client = config.redis.instance.to_redis_client()?;
    let manager = RedisConnectionManager { redis_client };
    let pool = r2d2::Pool::builder()
        .max_size(config.redis.pool.max_size as u32)
        .min_idle(Some(config.redis.pool.mini_idle as u32))
        .connection_timeout(Duration::from_secs(
            config.redis.pool.connection_timeout as u64,
        ))
        .build(manager)?;
    Ok(pool)
}

gen_redis_conn_pool 函數用來生成一個 redis 的連接池，根據配置文件來指定連接池的最大連接數，最小閑置連接以及連接超時時長。

三、連接池單例實現

在後端開發中，對於單一資源一般採取單例模式避免重覆產生實例的開銷。下麵來聊一聊如果構建一個全局的 redis 資源。

這一部分代碼在​ ​https://github.com/jiashiwen/fullstack-rs/blob/main/backend/src/resources/init_resources.rs​​ 文件中。

pub static GLOBAL_REDIS_POOL: OnceCell<r2d2::Pool<RedisConnectionManager>> = OnceCell::new();

利用 OnceCell 構建全局靜態變數。

fn init_global_redis() {
    GLOBAL_REDIS_POOL.get_or_init(|| {
        let pool = match gen_redis_conn_pool() {
            Ok(it) => it,
            Err(err) => panic!("{}", err.to_string()),
        };
        pool
    });
}

init_global_redis 函數，用來初始化 GLOBAL_REDIS_POOL 全局靜態變數。在一般的後端程式中，資源是強依賴，所以，初始化簡單粗暴，要麼成功要麼 panic。

四、資源調用

準備好 redis 資源後，我們聊聊如何調用。

調用例子在這裡​ ​https://github.com/jiashiwen/fullstack-rs/blob/main/backend/src/httpserver/service/service_redis.rs​​

pub fn put(kv: KV) -> Result<()> {
    let conn = GLOBAL_REDIS_POOL.get();
    return match conn {
        Some(c) => {
            c.get()?
                .query(redis::cmd("set").arg(kv.Key).arg(kv.Value))?;
            Ok(())
        }
        None => Err(anyhow!("redis pool not init")),
    };
}

​ ​https://github.com/jiashiwen/fullstack-rs/tree/main/backend​​ 這個工程里有從http入口開始到寫入redis的完整流程，http server 不在本文討論之列，就不贅述了，有興趣的同學可以去github看看。

咱們下期見。