f2e745a
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use config::identity::redis::RedisConfig;
use redis::ConnectionInfo;
use redis::aio::{ConnectionManager, ConnectionManagerConfig};
use redis::cluster_async::ClusterConnection;
use crate::error::{CacheError, CacheResult};
#[derive(Clone)]
pub struct CacheConn {
connections: Arc<Vec<CacheConnection>>,
cursor: Arc<AtomicUsize>,
db: i64,
cluster: bool,
}
#[derive(Clone)]
pub enum CacheConnection {
Standard(ConnectionManager),
Cluster(ClusterConnection),
}
impl CacheConn {
pub async fn init(cfg: &RedisConfig) -> 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 pool_size = cfg.pool.max_size.max(1);
let connections = if cluster {
let nodes = seeds
.iter()
.map(|(host, port)| connection_info(host, *port, cfg))
.collect::<CacheResult<Vec<_>>>()?;
let client = redis::cluster::ClusterClient::builder(nodes)
.connection_timeout(cfg.pool.create_timeout())
.build()
.map_err(|error| CacheError::config(error.to_string()))?;
let mut connections = Vec::with_capacity(pool_size);
for _ in 0..pool_size {
let connection = client
.get_async_connection()
.await
.map_err(CacheError::from)?;
connections.push(CacheConnection::Cluster(connection));
}
connections
} else {
let (host, port) = seeds
.first()
.cloned()
.ok_or_else(|| CacheError::config("redis requires at least one node"))?;
let info = connection_info(&host, port, cfg)?;
let client = redis::Client::open(info).map_err(CacheError::from)?;
let manager_config = ConnectionManagerConfig::new()
.set_connection_timeout(Some(cfg.pool.create_timeout()));
let mut connections = Vec::with_capacity(pool_size);
for _ in 0..pool_size {
let manager =
ConnectionManager::new_with_config(client.clone(), manager_config.clone())
.await
.map_err(CacheError::from)?;
connections.push(CacheConnection::Standard(manager));
}
connections
};
let conn = Self {
connections: Arc::new(connections),
cursor: Arc::new(AtomicUsize::new(0)),
db,
cluster,
};
Ok(conn)
}
pub async fn acquire(&self) -> CacheResult<CacheConnection> {
if self.connections.is_empty() {
return Err(CacheError::config("redis has no connections"));
}
let index = self.cursor.fetch_add(1, Ordering::Relaxed) % self.connections.len();
Ok(self.connections[index].clone())
}
pub async fn ping(&self) -> CacheResult<()> {
let mut conn = self.acquire().await?;
redis::cmd("PING")
.query_async::<String>(&mut conn)
.await
.map(|_| ())
.map_err(CacheError::from)
}
pub fn database(&self) -> i64 {
self.db
}
pub fn is_cluster(&self) -> bool {
self.cluster
}
pub fn pool_size(&self) -> usize {
self.connections.len()
}
pub fn available(&self) -> usize {
self.connections.len()
}
}
fn connection_info(host: &str, port: u16, cfg: &RedisConfig) -> CacheResult<ConnectionInfo> {
let host = if host.contains(':') && !host.starts_with('[') {
format!("[{host}]")
} else {
host.to_string()
};
let url = format!("redis://{host}:{port}");
let info = redis::Client::open(url)
.map_err(CacheError::from)?
.get_connection_info()
.clone();
let mut settings = info.redis_settings().clone();
settings = settings.set_db(cfg.database());
if let Some(username) = cfg.username.as_deref().filter(|value| !value.is_empty()) {
settings = settings.set_username(username);
}
if let Some(password) = cfg.password.as_deref().filter(|value| !value.is_empty()) {
settings = settings.set_password(password);
}
Ok(info.set_redis_settings(settings))
}