b7cd2b1
use crate::error::{CacheError, CacheResult};
use config::identity::redis::RedisConfig;
use deadpool_redis::{
ConnectionAddr, ConnectionInfo, PoolConfig, RedisConnectionInfo, Runtime, Status, Timeouts,
};
#[derive(Clone)]
pub struct CacheConn {
pool: CachePool,
db: i64,
}
#[derive(Clone)]
enum CachePool {
Standard(deadpool_redis::Pool),
Cluster(deadpool_redis::cluster::Pool),
}
pub enum CacheConnection {
Standard(deadpool_redis::Connection),
Cluster(deadpool_redis::cluster::Connection),
}
impl CacheConn {
pub async fn init(cfg: &RedisConfig) -> CacheResult<Self> {
let pool = &cfg.pool;
let pool_config = PoolConfig {
max_size: pool.max_size,
timeouts: Timeouts {
wait: Some(pool.wait_timeout()),
create: Some(pool.create_timeout()),
recycle: Some(pool.recycle_timeout()),
},
..PoolConfig::default()
};
Self::init_with_pool_config(cfg, pool_config).await
}
pub async fn init_with_pool_config(
cfg: &RedisConfig,
pool_config: PoolConfig,
) -> CacheResult<Self> {
cfg.validate().map_err(CacheError::Config)?;
let db = cfg.database();
let cluster = cfg.is_cluster();
let seeds = cfg.seeds().map_err(CacheError::Config)?;
let connections: Vec<ConnectionInfo> = seeds
.into_iter()
.map(|(host, port)| ConnectionInfo {
addr: ConnectionAddr::Tcp(host, port),
redis: RedisConnectionInfo {
db,
username: cfg.username.clone(),
password: cfg.password.clone(),
..RedisConnectionInfo::default()
},
})
.collect();
let pool = if cluster {
let manager = deadpool_redis::cluster::Manager::new(connections, false)?;
CachePool::Cluster(
deadpool_redis::cluster::Pool::builder(manager)
.config(pool_config)
.runtime(Runtime::Tokio1)
.build()
.map_err(|error| CacheError::config(error.to_string()))?,
)
} else {
let connection = single_connection(connections)?;
let manager = deadpool_redis::Manager::new(connection)?;
CachePool::Standard(
deadpool_redis::Pool::builder(manager)
.config(pool_config)
.runtime(Runtime::Tokio1)
.build()
.map_err(|error| CacheError::config(error.to_string()))?,
)
};
let conn = Self { pool, db };
if let Err(error) = conn.acquire().await {
conn.close();
return Err(error);
}
Ok(conn)
}
pub async fn acquire(&self) -> CacheResult<CacheConnection> {
match &self.pool {
CachePool::Standard(pool) => pool.get().await.map(CacheConnection::Standard),
CachePool::Cluster(pool) => pool.get().await.map(CacheConnection::Cluster),
}
.map_err(CacheError::from)
}
pub fn database(&self) -> i64 {
self.db
}
pub fn is_cluster(&self) -> bool {
matches!(self.pool, CachePool::Cluster(_))
}
pub fn status(&self) -> Status {
match &self.pool {
CachePool::Standard(pool) => pool.status(),
CachePool::Cluster(pool) => pool.status(),
}
}
pub fn close(&self) {
match &self.pool {
CachePool::Standard(pool) => pool.close(),
CachePool::Cluster(pool) => pool.close(),
}
}
pub fn is_closed(&self) -> bool {
match &self.pool {
CachePool::Standard(pool) => pool.is_closed(),
CachePool::Cluster(pool) => pool.is_closed(),
}
}
pub fn pool_size(&self) -> usize {
match &self.pool {
CachePool::Standard(pool) => pool.status().size,
CachePool::Cluster(pool) => pool.status().size,
}
}
pub fn available(&self) -> usize {
match &self.pool {
CachePool::Standard(pool) => pool.status().available,
CachePool::Cluster(pool) => pool.status().available,
}
}
}
impl CacheConnection {
pub fn invalidate(self) {
match self {
Self::Standard(connection) => drop(deadpool_redis::Connection::take(connection)),
Self::Cluster(connection) => {
drop(deadpool_redis::cluster::Connection::take(connection));
}
}
}
}
fn single_connection(connections: Vec<ConnectionInfo>) -> CacheResult<ConnectionInfo> {
connections
.into_iter()
.next()
.ok_or_else(|| CacheError::config("redis requires at least one node"))
}