main
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 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 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
"#;
static RELEASE: LazyLock<Script> = LazyLock::new(|| Script::new(RELEASE_CODE));
static REFRESH: LazyLock<Script> = LazyLock::new(|| Script::new(REFRESH_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 LOCK_MANY: LazyLock<Script> = LazyLock::new(|| Script::new(LOCK_MANY_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 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 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 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)
}
}