use async_trait::async_trait;
use bson::Document;
use deadpool_redis::{Config as PoolConfig, Pool, Runtime};
use redis::{AsyncCommands, RedisError};
use rigatoni_core::state::{StateStore, StateStoreError};
use std::collections::HashMap;
use std::time::Duration;
use tracing::{debug, error, warn};
const KEY_PREFIX: &str = "rigatoni:resume_token";
const MAX_RETRIES: u32 = 3;
const BASE_RETRY_DELAY_MS: u64 = 100;
#[derive(Clone)]
pub struct RedisConfig {
pub url: String,
pub pool_size: usize,
pub ttl: Option<Duration>,
pub cluster_mode: bool,
pub connection_timeout: Duration,
pub max_retries: u32,
}
impl Default for RedisConfig {
fn default() -> Self {
Self {
url: "redis://localhost:6379".to_string(),
pool_size: 10,
ttl: None,
cluster_mode: false,
connection_timeout: Duration::from_secs(5),
max_retries: MAX_RETRIES,
}
}
}
impl std::fmt::Debug for RedisConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let masked_url = Self::mask_credentials(&self.url);
f.debug_struct("RedisConfig")
.field("url", &masked_url)
.field("pool_size", &self.pool_size)
.field("ttl", &self.ttl)
.field("cluster_mode", &self.cluster_mode)
.field("connection_timeout", &self.connection_timeout)
.field("max_retries", &self.max_retries)
.finish()
}
}
impl RedisConfig {
#[must_use]
pub fn builder() -> RedisConfigBuilder {
RedisConfigBuilder::new()
}
fn mask_credentials(url: &str) -> String {
if let Ok(parsed) = url::Url::parse(url) {
let mut masked = parsed.clone();
if !parsed.username().is_empty() {
let _ = masked.set_username("***");
}
if parsed.password().is_some() {
let _ = masked.set_password(Some("***"));
}
masked.to_string()
} else {
if url.contains("://") {
let parts: Vec<&str> = url.split("://").collect();
if parts.len() == 2 {
format!("{}://***.***", parts[0])
} else {
"***.***".to_string()
}
} else {
"***.***".to_string()
}
}
}
}
#[derive(Debug, Default)]
pub struct RedisConfigBuilder {
url: Option<String>,
pool_size: Option<usize>,
ttl: Option<Duration>,
cluster_mode: Option<bool>,
connection_timeout: Option<Duration>,
max_retries: Option<u32>,
}
impl RedisConfigBuilder {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn url(mut self, url: impl Into<String>) -> Self {
self.url = Some(url.into());
self
}
#[must_use]
pub fn pool_size(mut self, size: usize) -> Self {
self.pool_size = Some(size);
self
}
#[must_use]
pub fn ttl(mut self, ttl: Duration) -> Self {
self.ttl = Some(ttl);
self
}
#[must_use]
pub fn cluster_mode(mut self, enabled: bool) -> Self {
if enabled {
warn!(
"Redis Cluster mode is not implemented yet. \
This flag will be ignored and the connection will use standalone mode. \
Use Redis Sentinel for high availability."
);
}
self.cluster_mode = Some(enabled);
self
}
#[must_use]
pub fn connection_timeout(mut self, timeout: Duration) -> Self {
self.connection_timeout = Some(timeout);
self
}
#[must_use]
pub fn max_retries(mut self, retries: u32) -> Self {
self.max_retries = Some(retries);
self
}
pub fn build(self) -> Result<RedisConfig, StateStoreError> {
let url = self
.url
.ok_or_else(|| StateStoreError::Other("Redis URL is required".to_string()))?;
let pool_size = self.pool_size.unwrap_or(10);
if pool_size == 0 {
return Err(StateStoreError::Other(
"Pool size must be greater than 0".to_string(),
));
}
Ok(RedisConfig {
url,
pool_size,
ttl: self.ttl,
cluster_mode: self.cluster_mode.unwrap_or(false),
connection_timeout: self.connection_timeout.unwrap_or(Duration::from_secs(5)),
max_retries: self.max_retries.unwrap_or(MAX_RETRIES),
})
}
}
#[derive(Clone)]
pub struct RedisStore {
pool: Pool,
config: RedisConfig,
}
impl RedisStore {
pub async fn new(config: RedisConfig) -> Result<Self, StateStoreError> {
debug!("Initializing Redis state store with config: {:?}", config);
let mut pool_config = PoolConfig::from_url(&config.url);
if let Some(pool) = pool_config.pool.as_mut() {
pool.max_size = config.pool_size;
pool.timeouts.wait = Some(config.connection_timeout);
pool.timeouts.create = Some(config.connection_timeout);
pool.timeouts.recycle = Some(config.connection_timeout);
}
let pool = pool_config
.create_pool(Some(Runtime::Tokio1))
.map_err(|e| {
error!("Failed to create Redis connection pool: {}", e);
StateStoreError::Connection(format!("Failed to create pool: {e}"))
})?;
let mut conn = pool.get().await.map_err(|e| {
error!("Failed to get connection from pool: {}", e);
StateStoreError::Connection(format!("Failed to connect to Redis: {e}"))
})?;
redis::cmd("PING")
.query_async::<()>(&mut *conn)
.await
.map_err(|e| {
error!("Redis PING failed: {}", e);
StateStoreError::Connection(format!("Redis connection test failed: {e}"))
})?;
debug!("Redis state store initialized successfully");
Ok(Self { pool, config })
}
fn make_key(collection: &str) -> String {
format!("{KEY_PREFIX}:{collection}")
}
async fn with_retry<F, T, Fut>(&self, operation: F) -> Result<T, StateStoreError>
where
F: Fn() -> Fut,
Fut: std::future::Future<Output = Result<T, RedisError>>,
{
let mut retries = 0;
loop {
match operation().await {
Ok(result) => return Ok(result),
Err(e) if Self::is_retryable(&e) && retries < self.config.max_retries => {
retries += 1;
let delay = Duration::from_millis(BASE_RETRY_DELAY_MS * 2_u64.pow(retries - 1));
warn!(
"Redis operation failed (attempt {}/{}), retrying in {:?}: {}",
retries, self.config.max_retries, delay, e
);
tokio::time::sleep(delay).await;
}
Err(e) => {
error!("Redis operation failed after {} retries: {}", retries, e);
return Err(StateStoreError::Connection(format!(
"Redis operation failed: {e}"
)));
}
}
}
}
fn is_retryable(error: &RedisError) -> bool {
matches!(
error.kind(),
redis::ErrorKind::IoError | redis::ErrorKind::ResponseError
)
}
fn serialize_token(token: &Document) -> Result<Vec<u8>, StateStoreError> {
bson::to_vec(token).map_err(|e| {
StateStoreError::Serialization(format!("Failed to serialize resume token: {e}"))
})
}
fn deserialize_token(bytes: &[u8]) -> Result<Document, StateStoreError> {
bson::from_slice(bytes).map_err(|e| {
StateStoreError::Serialization(format!("Failed to deserialize resume token: {e}"))
})
}
}
const REFRESH_LOCK_SCRIPT: &str = r#"
if redis.call("GET", KEYS[1]) == ARGV[1] then
return redis.call("EXPIRE", KEYS[1], ARGV[2])
else
return 0
end
"#;
const RELEASE_LOCK_SCRIPT: &str = r#"
if redis.call("GET", KEYS[1]) == ARGV[1] then
return redis.call("DEL", KEYS[1])
else
return 0
end
"#;
#[async_trait]
impl StateStore for RedisStore {
async fn save_resume_token(
&self,
collection: &str,
token: &Document,
) -> Result<(), StateStoreError> {
let key = Self::make_key(collection);
let value = Self::serialize_token(token)?;
debug!(
"Saving resume token for collection '{}' to Redis key '{}'",
collection, key
);
let pool = self.pool.clone();
let ttl = self.config.ttl;
self.with_retry::<_, (), _>(|| async {
let mut conn = pool.get().await.map_err(|e| {
RedisError::from((
redis::ErrorKind::IoError,
"Failed to get connection from pool",
e.to_string(),
))
})?;
if let Some(ttl_duration) = ttl {
let ttl_secs = ttl_duration.as_secs();
conn.set_ex(&key, &value, ttl_secs).await
} else {
conn.set(&key, &value).await
}
})
.await?;
debug!(
"Successfully saved resume token for collection '{}'",
collection
);
Ok(())
}
async fn get_resume_token(
&self,
collection: &str,
) -> Result<Option<Document>, StateStoreError> {
let key = Self::make_key(collection);
debug!(
"Retrieving resume token for collection '{}' from Redis key '{}'",
collection, key
);
let pool = self.pool.clone();
let bytes: Option<Vec<u8>> = self
.with_retry(|| async {
let mut conn = pool.get().await.map_err(|e| {
RedisError::from((
redis::ErrorKind::IoError,
"Failed to get connection from pool",
e.to_string(),
))
})?;
conn.get(&key).await
})
.await?;
if let Some(data) = bytes {
let token = Self::deserialize_token(&data)?;
debug!(
"Successfully retrieved resume token for collection '{}'",
collection
);
Ok(Some(token))
} else {
debug!("No resume token found for collection '{}'", collection);
Ok(None)
}
}
async fn delete_resume_token(&self, collection: &str) -> Result<(), StateStoreError> {
let key = Self::make_key(collection);
debug!(
"Deleting resume token for collection '{}' from Redis key '{}'",
collection, key
);
let pool = self.pool.clone();
self.with_retry::<_, (), _>(|| async {
let mut conn = pool.get().await.map_err(|e| {
RedisError::from((
redis::ErrorKind::IoError,
"Failed to get connection from pool",
e.to_string(),
))
})?;
conn.del(&key).await
})
.await?;
debug!(
"Successfully deleted resume token for collection '{}'",
collection
);
Ok(())
}
async fn list_resume_tokens(&self) -> Result<HashMap<String, Document>, StateStoreError> {
let pattern = format!("{KEY_PREFIX}:*");
debug!("Listing all resume tokens with pattern '{}'", pattern);
let pool = self.pool.clone();
let prefix_len = KEY_PREFIX.len() + 1; let mut result = HashMap::new();
let mut cursor: u64 = 0;
loop {
let pool_clone = pool.clone();
let (next_cursor, batch_keys): (u64, Vec<String>) = self
.with_retry(|| async {
let mut conn = pool_clone.get().await.map_err(|e| {
RedisError::from((
redis::ErrorKind::IoError,
"Failed to get connection from pool",
e.to_string(),
))
})?;
redis::cmd("SCAN")
.arg(cursor)
.arg("MATCH")
.arg(&pattern)
.arg("COUNT")
.arg(100) .query_async(&mut *conn)
.await
})
.await?;
if !batch_keys.is_empty() {
let pool_clone = pool.clone();
let values: Vec<Option<Vec<u8>>> = self
.with_retry(|| async {
let mut conn = pool_clone.get().await.map_err(|e| {
RedisError::from((
redis::ErrorKind::IoError,
"Failed to get connection from pool",
e.to_string(),
))
})?;
redis::cmd("MGET")
.arg(&batch_keys)
.query_async(&mut *conn)
.await
})
.await?;
for (key, value) in batch_keys.into_iter().zip(values) {
if let Some(bytes) = value {
let collection = key[prefix_len..].to_string();
let token = Self::deserialize_token(&bytes)?;
result.insert(collection, token);
}
}
}
cursor = next_cursor;
if cursor == 0 {
break;
}
}
debug!("Successfully listed {} resume tokens", result.len());
Ok(result)
}
async fn close(&self) -> Result<(), StateStoreError> {
debug!("Closing Redis state store");
debug!("Redis state store closed");
Ok(())
}
async fn try_acquire_lock(
&self,
key: &str,
owner_id: &str,
ttl: Duration,
) -> Result<bool, StateStoreError> {
debug!(
"Attempting to acquire lock '{}' for owner '{}' with TTL {:?}",
key, owner_id, ttl
);
let pool = self.pool.clone();
let key = key.to_string();
let owner_id = owner_id.to_string();
let ttl_secs = ttl.as_secs();
let result: Option<String> = self
.with_retry(|| async {
let mut conn = pool.get().await.map_err(|e| {
RedisError::from((
redis::ErrorKind::IoError,
"Failed to get connection from pool",
e.to_string(),
))
})?;
redis::cmd("SET")
.arg(&key)
.arg(&owner_id)
.arg("NX")
.arg("EX")
.arg(ttl_secs)
.query_async(&mut *conn)
.await
})
.await?;
let acquired = result.is_some();
if acquired {
debug!("Lock '{}' acquired by owner '{}'", key, owner_id);
} else {
debug!(
"Lock '{}' not acquired (already held by another owner)",
key
);
}
Ok(acquired)
}
async fn refresh_lock(
&self,
key: &str,
owner_id: &str,
ttl: Duration,
) -> Result<bool, StateStoreError> {
debug!(
"Refreshing lock '{}' for owner '{}' with TTL {:?}",
key, owner_id, ttl
);
let pool = self.pool.clone();
let key = key.to_string();
let owner_id = owner_id.to_string();
let ttl_secs = ttl.as_secs();
let result: i32 = self
.with_retry(|| async {
let mut conn = pool.get().await.map_err(|e| {
RedisError::from((
redis::ErrorKind::IoError,
"Failed to get connection from pool",
e.to_string(),
))
})?;
redis::Script::new(REFRESH_LOCK_SCRIPT)
.key(&key)
.arg(&owner_id)
.arg(ttl_secs)
.invoke_async(&mut *conn)
.await
})
.await?;
let refreshed = result == 1;
if refreshed {
debug!("Lock '{}' refreshed for owner '{}'", key, owner_id);
} else {
warn!(
"Lock '{}' NOT refreshed (not owned by '{}' or expired)",
key, owner_id
);
}
Ok(refreshed)
}
async fn release_lock(&self, key: &str, owner_id: &str) -> Result<bool, StateStoreError> {
debug!("Releasing lock '{}' for owner '{}'", key, owner_id);
let pool = self.pool.clone();
let key = key.to_string();
let owner_id = owner_id.to_string();
let result: i32 = self
.with_retry(|| async {
let mut conn = pool.get().await.map_err(|e| {
RedisError::from((
redis::ErrorKind::IoError,
"Failed to get connection from pool",
e.to_string(),
))
})?;
redis::Script::new(RELEASE_LOCK_SCRIPT)
.key(&key)
.arg(&owner_id)
.invoke_async(&mut *conn)
.await
})
.await?;
let released = result == 1;
if released {
debug!("Lock '{}' released by owner '{}'", key, owner_id);
} else {
debug!(
"Lock '{}' NOT released (not owned by '{}' or already released)",
key, owner_id
);
}
Ok(released)
}
async fn is_locked(&self, key: &str) -> Result<bool, StateStoreError> {
debug!("Checking if lock '{}' is held", key);
let pool = self.pool.clone();
let key = key.to_string();
let exists: bool = self
.with_retry(|| async {
let mut conn = pool.get().await.map_err(|e| {
RedisError::from((
redis::ErrorKind::IoError,
"Failed to get connection from pool",
e.to_string(),
))
})?;
conn.exists(&key).await
})
.await?;
debug!("Lock '{}' is_locked: {}", key, exists);
Ok(exists)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_redis_config_builder_defaults() {
let config = RedisConfig::builder()
.url("redis://localhost:6379")
.build()
.expect("Failed to build config");
assert_eq!(config.url, "redis://localhost:6379");
assert_eq!(config.pool_size, 10);
assert!(config.ttl.is_none());
assert!(!config.cluster_mode);
assert_eq!(config.connection_timeout, Duration::from_secs(5));
assert_eq!(config.max_retries, MAX_RETRIES);
}
#[test]
fn test_redis_config_builder_custom_values() {
let config = RedisConfig::builder()
.url("redis://custom:6380")
.pool_size(20)
.ttl(Duration::from_secs(3600))
.connection_timeout(Duration::from_secs(10))
.max_retries(5)
.build()
.expect("Failed to build config");
assert_eq!(config.url, "redis://custom:6380");
assert_eq!(config.pool_size, 20);
assert_eq!(config.ttl, Some(Duration::from_secs(3600)));
assert_eq!(config.connection_timeout, Duration::from_secs(10));
assert_eq!(config.max_retries, 5);
}
#[test]
fn test_redis_config_builder_missing_url() {
let result = RedisConfig::builder().pool_size(10).build();
assert!(result.is_err());
}
#[test]
fn test_redis_config_builder_zero_pool_size() {
let result = RedisConfig::builder()
.url("redis://localhost:6379")
.pool_size(0)
.build();
assert!(result.is_err());
}
#[test]
fn test_redis_config_mask_credentials() {
let masked = RedisConfig::mask_credentials("redis://:mypassword@localhost:6379");
assert!(!masked.contains("mypassword"));
assert!(masked.contains("***"));
let masked = RedisConfig::mask_credentials("redis://user:pass@localhost:6379");
assert!(!masked.contains("user"));
assert!(!masked.contains("pass"));
assert!(masked.contains("***"));
let masked = RedisConfig::mask_credentials("redis://localhost:6379");
assert!(!masked.contains("***@"));
}
#[test]
fn test_make_key() {
let key = RedisStore::make_key("users");
assert_eq!(key, "rigatoni:resume_token:users");
let key = RedisStore::make_key("my_database.orders");
assert_eq!(key, "rigatoni:resume_token:my_database.orders");
}
}