030e59e
zhenyi
refactor(lib): migrate storage and security tier to central lib/config (auth, cache, db, email, jwt, payment, queue, rate_limit, storage)
030e59e
mod conn;
pub use conn::RedisConnClone;
use crate::metrics;
use config::cache::{AppRedisMode, AppRedisSettings};
use conn::RedisConn;
use redis::{
AsyncCommands, ConnectionAddr, IntoConnectionInfo, RedisError as RedisDriverError,
cluster::ClusterClientBuilder, sentinel::SentinelClientBuilder,
};
use tracing::debug;
const CLEAR_BATCH_SIZE: usize = 256;
const COMPARE_AND_SET_BYTES_SCRIPT: &str = r#"
if redis.call('GET', KEYS[1]) ~= ARGV[1] then
return 0
end
if tonumber(ARGV[3]) > 0 then
redis.call('SET', KEYS[1], ARGV[2], 'EX', ARGV[3])
else
redis.call('SET', KEYS[1], ARGV[2])
end
return 1
"#;
const COMPARE_AND_DELETE_SCRIPT: &str = r#"
if redis.call('GET', KEYS[1]) ~= ARGV[1] then
return 0
end
return redis.call('DEL', KEYS[1])
"#;
const REFRESH_LOCK_SCRIPT: &str = r#"
if redis.call('GET', KEYS[1]) ~= ARGV[1] then
return 0
end
return redis.call('EXPIRE', KEYS[1], ARGV[2])
"#;
#[derive(Debug, thiserror::Error)]
pub enum RedisError {
#[error("redis error: {0}")]
Redis(#[from] RedisDriverError),
#[error("no redis url configured")]
NoUrl,
}
pub struct RedisCache {
conn: RedisConn,
key_prefix: String,
}
impl RedisCache {
pub async fn new(cfg: &AppRedisSettings) -> Result<Self, RedisError> {
let conn = match cfg.mode {
AppRedisMode::Single => {
let url = cfg.urls.first().ok_or(RedisError::NoUrl)?;
let client = redis::Client::open(url.as_str())?;
let c = client.get_multiplexed_async_connection().await?;
debug!(url = %url, "Redis single-node connected");
RedisConn::Single(c)
}
AppRedisMode::Cluster => {
let urls: Vec<&str> = cfg.urls.iter().map(|s| s.as_str()).collect();
let mut builder = ClusterClientBuilder::new(urls)
.connection_timeout(std::time::Duration::from_millis(
cfg.connection_timeout_ms.unwrap_or(2000),
))
.response_timeout(std::time::Duration::from_millis(
cfg.response_timeout_ms.unwrap_or(5000),
));
if let Some(ref password) = cfg.password
&& !password.is_empty()
{
builder = builder.password(password.as_str());
}
let client = builder.build()?;
let c = client.get_async_connection().await?;
debug!(nodes = ?cfg.urls, "Redis cluster connected");
RedisConn::Cluster(c)
}
AppRedisMode::Sentinel => {
let service_name = cfg
.urls
.first()
.and_then(|u| {
u.split('?').nth(1).and_then(|q| {
q.split('&').find_map(|kv| kv.strip_prefix("service_name="))
})
})
.unwrap_or("mymaster")
.to_string();
let addrs = cfg
.urls
.iter()
.filter_map(|u| {
u.split('?')
.next()
.unwrap_or(u)
.into_connection_info()
.ok()
.map(|info| info.addr().clone())
})
.collect::<Vec<ConnectionAddr>>();
let mut sentinel_builder = SentinelClientBuilder::new(
addrs,
&service_name,
redis::sentinel::SentinelServerType::Master,
)?
.set_client_to_redis_db(cfg.database.unwrap_or(0));
if let Some(ref pw) = cfg.password {
sentinel_builder = sentinel_builder.set_client_to_redis_password(pw.as_str());
}
let mut sentinel_client = sentinel_builder.build()?;
let client = sentinel_client.async_get_client().await?;
let c = client.get_multiplexed_async_connection().await?;
debug!(service = %service_name, nodes = ?cfg.urls, "Redis sentinel connected");
RedisConn::Sentinel(c)
}
};
Ok(Self {
conn,
key_prefix: "gta:cache:".to_string(),
})
}
pub fn conn(&mut self) -> &mut RedisConn {
&mut self.conn
}
pub async fn get_bytes(&self, key: &str) -> Result<Option<Vec<u8>>, RedisError> {
let op = "get_bytes";
let result: Result<Option<Vec<u8>>, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
Ok(conn.get(&*full_key).await?)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub(crate) async fn get_bytes_with_ttl(
&self,
key: &str,
) -> Result<Option<(Vec<u8>, Option<u64>)>, RedisError> {
let op = "get_bytes_with_ttl";
let result = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let Some(value) = conn.get(&*full_key).await? else {
return Ok(None);
};
let ttl_millis: i64 = conn.pttl(&*full_key).await?;
let ttl_millis = match ttl_millis {
-1 => None,
value if value >= 0 => Some(value as u64),
_ => Some(0),
};
Ok(Some((value, ttl_millis)))
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn set_bytes(
&self,
key: &str,
value: Vec<u8>,
ttl_seconds: u64,
) -> Result<(), RedisError> {
let op = "set_bytes";
let result: Result<(), RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
if ttl_seconds > 0 {
let _: () = conn.set_ex(&*full_key, value, ttl_seconds).await?;
} else {
let _: () = conn.set(&*full_key, value).await?;
}
Ok(())
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn compare_and_set_bytes(
&self,
key: &str,
expected: &[u8],
value: Vec<u8>,
ttl_seconds: u64,
) -> Result<bool, RedisError> {
let op = "compare_and_set_bytes";
let result: Result<bool, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let changed: i64 = redis::Script::new(COMPARE_AND_SET_BYTES_SCRIPT)
.key(full_key)
.arg(expected)
.arg(value)
.arg(ttl_seconds)
.invoke_async(&mut conn)
.await?;
Ok(changed == 1)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn delete(&self, key: &str) -> Result<(), RedisError> {
let op = "delete";
let result: Result<(), RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let _: () = conn.del(&*full_key).await?;
Ok(())
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn set_nx(
&self,
key: &str,
value: Vec<u8>,
ttl_seconds: u64,
) -> Result<bool, RedisError> {
let op = "set_nx";
let result: Result<bool, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let set: Option<()> = redis::cmd("SET")
.arg(&*full_key)
.arg(value)
.arg("NX")
.arg("EX")
.arg(ttl_seconds as i64)
.query_async(&mut conn)
.await?;
Ok(set.is_some())
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn compare_and_delete(&self, key: &str, expected: &[u8]) -> Result<bool, RedisError> {
let op = "compare_and_delete";
let result: Result<bool, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let deleted: i64 = redis::Script::new(COMPARE_AND_DELETE_SCRIPT)
.key(full_key)
.arg(expected)
.invoke_async(&mut conn)
.await?;
Ok(deleted == 1)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn refresh_lock(
&self,
key: &str,
value: &[u8],
ttl_seconds: u64,
) -> Result<bool, RedisError> {
let op = "refresh_lock";
let result: Result<bool, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let refreshed: i64 = redis::Script::new(REFRESH_LOCK_SCRIPT)
.key(full_key)
.arg(value)
.arg(ttl_seconds as i64)
.invoke_async(&mut conn)
.await?;
Ok(refreshed == 1)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn rpush(&self, key: &str, value: Vec<u8>) -> Result<u64, RedisError> {
let op = "rpush";
let result: Result<u64, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let len: u64 = conn.rpush(&*full_key, value).await?;
Ok(len)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn lpush(&self, key: &str, value: Vec<u8>) -> Result<u64, RedisError> {
let op = "lpush";
let result: Result<u64, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let len: u64 = conn.lpush(&*full_key, value).await?;
Ok(len)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn lpop(&self, key: &str) -> Result<Option<Vec<u8>>, RedisError> {
let op = "lpop";
let result: Result<Option<Vec<u8>>, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
Ok(redis::cmd("LPOP")
.arg(&*full_key)
.query_async(&mut conn)
.await?)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn lrem(&self, key: &str, count: i64, member: &[u8]) -> Result<u64, RedisError> {
let op = "lrem";
let result: Result<u64, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let removed: u64 = conn.lrem(&*full_key, count as isize, member).await?;
Ok(removed)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn llen(&self, key: &str) -> Result<u64, RedisError> {
let op = "llen";
let result: Result<u64, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let len: u64 = conn.llen(&*full_key).await?;
Ok(len)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn getdel_bytes(&self, key: &str) -> Result<Option<Vec<u8>>, RedisError> {
let op = "getdel_bytes";
let result: Result<Option<Vec<u8>>, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let value: Option<Vec<u8>> = redis::cmd("GETDEL")
.arg(&*full_key)
.query_async(&mut conn)
.await?;
Ok(value)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn incr_with_ttl(&self, key: &str, ttl_seconds: u64) -> Result<u64, RedisError> {
let op = "incr_with_ttl";
let result: Result<u64, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let (count,): (u64,) = redis::pipe()
.atomic()
.incr(&*full_key, 1u64)
.expire(&*full_key, ttl_seconds as i64)
.ignore()
.query_async(&mut conn)
.await?;
Ok(count)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn get_count(&self, key: &str) -> Result<u64, RedisError> {
let op = "get_count";
let result: Result<u64, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let value: Option<u64> = conn.get(&*full_key).await?;
Ok(value.unwrap_or(0))
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn add_to_set_with_ttl(
&self,
key: &str,
member: &str,
ttl_seconds: u64,
) -> Result<(), RedisError> {
let op = "add_to_set_with_ttl";
let result: Result<(), RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let _: () = redis::pipe()
.atomic()
.sadd(&*full_key, member)
.expire(&*full_key, ttl_seconds as i64)
.ignore()
.query_async(&mut conn)
.await?;
Ok(())
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn remove_from_set(&self, key: &str, member: &str) -> Result<(), RedisError> {
let op = "remove_from_set";
let result: Result<(), RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let _: () = conn.srem(&*full_key, member).await?;
Ok(())
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn set_members(&self, key: &str) -> Result<Vec<String>, RedisError> {
let op = "set_members";
let result: Result<Vec<String>, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let members: Vec<String> = conn.smembers(&*full_key).await?;
Ok(members)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn delete_set_members(&self, key: &str) -> Result<(), RedisError> {
let op = "delete_set_members";
let result: Result<(), RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let _: () = conn.del(&*full_key).await?;
Ok(())
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn zadd(
&self,
key: &str,
member: &str,
score: f64,
ttl_seconds: u64,
) -> Result<(), RedisError> {
let op = "zadd";
let result: Result<(), RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let _: () = redis::pipe()
.atomic()
.zadd(&full_key, member, score)
.expire(&full_key, ttl_seconds as i64)
.ignore()
.query_async(&mut conn)
.await?;
Ok(())
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn zremrange_by_score_lt(
&self,
key: &str,
max_score: f64,
) -> Result<u64, RedisError> {
let op = "zremrange_by_score_lt";
let result: Result<u64, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let removed: u64 = conn.zrembyscore(&*full_key, "-inf", max_score).await?;
Ok(removed)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn zcard(&self, key: &str) -> Result<u64, RedisError> {
let op = "zcard";
let result: Result<u64, RedisError> = async {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let count: u64 = conn.zcard(&*full_key).await?;
Ok(count)
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub async fn clear(&self) -> Result<(), RedisError> {
let op = "clear";
let result: Result<(), RedisError> = async {
let pattern = format!("{}*", self.key_prefix);
let mut scan_conn = self.spawn_conn();
let mut delete_conn = self.spawn_conn();
let mut keys: Vec<String> = Vec::with_capacity(CLEAR_BATCH_SIZE);
{
let mut iter = scan_conn.scan_match::<_, String>(&*pattern).await?;
while let Some(key) = iter.next_item().await {
keys.push(key?);
if keys.len() == CLEAR_BATCH_SIZE {
let _: usize = delete_conn.unlink(&keys).await?;
keys.clear();
}
}
}
if !keys.is_empty() {
let _: usize = delete_conn.unlink(&keys).await?;
}
Ok(())
}
.await;
metrics::record_redis_op(op, result.is_ok());
result
}
pub fn connection(&self) -> RedisConnClone {
self.spawn_conn()
}
fn prefixed(&self, key: &str) -> String {
format!("{}{}", self.key_prefix, key)
}
fn spawn_conn(&self) -> RedisConnClone {
match &self.conn {
RedisConn::Single(c) => RedisConnClone::Single(c.clone()),
RedisConn::Cluster(c) => RedisConnClone::Cluster(c.clone()),
RedisConn::Sentinel(c) => RedisConnClone::Sentinel(c.clone()),
}
}
}