fad06fd
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::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>,
/// Web Push 是否启用；入队点用它避免往无人消费的队列里写任务。
push_enabled: bool,
}
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(|node| connection_info(node, &config))
.collect::<Result<Vec<_>, _>>()?;
let node_info = nodes.first().cloned();
let client = redis::cluster::ClusterClient::builder(nodes)
.connection_timeout(config.connection_timeout())
.build()
.map_err(ExportError::from)?;
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, node_info)
}
RedisMode::Standalone | RedisMode::Sentinel | RedisMode::ValkeyStandalone => {
let node_info = connection_info(url, &config)?;
let client = redis::Client::open(node_info.clone()).map_err(ExportError::from)?;
let manager_config = ConnectionManagerConfig::new()
.set_connection_timeout(Some(config.connection_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(ExportError::from)?;
connections.push(RedisConnection::Standalone(manager));
}
(connections, Some(node_info))
}
};
Ok(Self {
config,
connections,
cursor: Arc::new(AtomicUsize::new(0)),
node_info,
push_enabled: false,
})
}
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()
}
pub async fn publish(&self, channel: &str, payload: &str) -> Result<(), ExportError> {
let mut conn = self.connection_clone();
let result: Result<i64, redis::RedisError> = match &mut conn {
RedisConnection::Standalone(c) => {
redis::cmd("PUBLISH")
.arg(channel)
.arg(payload)
.query_async(c)
.await
}
RedisConnection::Cluster(c) => {
redis::cmd("PUBLISH")
.arg(channel)
.arg(payload)
.query_async(c)
.await
}
};
result.map(|_| ()).map_err(ExportError::from)
}
/// 列表左端入队（`LPUSH`）；与 `AppRedis::publish` 一样按 standalone/cluster
/// 双分支实现，payload 是二进制安全的字节串。
pub async fn lpush_bytes(&self, key: &str, payload: &[u8]) -> Result<i64, ExportError> {
let mut conn = self.connection_clone();
let result: Result<i64, redis::RedisError> = match &mut conn {
RedisConnection::Standalone(c) => {
redis::cmd("LPUSH")
.arg(key)
.arg(payload)
.query_async(c)
.await
}
RedisConnection::Cluster(c) => {
redis::cmd("LPUSH")
.arg(key)
.arg(payload)
.query_async(c)
.await
}
};
result.map_err(ExportError::from)
}
/// 推送队列是否有人消费；由 `AppState` 在构建时从 `[push] enabled` 写入。
///
/// 放在 `AppRedis` 上是因为它是唯一"已经流到每个入队点"的句柄：否则要么
/// 每个业务 crate 都拿一次配置，要么关闭推送的部署会无限堆积死任务。
pub fn set_push_enabled(&mut self, enabled: bool) {
self.push_enabled = enabled;
}
#[must_use]
pub fn push_enabled(&self) -> bool {
self.push_enabled
}
/// `LPUSH` 后把列表裁到 `max_len`，超出部分丢最旧的任务。
pub async fn lpush_capped_bytes(
&self,
key: &str,
payload: &[u8],
max_len: i64,
) -> Result<i64, ExportError> {
let mut conn = self.connection_clone();
let result: Result<i64, redis::RedisError> = match &mut conn {
RedisConnection::Standalone(c) => {
let len: i64 = redis::cmd("LPUSH")
.arg(key)
.arg(payload)
.query_async(c)
.await?;
if len > max_len {
let _: Result<redis::Value, redis::RedisError> = redis::cmd("LTRIM")
.arg(key)
.arg(0)
.arg(max_len - 1)
.query_async(c)
.await;
}
Ok(len)
}
RedisConnection::Cluster(c) => {
let len: i64 = redis::cmd("LPUSH")
.arg(key)
.arg(payload)
.query_async(c)
.await?;
if len > max_len {
let _: Result<redis::Value, redis::RedisError> = redis::cmd("LTRIM")
.arg(key)
.arg(0)
.arg(max_len - 1)
.query_async(c)
.await;
}
Ok(len)
}
};
result.map_err(ExportError::from)
}
/// 阻塞式右端出队（`BRPOP key timeout`）；超时返回 `None`。
///
/// 走独立连接而不是连接池：BRPOP 会独占 socket 整个超时窗口，池化连接被
/// 它占住会让同一进程里的其它 Redis 调用一起排长队。
pub async fn brpop_bytes(
&self,
key: &str,
timeout_secs: u64,
) -> Result<Option<Vec<u8>>, ExportError> {
let mut conn = self.blocking_consumer_connection(key).await?;
let result: Option<(String, Vec<u8>)> = redis::cmd("BRPOP")
.arg(key)
.arg(timeout_secs.max(1))
.query_async(&mut conn)
.await
.map_err(ExportError::from)?;
Ok(result.map(|(_, payload)| payload))
}
/// 打开一条专供阻塞命令的连接。
///
/// cluster 下 BRPOP 落在持有该 key slot 的 master 上，否则会收到 MOVED；
/// standalone/sentinel 直接用首个节点。
async fn blocking_consumer_connection(
&self,
key: &str,
) -> Result<MultiplexedConnection, ExportError> {
let client = if self.is_cluster() {
let node_info = self.node_info.clone().ok_or_else(|| {
ExportError::InvalidConfig("cluster connection info is unavailable".into())
})?;
let mut conn = self.connection_clone();
let RedisConnection::Cluster(cluster) = &mut conn else {
return Err(ExportError::InvalidConfig(
"cluster connection missing".into(),
));
};
let slot: u16 = redis::cmd("CLUSTER")
.arg("KEYSLOT")
.arg(key)
.query_async(cluster)
.await
.map_err(ExportError::from)?;
let slots: Value = redis::cmd("CLUSTER")
.arg("SLOTS")
.query_async(cluster)
.await
.map_err(ExportError::from)?;
let (host, port) = cluster_master_for_slot(&slots, slot)
.ok_or_else(|| ExportError::InvalidConfig(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(ExportError::InvalidConfig(
"unix sockets are not supported for cluster connections".into(),
));
}
_ => {
return Err(ExportError::InvalidConfig(
"unsupported connection address type for cluster connections".into(),
));
}
};
redis::Client::open(node_info.set_addr(addr)).map_err(ExportError::from)?
} else {
let url = self
.config
.url()
.split(',')
.map(str::trim)
.find(|node| !node.is_empty())
.ok_or_else(|| ExportError::InvalidConfig("redis url contains no nodes".into()))?;
redis::Client::open(connection_info(url, &self.config)?).map_err(ExportError::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(ExportError::from)
}
/// `SET key 1 NX EX ttl`; returns false when the key already exists
/// (already consumed / inside an active window) or redis fails.
pub async fn set_nx_ex(&self, key: &str, ttl_secs: u64) -> bool {
let mut conn = self.connection_clone();
let result: Result<Option<String>, redis::RedisError> = match &mut conn {
RedisConnection::Standalone(c) => {
redis::cmd("SET")
.arg(key)
.arg("1")
.arg("NX")
.arg("EX")
.arg(ttl_secs.max(1))
.query_async(c)
.await
}
RedisConnection::Cluster(c) => {
redis::cmd("SET")
.arg(key)
.arg("1")
.arg("NX")
.arg("EX")
.arg(ttl_secs.max(1))
.query_async(c)
.await
}
};
matches!(result, Ok(Some(_)))
}
/// `INCR key` with a window TTL set on first increment; returns the count.
pub async fn incr_window(&self, key: &str, window_secs: u64) -> i64 {
let mut conn = self.connection_clone();
let result: Result<i64, redis::RedisError> = match &mut conn {
RedisConnection::Standalone(c) => {
let count: i64 = redis::cmd("INCR")
.arg(key)
.query_async(c)
.await
.unwrap_or(0);
if count == 1 {
let _: Result<redis::Value, redis::RedisError> = redis::cmd("EXPIRE")
.arg(key)
.arg(window_secs)
.query_async(c)
.await;
}
Ok(count)
}
RedisConnection::Cluster(c) => {
let count: i64 = redis::cmd("INCR")
.arg(key)
.query_async(c)
.await
.unwrap_or(0);
if count == 1 {
let _: Result<redis::Value, redis::RedisError> = redis::cmd("EXPIRE")
.arg(key)
.arg(window_secs)
.query_async(c)
.await;
}
Ok(count)
}
};
result.unwrap_or(0)
}
/// Rolling-window quota check; `max == 0` disables the quota.
pub async fn quota_allow(&self, key: &str, max: i64, window_secs: u64) -> bool {
if max <= 0 {
return true;
}
self.incr_window(key, window_secs).await <= max
}
/// 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)?;
// BRPOP 之类的阻塞命令会持有连接直到超时；必须关闭响应超时，
// 否则默认超时会让阻塞读被误判为 timeout 错误。
let conn_config = redis::AsyncConnectionConfig::new()
.set_response_timeout(None)
.set_connection_timeout(Some(self.config.connection_timeout()));
let connection = client
.get_multiplexed_async_connection_with_config(&conn_config)
.await
.map_err(ExportError::from)?;
connections.push(connection);
}
Ok(connections)
}
/// Returns `(slot_start, slot_end, master_index)` triples where `master_index`
/// matches the position of the owning connection in [`cluster_master_connections`].
///
/// Empty for standalone connections. The ordering mirrors the dedup logic used
/// when collecting master addresses, so callers can map a slot to the right
/// pre-established connection without reconnecting.
pub async fn cluster_slot_ranges(&self) -> Result<Vec<(i64, i64, usize)>, 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)?;
Ok(cluster_slot_ranges(&slots))
}
}
/// 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.
/// Resolves the master node that owns `slot` straight from `CLUSTER SLOTS`.
///
/// Unlike [`cluster_master_addrs`] this needs no de-duplication: the range that
/// contains the slot names its own master in the same group.
fn cluster_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
}
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, config: &RedisConfig) -> 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()))?;
let url = if config.mode() == RedisMode::Cluster
&& !(first.starts_with("redis://") || first.starts_with("rediss://"))
{
format!("redis://{first}")
} else {
first.to_string()
};
let info = redis::Client::open(url)?.get_connection_info().clone();
let mut settings = info.redis_settings().clone();
if let Some(database) = config.database {
settings = settings.set_db(i64::from(database));
}
if let Some(username) = config.username.as_deref().filter(|value| !value.is_empty()) {
settings = settings.set_username(username);
}
if let Some(password) = config.password.as_deref().filter(|value| !value.is_empty()) {
settings = settings.set_password(password);
}
Ok(info.set_redis_settings(settings))
}
macro_rules! dispatch {
($conn:expr, |$c:ident| $body:expr) => {
match &mut $conn {
RedisConnection::Standalone($c) => $body,
RedisConnection::Cluster($c) => $body,
}
};
}
pub async fn get_string(redis: &AppRedis, key: &str) -> Result<Option<String>, redis::RedisError> {
let mut conn = redis.connection_clone();
dispatch!(conn, |c| redis::cmd("GET").arg(key).query_async(c).await)
}
pub async fn set_string(
redis: &AppRedis,
key: &str,
value: &str,
ttl_secs: u64,
) -> Result<(), redis::RedisError> {
let mut conn = redis.connection_clone();
dispatch!(conn, |c| {
redis::cmd("SET")
.arg(key)
.arg(value)
.arg("EX")
.arg(ttl_secs)
.exec_async(c)
.await
})
}
pub async fn set_nx_ex(
redis: &AppRedis,
key: &str,
value: &str,
ttl_secs: u64,
) -> Result<bool, redis::RedisError> {
let mut conn = redis.connection_clone();
dispatch!(conn, |c| {
redis::cmd("SET")
.arg(key)
.arg(value)
.arg("NX")
.arg("EX")
.arg(ttl_secs)
.query_async(c)
.await
})
}
pub async fn delete(redis: &AppRedis, key: &str) -> Result<(), redis::RedisError> {
let mut conn = redis.connection_clone();
dispatch!(conn, |c| redis::cmd("DEL").arg(key).exec_async(c).await)
}
const INCR_TOUCH_SCRIPT: &str = r#"local c = redis.call('INCR', KEYS[1])
redis.call('EXPIRE', KEYS[1], ARGV[1])
return c"#;
/// `INCR` 并把存活时间重置为 `ttl_secs`，返回递增后的计数。
///
/// 与 [`incr_expire`] 的区别是**每次**都续期：那个是为限流窗口写的（窗口
/// 从第一次计数开始算），这个是给心跳占位写的（每次心跳都要把到期时间往后推）。
pub async fn incr_touch(
redis: &AppRedis,
key: &str,
ttl_secs: u64,
) -> Result<i64, redis::RedisError> {
let mut conn = redis.connection_clone();
dispatch!(conn, |c| {
redis::cmd("EVAL")
.arg(INCR_TOUCH_SCRIPT)
.arg(1)
.arg(key)
.arg(ttl_secs.max(1))
.query_async(c)
.await
})
}
const DECR_DELETE_SCRIPT: &str = r#"local c = redis.call('DECR', KEYS[1])
if c <= 0 then
redis.call('DEL', KEYS[1])
return 0
end
return c"#;
/// 引用计数递减；归零时删键，返回剩余计数。
///
/// 一个用户可以同时开多个连接（多标签页、多设备同源会话），而"在线"是以键
/// 存在与否表达的二值状态，所以计数减到 0 时必须把键删掉——否则最后一条连接
/// 关闭后用户会一直显示在线，推送也就一直跳过。
///
/// 递减与删除放在一段脚本里：分开做的话，两次调用之间其它连接的心跳可能把
/// 计数从 0 又推到 1，用户就会永远停在"在线"。
pub async fn decr_delete(redis: &AppRedis, key: &str) -> Result<i64, redis::RedisError> {
let mut conn = redis.connection_clone();
dispatch!(conn, |c| {
redis::cmd("EVAL")
.arg(DECR_DELETE_SCRIPT)
.arg(1)
.arg(key)
.query_async(c)
.await
})
}
pub async fn exists(redis: &AppRedis, key: &str) -> Result<bool, redis::RedisError> {
let mut conn = redis.connection_clone();
dispatch!(conn, |c| redis::cmd("EXISTS").arg(key).query_async(c).await)
}
pub async fn getdel_string(
redis: &AppRedis,
key: &str,
) -> Result<Option<String>, redis::RedisError> {
let mut conn = redis.connection_clone();
dispatch!(conn, |c| redis::cmd("GETDEL").arg(key).query_async(c).await)
}
const INCR_EXPIRE_SCRIPT: &str = r#"local c = redis.call('INCR', KEYS[1])
if c == 1 then
redis.call('EXPIRE', KEYS[1], ARGV[1])
end
return c"#;
pub async fn incr_expire(
redis: &AppRedis,
key: &str,
ttl_secs: u64,
) -> Result<i64, redis::RedisError> {
let mut conn = redis.connection_clone();
dispatch!(conn, |c| {
redis::cmd("EVAL")
.arg(INCR_EXPIRE_SCRIPT)
.arg(1)
.arg(key)
.arg(ttl_secs)
.query_async(c)
.await
})
}
/// Extracts `(start, end, master_index)` for every slot range.
///
/// `master_index` counts unique master addresses in first-seen order, matching
/// [`AppRedis::cluster_master_connections`].
fn cluster_slot_ranges(slots: &Value) -> Vec<(i64, i64, usize)> {
let Value::Array(groups) = slots else {
return Vec::new();
};
let mut seen: std::collections::HashMap<(String, u16), usize> =
std::collections::HashMap::new();
let mut ranges = Vec::new();
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;
};
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;
};
let addr = (String::from_utf8_lossy(host).to_string(), *port as u16);
let next_index = seen.len();
let index = *seen.entry(addr).or_insert(next_index);
ranges.push((*start, *end, index));
}
ranges
}