main
use config::redis::{RedisConfig, RedisMode};
use crate::conn::{BlockingConnection, CacheConn, CacheConnection};
use crate::error::CacheError;
pub type RedisConnection = CacheConnection;
/// Raw-key facade over the shared Redis pool.
///
/// Unlike [`crate::Cache`] it does not apply `redis.key_prefix`, because queue,
/// quota and lock keys are consumed verbatim by other processes and workers.
#[derive(Clone)]
pub struct AppRedis {
conn: CacheConn,
}
impl AppRedis {
pub async fn connect(config: RedisConfig) -> Result<Self, CacheError> {
Ok(Self::from_conn(CacheConn::init(&config).await?))
}
pub fn from_conn(conn: CacheConn) -> Self {
Self { conn }
}
pub fn config(&self) -> &RedisConfig {
self.conn.config()
}
pub fn mode(&self) -> RedisMode {
self.conn.mode()
}
pub fn is_cluster(&self) -> bool {
self.conn.is_cluster()
}
pub fn connection(&self) -> &RedisConnection {
self.conn.connection()
}
pub fn connection_clone(&self) -> RedisConnection {
self.conn.connection_clone()
}
pub async fn ping(&self) -> Result<(), CacheError> {
self.conn.ping().await
}
pub async fn publish(&self, channel: &str, payload: &str) -> Result<(), CacheError> {
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(CacheError::from)
}
/// 列表左端入队（`LPUSH`）；与 `AppRedis::publish` 一样按 standalone/cluster
/// 双分支实现，payload 是二进制安全的字节串。
pub async fn lpush_bytes(&self, key: &str, payload: &[u8]) -> Result<i64, CacheError> {
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(CacheError::from)
}
/// 队列积压深度（`LLEN`）；供指标采样使用。
pub async fn llen(&self, key: &str) -> Result<i64, CacheError> {
let mut conn = self.connection_clone();
let result: Result<i64, redis::RedisError> = 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(CacheError::from)
}
/// `LPUSH` 后把列表裁到 `max_len`，超出部分丢最旧的任务。
pub async fn lpush_capped_bytes(
&self,
key: &str,
payload: &[u8],
max_len: i64,
) -> Result<i64, CacheError> {
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(CacheError::from)
}
/// 阻塞式右端出队（`BRPOP key timeout`）；超时返回 `None`。
///
/// 每次调用新建并释放一条独立连接。需要长期轮询的消费者应改用
/// [`AppRedis::blocking`]，复用同一条 socket。
pub async fn brpop_bytes(
&self,
key: &str,
timeout_secs: u64,
) -> Result<Option<Vec<u8>>, CacheError> {
let mut conn = self.conn.blocking(key).await?;
conn.brpop_bytes(key, timeout_secs).await
}
/// Opens a long-lived blocking connection for a queue consumer.
pub async fn blocking(&self, key: &str) -> Result<BlockingConnection, CacheError> {
self.conn.blocking(key).await
}
/// `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
}
}
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 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<i64, redis::RedisError> {
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<i64, redis::RedisError> {
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<bool, redis::RedisError> {
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<bool, redis::RedisError> {
let mut conn = redis.connection_clone();
dispatch!(conn, |c| redis::cmd("EXISTS").arg(key).query_async(c).await)
}