main
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use config::redis::{RedisConfig, RedisMode};
use redis::aio::{ConnectionManager, ConnectionManagerConfig, MultiplexedConnection};
use redis::cluster_async::ClusterConnection;
use redis::{ConnectionAddr, ConnectionInfo, Value};
use crate::error::{CacheError, CacheResult};
/// The process-wide Redis connection layer.
///
/// One instance owns `pool_size` multiplexed connections and is cloned into every
/// consumer ([`crate::AppRedis`], [`crate::Cache`], session, actor, subscriptions)
/// so a process opens exactly one pool instead of one per consumer.
#[derive(Clone)]
pub struct CacheConn {
connections: Arc<Vec<CacheConnection>>,
cursor: Arc<AtomicUsize>,
config: RedisConfig,
/// Connection parameters (scheme, auth, tls) of the first configured node,
/// reused to open per-node connections for cluster-wide operations.
node_info: Option<ConnectionInfo>,
db: i64,
cluster: bool,
}
#[derive(Clone)]
pub enum CacheConnection {
Standalone(ConnectionManager),
Cluster(ClusterConnection),
}
impl CacheConnection {
pub fn is_cluster(&self) -> bool {
matches!(self, CacheConnection::Cluster(_))
}
}
impl CacheConn {
pub async fn init(cfg: &RedisConfig) -> CacheResult<Self> {
cfg.validate().map_err(CacheError::Config)?;
let cluster = cfg.is_cluster();
let seeds = cfg.node_urls().map_err(CacheError::Config)?;
let pool_size = cfg.pool_size().max(1);
let (connections, node_info) = if cluster {
let nodes = seeds
.iter()
.map(|url| connection_info(url, cfg))
.collect::<CacheResult<Vec<_>>>()?;
let node_info = nodes.first().cloned();
let client = redis::cluster::ClusterClient::builder(nodes)
.connection_timeout(cfg.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, node_info)
} else {
let url = seeds
.first()
.ok_or_else(|| CacheError::config("redis requires at least one node"))?;
let node_info = connection_info(url, cfg)?;
let client = redis::Client::open(node_info.clone()).map_err(CacheError::from)?;
// 阻塞命令（如 XREADGROUP BLOCK）的耗时会超过 redis-rs 默认的 500ms
// 响应超时，这里与 get_multiplexed_async_connection_with_config 保持一致。
let manager_config = ConnectionManagerConfig::new()
.set_connection_timeout(Some(cfg.create_timeout()))
.set_response_timeout(None);
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::Standalone(manager));
}
(connections, Some(node_info))
};
let db = node_info
.as_ref()
.map(|info| info.redis_settings().db())
.unwrap_or_else(|| cfg.database());
let conn = Self {
connections: Arc::new(connections),
cursor: Arc::new(AtomicUsize::new(0)),
config: cfg.clone(),
node_info,
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 config(&self) -> &RedisConfig {
&self.config
}
pub fn mode(&self) -> RedisMode {
self.config.mode()
}
pub fn connection(&self) -> &CacheConnection {
&self.connections[0]
}
pub fn connection_clone(&self) -> CacheConnection {
let index = self.cursor.fetch_add(1, Ordering::Relaxed) % self.connections.len();
self.connections[index].clone()
}
/// Opens a connection dedicated to a blocking command.
///
/// `BRPOP` holds its socket for the whole timeout window, so it must not take
/// a pooled connection. Under cluster the command has to reach the master
/// owning the key's slot, otherwise it receives `MOVED`; standalone uses the
/// configured node directly.
pub async fn blocking_connection(&self, key: &str) -> CacheResult<MultiplexedConnection> {
let client = if self.cluster {
let node_info = self
.node_info
.clone()
.ok_or_else(|| CacheError::config("cluster connection info is unavailable"))?;
let mut conn = self.connection_clone();
let CacheConnection::Cluster(cluster) = &mut conn else {
return Err(CacheError::config("cluster connection missing"));
};
let slot: u16 = redis::cmd("CLUSTER")
.arg("KEYSLOT")
.arg(key)
.query_async(cluster)
.await
.map_err(CacheError::from)?;
let slots: Value = redis::cmd("CLUSTER")
.arg("SLOTS")
.query_async(cluster)
.await
.map_err(CacheError::from)?;
let (host, port) = master_for_slot(&slots, slot)
.ok_or_else(|| CacheError::config(format!("no master owns slot {slot}")))?;
let addr = match node_info.addr() {
ConnectionAddr::Tcp(_, _) => ConnectionAddr::Tcp(host, port),
ConnectionAddr::TcpTls {
insecure,
tls_params,
..
} => ConnectionAddr::TcpTls {
host,
port,
insecure: *insecure,
tls_params: tls_params.clone(),
},
ConnectionAddr::Unix(_) => {
return Err(CacheError::config(
"unix sockets are not supported for cluster connections",
));
}
_ => {
return Err(CacheError::config(
"unsupported connection address type for cluster connections",
));
}
};
redis::Client::open(node_info.set_addr(addr)).map_err(CacheError::from)?
} else {
let node_info = self
.node_info
.clone()
.ok_or_else(|| CacheError::config("redis connection info is unavailable"))?;
redis::Client::open(node_info).map_err(CacheError::from)?
};
let config = redis::AsyncConnectionConfig::new()
.set_response_timeout(None)
.set_connection_timeout(Some(self.config.connection_timeout()));
client
.get_multiplexed_async_connection_with_config(&config)
.await
.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()
}
/// Opens a long-lived blocking connection for a queue consumer.
pub async fn blocking(&self, key: &str) -> CacheResult<BlockingConnection> {
Ok(BlockingConnection {
connection: self.blocking_connection(key).await?,
})
}
}
/// A socket dedicated to blocking commands, held for the consumer's lifetime.
///
/// Reusing one connection avoids a fresh TCP/TLS handshake on every poll cycle.
pub struct BlockingConnection {
connection: MultiplexedConnection,
}
impl BlockingConnection {
/// `BRPOP key timeout`; returns `None` when the timeout elapses.
pub async fn brpop_bytes(
&mut self,
key: &str,
timeout_secs: u64,
) -> CacheResult<Option<Vec<u8>>> {
let result: Option<(String, Vec<u8>)> = redis::cmd("BRPOP")
.arg(key)
.arg(timeout_secs.max(1))
.query_async(&mut self.connection)
.await
.map_err(CacheError::from)?;
Ok(result.map(|(_, payload)| payload))
}
}
/// Builds a node's connection info from its configured URL.
///
/// The URL is authoritative for scheme, credentials and database; the explicit
/// config fields override it when set. This is the single place that decides how
/// a seed string becomes a connection, so both the keyed and the raw-key
/// consumers see identical semantics.
fn connection_info(url: &str, cfg: &RedisConfig) -> CacheResult<ConnectionInfo> {
let info = redis::Client::open(url)
.map_err(CacheError::from)?
.get_connection_info()
.clone();
let mut settings = info.redis_settings().clone();
if let Some(database) = cfg.database {
settings = settings.set_db(i64::from(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))
}
fn master_for_slot(slots: &Value, slot: u16) -> Option<(String, u16)> {
let Value::Array(groups) = slots else {
return None;
};
for group in groups {
let Value::Array(fields) = group else {
continue;
};
let (Some(Value::Int(start)), Some(Value::Int(end))) = (fields.first(), fields.get(1))
else {
continue;
};
if !(*start..=*end).contains(&(slot as i64)) {
continue;
}
let Some(Value::Array(nodes)) = fields.get(2) else {
continue;
};
let (Some(Value::BulkString(host)), Some(Value::Int(port))) = (nodes.first(), nodes.get(1))
else {
continue;
};
return Some((String::from_utf8_lossy(host).into_owned(), *port as u16));
}
None
}