mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
refactor: extract runtime state backends
This commit is contained in:
15
crates/aether-runtime-state/src/error.rs
Normal file
15
crates/aether-runtime-state/src/error.rs
Normal file
@@ -0,0 +1,15 @@
|
||||
pub use aether_data_contracts::DataLayerError;
|
||||
|
||||
pub(crate) fn redis_error(error: impl std::fmt::Display) -> DataLayerError {
|
||||
DataLayerError::redis(error)
|
||||
}
|
||||
|
||||
pub(crate) trait RedisResultExt<T> {
|
||||
fn map_redis_err(self) -> Result<T, DataLayerError>;
|
||||
}
|
||||
|
||||
impl<T> RedisResultExt<T> for Result<T, redis::RedisError> {
|
||||
fn map_redis_err(self) -> Result<T, DataLayerError> {
|
||||
self.map_err(redis_error)
|
||||
}
|
||||
}
|
||||
1879
crates/aether-runtime-state/src/lib.rs
Normal file
1879
crates/aether-runtime-state/src/lib.rs
Normal file
File diff suppressed because it is too large
Load Diff
537
crates/aether-runtime-state/src/memory.rs
Normal file
537
crates/aether-runtime-state/src/memory.rs
Normal file
@@ -0,0 +1,537 @@
|
||||
use std::collections::{BTreeMap, BTreeSet, HashMap, VecDeque};
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::{RuntimeQueueEntry, RuntimeQueueReclaimConfig};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct MemoryRuntimeStateConfig {
|
||||
pub max_kv_entries: usize,
|
||||
}
|
||||
|
||||
impl Default for MemoryRuntimeStateConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_kv_entries: 10_000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct MemoryKvEntry {
|
||||
pub(crate) value: String,
|
||||
pub(crate) inserted_at: Instant,
|
||||
pub(crate) expires_at: Option<Instant>,
|
||||
}
|
||||
|
||||
impl MemoryKvEntry {
|
||||
fn is_expired(&self, now: Instant) -> bool {
|
||||
self.expires_at.is_some_and(|expires_at| now >= expires_at)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub(crate) struct MemoryRuntimeBackend {
|
||||
config: MemoryRuntimeStateConfig,
|
||||
kv: Mutex<HashMap<String, MemoryKvEntry>>,
|
||||
counters: Mutex<HashMap<String, MemoryCounterEntry>>,
|
||||
sets: Mutex<HashMap<String, BTreeSet<String>>>,
|
||||
scores: Mutex<HashMap<String, BTreeMap<String, f64>>>,
|
||||
queues: Mutex<HashMap<String, VecDeque<RuntimeQueueEntry>>>,
|
||||
queue_seq: AtomicU64,
|
||||
locks: Mutex<HashMap<String, MemoryLockEntry>>,
|
||||
semaphores: Mutex<HashMap<String, BTreeMap<String, u64>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct MemoryCounterEntry {
|
||||
value: u32,
|
||||
bucket: u64,
|
||||
expires_at: Instant,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct MemoryLockEntry {
|
||||
pub(crate) token: String,
|
||||
#[allow(dead_code)]
|
||||
pub(crate) owner: String,
|
||||
pub(crate) expires_at: Instant,
|
||||
}
|
||||
|
||||
impl MemoryRuntimeBackend {
|
||||
pub(crate) fn new(config: MemoryRuntimeStateConfig) -> Self {
|
||||
Self {
|
||||
config,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn kv_set(&self, key: &str, value: String, ttl: Option<Duration>) {
|
||||
let mut kv = self.kv.lock().await;
|
||||
let now = Instant::now();
|
||||
if ttl.is_some_and(|ttl| ttl.is_zero()) {
|
||||
kv.remove(key);
|
||||
return;
|
||||
}
|
||||
prune_kv(&mut kv, now);
|
||||
while kv.len() >= self.config.max_kv_entries.max(1) {
|
||||
let Some(oldest_key) = kv
|
||||
.iter()
|
||||
.min_by_key(|(_, entry)| entry.inserted_at)
|
||||
.map(|(key, _)| key.clone())
|
||||
else {
|
||||
break;
|
||||
};
|
||||
kv.remove(&oldest_key);
|
||||
}
|
||||
kv.insert(
|
||||
key.to_string(),
|
||||
MemoryKvEntry {
|
||||
value,
|
||||
inserted_at: now,
|
||||
expires_at: ttl.map(|ttl| now + ttl),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn kv_set_nowait(&self, key: &str, value: String, ttl: Option<Duration>) -> bool {
|
||||
let Ok(mut kv) = self.kv.try_lock() else {
|
||||
return false;
|
||||
};
|
||||
let now = Instant::now();
|
||||
if ttl.is_some_and(|ttl| ttl.is_zero()) {
|
||||
kv.remove(key);
|
||||
return true;
|
||||
}
|
||||
prune_kv(&mut kv, now);
|
||||
while kv.len() >= self.config.max_kv_entries.max(1) {
|
||||
let Some(oldest_key) = kv
|
||||
.iter()
|
||||
.min_by_key(|(_, entry)| entry.inserted_at)
|
||||
.map(|(key, _)| key.clone())
|
||||
else {
|
||||
break;
|
||||
};
|
||||
kv.remove(&oldest_key);
|
||||
}
|
||||
kv.insert(
|
||||
key.to_string(),
|
||||
MemoryKvEntry {
|
||||
value,
|
||||
inserted_at: now,
|
||||
expires_at: ttl.map(|ttl| now + ttl),
|
||||
},
|
||||
);
|
||||
true
|
||||
}
|
||||
|
||||
pub(crate) async fn kv_get(&self, key: &str) -> Option<String> {
|
||||
let mut kv = self.kv.lock().await;
|
||||
get_fresh_locked(&mut kv, key, Instant::now())
|
||||
}
|
||||
|
||||
pub(crate) async fn kv_take(&self, key: &str) -> Option<String> {
|
||||
let mut kv = self.kv.lock().await;
|
||||
let now = Instant::now();
|
||||
let entry = kv.remove(key)?;
|
||||
if entry.is_expired(now) {
|
||||
return None;
|
||||
}
|
||||
Some(entry.value)
|
||||
}
|
||||
|
||||
pub(crate) async fn kv_delete(&self, key: &str) -> bool {
|
||||
self.kv.lock().await.remove(key).is_some()
|
||||
}
|
||||
|
||||
pub(crate) async fn kv_delete_many(&self, keys: &[String]) -> usize {
|
||||
let mut kv = self.kv.lock().await;
|
||||
keys.iter().filter(|key| kv.remove(*key).is_some()).count()
|
||||
}
|
||||
|
||||
pub(crate) async fn kv_exists(&self, key: &str) -> bool {
|
||||
self.kv_get(key).await.is_some()
|
||||
}
|
||||
|
||||
pub(crate) async fn kv_ttl_seconds(&self, key: &str) -> Option<i64> {
|
||||
let mut kv = self.kv.lock().await;
|
||||
let now = Instant::now();
|
||||
let entry = kv.get(key).cloned()?;
|
||||
if entry.is_expired(now) {
|
||||
kv.remove(key);
|
||||
return None;
|
||||
}
|
||||
Some(
|
||||
entry
|
||||
.expires_at
|
||||
.map(|expires_at| {
|
||||
expires_at
|
||||
.saturating_duration_since(now)
|
||||
.as_secs()
|
||||
.try_into()
|
||||
.unwrap_or(i64::MAX)
|
||||
})
|
||||
.unwrap_or(-1),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) async fn kv_scan(&self, pattern: &str) -> Vec<String> {
|
||||
let mut kv = self.kv.lock().await;
|
||||
prune_kv(&mut kv, Instant::now());
|
||||
let mut keys = kv
|
||||
.keys()
|
||||
.filter(|key| key_matches_pattern(key, pattern))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
keys.sort();
|
||||
keys
|
||||
}
|
||||
|
||||
pub(crate) async fn check_and_consume_rate_limit(
|
||||
&self,
|
||||
user_key: &str,
|
||||
key_key: &str,
|
||||
bucket: u64,
|
||||
user_limit: u32,
|
||||
key_limit: u32,
|
||||
ttl: Duration,
|
||||
) -> Result<crate::RateLimitCheck, crate::DataLayerError> {
|
||||
let mut counters = self.counters.lock().await;
|
||||
let now = Instant::now();
|
||||
counters.retain(|_, entry| entry.expires_at > now && entry.bucket >= bucket);
|
||||
|
||||
if user_limit > 0 {
|
||||
let user_count = counters
|
||||
.get(user_key)
|
||||
.filter(|entry| entry.bucket == bucket)
|
||||
.map(|entry| entry.value)
|
||||
.unwrap_or_default();
|
||||
if user_count >= user_limit {
|
||||
return Ok(crate::RateLimitCheck::Rejected {
|
||||
scope: crate::RateLimitScope::User,
|
||||
limit: user_limit,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if key_limit > 0 {
|
||||
let key_count = counters
|
||||
.get(key_key)
|
||||
.filter(|entry| entry.bucket == bucket)
|
||||
.map(|entry| entry.value)
|
||||
.unwrap_or_default();
|
||||
if key_count >= key_limit {
|
||||
return Ok(crate::RateLimitCheck::Rejected {
|
||||
scope: crate::RateLimitScope::Key,
|
||||
limit: key_limit,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let mut remaining = None::<u32>;
|
||||
let expires_at = now + ttl;
|
||||
if user_limit > 0 {
|
||||
let next = counters
|
||||
.entry(user_key.to_string())
|
||||
.and_modify(|entry| {
|
||||
entry.bucket = bucket;
|
||||
entry.value = entry.value.saturating_add(1);
|
||||
entry.expires_at = expires_at;
|
||||
})
|
||||
.or_insert(MemoryCounterEntry {
|
||||
value: 1,
|
||||
bucket,
|
||||
expires_at,
|
||||
})
|
||||
.value;
|
||||
remaining = Some(user_limit.saturating_sub(next));
|
||||
}
|
||||
if key_limit > 0 {
|
||||
let next = counters
|
||||
.entry(key_key.to_string())
|
||||
.and_modify(|entry| {
|
||||
entry.bucket = bucket;
|
||||
entry.value = entry.value.saturating_add(1);
|
||||
entry.expires_at = expires_at;
|
||||
})
|
||||
.or_insert(MemoryCounterEntry {
|
||||
value: 1,
|
||||
bucket,
|
||||
expires_at,
|
||||
})
|
||||
.value;
|
||||
let key_remaining = key_limit.saturating_sub(next);
|
||||
remaining = Some(remaining.map_or(key_remaining, |value| value.min(key_remaining)));
|
||||
}
|
||||
Ok(crate::RateLimitCheck::Allowed {
|
||||
remaining: remaining.unwrap_or(0),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn set_add(&self, key: &str, member: &str) -> bool {
|
||||
self.sets
|
||||
.lock()
|
||||
.await
|
||||
.entry(key.to_string())
|
||||
.or_default()
|
||||
.insert(member.to_string())
|
||||
}
|
||||
|
||||
pub(crate) fn set_add_nowait(&self, key: &str, member: &str) -> bool {
|
||||
let Ok(mut sets) = self.sets.try_lock() else {
|
||||
return false;
|
||||
};
|
||||
sets.entry(key.to_string())
|
||||
.or_default()
|
||||
.insert(member.to_string())
|
||||
}
|
||||
|
||||
pub(crate) async fn set_remove(&self, key: &str, member: &str) -> bool {
|
||||
self.sets
|
||||
.lock()
|
||||
.await
|
||||
.get_mut(key)
|
||||
.is_some_and(|set| set.remove(member))
|
||||
}
|
||||
|
||||
pub(crate) async fn set_members(&self, key: &str) -> Vec<String> {
|
||||
self.sets
|
||||
.lock()
|
||||
.await
|
||||
.get(key)
|
||||
.map(|set| set.iter().cloned().collect())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub(crate) async fn set_len(&self, key: &str) -> usize {
|
||||
self.sets.lock().await.get(key).map_or(0, BTreeSet::len)
|
||||
}
|
||||
|
||||
pub(crate) async fn score_set(&self, key: &str, member: &str, score: f64) {
|
||||
self.scores
|
||||
.lock()
|
||||
.await
|
||||
.entry(key.to_string())
|
||||
.or_default()
|
||||
.insert(member.to_string(), score);
|
||||
}
|
||||
|
||||
pub(crate) async fn score_many(&self, key: &str, members: &[String]) -> Vec<Option<f64>> {
|
||||
let scores = self.scores.lock().await;
|
||||
members
|
||||
.iter()
|
||||
.map(|member| scores.get(key).and_then(|set| set.get(member)).copied())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) async fn score_range_by_min(&self, key: &str, min_score: f64) -> Vec<String> {
|
||||
let scores = self.scores.lock().await;
|
||||
scores
|
||||
.get(key)
|
||||
.map(|set| {
|
||||
set.iter()
|
||||
.filter(|(_, score)| **score >= min_score)
|
||||
.map(|(member, _)| member.clone())
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub(crate) async fn score_remove_by_score(&self, key: &str, max_score: f64) -> usize {
|
||||
let mut scores = self.scores.lock().await;
|
||||
let Some(set) = scores.get_mut(key) else {
|
||||
return 0;
|
||||
};
|
||||
let before = set.len();
|
||||
set.retain(|_, score| *score > max_score);
|
||||
before.saturating_sub(set.len())
|
||||
}
|
||||
|
||||
pub(crate) async fn score_len(&self, key: &str) -> usize {
|
||||
self.scores.lock().await.get(key).map_or(0, BTreeMap::len)
|
||||
}
|
||||
|
||||
pub(crate) async fn queue_append(
|
||||
&self,
|
||||
stream: &str,
|
||||
fields: BTreeMap<String, String>,
|
||||
maxlen: Option<usize>,
|
||||
) -> String {
|
||||
let id = format!(
|
||||
"{}-0",
|
||||
self.queue_seq
|
||||
.fetch_add(1, Ordering::Relaxed)
|
||||
.saturating_add(1)
|
||||
);
|
||||
let mut queues = self.queues.lock().await;
|
||||
let queue = queues.entry(stream.to_string()).or_default();
|
||||
queue.push_back(RuntimeQueueEntry {
|
||||
id: id.clone(),
|
||||
fields,
|
||||
});
|
||||
if let Some(maxlen) = maxlen.filter(|value| *value > 0) {
|
||||
while queue.len() > maxlen {
|
||||
queue.pop_front();
|
||||
}
|
||||
}
|
||||
id
|
||||
}
|
||||
|
||||
pub(crate) async fn queue_read(&self, stream: &str, count: usize) -> Vec<RuntimeQueueEntry> {
|
||||
let mut queues = self.queues.lock().await;
|
||||
let Some(queue) = queues.get_mut(stream) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let mut entries = Vec::new();
|
||||
for _ in 0..count.max(1) {
|
||||
let Some(entry) = queue.pop_front() else {
|
||||
break;
|
||||
};
|
||||
entries.push(entry);
|
||||
}
|
||||
entries
|
||||
}
|
||||
|
||||
pub(crate) async fn queue_claim_stale(
|
||||
&self,
|
||||
_stream: &str,
|
||||
_config: RuntimeQueueReclaimConfig,
|
||||
) -> Vec<RuntimeQueueEntry> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
pub(crate) async fn queue_delete(&self, _stream: &str, _ids: &[String]) -> usize {
|
||||
0
|
||||
}
|
||||
|
||||
pub(crate) async fn lock_try_acquire(
|
||||
&self,
|
||||
key: &str,
|
||||
owner: &str,
|
||||
token: String,
|
||||
ttl: Duration,
|
||||
) -> bool {
|
||||
let mut locks = self.locks.lock().await;
|
||||
let now = Instant::now();
|
||||
locks.retain(|_, entry| entry.expires_at > now);
|
||||
if locks.contains_key(key) {
|
||||
return false;
|
||||
}
|
||||
locks.insert(
|
||||
key.to_string(),
|
||||
MemoryLockEntry {
|
||||
token,
|
||||
owner: owner.to_string(),
|
||||
expires_at: now + ttl,
|
||||
},
|
||||
);
|
||||
true
|
||||
}
|
||||
|
||||
pub(crate) async fn lock_release(&self, key: &str, token: &str) -> bool {
|
||||
let mut locks = self.locks.lock().await;
|
||||
if locks.get(key).is_some_and(|entry| entry.token == token) {
|
||||
locks.remove(key);
|
||||
return true;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
pub(crate) async fn lock_renew(&self, key: &str, token: &str, ttl: Duration) -> bool {
|
||||
let mut locks = self.locks.lock().await;
|
||||
if let Some(entry) = locks.get_mut(key) {
|
||||
if entry.token == token {
|
||||
entry.expires_at = Instant::now() + ttl;
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
pub(crate) async fn semaphore_try_acquire(
|
||||
&self,
|
||||
key: &str,
|
||||
token: String,
|
||||
limit: usize,
|
||||
ttl_ms: u64,
|
||||
) -> Result<usize, usize> {
|
||||
let now_ms = unix_time_ms();
|
||||
let expires_at = now_ms.saturating_add(ttl_ms);
|
||||
let mut semaphores = self.semaphores.lock().await;
|
||||
let holders = semaphores.entry(key.to_string()).or_default();
|
||||
holders.retain(|_, expires| *expires > now_ms);
|
||||
let count = holders.len();
|
||||
if count >= limit {
|
||||
return Err(count);
|
||||
}
|
||||
holders.insert(token, expires_at);
|
||||
Ok(holders.len())
|
||||
}
|
||||
|
||||
pub(crate) async fn semaphore_renew(&self, key: &str, token: &str, ttl_ms: u64) -> bool {
|
||||
let now_ms = unix_time_ms();
|
||||
let mut semaphores = self.semaphores.lock().await;
|
||||
let Some(holders) = semaphores.get_mut(key) else {
|
||||
return false;
|
||||
};
|
||||
holders.retain(|_, expires| *expires > now_ms);
|
||||
if let Some(expires) = holders.get_mut(token) {
|
||||
*expires = now_ms.saturating_add(ttl_ms);
|
||||
return true;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
pub(crate) async fn semaphore_release(&self, key: &str, token: &str) {
|
||||
let mut semaphores = self.semaphores.lock().await;
|
||||
if let Some(holders) = semaphores.get_mut(key) {
|
||||
holders.remove(token);
|
||||
if holders.is_empty() {
|
||||
semaphores.remove(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn semaphore_live_count(&self, key: &str) -> usize {
|
||||
let now_ms = unix_time_ms();
|
||||
let mut semaphores = self.semaphores.lock().await;
|
||||
let Some(holders) = semaphores.get_mut(key) else {
|
||||
return 0;
|
||||
};
|
||||
holders.retain(|_, expires| *expires > now_ms);
|
||||
holders.len()
|
||||
}
|
||||
}
|
||||
|
||||
fn get_fresh_locked(
|
||||
kv: &mut HashMap<String, MemoryKvEntry>,
|
||||
key: &str,
|
||||
now: Instant,
|
||||
) -> Option<String> {
|
||||
let entry = kv.get(key).cloned()?;
|
||||
if entry.is_expired(now) {
|
||||
kv.remove(key);
|
||||
return None;
|
||||
}
|
||||
Some(entry.value)
|
||||
}
|
||||
|
||||
fn prune_kv(kv: &mut HashMap<String, MemoryKvEntry>, now: Instant) {
|
||||
kv.retain(|_, entry| !entry.is_expired(now));
|
||||
}
|
||||
|
||||
pub(crate) fn key_matches_pattern(key: &str, pattern: &str) -> bool {
|
||||
match pattern.strip_suffix('*') {
|
||||
Some(prefix) => key.starts_with(prefix),
|
||||
None => key == pattern,
|
||||
}
|
||||
}
|
||||
|
||||
fn unix_time_ms() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_millis() as u64
|
||||
}
|
||||
69
crates/aether-runtime-state/src/redis/client.rs
Normal file
69
crates/aether-runtime-state/src/redis/client.rs
Normal file
@@ -0,0 +1,69 @@
|
||||
use crate::error::RedisResultExt;
|
||||
use crate::redis::RedisKeyspace;
|
||||
use crate::DataLayerError;
|
||||
|
||||
pub type RedisClient = redis::Client;
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
|
||||
pub struct RedisClientConfig {
|
||||
pub url: String,
|
||||
pub key_prefix: Option<String>,
|
||||
}
|
||||
|
||||
impl RedisClientConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
let raw = self.url.trim();
|
||||
if raw.is_empty() {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis url cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
url::Url::parse(raw).map_err(|err| {
|
||||
DataLayerError::InvalidConfiguration(format!("invalid redis url: {err}"))
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn keyspace(&self) -> RedisKeyspace {
|
||||
RedisKeyspace::new(self.key_prefix.as_deref())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RedisClientFactory {
|
||||
config: RedisClientConfig,
|
||||
}
|
||||
|
||||
impl RedisClientFactory {
|
||||
pub fn new(config: RedisClientConfig) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
Ok(Self { config })
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &RedisClientConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
pub fn connect_lazy(&self) -> Result<RedisClient, DataLayerError> {
|
||||
RedisClient::open(self.config.url.clone()).map_redis_err()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{RedisClientConfig, RedisClientFactory};
|
||||
|
||||
#[test]
|
||||
fn factory_builds_lazy_client_from_valid_config() {
|
||||
let config = RedisClientConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
};
|
||||
let factory = RedisClientFactory::new(config.clone()).expect("factory should build");
|
||||
|
||||
assert_eq!(factory.config(), &config);
|
||||
let _client = factory
|
||||
.connect_lazy()
|
||||
.expect("lazy redis client should build");
|
||||
}
|
||||
}
|
||||
235
crates/aether-runtime-state/src/redis/kv.rs
Normal file
235
crates/aether-runtime-state/src/redis/kv.rs
Normal file
@@ -0,0 +1,235 @@
|
||||
use std::future::Future;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::error::RedisResultExt;
|
||||
use crate::redis::{RedisClient, RedisKeyspace};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct RedisKvRunnerConfig {
|
||||
pub command_timeout_ms: Option<u64>,
|
||||
pub default_ttl_seconds: u64,
|
||||
}
|
||||
|
||||
impl Default for RedisKvRunnerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
command_timeout_ms: Some(1_000),
|
||||
default_ttl_seconds: 300,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisKvRunnerConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if let Some(timeout) = self.command_timeout_ms {
|
||||
if timeout == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis kv command_timeout_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
if self.default_ttl_seconds == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis kv default_ttl_seconds must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RedisKvRunner {
|
||||
client: RedisClient,
|
||||
keyspace: RedisKeyspace,
|
||||
config: RedisKvRunnerConfig,
|
||||
}
|
||||
|
||||
impl RedisKvRunner {
|
||||
pub fn new(
|
||||
client: RedisClient,
|
||||
keyspace: RedisKeyspace,
|
||||
config: RedisKvRunnerConfig,
|
||||
) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
Ok(Self {
|
||||
client,
|
||||
keyspace,
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn client(&self) -> &RedisClient {
|
||||
&self.client
|
||||
}
|
||||
|
||||
pub fn keyspace(&self) -> &RedisKeyspace {
|
||||
&self.keyspace
|
||||
}
|
||||
|
||||
pub fn config(&self) -> RedisKvRunnerConfig {
|
||||
self.config
|
||||
}
|
||||
|
||||
pub async fn setex(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &str,
|
||||
ttl_seconds: Option<u64>,
|
||||
) -> Result<String, DataLayerError> {
|
||||
let resolved_ttl = ttl_seconds.unwrap_or(self.config.default_ttl_seconds);
|
||||
let namespaced_key = self.keyspace.key(key);
|
||||
self.run_with_timeout("redis kv setex", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
redis::cmd("SETEX")
|
||||
.arg(&namespaced_key)
|
||||
.arg(resolved_ttl)
|
||||
.arg(value)
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.map_redis_err()
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn get(&self, key: &str) -> Result<Option<String>, DataLayerError> {
|
||||
let namespaced_key = self.keyspace.key(key);
|
||||
self.run_with_timeout("redis kv get", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
redis::cmd("GET")
|
||||
.arg(&namespaced_key)
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.map_redis_err()
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn getdel(&self, key: &str) -> Result<Option<String>, DataLayerError> {
|
||||
let namespaced_key = self.keyspace.key(key);
|
||||
self.run_with_timeout("redis kv getdel", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
redis::cmd("GETDEL")
|
||||
.arg(&namespaced_key)
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.map_redis_err()
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn exists(&self, key: &str) -> Result<bool, DataLayerError> {
|
||||
let namespaced_key = self.keyspace.key(key);
|
||||
let exists = self
|
||||
.run_with_timeout("redis kv exists", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
redis::cmd("EXISTS")
|
||||
.arg(&namespaced_key)
|
||||
.query_async::<i64>(&mut connection)
|
||||
.await
|
||||
.map_redis_err()
|
||||
})
|
||||
.await?;
|
||||
Ok(exists > 0)
|
||||
}
|
||||
|
||||
pub async fn del(&self, key: &str) -> Result<i64, DataLayerError> {
|
||||
let namespaced_key = self.keyspace.key(key);
|
||||
self.run_with_timeout("redis kv del", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
redis::cmd("DEL")
|
||||
.arg(&namespaced_key)
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.map_redis_err()
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_with_timeout<T, F>(
|
||||
&self,
|
||||
operation: &'static str,
|
||||
future: F,
|
||||
) -> Result<T, DataLayerError>
|
||||
where
|
||||
F: Future<Output = Result<T, DataLayerError>>,
|
||||
{
|
||||
if let Some(timeout_ms) = self.config.command_timeout_ms {
|
||||
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
|
||||
})?
|
||||
} else {
|
||||
future.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{RedisKvRunner, RedisKvRunnerConfig};
|
||||
use crate::redis::{RedisClientConfig, RedisClientFactory, RedisKeyspace};
|
||||
|
||||
fn build_runner() -> RedisKvRunner {
|
||||
let config = RedisClientConfig {
|
||||
url: "redis://localhost/0".to_string(),
|
||||
key_prefix: Some("aether-test".to_string()),
|
||||
};
|
||||
let factory = RedisClientFactory::new(config).expect("redis factory");
|
||||
let client = factory.connect_lazy().expect("connect");
|
||||
let keyspace = factory.config().keyspace();
|
||||
RedisKvRunner::new(client, keyspace, RedisKvRunnerConfig::default()).expect("runner build")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runner_reuses_client_keyspace_and_config() {
|
||||
let runner = build_runner();
|
||||
assert_eq!(
|
||||
runner.keyspace().key("kv:setex:1"),
|
||||
"aether-test:kv:setex:1"
|
||||
);
|
||||
assert_eq!(runner.config(), RedisKvRunnerConfig::default());
|
||||
let _client = runner.client();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_zero_default_ttl() {
|
||||
let config = RedisKvRunnerConfig {
|
||||
command_timeout_ms: Some(100),
|
||||
default_ttl_seconds: 0,
|
||||
};
|
||||
assert!(RedisKvRunner::new(
|
||||
RedisClientFactory::new(RedisClientConfig {
|
||||
url: "redis://localhost/0".to_string(),
|
||||
key_prefix: Some("aether-test".to_string()),
|
||||
})
|
||||
.expect("redis factory")
|
||||
.connect_lazy()
|
||||
.expect("redis client"),
|
||||
RedisKeyspace::new(Some("aether-test")),
|
||||
config,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
323
crates/aether-runtime-state/src/redis/lock.rs
Normal file
323
crates/aether-runtime-state/src/redis/lock.rs
Normal file
@@ -0,0 +1,323 @@
|
||||
use std::future::Future;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::error::RedisResultExt;
|
||||
use crate::redis::{RedisClient, RedisKeyspace};
|
||||
use crate::DataLayerError;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct RedisLockKey(pub String);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RedisLockLease {
|
||||
pub key: RedisLockKey,
|
||||
pub owner: String,
|
||||
pub token: String,
|
||||
pub ttl_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct RedisLockRunnerConfig {
|
||||
pub command_timeout_ms: Option<u64>,
|
||||
pub default_ttl_ms: u64,
|
||||
}
|
||||
|
||||
impl Default for RedisLockRunnerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
command_timeout_ms: Some(1_000),
|
||||
default_ttl_ms: 15_000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisLockRunnerConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if matches!(self.command_timeout_ms, Some(0)) {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis lock command_timeout_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.default_ttl_ms == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis lock default_ttl_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RedisLockRunner {
|
||||
client: RedisClient,
|
||||
keyspace: RedisKeyspace,
|
||||
config: RedisLockRunnerConfig,
|
||||
}
|
||||
|
||||
impl RedisLockRunner {
|
||||
pub fn new(
|
||||
client: RedisClient,
|
||||
keyspace: RedisKeyspace,
|
||||
config: RedisLockRunnerConfig,
|
||||
) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
Ok(Self {
|
||||
client,
|
||||
keyspace,
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn client(&self) -> &RedisClient {
|
||||
&self.client
|
||||
}
|
||||
|
||||
pub fn keyspace(&self) -> &RedisKeyspace {
|
||||
&self.keyspace
|
||||
}
|
||||
|
||||
pub fn config(&self) -> RedisLockRunnerConfig {
|
||||
self.config
|
||||
}
|
||||
|
||||
pub async fn try_acquire(
|
||||
&self,
|
||||
key: &RedisLockKey,
|
||||
owner: &str,
|
||||
ttl_ms: Option<u64>,
|
||||
) -> Result<Option<RedisLockLease>, DataLayerError> {
|
||||
validate_owner(owner)?;
|
||||
validate_key(key)?;
|
||||
let ttl_ms = self.resolve_ttl_ms(ttl_ms)?;
|
||||
let token = format!("{owner}:{}", Uuid::new_v4());
|
||||
|
||||
self.run_with_timeout("redis lock acquire", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
let status = redis::cmd("SET")
|
||||
.arg(&key.0)
|
||||
.arg(&token)
|
||||
.arg("NX")
|
||||
.arg("PX")
|
||||
.arg(ttl_ms)
|
||||
.query_async::<Option<String>>(&mut connection)
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
|
||||
Ok(status.map(|_| RedisLockLease {
|
||||
key: key.clone(),
|
||||
owner: owner.to_string(),
|
||||
token,
|
||||
ttl_ms,
|
||||
}))
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn release(&self, lease: &RedisLockLease) -> Result<bool, DataLayerError> {
|
||||
validate_lease(lease)?;
|
||||
self.run_with_timeout("redis lock release", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
let deleted = redis::Script::new(
|
||||
"if redis.call('get', KEYS[1]) == ARGV[1] then \
|
||||
return redis.call('del', KEYS[1]) \
|
||||
else \
|
||||
return 0 \
|
||||
end",
|
||||
)
|
||||
.key(&lease.key.0)
|
||||
.arg(&lease.token)
|
||||
.invoke_async::<i32>(&mut connection)
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
Ok(deleted > 0)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn renew(
|
||||
&self,
|
||||
lease: &RedisLockLease,
|
||||
ttl_ms: Option<u64>,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
validate_lease(lease)?;
|
||||
let ttl_ms = self.resolve_ttl_ms(ttl_ms)?;
|
||||
|
||||
self.run_with_timeout("redis lock renew", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
let renewed = redis::Script::new(
|
||||
"if redis.call('get', KEYS[1]) == ARGV[1] then \
|
||||
return redis.call('pexpire', KEYS[1], ARGV[2]) \
|
||||
else \
|
||||
return 0 \
|
||||
end",
|
||||
)
|
||||
.key(&lease.key.0)
|
||||
.arg(&lease.token)
|
||||
.arg(ttl_ms)
|
||||
.invoke_async::<i32>(&mut connection)
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
Ok(renewed > 0)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_with_timeout<T, F>(
|
||||
&self,
|
||||
operation: &'static str,
|
||||
future: F,
|
||||
) -> Result<T, DataLayerError>
|
||||
where
|
||||
F: Future<Output = Result<T, DataLayerError>>,
|
||||
{
|
||||
if let Some(timeout_ms) = self.config.command_timeout_ms {
|
||||
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
|
||||
})?
|
||||
} else {
|
||||
future.await
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_ttl_ms(&self, ttl_ms: Option<u64>) -> Result<u64, DataLayerError> {
|
||||
let ttl_ms = ttl_ms.unwrap_or(self.config.default_ttl_ms);
|
||||
if ttl_ms == 0 {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis lock ttl_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(ttl_ms)
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_owner(owner: &str) -> Result<(), DataLayerError> {
|
||||
if owner.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis lock owner cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_key(key: &RedisLockKey) -> Result<(), DataLayerError> {
|
||||
if key.0.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis lock key cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_lease(lease: &RedisLockLease) -> Result<(), DataLayerError> {
|
||||
validate_key(&lease.key)?;
|
||||
validate_owner(&lease.owner)?;
|
||||
if lease.token.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis lock token cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if lease.ttl_ms == 0 {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis lock ttl_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
|
||||
use crate::redis::{RedisClientConfig, RedisClientFactory};
|
||||
|
||||
fn sample_runner() -> RedisLockRunner {
|
||||
let client = RedisClientFactory::new(RedisClientConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
})
|
||||
.expect("factory should build")
|
||||
.connect_lazy()
|
||||
.expect("client should build");
|
||||
|
||||
RedisLockRunner::new(
|
||||
client,
|
||||
RedisClientConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
}
|
||||
.keyspace(),
|
||||
RedisLockRunnerConfig::default(),
|
||||
)
|
||||
.expect("runner should build")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_runner_config() {
|
||||
assert!(RedisLockRunnerConfig {
|
||||
command_timeout_ms: Some(0),
|
||||
..RedisLockRunnerConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
assert!(RedisLockRunnerConfig {
|
||||
default_ttl_ms: 0,
|
||||
..RedisLockRunnerConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runner_reuses_client_and_keyspace() {
|
||||
let runner = sample_runner();
|
||||
|
||||
assert_eq!(runner.config(), RedisLockRunnerConfig::default());
|
||||
assert_eq!(runner.keyspace().lock_key("poller").0, "aether:lock:poller");
|
||||
let _client_ref = runner.client();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_invalid_owner_or_lease_before_network() {
|
||||
let runner = sample_runner();
|
||||
|
||||
assert!(runner
|
||||
.try_acquire(&RedisLockKey("aether:lock:poller".to_string()), "", None)
|
||||
.await
|
||||
.is_err());
|
||||
assert!(runner
|
||||
.release(&RedisLockLease {
|
||||
key: RedisLockKey("aether:lock:poller".to_string()),
|
||||
owner: "worker-1".to_string(),
|
||||
token: String::new(),
|
||||
ttl_ms: 1_000,
|
||||
})
|
||||
.await
|
||||
.is_err());
|
||||
assert!(runner
|
||||
.renew(
|
||||
&RedisLockLease {
|
||||
key: RedisLockKey("aether:lock:poller".to_string()),
|
||||
owner: "worker-1".to_string(),
|
||||
token: "token-1".to_string(),
|
||||
ttl_ms: 1_000,
|
||||
},
|
||||
Some(0),
|
||||
)
|
||||
.await
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
25
crates/aether-runtime-state/src/redis/mod.rs
Normal file
25
crates/aether-runtime-state/src/redis/mod.rs
Normal file
@@ -0,0 +1,25 @@
|
||||
mod client;
|
||||
mod kv;
|
||||
mod lock;
|
||||
mod namespace;
|
||||
mod stream;
|
||||
|
||||
pub use client::{RedisClient, RedisClientConfig, RedisClientFactory};
|
||||
pub use kv::{RedisKvRunner, RedisKvRunnerConfig};
|
||||
pub use lock::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
|
||||
pub use namespace::RedisKeyspace;
|
||||
pub use stream::{
|
||||
RedisConsumerGroup, RedisConsumerName, RedisStreamEntry, RedisStreamName,
|
||||
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner, RedisStreamRunnerConfig,
|
||||
};
|
||||
|
||||
pub(crate) type RedisCmd = redis::Cmd;
|
||||
pub(crate) type RedisScript = redis::Script;
|
||||
|
||||
pub(crate) fn cmd(name: &str) -> RedisCmd {
|
||||
redis::cmd(name)
|
||||
}
|
||||
|
||||
pub(crate) fn script(source: &str) -> RedisScript {
|
||||
redis::Script::new(source)
|
||||
}
|
||||
43
crates/aether-runtime-state/src/redis/namespace.rs
Normal file
43
crates/aether-runtime-state/src/redis/namespace.rs
Normal file
@@ -0,0 +1,43 @@
|
||||
use aether_cache::CacheKeyNamespace;
|
||||
|
||||
use crate::redis::{RedisLockKey, RedisStreamName};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RedisKeyspace {
|
||||
namespace: CacheKeyNamespace,
|
||||
}
|
||||
|
||||
impl RedisKeyspace {
|
||||
pub fn new(prefix: Option<&str>) -> Self {
|
||||
let normalized = prefix.unwrap_or_default().trim().trim_matches(':');
|
||||
Self {
|
||||
namespace: CacheKeyNamespace::new(normalized),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn key(&self, raw_key: &str) -> String {
|
||||
self.namespace.key(raw_key)
|
||||
}
|
||||
|
||||
pub fn lock_key(&self, raw_key: &str) -> RedisLockKey {
|
||||
RedisLockKey(self.namespace.child("lock").key(raw_key))
|
||||
}
|
||||
|
||||
pub fn stream_name(&self, raw_name: &str) -> RedisStreamName {
|
||||
RedisStreamName(self.namespace.child("stream").key(raw_name))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::RedisKeyspace;
|
||||
|
||||
#[test]
|
||||
fn composes_prefixed_lock_and_stream_names() {
|
||||
let keyspace = RedisKeyspace::new(Some("aether"));
|
||||
|
||||
assert_eq!(keyspace.key("auth:user"), "aether:auth:user");
|
||||
assert_eq!(keyspace.lock_key("poller").0, "aether:lock:poller");
|
||||
assert_eq!(keyspace.stream_name("audit").0, "aether:stream:audit");
|
||||
}
|
||||
}
|
||||
747
crates/aether-runtime-state/src/redis/stream.rs
Normal file
747
crates/aether-runtime-state/src/redis/stream.rs
Normal file
@@ -0,0 +1,747 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::future::Future;
|
||||
use std::time::Duration;
|
||||
|
||||
use redis::from_redis_value;
|
||||
use redis::streams::StreamReadReply;
|
||||
use redis::Value as RedisValue;
|
||||
|
||||
use crate::error::{redis_error, RedisResultExt};
|
||||
use crate::redis::{RedisClient, RedisKeyspace};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct RedisStreamName(pub String);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct RedisConsumerGroup(pub String);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct RedisConsumerName(pub String);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RedisStreamEntry {
|
||||
pub id: String,
|
||||
pub fields: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RedisStreamReclaimResult {
|
||||
pub next_start_id: String,
|
||||
pub entries: Vec<RedisStreamEntry>,
|
||||
pub deleted_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct RedisStreamReclaimConfig {
|
||||
pub min_idle_ms: u64,
|
||||
pub count: usize,
|
||||
}
|
||||
|
||||
impl Default for RedisStreamReclaimConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
min_idle_ms: 60_000,
|
||||
count: 32,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisStreamReclaimConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if self.min_idle_ms == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis stream reclaim min_idle_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.count == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis stream reclaim count must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct RedisStreamRunnerConfig {
|
||||
pub command_timeout_ms: Option<u64>,
|
||||
pub read_block_ms: Option<u64>,
|
||||
pub read_count: usize,
|
||||
}
|
||||
|
||||
impl Default for RedisStreamRunnerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
command_timeout_ms: Some(2_000),
|
||||
read_block_ms: Some(1_000),
|
||||
read_count: 32,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisStreamRunnerConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if matches!(self.command_timeout_ms, Some(0)) {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis stream command_timeout_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if matches!(self.read_block_ms, Some(0)) {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis stream read_block_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if let (Some(command_timeout_ms), Some(read_block_ms)) =
|
||||
(self.command_timeout_ms, self.read_block_ms)
|
||||
{
|
||||
if command_timeout_ms <= read_block_ms {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis stream command_timeout_ms must be greater than read_block_ms"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
if self.read_count == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis stream read_count must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RedisStreamRunner {
|
||||
client: RedisClient,
|
||||
keyspace: RedisKeyspace,
|
||||
config: RedisStreamRunnerConfig,
|
||||
}
|
||||
|
||||
impl RedisStreamRunner {
|
||||
pub fn new(
|
||||
client: RedisClient,
|
||||
keyspace: RedisKeyspace,
|
||||
config: RedisStreamRunnerConfig,
|
||||
) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
Ok(Self {
|
||||
client,
|
||||
keyspace,
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn client(&self) -> &RedisClient {
|
||||
&self.client
|
||||
}
|
||||
|
||||
pub fn keyspace(&self) -> &RedisKeyspace {
|
||||
&self.keyspace
|
||||
}
|
||||
|
||||
pub fn config(&self) -> RedisStreamRunnerConfig {
|
||||
self.config
|
||||
}
|
||||
|
||||
pub async fn ensure_consumer_group(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
group: &RedisConsumerGroup,
|
||||
start_id: &str,
|
||||
) -> Result<(), DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
validate_group(group)?;
|
||||
validate_stream_position(start_id)?;
|
||||
|
||||
self.run_with_timeout("redis stream ensure consumer group", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
let result = redis::cmd("XGROUP")
|
||||
.arg("CREATE")
|
||||
.arg(&stream.0)
|
||||
.arg(&group.0)
|
||||
.arg(start_id)
|
||||
.arg("MKSTREAM")
|
||||
.query_async::<String>(&mut connection)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(_) => Ok(()),
|
||||
Err(err) if err.code() == Some("BUSYGROUP") => Ok(()),
|
||||
Err(err) => Err(redis_error(err)),
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn append_fields(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
fields: &BTreeMap<String, String>,
|
||||
) -> Result<String, DataLayerError> {
|
||||
self.append_fields_with_maxlen(stream, fields, None).await
|
||||
}
|
||||
|
||||
pub async fn append_fields_with_maxlen(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
fields: &BTreeMap<String, String>,
|
||||
maxlen: Option<usize>,
|
||||
) -> Result<String, DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
if fields.is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis stream fields cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
self.run_with_timeout("redis stream append", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
let mut command = redis::cmd("XADD");
|
||||
command.arg(&stream.0);
|
||||
if let Some(maxlen) = maxlen.filter(|value| *value > 0) {
|
||||
command.arg("MAXLEN").arg("~").arg(maxlen);
|
||||
}
|
||||
command.arg("*");
|
||||
for (key, value) in fields {
|
||||
command.arg(key).arg(value);
|
||||
}
|
||||
command
|
||||
.query_async::<String>(&mut connection)
|
||||
.await
|
||||
.map_redis_err()
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn append_json(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
field: &str,
|
||||
payload: &serde_json::Value,
|
||||
) -> Result<String, DataLayerError> {
|
||||
if field.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis stream json field cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let mut fields = BTreeMap::new();
|
||||
fields.insert(
|
||||
field.to_string(),
|
||||
serde_json::to_string(payload).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"failed to serialize redis stream payload: {err}"
|
||||
))
|
||||
})?,
|
||||
);
|
||||
self.append_fields(stream, &fields).await
|
||||
}
|
||||
|
||||
pub async fn read_group(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
group: &RedisConsumerGroup,
|
||||
consumer: &RedisConsumerName,
|
||||
) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
validate_group(group)?;
|
||||
validate_consumer(consumer)?;
|
||||
|
||||
self.run_with_timeout("redis stream read group", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
let mut command = redis::cmd("XREADGROUP");
|
||||
command
|
||||
.arg("GROUP")
|
||||
.arg(&group.0)
|
||||
.arg(&consumer.0)
|
||||
.arg("COUNT")
|
||||
.arg(self.config.read_count);
|
||||
if let Some(block_ms) = self.config.read_block_ms {
|
||||
command.arg("BLOCK").arg(block_ms);
|
||||
}
|
||||
command.arg("STREAMS").arg(&stream.0).arg(">");
|
||||
|
||||
let reply = command
|
||||
.query_async::<StreamReadReply>(&mut connection)
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
|
||||
Ok(reply
|
||||
.keys
|
||||
.into_iter()
|
||||
.flat_map(|key| key.ids.into_iter())
|
||||
.map(|id| RedisStreamEntry {
|
||||
id: id.id,
|
||||
fields: id
|
||||
.map
|
||||
.into_iter()
|
||||
.filter_map(|(field, value)| {
|
||||
redis::from_redis_value::<String>(&value)
|
||||
.ok()
|
||||
.map(|text| (field, text))
|
||||
})
|
||||
.collect(),
|
||||
})
|
||||
.collect())
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn ack(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
group: &RedisConsumerGroup,
|
||||
ids: &[String],
|
||||
) -> Result<usize, DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
validate_group(group)?;
|
||||
if ids.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
self.run_with_timeout("redis stream ack", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
let mut command = redis::cmd("XACK");
|
||||
command.arg(&stream.0).arg(&group.0);
|
||||
for id in ids {
|
||||
command.arg(id);
|
||||
}
|
||||
command
|
||||
.query_async::<usize>(&mut connection)
|
||||
.await
|
||||
.map_redis_err()
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn delete(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
ids: &[String],
|
||||
) -> Result<usize, DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
if ids.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
self.run_with_timeout("redis stream delete", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
let mut command = redis::cmd("XDEL");
|
||||
command.arg(&stream.0);
|
||||
for id in ids {
|
||||
command.arg(id);
|
||||
}
|
||||
command
|
||||
.query_async::<usize>(&mut connection)
|
||||
.await
|
||||
.map_redis_err()
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn claim_stale(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
group: &RedisConsumerGroup,
|
||||
consumer: &RedisConsumerName,
|
||||
start_id: &str,
|
||||
config: RedisStreamReclaimConfig,
|
||||
) -> Result<RedisStreamReclaimResult, DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
validate_group(group)?;
|
||||
validate_consumer(consumer)?;
|
||||
validate_stream_position(start_id)?;
|
||||
config.validate()?;
|
||||
|
||||
self.run_with_timeout("redis stream reclaim", async {
|
||||
let mut connection = self
|
||||
.client
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
let reply = redis::cmd("XAUTOCLAIM")
|
||||
.arg(&stream.0)
|
||||
.arg(&group.0)
|
||||
.arg(&consumer.0)
|
||||
.arg(config.min_idle_ms)
|
||||
.arg(start_id)
|
||||
.arg("COUNT")
|
||||
.arg(config.count)
|
||||
.query_async::<RedisValue>(&mut connection)
|
||||
.await
|
||||
.map_redis_err()?;
|
||||
|
||||
parse_reclaim_result(reply)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_with_timeout<T, F>(
|
||||
&self,
|
||||
operation: &'static str,
|
||||
future: F,
|
||||
) -> Result<T, DataLayerError>
|
||||
where
|
||||
F: Future<Output = Result<T, DataLayerError>>,
|
||||
{
|
||||
if let Some(timeout_ms) = self.config.command_timeout_ms {
|
||||
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
|
||||
})?
|
||||
} else {
|
||||
future.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_stream_name(stream: &RedisStreamName) -> Result<(), DataLayerError> {
|
||||
if stream.0.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis stream name cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_group(group: &RedisConsumerGroup) -> Result<(), DataLayerError> {
|
||||
if group.0.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis consumer group cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_consumer(consumer: &RedisConsumerName) -> Result<(), DataLayerError> {
|
||||
if consumer.0.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis consumer name cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_stream_position(position: &str) -> Result<(), DataLayerError> {
|
||||
if position.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis stream position cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_reclaim_result(value: RedisValue) -> Result<RedisStreamReclaimResult, DataLayerError> {
|
||||
let RedisValue::Array(parts) = value else {
|
||||
return Err(DataLayerError::UnexpectedValue(
|
||||
"redis xautoclaim returned non-array payload".to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
if parts.len() < 2 || parts.len() > 3 {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"redis xautoclaim returned {} top-level fields, expected 2 or 3",
|
||||
parts.len()
|
||||
)));
|
||||
}
|
||||
|
||||
let next_start_id = parse_string_value(&parts[0], "redis xautoclaim next_start_id")?;
|
||||
let entries = parse_reclaim_entries(&parts[1])?;
|
||||
let deleted_ids = match parts.get(2) {
|
||||
Some(value) => parse_string_array(value, "redis xautoclaim deleted_ids")?,
|
||||
None => Vec::new(),
|
||||
};
|
||||
|
||||
Ok(RedisStreamReclaimResult {
|
||||
next_start_id,
|
||||
entries,
|
||||
deleted_ids,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_reclaim_entries(value: &RedisValue) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
|
||||
match value {
|
||||
RedisValue::Array(entries) => entries.iter().map(parse_reclaim_entry).collect(),
|
||||
RedisValue::Nil => Ok(Vec::new()),
|
||||
_ => Err(DataLayerError::UnexpectedValue(
|
||||
"redis xautoclaim entries payload was not an array".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_reclaim_entry(value: &RedisValue) -> Result<RedisStreamEntry, DataLayerError> {
|
||||
let RedisValue::Array(parts) = value else {
|
||||
return Err(DataLayerError::UnexpectedValue(
|
||||
"redis xautoclaim entry was not an array".to_string(),
|
||||
));
|
||||
};
|
||||
if parts.len() != 2 {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"redis xautoclaim entry had {} fields, expected 2",
|
||||
parts.len()
|
||||
)));
|
||||
}
|
||||
|
||||
let id = parse_string_value(&parts[0], "redis xautoclaim entry id")?;
|
||||
let fields = parse_string_map(&parts[1], "redis xautoclaim entry fields")?;
|
||||
Ok(RedisStreamEntry { id, fields })
|
||||
}
|
||||
|
||||
fn parse_string_map(
|
||||
value: &RedisValue,
|
||||
context: &str,
|
||||
) -> Result<BTreeMap<String, String>, DataLayerError> {
|
||||
match value {
|
||||
RedisValue::Array(values) => {
|
||||
if values.len() % 2 != 0 {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"{context} expected an even number of field elements, got {}",
|
||||
values.len()
|
||||
)));
|
||||
}
|
||||
let mut fields = BTreeMap::new();
|
||||
for pair in values.chunks(2) {
|
||||
let key = parse_string_value(&pair[0], context)?;
|
||||
let value = parse_string_value(&pair[1], context)?;
|
||||
fields.insert(key, value);
|
||||
}
|
||||
Ok(fields)
|
||||
}
|
||||
RedisValue::Map(entries) => entries
|
||||
.iter()
|
||||
.map(|(key, value)| {
|
||||
Ok((
|
||||
parse_string_value(key, context)?,
|
||||
parse_string_value(value, context)?,
|
||||
))
|
||||
})
|
||||
.collect(),
|
||||
RedisValue::Nil => Ok(BTreeMap::new()),
|
||||
_ => Err(DataLayerError::UnexpectedValue(format!(
|
||||
"{context} expected a redis array/map payload"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_string_array(value: &RedisValue, context: &str) -> Result<Vec<String>, DataLayerError> {
|
||||
match value {
|
||||
RedisValue::Array(values) => values
|
||||
.iter()
|
||||
.map(|value| parse_string_value(value, context))
|
||||
.collect(),
|
||||
RedisValue::Nil => Ok(Vec::new()),
|
||||
_ => Err(DataLayerError::UnexpectedValue(format!(
|
||||
"{context} expected a redis array payload"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_string_value(value: &RedisValue, context: &str) -> Result<String, DataLayerError> {
|
||||
from_redis_value::<String>(value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"{context} was not a string-compatible redis value: {err}"
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{
|
||||
parse_reclaim_result, RedisConsumerGroup, RedisConsumerName, RedisStreamName,
|
||||
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner,
|
||||
RedisStreamRunnerConfig,
|
||||
};
|
||||
use crate::redis::{RedisClientConfig, RedisClientFactory};
|
||||
use redis::Value as RedisValue;
|
||||
|
||||
fn sample_runner() -> RedisStreamRunner {
|
||||
let config = RedisClientConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
};
|
||||
let client = RedisClientFactory::new(config.clone())
|
||||
.expect("factory should build")
|
||||
.connect_lazy()
|
||||
.expect("client should build");
|
||||
|
||||
RedisStreamRunner::new(
|
||||
client,
|
||||
config.keyspace(),
|
||||
RedisStreamRunnerConfig::default(),
|
||||
)
|
||||
.expect("runner should build")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_stream_runner_config() {
|
||||
assert!(RedisStreamRunnerConfig {
|
||||
command_timeout_ms: Some(0),
|
||||
..RedisStreamRunnerConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
assert!(RedisStreamRunnerConfig {
|
||||
read_block_ms: Some(0),
|
||||
..RedisStreamRunnerConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
assert!(RedisStreamRunnerConfig {
|
||||
command_timeout_ms: Some(1_000),
|
||||
read_block_ms: Some(1_000),
|
||||
..RedisStreamRunnerConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
assert!(RedisStreamRunnerConfig {
|
||||
read_count: 0,
|
||||
..RedisStreamRunnerConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_reclaim_config() {
|
||||
assert!(RedisStreamReclaimConfig {
|
||||
min_idle_ms: 0,
|
||||
..RedisStreamReclaimConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
assert!(RedisStreamReclaimConfig {
|
||||
count: 0,
|
||||
..RedisStreamReclaimConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runner_reuses_client_and_keyspace() {
|
||||
let runner = sample_runner();
|
||||
|
||||
assert_eq!(runner.config(), RedisStreamRunnerConfig::default());
|
||||
assert_eq!(
|
||||
runner.keyspace().stream_name("audit").0,
|
||||
"aether:stream:audit"
|
||||
);
|
||||
let _client_ref = runner.client();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_invalid_inputs_before_network() {
|
||||
let runner = sample_runner();
|
||||
let stream = RedisStreamName("aether:stream:audit".to_string());
|
||||
let group = RedisConsumerGroup("audit-workers".to_string());
|
||||
let consumer = RedisConsumerName("worker-1".to_string());
|
||||
|
||||
assert!(runner
|
||||
.ensure_consumer_group(&stream, &group, "")
|
||||
.await
|
||||
.is_err());
|
||||
assert!(runner
|
||||
.append_fields(&stream, &BTreeMap::new())
|
||||
.await
|
||||
.is_err());
|
||||
assert!(runner
|
||||
.append_json(&stream, "", &serde_json::json!({"ok": true}))
|
||||
.await
|
||||
.is_err());
|
||||
assert!(runner
|
||||
.read_group(&stream, &group, &RedisConsumerName(String::new()))
|
||||
.await
|
||||
.is_err());
|
||||
assert_eq!(
|
||||
runner.ack(&stream, &group, &[]).await.expect("empty ack"),
|
||||
0
|
||||
);
|
||||
assert_eq!(runner.delete(&stream, &[]).await.expect("empty delete"), 0);
|
||||
assert!(runner
|
||||
.claim_stale(
|
||||
&stream,
|
||||
&group,
|
||||
&consumer,
|
||||
"",
|
||||
RedisStreamReclaimConfig::default()
|
||||
)
|
||||
.await
|
||||
.is_err());
|
||||
let _ = consumer;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_reclaim_result_with_deleted_ids() {
|
||||
let parsed = parse_reclaim_result(RedisValue::Array(vec![
|
||||
RedisValue::BulkString(b"0-0".to_vec()),
|
||||
RedisValue::Array(vec![RedisValue::Array(vec![
|
||||
RedisValue::BulkString(b"1710000000000-0".to_vec()),
|
||||
RedisValue::Array(vec![
|
||||
RedisValue::BulkString(b"payload".to_vec()),
|
||||
RedisValue::BulkString(br#"{"ok":true}"#.to_vec()),
|
||||
RedisValue::BulkString(b"kind".to_vec()),
|
||||
RedisValue::BulkString(b"shadow".to_vec()),
|
||||
]),
|
||||
])]),
|
||||
RedisValue::Array(vec![RedisValue::BulkString(b"1709999999999-0".to_vec())]),
|
||||
]))
|
||||
.expect("reclaim result should parse");
|
||||
|
||||
assert_eq!(
|
||||
parsed,
|
||||
RedisStreamReclaimResult {
|
||||
next_start_id: "0-0".to_string(),
|
||||
entries: vec![super::RedisStreamEntry {
|
||||
id: "1710000000000-0".to_string(),
|
||||
fields: BTreeMap::from([
|
||||
("kind".to_string(), "shadow".to_string()),
|
||||
("payload".to_string(), r#"{"ok":true}"#.to_string()),
|
||||
]),
|
||||
}],
|
||||
deleted_ids: vec!["1709999999999-0".to_string()],
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_reclaim_result_without_deleted_ids() {
|
||||
let parsed = parse_reclaim_result(RedisValue::Array(vec![
|
||||
RedisValue::BulkString(b"0-0".to_vec()),
|
||||
RedisValue::Array(vec![]),
|
||||
]))
|
||||
.expect("reclaim result should parse");
|
||||
|
||||
assert_eq!(
|
||||
parsed,
|
||||
RedisStreamReclaimResult {
|
||||
next_start_id: "0-0".to_string(),
|
||||
entries: Vec::new(),
|
||||
deleted_ids: Vec::new(),
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user