i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(runtime-state): govern redis connections

Browse Source
This commit is contained in:
fawney19
2026-05-21 22:51:57 +08:00
parent b8a65cbdec
commit ab0a90de97
24 changed files with 3155 additions and 1064 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
Cargo.lock generated
View File

@@ -858,6 +858,15 @@ dependencies = [
"tracing",
]
[[package]]
name = "backon"
version = "1.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef"
dependencies = [
"fastrand",
]
[[package]]
name = "base64"
version = "0.21.7"
@@ -1637,6 +1646,12 @@ dependencies = [
"regex",
]
[[package]]
name = "fastrand"
version = "2.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6"
[[package]]
name = "filedescriptor"
version = "0.8.3"
@@ -3389,8 +3404,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e37ec3fd44bea2ec947ba6cc7634d7999a6590aca7c35827c250bc0de502bda6"
dependencies = [
"arc-swap",
"backon",
"bytes",
"combine",
"futures-channel",
"futures-util",
"itoa",
"num-bigint",

2
Cargo.toml
View File

@@ -82,7 +82,7 @@ hmac = "0.12"
http = "1"
pbkdf2 = { version = "0.12", default-features = false, features = ["hmac"] }
reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls", "http2", "socks"] }
redis = { version = "0.28", default-features = false, features = ["tokio-comp", "script", "streams"] }
redis = { version = "0.28", default-features = false, features = ["tokio-comp", "script", "streams", "connection-manager"] }
regex = "1"
rustls = { version = "0.23", features = ["ring"] }
semver = "1"

6
apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_model_mapping.rs
View File

@@ -59,6 +59,11 @@ pub(super) async fn build_admin_monitoring_redis_cache_categories_response(
) -> Result<Response<Body>, GatewayError> {
let mut categories = Vec::with_capacity(ADMIN_MONITORING_REDIS_CACHE_CATEGORIES.len());
let mut total_keys = 0usize;
let diagnostics = state
.runtime_state()
.redis_diagnostics()
.await
.map_err(|err| GatewayError::Internal(format!("redis diagnostics failed: {err}")))?;
for (key, name, pattern, description) in ADMIN_MONITORING_REDIS_CACHE_CATEGORIES {
let count = list_admin_monitoring_namespaced_keys(state, pattern)
@@ -81,6 +86,7 @@ pub(super) async fn build_admin_monitoring_redis_cache_categories_response(
"backend": state.runtime_state().backend_kind().as_str(),
"categories": categories,
"total_keys": total_keys,
"diagnostics": diagnostics,
}
}))
.into_response())

1
apps/aether-gateway/src/handlers/admin/observability/monitoring/tests/mod.rs
View File

@@ -1194,6 +1194,7 @@ async fn admin_monitoring_redis_keys_returns_local_payload_without_redis() {
assert_eq!(payload["data"]["available"], json!(true));
assert_eq!(payload["data"]["backend"], json!("memory"));
assert_eq!(payload["data"]["total_keys"], json!(0));
assert_eq!(payload["data"]["diagnostics"], serde_json::Value::Null);
}
#[tokio::test]

1
apps/aether-gateway/src/main.rs
View File

