use std::sync::Arc;
use redis::AsyncCommands;
use tracing::info;
use crate::server::config::StoreConfig;
use crate::server::error::AppError;
use super::{BatchOps, BoxFuture, KeyspaceOps, RawKvPair, StorageBackend};
pub struct RedisBackend {
conn: redis::aio::MultiplexedConnection,
}
impl RedisBackend {
pub async fn open(config: &StoreConfig) -> Result<Box<dyn StorageBackend>, AppError> {
let url = config
.redis_url
.as_deref()
.unwrap_or("redis://127.0.0.1:6379");
info!(url, "opening redis store");
let client =
redis::Client::open(url).map_err(|e| AppError::Store(format!("redis connect: {e}")))?;
let mut conn = client
.get_multiplexed_async_connection()
.await
.map_err(|e| AppError::Store(format!("redis connect: {e}")))?;
redis::cmd("PING")
.query_async::<String>(&mut conn)
.await
.map_err(|e| AppError::Store(format!("redis ping: {e}")))?;
Ok(Box::new(Self { conn }))
}
}
impl StorageBackend for RedisBackend {
fn keyspace(&self, name: &str) -> Result<(String, Arc<dyn KeyspaceOps>), AppError> {
Ok((
name.to_string(),
Arc::new(RedisKeyspace {
conn: self.conn.clone(),
prefix: format!("{name}:"),
}),
))
}
fn batch(&self) -> Box<dyn BatchOps> {
Box::new(RedisBatch {
conn: self.conn.clone(),
ops: Vec::new(),
})
}
fn persist(&self) -> BoxFuture<'_, Result<(), AppError>> {
Box::pin(async { Ok(()) })
}
}
struct RedisKeyspace {
conn: redis::aio::MultiplexedConnection,
prefix: String,
}
impl RedisKeyspace {
fn full_key(&self, key: &[u8]) -> Vec<u8> {
let mut fk = self.prefix.as_bytes().to_vec();
fk.extend_from_slice(key);
fk
}
}
impl KeyspaceOps for RedisKeyspace {
fn insert_raw(&self, key: Vec<u8>, value: Vec<u8>) -> BoxFuture<'_, Result<(), AppError>> {
Box::pin(async move {
let fk = self.full_key(&key);
let mut conn = self.conn.clone();
conn.set::<_, _, ()>(fk, value)
.await
.map_err(|e| AppError::Store(format!("redis SET: {e}")))?;
Ok(())
})
}
fn get_raw(&self, key: Vec<u8>) -> BoxFuture<'_, Result<Option<Vec<u8>>, AppError>> {
Box::pin(async move {
let fk = self.full_key(&key);
let mut conn = self.conn.clone();
let result: Option<Vec<u8>> = conn
.get(fk)
.await
.map_err(|e| AppError::Store(format!("redis GET: {e}")))?;
Ok(result)
})
}
fn remove(&self, key: Vec<u8>) -> BoxFuture<'_, Result<(), AppError>> {
Box::pin(async move {
let fk = self.full_key(&key);
let mut conn = self.conn.clone();
conn.del::<_, ()>(fk)
.await
.map_err(|e| AppError::Store(format!("redis DEL: {e}")))?;
Ok(())
})
}
fn contains_key(&self, key: Vec<u8>) -> BoxFuture<'_, Result<bool, AppError>> {
Box::pin(async move {
let fk = self.full_key(&key);
let mut conn = self.conn.clone();
let exists: bool = conn
.exists(fk)
.await
.map_err(|e| AppError::Store(format!("redis EXISTS: {e}")))?;
Ok(exists)
})
}
fn take_raw_atomic(&self, key: Vec<u8>) -> BoxFuture<'_, Result<Option<Vec<u8>>, AppError>> {
Box::pin(async move {
let fk = self.full_key(&key);
let mut conn = self.conn.clone();
let result: Option<Vec<u8>> = redis::cmd("GETDEL")
.arg(fk)
.query_async(&mut conn)
.await
.map_err(|e| AppError::Store(format!("redis GETDEL: {e}")))?;
Ok(result)
})
}
fn incr_raw(&self, key: Vec<u8>) -> BoxFuture<'_, Result<u64, AppError>> {
Box::pin(async move {
let fk = self.full_key(&key);
let mut conn = self.conn.clone();
let next: u64 = conn
.incr(fk, 1i64)
.await
.map_err(|e| AppError::Store(format!("redis INCR: {e}")))?;
Ok(next)
})
}
fn prefix_iter_raw(&self, prefix: Vec<u8>) -> BoxFuture<'_, Result<Vec<RawKvPair>, AppError>> {
Box::pin(async move {
let mut pattern = self.full_key(&prefix);
pattern.extend_from_slice(b"*");
let mut conn = self.conn.clone();
let keys: Vec<Vec<u8>> = {
let mut collected = Vec::new();
let mut iter: redis::AsyncIter<Vec<u8>> = conn
.scan_match(&pattern)
.await
.map_err(|e| AppError::Store(format!("redis SCAN: {e}")))?;
while let Some(key) = iter.next_item().await {
collected
.push(key.map_err(|e| AppError::Store(format!("redis SCAN iter: {e}")))?);
}
collected
};
if keys.is_empty() {
return Ok(Vec::new());
}
let values: Vec<Option<Vec<u8>>> = conn
.mget(&keys)
.await
.map_err(|e| AppError::Store(format!("redis MGET: {e}")))?;
let prefix_len = self.prefix.len();
let results = keys
.into_iter()
.zip(values)
.filter_map(|(k, v)| {
v.map(|val| {
let raw_key = k[prefix_len..].to_vec();
(raw_key, val)
})
})
.collect();
Ok(results)
})
}
}
enum RedisBatchOp {
Insert { full_key: Vec<u8>, value: Vec<u8> },
Remove { full_key: Vec<u8> },
}
struct RedisBatch {
conn: redis::aio::MultiplexedConnection,
ops: Vec<RedisBatchOp>,
}
impl BatchOps for RedisBatch {
fn insert_raw(&mut self, keyspace: &str, key: Vec<u8>, value: Vec<u8>) {
let mut full_key = format!("{keyspace}:").into_bytes();
full_key.extend_from_slice(&key);
self.ops.push(RedisBatchOp::Insert { full_key, value });
}
fn remove(&mut self, keyspace: &str, key: Vec<u8>) {
let mut full_key = format!("{keyspace}:").into_bytes();
full_key.extend_from_slice(&key);
self.ops.push(RedisBatchOp::Remove { full_key });
}
fn commit(self: Box<Self>) -> BoxFuture<'static, Result<(), AppError>> {
Box::pin(async move {
let mut conn = self.conn.clone();
let mut pipe = redis::pipe();
pipe.atomic();
for op in &self.ops {
match op {
RedisBatchOp::Insert { full_key, value } => {
pipe.set(full_key.as_slice(), value.as_slice());
}
RedisBatchOp::Remove { full_key } => {
pipe.del(full_key.as_slice());
}
}
}
pipe.query_async::<()>(&mut conn)
.await
.map_err(|e| AppError::Store(format!("redis pipeline: {e}")))?;
Ok(())
})
}
}