4c83001
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use config::redis::{RedisConfig, RedisMode};
use redis::aio::{ConnectionManager, MultiplexedConnection};
use redis::cluster_async::ClusterConnection;
use redis::{ConnectionAddr, ConnectionInfo, Value};
use crate::error::ExportError;
#[derive(Clone)]
pub enum RedisConnection {
Standalone(ConnectionManager),
Cluster(ClusterConnection),
}
impl RedisConnection {
pub fn is_cluster(&self) -> bool {
matches!(self, RedisConnection::Cluster(_))
}
}
#[derive(Clone)]
pub struct AppRedis {
config: RedisConfig,
/// One multiplexed connection per pool slot; every slot is an independent
/// connection (not a clone of another slot) so requests spread over
/// multiple sockets.
connections: Vec<RedisConnection>,
/// Round-robin cursor shared across clones of this handle.
cursor: Arc<AtomicUsize>,
/// Connection parameters (auth, tls) of the first configured node,
/// reused to open per-node connections for cluster-wide operations.
node_info: Option<ConnectionInfo>,
}
impl AppRedis {
pub async fn connect(config: RedisConfig) -> Result<Self, ExportError> {
let url = config.validate().map_err(ExportError::InvalidConfig)?;
let pool_size = config.pool_size().max(1);
let (connections, node_info) = match config.mode() {
RedisMode::Cluster => {
let nodes = url
.split(',')
.map(str::trim)
.filter(|node| !node.is_empty())
.map(String::from)
.collect::<Vec<_>>();
let client = redis::cluster::ClusterClient::new(nodes).map_err(ExportError::from)?;
let node_info = connection_info(url)?;
let mut connections = Vec::with_capacity(pool_size);
for _ in 0..pool_size {
let connection = client
.get_async_connection()
.await
.map_err(ExportError::from)?;
connections.push(RedisConnection::Cluster(connection));
}
(connections, Some(node_info))
}
RedisMode::Standalone | RedisMode::Sentinel | RedisMode::ValkeyStandalone => {
let mut connections = Vec::with_capacity(pool_size);
for _ in 0..pool_size {
let client = redis::Client::open(url).map_err(ExportError::from)?;
let manager = ConnectionManager::new(client).await.map_err(ExportError::from)?;
connections.push(RedisConnection::Standalone(manager));
}
let node_info = connection_info(url)?;
(connections, Some(node_info))
}
};
Ok(Self {
config,
connections,
cursor: Arc::new(AtomicUsize::new(0)),
node_info,
})
}
pub fn config(&self) -> &RedisConfig {
&self.config
}
pub fn mode(&self) -> RedisMode {
self.config.mode()
}
pub fn is_cluster(&self) -> bool {
self.config.mode() == RedisMode::Cluster
}
pub fn connection(&self) -> &RedisConnection {
&self.connections[0]
}
pub fn connection_clone(&self) -> RedisConnection {
let index = self.cursor.fetch_add(1, Ordering::Relaxed) % self.connections.len();
self.connections[index].clone()
}
/// Opens a multiplexed connection to every master node of the cluster.
///
/// Masters are discovered at call time via `CLUSTER SLOTS`, so this keeps
/// working after topology changes. Returns an empty list for standalone
/// connections.
pub async fn cluster_master_connections(
&self,
) -> Result<Vec<MultiplexedConnection>, ExportError> {
let mut conn = self.connection_clone();
let RedisConnection::Cluster(cluster) = &mut conn else {
return Ok(Vec::new());
};
let slots: Value = redis::cmd("CLUSTER")
.arg("SLOTS")
.query_async(cluster)
.await
.map_err(ExportError::from)?;
let node_info = self.node_info.clone().ok_or_else(|| {
ExportError::InvalidConfig("cluster connection info is unavailable".into())
})?;
let mut seen = std::collections::HashSet::new();
let mut connections = Vec::new();
for (host, port) in cluster_master_addrs(&slots) {
if !seen.insert((host.clone(), port)) {
continue;
}
let addr = match node_info.addr() {
ConnectionAddr::Tcp(_, _) => ConnectionAddr::Tcp(host.clone(), port),
ConnectionAddr::TcpTls { insecure, tls_params, .. } => ConnectionAddr::TcpTls {
host: host.clone(),
port,
insecure: *insecure,
tls_params: tls_params.clone(),
},
ConnectionAddr::Unix(_) => {
return Err(ExportError::InvalidConfig(
"unix sockets are not supported for cluster connections".into(),
));
}
_ => {
return Err(ExportError::InvalidConfig(
"unsupported connection address type for cluster connections".into(),
));
}
};
let mut info = node_info.clone();
info = info.set_addr(addr);
let client = redis::Client::open(info).map_err(ExportError::from)?;
let connection = client
.get_multiplexed_async_connection()
.await
.map_err(ExportError::from)?;
connections.push(connection);
}
Ok(connections)
}
}
/// Extracts the host/port of every master from a `CLUSTER SLOTS` response.
///
/// Each slot group is `[start, end, [ip, port, node_id, ...], replicas...]`,
/// where the first entry of the node list is the master.
fn cluster_master_addrs(slots: &Value) -> Vec<(String, u16)> {
let Value::Array(groups) = slots else {
return Vec::new();
};
let mut masters = Vec::new();
for group in groups {
let Value::Array(fields) = group else {
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;
};
masters.push((String::from_utf8_lossy(host).into_owned(), *port as u16));
}
masters
}
/// Builds the connection info of the first configured node, used as a template
/// (scheme, auth, tls settings) for connections to other cluster nodes.
fn connection_info(urls: &str) -> Result<ConnectionInfo, ExportError> {
let first = urls
.split(',')
.map(str::trim)
.find(|node| !node.is_empty())
.ok_or_else(|| ExportError::InvalidConfig("redis url contains no nodes".into()))?;
Ok(redis::Client::open(first)?
.get_connection_info()
.clone())
}