d5939db
mod conn;
pub use conn::RedisConnClone;
use crate::config::redis::{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
"#;
#[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 full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
Ok(conn.get(&*full_key).await?)
}
pub(crate) async fn get_bytes_with_ttl(
&self,
key: &str,
) -> Result<Option<(Vec<u8>, Option<u64>)>, RedisError> {
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)))
}
pub async fn set_bytes(
&self,
key: &str,
value: Vec<u8>,
ttl_seconds: u64,
) -> Result<(), RedisError> {
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(())
}
pub async fn compare_and_set_bytes(
&self,
key: &str,
expected: &[u8],
value: Vec<u8>,
ttl_seconds: u64,
) -> Result<bool, RedisError> {
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)
}
pub async fn delete(&self, key: &str) -> Result<(), RedisError> {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let _: () = conn.del(&*full_key).await?;
Ok(())
}
/// Atomically fetch-and-delete a key (GETDEL). Returns the previous value
/// if present, enabling single-use consumption of one-time tokens.
pub async fn getdel_bytes(&self, key: &str) -> Result<Option<Vec<u8>>, RedisError> {
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)
}
pub async fn incr_with_ttl(&self, key: &str, ttl_seconds: u64) -> Result<u64, RedisError> {
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)
}
pub async fn get_count(&self, key: &str) -> Result<u64, RedisError> {
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))
}
pub async fn add_to_set_with_ttl(
&self,
key: &str,
member: &str,
ttl_seconds: u64,
) -> Result<(), RedisError> {
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(())
}
pub async fn remove_from_set(&self, key: &str, member: &str) -> Result<(), RedisError> {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let _: () = conn.srem(&*full_key, member).await?;
Ok(())
}
pub async fn set_members(&self, key: &str) -> Result<Vec<String>, RedisError> {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let members: Vec<String> = conn.smembers(&*full_key).await?;
Ok(members)
}
pub async fn delete_set_members(&self, key: &str) -> Result<(), RedisError> {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let _: () = conn.del(&*full_key).await?;
Ok(())
}
pub async fn zadd(
&self,
key: &str,
member: &str,
score: f64,
ttl_seconds: u64,
) -> Result<(), RedisError> {
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(())
}
pub async fn zremrange_by_score_lt(&self, key: &str, max_score: f64) -> Result<u64, RedisError> {
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)
}
pub async fn zcard(&self, key: &str) -> Result<u64, RedisError> {
let full_key = self.prefixed(key);
let mut conn = self.spawn_conn();
let count: u64 = conn.zcard(&*full_key).await?;
Ok(count)
}
pub async fn clear(&self) -> Result<(), RedisError> {
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(())
}
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()),
}
}
}