@@ -497,6 +497,7 @@ impl GatewayUsageArgs {
fn to_config(&self) -> UsageRuntimeConfig {
UsageRuntimeConfig {
enabled: true,
queue_terminal_events: true,
stream_key: self.queue_stream_key.trim().to_string(),
consumer_group: self.queue_group.trim().to_string(),
dlq_stream_key: self.queue_dlq_stream_key.trim().to_string(),

74
apps/aether-gateway/src/rate_limit.rs
View File

@@ -150,11 +150,20 @@ impl FrontdoorUserRpmLimiter {
scope_key: &str,
bucket: u64,
) -> Result<u32, GatewayError> {
if !state.runtime_state.is_memory() {
let raw = state.runtime_state.kv_get(scope_key).await.map_err(|err| {
GatewayError::Internal(format!("frontdoor user rpm runtime read failed: {err}"))
})?;
return Ok(raw.and_then(|value| value.parse::<u32>().ok()).unwrap_or(0));
match state.runtime_state.kv_get(scope_key).await {
Ok(raw) => return Ok(raw.and_then(|value| value.parse::<u32>().ok()).unwrap_or(0)),
Err(err) if !self.config.allow_local_fallback() => {
return Err(GatewayError::Internal(format!(
"frontdoor user rpm runtime read failed: {err}"
)));
}
Err(err) => {
warn!(
error = ?err,
scope_key = %scope_key,
"frontdoor user rpm runtime count read failed; using local fallback"
);
}
}
let counts = self.memory_counts.lock().await;
@@ -187,24 +196,22 @@ impl FrontdoorUserRpmLimiter {
return Ok(FrontdoorUserRpmOutcome::Allowed);
}
if !state.runtime_state.is_memory() {
match self.check_and_consume_runtime(state, &plan).await {
Ok(outcome) => return Ok(outcome),
Err(err) => {
warn!(
error = ?err,
user_rpm_key = %plan.user_rpm_key,
key_rpm_key = %plan.key_rpm_key,
"frontdoor user rpm redis check failed"
);
if self.config.fail_open() {
return Ok(FrontdoorUserRpmOutcome::NotApplicable);
}
if !self.config.allow_local_fallback() {
return Err(GatewayError::Internal(
"frontdoor user rpm runtime backend is unavailable and local fallback is disabled for the current deployment mode".to_string(),
));
}
match self.check_and_consume_runtime(state, &plan).await {
Ok(outcome) => return Ok(outcome),
Err(err) => {
warn!(
error = ?err,
user_rpm_key = %plan.user_rpm_key,
key_rpm_key = %plan.key_rpm_key,
"frontdoor user rpm runtime check failed"
);
if self.config.fail_open() {
return Ok(FrontdoorUserRpmOutcome::NotApplicable);
}
if !self.config.allow_local_fallback() {
return Err(GatewayError::Internal(
"frontdoor user rpm runtime backend is unavailable and local fallback is disabled for the current deployment mode".to_string(),
));
}
}
}
@@ -605,7 +612,7 @@ mod tests {
}
#[tokio::test]
async fn limiter_rejects_missing_shared_runtime_when_local_fallback_disabled() {
async fn limiter_uses_runtime_state_when_local_fallback_disabled() {
let limiter = FrontdoorUserRpmLimiter::new(
FrontdoorUserRpmConfig::new(60, 120, false).with_local_fallback(false),
);
@@ -626,15 +633,22 @@ mod tests {
});
let state = AppState::new().expect("state should build for tests");
let err = limiter
let first = limiter
.check_and_consume(&state, Some(&decision))
.await
.expect_err("missing shared runtime should fail in strict mode");
match err {
crate::GatewayError::Internal(message) => {
assert!(message.contains("requires shared runtime state"));
.expect("runtime check should succeed");
assert_eq!(first, FrontdoorUserRpmOutcome::Allowed);
let second = limiter
.check_and_consume(&state, Some(&decision))
.await
.expect("runtime check should succeed");
match second {
FrontdoorUserRpmOutcome::Rejected(rejection) => {
assert_eq!(rejection.scope, "user");
assert_eq!(rejection.limit, 1);
}
other => panic!("expected internal error, got {other:?}"),
other => panic!("expected rejection, got {other:?}"),
}
}
}

15
apps/aether-gateway/src/state/core.rs
View File

@@ -89,24 +89,17 @@ impl AppState {
fn usage_worker_queue_for(
runtime_state: &Arc<RuntimeState>,
) -> Option<Arc<dyn RuntimeQueueStore>> {
if runtime_state.is_redis() {
let queue: Arc<dyn RuntimeQueueStore> = runtime_state.clone();
Some(queue)
} else {
None
}
let queue: Arc<dyn RuntimeQueueStore> = runtime_state.clone();
Some(queue)
}
fn spawn_scheduler_affinity_redis_write(
fn spawn_scheduler_affinity_runtime_write(
&self,
cache_key: &str,
target: &SchedulerAffinityTarget,
ttl: Duration,
epoch: u64,
) {
if self.runtime_state.is_memory() {
return;
}
let Ok(handle) = tokio::runtime::Handle::try_current() else {
return;
};
@@ -1088,7 +1081,7 @@ impl AppState {
if self.scheduler_affinity_epoch() != epoch {
return false;
}
self.spawn_scheduler_affinity_redis_write(cache_key, &target, ttl, epoch);
self.spawn_scheduler_affinity_runtime_write(cache_key, &target, ttl, epoch);
self.scheduler_affinity_cache.insert_for_epoch(
cache_key.to_string(),
target,

50
apps/aether-gateway/src/tests/architecture/runtime_and_security.rs
View File

@@ -104,9 +104,16 @@ fn runtime_state_owns_redis_runtime_boundaries() {
let mut violations = Vec::new();
for root in [
"apps/aether-gateway/src",
"apps/aether-tunnel/src",
"crates/aether-admin/src",
"crates/aether-billing/src",
"crates/aether-model-fetch/src",
"crates/aether-provider-pool/src",
"crates/aether-runtime/src",
"crates/aether-task-runtime/src",
"crates/aether-usage-runtime/src",
"crates/aether-provider-transport/src",
"crates/aether-wallet/src",
] {
for path in collect_workspace_rust_files(root) {
if path
@@ -132,6 +139,33 @@ fn runtime_state_owns_redis_runtime_boundaries() {
violations.join("\n")
);
let mut dependency_violations = Vec::new();
for manifest in [
"apps/aether-gateway/Cargo.toml",
"apps/aether-tunnel/Cargo.toml",
"crates/aether-admin/Cargo.toml",
"crates/aether-billing/Cargo.toml",
"crates/aether-model-fetch/Cargo.toml",
"crates/aether-provider-pool/Cargo.toml",
"crates/aether-provider-transport/Cargo.toml",
"crates/aether-runtime/Cargo.toml",
"crates/aether-task-runtime/Cargo.toml",
"crates/aether-usage-runtime/Cargo.toml",
"crates/aether-wallet/Cargo.toml",
] {
let cargo = read_workspace_file(manifest);
for forbidden in ["redis.workspace", "redis ="] {
if cargo.contains(forbidden) {
dependency_violations.push(format!("{manifest} -> {forbidden}"));
}
}
}
assert!(
dependency_violations.is_empty(),
"business/runtime crates must not depend on redis directly:\n{}",
dependency_violations.join("\n")
);
let mut runtime_state_violations = Vec::new();
for path in collect_workspace_rust_files("crates/aether-runtime-state/src") {
if path
@@ -160,6 +194,22 @@ fn runtime_state_owns_redis_runtime_boundaries() {
"only crates/aether-runtime-state/src/redis may depend on the redis crate directly:\n{}",
runtime_state_violations.join("\n")
);
let mut runtime_connection_violations = Vec::new();
for path in collect_workspace_rust_files("crates/aether-runtime-state/src") {
if path.ends_with("crates/aether-runtime-state/src/redis/client.rs") {
continue;
}
let source = production_workspace_source(&path);
if source.contains("get_multiplexed_async_connection") {
runtime_connection_violations.push(path.display().to_string());
}
}
assert!(
runtime_connection_violations.is_empty(),
"runtime Redis connections must be initialized only by redis/client.rs:\n{}",
runtime_connection_violations.join("\n")
);
}
#[test]

1278
crates/aether-runtime-state/src/lib.rs
View File

File diff suppressed because it is too large Load Diff

632
crates/aether-runtime-state/src/memory.rs
View File

@@ -4,7 +4,7 @@ use std::time::{Duration, Instant};
use tokio::sync::Mutex;
use crate::{RuntimeQueueEntry, RuntimeQueueReclaimConfig};
use crate::{DataLayerError, RuntimeQueueEntry, RuntimeQueueReclaimConfig};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MemoryRuntimeStateConfig {
@@ -37,9 +37,9 @@ 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>>>,
sets: Mutex<HashMap<String, MemorySetEntry>>,
scores: Mutex<HashMap<String, MemoryScoreEntry>>,
queues: Mutex<HashMap<String, MemoryQueueStream>>,
queue_seq: AtomicU64,
locks: Mutex<HashMap<String, MemoryLockEntry>>,
semaphores: Mutex<HashMap<String, BTreeMap<String, u64>>>,
@@ -52,6 +52,80 @@ struct MemoryCounterEntry {
expires_at: Instant,
}
#[derive(Debug, Default)]
struct MemorySetEntry {
members: BTreeSet<String>,
expires_at: Option<Instant>,
}
#[derive(Debug, Default)]
struct MemoryScoreEntry {
scores: BTreeMap<String, f64>,
expires_at: Option<Instant>,
}
#[derive(Debug, Default)]
struct MemoryQueueStream {
entries: VecDeque<MemoryQueuedEntry>,
groups: HashMap<String, MemoryConsumerGroup>,
expires_at: Option<Instant>,
}
trait MemoryExpiringKey {
fn is_expired(&self, now: Instant) -> bool;
fn set_expires_at(&mut self, expires_at: Instant);
}
impl MemoryExpiringKey for MemorySetEntry {
fn is_expired(&self, now: Instant) -> bool {
self.expires_at.is_some_and(|expires_at| now >= expires_at)
}
fn set_expires_at(&mut self, expires_at: Instant) {
self.expires_at = Some(expires_at);
}
}
impl MemoryExpiringKey for MemoryScoreEntry {
fn is_expired(&self, now: Instant) -> bool {
self.expires_at.is_some_and(|expires_at| now >= expires_at)
}
fn set_expires_at(&mut self, expires_at: Instant) {
self.expires_at = Some(expires_at);
}
}
impl MemoryExpiringKey for MemoryQueueStream {
fn is_expired(&self, now: Instant) -> bool {
self.expires_at.is_some_and(|expires_at| now >= expires_at)
}
fn set_expires_at(&mut self, expires_at: Instant) {
self.expires_at = Some(expires_at);
}
}
#[derive(Debug, Clone)]
struct MemoryQueuedEntry {
sequence: u64,
entry: RuntimeQueueEntry,
}
#[derive(Debug, Default)]
struct MemoryConsumerGroup {
last_delivered_sequence: u64,
pending: BTreeMap<String, MemoryPendingQueueEntry>,
}
#[derive(Debug, Clone)]
struct MemoryPendingQueueEntry {
sequence: u64,
entry: RuntimeQueueEntry,
consumer: String,
delivered_at: Instant,
}
#[derive(Debug, Clone)]
pub(crate) struct MemoryLockEntry {
pub(crate) token: String,
@@ -143,16 +217,66 @@ impl MemoryRuntimeBackend {
}
pub(crate) async fn kv_delete(&self, key: &str) -> bool {
self.kv.lock().await.remove(key).is_some()
let kv_deleted = self.kv.lock().await.remove(key).is_some();
let set_deleted = self.sets.lock().await.remove(key).is_some();
let score_deleted = self.scores.lock().await.remove(key).is_some();
let queue_deleted = self.queues.lock().await.remove(key).is_some();
kv_deleted || set_deleted || score_deleted || queue_deleted
}
pub(crate) async fn kv_delete_many(&self, keys: &[String]) -> usize {
let keys = keys.iter().cloned().collect::<BTreeSet<_>>();
let mut deleted = BTreeSet::new();
let mut kv = self.kv.lock().await;
keys.iter().filter(|key| kv.remove(*key).is_some()).count()
for key in &keys {
if kv.remove(key).is_some() {
deleted.insert(key.clone());
}
}
drop(kv);
let mut sets = self.sets.lock().await;
for key in &keys {
if sets.remove(key).is_some() {
deleted.insert(key.clone());
}
}
drop(sets);
let mut scores = self.scores.lock().await;
for key in &keys {
if scores.remove(key).is_some() {
deleted.insert(key.clone());
}
}
drop(scores);
let mut queues = self.queues.lock().await;
for key in &keys {
if queues.remove(key).is_some() {
deleted.insert(key.clone());
}
}
deleted.len()
}
pub(crate) async fn kv_exists(&self, key: &str) -> bool {
self.kv_get(key).await.is_some()
if self.kv_get(key).await.is_some() {
return true;
}
let now = Instant::now();
let mut sets = self.sets.lock().await;
prune_memory_key(&mut sets, key, now);
if sets.contains_key(key) {
return true;
}
drop(sets);
let mut scores = self.scores.lock().await;
prune_memory_key(&mut scores, key, now);
if scores.contains_key(key) {
return true;
}
drop(scores);
let mut queues = self.queues.lock().await;
prune_memory_key(&mut queues, key, now);
queues.contains_key(key)
}
pub(crate) async fn kv_ttl_seconds(&self, key: &str) -> Option<i64> {
@@ -177,16 +301,77 @@ impl MemoryRuntimeBackend {
)
}
pub(crate) async fn key_expire(&self, key: &str, ttl: Duration) -> bool {
let now = Instant::now();
if ttl.is_zero() {
let kv_deleted = self.kv.lock().await.remove(key).is_some();
let set_deleted = self.sets.lock().await.remove(key).is_some();
let score_deleted = self.scores.lock().await.remove(key).is_some();
let queue_deleted = self.queues.lock().await.remove(key).is_some();
return kv_deleted || set_deleted || score_deleted || queue_deleted;
}
let expires_at = now + ttl;
{
let mut kv = self.kv.lock().await;
if let Some(entry) = kv.get_mut(key) {
if entry.is_expired(now) {
kv.remove(key);
} else {
entry.expires_at = Some(expires_at);
return true;
}
}
}
if set_memory_key_expiry(&self.sets, key, expires_at, now).await {
return true;
}
if set_memory_key_expiry(&self.scores, key, expires_at, now).await {
return true;
}
if set_memory_key_expiry(&self.queues, key, expires_at, now).await {
return true;
}
false
}
pub(crate) async fn kv_scan(&self, pattern: &str) -> Vec<String> {
let now = Instant::now();
let mut keys = BTreeSet::new();
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
prune_kv(&mut kv, now);
keys.extend(
kv.keys()
.filter(|key| key_matches_pattern(key, pattern))
.cloned(),
);
drop(kv);
let mut sets = self.sets.lock().await;
prune_expiring_map(&mut sets, now);
keys.extend(
sets.keys()
.filter(|key| key_matches_pattern(key, pattern))
.cloned(),
);
drop(sets);
let mut scores = self.scores.lock().await;
prune_expiring_map(&mut scores, now);
keys.extend(
scores
.keys()
.filter(|key| key_matches_pattern(key, pattern))
.cloned(),
);
drop(scores);
let mut queues = self.queues.lock().await;
prune_expiring_map(&mut queues, now);
keys.extend(
queues
.keys()
.filter(|key| key_matches_pattern(key, pattern))
.cloned(),
);
keys.into_iter().collect()
}
pub(crate) async fn check_and_consume_rate_limit(
@@ -231,6 +416,8 @@ impl MemoryRuntimeBackend {
}
let mut remaining = None::<u32>;
let mut user_next = None::<u32>;
let mut key_next = None::<u32>;
let expires_at = now + ttl;
if user_limit > 0 {
let next = counters
@@ -247,6 +434,7 @@ impl MemoryRuntimeBackend {
})
.value;
remaining = Some(user_limit.saturating_sub(next));
user_next = Some(next);
}
if key_limit > 0 {
let next = counters
@@ -264,6 +452,15 @@ impl MemoryRuntimeBackend {
.value;
let key_remaining = key_limit.saturating_sub(next);
remaining = Some(remaining.map_or(key_remaining, |value| value.min(key_remaining)));
key_next = Some(next);
}
drop(counters);
if let Some(next) = user_next {
self.kv_set(user_key, next.to_string(), Some(ttl)).await;
}
if let Some(next) = key_next {
self.kv_set(key_key, next.to_string(), Some(ttl)).await;
}
Ok(crate::RateLimitCheck::Allowed {
remaining: remaining.unwrap_or(0),
@@ -271,11 +468,11 @@ impl MemoryRuntimeBackend {
}
pub(crate) async fn set_add(&self, key: &str, member: &str) -> bool {
self.sets
.lock()
.await
.entry(key.to_string())
let mut sets = self.sets.lock().await;
prune_memory_key(&mut sets, key, Instant::now());
sets.entry(key.to_string())
.or_default()
.members
.insert(member.to_string())
}
@@ -283,82 +480,113 @@ impl MemoryRuntimeBackend {
let Ok(mut sets) = self.sets.try_lock() else {
return false;
};
prune_memory_key(&mut sets, key, Instant::now());
sets.entry(key.to_string())
.or_default()
.members
.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))
let mut sets = self.sets.lock().await;
prune_memory_key(&mut sets, key, Instant::now());
sets.get_mut(key)
.is_some_and(|entry| entry.members.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())
let mut sets = self.sets.lock().await;
prune_memory_key(&mut sets, key, Instant::now());
sets.get(key)
.map(|entry| entry.members.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)
let mut sets = self.sets.lock().await;
prune_memory_key(&mut sets, key, Instant::now());
sets.get(key).map_or(0, |entry| entry.members.len())
}
pub(crate) async fn score_set(&self, key: &str, member: &str, score: f64) {
self.scores
.lock()
.await
let mut scores = self.scores.lock().await;
prune_memory_key(&mut scores, key, Instant::now());
scores
.entry(key.to_string())
.or_default()
.scores
.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;
let mut scores = self.scores.lock().await;
prune_memory_key(&mut scores, key, Instant::now());
members
.iter()
.map(|member| scores.get(key).and_then(|set| set.get(member)).copied())
.map(|member| {
scores
.get(key)
.and_then(|entry| entry.scores.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;
let mut scores = self.scores.lock().await;
prune_memory_key(&mut scores, key, Instant::now());
scores
.get(key)
.map(|set| {
set.iter()
.filter(|(_, score)| **score >= min_score)
.map(|(member, _)| member.clone())
.collect()
})
.map(|entry| sorted_score_members(&entry.scores, |score| score >= min_score))
.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 {
prune_memory_key(&mut scores, key, Instant::now());
let Some(entry) = scores.get_mut(key) else {
return 0;
};
let before = set.len();
set.retain(|_, score| *score > max_score);
before.saturating_sub(set.len())
let before = entry.scores.len();
entry.scores.retain(|_, score| *score > max_score);
before.saturating_sub(entry.scores.len())
}
pub(crate) async fn score_remove(&self, key: &str, member: &str) -> bool {
self.scores
.lock()
.await
let mut scores = self.scores.lock().await;
prune_memory_key(&mut scores, key, Instant::now());
scores
.get_mut(key)
.is_some_and(|set| set.remove(member).is_some())
.is_some_and(|entry| entry.scores.remove(member).is_some())
}
pub(crate) async fn score_remove_by_rank(&self, key: &str, start: i64, stop: i64) -> usize {
let mut scores = self.scores.lock().await;
prune_memory_key(&mut scores, key, Instant::now());
let Some(entry) = scores.get_mut(key) else {
return 0;
};
let Some((start, stop)) = normalize_redis_rank_range(entry.scores.len(), start, stop)
else {
return 0;
};
let members = sorted_score_members(&entry.scores, |_| true);
let remove = members
.into_iter()
.enumerate()
.filter_map(|(index, member)| (index >= start && index <= stop).then_some(member))
.collect::<Vec<_>>();
let before = entry.scores.len();
for member in remove {
entry.scores.remove(&member);
}
before.saturating_sub(entry.scores.len())
}
pub(crate) async fn score_len(&self, key: &str) -> usize {
self.scores.lock().await.get(key).map_or(0, BTreeMap::len)
let mut scores = self.scores.lock().await;
prune_memory_key(&mut scores, key, Instant::now());
scores.get(key).map_or(0, |entry| entry.scores.len())
}
pub(crate) async fn queue_append(
@@ -367,51 +595,206 @@ impl MemoryRuntimeBackend {
fields: BTreeMap<String, String>,
maxlen: Option<usize>,
) -> String {
let id = format!(
"{}-0",
self.queue_seq
.fetch_add(1, Ordering::Relaxed)
.saturating_add(1)
);
let sequence = self
.queue_seq
.fetch_add(1, Ordering::Relaxed)
.saturating_add(1);
let id = format!("{sequence}-0");
let mut queues = self.queues.lock().await;
let queue = queues.entry(stream.to_string()).or_default();
queue.push_back(RuntimeQueueEntry {
id: id.clone(),
fields,
prune_memory_key(&mut queues, stream, Instant::now());
let stream_state = queues.entry(stream.to_string()).or_default();
stream_state.entries.push_back(MemoryQueuedEntry {
sequence,
entry: RuntimeQueueEntry {
id: id.clone(),
fields,
},
});
if let Some(maxlen) = maxlen.filter(|value| *value > 0) {
while queue.len() > maxlen {
queue.pop_front();
while stream_state.entries.len() > maxlen {
let Some(removed) = stream_state.entries.pop_front() else {
break;
};
remove_pending_from_all_groups(stream_state, &removed.entry.id);
}
}
id
}
pub(crate) async fn queue_read(&self, stream: &str, count: usize) -> Vec<RuntimeQueueEntry> {
pub(crate) async fn queue_ensure_consumer_group(
&self,
stream: &str,
group: &str,
start_id: &str,
) -> Result<(), DataLayerError> {
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);
prune_memory_key(&mut queues, stream, Instant::now());
let stream_state = queues.entry(stream.to_string()).or_default();
if stream_state.groups.contains_key(group) {
return Ok(());
}
let last_delivered_sequence = match start_id {
"$" => stream_state
.entries
.back()
.map(|entry| entry.sequence)
.unwrap_or_default(),
_ => parse_memory_stream_sequence(start_id)?,
};
stream_state.groups.insert(
group.to_string(),
MemoryConsumerGroup {
last_delivered_sequence,
pending: BTreeMap::new(),
},
);
Ok(())
}
pub(crate) async fn queue_read(
&self,
stream: &str,
group: &str,
consumer: &str,
count: usize,
block_ms: Option<u64>,
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
let deadline = block_ms.map(|value| Instant::now() + Duration::from_millis(value.max(1)));
loop {
let entries = {
let mut queues = self.queues.lock().await;
prune_memory_key(&mut queues, stream, Instant::now());
let Some(stream_state) = queues.get_mut(stream) else {
return Err(DataLayerError::InvalidInput(format!(
"runtime queue stream {stream} does not exist"
)));
};
let Some(group_state) = stream_state.groups.get_mut(group) else {
return Err(DataLayerError::InvalidInput(format!(
"runtime queue group {group} does not exist for stream {stream}"
)));
};
let now = Instant::now();
let mut delivered = Vec::new();
let last_delivered_sequence = group_state.last_delivered_sequence;
let queued_entries = stream_state
.entries
.iter()
.filter(|entry| entry.sequence > last_delivered_sequence)
.take(count.max(1))
.cloned()
.collect::<Vec<_>>();
for queued in queued_entries {
group_state.last_delivered_sequence = queued.sequence;
group_state.pending.insert(
queued.entry.id.clone(),
MemoryPendingQueueEntry {
sequence: queued.sequence,
entry: queued.entry.clone(),
consumer: consumer.to_string(),
delivered_at: now,
},
);
delivered.push(queued.entry.clone());
}
delivered
};
if !entries.is_empty() {
return Ok(entries);
}
let Some(deadline) = deadline else {
return Ok(Vec::new());
};
if Instant::now() >= deadline {
return Ok(Vec::new());
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
entries
}
pub(crate) async fn queue_claim_stale(
&self,
_stream: &str,
_config: RuntimeQueueReclaimConfig,
) -> Vec<RuntimeQueueEntry> {
Vec::new()
stream: &str,
group: &str,
consumer: &str,
start_id: &str,
config: RuntimeQueueReclaimConfig,
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
let start_sequence = parse_memory_stream_sequence(start_id)?;
let min_idle = Duration::from_millis(config.min_idle_ms.max(1));
let now = Instant::now();
let mut queues = self.queues.lock().await;
prune_memory_key(&mut queues, stream, now);
let Some(stream_state) = queues.get_mut(stream) else {
return Err(DataLayerError::InvalidInput(format!(
"runtime queue stream {stream} does not exist"
)));
};
let Some(group_state) = stream_state.groups.get_mut(group) else {
return Err(DataLayerError::InvalidInput(format!(
"runtime queue group {group} does not exist for stream {stream}"
)));
};
let ids = group_state
.pending
.values()
.filter(|entry| entry.sequence >= start_sequence)
.filter(|entry| now.saturating_duration_since(entry.delivered_at) >= min_idle)
.map(|entry| (entry.sequence, entry.entry.id.clone()))
.collect::<Vec<_>>();
let mut ids = ids;
ids.sort_by_key(|(sequence, _)| *sequence);
let mut claimed = Vec::new();
for (_, id) in ids.into_iter().take(config.count.max(1)) {
if let Some(pending) = group_state.pending.get_mut(&id) {
pending.consumer = consumer.to_string();
pending.delivered_at = now;
claimed.push(pending.entry.clone());
}
}
Ok(claimed)
}
pub(crate) async fn queue_delete(&self, _stream: &str, _ids: &[String]) -> usize {
0
pub(crate) async fn queue_ack(
&self,
stream: &str,
group: &str,
ids: &[String],
) -> Result<usize, DataLayerError> {
let mut queues = self.queues.lock().await;
prune_memory_key(&mut queues, stream, Instant::now());
let Some(stream_state) = queues.get_mut(stream) else {
return Err(DataLayerError::InvalidInput(format!(
"runtime queue stream {stream} does not exist"
)));
};
let Some(group_state) = stream_state.groups.get_mut(group) else {
return Err(DataLayerError::InvalidInput(format!(
"runtime queue group {group} does not exist for stream {stream}"
)));
};
Ok(ids
.iter()
.filter(|id| group_state.pending.remove(*id).is_some())
.count())
}
pub(crate) async fn queue_delete(&self, stream: &str, ids: &[String]) -> usize {
let mut queues = self.queues.lock().await;
prune_memory_key(&mut queues, stream, Instant::now());
let Some(stream_state) = queues.get_mut(stream) else {
return 0;
};
let ids = ids.iter().cloned().collect::<BTreeSet<_>>();
let before = stream_state.entries.len();
stream_state
.entries
.retain(|entry| !ids.contains(&entry.entry.id));
for id in &ids {
remove_pending_from_all_groups(stream_state, id);
}
before.saturating_sub(stream_state.entries.len())
}
pub(crate) async fn lock_try_acquire(
@@ -530,6 +913,43 @@ fn prune_kv(kv: &mut HashMap<String, MemoryKvEntry>, now: Instant) {
kv.retain(|_, entry| !entry.is_expired(now));
}
fn prune_memory_key<T>(values: &mut HashMap<String, T>, key: &str, now: Instant)
where
T: MemoryExpiringKey,
{
if values.get(key).is_some_and(|entry| entry.is_expired(now)) {
values.remove(key);
}
}
fn prune_expiring_map<T>(values: &mut HashMap<String, T>, now: Instant)
where
T: MemoryExpiringKey,
{
values.retain(|_, entry| !entry.is_expired(now));
}
async fn set_memory_key_expiry<T>(
values: &Mutex<HashMap<String, T>>,
key: &str,
expires_at: Instant,
now: Instant,
) -> bool
where
T: MemoryExpiringKey,
{
let mut values = values.lock().await;
if values.get(key).is_some_and(|entry| entry.is_expired(now)) {
values.remove(key);
return false;
}
let Some(entry) = values.get_mut(key) else {
return false;
};
entry.set_expires_at(expires_at);
true
}
pub(crate) fn key_matches_pattern(key: &str, pattern: &str) -> bool {
match pattern.strip_suffix('*') {
Some(prefix) => key.starts_with(prefix),
@@ -537,6 +957,58 @@ pub(crate) fn key_matches_pattern(key: &str, pattern: &str) -> bool {
}
}
fn sorted_score_members<F>(scores: &BTreeMap<String, f64>, include: F) -> Vec<String>
where
F: Fn(f64) -> bool,
{
let mut entries = scores
.iter()
.filter_map(|(member, score)| include(*score).then_some((member.clone(), *score)))
.collect::<Vec<_>>();
entries.sort_by(|(left_member, left_score), (right_member, right_score)| {
left_score
.total_cmp(right_score)
.then_with(|| left_member.cmp(right_member))
});
entries.into_iter().map(|(member, _)| member).collect()
}
fn normalize_redis_rank_range(len: usize, start: i64, stop: i64) -> Option<(usize, usize)> {
if len == 0 {
return None;
}
let len = i64::try_from(len).ok()?;
let mut start = if start < 0 { len + start } else { start };
let mut stop = if stop < 0 { len + stop } else { stop };
if start < 0 {
start = 0;
}
if stop < 0 || start >= len || start > stop {
return None;
}
if stop >= len {
stop = len - 1;
}
Some((usize::try_from(start).ok()?, usize::try_from(stop).ok()?))
}
fn remove_pending_from_all_groups(stream: &mut MemoryQueueStream, id: &str) {
for group in stream.groups.values_mut() {
group.pending.remove(id);
}
}
fn parse_memory_stream_sequence(id: &str) -> Result<u64, DataLayerError> {
let Some((sequence, _)) = id.split_once('-') else {
return Err(DataLayerError::InvalidInput(format!(
"runtime queue stream id {id} must use redis stream id format"
)));
};
sequence.parse::<u64>().map_err(|err| {
DataLayerError::InvalidInput(format!("runtime queue stream id {id} is invalid: {err}"))
})
}
fn unix_time_ms() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)

201
crates/aether-runtime-state/src/redis/client.rs
View File

@@ -1,8 +1,13 @@
use crate::error::RedisResultExt;
use crate::redis::RedisKeyspace;
use crate::DataLayerError;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tracing::info;
pub type RedisClient = redis::Client;
pub(crate) type RedisClient = redis::Client;
pub(crate) type RedisManagedConnection = redis::aio::ConnectionManager;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct RedisClientConfig {
@@ -30,23 +35,209 @@ impl RedisClientConfig {
}
#[derive(Debug, Clone)]
pub struct RedisClientFactory {
pub(crate) struct RedisClientFactory {
config: RedisClientConfig,
}
impl RedisClientFactory {
pub fn new(config: RedisClientConfig) -> Result<Self, DataLayerError> {
pub(crate) fn new(config: RedisClientConfig) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self { config })
}
pub fn config(&self) -> &RedisClientConfig {
pub(crate) fn config(&self) -> &RedisClientConfig {
&self.config
}
pub fn connect_lazy(&self) -> Result<RedisClient, DataLayerError> {
pub(crate) fn connect_lazy(&self) -> Result<RedisClient, DataLayerError> {
RedisClient::open(self.config.url.clone()).map_redis_err()
}
pub(crate) async fn connect_router(
&self,
command_timeout_ms: Option<u64>,
) -> Result<RedisConnectionRouter, DataLayerError> {
RedisConnectionRouter::connect(self.connect_lazy()?, command_timeout_ms).await
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RedisConnectionLane {
Fast,
Stream,
BlockingStream,
Admin,
}
impl RedisConnectionLane {
pub(crate) const fn as_str(self) -> &'static str {
match self {
Self::Fast => "fast",
Self::Stream => "stream",
Self::BlockingStream => "blocking_stream",
Self::Admin => "admin",
}
}
}
#[derive(Clone)]
pub(crate) struct RedisConnectionRouter {
fast: RedisManagedConnection,
stream: RedisManagedConnection,
blocking_stream: RedisManagedConnection,
admin: RedisManagedConnection,
metrics: Arc<RedisConnectionMetrics>,
}
impl std::fmt::Debug for RedisConnectionRouter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RedisConnectionRouter")
.field("lanes", &["fast", "stream", "blocking_stream", "admin"])
.finish()
}
}
impl RedisConnectionRouter {
pub(crate) async fn connect(
client: RedisClient,
command_timeout_ms: Option<u64>,
) -> Result<Self, DataLayerError> {
let fast = connect_lane(
&client,
connection_manager_config(command_timeout_ms),
RedisConnectionLane::Fast,
)
.await?;
let stream = connect_lane(
&client,
connection_manager_config(command_timeout_ms),
RedisConnectionLane::Stream,
)
.await?;
let blocking_stream = connect_lane(
&client,
connection_manager_config(command_timeout_ms),
RedisConnectionLane::BlockingStream,
)
.await?;
let admin = connect_lane(
&client,
connection_manager_config(command_timeout_ms),
RedisConnectionLane::Admin,
)
.await?;
info!(
redis_lanes = "fast,stream,blocking_stream,admin",
"runtime redis connection lanes initialized"
);
Ok(Self {
fast,
stream,
blocking_stream,
admin,
metrics: Arc::new(RedisConnectionMetrics::default()),
})
}
pub(crate) fn connection(&self, lane: RedisConnectionLane) -> RedisManagedConnection {
match lane {
RedisConnectionLane::Fast => self.fast.clone(),
RedisConnectionLane::Stream => self.stream.clone(),
RedisConnectionLane::BlockingStream => self.blocking_stream.clone(),
RedisConnectionLane::Admin => self.admin.clone(),
}
}
pub(crate) fn record_error(&self, lane: RedisConnectionLane) {
self.metrics
.for_lane(lane)
.errors
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn record_timeout(&self, lane: RedisConnectionLane) {
self.metrics
.for_lane(lane)
.timeouts
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn lane_diagnostics(&self) -> Vec<RedisLaneDiagnostics> {
[
RedisConnectionLane::Fast,
RedisConnectionLane::Stream,
RedisConnectionLane::BlockingStream,
RedisConnectionLane::Admin,
]
.into_iter()
.map(|lane| {
let metrics = self.metrics.for_lane(lane);
RedisLaneDiagnostics {
lane: lane.as_str(),
command_errors: metrics.errors.load(Ordering::Relaxed),
command_timeouts: metrics.timeouts.load(Ordering::Relaxed),
}
})
.collect()
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct RedisLaneDiagnostics {
pub lane: &'static str,
pub command_errors: u64,
pub command_timeouts: u64,
}
#[derive(Default)]
struct RedisConnectionMetrics {
fast: RedisLaneMetrics,
stream: RedisLaneMetrics,
blocking_stream: RedisLaneMetrics,
admin: RedisLaneMetrics,
}
impl RedisConnectionMetrics {
fn for_lane(&self, lane: RedisConnectionLane) -> &RedisLaneMetrics {
match lane {
RedisConnectionLane::Fast => &self.fast,
RedisConnectionLane::Stream => &self.stream,
RedisConnectionLane::BlockingStream => &self.blocking_stream,
RedisConnectionLane::Admin => &self.admin,
}
}
}
#[derive(Default)]
struct RedisLaneMetrics {
errors: AtomicU64,
timeouts: AtomicU64,
}
fn connection_manager_config(
command_timeout_ms: Option<u64>,
) -> redis::aio::ConnectionManagerConfig {
let mut config = redis::aio::ConnectionManagerConfig::new();
if let Some(timeout_ms) = command_timeout_ms {
config = config.set_connection_timeout(Duration::from_millis(timeout_ms));
}
config
}
async fn connect_lane(
client: &RedisClient,
config: redis::aio::ConnectionManagerConfig,
lane: RedisConnectionLane,
) -> Result<RedisManagedConnection, DataLayerError> {
client
.get_connection_manager_with_config(config)
.await
.map_err(|err| {
DataLayerError::Redis(format!(
"failed to initialize runtime redis {} lane: {err}",
lane.as_str()
))
})
}
#[cfg(test)]

123
crates/aether-runtime-state/src/redis/kv.rs
View File

@@ -1,8 +1,10 @@
use std::future::Future;
use std::time::Duration;
use crate::error::RedisResultExt;
use crate::redis::{RedisClient, RedisKeyspace};
use crate::redis::{
run_lane_with_timeout, RedisClientConfig, RedisClientFactory, RedisConnectionLane,
RedisConnectionRouter, RedisKeyspace,
};
use crate::DataLayerError;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -40,27 +42,35 @@ impl RedisKvRunnerConfig {
#[derive(Debug, Clone)]
pub struct RedisKvRunner {
client: RedisClient,
connections: RedisConnectionRouter,
keyspace: RedisKeyspace,
config: RedisKvRunnerConfig,
}
impl RedisKvRunner {
pub fn new(
client: RedisClient,
pub(crate) fn new(
connections: RedisConnectionRouter,
keyspace: RedisKeyspace,
config: RedisKvRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
client,
connections,
keyspace,
config,
})
}
pub fn client(&self) -> &RedisClient {
&self.client
pub async fn from_config(
config: RedisClientConfig,
runner_config: RedisKvRunnerConfig,
) -> Result<Self, DataLayerError> {
let factory = RedisClientFactory::new(config)?;
let keyspace = factory.config().keyspace();
let connections = factory
.connect_router(runner_config.command_timeout_ms)
.await?;
Self::new(connections, keyspace, runner_config)
}
pub fn keyspace(&self) -> &RedisKeyspace {
@@ -79,12 +89,8 @@ impl RedisKvRunner {
) -> 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()?;
self.run_with_timeout(RedisConnectionLane::Fast, "redis kv setex", async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
redis::cmd("SETEX")
.arg(&namespaced_key)
.arg(resolved_ttl)
@@ -98,12 +104,8 @@ impl RedisKvRunner {
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()?;
self.run_with_timeout(RedisConnectionLane::Fast, "redis kv get", async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
redis::cmd("GET")
.arg(&namespaced_key)
.query_async(&mut connection)
@@ -115,12 +117,8 @@ impl RedisKvRunner {
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()?;
self.run_with_timeout(RedisConnectionLane::Fast, "redis kv getdel", async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
redis::cmd("GETDEL")
.arg(&namespaced_key)
.query_async(&mut connection)
@@ -133,12 +131,8 @@ impl RedisKvRunner {
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()?;
.run_with_timeout(RedisConnectionLane::Fast, "redis kv exists", async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
redis::cmd("EXISTS")
.arg(&namespaced_key)
.query_async::<i64>(&mut connection)
@@ -151,12 +145,8 @@ impl RedisKvRunner {
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()?;
self.run_with_timeout(RedisConnectionLane::Fast, "redis kv del", async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
redis::cmd("DEL")
.arg(&namespaced_key)
.query_async(&mut connection)
@@ -168,49 +158,33 @@ impl RedisKvRunner {
async fn run_with_timeout<T, F>(
&self,
lane: RedisConnectionLane,
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
}
run_lane_with_timeout(
&self.connections,
lane,
self.config.command_timeout_ms,
operation,
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")
}
use super::RedisKvRunnerConfig;
#[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();
fn validates_default_config() {
RedisKvRunnerConfig::default()
.validate()
.expect("default kv config should be valid");
}
#[test]
@@ -219,17 +193,6 @@ mod tests {
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());
assert!(config.validate().is_err());
}
}

143
crates/aether-runtime-state/src/redis/lock.rs
View File

@@ -1,8 +1,10 @@
use std::future::Future;
use std::time::Duration;
use crate::error::RedisResultExt;
use crate::redis::{RedisClient, RedisKeyspace};
use crate::redis::{
run_lane_with_timeout, RedisClientConfig, RedisClientFactory, RedisConnectionLane,
RedisConnectionRouter, RedisKeyspace,
};
use crate::DataLayerError;
use uuid::Uuid;
@@ -50,27 +52,35 @@ impl RedisLockRunnerConfig {
#[derive(Debug, Clone)]
pub struct RedisLockRunner {
client: RedisClient,
connections: RedisConnectionRouter,
keyspace: RedisKeyspace,
config: RedisLockRunnerConfig,
}
impl RedisLockRunner {
pub fn new(
client: RedisClient,
pub(crate) fn new(
connections: RedisConnectionRouter,
keyspace: RedisKeyspace,
config: RedisLockRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
client,
connections,
keyspace,
config,
})
}
pub fn client(&self) -> &RedisClient {
&self.client
pub async fn from_config(
config: RedisClientConfig,
runner_config: RedisLockRunnerConfig,
) -> Result<Self, DataLayerError> {
let factory = RedisClientFactory::new(config)?;
let keyspace = factory.config().keyspace();
let connections = factory
.connect_router(runner_config.command_timeout_ms)
.await?;
Self::new(connections, keyspace, runner_config)
}
pub fn keyspace(&self) -> &RedisKeyspace {
@@ -92,12 +102,8 @@ impl RedisLockRunner {
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()?;
self.run_with_timeout(RedisConnectionLane::Fast, "redis lock acquire", async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
let status = redis::cmd("SET")
.arg(&key.0)
.arg(&token)
@@ -120,12 +126,8 @@ impl RedisLockRunner {
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()?;
self.run_with_timeout(RedisConnectionLane::Fast, "redis lock release", async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
let deleted = redis::Script::new(
"if redis.call('get', KEYS[1]) == ARGV[1] then \
return redis.call('del', KEYS[1]) \
@@ -151,12 +153,8 @@ impl RedisLockRunner {
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()?;
self.run_with_timeout(RedisConnectionLane::Fast, "redis lock renew", async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
let renewed = redis::Script::new(
"if redis.call('get', KEYS[1]) == ARGV[1] then \
return redis.call('pexpire', KEYS[1], ARGV[2]) \
@@ -177,21 +175,21 @@ impl RedisLockRunner {
async fn run_with_timeout<T, F>(
&self,
lane: RedisConnectionLane,
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
}
run_lane_with_timeout(
&self.connections,
lane,
self.config.command_timeout_ms,
operation,
future,
)
.await
}
fn resolve_ttl_ms(&self, ttl_ms: Option<u64>) -> Result<u64, DataLayerError> {
@@ -241,29 +239,10 @@ fn validate_lease(lease: &RedisLockLease) -> Result<(), DataLayerError> {
#[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")
}
use super::{
validate_key, validate_lease, validate_owner, RedisLockKey, RedisLockLease,
RedisLockRunnerConfig,
};
#[test]
fn validates_runner_config() {
@@ -283,41 +262,21 @@ mod tests {
#[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();
RedisLockRunnerConfig::default()
.validate()
.expect("default lock config should be valid");
}
#[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());
#[test]
fn rejects_invalid_owner_or_lease_before_network() {
assert!(validate_owner("").is_err());
assert!(validate_key(&RedisLockKey(String::new())).is_err());
assert!(validate_lease(&RedisLockLease {
key: RedisLockKey("aether:lock:poller".to_string()),
owner: "worker-1".to_string(),
token: String::new(),
ttl_ms: 1_000,
})
.is_err());
}
}

36
crates/aether-runtime-state/src/redis/mod.rs
View File

@@ -2,12 +2,14 @@ mod client;
mod kv;
mod lock;
mod namespace;
mod runtime;
mod stream;
pub use client::{RedisClient, RedisClientConfig, RedisClientFactory};
pub use client::{RedisClientConfig, RedisLaneDiagnostics};
pub use kv::{RedisKvRunner, RedisKvRunnerConfig};
pub use lock::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
pub use namespace::RedisKeyspace;
pub use runtime::RedisRuntimeDiagnostics;
pub use stream::{
RedisConsumerGroup, RedisConsumerName, RedisStreamEntry, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner, RedisStreamRunnerConfig,
@@ -16,6 +18,9 @@ pub use stream::{
pub(crate) type RedisCmd = redis::Cmd;
pub(crate) type RedisScript = redis::Script;
pub(crate) use client::{RedisClientFactory, RedisConnectionLane, RedisConnectionRouter};
pub(crate) use runtime::RedisRuntimeRunner;
pub(crate) fn cmd(name: &str) -> RedisCmd {
redis::cmd(name)
}
@@ -23,3 +28,32 @@ pub(crate) fn cmd(name: &str) -> RedisCmd {
pub(crate) fn script(source: &str) -> RedisScript {
redis::Script::new(source)
}
pub(crate) async fn run_lane_with_timeout<T, F>(
connections: &RedisConnectionRouter,
lane: RedisConnectionLane,
timeout_ms: Option<u64>,
operation: &'static str,
future: F,
) -> Result<T, crate::DataLayerError>
where
F: std::future::Future<Output = Result<T, crate::DataLayerError>>,
{
let result = if let Some(timeout_ms) = timeout_ms {
match tokio::time::timeout(std::time::Duration::from_millis(timeout_ms), future).await {
Ok(result) => result,
Err(_) => {
connections.record_timeout(lane);
return Err(crate::DataLayerError::TimedOut(format!(
"{operation} exceeded {timeout_ms}ms timeout"
)));
}
}
} else {
future.await
};
if result.is_err() {
connections.record_error(lane);
}
result
}

699
crates/aether-runtime-state/src/redis/runtime.rs Normal file
View File

@@ -0,0 +1,699 @@
use std::time::Duration;
use crate::error::RedisResultExt;
use crate::redis::{
cmd, run_lane_with_timeout, script, RedisCmd, RedisConnectionLane, RedisConnectionRouter,
RedisKeyspace, RedisLaneDiagnostics,
};
use crate::{
DataLayerError, RateLimitCheck, RateLimitInput, RateLimitScope, RuntimeSemaphoreError,
};
const RATE_LIMIT_CHECK_AND_CONSUME_SCRIPT: &str = r#"
local user_key = KEYS[1]
local key_key = KEYS[2]
local user_limit = tonumber(ARGV[1])
local key_limit = tonumber(ARGV[2])
local ttl = tonumber(ARGV[3])
local user_count = 0
if user_limit > 0 then
user_count = tonumber(redis.call('GET', user_key) or '0')
if user_count >= user_limit then
return {0, 1, user_limit, 0}
end
end
local key_count = 0
if key_limit > 0 then
key_count = tonumber(redis.call('GET', key_key) or '0')
if key_count >= key_limit then
return {0, 2, key_limit, 0}
end
end
local remaining = -1
if user_limit > 0 then
user_count = redis.call('INCR', user_key)
redis.call('EXPIRE', user_key, ttl)
remaining = user_limit - user_count
end
if key_limit > 0 then
key_count = redis.call('INCR', key_key)
redis.call('EXPIRE', key_key, ttl)
local key_remaining = key_limit - key_count
if remaining == -1 or key_remaining < remaining then
remaining = key_remaining
end
end
return {1, 0, 0, remaining}
"#;
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct RedisRuntimeDiagnostics {
pub connected_clients: Option<u64>,
pub total_connections_received: Option<u64>,
pub lanes: Vec<RedisLaneDiagnostics>,
}
#[derive(Debug, Clone)]
pub(crate) struct RedisRuntimeRunner {
connections: RedisConnectionRouter,
keyspace: RedisKeyspace,
command_timeout_ms: Option<u64>,
}
impl RedisRuntimeRunner {
pub(crate) fn new(
connections: RedisConnectionRouter,
keyspace: RedisKeyspace,
command_timeout_ms: Option<u64>,
) -> Self {
Self {
connections,
keyspace,
command_timeout_ms,
}
}
pub(crate) async fn ping(&self) -> Result<(), DataLayerError> {
let pong = self
.query_string(RedisConnectionLane::Fast, "runtime redis ping", cmd("PING"))
.await?;
if pong.eq_ignore_ascii_case("PONG") {
Ok(())
} else {
Err(DataLayerError::UnexpectedValue(format!(
"unexpected runtime redis ping response {pong}"
)))
}
}
pub(crate) async fn diagnostics(&self) -> Result<RedisRuntimeDiagnostics, DataLayerError> {
let info = self
.query_string(
RedisConnectionLane::Admin,
"runtime redis diagnostics",
cmd("INFO"),
)
.await?;
Ok(parse_diagnostics(
&info,
self.connections.lane_diagnostics(),
))
}
pub(crate) async fn kv_set_plain(
&self,
key: &str,
value: String,
) -> Result<(), DataLayerError> {
let namespaced_key = self.keyspace.key(key);
let mut command = cmd("SET");
command.arg(namespaced_key).arg(value);
self.query_string(RedisConnectionLane::Fast, "runtime kv set", command)
.await?;
Ok(())
}
pub(crate) async fn kv_set_with_ttl(
&self,
key: &str,
value: String,
ttl: Duration,
) -> Result<(), DataLayerError> {
let namespaced_key = self.keyspace.key(key);
let mut command = cmd("PSETEX");
command
.arg(namespaced_key)
.arg(u64::try_from(ttl.as_millis().max(1)).unwrap_or(u64::MAX))
.arg(value);
self.query_string(RedisConnectionLane::Fast, "runtime kv set ttl", command)
.await?;
Ok(())
}
pub(crate) async fn kv_get_many(
&self,
keys: &[String],
) -> Result<Vec<Option<String>>, DataLayerError> {
let namespaced = keys
.iter()
.map(|key| self.keyspace.key(key))
.collect::<Vec<_>>();
let mut command = cmd("MGET");
command.arg(&namespaced);
self.query(RedisConnectionLane::Fast, "runtime kv mget", command)
.await
}
pub(crate) async fn kv_delete_many(&self, keys: &[String]) -> Result<usize, DataLayerError> {
let prefix = self.keyspace.key("");
let namespaced = keys
.iter()
.map(|key| {
if key_belongs_to_prefix(key, &prefix) {
key.clone()
} else {
self.keyspace.key(key)
}
})
.collect::<Vec<_>>();
let mut command = cmd("DEL");
command.arg(&namespaced);
let deleted = self
.query_i64(
RedisConnectionLane::Admin,
"runtime kv delete many",
command,
)
.await?;
Ok(usize::try_from(deleted).unwrap_or(0))
}
pub(crate) async fn kv_ttl_seconds(&self, key: &str) -> Result<Option<i64>, DataLayerError> {
let namespaced_key = self.keyspace.key(key);
let mut command = cmd("TTL");
command.arg(&namespaced_key);
let ttl = self
.query_i64(RedisConnectionLane::Fast, "runtime kv ttl", command)
.await?;
Ok((ttl >= -1).then_some(ttl))
}
pub(crate) async fn scan_keys(
&self,
pattern: &str,
count: usize,
) -> Result<Vec<String>, DataLayerError> {
let pattern = self.keyspace.key(pattern);
run_lane_with_timeout(
&self.connections,
RedisConnectionLane::Admin,
self.command_timeout_ms,
"runtime scan keys",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Admin);
let mut cursor = 0u64;
let mut keys = Vec::new();
loop {
let (next_cursor, mut batch) = cmd("SCAN")
.arg(cursor)
.arg("MATCH")
.arg(&pattern)
.arg("COUNT")
.arg(count.max(1))
.query_async::<(u64, Vec<String>)>(&mut connection)
.await
.map_redis_err()?;
keys.append(&mut batch);
if next_cursor == 0 {
break;
}
cursor = next_cursor;
}
keys.sort();
Ok(keys)
},
)
.await
}
pub(crate) async fn check_and_consume_rate_limit(
&self,
input: RateLimitInput<'_>,
) -> Result<RateLimitCheck, DataLayerError> {
let user_key = self.keyspace.key(input.user_key);
let key_key = self.keyspace.key(input.key_key);
let raw = run_lane_with_timeout(
&self.connections,
RedisConnectionLane::Fast,
self.command_timeout_ms,
"runtime rate limit check",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
script(RATE_LIMIT_CHECK_AND_CONSUME_SCRIPT)
.key(user_key)
.key(key_key)
.arg(i64::from(input.user_limit))
.arg(i64::from(input.key_limit))
.arg(i64::try_from(input.ttl_seconds.max(1)).unwrap_or(i64::MAX))
.invoke_async::<Vec<i64>>(&mut connection)
.await
.map_redis_err()
},
)
.await?;
if raw.first().copied().unwrap_or_default() == 1 {
return Ok(RateLimitCheck::Allowed {
remaining: raw
.get(3)
.copied()
.and_then(|value| u32::try_from(value).ok())
.unwrap_or_default(),
});
}
let scope = match raw.get(1).copied().unwrap_or_default() {
2 => RateLimitScope::Key,
_ => RateLimitScope::User,
};
let limit = raw
.get(2)
.copied()
.and_then(|value| u32::try_from(value).ok())
.unwrap_or(match scope {
RateLimitScope::User => input.user_limit,
RateLimitScope::Key => input.key_limit,
});
Ok(RateLimitCheck::Rejected { scope, limit })
}
pub(crate) async fn set_add(&self, key: &str, member: &str) -> Result<bool, DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("SADD");
command.arg(&key).arg(member);
Ok(self
.query_i64(RedisConnectionLane::Fast, "runtime set add", command)
.await?
> 0)
}
pub(crate) async fn set_remove(&self, key: &str, member: &str) -> Result<bool, DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("SREM");
command.arg(&key).arg(member);
Ok(self
.query_i64(RedisConnectionLane::Fast, "runtime set remove", command)
.await?
> 0)
}
pub(crate) async fn set_members(&self, key: &str) -> Result<Vec<String>, DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("SMEMBERS");
command.arg(&key);
let mut values = self
.query::<Vec<String>>(RedisConnectionLane::Admin, "runtime set members", command)
.await?;
values.sort();
Ok(values)
}
pub(crate) async fn set_len(&self, key: &str) -> Result<usize, DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("SCARD");
command.arg(&key);
let len = self
.query_i64(RedisConnectionLane::Fast, "runtime set len", command)
.await?;
Ok(usize::try_from(len).unwrap_or(0))
}
pub(crate) async fn score_set(
&self,
key: &str,
member: &str,
score: f64,
) -> Result<(), DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("ZADD");
command.arg(&key).arg(score).arg(member);
self.query_i64(RedisConnectionLane::Fast, "runtime score set", command)
.await?;
Ok(())
}
pub(crate) async fn score_many(
&self,
key: &str,
members: &[String],
) -> Result<Vec<Option<f64>>, DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("ZMSCORE");
command.arg(&key);
for member in members {
command.arg(member);
}
self.query(RedisConnectionLane::Fast, "runtime score many", command)
.await
}
pub(crate) async fn score_range_by_min(
&self,
key: &str,
min_score: f64,
) -> Result<Vec<String>, DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("ZRANGEBYSCORE");
command.arg(&key).arg(min_score).arg("+inf");
self.query(RedisConnectionLane::Admin, "runtime score range", command)
.await
}
pub(crate) async fn score_remove_by_score(
&self,
key: &str,
max_score: f64,
) -> Result<usize, DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("ZREMRANGEBYSCORE");
command.arg(&key).arg("-inf").arg(max_score);
let removed = self
.query_i64(RedisConnectionLane::Admin, "runtime score trim", command)
.await?;
Ok(usize::try_from(removed).unwrap_or(0))
}
pub(crate) async fn score_remove(
&self,
key: &str,
member: &str,
) -> Result<bool, DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("ZREM");
command.arg(&key).arg(member);
Ok(self
.query_i64(RedisConnectionLane::Fast, "runtime score remove", command)
.await?
> 0)
}
pub(crate) async fn score_remove_by_rank(
&self,
key: &str,
start: i64,
stop: i64,
) -> Result<usize, DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("ZREMRANGEBYRANK");
command.arg(&key).arg(start).arg(stop);
let removed = self
.query_i64(
RedisConnectionLane::Admin,
"runtime score rank trim",
command,
)
.await?;
Ok(usize::try_from(removed).unwrap_or(0))
}
pub(crate) async fn score_len(&self, key: &str) -> Result<usize, DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("ZCARD");
command.arg(&key);
let len = self
.query_i64(RedisConnectionLane::Fast, "runtime score len", command)
.await?;
Ok(usize::try_from(len).unwrap_or(0))
}
pub(crate) async fn key_expire(
&self,
key: &str,
ttl: Duration,
) -> Result<bool, DataLayerError> {
let key = self.keyspace.key(key);
let mut command = cmd("PEXPIRE");
command
.arg(&key)
.arg(u64::try_from(ttl.as_millis()).unwrap_or(u64::MAX));
Ok(self
.query_i64(RedisConnectionLane::Fast, "runtime key expire", command)
.await?
> 0)
}
pub(crate) async fn semaphore_try_acquire(
&self,
gate: &'static str,
limit: usize,
key: &str,
token: &str,
lease_ttl_ms: u64,
timeout_ms: Option<u64>,
) -> Result<(i64, i64), RuntimeSemaphoreError> {
let now_ms = crate::unix_time_ms();
let expires_at_ms = now_ms.saturating_add(lease_ttl_ms);
let key = self.keyspace.key(key);
let timeout_ms = timeout_ms.or(self.command_timeout_ms);
run_lane_with_timeout(
&self.connections,
RedisConnectionLane::Fast,
timeout_ms,
"runtime semaphore acquire",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
script(
"redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]); \
local count = redis.call('ZCARD', KEYS[1]); \
if count >= tonumber(ARGV[3]) then \
redis.call('PEXPIRE', KEYS[1], ARGV[5]); \
return {0, count}; \
end; \
redis.call('ZADD', KEYS[1], ARGV[2], ARGV[4]); \
count = redis.call('ZCARD', KEYS[1]); \
redis.call('PEXPIRE', KEYS[1], ARGV[5]); \
return {1, count};",
)
.key(&key)
.arg(now_ms as i64)
.arg(expires_at_ms as i64)
.arg(limit as i64)
.arg(token)
.arg(lease_ttl_ms as i64)
.invoke_async::<(i64, i64)>(&mut connection)
.await
.map_redis_err()
},
)
.await
.map_err(|err| RuntimeSemaphoreError::Unavailable {
gate,
limit,
message: format!("acquire failed: {err}"),
})
}
pub(crate) async fn semaphore_renew(
&self,
gate: &'static str,
limit: usize,
key: &str,
token: &str,
lease_ttl_ms: u64,
timeout_ms: Option<u64>,
) -> Result<i64, RuntimeSemaphoreError> {
let now_ms = crate::unix_time_ms();
let expires_at_ms = now_ms.saturating_add(lease_ttl_ms);
let key = self.keyspace.key(key);
let timeout_ms = timeout_ms.or(self.command_timeout_ms);
run_lane_with_timeout(
&self.connections,
RedisConnectionLane::Fast,
timeout_ms,
"runtime semaphore renew",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
script(
"redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]); \
local score = redis.call('ZSCORE', KEYS[1], ARGV[2]); \
if not score then return 0; end; \
redis.call('ZADD', KEYS[1], 'XX', ARGV[3], ARGV[2]); \
redis.call('PEXPIRE', KEYS[1], ARGV[4]); \
return 1;",
)
.key(&key)
.arg(now_ms as i64)
.arg(token)
.arg(expires_at_ms as i64)
.arg(lease_ttl_ms as i64)
.invoke_async::<i64>(&mut connection)
.await
.map_redis_err()
},
)
.await
.map_err(|err| RuntimeSemaphoreError::Unavailable {
gate,
limit,
message: format!("renew failed: {err}"),
})
}
pub(crate) async fn semaphore_release(
&self,
gate: &'static str,
limit: usize,
key: &str,
token: &str,
timeout_ms: Option<u64>,
) -> Result<(), RuntimeSemaphoreError> {
let key = self.keyspace.key(key);
let timeout_ms = timeout_ms.or(self.command_timeout_ms);
run_lane_with_timeout(
&self.connections,
RedisConnectionLane::Fast,
timeout_ms,
"runtime semaphore release",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
script(
"local removed = redis.call('ZREM', KEYS[1], ARGV[1]); \
if removed > 0 and redis.call('ZCARD', KEYS[1]) == 0 then \
redis.call('DEL', KEYS[1]); \
end; \
return removed;",
)
.key(&key)
.arg(token)
.invoke_async::<i64>(&mut connection)
.await
.map(|_| ())
.map_redis_err()
},
)
.await
.map_err(|err| RuntimeSemaphoreError::Unavailable {
gate,
limit,
message: format!("release failed: {err}"),
})
}
pub(crate) async fn semaphore_live_count(
&self,
gate: &'static str,
limit: usize,
key: &str,
timeout_ms: Option<u64>,
) -> Result<usize, RuntimeSemaphoreError> {
let now_ms = crate::unix_time_ms();
let key = self.keyspace.key(key);
let timeout_ms = timeout_ms.or(self.command_timeout_ms);
run_lane_with_timeout(
&self.connections,
RedisConnectionLane::Fast,
timeout_ms,
"runtime semaphore snapshot",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
script(
"redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]); \
return redis.call('ZCARD', KEYS[1]);",
)
.key(&key)
.arg(now_ms as i64)
.invoke_async::<i64>(&mut connection)
.await
.map(|value| value.max(0) as usize)
.map_redis_err()
},
)
.await
.map_err(|err| RuntimeSemaphoreError::Unavailable {
gate,
limit,
message: format!("snapshot failed: {err}"),
})
}
async fn query<T>(
&self,
lane: RedisConnectionLane,
operation: &'static str,
command: RedisCmd,
) -> Result<T, DataLayerError>
where
T: redis::FromRedisValue,
{
run_lane_with_timeout(
&self.connections,
lane,
self.command_timeout_ms,
operation,
async {
let mut connection = self.connections.connection(lane);
command
.query_async::<T>(&mut connection)
.await
.map_redis_err()
},
)
.await
}
async fn query_i64(
&self,
lane: RedisConnectionLane,
operation: &'static str,
command: RedisCmd,
) -> Result<i64, DataLayerError> {
self.query(lane, operation, command).await
}
async fn query_string(
&self,
lane: RedisConnectionLane,
operation: &'static str,
command: RedisCmd,
) -> Result<String, DataLayerError> {
self.query(lane, operation, command).await
}
}
fn parse_diagnostics(info: &str, lanes: Vec<RedisLaneDiagnostics>) -> RedisRuntimeDiagnostics {
RedisRuntimeDiagnostics {
connected_clients: parse_info_u64(info, "connected_clients"),
total_connections_received: parse_info_u64(info, "total_connections_received"),
lanes,
}
}
fn parse_info_u64(info: &str, key: &str) -> Option<u64> {
info.lines().find_map(|line| {
let (name, value) = line.split_once(':')?;
(name == key)
.then(|| value.trim().parse::<u64>().ok())
.flatten()
})
}
fn key_belongs_to_prefix(key: &str, prefix: &str) -> bool {
prefix.is_empty()
|| key == prefix
|| key
.strip_prefix(prefix)
.is_some_and(|rest| rest.starts_with(':'))
}
#[cfg(test)]
mod tests {
use super::{key_belongs_to_prefix, parse_diagnostics, RedisRuntimeDiagnostics};
#[test]
fn parses_runtime_diagnostics_from_info() {
let parsed = parse_diagnostics(
"# Clients\r\nconnected_clients:5\r\n# Stats\r\ntotal_connections_received:42\r\n",
Vec::new(),
);
assert_eq!(
parsed,
RedisRuntimeDiagnostics {
connected_clients: Some(5),
total_connections_received: Some(42),
lanes: Vec::new(),
}
);
}
#[test]
fn detects_namespaced_key_prefix_on_boundary() {
assert!(key_belongs_to_prefix("aether:cache:item", "aether"));
assert!(key_belongs_to_prefix("aether", "aether"));
assert!(key_belongs_to_prefix("raw:key", ""));
assert!(!key_belongs_to_prefix("aetherish:cache:item", "aether"));
}
}

217
crates/aether-runtime-state/src/redis/stream.rs
View File

@@ -1,13 +1,15 @@
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::redis::{
run_lane_with_timeout, RedisClientConfig, RedisClientFactory, RedisConnectionLane,
RedisConnectionRouter, RedisKeyspace,
};
use crate::DataLayerError;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -113,27 +115,35 @@ impl RedisStreamRunnerConfig {
#[derive(Debug, Clone)]
pub struct RedisStreamRunner {
client: RedisClient,
connections: RedisConnectionRouter,
keyspace: RedisKeyspace,
config: RedisStreamRunnerConfig,
}
impl RedisStreamRunner {
pub fn new(
client: RedisClient,
pub(crate) fn new(
connections: RedisConnectionRouter,
keyspace: RedisKeyspace,
config: RedisStreamRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
client,
connections,
keyspace,
config,
})
}
pub fn client(&self) -> &RedisClient {
&self.client
pub async fn from_config(
config: RedisClientConfig,
runner_config: RedisStreamRunnerConfig,
) -> Result<Self, DataLayerError> {
let factory = RedisClientFactory::new(config)?;
let keyspace = factory.config().keyspace();
let connections = factory
.connect_router(runner_config.command_timeout_ms)
.await?;
Self::new(connections, keyspace, runner_config)
}
pub fn keyspace(&self) -> &RedisKeyspace {
@@ -144,6 +154,13 @@ impl RedisStreamRunner {
self.config
}
pub(crate) fn with_config(
&self,
config: RedisStreamRunnerConfig,
) -> Result<Self, DataLayerError> {
Self::new(self.connections.clone(), self.keyspace.clone(), config)
}
pub async fn ensure_consumer_group(
&self,
stream: &RedisStreamName,
@@ -154,27 +171,27 @@ impl RedisStreamRunner {
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;
self.run_with_timeout(
RedisConnectionLane::Stream,
"redis stream ensure consumer group",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Stream);
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)),
}
})
match result {
Ok(_) => Ok(()),
Err(err) if err.code() == Some("BUSYGROUP") => Ok(()),
Err(err) => Err(redis_error(err)),
}
},
)
.await
}
@@ -199,12 +216,8 @@ impl RedisStreamRunner {
));
}
self.run_with_timeout("redis stream append", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
self.run_with_timeout(RedisConnectionLane::Stream, "redis stream append", async {
let mut connection = self.connections.connection(RedisConnectionLane::Stream);
let mut command = redis::cmd("XADD");
command.arg(&stream.0);
if let Some(maxlen) = maxlen.filter(|value| *value > 0) {
@@ -256,12 +269,13 @@ impl RedisStreamRunner {
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 lane = if self.config.read_block_ms.is_some() {
RedisConnectionLane::BlockingStream
} else {
RedisConnectionLane::Stream
};
self.run_with_timeout(lane, "redis stream read group", async {
let mut connection = self.connections.connection(lane);
let mut command = redis::cmd("XREADGROUP");
command
.arg("GROUP")
@@ -312,12 +326,8 @@ impl RedisStreamRunner {
return Ok(0);
}
self.run_with_timeout("redis stream ack", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
self.run_with_timeout(RedisConnectionLane::Stream, "redis stream ack", async {
let mut connection = self.connections.connection(RedisConnectionLane::Stream);
let mut command = redis::cmd("XACK");
command.arg(&stream.0).arg(&group.0);
for id in ids {
@@ -341,12 +351,8 @@ impl RedisStreamRunner {
return Ok(0);
}
self.run_with_timeout("redis stream delete", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
self.run_with_timeout(RedisConnectionLane::Stream, "redis stream delete", async {
let mut connection = self.connections.connection(RedisConnectionLane::Stream);
let mut command = redis::cmd("XDEL");
command.arg(&stream.0);
for id in ids {
@@ -374,12 +380,8 @@ impl RedisStreamRunner {
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()?;
self.run_with_timeout(RedisConnectionLane::Stream, "redis stream reclaim", async {
let mut connection = self.connections.connection(RedisConnectionLane::Stream);
let reply = redis::cmd("XAUTOCLAIM")
.arg(&stream.0)
.arg(&group.0)
@@ -399,21 +401,21 @@ impl RedisStreamRunner {
async fn run_with_timeout<T, F>(
&self,
lane: RedisConnectionLane,
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
}
run_lane_with_timeout(
&self.connections,
lane,
self.config.command_timeout_ms,
operation,
future,
)
.await
}
}
@@ -571,31 +573,12 @@ mod tests {
use std::collections::BTreeMap;
use super::{
parse_reclaim_result, RedisConsumerGroup, RedisConsumerName, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner,
RedisStreamRunnerConfig,
parse_reclaim_result, validate_consumer, validate_group, validate_stream_name,
validate_stream_position, RedisConsumerName, RedisStreamName, RedisStreamReclaimConfig,
RedisStreamReclaimResult, 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 {
@@ -643,55 +626,17 @@ mod tests {
#[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();
RedisStreamRunnerConfig::default()
.validate()
.expect("default stream config should be valid");
}
#[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 rejects_invalid_inputs_before_network() {
assert!(validate_stream_name(&RedisStreamName(String::new())).is_err());
assert!(validate_group(&super::RedisConsumerGroup(String::new())).is_err());
assert!(validate_consumer(&RedisConsumerName(String::new())).is_err());
assert!(validate_stream_position("").is_err());
}
#[test]

29
crates/aether-testkit/src/bin/dependency_pressure_baseline.rs
View File

@@ -8,9 +8,9 @@ use aether_data::driver::postgres::{
};
use aether_data::PostgresBackend;
use aether_runtime_state::{
RedisClientConfig, RedisClientFactory, RedisConsumerGroup, RedisConsumerName, RedisKeyspace,
RedisLockLease, RedisLockRunner, RedisLockRunnerConfig, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamRunner, RedisStreamRunnerConfig,
RedisClientConfig, RedisConsumerGroup, RedisConsumerName, RedisKeyspace, RedisLockLease,
RedisLockRunner, RedisLockRunnerConfig, RedisStreamName, RedisStreamReclaimConfig,
RedisStreamRunner, RedisStreamRunnerConfig,
};
use aether_testkit::{init_test_runtime_for, ManagedPostgresServer, ManagedRedisServer};
use futures_util::stream::{self, StreamExt};
@@ -189,12 +189,11 @@ async fn run_suite(
})
.expect("postgres url should resolve");
let redis_factory = RedisClientFactory::new(RedisClientConfig {
let redis_config = RedisClientConfig {
url: redis_url.clone(),
key_prefix: Some(format!("aether-dependency-pressure-{}", std::process::id())),
})?;
let redis_client = redis_factory.connect_lazy()?;
let redis_keyspace = redis_factory.config().keyspace();
};
let redis_keyspace = redis_config.keyspace();
let postgres_backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: postgres_url.clone(),
min_connections: 1,
@@ -208,23 +207,23 @@ async fn run_suite(
bootstrap_postgres_lease_table(postgres_backend.pool_clone(), config).await?;
let lock_runner = RedisLockRunner::new(
redis_client.clone(),
redis_keyspace.clone(),
let lock_runner = RedisLockRunner::from_config(
redis_config.clone(),
RedisLockRunnerConfig {
command_timeout_ms: Some(config.timeout.as_millis() as u64),
default_ttl_ms: 5_000,
},
)?;
let stream_runner = RedisStreamRunner::new(
redis_client.clone(),
redis_keyspace.clone(),
)
.await?;
let stream_runner = RedisStreamRunner::from_config(
redis_config,
RedisStreamRunnerConfig {
command_timeout_ms: Some(config.timeout.as_millis() as u64),
read_block_ms: Some(10),
read_count: 64,
},
)?;
)
.await?;
let lease_runner = postgres_backend.lease_runner(PostgresLeaseRunnerConfig {
statement_timeout_ms: Some(config.timeout.as_millis() as u64),
lock_timeout_ms: Some(1_000),

19
crates/aether-testkit/src/bin/failure_recovery_baseline.rs
View File

@@ -8,9 +8,7 @@ use aether_data::driver::postgres::{
PostgresPoolConfig, PostgresTransactionOptions,
};
use aether_data::{DataLayerError, PostgresBackend};
use aether_runtime_state::{
RedisClientConfig, RedisClientFactory, RedisLockRunner, RedisLockRunnerConfig,
};
use aether_runtime_state::{RedisClientConfig, RedisLockRunner, RedisLockRunnerConfig};
use aether_testkit::{
init_test_runtime_for, reserve_local_port, BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot,
ManagedPostgresServer, ManagedRedisServer, TunnelHarness, TunnelHarnessConfig,
@@ -247,20 +245,19 @@ async fn benchmark_redis_restart_recovery(
) -> Result<RecoverySummary, Box<dyn std::error::Error>> {
let mut runtime_sampler = BenchmarkRuntimeSampler::new();
let redis_url = redis_server.lock().await.redis_url().to_string();
let factory = RedisClientFactory::new(RedisClientConfig {
let redis_config = RedisClientConfig {
url: redis_url,
key_prefix: Some(format!("aether-failure-recovery-{}", std::process::id())),
})?;
let client = factory.connect_lazy()?;
let keyspace = factory.config().keyspace();
let runner = RedisLockRunner::new(
client,
keyspace.clone(),
};
let keyspace = redis_config.keyspace();
let runner = RedisLockRunner::from_config(
redis_config,
RedisLockRunnerConfig {
command_timeout_ms: Some(250),
default_ttl_ms: 1_000,
},
)?;
)
.await?;
let collector = Arc::new(RecoveryCollector::default());
let next_attempt = Arc::new(AtomicUsize::new(0));
let phase = Arc::new(AtomicUsize::new(0));

17
crates/aether-testkit/src/bin/redis_worker_baseline.rs
View File

@@ -2,7 +2,7 @@ use std::path::PathBuf;
use std::time::{Duration, Instant};
use aether_runtime_state::{
RedisClientConfig, RedisClientFactory, RedisConsumerGroup, RedisConsumerName, RedisStreamName,
RedisClientConfig, RedisConsumerGroup, RedisConsumerName, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamRunner, RedisStreamRunnerConfig,
};
use aether_testkit::{init_test_runtime_for, ManagedRedisServer};
@@ -86,25 +86,24 @@ async fn run_suite(
})
.expect("redis url should be resolved");
let factory = RedisClientFactory::new(RedisClientConfig {
let redis_config = RedisClientConfig {
url: redis_url.clone(),
key_prefix: Some(format!("aether-baseline-{}", std::process::id())),
})?;
let client = factory.connect_lazy()?;
let keyspace = factory.config().keyspace();
};
let keyspace = redis_config.keyspace();
let stream = keyspace.stream_name("worker-baseline");
let group = RedisConsumerGroup("worker-group".to_string());
let consumer_a = RedisConsumerName("consumer-a".to_string());
let consumer_b = RedisConsumerName("consumer-b".to_string());
let runner = RedisStreamRunner::new(
client,
keyspace,
let runner = RedisStreamRunner::from_config(
redis_config,
RedisStreamRunnerConfig {
command_timeout_ms: Some(2_000),
read_block_ms: Some(10),
read_count: 64,
},
)?;
)
.await?;
runner
.ensure_consumer_group(&stream, &group, "0-0")
.await

549
crates/aether-testkit/src/bin/runtime_redis_pressure.rs Normal file
View File

@@ -0,0 +1,549 @@
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use aether_runtime_state::{
DataLayerError, RedisClientConfig, RedisRuntimeDiagnostics, RuntimeQueueStore,
RuntimeSemaphoreConfig, RuntimeState,
};
use aether_testkit::{init_test_runtime_for, ManagedRedisServer};
use futures_util::stream::{self, StreamExt};
use serde::Serialize;
#[derive(Debug, Clone)]
struct RuntimeRedisPressureConfig {
kv_total: usize,
kv_concurrency: usize,
lock_total: usize,
lock_concurrency: usize,
semaphore_total: usize,
semaphore_concurrency: usize,
stream_total: usize,
stream_concurrency: usize,
blocking_probe_total: usize,
blocking_probe_concurrency: usize,
command_timeout_ms: u64,
output_path: Option<PathBuf>,
redis_url: Option<String>,
}
impl Default for RuntimeRedisPressureConfig {
fn default() -> Self {
Self {
kv_total: 20_000,
kv_concurrency: 200,
lock_total: 10_000,
lock_concurrency: 100,
semaphore_total: 5_000,
semaphore_concurrency: 100,
stream_total: 10_000,
stream_concurrency: 100,
blocking_probe_total: 1_000,
blocking_probe_concurrency: 100,
command_timeout_ms: 2_000,
output_path: None,
redis_url: None,
}
}
}
#[derive(Debug, Clone, Serialize)]
struct OperationSummary {
total_calls: usize,
total_items: usize,
failed_calls: usize,
p50_ms: u64,
p95_ms: u64,
p99_ms: u64,
max_ms: u64,
mean_ms: u64,
}
#[derive(Debug, Clone, Serialize)]
struct RuntimeRedisPressureReport {
suite: &'static str,
redis_url: String,
total_connections_before: Option<u64>,
total_connections_after: Option<u64>,
total_connections_delta: Option<i64>,
connected_clients_after: Option<u64>,
diagnostics_after: RedisRuntimeDiagnostics,
kv: OperationSummary,
lock: OperationSummary,
semaphore: OperationSummary,
stream_append: OperationSummary,
stream_read_group: OperationSummary,
stream_ack: OperationSummary,
blocking_fast_lane_probe: OperationSummary,
}
#[derive(Default)]
struct SummaryCollector {
latencies_ms: tokio::sync::Mutex<Vec<u64>>,
total_calls: AtomicUsize,
total_items: AtomicUsize,
failed_calls: AtomicUsize,
}
impl SummaryCollector {
async fn record(&self, elapsed: Duration, items: usize, failed: bool) {
self.latencies_ms
.lock()
.await
.push(elapsed.as_millis() as u64);
self.total_calls.fetch_add(1, Ordering::AcqRel);
self.total_items.fetch_add(items, Ordering::AcqRel);
if failed {
self.failed_calls.fetch_add(1, Ordering::AcqRel);
}
}
async fn summarize(&self) -> OperationSummary {
let mut latencies = self.latencies_ms.lock().await.clone();
latencies.sort_unstable();
let total_calls = self.total_calls.load(Ordering::Acquire);
OperationSummary {
total_calls,
total_items: self.total_items.load(Ordering::Acquire),
failed_calls: self.failed_calls.load(Ordering::Acquire),
p50_ms: percentile(&latencies, 50),
p95_ms: percentile(&latencies, 95),
p99_ms: percentile(&latencies, 99),
max_ms: latencies.last().copied().unwrap_or_default(),
mean_ms: if total_calls == 0 {
0
} else {
latencies.iter().sum::<u64>() / total_calls as u64
},
}
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
init_test_runtime_for("runtime-redis-pressure");
let config = parse_args(std::env::args().skip(1).collect())?;
let report = run_suite(&config).await?;
let raw = serde_json::to_string_pretty(&report)?;
println!("{raw}");
if let Some(path) = config.output_path.as_ref() {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, format!("{raw}\n"))?;
}
Ok(())
}
async fn run_suite(
config: &RuntimeRedisPressureConfig,
) -> Result<RuntimeRedisPressureReport, Box<dyn std::error::Error>> {
let managed_redis = if config.redis_url.is_none() {
Some(ManagedRedisServer::start().await?)
} else {
None
};
let redis_url = config
.redis_url
.clone()
.or_else(|| {
managed_redis
.as_ref()
.map(|server| server.redis_url().to_string())
})
.expect("redis url should resolve");
let runtime = Arc::new(
RuntimeState::redis(
RedisClientConfig {
url: redis_url.clone(),
key_prefix: Some(format!("aether-runtime-pressure-{}", std::process::id())),
},
Some(config.command_timeout_ms),
)
.await?,
);
let before = runtime
.redis_diagnostics()
.await?
.expect("redis diagnostics should be available");
let kv = benchmark_kv(runtime.clone(), config).await;
let lock = benchmark_lock(runtime.clone(), config).await;
let semaphore = benchmark_semaphore(runtime.clone(), config).await?;
let (stream_append, stream_read_group, stream_ack) =
benchmark_stream(runtime.clone(), config).await;
let blocking_fast_lane_probe = benchmark_blocking_fast_lane(runtime.clone(), config).await?;
tokio::time::sleep(Duration::from_millis(200)).await;
let after = runtime
.redis_diagnostics()
.await?
.expect("redis diagnostics should be available");
let total_connections_delta = match (
before.total_connections_received,
after.total_connections_received,
) {
(Some(before), Some(after)) => Some(after as i64 - before as i64),
_ => None,
};
Ok(RuntimeRedisPressureReport {
suite: "runtime_redis_pressure",
redis_url,
total_connections_before: before.total_connections_received,
total_connections_after: after.total_connections_received,
total_connections_delta,
connected_clients_after: after.connected_clients,
diagnostics_after: after,
kv,
lock,
semaphore,
stream_append,
stream_read_group,
stream_ack,
blocking_fast_lane_probe,
})
}
async fn benchmark_kv(
runtime: Arc<RuntimeState>,
config: &RuntimeRedisPressureConfig,
) -> OperationSummary {
let collector = Arc::new(SummaryCollector::default());
stream::iter(0..config.kv_total)
.for_each_concurrent(config.kv_concurrency, |index| {
let runtime = runtime.clone();
let collector = collector.clone();
async move {
let started = Instant::now();
let value = format!("value-{index}");
let key = format!("pressure:kv:{index}");
let result = async {
runtime
.kv_set(&key, value.clone(), Some(Duration::from_secs(60)))
.await?;
let actual = runtime.kv_get(&key).await?;
if actual.as_deref() != Some(value.as_str()) {
return Err(DataLayerError::UnexpectedValue(format!(
"kv pressure mismatch for {key}"
)));
}
runtime.kv_delete(&key).await?;
Ok::<usize, DataLayerError>(3)
}
.await;
let failed = result.is_err();
collector
.record(started.elapsed(), result.unwrap_or(0), failed)
.await;
}
})
.await;
collector.summarize().await
}
async fn benchmark_lock(
runtime: Arc<RuntimeState>,
config: &RuntimeRedisPressureConfig,
) -> OperationSummary {
let collector = Arc::new(SummaryCollector::default());
stream::iter(0..config.lock_total)
.for_each_concurrent(config.lock_concurrency, |index| {
let runtime = runtime.clone();
let collector = collector.clone();
async move {
let started = Instant::now();
let key = format!("pressure:lock:{index}");
let owner = format!("owner-{index}");
let result = async {
let Some(lease) = runtime
.lock_try_acquire(&key, &owner, Duration::from_secs(30))
.await?
else {
return Err(DataLayerError::UnexpectedValue(format!(
"lock pressure acquire returned none for {key}"
)));
};
if !runtime.lock_renew(&lease, Duration::from_secs(30)).await? {
return Err(DataLayerError::UnexpectedValue(format!(
"lock pressure renew returned false for {key}"
)));
}
if !runtime.lock_release(&lease).await? {
return Err(DataLayerError::UnexpectedValue(format!(
"lock pressure release returned false for {key}"
)));
}
Ok::<usize, DataLayerError>(3)
}
.await;
let failed = result.is_err();
collector
.record(started.elapsed(), result.unwrap_or(0), failed)
.await;
}
})
.await;
collector.summarize().await
}
async fn benchmark_semaphore(
runtime: Arc<RuntimeState>,
config: &RuntimeRedisPressureConfig,
) -> Result<OperationSummary, Box<dyn std::error::Error>> {
let collector = Arc::new(SummaryCollector::default());
let semaphore = Arc::new(runtime.semaphore(
"redis_pressure",
config.semaphore_concurrency.saturating_mul(4).max(1),
RuntimeSemaphoreConfig {
lease_ttl_ms: 10_000,
renew_interval_ms: 5_000,
command_timeout_ms: Some(config.command_timeout_ms),
},
)?);
stream::iter(0..config.semaphore_total)
.for_each_concurrent(config.semaphore_concurrency, |_| {
let semaphore = semaphore.clone();
let collector = collector.clone();
async move {
let started = Instant::now();
let result = semaphore.try_acquire().await;
let failed = result.is_err();
drop(result);
collector
.record(started.elapsed(), usize::from(!failed), failed)
.await;
}
})
.await;
tokio::time::sleep(Duration::from_millis(200)).await;
Ok(collector.summarize().await)
}
async fn benchmark_stream(
runtime: Arc<RuntimeState>,
config: &RuntimeRedisPressureConfig,
) -> (OperationSummary, OperationSummary, OperationSummary) {
let stream_name = "pressure-stream";
let group = "pressure-workers";
let consumer = "consumer-a";
RuntimeQueueStore::ensure_consumer_group(runtime.as_ref(), stream_name, group, "0-0")
.await
.expect("stream consumer group should initialize");
let append = Arc::new(SummaryCollector::default());
stream::iter(0..config.stream_total)
.for_each_concurrent(config.stream_concurrency, |index| {
let runtime = runtime.clone();
let append = append.clone();
async move {
let started = Instant::now();
let mut fields = BTreeMap::new();
fields.insert("payload".to_string(), format!("stream-value-{index}"));
let result = RuntimeQueueStore::append_fields_with_maxlen(
runtime.as_ref(),
stream_name,
&fields,
Some(config.stream_total.saturating_mul(2)),
)
.await;
append
.record(
started.elapsed(),
usize::from(result.is_ok()),
result.is_err(),
)
.await;
}
})
.await;
let read = SummaryCollector::default();
let mut ids = Vec::with_capacity(config.stream_total);
while ids.len() < config.stream_total {
let started = Instant::now();
match RuntimeQueueStore::read_group(
runtime.as_ref(),
stream_name,
group,
consumer,
128,
Some(10),
)
.await
{
Ok(entries) => {
let item_count = entries.len();
ids.extend(entries.into_iter().map(|entry| entry.id));
read.record(started.elapsed(), item_count, false).await;
}
Err(_) => {
read.record(started.elapsed(), 0, true).await;
}
}
}
let ack = SummaryCollector::default();
for chunk in ids.chunks(128) {
let started = Instant::now();
match RuntimeQueueStore::ack(runtime.as_ref(), stream_name, group, chunk).await {
Ok(count) => ack.record(started.elapsed(), count, false).await,
Err(_) => ack.record(started.elapsed(), 0, true).await,
}
}
(
append.summarize().await,
read.summarize().await,
ack.summarize().await,
)
}
async fn benchmark_blocking_fast_lane(
runtime: Arc<RuntimeState>,
config: &RuntimeRedisPressureConfig,
) -> Result<OperationSummary, Box<dyn std::error::Error>> {
let stream_name = "pressure-blocking-empty";
let group = "pressure-blocking-workers";
RuntimeQueueStore::ensure_consumer_group(runtime.as_ref(), stream_name, group, "0-0").await?;
let blocking_runtime = runtime.clone();
let blocking = tokio::spawn(async move {
RuntimeQueueStore::read_group(
blocking_runtime.as_ref(),
stream_name,
group,
"blocked-consumer",
1,
Some(1_000),
)
.await
});
tokio::time::sleep(Duration::from_millis(50)).await;
let summary = benchmark_fast_lane_probe(runtime, config).await;
let _ = blocking.await?;
Ok(summary)
}
async fn benchmark_fast_lane_probe(
runtime: Arc<RuntimeState>,
config: &RuntimeRedisPressureConfig,
) -> OperationSummary {
let collector = Arc::new(SummaryCollector::default());
stream::iter(0..config.blocking_probe_total)
.for_each_concurrent(config.blocking_probe_concurrency, |index| {
let runtime = runtime.clone();
let collector = collector.clone();
async move {
let started = Instant::now();
let key = format!("pressure:blocking-probe:{index}");
let result = async {
runtime
.kv_set(&key, "ok", Some(Duration::from_secs(30)))
.await?;
let ok = runtime.kv_get(&key).await?.as_deref() == Some("ok");
if !ok {
return Err(DataLayerError::UnexpectedValue(format!(
"blocking fast lane probe mismatch for {key}"
)));
}
Ok::<usize, DataLayerError>(2)
}
.await;
let failed = result.is_err();
collector
.record(started.elapsed(), result.unwrap_or(0), failed)
.await;
}
})
.await;
collector.summarize().await
}
fn percentile(latencies: &[u64], percentile: u8) -> u64 {
if latencies.is_empty() {
return 0;
}
let last_index = latencies.len() - 1;
let rank = ((last_index as f64) * (percentile as f64 / 100.0)).round() as usize;
latencies[rank.min(last_index)]
}
fn parse_args(args: Vec<String>) -> Result<RuntimeRedisPressureConfig, Box<dyn std::error::Error>> {
let mut config = RuntimeRedisPressureConfig::default();
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--kv-total" => config.kv_total = next_value(&mut iter, "--kv-total")?.parse()?,
"--kv-concurrency" => {
config.kv_concurrency = next_value(&mut iter, "--kv-concurrency")?.parse()?
}
"--lock-total" => config.lock_total = next_value(&mut iter, "--lock-total")?.parse()?,
"--lock-concurrency" => {
config.lock_concurrency = next_value(&mut iter, "--lock-concurrency")?.parse()?
}
"--semaphore-total" => {
config.semaphore_total = next_value(&mut iter, "--semaphore-total")?.parse()?
}
"--semaphore-concurrency" => {
config.semaphore_concurrency =
next_value(&mut iter, "--semaphore-concurrency")?.parse()?
}
"--stream-total" => {
config.stream_total = next_value(&mut iter, "--stream-total")?.parse()?
}
"--stream-concurrency" => {
config.stream_concurrency =
next_value(&mut iter, "--stream-concurrency")?.parse()?
}
"--blocking-probe-total" => {
config.blocking_probe_total =
next_value(&mut iter, "--blocking-probe-total")?.parse()?
}
"--blocking-probe-concurrency" => {
config.blocking_probe_concurrency =
next_value(&mut iter, "--blocking-probe-concurrency")?.parse()?
}
"--command-timeout-ms" => {
config.command_timeout_ms =
next_value(&mut iter, "--command-timeout-ms")?.parse()?
}
"--redis-url" => config.redis_url = Some(next_value(&mut iter, "--redis-url")?),
"--output" => {
config.output_path = Some(PathBuf::from(next_value(&mut iter, "--output")?))
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
}
other => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("unknown argument: {other}"),
)
.into());
}
}
}
Ok(config)
}
fn next_value(
iter: &mut impl Iterator<Item = String>,
flag: &str,
) -> Result<String, Box<dyn std::error::Error>> {
iter.next().ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("missing value for {flag}"),
)
.into()
})
}
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin runtime_redis_pressure -- [--kv-total 20000] [--kv-concurrency 200] [--lock-total 10000] [--lock-concurrency 100] [--semaphore-total 5000] [--semaphore-concurrency 100] [--stream-total 10000] [--stream-concurrency 100] [--blocking-probe-total 1000] [--blocking-probe-concurrency 100] [--command-timeout-ms 2000] [--redis-url redis://127.0.0.1:6379/0] [--output /tmp/runtime_redis_pressure.json]"
);
}

2
crates/aether-usage-runtime/src/config.rs
View File

@@ -3,6 +3,7 @@ use aether_data_contracts::DataLayerError;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UsageRuntimeConfig {
pub enabled: bool,
pub queue_terminal_events: bool,
pub stream_key: String,
pub consumer_group: String,
pub dlq_stream_key: String,
@@ -18,6 +19,7 @@ impl Default for UsageRuntimeConfig {
fn default() -> Self {
Self {
enabled: false,
queue_terminal_events: false,
stream_key: "usage:events".to_string(),
consumer_group: "usage_consumers".to_string(),
dlq_stream_key: "usage:events:dlq".to_string(),

39
crates/aether-usage-runtime/src/runtime.rs
View File

@@ -104,7 +104,10 @@ impl UsageRuntime {
where
T: UsageRuntimeAccess,
{
self.is_enabled() && data.has_usage_writer() && data.has_usage_worker_queue()
self.is_enabled()
&& self.config.queue_terminal_events
&& data.has_usage_writer()
&& data.has_usage_worker_queue()
}
pub fn spawn_worker<T>(&self, data: Arc<T>) -> Option<tokio::task::JoinHandle<()>>
@@ -352,30 +355,32 @@ impl UsageRuntime {
where
T: UsageRuntimeAccess,
{
if let Some(runner) = data.usage_worker_queue() {
match UsageQueue::new(runner, self.config.clone()) {
Ok(queue) => match queue.enqueue(&event).await {
Ok(_) => return,
if self.config.queue_terminal_events {
if let Some(runner) = data.usage_worker_queue() {
match UsageQueue::new(runner, self.config.clone()) {
Ok(queue) => match queue.enqueue(&event).await {
Ok(_) => return,
Err(err) => {
warn!(
event_name = "usage_terminal_enqueue_failed",
log_type = "event",
request_id = %event.request_id,
fallback = "direct_write",
error = %err,
"usage runtime failed to enqueue terminal usage event; falling back to direct write"
)
}
},
Err(err) => {
warn!(
event_name = "usage_terminal_enqueue_failed",
event_name = "usage_terminal_queue_init_failed",
log_type = "event",
request_id = %event.request_id,
fallback = "direct_write",
error = %err,
"usage runtime failed to enqueue terminal usage event; falling back to direct write"
"usage runtime failed to build queue; falling back to direct write"
)
}
},
Err(err) => {
warn!(
event_name = "usage_terminal_queue_init_failed",
log_type = "event",
request_id = %event.request_id,
fallback = "direct_write",
error = %err,
"usage runtime failed to build queue; falling back to direct write"
)
}
}
}

2
docker-compose.yml
View File

@@ -36,7 +36,7 @@ services:
redis:
image: redis:7-alpine
container_name: aether-redis
command: redis-server --appendonly yes --appendfsync everysec --save 60 1000 --requirepass ${REDIS_PASSWORD}
command: redis-server --appendonly yes --appendfsync everysec --save 60 1000 --requirepass ${REDIS_PASSWORD} --maxclients ${REDIS_MAXCLIENTS:-10000}
volumes:
- redis_data:/data
ports:

67
docs/operations/redis-runtime-runbook.md Normal file
View File

@@ -0,0 +1,67 @@
# Runtime Redis Operations Runbook
This runbook covers Aether runtime Redis connection pressure incidents. It is
not a substitute for fixing application-level connection churn.
## Normal Expectations
- Each `RuntimeState` Redis backend initializes a fixed set of long-lived
connection lanes: fast, stream, blocking stream, and admin.
- `connected_clients` should stay near a small fixed number per app instance,
plus health checks and ad hoc admin clients.
- `total_connections_received` should not grow linearly with request volume.
- Large TIME_WAIT spikes between app and Redis indicate a regression or a
separate process repeatedly opening Redis connections.
## Emergency Mitigation
1. Disable the retry source first, such as expired Codex/OAuth keys causing a
retry storm.
2. Restart the app to stop continued connection creation:
```sh
docker compose restart app
```
3. On a Linux host, temporarily widen the ephemeral port range and enable safe
TIME_WAIT reuse:
```sh
sudo sysctl -w net.ipv4.ip_local_port_range="10000 65535"
sudo sysctl -w net.ipv4.tcp_tw_reuse=1
```
4. Do not enable `tcp_tw_recycle`; it is obsolete and unsafe with NAT.
Docker Desktop on macOS runs containers inside a Linux VM. Host-level macOS
`sysctl` changes do not necessarily affect the VM network namespace.
## Checks
Use Redis `INFO clients` and `INFO stats` to inspect:
- `connected_clients`
- `total_connections_received`
Use OS socket tooling on the Redis host or container namespace to inspect
TIME_WAIT counts. Persistent growth after the runtime Redis refactor means a
different code path or process is still opening short-lived Redis connections.
## File Descriptor Limits
Aether's compose files intentionally do not set container `ulimits.nofile`.
Redis connection churn must be fixed in application code, not hidden by larger
file descriptor limits.
For high-concurrency production hosts, set file descriptor policy at the
runtime or service-manager layer instead:
- Docker daemon default ulimit, for example `default-ulimits` in
`/etc/docker/daemon.json`.
- systemd service limits such as `LimitNOFILE=` for Docker or the process
supervisor.
- Managed container platform resource settings, when Docker daemon settings are
not available.
Keep Redis `maxclients` below the effective Redis process `nofile` limit with
room for persistence files, replicas, and admin connections.