ef0b83a
use std::sync::LazyLock;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use redis::Script;
use crate::error::{CacheError, CacheResult};
use crate::keys;
use crate::kv::{Cache, millis};
const RELEASE_CODE: &str = r#"
if redis.call('GET', KEYS[1]) == ARGV[1] then
return redis.call('DEL', KEYS[1])
end
return 0
"#;
const REFRESH_CODE: &str = r#"
if redis.call('GET', KEYS[1]) == ARGV[1] then
return redis.call('PEXPIRE', KEYS[1], ARGV[2])
end
return 0
"#;
const COMPARE_AND_SET_CODE: &str = r#"
local current = redis.call('GET', KEYS[1])
local matches = current == ARGV[1]
if ARGV[4] == '1' and current == false then
matches = true
end
if not matches then
return 0
end
redis.call('SET', KEYS[1], ARGV[2], 'PX', ARGV[3])
return 1
"#;
const INCR_EX_CODE: &str = r#"
local value = redis.call('INCR', KEYS[1])
if value == 1 then
redis.call('PEXPIRE', KEYS[1], ARGV[1])
end
return value
"#;
/// 验证码原子消费：值匹配时同时删除验证码与尝试计数，匹配失败不动键。
const CONSUME_MATCHING_CODE: &str = r#"
if redis.call('GET', KEYS[1]) == ARGV[1] then
redis.call('DEL', KEYS[1])
redis.call('DEL', KEYS[2])
return 1
end
return 0
"#;
const GETSET_EX_CODE: &str = r#"
local previous = redis.call('GET', KEYS[1])
redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[2])
return previous
"#;
const LPUSH_CAPPED_CODE: &str = r#"
local length = redis.call('LPUSH', KEYS[1], ARGV[1])
local limit = tonumber(ARGV[2])
if limit > 0 and length > limit then
redis.call('LTRIM', KEYS[1], 0, limit - 1)
return limit
end
return length
"#;
const HSET_EX_CODE: &str = r#"
redis.call('HSET', KEYS[1], ARGV[1], ARGV[2])
local ttl = tonumber(ARGV[3])
if ttl > 0 then
redis.call('PEXPIRE', KEYS[1], ttl)
end
return 1
"#;
const LOCK_MANY_CODE: &str = r#"
for index = 1, #KEYS do
if redis.call('SET', KEYS[index], ARGV[1], 'PX', ARGV[2], 'NX') == false then
for released = 1, index - 1 do
redis.call('DEL', KEYS[released])
end
return 0
end
end
return 1
"#;
const RATE_LIMIT_CODE: &str = r#"
-- 时间取自 Redis 服务器（TIME），避免多实例应用时钟漂移破坏窗口。
local t = redis.call('TIME')
local now = t[1] * 1000 + math.floor(t[2] / 1000)
local window = tonumber(ARGV[1])
local limit = tonumber(ARGV[2])
redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, now - window)
local count = redis.call('ZCARD', KEYS[1])
if count >= limit then
return {0, count}
end
redis.call('ZADD', KEYS[1], now, ARGV[3])
redis.call('PEXPIRE', KEYS[1], window)
return {1, count + 1}
"#;
static RELEASE: LazyLock<Script> = LazyLock::new(|| Script::new(RELEASE_CODE));
static REFRESH: LazyLock<Script> = LazyLock::new(|| Script::new(REFRESH_CODE));
static COMPARE_AND_SET: LazyLock<Script> = LazyLock::new(|| Script::new(COMPARE_AND_SET_CODE));
static INCR_EX: LazyLock<Script> = LazyLock::new(|| Script::new(INCR_EX_CODE));
static CONSUME_MATCHING: LazyLock<Script> = LazyLock::new(|| Script::new(CONSUME_MATCHING_CODE));
static GETSET_EX: LazyLock<Script> = LazyLock::new(|| Script::new(GETSET_EX_CODE));
static LPUSH_CAPPED: LazyLock<Script> = LazyLock::new(|| Script::new(LPUSH_CAPPED_CODE));
static HSET_EX: LazyLock<Script> = LazyLock::new(|| Script::new(HSET_EX_CODE));
static LOCK_MANY: LazyLock<Script> = LazyLock::new(|| Script::new(LOCK_MANY_CODE));
static RATE_LIMIT: LazyLock<Script> = LazyLock::new(|| Script::new(RATE_LIMIT_CODE));
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RateLimit {
pub allowed: bool,
pub count: i64,
}
#[derive(Clone)]
pub struct LuaScripts {
cache: Cache,
}
pub fn unix_millis() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|elapsed| elapsed.as_millis() as i64)
.unwrap_or_default()
}
impl LuaScripts {
pub(crate) fn new(cache: Cache) -> Self {
Self { cache }
}
pub fn cache(&self) -> &Cache {
&self.cache
}
pub async fn release_lock(&self, key: &str, token: &str) -> CacheResult<bool> {
let mut conn = self.cache.conn().clone();
RELEASE
.key(key)
.arg(token)
.invoke_async(&mut conn)
.await
.map_err(CacheError::from)
}
pub async fn compare_and_delete(&self, key: &str, expected: &[u8]) -> CacheResult<bool> {
let mut conn = self.cache.conn().clone();
RELEASE
.key(key)
.arg(expected)
.invoke_async(&mut conn)
.await
.map_err(CacheError::from)
}
pub async fn refresh_lock(&self, key: &str, token: &str, ttl: Duration) -> CacheResult<bool> {
let mut conn = self.cache.conn().clone();
REFRESH
.key(key)
.arg(token)
.arg(millis(ttl))
.invoke_async(&mut conn)
.await
.map_err(CacheError::from)
}
pub async fn compare_and_set(
&self,
key: &str,
expected: &[u8],
value: &[u8],
ttl: Duration,
allow_missing: bool,
) -> CacheResult<bool> {
let mut conn = self.cache.conn().clone();
COMPARE_AND_SET
.key(key)
.arg(expected)
.arg(value)
.arg(millis(ttl))
.arg(if allow_missing { "1" } else { "0" })
.invoke_async(&mut conn)
.await
.map_err(CacheError::from)
}
pub async fn incr_ex(&self, key: &str, ttl: Duration) -> CacheResult<i64> {
let mut conn = self.cache.conn().clone();
INCR_EX
.key(key)
.arg(millis(ttl))
.invoke_async(&mut conn)
.await
.map_err(CacheError::from)
}
/// 值匹配时原子删除 `key` 与 `key+""` 附带键（验证码+尝试计数），返回是否消费成功。
pub async fn consume_matching(&self, key: &str, expected: &[u8]) -> CacheResult<bool> {
let mut conn = self.cache.conn().clone();
CONSUME_MATCHING
.key(key)
.key(format!("{key}:attempts"))
.arg(expected)
.invoke_async(&mut conn)
.await
.map_err(CacheError::from)
}
pub async fn getset_ex(
&self,
key: &str,
value: &[u8],
ttl: Duration,
) -> CacheResult<Option<Vec<u8>>> {
let mut conn = self.cache.conn().clone();
GETSET_EX
.key(key)
.arg(value)
.arg(millis(ttl))
.invoke_async(&mut conn)
.await
.map_err(CacheError::from)
}
pub async fn lpush_capped(&self, key: &str, value: &[u8], limit: usize) -> CacheResult<i64> {
let mut conn = self.cache.conn().clone();
LPUSH_CAPPED
.key(key)
.arg(value)
.arg(limit)
.invoke_async(&mut conn)
.await
.map_err(CacheError::from)
}
pub async fn hset_ex(
&self,
key: &str,
field: &str,
value: &[u8],
ttl: Option<Duration>,
) -> CacheResult<bool> {
let mut conn = self.cache.conn().clone();
HSET_EX
.key(key)
.arg(field)
.arg(value)
.arg(ttl.map(millis).unwrap_or_default())
.invoke_async(&mut conn)
.await
.map_err(CacheError::from)
}
pub async fn lock_many(
&self,
keys: &[String],
token: &str,
ttl: Duration,
) -> CacheResult<bool> {
if self.cache.is_cluster() && !keys::same_slot(keys) {
return Err(CacheError::unsupported(
"lock_many",
"every key must carry the same {hashtag} on Redis Cluster",
));
}
let mut invocation = LOCK_MANY.prepare_invoke();
for key in keys {
invocation.key(key);
}
invocation.arg(token).arg(millis(ttl));
let mut conn = self.cache.conn().clone();
invocation
.invoke_async(&mut conn)
.await
.map_err(CacheError::from)
}
pub async fn rate_limit(
&self,
key: &str,
limit: usize,
window: Duration,
member: &str,
) -> CacheResult<RateLimit> {
let mut conn = self.cache.conn().clone();
let reply: Vec<i64> = RATE_LIMIT
.key(key)
.arg(millis(window))
.arg(limit)
.arg(member)
.invoke_async(&mut conn)
.await
.map_err(CacheError::from)?;
Ok(RateLimit {
allowed: reply.first().copied().unwrap_or_default() == 1,
count: reply.get(1).copied().unwrap_or_default(),
})
}
}