zhenyi
118 lines
3.8 KiB
Rust
118 lines
3.8 KiB
Rust
use crate::config::AppConfig;
|
|
use crate::error::AppError;
|
|
use crate::error::AppResult;
|
|
use r2d2::Pool;
|
|
use redis::cluster::ClusterClient;
|
|
use redis::{Client, ConnectionLike, RedisError};
|
|
use std::time::Duration;
|
|
|
|
#[derive(Clone)]
|
|
enum RedisBackend {
|
|
Single(Pool<Client>),
|
|
Cluster(Pool<ClusterClient>),
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct AppRedis {
|
|
backend: RedisBackend,
|
|
}
|
|
|
|
impl AppRedis {
|
|
pub fn from_config(config: &AppConfig) -> AppResult<Self> {
|
|
let backend = if config.redis_cluster_enabled()? {
|
|
let nodes = config.redis_cluster_nodes()?;
|
|
let cluster_client =
|
|
ClusterClient::new(nodes.iter().map(|s| s.as_str()).collect::<Vec<_>>())?;
|
|
let pool = Self::build_pool(config, cluster_client)?;
|
|
RedisBackend::Cluster(pool)
|
|
} else {
|
|
let url = config
|
|
.redis_url()?
|
|
.ok_or_else(|| AppError::Config("APP_REDIS_URL is not set".into()))?;
|
|
let client = Client::open(url.as_str())?;
|
|
let pool = Self::build_pool(config, client)?;
|
|
RedisBackend::Single(pool)
|
|
};
|
|
|
|
Ok(Self { backend })
|
|
}
|
|
|
|
fn build_pool<M: r2d2::ManageConnection>(config: &AppConfig, manager: M) -> AppResult<Pool<M>> {
|
|
let max_conn = config.redis_max_connections()?;
|
|
let min_conn = config.redis_min_connections()?;
|
|
let idle_timeout = config.redis_idle_timeout()?;
|
|
let conn_timeout = config.redis_connection_timeout()?;
|
|
|
|
Ok(r2d2::Builder::new()
|
|
.max_size(max_conn)
|
|
.min_idle(Some(min_conn))
|
|
.idle_timeout(Some(Duration::from_secs(idle_timeout)))
|
|
.connection_timeout(Duration::from_secs(conn_timeout))
|
|
.build(manager)?)
|
|
}
|
|
|
|
pub fn get_connection(&self) -> Result<PooledRedisConnection, r2d2::Error> {
|
|
match &self.backend {
|
|
RedisBackend::Single(pool) => pool.get().map(PooledRedisConnection::Single),
|
|
RedisBackend::Cluster(pool) => pool.get().map(PooledRedisConnection::Cluster),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[allow(clippy::large_enum_variant)]
|
|
pub enum PooledRedisConnection {
|
|
Single(r2d2::PooledConnection<Client>),
|
|
Cluster(r2d2::PooledConnection<ClusterClient>),
|
|
}
|
|
|
|
impl PooledRedisConnection {
|
|
pub fn inner_mut(&mut self) -> &mut dyn ConnectionLike {
|
|
match self {
|
|
PooledRedisConnection::Single(conn) => conn,
|
|
PooledRedisConnection::Cluster(conn) => conn,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl ConnectionLike for PooledRedisConnection {
|
|
fn req_packed_command(&mut self, cmd: &[u8]) -> Result<redis::Value, RedisError> {
|
|
match self {
|
|
PooledRedisConnection::Single(conn) => conn.req_packed_command(cmd),
|
|
PooledRedisConnection::Cluster(conn) => conn.req_packed_command(cmd),
|
|
}
|
|
}
|
|
|
|
fn req_packed_commands(
|
|
&mut self,
|
|
cmd: &[u8],
|
|
offset: usize,
|
|
count: usize,
|
|
) -> Result<Vec<redis::Value>, RedisError> {
|
|
match self {
|
|
PooledRedisConnection::Single(conn) => conn.req_packed_commands(cmd, offset, count),
|
|
PooledRedisConnection::Cluster(conn) => conn.req_packed_commands(cmd, offset, count),
|
|
}
|
|
}
|
|
|
|
fn get_db(&self) -> i64 {
|
|
match self {
|
|
PooledRedisConnection::Single(conn) => conn.get_db(),
|
|
PooledRedisConnection::Cluster(conn) => conn.get_db(),
|
|
}
|
|
}
|
|
|
|
fn check_connection(&mut self) -> bool {
|
|
match self {
|
|
PooledRedisConnection::Single(conn) => conn.check_connection(),
|
|
PooledRedisConnection::Cluster(conn) => conn.check_connection(),
|
|
}
|
|
}
|
|
|
|
fn is_open(&self) -> bool {
|
|
match self {
|
|
PooledRedisConnection::Single(conn) => conn.is_open(),
|
|
PooledRedisConnection::Cluster(conn) => conn.is_open(),
|
|
}
|
|
}
|
|
}
|