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, /// Round-robin cursor shared across clones of this handle. cursor: Arc, /// Connection parameters (auth, tls) of the first configured node, /// reused to open per-node connections for cluster-wide operations. node_info: Option, /// Web Push 是否启用;入队点用它避免往无人消费的队列里写任务。 push_enabled: bool, } impl AppRedis { pub async fn connect(config: RedisConfig) -> Result { 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::, _>>()?; 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 = 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 { let mut conn = self.connection_clone(); let result: Result = 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) } /// 队列积压深度(`LLEN`);供指标采样使用。 pub async fn llen(&self, key: &str) -> Result { let mut conn = self.connection_clone(); let result: Result = match &mut conn { RedisConnection::Standalone(c) => redis::cmd("LLEN").arg(key).query_async(c).await, RedisConnection::Cluster(c) => redis::cmd("LLEN").arg(key).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 { let mut conn = self.connection_clone(); let result: Result = 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::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::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>, ExportError> { let mut conn = self.blocking_consumer_connection(key).await?; let result: Option<(String, Vec)> = 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 { 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, 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 = 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::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::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, 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, 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 { 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, 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 { 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 ZSET_TOUCH_SCRIPT: &str = r#"local t = redis.call('TIME') local now = tonumber(t[1]) redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now) redis.call('ZADD', KEYS[1], now + tonumber(ARGV[2]), ARGV[1]) redis.call('EXPIRE', KEYS[1], tonumber(ARGV[2]) * 3) return 1"#; /// 以"带到期时刻的集合成员"登记一份占用并续期;成员分数是到期时刻, /// 过期成员在每次交互时被惰性清掉。 /// /// 计数型占位(INCR/DECR)在进程被强杀时会留下永不衰减的陈旧计数——只要 /// 还有任何心跳在续期键,陈旧计数就永远不会归零。每个连接一个成员、各自 /// 携带到期时刻的集合没有这个问题:强杀只是留下一个分数过期的成员,之后 /// 的任意一次 touch/release/读取都会把它清掉。 pub async fn zset_touch( redis: &AppRedis, key: &str, member: &str, ttl_secs: u64, ) -> Result<(), redis::RedisError> { let mut conn = redis.connection_clone(); let _: i64 = dispatch!(conn, |c| { redis::cmd("EVAL") .arg(ZSET_TOUCH_SCRIPT) .arg(1) .arg(key) .arg(member) .arg(ttl_secs.max(1)) .query_async(c) .await })?; Ok(()) } const ZSET_RELEASE_SCRIPT: &str = r#"local t = redis.call('TIME') local now = tonumber(t[1]) redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now) redis.call('ZREM', KEYS[1], ARGV[1]) local n = redis.call('ZCARD', KEYS[1]) if n == 0 then redis.call('DEL', KEYS[1]) end return n"#; /// 注销一份占用,返回剩余**存活**成员数;归零时删键。计数前先清过期成员。 pub async fn zset_release( redis: &AppRedis, key: &str, member: &str, ) -> Result { let mut conn = redis.connection_clone(); dispatch!(conn, |c| { redis::cmd("EVAL") .arg(ZSET_RELEASE_SCRIPT) .arg(1) .arg(key) .arg(member) .query_async(c) .await }) } const ZSET_LIVE_COUNT_SCRIPT: &str = r#"local t = redis.call('TIME') local now = tonumber(t[1]) redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now) return redis.call('ZCARD', KEYS[1])"#; /// 当前**存活**成员数(先清过期成员);键不存在时为 0。 pub async fn zset_live_count(redis: &AppRedis, key: &str) -> Result { let mut conn = redis.connection_clone(); dispatch!(conn, |c| { redis::cmd("EVAL") .arg(ZSET_LIVE_COUNT_SCRIPT) .arg(1) .arg(key) .query_async(c) .await }) } const TOUCH_EXPIRE_SCRIPT: &str = r#"if redis.call('EXISTS', KEYS[1]) == 1 then redis.call('EXPIRE', KEYS[1], ARGV[1]) return 1 end return 0"#; /// 键在则续期并返回 `true`;键不在返回 `false`,由调用方决定是否补写。 pub async fn touch_expire( redis: &AppRedis, key: &str, ttl_secs: u64, ) -> Result { let mut conn = redis.connection_clone(); let renewed: i64 = dispatch!(conn, |c| { redis::cmd("EVAL") .arg(TOUCH_EXPIRE_SCRIPT) .arg(1) .arg(key) .arg(ttl_secs.max(1)) .query_async(c) .await })?; Ok(renewed == 1) } pub async fn exists(redis: &AppRedis, key: &str) -> Result { 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, 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 { 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 }