fix: harden concurrency limits and high-RPM runtime paths

Bound request, stream, queue, and shutdown resource lifetimes. Reduce scheduler and Redis hot-path work and isolate database maintenance. Include regression coverage, load probes, and concurrency audit results.
This commit is contained in:
elky
2026-09-10 08:14:58 +08:00
parent 361952ada9
commit ecc16673eb
149 changed files with 27963 additions and 1926 deletions
+598 -20
View File
@@ -1,6 +1,7 @@
mod error;
mod memory;
pub mod redis;
mod score_window;
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
@@ -17,6 +18,7 @@ use async_trait::async_trait;
pub use error::DataLayerError;
use memory::MemoryRuntimeBackend;
pub use memory::MemoryRuntimeStateConfig;
pub use score_window::{ScoreWindowU64Stats, SCORE_WINDOW_AGGREGATION_MEMBER_LIMIT};
use tokio::task::JoinHandle;
use tracing::warn;
use uuid::Uuid;
@@ -621,6 +623,32 @@ impl RuntimeState {
}
}
/// Aggregate at most 512 timestamped `prefix:u64` members per key without
/// transferring their history. `None` requires an exact full-range fallback;
/// it never represents an empty or cached window.
pub async fn score_window_u64_stats_by_min(
&self,
keys: &[String],
min_score: f64,
) -> Result<Vec<Option<ScoreWindowU64Stats>>, DataLayerError> {
if !min_score.is_finite() {
return Err(DataLayerError::InvalidInput(
"runtime window minimum score must be finite".to_string(),
));
}
match self.backend.as_ref() {
RuntimeStateBackend::Memory(memory) => {
Ok(memory.score_window_u64_stats_by_min(keys, min_score).await)
}
RuntimeStateBackend::Redis(redis) => {
redis
.runtime
.score_window_u64_stats_by_min(keys, min_score)
.await
}
}
}
pub async fn score_remove_by_score(
&self,
key: &str,
@@ -1046,6 +1074,26 @@ pub struct RuntimeQueueEntry {
pub fields: BTreeMap<String, String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RuntimeQueueReclaimPage {
/// Resume the next reclaim scan here; `0-0` marks the end of the current scan.
pub next_start_id: String,
pub entries: Vec<RuntimeQueueEntry>,
pub deleted_ids: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RuntimeQueueTransferOutcome {
Transferred {
destination_id: String,
acked: usize,
deleted: usize,
},
/// No pending entry was present. This does not assert that it was archived:
/// another consumer, deletion, or retention policy may have removed it.
NotPending,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct RuntimeQueueStats {
pub stream_length: u64,
@@ -1085,6 +1133,45 @@ fn validate_runtime_queue_reclaim_config(
Ok(())
}
pub(crate) fn validate_runtime_queue_transfer(
source: &str,
group: &str,
entry_id: &str,
destination: &str,
destination_fields: &BTreeMap<String, String>,
) -> Result<(), DataLayerError> {
validate_runtime_queue_name(source, "runtime queue source stream")?;
validate_runtime_queue_name(group, "runtime queue group")?;
validate_runtime_queue_name(destination, "runtime queue destination stream")?;
if source == destination {
return Err(DataLayerError::InvalidInput(
"runtime queue transfer source and destination must differ".to_string(),
));
}
if destination_fields.is_empty() {
return Err(DataLayerError::InvalidInput(
"runtime queue transfer destination fields cannot be empty".to_string(),
));
}
let canonical_u64 = |value: &str| {
!value.is_empty()
&& value.bytes().all(|byte| byte.is_ascii_digit())
&& (value.len() == 1 || !value.starts_with('0'))
&& value.parse::<u64>().is_ok()
};
if !entry_id
.split_once('-')
.is_some_and(|(milliseconds, sequence)| {
canonical_u64(milliseconds) && canonical_u64(sequence)
})
{
return Err(DataLayerError::InvalidInput(
"runtime queue transfer entry id must be a canonical u64-u64 stream id".to_string(),
));
}
Ok(())
}
#[async_trait]
pub trait RuntimeQueueStore: Send + Sync {
async fn ensure_consumer_group(
@@ -1119,6 +1206,39 @@ pub trait RuntimeQueueStore: Send + Sync {
config: RuntimeQueueReclaimConfig,
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError>;
/// Existing queue backends can retain their complete-scan behavior without implementing paging.
async fn claim_stale_page(
&self,
stream: &str,
group: &str,
consumer: &str,
start_id: &str,
config: RuntimeQueueReclaimConfig,
) -> Result<RuntimeQueueReclaimPage, DataLayerError> {
Ok(RuntimeQueueReclaimPage {
next_start_id: "0-0".to_string(),
entries: self
.claim_stale(stream, group, consumer, start_id, config)
.await?,
deleted_ids: Vec::new(),
})
}
/// Atomically append caller-supplied fields, acknowledge the pending source entry, and
/// delete that source ID. Repeated calls must not append when the entry is no longer pending.
/// `None` means unsupported and has no side effects; callers may explicitly retain their
/// existing non-atomic fallback for third-party queue implementations.
async fn try_transfer_pending_to_stream(
&self,
_source: &str,
_group: &str,
_entry_id: &str,
_destination: &str,
_destination_fields: &BTreeMap<String, String>,
) -> Result<Option<RuntimeQueueTransferOutcome>, DataLayerError> {
Ok(None)
}
async fn ack(&self, stream: &str, group: &str, ids: &[String])
-> Result<usize, DataLayerError>;
@@ -1242,6 +1362,20 @@ impl RuntimeQueueStore for RuntimeState {
start_id: &str,
config: RuntimeQueueReclaimConfig,
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
Ok(self
.claim_stale_page(stream, group, consumer, start_id, config)
.await?
.entries)
}
async fn claim_stale_page(
&self,
stream: &str,
group: &str,
consumer: &str,
start_id: &str,
config: RuntimeQueueReclaimConfig,
) -> Result<RuntimeQueueReclaimPage, DataLayerError> {
validate_runtime_queue_name(stream, "runtime queue stream")?;
validate_runtime_queue_name(group, "runtime queue group")?;
validate_runtime_queue_name(consumer, "runtime queue consumer")?;
@@ -1250,32 +1384,76 @@ impl RuntimeQueueStore for RuntimeState {
match self.backend.as_ref() {
RuntimeStateBackend::Memory(memory) => {
memory
.queue_claim_stale(stream, group, consumer, start_id, config)
.queue_claim_stale_page(stream, group, consumer, start_id, config)
.await
}
RuntimeStateBackend::Redis(redis) => Ok(redis
.stream
.claim_stale(
&RedisStreamName(stream.to_string()),
&RedisConsumerGroup(group.to_string()),
&RedisConsumerName(consumer.to_string()),
start_id,
RedisStreamReclaimConfig {
min_idle_ms: config.min_idle_ms,
count: config.count,
},
)
.await?
.entries
.into_iter()
.map(|entry| RuntimeQueueEntry {
id: entry.id,
fields: entry.fields,
RuntimeStateBackend::Redis(redis) => {
let page = redis
.stream
.claim_stale(
&RedisStreamName(stream.to_string()),
&RedisConsumerGroup(group.to_string()),
&RedisConsumerName(consumer.to_string()),
start_id,
RedisStreamReclaimConfig {
min_idle_ms: config.min_idle_ms,
count: config.count,
},
)
.await?;
Ok(RuntimeQueueReclaimPage {
next_start_id: page.next_start_id,
entries: page
.entries
.into_iter()
.map(|entry| RuntimeQueueEntry {
id: entry.id,
fields: entry.fields,
})
.collect(),
deleted_ids: page.deleted_ids,
})
.collect()),
}
}
}
async fn try_transfer_pending_to_stream(
&self,
source: &str,
group: &str,
entry_id: &str,
destination: &str,
destination_fields: &BTreeMap<String, String>,
) -> Result<Option<RuntimeQueueTransferOutcome>, DataLayerError> {
validate_runtime_queue_transfer(source, group, entry_id, destination, destination_fields)?;
let outcome = match self.backend.as_ref() {
RuntimeStateBackend::Memory(memory) => {
memory
.queue_transfer_pending_to_stream(
source,
group,
entry_id,
destination,
destination_fields,
)
.await?
}
RuntimeStateBackend::Redis(redis) => {
redis
.stream
.try_transfer_pending_to_stream(
source,
group,
entry_id,
destination,
destination_fields,
)
.await?
}
};
Ok(Some(outcome))
}
async fn ack(
&self,
stream: &str,
@@ -1817,6 +1995,18 @@ mod tests {
use std::process::{Child, Command, Stdio};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
mod stream_receive {
include!("redis/stream_receive_tests.rs");
}
mod dead_letter_transfer {
include!("redis/dead_letter_transfer_tests.rs");
}
mod usage_limit_cleanup {
include!("redis/usage_limit_cleanup_tests.rs");
}
#[tokio::test]
async fn memory_kv_expires_entries() {
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
@@ -2722,6 +2912,131 @@ mod tests {
}
}
#[tokio::test]
async fn redis_large_stream_batches_preserve_fields_across_read_reclaim_and_ack() {
let Some(redis) = TestRedisServer::start().await else {
return;
};
for protocol in ["resp2", "resp3"] {
let runtime = RuntimeState::redis(
RedisClientConfig {
url: format!("{}?protocol={protocol}", redis.redis_url),
key_prefix: Some(format!("large-batch-{protocol}")),
},
Some(5_000),
)
.await
.expect("large batch runtime should connect");
let stream = "usage:large-batch";
let group = "workers";
RuntimeQueueStore::ensure_consumer_group(&runtime, stream, group, "0-0")
.await
.unwrap();
let payload = format!(
"{}\r\n\"escaped\"\\\u{4e2d}\u{6587}",
"x".repeat(512 * 1024)
);
let mut expected = BTreeMap::new();
for sequence in 0..24 {
let fields = BTreeMap::from([
("payload".to_string(), payload.clone()),
("sequence".to_string(), sequence.to_string()),
("legacy_marker".to_string(), "preserve exactly".to_string()),
]);
let id =
RuntimeQueueStore::append_fields_with_maxlen(&runtime, stream, &fields, None)
.await
.unwrap();
expected.insert(id, sequence.to_string());
}
let mut readers = tokio::task::JoinSet::new();
for index in 0..3 {
let runtime = runtime.clone();
readers.spawn(async move {
RuntimeQueueStore::read_group(
&runtime,
stream,
group,
&format!("reader-{index}"),
8,
Some(1),
)
.await
.unwrap()
});
}
let mut delivered = std::collections::BTreeSet::new();
while let Some(entries) = readers.join_next().await {
let entries = entries.unwrap();
assert_eq!(entries.len(), 8);
for entry in entries {
assert_eq!(entry.fields.len(), 3);
assert_eq!(entry.fields["payload"].as_bytes(), payload.as_bytes());
assert_eq!(entry.fields["sequence"], expected[&entry.id]);
assert_eq!(entry.fields["legacy_marker"], "preserve exactly");
assert!(delivered.insert(entry.id));
}
}
assert_eq!(delivered.len(), 24);
let stats = RuntimeQueueStore::stats(&runtime, stream, Some(group))
.await
.unwrap();
assert_eq!(stats.group_pending, 24);
assert_eq!(stats.group_lag, Some(0));
tokio::time::sleep(Duration::from_millis(20)).await;
let mut reclaimed = std::collections::BTreeSet::new();
while reclaimed.len() < 24 {
let entries = RuntimeQueueStore::claim_stale(
&runtime,
stream,
group,
"retry-consumer",
"0-0",
RuntimeQueueReclaimConfig {
min_idle_ms: 1,
count: 5,
},
)
.await
.unwrap();
assert!(!entries.is_empty());
assert!(entries.len() <= 5);
let mut ids = Vec::new();
for entry in entries {
assert_eq!(entry.fields.len(), 3);
assert_eq!(entry.fields["payload"].as_bytes(), payload.as_bytes());
assert_eq!(entry.fields["sequence"], expected[&entry.id]);
assert_eq!(entry.fields["legacy_marker"], "preserve exactly");
assert!(reclaimed.insert(entry.id.clone()));
ids.push(entry.id);
}
assert_eq!(
RuntimeQueueStore::ack(&runtime, stream, group, &ids)
.await
.unwrap(),
ids.len()
);
assert_eq!(
RuntimeQueueStore::delete(&runtime, stream, &ids)
.await
.unwrap(),
ids.len()
);
}
assert_eq!(reclaimed, delivered);
let stats = RuntimeQueueStore::stats(&runtime, stream, Some(group))
.await
.unwrap();
assert_eq!(stats.stream_length, 0);
assert_eq!(stats.group_pending, 0);
assert_eq!(stats.group_lag, Some(0));
eprintln!(
"verified {protocol}: 24 large records, 3 readers, read/reclaim/ack complete"
);
}
}
#[tokio::test]
async fn redis_connection_manager_recovers_after_restart() {
let Some(mut redis) = TestRedisServer::start().await else {
@@ -2776,6 +3091,269 @@ mod tests {
assert_kv_score_and_queue_contract(&redis_runtime).await;
}
#[tokio::test]
async fn runtime_backends_share_bounded_score_window_aggregation() {
let memory = RuntimeState::memory(MemoryRuntimeStateConfig::default());
assert_bounded_score_window_aggregation(&memory).await;
let Some((_server, runtime)) = redis_runtime_for_test("score-window").await else {
return;
};
assert_bounded_score_window_aggregation(&runtime).await;
}
async fn assert_bounded_score_window_aggregation(runtime: &RuntimeState) {
let keys = (0..35)
.map(|index| format!("window:{index}"))
.collect::<Vec<_>>();
for (member, score) in [
("expired:999", 99.999),
("boundary:7", 100.0),
("recent:9007199254740993", 101.0),
("nested:prefix:+00012", 102.0),
("zero:0", 103.0),
("invalid:1.5", 104.0),
("invalid:-1", 105.0),
("invalid:18446744073709551616", 106.0),
("invalid: 12", 107.0),
("missing-separator", 108.0),
] {
runtime
.score_set(&keys[0], member, score)
.await
.expect("seed values");
}
runtime
.score_set(&keys[1], "max:18446744073709551615", 100.0)
.await
.expect("seed max");
runtime
.score_set(&keys[2], "max:18446744073709551615", 100.0)
.await
.expect("seed overflow");
runtime
.score_set(&keys[2], "additional:2", 100.0)
.await
.expect("seed overflow addition");
for index in 0..SCORE_WINDOW_AGGREGATION_MEMBER_LIMIT {
runtime
.score_set(&keys[3], &format!("{index}:3"), 100.0)
.await
.expect("seed bounded window");
}
runtime
.score_set(&keys[3], "expired:9999", 0.0)
.await
.expect("seed expired sample");
let stats = runtime
.score_window_u64_stats_by_min(&keys, 100.0)
.await
.expect("aggregate");
assert_eq!(
stats.len(),
keys.len(),
"pipeline batches preserve key order"
);
assert_eq!(
stats[0],
Some(ScoreWindowU64Stats {
sum: 9_007_199_254_741_012,
positive_count: 3
})
);
assert_eq!(
stats[1],
Some(ScoreWindowU64Stats {
sum: u64::MAX,
positive_count: 1
})
);
assert_eq!(
stats[2],
Some(ScoreWindowU64Stats {
sum: u64::MAX,
positive_count: 2
})
);
assert_eq!(
stats[3],
Some(ScoreWindowU64Stats {
sum: 1536,
positive_count: 512
})
);
assert!(stats[4..]
.iter()
.all(|stats| *stats == Some(ScoreWindowU64Stats::default())));
runtime
.score_set(&keys[3], "overflowing-window:11", 101.0)
.await
.expect("exceed server limit");
let stats = runtime
.score_window_u64_stats_by_min(&keys[3..4], 100.0)
.await
.expect("bounded fallback");
assert_eq!(
stats,
vec![None],
"oversized windows require the full exact read"
);
let members = runtime
.score_range_by_min(&keys[3], 100.0)
.await
.expect("full window");
assert_eq!(
ScoreWindowU64Stats::from_members(members.iter().map(String::as_str)).sum,
1547
);
runtime
.score_remove(&keys[3], "overflowing-window:11")
.await
.expect("remove newest");
assert_eq!(
runtime
.score_window_u64_stats_by_min(&keys[3..4], 100.0)
.await
.expect("read after remove")[0]
.unwrap()
.sum,
1536
);
runtime
.score_set(&keys[3], "0:3", 99.0)
.await
.expect("move sample outside window");
assert_eq!(
runtime
.score_window_u64_stats_by_min(&keys[3..4], 100.0)
.await
.expect("read changed score")[0]
.unwrap()
.sum,
1533
);
assert!(runtime
.score_window_u64_stats_by_min(&[], 100.0)
.await
.expect("empty query")
.is_empty());
assert!(runtime
.score_window_u64_stats_by_min(&keys, f64::NAN)
.await
.is_err());
}
#[tokio::test]
async fn redis_score_window_aggregation_reloads_scripts_without_caching_old_cost() {
let Some((server, runtime)) = redis_runtime_for_test("score-window-reload").await else {
return;
};
let keys = vec!["reload:cost".to_string()];
runtime
.score_set(&keys[0], "first:7", 100.0)
.await
.expect("first cost");
assert_eq!(
runtime
.score_window_u64_stats_by_min(&keys, 100.0)
.await
.expect("first aggregate")[0]
.unwrap()
.sum,
7
);
let client = ::redis::Client::open(server.redis_url.as_str()).expect("test Redis client");
let mut connection = client
.get_multiplexed_async_connection()
.await
.expect("test connection");
::redis::cmd("SCRIPT")
.arg("FLUSH")
.query_async::<()>(&mut connection)
.await
.expect("flush scripts");
runtime
.score_set(&keys[0], "second:11", 101.0)
.await
.expect("new cost");
assert_eq!(
runtime
.score_window_u64_stats_by_min(&keys, 100.0)
.await
.expect("reload aggregate")[0]
.unwrap()
.sum,
18
);
assert_eq!(
runtime
.score_window_u64_stats_by_min(&keys, 101.0)
.await
.expect("changed window")[0]
.unwrap()
.sum,
11
);
runtime
.key_expire(&keys[0], Duration::ZERO)
.await
.expect("expire window");
assert_eq!(
runtime
.score_window_u64_stats_by_min(&keys, 100.0)
.await
.expect("expired aggregate"),
vec![Some(ScoreWindowU64Stats::default())]
);
}
#[tokio::test]
async fn redis_score_window_aggregation_observes_completed_concurrent_writes() {
let Some((_server, runtime)) = redis_runtime_for_test("score-window-concurrent").await
else {
return;
};
let writer_runtime = runtime.clone();
let (written_tx, mut written_rx) = tokio::sync::mpsc::channel(8);
let writer = tokio::spawn(async move {
for index in 1..=128_u64 {
writer_runtime
.score_set("concurrent:cost", &format!("{index}:2"), 100.0)
.await
.expect("concurrent write");
written_tx.send(index).await.expect("notify reader");
}
});
let keys = vec!["concurrent:cost".to_string()];
while let Some(written) = written_rx.recv().await {
let stats = runtime
.score_window_u64_stats_by_min(&keys, 100.0)
.await
.expect("concurrent aggregate")[0]
.unwrap();
assert!(
stats.positive_count >= written,
"completed writes must not be hidden by a stale aggregate"
);
assert_eq!(
stats.sum,
stats.positive_count * 2,
"one script observes one consistent window"
);
}
writer.await.expect("writer task");
assert_eq!(
runtime
.score_window_u64_stats_by_min(&keys, 100.0)
.await
.expect("final aggregate")[0]
.unwrap()
.sum,
256
);
}
#[tokio::test]
async fn runtime_backends_reject_invalid_shared_inputs() {
let memory = RuntimeState::memory(MemoryRuntimeStateConfig::default());
File diff suppressed because it is too large Load Diff
+255 -21
View File
@@ -1,9 +1,12 @@
use crate::error::RedisResultExt;
use crate::redis::RedisKeyspace;
use crate::DataLayerError;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use tracing::info;
pub(crate) type RedisClient = redis::Client;
@@ -141,8 +144,8 @@ pub(crate) struct RedisConnectionRouter {
fast: RedisManagedConnection,
stream: Arc<Vec<RedisManagedConnection>>,
stream_next: Arc<AtomicUsize>,
blocking_stream: Arc<Vec<RedisManagedConnection>>,
blocking_stream_next: Arc<AtomicUsize>,
blocking_stream: Arc<RedisBlockingStreamPool>,
usage_cleanup: Arc<RedisBlockingStreamPool>,
admin: RedisManagedConnection,
metrics: Arc<RedisConnectionMetrics>,
}
@@ -152,7 +155,7 @@ impl std::fmt::Debug for RedisConnectionRouter {
f.debug_struct("RedisConnectionRouter")
.field("lanes", &["fast", "stream", "blocking_stream", "admin"])
.field("stream_lanes", &self.stream.len())
.field("blocking_stream_lanes", &self.blocking_stream.len())
.field("blocking_stream_lanes", &self.blocking_stream.capacity)
.finish()
}
}
@@ -182,7 +185,7 @@ impl RedisConnectionRouter {
)
.await?;
let stream_lanes = stream.len();
let blocking_stream_lanes = blocking_stream.len();
let blocking_stream_lanes = blocking_stream.capacity;
info!(
redis_lanes = "fast,stream,blocking_stream,admin",
redis_stream_lanes = stream_lanes,
@@ -194,7 +197,11 @@ impl RedisConnectionRouter {
stream: Arc::new(stream),
stream_next: Arc::new(AtomicUsize::new(0)),
blocking_stream: Arc::new(blocking_stream),
blocking_stream_next: Arc::new(AtomicUsize::new(0)),
usage_cleanup: Arc::new(RedisBlockingStreamPool::new(
client,
command_timeout_ms,
vec![None, None],
)),
admin,
metrics: Arc::new(RedisConnectionMetrics::default()),
})
@@ -208,13 +215,26 @@ impl RedisConnectionRouter {
self.stream[index].clone()
}
RedisConnectionLane::BlockingStream => {
let index = next_lane_index(&self.blocking_stream_next, self.blocking_stream.len());
self.blocking_stream[index].clone()
unreachable!("blocking stream commands require an exclusive connection lease")
}
RedisConnectionLane::Admin => self.admin.clone(),
}
}
pub(crate) async fn blocking_stream_connection(
&self,
) -> Result<RedisBlockingStreamLease, DataLayerError> {
self.blocking_stream.checkout().await
}
// WATCH/MULTI state must never share a multiplexed connection with other callers.
// These lazy leases own their drivers, so cancellation closes an unfinished transaction.
pub(crate) async fn usage_cleanup_connection(
&self,
) -> Result<RedisBlockingStreamLease, DataLayerError> {
self.usage_cleanup.checkout().await
}
pub(crate) fn record_error(&self, lane: RedisConnectionLane) {
self.metrics
.for_lane(lane)
@@ -257,6 +277,114 @@ impl RedisConnectionRouter {
}
}
struct RedisBlockingConnection {
connection: redis::aio::MultiplexedConnection,
driver: Pin<Box<dyn Future<Output = ()> + Send>>,
}
impl RedisBlockingConnection {
async fn query(&mut self, command: &redis::Cmd) -> Result<redis::Value, DataLayerError> {
// Drive the connection inside its owning query, so cancellation drops the
// socket immediately instead of leaving a spawned driver with an old BLOCK.
tokio::select! {
result = command.query_async(&mut self.connection) => result.map_redis_err(),
() = self.driver.as_mut() => Err(DataLayerError::Redis(
"runtime redis blocking stream connection driver terminated".to_string(),
)),
}
}
}
struct RedisBlockingStreamPool {
client: RedisClient,
command_timeout_ms: Option<u64>,
capacity: usize,
available: Mutex<Vec<Option<RedisBlockingConnection>>>,
permits: Arc<Semaphore>,
}
impl RedisBlockingStreamPool {
fn new(
client: RedisClient,
command_timeout_ms: Option<u64>,
available: Vec<Option<RedisBlockingConnection>>,
) -> Self {
let capacity = available.len();
Self {
client,
command_timeout_ms,
capacity,
available: Mutex::new(available),
permits: Arc::new(Semaphore::new(capacity)),
}
}
async fn checkout(self: &Arc<Self>) -> Result<RedisBlockingStreamLease, DataLayerError> {
let permit = Arc::clone(&self.permits)
.acquire_owned()
.await
.map_err(|_| {
DataLayerError::Redis("runtime redis blocking stream pool closed".to_string())
})?;
let connection = self
.available
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.pop()
.expect("blocking stream permit must have an available slot");
Ok(RedisBlockingStreamLease {
pool: Arc::clone(self),
connection,
reusable: false,
_permit: permit,
})
}
}
pub(crate) struct RedisBlockingStreamLease {
pool: Arc<RedisBlockingStreamPool>,
connection: Option<RedisBlockingConnection>,
reusable: bool,
_permit: OwnedSemaphorePermit,
}
impl RedisBlockingStreamLease {
pub(crate) async fn query(
&mut self,
command: &redis::Cmd,
) -> Result<redis::Value, DataLayerError> {
self.reusable = false;
if self.connection.is_none() {
self.connection = Some(
connect_blocking_stream_lane(&self.pool.client, self.pool.command_timeout_ms)
.await?,
);
}
self.connection
.as_mut()
.expect("blocking stream connection initialized")
.query(command)
.await
}
pub(crate) fn recycle(&mut self) {
self.reusable = true;
}
}
impl Drop for RedisBlockingStreamLease {
fn drop(&mut self) {
let connection = self.connection.take().filter(|_| self.reusable);
self.pool
.available
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.push(connection);
// The permit is released after the slot is restored. An uncompleted
// query has already dropped both its connection and its owned driver.
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct RedisLaneDiagnostics {
pub lane: &'static str,
@@ -406,21 +534,46 @@ async fn connect_blocking_stream_lanes(
client: &RedisClient,
command_timeout_ms: Option<u64>,
requested_lanes: Option<usize>,
) -> Result<Vec<RedisManagedConnection>, DataLayerError> {
) -> Result<RedisBlockingStreamPool, DataLayerError> {
let lane_count = blocking_stream_lane_count(requested_lanes)?;
let mut lanes = Vec::with_capacity(lane_count);
for _ in 0..lane_count {
lanes.push(
connect_lane(
client,
connection_manager_config(command_timeout_ms),
RedisConnectionLane::BlockingStream,
command_timeout_ms,
)
.await?,
);
lanes.push(Some(
connect_blocking_stream_lane(client, command_timeout_ms).await?,
));
}
Ok(lanes)
Ok(RedisBlockingStreamPool::new(
client.clone(),
command_timeout_ms,
lanes,
))
}
async fn connect_blocking_stream_lane(
client: &RedisClient,
command_timeout_ms: Option<u64>,
) -> Result<RedisBlockingConnection, DataLayerError> {
let connect = client.create_multiplexed_tokio_connection();
let result = if let Some(timeout_ms) = command_timeout_ms {
tokio::time::timeout(Duration::from_millis(timeout_ms), connect)
.await
.map_err(|_| {
DataLayerError::TimedOut(format!(
"runtime redis blocking_stream lane connection exceeded {timeout_ms}ms timeout"
))
})?
} else {
connect.await
};
let (connection, driver) = result.map_err(|err| {
DataLayerError::Redis(format!(
"failed to initialize runtime redis blocking_stream lane: {err}"
))
})?;
Ok(RedisBlockingConnection {
connection,
driver: Box::pin(driver),
})
}
fn blocking_stream_lane_count(requested_lanes: Option<usize>) -> Result<usize, DataLayerError> {
@@ -480,10 +633,12 @@ async fn connect_lane(
mod tests {
use super::{
blocking_stream_lane_count, default_blocking_stream_lane_count, next_lane_index,
stream_lane_count, RedisClientConfig, RedisClientFactory, RedisLaneMetrics,
DEFAULT_STREAM_LANES, MAX_BLOCKING_STREAM_LANES_CAP, REDIS_COMMAND_LATENCY_BUCKETS_MS,
stream_lane_count, RedisBlockingStreamPool, RedisClientConfig, RedisClientFactory,
RedisLaneMetrics, DEFAULT_STREAM_LANES, MAX_BLOCKING_STREAM_LANES_CAP,
REDIS_COMMAND_LATENCY_BUCKETS_MS,
};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
#[test]
@@ -543,6 +698,85 @@ mod tests {
assert!(blocking_stream_lane_count(Some(0)).is_err());
}
fn empty_blocking_pool(capacity: usize) -> Arc<RedisBlockingStreamPool> {
Arc::new(RedisBlockingStreamPool::new(
redis::Client::open("redis://127.0.0.1/0").expect("lazy client"),
None,
(0..capacity).map(|_| None).collect(),
))
}
#[tokio::test]
async fn blocking_stream_pool_cancelled_checkout_preserves_owner_and_capacity() {
use std::future::Future;
use std::task::Poll;
let pool = empty_blocking_pool(1);
let owner = pool.checkout().await.expect("first lease");
let mut waiting = Box::pin(pool.checkout());
std::future::poll_fn(|cx| {
assert!(matches!(waiting.as_mut().poll(cx), Poll::Pending));
Poll::Ready(())
})
.await;
drop(waiting);
assert_eq!(pool.permits.available_permits(), 0);
assert!(pool.available.lock().unwrap().is_empty());
drop(owner);
assert_eq!(pool.permits.available_permits(), 1);
let replacement = pool
.checkout()
.await
.expect("cancelled owner slot restored");
assert!(replacement.connection.is_none());
let panic = tokio::spawn(async move {
let _lease = replacement;
panic!("test owner panic");
});
assert!(panic.await.expect_err("owner panicked").is_panic());
assert_eq!(pool.permits.available_permits(), 1);
assert_eq!(pool.available.lock().unwrap().len(), 1);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn blocking_stream_pool_concurrent_checkouts_stay_within_capacity() {
let pool = empty_blocking_pool(3);
let active = Arc::new(AtomicUsize::new(0));
let peak = Arc::new(AtomicUsize::new(0));
let barrier = Arc::new(tokio::sync::Barrier::new(32));
let mut tasks = tokio::task::JoinSet::new();
for _ in 0..32 {
let pool = Arc::clone(&pool);
let active = Arc::clone(&active);
let peak = Arc::clone(&peak);
let barrier = Arc::clone(&barrier);
tasks.spawn(async move {
barrier.wait().await;
for _ in 0..32 {
let lease = pool.checkout().await.expect("bounded lease");
let concurrent = active.fetch_add(1, Ordering::SeqCst) + 1;
peak.fetch_max(concurrent, Ordering::SeqCst);
assert!(concurrent <= 3);
tokio::task::yield_now().await;
active.fetch_sub(1, Ordering::SeqCst);
drop(lease);
}
});
}
tokio::time::timeout(Duration::from_secs(5), async {
while let Some(result) = tasks.join_next().await {
result.expect("checkout task");
}
})
.await
.expect("all checkouts complete without losing capacity");
assert!(peak.load(Ordering::SeqCst) <= 3);
assert_eq!(active.load(Ordering::SeqCst), 0);
assert_eq!(pool.permits.available_permits(), 3);
assert_eq!(pool.available.lock().unwrap().len(), 3);
}
#[test]
fn stream_lane_count_uses_fixed_default() {
assert_eq!(stream_lane_count(), DEFAULT_STREAM_LANES);
@@ -0,0 +1,44 @@
-- Redis scripts do not roll back errors. Complete all predictable checks before XADD.
if #KEYS ~= 2 or #ARGV < 4 or (#ARGV - 2) % 2 ~= 0 then
return redis.error_reply('ERR invalid pending transfer arguments')
end
if KEYS[1] == KEYS[2] then
return redis.error_reply('ERR pending transfer source and destination must differ')
end
if type(redis.acl_check_cmd) ~= 'function' then
return redis.error_reply('ERR pending transfer requires Redis 7 or later for ACL preflight')
end
-- The exact range bounds keep work independent of the size of the PEL.
local pending = redis.call('XPENDING', KEYS[1], ARGV[1], ARGV[2], ARGV[2], 1)
if #pending == 0 then
return {0, '', 0, 0}
end
if pending[1][1] ~= ARGV[2] then
return redis.error_reply('ERR pending transfer requires an exact canonical entry ID')
end
local destination_type = redis.call('TYPE', KEYS[2]).ok
if destination_type ~= 'none' and destination_type ~= 'stream' then
return redis.error_reply('WRONGTYPE pending transfer destination must be a stream')
end
local append = {'XADD', KEYS[2], '*'}
for index = 3, #ARGV do
append[#append + 1] = ARGV[index]
end
if not redis.acl_check_cmd(unpack(append)) then
return redis.error_reply('NOPERM pending transfer requires XADD permission')
end
if not redis.acl_check_cmd('XACK', KEYS[1], ARGV[1], ARGV[2]) then
return redis.error_reply('NOPERM pending transfer requires XACK permission')
end
if not redis.acl_check_cmd('XDEL', KEYS[1], ARGV[2]) then
return redis.error_reply('NOPERM pending transfer requires XDEL permission')
end
-- PEL membership, stream types and ACLs cannot change between these commands.
-- A trimmed source body is still recoverable from the caller's retained fields.
local destination_id = redis.call(unpack(append))
local acked = redis.call('XACK', KEYS[1], ARGV[1], ARGV[2])
local deleted = redis.call('XDEL', KEYS[1], ARGV[2])
return {1, destination_id, acked, deleted}
@@ -0,0 +1,500 @@
use super::*;
type TransferTestConnection = ::redis::aio::MultiplexedConnection;
const TRANSFER_GROUP: &str = "transfer-workers";
const TRANSFER_USER: &str = "transfer-worker";
async fn transfer_runtime(
protocol: &str,
) -> Option<(TestRedisServer, RuntimeState, TransferTestConnection)> {
let Some(server) = TestRedisServer::start().await else {
eprintln!(
"dead letter transfer {protocol} skipped: isolated Redis fixture unavailable; check AETHER_REDIS_SERVER_BIN"
);
return None;
};
let mut admin = ::redis::Client::open(server.redis_url.clone())
.expect("transfer admin client")
.get_multiplexed_async_connection()
.await
.expect("transfer admin connection");
::redis::cmd("ACL")
.arg("SETUSER")
.arg(TRANSFER_USER)
.arg("on")
.arg(">transfer-test-password")
.arg("~*")
.arg("+@all")
.query_async::<()>(&mut admin)
.await
.expect("transfer test user");
let runtime = RuntimeState::redis_with_blocking_stream_lanes(
RedisClientConfig {
url: format!(
"redis://{TRANSFER_USER}:[email protected]:{}/5?protocol={protocol}",
server.port
),
key_prefix: Some(format!("transfer-{protocol}")),
},
Some(5_000),
Some(4),
)
.await
.expect("authenticated transfer runtime");
::redis::cmd("SELECT")
.arg(5)
.query_async::<()>(&mut admin)
.await
.expect("transfer admin database");
eprintln!(
"dead letter transfer fixture ready: protocol={protocol} db=5 port={} authenticated=true",
server.port
);
Some((server, runtime, admin))
}
fn transfer_source_fields() -> BTreeMap<String, String> {
BTreeMap::from([
(
"payload".to_string(),
"malformed\r\n\"quoted\"\\\u{4e2d}\u{6587}\0".to_string(),
),
("legacy".to_string(), "retain every field".to_string()),
(
String::new(),
"empty field name is valid Redis data".to_string(),
),
])
}
fn transfer_archive_fields(entry: &RuntimeQueueEntry) -> BTreeMap<String, String> {
BTreeMap::from([
(
"payload".to_string(),
serde_json::json!({
"entry_id": entry.id,
"fields": entry.fields,
"error": "invalid record\r\n\"details\"\\\u{4e2d}\u{6587}"
})
.to_string(),
),
("archive_version".to_string(), "1".to_string()),
])
}
async fn seed_transfer_entry(runtime: &RuntimeState, source: &str) -> RuntimeQueueEntry {
RuntimeQueueStore::ensure_consumer_group(runtime, source, TRANSFER_GROUP, "0-0")
.await
.expect("source consumer group");
let id = RuntimeQueueStore::append_fields_with_maxlen(
runtime,
source,
&transfer_source_fields(),
None,
)
.await
.expect("source append");
let mut entries =
RuntimeQueueStore::read_group(runtime, source, TRANSFER_GROUP, "owner", 1, None)
.await
.expect("pending source entry");
assert_eq!(entries.len(), 1);
let entry = entries.pop().expect("one entry");
assert_eq!(entry.id, id);
assert_eq!(entry.fields, transfer_source_fields());
entry
}
async fn transfer_entries(
admin: &mut TransferTestConnection,
stream: &str,
) -> Vec<RuntimeQueueEntry> {
let rows = ::redis::cmd("XRANGE")
.arg(stream)
.arg("-")
.arg("+")
.query_async::<::redis::streams::StreamRangeReply>(admin)
.await
.expect("inspect transfer stream");
rows.ids
.into_iter()
.map(|row| RuntimeQueueEntry {
id: row.id,
fields: row
.map
.into_iter()
.map(|(field, value)| {
let value =
::redis::from_redis_value::<String>(&value).expect("string field value");
(field, value)
})
.collect(),
})
.collect()
}
async fn transfer_pending(
admin: &mut TransferTestConnection,
source: &str,
) -> Vec<(String, String, u64, u64)> {
::redis::cmd("XPENDING")
.arg(source)
.arg(TRANSFER_GROUP)
.arg("-")
.arg("+")
.arg(16)
.query_async(admin)
.await
.expect("inspect transfer pending entries")
}
async fn transfer_entry(
runtime: &RuntimeState,
source: &str,
entry: &RuntimeQueueEntry,
destination: &str,
fields: &BTreeMap<String, String>,
) -> Result<RuntimeQueueTransferOutcome, DataLayerError> {
RuntimeQueueStore::try_transfer_pending_to_stream(
runtime,
source,
TRANSFER_GROUP,
&entry.id,
destination,
fields,
)
.await
.map(|outcome| outcome.expect("Redis implements atomic transfer"))
}
async fn assert_transfer_source_unchanged(
admin: &mut TransferTestConnection,
source: &str,
entry: &RuntimeQueueEntry,
) {
assert_eq!(transfer_entries(admin, source).await, [entry.clone()]);
let pending = transfer_pending(admin, source).await;
assert_eq!(pending.len(), 1);
assert_eq!(pending[0].0, entry.id);
assert_eq!(pending[0].1, "owner");
}
async fn assert_transfer_completed(
admin: &mut TransferTestConnection,
source: &str,
destination: &str,
fields: &BTreeMap<String, String>,
) {
assert!(transfer_entries(admin, source).await.is_empty());
assert!(transfer_pending(admin, source).await.is_empty());
let archived = transfer_entries(admin, destination).await;
assert_eq!(archived.len(), 1);
assert_eq!(&archived[0].fields, fields);
}
#[tokio::test]
async fn redis_dead_letter_transfer_concurrent_consumers_archive_exactly_once() {
for protocol in ["resp2", "resp3"] {
let Some((_server, runtime, mut admin)) = transfer_runtime(protocol).await else {
return;
};
let source = "usage:{transfer}:concurrent";
let destination = "usage:{transfer}:concurrent:dlq";
let entry = seed_transfer_entry(&runtime, source).await;
let fields = transfer_archive_fields(&entry);
let barrier = Arc::new(tokio::sync::Barrier::new(16));
let mut tasks = tokio::task::JoinSet::new();
for _ in 0..16 {
let runtime = runtime.clone();
let entry = entry.clone();
let fields = fields.clone();
let barrier = Arc::clone(&barrier);
tasks.spawn(async move {
barrier.wait().await;
transfer_entry(&runtime, source, &entry, destination, &fields)
.await
.expect("concurrent transfer")
});
}
let mut transferred_ids = Vec::new();
let mut not_pending = 0;
while let Some(result) = tasks.join_next().await {
match result.expect("transfer task") {
RuntimeQueueTransferOutcome::Transferred {
destination_id,
acked,
deleted,
} => {
assert_eq!((acked, deleted), (1, 1));
transferred_ids.push(destination_id);
}
RuntimeQueueTransferOutcome::NotPending => not_pending += 1,
}
}
assert_eq!(not_pending, 15);
assert_eq!(transferred_ids.len(), 1);
assert_transfer_completed(&mut admin, source, destination, &fields).await;
assert_eq!(
transfer_entries(&mut admin, destination).await[0].id,
transferred_ids[0]
);
let diagnostics = runtime.redis_diagnostics().await.unwrap().unwrap();
let lane = diagnostics
.lanes
.iter()
.find(|lane| lane.lane == "blocking_stream")
.expect("exclusive stream lane diagnostics");
assert!(lane.command_count >= 16);
}
}
#[tokio::test]
async fn redis_dead_letter_transfer_retry_after_ignored_success_does_not_archive_again() {
let Some((_server, runtime, mut admin)) = transfer_runtime("resp2").await else {
return;
};
let source = "usage:{transfer}:retry";
let destination = "usage:{transfer}:retry:dlq";
let entry = seed_transfer_entry(&runtime, source).await;
let fields = transfer_archive_fields(&entry);
// Commit the operation, but discard its result as a caller missing the reply would.
let _ = transfer_entry(&runtime, source, &entry, destination, &fields)
.await
.expect("first transfer commits");
let changed_fields = BTreeMap::from([("payload".to_string(), "retry value".to_string())]);
assert_eq!(
transfer_entry(&runtime, source, &entry, destination, &changed_fields)
.await
.expect("retry is successful"),
RuntimeQueueTransferOutcome::NotPending
);
assert_transfer_completed(&mut admin, source, destination, &fields).await;
}
#[tokio::test]
async fn redis_dead_letter_transfer_acl_preflight_prevents_partial_writes() {
let Some((_server, runtime, mut admin)) = transfer_runtime("resp3").await else {
return;
};
for forbidden in ["XADD", "XACK", "XDEL"] {
let source = format!("usage:{{transfer}}:acl-{forbidden}");
let destination = format!("{source}:dlq");
let entry = seed_transfer_entry(&runtime, &source).await;
let fields = transfer_archive_fields(&entry);
::redis::cmd("ACL")
.arg("SETUSER")
.arg(TRANSFER_USER)
.arg(format!("-{forbidden}"))
.query_async::<()>(&mut admin)
.await
.expect("deny one write command");
let error = transfer_entry(&runtime, &source, &entry, &destination, &fields)
.await
.expect_err("denied write must fail before archiving");
assert!(
error
.to_string()
.contains(&format!("requires {forbidden} permission")),
"expected the {forbidden} preflight error, got {error}"
);
assert_transfer_source_unchanged(&mut admin, &source, &entry).await;
assert!(transfer_entries(&mut admin, &destination).await.is_empty());
::redis::cmd("ACL")
.arg("SETUSER")
.arg(TRANSFER_USER)
.arg(format!("+{forbidden}"))
.query_async::<()>(&mut admin)
.await
.expect("restore one write command");
assert!(matches!(
transfer_entry(&runtime, &source, &entry, &destination, &fields)
.await
.expect("retry after restoring permission"),
RuntimeQueueTransferOutcome::Transferred {
acked: 1,
deleted: 1,
..
}
));
assert_eq!(
transfer_entry(&runtime, &source, &entry, &destination, &fields)
.await
.expect("idempotent retry"),
RuntimeQueueTransferOutcome::NotPending
);
assert_transfer_completed(&mut admin, &source, &destination, &fields).await;
}
}
#[tokio::test]
async fn redis_dead_letter_transfer_invalid_state_preserves_source_until_repaired() {
let Some((_server, runtime, mut admin)) = transfer_runtime("resp2").await else {
return;
};
let source = "usage:{transfer}:invalid";
let destination = "usage:{transfer}:invalid:dlq";
let entry = seed_transfer_entry(&runtime, source).await;
let fields = transfer_archive_fields(&entry);
::redis::cmd("SET")
.arg(destination)
.arg("existing non-stream data")
.query_async::<()>(&mut admin)
.await
.expect("wrong-type destination");
let error = transfer_entry(&runtime, source, &entry, destination, &fields)
.await
.expect_err("wrong type must not acknowledge source");
assert!(error.to_string().contains("WRONGTYPE"));
assert_transfer_source_unchanged(&mut admin, source, &entry).await;
assert_eq!(
::redis::cmd("GET")
.arg(destination)
.query_async::<String>(&mut admin)
.await
.unwrap(),
"existing non-stream data"
);
::redis::cmd("DEL")
.arg(destination)
.query_async::<usize>(&mut admin)
.await
.expect("repair destination type");
let error = RuntimeQueueStore::try_transfer_pending_to_stream(
&runtime,
source,
"missing-group",
&entry.id,
destination,
&fields,
)
.await
.expect_err("missing group must fail before archiving");
assert!(error.to_string().contains("NOGROUP"));
for invalid_id in [
"",
"-",
"+",
"1",
"01-0",
"1-00",
"(1-0",
"1-+0",
"18446744073709551616-0",
] {
assert!(matches!(
RuntimeQueueStore::try_transfer_pending_to_stream(
&runtime,
source,
TRANSFER_GROUP,
invalid_id,
destination,
&fields,
)
.await,
Err(DataLayerError::InvalidInput(_))
));
}
assert!(matches!(
transfer_entry(&runtime, source, &entry, source, &fields).await,
Err(DataLayerError::InvalidInput(_))
));
assert!(matches!(
transfer_entry(&runtime, source, &entry, destination, &BTreeMap::new()).await,
Err(DataLayerError::InvalidInput(_))
));
assert_transfer_source_unchanged(&mut admin, source, &entry).await;
assert!(transfer_entries(&mut admin, destination).await.is_empty());
assert!(matches!(
transfer_entry(&runtime, source, &entry, destination, &fields)
.await
.expect("valid transfer after failed attempts"),
RuntimeQueueTransferOutcome::Transferred {
acked: 1,
deleted: 1,
..
}
));
assert_transfer_completed(&mut admin, source, destination, &fields).await;
}
#[tokio::test]
async fn redis_dead_letter_transfer_archives_retained_fields_after_source_trim() {
for protocol in ["resp2", "resp3"] {
let Some((_server, runtime, mut admin)) = transfer_runtime(protocol).await else {
return;
};
let source = "usage:{transfer}:trimmed";
let destination = "usage:{transfer}:trimmed:dlq";
let entry = seed_transfer_entry(&runtime, source).await;
let fields = transfer_archive_fields(&entry);
let trimmed = ::redis::cmd("XTRIM")
.arg(source)
.arg("MAXLEN")
.arg(0)
.query_async::<usize>(&mut admin)
.await
.expect("trim source body while retaining PEL");
assert_eq!(trimmed, 1);
assert!(transfer_entries(&mut admin, source).await.is_empty());
assert_eq!(transfer_pending(&mut admin, source).await[0].0, entry.id);
assert!(matches!(
transfer_entry(&runtime, source, &entry, destination, &fields)
.await
.expect("pending body remains recoverable from caller fields"),
RuntimeQueueTransferOutcome::Transferred {
acked: 1,
deleted: 0,
..
}
));
assert_eq!(
transfer_entry(&runtime, source, &entry, destination, &fields)
.await
.expect("trimmed entry retry"),
RuntimeQueueTransferOutcome::NotPending
);
assert_transfer_completed(&mut admin, source, destination, &fields).await;
}
}
#[tokio::test]
async fn redis_dead_letter_transfer_preserves_ids_larger_than_lua_integer_precision() {
let Some((_server, runtime, mut admin)) = transfer_runtime("resp3").await else {
return;
};
let source = "usage:{transfer}:large-id";
let destination = "usage:{transfer}:large-id:dlq";
let entry_id = "9007199254740993-18446744073709551614";
RuntimeQueueStore::ensure_consumer_group(&runtime, source, TRANSFER_GROUP, "0-0")
.await
.expect("large-ID source group");
::redis::cmd("XADD")
.arg(source)
.arg(entry_id)
.arg("payload")
.arg("retained value")
.query_async::<String>(&mut admin)
.await
.expect("large stream ID");
let mut entries =
RuntimeQueueStore::read_group(&runtime, source, TRANSFER_GROUP, "owner", 1, None)
.await
.expect("read large-ID entry");
assert_eq!(entries.len(), 1);
let entry = entries.pop().unwrap();
assert_eq!(entry.id, entry_id);
let fields = transfer_archive_fields(&entry);
assert!(matches!(
transfer_entry(&runtime, source, &entry, destination, &fields)
.await
.expect("transfer exact large ID"),
RuntimeQueueTransferOutcome::Transferred {
acked: 1,
deleted: 1,
..
}
));
assert_transfer_completed(&mut admin, source, destination, &fields).await;
}
@@ -4,6 +4,7 @@ mod lock;
mod namespace;
mod runtime;
mod stream;
mod usage_cleanup;
pub use client::{RedisClientConfig, RedisLaneDiagnostics};
pub use kv::{RedisKvRunner, RedisKvRunnerConfig};
+186 -25
View File
@@ -7,9 +7,13 @@ use crate::redis::{
};
use crate::{
DataLayerError, RateLimitCheck, RateLimitInput, RateLimitScope, RuntimeSemaphoreError,
UsageLimitCheck, UsageLimitInput, UsageLimitReleaseInput,
ScoreWindowU64Stats, UsageLimitCheck, UsageLimitInput, UsageLimitReleaseInput,
SCORE_WINDOW_AGGREGATION_MEMBER_LIMIT,
};
const SCORE_WINDOW_STATS_SCRIPT: &str = include_str!("score_window.lua");
const SCORE_WINDOW_STATS_PIPELINE_KEY_LIMIT: usize = 16;
const RATE_LIMIT_CHECK_AND_CONSUME_SCRIPT: &str = r#"
local user_key = KEYS[1]
local key_key = KEYS[2]
@@ -52,15 +56,119 @@ end
return {1, 0, 0, remaining}
"#;
const USAGE_LIMIT_CHECK_AND_CONSUME_SCRIPT: &str = r#"
pub(super) const USAGE_LIMIT_CHECK_AND_CONSUME_SCRIPT: &str = r#"
local count = #KEYS
local now = tonumber(ARGV[1])
local event_id = ARGV[2]
-- Large mixed windows are copied in bounded commands on an exclusive WATCH
-- connection. This read-only pass must finish before pruning any of the rules.
if ARGV[count * 3 + 3] ~= 'inline' and redis.acl_check_cmd then
local deferred = {2}
for i = 1, count do
local key = KEYS[i]
if redis.call('ZCARD', key) > 4096 then
local cutoff = now - tonumber(ARGV[(i - 1) * 3 + 4]) * 1000
local expired = redis.pcall('ZCOUNT', key, '-inf', cutoff)
if type(expired) == 'number' and expired > 4096 then
local live = redis.call('ZCARD', key) - expired
local temporary = key .. ':__usage_copy:acl'
if live > 256
and redis.acl_check_cmd('WATCH', key)
and redis.acl_check_cmd('UNWATCH')
and redis.acl_check_cmd('MULTI')
and redis.acl_check_cmd('EXEC')
and redis.acl_check_cmd('EVAL', 'return 1', 0)
and redis.acl_check_cmd('EXISTS', temporary)
and redis.acl_check_cmd('ZRANGE', key, 0, 511, 'WITHSCORES')
and redis.acl_check_cmd('ZADD', temporary, 0, 'acl')
and redis.acl_check_cmd('PTTL', key)
and redis.acl_check_cmd('PEXPIRE', temporary, 60000)
and redis.acl_check_cmd('PERSIST', temporary)
and redis.acl_check_cmd('UNLINK', key, temporary)
and redis.acl_check_cmd('RENAME', temporary, key) then
deferred[#deferred + 1] = i
deferred[#deferred + 1] = expired
deferred[#deferred + 1] = live
end
end
end
end
if #deferred > 1 then return deferred end
end
local function replace_with_survivors(key, live)
if not redis.acl_check_cmd then return false end
local temporary = key .. ':__usage_trim'
local exists = redis.pcall('EXISTS', temporary)
local ttl = redis.pcall('PTTL', key)
if exists ~= 0 or type(ttl) ~= 'number' or ttl == 0 or ttl < -1 then
return false
end
local rows = redis.call('ZRANGE', key, -live, -1, 'WITHSCORES')
local args = {}
for i = 1, #rows, 2 do
args[#args + 1] = rows[i + 1]
args[#args + 1] = rows[i]
end
-- Validate every write before detaching the original. The temporary key
-- is fully built first; restricted ACLs keep the original cleanup path.
if not redis.acl_check_cmd('ZADD', temporary, unpack(args))
or not redis.acl_check_cmd('UNLINK', key)
or not redis.acl_check_cmd('UNLINK', temporary)
or not redis.acl_check_cmd('RENAME', temporary, key)
or (ttl > 0 and not redis.acl_check_cmd('PEXPIRE', temporary, ttl)) then
return false
end
if type(redis.pcall('ZADD', temporary, unpack(args))) ~= 'number' then
return false
end
if ttl > 0 and redis.pcall('PEXPIRE', temporary, ttl) ~= 1 then
redis.call('UNLINK', temporary)
return false
end
if type(redis.pcall('UNLINK', key)) ~= 'number' then
redis.call('UNLINK', temporary)
return false
end
redis.call('RENAME', temporary, key)
return true
end
local function prune_window(key, cutoff)
local cardinality = redis.call('ZCARD', key)
if cardinality == 0 then return 0 end
if cardinality > 256 then
local earliest = redis.call('ZRANGE', key, 0, 0, 'WITHSCORES')
if #earliest >= 2 and tonumber(earliest[2]) > cutoff then
return cardinality
end
local latest = redis.call('ZRANGE', key, -1, -1, 'WITHSCORES')
if #latest >= 2 and tonumber(latest[2]) <= cutoff then
-- The entire window is expired. Redis can free the detached object
-- off its command thread while this key is reused.
local result = redis.pcall('UNLINK', key)
if type(result) == 'number' then return 0 end
elseif cardinality > 1024 then
local expired = redis.pcall('ZCOUNT', key, '-inf', cutoff)
if type(expired) == 'number' then
local live = cardinality - expired
if live > 0 and live <= 256 and expired > live * 4
and replace_with_survivors(key, live) then
return live
end
end
end
end
-- Preserve the existing path for large live windows and restricted ACLs.
-- Partial deferred deletion could resurrect entries on out-of-order calls.
return cardinality - redis.call('ZREMRANGEBYSCORE', key, '-inf', cutoff)
end
local current_counts = {}
for i = 1, count do
local window_ms = tonumber(ARGV[(i - 1) * 3 + 4]) * 1000
local cutoff = now - window_ms
redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', cutoff)
current_counts[i] = prune_window(KEYS[i], now - window_ms)
end
for i = 1, count do
@@ -68,7 +176,7 @@ for i = 1, count do
local window_ms = tonumber(ARGV[(i - 1) * 3 + 4]) * 1000
local already_consumed = redis.call('ZSCORE', KEYS[i], event_id)
if not already_consumed then
local current = redis.call('ZCARD', KEYS[i])
local current = current_counts[i]
if current >= limit then
local earliest = redis.call('ZRANGE', KEYS[i], 0, 0, 'WITHSCORES')
local retry_after = 1
@@ -362,30 +470,17 @@ impl RedisRuntimeRunner {
&self,
input: UsageLimitInput<'_>,
) -> Result<UsageLimitCheck, DataLayerError> {
let script = script(USAGE_LIMIT_CHECK_AND_CONSUME_SCRIPT);
let mut invocation = script.prepare_invoke();
for rule in input.rules {
invocation.key(self.keyspace.key(rule.key));
}
invocation.arg(input.now_unix_ms as i64);
invocation.arg(input.event_id);
for rule in input.rules {
invocation.arg(rule.limit as i64);
invocation.arg(rule.window_seconds as i64);
invocation.arg(rule.retention_seconds as i64);
}
let keys = input
.rules
.iter()
.map(|rule| self.keyspace.key(rule.key))
.collect::<Vec<_>>();
let raw = run_lane_with_timeout(
&self.connections,
RedisConnectionLane::Fast,
self.command_timeout_ms,
Some(self.command_timeout_ms.unwrap_or(30_000)),
"runtime usage limit check",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Fast);
invocation
.invoke_async::<Vec<i64>>(&mut connection)
.await
.map_redis_err()
},
super::usage_cleanup::check_and_consume(&self.connections, &keys, &input),
)
.await?;
match raw.first().copied() {
@@ -539,6 +634,72 @@ impl RedisRuntimeRunner {
.await
}
pub(crate) async fn score_window_u64_stats_by_min(
&self,
keys: &[String],
min_score: f64,
) -> Result<Vec<Option<ScoreWindowU64Stats>>, DataLayerError> {
let script = script(SCORE_WINDOW_STATS_SCRIPT);
let mut output = Vec::with_capacity(keys.len());
for batch in keys.chunks(SCORE_WINDOW_STATS_PIPELINE_KEY_LIMIT) {
let values: Vec<(u8, String, u64)> = run_lane_with_timeout(
&self.connections,
RedisConnectionLane::Admin,
self.command_timeout_ms,
"runtime score window stats",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Admin);
let mut pipeline = redis::pipe();
for key in batch {
pipeline
.cmd("EVALSHA")
.arg(script.get_hash())
.arg(1)
.arg(self.keyspace.key(key))
.arg(min_score)
.arg(SCORE_WINDOW_AGGREGATION_MEMBER_LIMIT);
}
match pipeline.query_async(&mut connection).await {
Err(err) if err.kind() == redis::ErrorKind::NoScriptError => {
script
.prepare_invoke()
.load_async(&mut connection)
.await
.map_redis_err()?;
pipeline.query_async(&mut connection).await.map_redis_err()
}
result => result.map_redis_err(),
}
},
)
.await?;
if values.len() != batch.len() {
return Err(DataLayerError::UnexpectedValue(
"runtime score window stats result count mismatch".to_string(),
));
}
for (aggregated, sum, positive_count) in values {
output.push(match aggregated {
0 => None,
1 => Some(ScoreWindowU64Stats {
sum: sum.parse().map_err(|_| {
DataLayerError::UnexpectedValue(
"runtime score window stats returned an invalid sum".to_string(),
)
})?,
positive_count,
}),
_ => {
return Err(DataLayerError::UnexpectedValue(
"runtime score window stats returned an invalid status".to_string(),
))
}
});
}
}
Ok(output)
}
pub(crate) async fn score_remove_by_score(
&self,
key: &str,
@@ -0,0 +1,36 @@
-- Count before reading members so a large window cannot run an unbounded Lua loop.
local count = redis.call('ZCOUNT', KEYS[1], ARGV[1], '+inf')
if count > tonumber(ARGV[2]) then
return {0, '0', 0}
end
local members = redis.call('ZRANGEBYSCORE', KEYS[1], ARGV[1], '+inf')
local high, low, positive = 0, 0, 0
local max_high, max_low = 18446744073, 709551615
for _, member in ipairs(members) do
local value = string.match(member, ':([^:]*)$')
if value and string.match(value, '^%+?%d+$') then
value = string.gsub(value, '^%+', '')
value = string.gsub(value, '^0+', '')
if #value > 0 and (#value < 20 or (#value == 20 and value <= '18446744073709551615')) then
positive = positive + 1
-- Two base-1e9 limbs keep every integer operation exactly representable
-- in Redis Lua's doubles, including values above 2^53 and u64::MAX.
local split = math.max(0, #value - 9)
local value_high = tonumber(string.sub(value, 1, split)) or 0
local value_low = tonumber(string.sub(value, split + 1))
low = low + value_low
high = high + value_high + math.floor(low / 1000000000)
low = low % 1000000000
if high > max_high or (high == max_high and low > max_low) then
high, low = max_high, max_low
end
end
end
end
local total = string.format('%.0f', low)
if high > 0 then
total = string.format('%.0f', high) .. string.format('%09d', low)
end
return {1, total, positive}
+190 -50
View File
@@ -1,16 +1,19 @@
use std::collections::BTreeMap;
use std::collections::{BTreeMap, HashMap};
use std::future::Future;
use redis::from_redis_value;
use redis::streams::StreamReadReply;
use redis::Value as RedisValue;
use redis::{from_owned_redis_value, from_redis_value};
use crate::error::{redis_error, RedisResultExt};
use crate::redis::{
run_lane_with_timeout, RedisClientConfig, RedisClientFactory, RedisConnectionLane,
RedisConnectionRouter, RedisKeyspace,
};
use crate::{DataLayerError, RuntimeQueueStats};
use crate::{
validate_runtime_queue_transfer, DataLayerError, RuntimeQueueStats, RuntimeQueueTransferOutcome,
};
const DEAD_LETTER_TRANSFER_SCRIPT: &str = include_str!("dead_letter_transfer.lua");
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisStreamName(pub String);
@@ -259,6 +262,41 @@ impl RedisStreamRunner {
self.append_fields(stream, &fields).await
}
/// Atomically archives a pending entry and removes it from its source group and stream.
/// Requires Redis 7+ for ACL preflight. Both stream keys must share a slot on Redis Cluster.
pub async fn try_transfer_pending_to_stream(
&self,
source: &str,
group: &str,
entry_id: &str,
destination: &str,
destination_fields: &BTreeMap<String, String>,
) -> Result<RuntimeQueueTransferOutcome, DataLayerError> {
validate_runtime_queue_transfer(source, group, entry_id, destination, destination_fields)?;
self.run_with_timeout(
RedisConnectionLane::BlockingStream,
"redis stream pending transfer",
async {
let mut lease = self.connections.blocking_stream_connection().await?;
let mut command = redis::cmd("EVAL");
command
.arg(DEAD_LETTER_TRANSFER_SCRIPT)
.arg(2)
.arg(source)
.arg(destination)
.arg(group)
.arg(entry_id);
for (field, value) in destination_fields {
command.arg(field).arg(value);
}
let outcome = parse_transfer_result(lease.query(&command).await?)?;
lease.recycle();
Ok(outcome)
},
)
.await
}
pub async fn read_group(
&self,
stream: &RedisStreamName,
@@ -275,7 +313,6 @@ impl RedisStreamRunner {
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")
@@ -288,6 +325,14 @@ impl RedisStreamRunner {
}
command.arg("STREAMS").arg(&stream.0).arg(">");
if lane == RedisConnectionLane::BlockingStream {
let mut lease = self.connections.blocking_stream_connection().await?;
let entries = parse_stream_read_entries(lease.query(&command).await?)?;
lease.recycle();
return Ok(entries);
}
let mut connection = self.connections.connection(lane);
let reply = command
.query_async::<RedisValue>(&mut connection)
.await
@@ -364,22 +409,25 @@ impl RedisStreamRunner {
validate_stream_position(start_id)?;
config.validate()?;
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)
.arg(&consumer.0)
.arg(config.min_idle_ms)
.arg(start_id)
.arg("COUNT")
.arg(config.count)
.query_async::<RedisValue>(&mut connection)
.await
.map_redis_err()?;
parse_reclaim_result(reply)
})
self.run_with_timeout(
RedisConnectionLane::BlockingStream,
"redis stream reclaim",
async {
let mut lease = self.connections.blocking_stream_connection().await?;
let mut command = redis::cmd("XAUTOCLAIM");
command
.arg(&stream.0)
.arg(&group.0)
.arg(&consumer.0)
.arg(config.min_idle_ms)
.arg(start_id)
.arg("COUNT")
.arg(config.count);
let result = parse_reclaim_result(lease.query(&command).await?)?;
lease.recycle();
Ok(result)
},
)
.await
}
@@ -536,18 +584,21 @@ fn parse_stream_read_entries(value: RedisValue) -> Result<Vec<RedisStreamEntry>,
return Ok(Vec::new());
}
let reply = from_redis_value::<StreamReadReply>(&value).map_err(redis_error)?;
Ok(reply
.keys
// StreamReadReply in redis 0.28 falls back to borrowed conversion even for an
// owned input. Its underlying containers support moving every payload buffer.
type StreamReadRows = Vec<HashMap<String, Vec<HashMap<String, HashMap<String, RedisValue>>>>>;
let rows = from_owned_redis_value::<StreamReadRows>(value).map_err(redis_error)?;
Ok(rows
.into_iter()
.flat_map(|key| key.ids.into_iter())
.map(|id| RedisStreamEntry {
id: id.id,
fields: id
.map
.flat_map(HashMap::into_values)
.flatten()
.flat_map(HashMap::into_iter)
.map(|(id, fields)| RedisStreamEntry {
id,
fields: fields
.into_iter()
.filter_map(|(field, value)| {
redis::from_redis_value::<String>(&value)
from_owned_redis_value::<String>(value)
.ok()
.map(|text| (field, text))
})
@@ -556,6 +607,31 @@ fn parse_stream_read_entries(value: RedisValue) -> Result<Vec<RedisStreamEntry>,
.collect())
}
fn parse_transfer_result(value: RedisValue) -> Result<RuntimeQueueTransferOutcome, DataLayerError> {
if !matches!(&value, RedisValue::Array(parts) if parts.len() == 4) {
return Err(DataLayerError::UnexpectedValue(
"redis stream pending transfer returned an invalid result shape".to_string(),
));
}
let (transferred, destination_id, acked, deleted) =
from_owned_redis_value::<(i64, String, usize, usize)>(value).map_err(redis_error)?;
match transferred {
0 if destination_id.is_empty() && acked == 0 && deleted == 0 => {
Ok(RuntimeQueueTransferOutcome::NotPending)
}
1 if !destination_id.is_empty() && acked == 1 && deleted <= 1 => {
Ok(RuntimeQueueTransferOutcome::Transferred {
destination_id,
acked,
deleted,
})
}
_ => Err(DataLayerError::UnexpectedValue(
"redis stream pending transfer returned inconsistent outcome fields".to_string(),
)),
}
}
fn parse_reclaim_result(value: RedisValue) -> Result<RedisStreamReclaimResult, DataLayerError> {
let RedisValue::Array(parts) = value else {
return Err(DataLayerError::UnexpectedValue(
@@ -570,9 +646,13 @@ fn parse_reclaim_result(value: RedisValue) -> Result<RedisStreamReclaimResult, D
)));
}
let next_start_id = parse_string_value(&parts[0], "redis xautoclaim next_start_id")?;
let entries = parse_reclaim_entries(&parts[1])?;
let deleted_ids = match parts.get(2) {
let mut parts = parts.into_iter();
let next_start_id = parse_owned_string_value(
parts.next().expect("validated reclaim result length"),
"redis xautoclaim next_start_id",
)?;
let entries = parse_reclaim_entries(parts.next().expect("validated reclaim result length"))?;
let deleted_ids = match parts.next() {
Some(value) => parse_string_array(value, "redis xautoclaim deleted_ids")?,
None => Vec::new(),
};
@@ -584,9 +664,9 @@ fn parse_reclaim_result(value: RedisValue) -> Result<RedisStreamReclaimResult, D
})
}
fn parse_reclaim_entries(value: &RedisValue) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
fn parse_reclaim_entries(value: RedisValue) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
match value {
RedisValue::Array(entries) => entries.iter().map(parse_reclaim_entry).collect(),
RedisValue::Array(entries) => entries.into_iter().map(parse_reclaim_entry).collect(),
RedisValue::Nil => Ok(Vec::new()),
_ => Err(DataLayerError::UnexpectedValue(
"redis xautoclaim entries payload was not an array".to_string(),
@@ -594,7 +674,7 @@ fn parse_reclaim_entries(value: &RedisValue) -> Result<Vec<RedisStreamEntry>, Da
}
}
fn parse_reclaim_entry(value: &RedisValue) -> Result<RedisStreamEntry, DataLayerError> {
fn parse_reclaim_entry(value: RedisValue) -> Result<RedisStreamEntry, DataLayerError> {
let RedisValue::Array(parts) = value else {
return Err(DataLayerError::UnexpectedValue(
"redis xautoclaim entry was not an array".to_string(),
@@ -607,13 +687,20 @@ fn parse_reclaim_entry(value: &RedisValue) -> Result<RedisStreamEntry, DataLayer
)));
}
let id = parse_string_value(&parts[0], "redis xautoclaim entry id")?;
let fields = parse_string_map(&parts[1], "redis xautoclaim entry fields")?;
let mut parts = parts.into_iter();
let id = parse_owned_string_value(
parts.next().expect("validated reclaim entry length"),
"redis xautoclaim entry id",
)?;
let fields = parse_string_map(
parts.next().expect("validated reclaim entry length"),
"redis xautoclaim entry fields",
)?;
Ok(RedisStreamEntry { id, fields })
}
fn parse_string_map(
value: &RedisValue,
value: RedisValue,
context: &str,
) -> Result<BTreeMap<String, String>, DataLayerError> {
match value {
@@ -625,19 +712,23 @@ fn parse_string_map(
)));
}
let mut fields = BTreeMap::new();
for pair in values.chunks(2) {
let key = parse_string_value(&pair[0], context)?;
let value = parse_string_value(&pair[1], context)?;
let mut values = values.into_iter();
while let Some(key) = values.next() {
let key = parse_owned_string_value(key, context)?;
let value = parse_owned_string_value(
values.next().expect("validated even number of fields"),
context,
)?;
fields.insert(key, value);
}
Ok(fields)
}
RedisValue::Map(entries) => entries
.iter()
.into_iter()
.map(|(key, value)| {
Ok((
parse_string_value(key, context)?,
parse_string_value(value, context)?,
parse_owned_string_value(key, context)?,
parse_owned_string_value(value, context)?,
))
})
.collect(),
@@ -648,11 +739,11 @@ fn parse_string_map(
}
}
fn parse_string_array(value: &RedisValue, context: &str) -> Result<Vec<String>, DataLayerError> {
fn parse_string_array(value: RedisValue, context: &str) -> Result<Vec<String>, DataLayerError> {
match value {
RedisValue::Array(values) => values
.iter()
.map(|value| parse_string_value(value, context))
.into_iter()
.map(|value| parse_owned_string_value(value, context))
.collect(),
RedisValue::Nil => Ok(Vec::new()),
_ => Err(DataLayerError::UnexpectedValue(format!(
@@ -669,6 +760,14 @@ fn parse_string_value(value: &RedisValue, context: &str) -> Result<String, DataL
})
}
fn parse_owned_string_value(value: RedisValue, context: &str) -> Result<String, DataLayerError> {
from_owned_redis_value::<String>(value).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"{context} was not a string-compatible redis value: {err}"
))
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
struct RedisXInfoGroupStats {
pending: Option<u64>,
@@ -788,19 +887,60 @@ fn parse_u64_value(value: &RedisValue, context: &str) -> Result<u64, DataLayerEr
})
}
#[cfg(test)]
#[path = "stream_owned_reply_tests.rs"]
mod owned_reply_tests;
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::{
parse_reclaim_result, parse_stream_read_entries, parse_xinfo_group_stats,
parse_xpending_oldest_idle_ms, redis_stream_stats_missing_stream,
parse_reclaim_result, parse_stream_read_entries, parse_transfer_result,
parse_xinfo_group_stats, parse_xpending_oldest_idle_ms, redis_stream_stats_missing_stream,
redis_stream_stats_missing_stream_or_group, validate_consumer, validate_group,
validate_stream_name, validate_stream_position, RedisConsumerName, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunnerConfig,
};
use redis::Value as RedisValue;
#[test]
fn pending_transfer_result_requires_complete_consistent_outcome() {
assert!(matches!(
parse_transfer_result(RedisValue::Array(vec![
RedisValue::Int(0),
RedisValue::BulkString(Vec::new()),
RedisValue::Int(0),
RedisValue::Int(0),
])),
Ok(crate::RuntimeQueueTransferOutcome::NotPending)
));
for values in [
Vec::new(),
vec![RedisValue::Int(0)],
vec![
RedisValue::Int(0),
RedisValue::BulkString(b"1-0".to_vec()),
RedisValue::Int(0),
RedisValue::Int(0),
],
vec![
RedisValue::Int(1),
RedisValue::BulkString(b"1-0".to_vec()),
RedisValue::Int(0),
RedisValue::Int(1),
],
vec![
RedisValue::Int(1),
RedisValue::BulkString(b"1-0".to_vec()),
RedisValue::Int(1),
RedisValue::Int(2),
],
] {
assert!(parse_transfer_result(RedisValue::Array(values)).is_err());
}
}
#[test]
fn validates_stream_runner_config() {
assert!(RedisStreamRunnerConfig {
@@ -0,0 +1,410 @@
use std::collections::BTreeMap;
use redis::streams::StreamReadReply;
use redis::{from_redis_value, Value as RedisValue, VerbatimFormat};
use super::{parse_reclaim_result, parse_stream_read_entries, RedisStreamEntry};
use crate::DataLayerError;
fn bulk(text: &str) -> RedisValue {
RedisValue::BulkString(text.as_bytes().to_vec())
}
fn fields_reply(fields: Vec<(RedisValue, RedisValue)>, resp3: bool) -> RedisValue {
if resp3 {
RedisValue::Map(fields)
} else {
RedisValue::Array(
fields
.into_iter()
.flat_map(|(key, value)| [key, value])
.collect(),
)
}
}
fn read_reply(id: RedisValue, fields: RedisValue, resp3: bool) -> RedisValue {
let entries = RedisValue::Array(vec![RedisValue::Array(vec![id, fields])]);
if resp3 {
RedisValue::Map(vec![(bulk("usage:events"), entries)])
} else {
RedisValue::Array(vec![RedisValue::Array(vec![bulk("usage:events"), entries])])
}
}
fn reclaim_reply(id: RedisValue, fields: RedisValue) -> RedisValue {
RedisValue::Array(vec![
bulk("0-0"),
RedisValue::Array(vec![RedisValue::Array(vec![id, fields])]),
RedisValue::Nil,
])
}
fn original_read_parser(value: &RedisValue) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
let reply = from_redis_value::<StreamReadReply>(value).map_err(crate::error::redis_error)?;
Ok(reply
.keys
.into_iter()
.flat_map(|key| key.ids)
.map(|entry| RedisStreamEntry {
id: entry.id,
fields: entry
.map
.into_iter()
.filter_map(|(field, value)| {
from_redis_value::<String>(&value)
.ok()
.map(|value| (field, value))
})
.collect(),
})
.collect())
}
#[test]
fn owned_read_moves_large_payload_and_id_buffers_in_resp2_and_resp3() {
for resp3 in [false, true] {
let mut payload = Vec::with_capacity(512 * 1024);
payload.extend_from_slice(b"{\"message\":\"");
payload.resize(256 * 1024, b'x');
payload.extend_from_slice(b"\",\"cache_read_input_tokens\":0}\n");
let expected = payload.clone();
let pointer = payload.as_ptr();
let capacity = payload.capacity();
let mut id = Vec::with_capacity(64);
id.extend_from_slice(b"1710000000000-0");
let id_pointer = id.as_ptr();
let id_capacity = id.capacity();
let fields = fields_reply(
vec![(bulk("payload"), RedisValue::BulkString(payload))],
resp3,
);
let parsed =
parse_stream_read_entries(read_reply(RedisValue::BulkString(id), fields, resp3))
.expect("read reply");
assert_eq!(parsed.len(), 1);
let body = &parsed[0].fields["payload"];
assert_eq!(body.as_bytes(), expected);
assert_eq!(
body.as_ptr(),
pointer,
"the RESP payload allocation must be reused"
);
assert_eq!(body.capacity(), capacity);
assert_eq!(parsed[0].id.as_ptr(), id_pointer);
assert_eq!(parsed[0].id.capacity(), id_capacity);
}
}
#[test]
fn owned_reclaim_moves_payload_and_all_id_buffers() {
for resp3 in [false, true] {
let mut payload = Vec::with_capacity(512 * 1024);
payload.resize(256 * 1024, b'x');
let pointer = payload.as_ptr();
let capacity = payload.capacity();
let mut id = Vec::with_capacity(64);
id.extend_from_slice(b"1710000000000-0");
let id_pointer = id.as_ptr();
let mut next_id = Vec::with_capacity(64);
next_id.extend_from_slice(b"1710000000001-0");
let next_pointer = next_id.as_ptr();
let mut deleted_id = Vec::with_capacity(64);
deleted_id.extend_from_slice(b"1709999999999-0");
let deleted_pointer = deleted_id.as_ptr();
let fields = fields_reply(
vec![(bulk("payload"), RedisValue::BulkString(payload))],
resp3,
);
let parsed = parse_reclaim_result(RedisValue::Array(vec![
RedisValue::BulkString(next_id),
RedisValue::Array(vec![RedisValue::Array(vec![
RedisValue::BulkString(id),
fields,
])]),
RedisValue::Array(vec![RedisValue::BulkString(deleted_id)]),
]))
.expect("reclaim reply");
let body = &parsed.entries[0].fields["payload"];
assert_eq!(body.len(), 256 * 1024);
assert!(body.bytes().all(|byte| byte == b'x'));
assert_eq!(body.as_ptr(), pointer);
assert_eq!(body.capacity(), capacity);
assert_eq!(parsed.entries[0].id.as_ptr(), id_pointer);
assert_eq!(parsed.next_start_id.as_ptr(), next_pointer);
assert_eq!(parsed.deleted_ids[0].as_ptr(), deleted_pointer);
}
}
#[test]
fn owned_read_matches_existing_redis_decoder_for_supported_reply_shapes() {
let mut replies = vec![
RedisValue::Nil,
RedisValue::Array(vec![]),
RedisValue::Map(vec![]),
];
for resp3 in [false, true] {
for fields in [
RedisValue::Nil,
fields_reply(
vec![(
bulk("payload"),
bulk("{ \"text\": \"caf\u{00e9}\", \"n\": 0 }\n"),
)],
resp3,
),
fields_reply(
vec![
(bulk("payload"), RedisValue::Nil),
(bulk("count"), RedisValue::Int(0)),
],
resp3,
),
fields_reply(
vec![(bulk("payload"), RedisValue::BulkString(vec![0xff]))],
resp3,
),
fields_reply(
vec![
(bulk("payload"), bulk("first")),
(bulk("payload"), bulk("last")),
(bulk("invalid"), RedisValue::Boolean(false)),
],
resp3,
),
fields_reply(
vec![(
bulk("payload"),
RedisValue::Attribute {
data: Box::new(bulk("annotated payload")),
attributes: vec![(bulk("encoding"), bulk("utf8"))],
},
)],
resp3,
),
] {
replies.push(read_reply(bulk("1-0"), fields, resp3));
}
replies.push(read_reply(RedisValue::Int(42), RedisValue::Nil, resp3));
}
for reply in replies {
let expected = original_read_parser(&reply).expect("baseline reply");
assert_eq!(
parse_stream_read_entries(reply).expect("owned reply"),
expected
);
}
}
#[test]
fn owned_read_preserves_duplicate_overwrite_before_value_filtering() {
for resp3 in [false, true] {
let reply = read_reply(
bulk("1-0"),
fields_reply(
vec![
(bulk("payload"), bulk("valid earlier payload")),
(bulk("payload"), RedisValue::BulkString(vec![0xff])),
(bulk("retry"), bulk("first")),
(bulk("retry"), bulk("last")),
],
resp3,
),
resp3,
);
let parsed = parse_stream_read_entries(reply)
.expect("invalid values are filtered after deduplication");
assert_eq!(
parsed[0].fields,
BTreeMap::from([("retry".to_string(), "last".to_string())])
);
}
}
#[test]
fn owned_read_keeps_invalid_id_key_and_shape_errors_in_redis_category() {
for reply in [
RedisValue::Int(7),
RedisValue::Array(vec![RedisValue::Array(vec![bulk("stream")])]),
read_reply(RedisValue::BulkString(vec![0xff]), RedisValue::Nil, false),
read_reply(bulk("1-0"), RedisValue::Array(vec![bulk("orphan")]), false),
read_reply(
bulk("1-0"),
fields_reply(
vec![(RedisValue::BulkString(vec![0xff]), bulk("value"))],
true,
),
true,
),
] {
assert!(matches!(
original_read_parser(&reply),
Err(DataLayerError::Redis(_))
));
assert!(matches!(
parse_stream_read_entries(reply),
Err(DataLayerError::Redis(_))
));
}
}
#[test]
fn owned_reclaim_preserves_string_types_nil_and_duplicate_fields() {
for resp3 in [false, true] {
let fields = fields_reply(
vec![
(bulk("payload"), bulk("first")),
(bulk("payload"), bulk("{\"text\":\"caf\u{00e9}\"}\n")),
(bulk("zero"), RedisValue::Int(0)),
(bulk("double"), RedisValue::Double(1.5)),
(
bulk("simple"),
RedisValue::SimpleString("simple".to_string()),
),
(bulk("okay"), RedisValue::Okay),
(
bulk("verbatim"),
RedisValue::VerbatimString {
format: VerbatimFormat::Text,
text: "verbatim".to_string(),
},
),
(
bulk("attribute"),
RedisValue::Attribute {
data: Box::new(bulk("annotated")),
attributes: vec![],
},
),
],
resp3,
);
let parsed =
parse_reclaim_result(reclaim_reply(bulk("1-0"), fields)).expect("reclaim reply");
assert_eq!(
parsed.entries[0].fields,
BTreeMap::from([
(
"payload".to_string(),
"{\"text\":\"caf\u{00e9}\"}\n".to_string()
),
("zero".to_string(), "0".to_string()),
("double".to_string(), "1.5".to_string()),
("simple".to_string(), "simple".to_string()),
("okay".to_string(), "OK".to_string()),
("verbatim".to_string(), "verbatim".to_string()),
("attribute".to_string(), "annotated".to_string()),
])
);
assert!(parsed.deleted_ids.is_empty());
}
let parsed =
parse_reclaim_result(reclaim_reply(bulk("1-0"), RedisValue::Nil)).expect("nil fields");
assert!(parsed.entries[0].fields.is_empty());
let parsed = parse_reclaim_result(RedisValue::Array(vec![bulk("0-0"), RedisValue::Nil]))
.expect("nil entries");
assert!(parsed.entries.is_empty());
assert!(parsed.deleted_ids.is_empty());
}
#[test]
fn owned_reclaim_preserves_strict_validation_and_error_context() {
for (reply, context) in [
(
RedisValue::Nil,
"redis xautoclaim returned non-array payload",
),
(
RedisValue::Array(vec![bulk("0-0")]),
"redis xautoclaim returned 1 top-level fields",
),
(
RedisValue::Array(vec![RedisValue::Nil, RedisValue::Nil]),
"redis xautoclaim next_start_id",
),
(
RedisValue::Array(vec![bulk("0-0"), RedisValue::Int(1)]),
"redis xautoclaim entries payload was not an array",
),
(
RedisValue::Array(vec![bulk("0-0"), RedisValue::Array(vec![RedisValue::Nil])]),
"redis xautoclaim entry was not an array",
),
(
RedisValue::Array(vec![
bulk("0-0"),
RedisValue::Array(vec![RedisValue::Array(vec![])]),
]),
"redis xautoclaim entry had 0 fields",
),
(
reclaim_reply(RedisValue::BulkString(vec![0xff]), RedisValue::Nil),
"redis xautoclaim entry id",
),
(
reclaim_reply(bulk("1-0"), RedisValue::Array(vec![bulk("orphan")])),
"redis xautoclaim entry fields expected an even number",
),
(
reclaim_reply(bulk("1-0"), RedisValue::Int(1)),
"redis xautoclaim entry fields expected a redis array/map payload",
),
(
reclaim_reply(
bulk("1-0"),
fields_reply(
vec![(bulk("payload"), RedisValue::BulkString(vec![0xff]))],
false,
),
),
"redis xautoclaim entry fields was not a string-compatible",
),
(
reclaim_reply(
bulk("1-0"),
fields_reply(
vec![(RedisValue::BulkString(vec![0xff]), bulk("value"))],
true,
),
),
"redis xautoclaim entry fields was not a string-compatible",
),
(
RedisValue::Array(vec![bulk("0-0"), RedisValue::Nil, RedisValue::Int(1)]),
"redis xautoclaim deleted_ids expected a redis array payload",
),
(
RedisValue::Array(vec![
bulk("0-0"),
RedisValue::Nil,
RedisValue::Array(vec![RedisValue::Nil]),
]),
"redis xautoclaim deleted_ids was not a string-compatible",
),
] {
let error = parse_reclaim_result(reply).expect_err("invalid reclaim reply");
let DataLayerError::UnexpectedValue(message) = error else {
panic!("reclaim parse error must keep its classification: {error}");
};
assert!(
message.starts_with(context),
"expected {context}, got {message}"
);
}
// Unlike read-group's filter, reclaim has always rejected an invalid value even if
// a later duplicate would overwrite it. Keep that validation order.
let reply = reclaim_reply(
bulk("1-0"),
fields_reply(
vec![
(bulk("payload"), RedisValue::Nil),
(bulk("payload"), bulk("later valid payload")),
],
false,
),
);
assert!(matches!(
parse_reclaim_result(reply),
Err(DataLayerError::UnexpectedValue(_))
));
}
@@ -0,0 +1,860 @@
use super::*;
use std::collections::BTreeSet;
use std::future::{poll_fn, Future};
use std::pin::Pin;
use std::task::Poll;
type TestConnection = ::redis::aio::MultiplexedConnection;
type ReadResult = Result<Vec<RuntimeQueueEntry>, DataLayerError>;
const TEST_GROUP: &str = "receive-workers";
const OWNER_BLOCK_MS: u64 = 60_000;
fn receive_lane_count() -> usize {
std::thread::available_parallelism()
.map(|value| value.get())
.unwrap_or(4)
.clamp(4, 16)
}
async fn receive_runtime(
protocol: &str,
command_timeout_ms: u64,
) -> Option<(TestRedisServer, RuntimeState, TestConnection)> {
let Some(server) = TestRedisServer::start().await else {
eprintln!(
"stream receive {protocol} skipped: isolated Redis fixture unavailable; check AETHER_REDIS_SERVER_BIN"
);
return None;
};
let mut admin = ::redis::Client::open(server.redis_url.clone())
.expect("test admin client")
.get_multiplexed_async_connection()
.await
.expect("test admin connection");
::redis::cmd("ACL")
.arg("SETUSER")
.arg("stream-reader")
.arg("on")
.arg(">stream-test-password")
.arg("~*")
.arg("+@all")
.query_async::<()>(&mut admin)
.await
.expect("test stream user");
let runtime = RuntimeState::redis_with_blocking_stream_lanes(
RedisClientConfig {
url: format!(
"redis://stream-reader:[email protected]:{}/7?protocol={protocol}",
server.port
),
key_prefix: Some(format!("receive-{protocol}")),
},
Some(command_timeout_ms),
Some(4),
)
.await
.expect("authenticated receive runtime in database 7");
::redis::cmd("SELECT")
.arg(7)
.query_async::<()>(&mut admin)
.await
.expect("admin selects test database");
eprintln!(
"stream receive fixture ready: protocol={protocol} db=7 port={} authenticated=true",
server.port
);
Some((server, runtime, admin))
}
async fn receive_group(runtime: &RuntimeState, stream: &str) {
RuntimeQueueStore::ensure_consumer_group(runtime, stream, TEST_GROUP, "0-0")
.await
.expect("receive consumer group");
}
fn receive_fields(sequence: usize) -> BTreeMap<String, String> {
BTreeMap::from([
(
"payload".to_string(),
format!("record-{sequence}\r\n\"quoted\"\\\u{4e2d}\u{6587}"),
),
("sequence".to_string(), sequence.to_string()),
("legacy_field".to_string(), "preserve exactly".to_string()),
])
}
async fn append_receive(runtime: &RuntimeState, stream: &str, sequence: usize) -> String {
RuntimeQueueStore::append_fields_with_maxlen(runtime, stream, &receive_fields(sequence), None)
.await
.expect("append receive entry")
}
async fn client_rows(admin: &mut TestConnection) -> Vec<BTreeMap<String, String>> {
let value = ::redis::cmd("CLIENT")
.arg("LIST")
.query_async::<String>(admin)
.await
.expect("Redis client list");
value
.lines()
.map(|line| {
line.split_whitespace()
.filter_map(|field| field.split_once('='))
.map(|(key, value)| (key.to_string(), value.to_string()))
.collect()
})
.collect()
}
fn blocked_rows(rows: &[BTreeMap<String, String>]) -> Vec<&BTreeMap<String, String>> {
rows.iter()
.filter(|row| row.get("flags").is_some_and(|flags| flags.contains('b')))
.collect()
}
async fn wait_for_blocked(admin: &mut TestConnection, expected: usize) {
tokio::time::timeout(Duration::from_secs(10), async {
loop {
if blocked_rows(&client_rows(admin).await).len() == expected {
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("Redis blocked clients reach expected count");
}
fn spawn_owner(
owners: &mut tokio::task::JoinSet<ReadResult>,
runtime: &RuntimeState,
stream: &'static str,
index: usize,
) {
let runtime = runtime.clone();
owners.spawn(async move {
RuntimeQueueStore::read_group(
&runtime,
stream,
TEST_GROUP,
&format!("owner-{index}"),
1,
Some(OWNER_BLOCK_MS),
)
.await
});
}
async fn abort_owners(owners: &mut tokio::task::JoinSet<ReadResult>) {
owners.abort_all();
while let Some(result) = owners.join_next().await {
assert!(result.expect_err("owner must be cancelled").is_cancelled());
}
}
async fn assert_receive_pending<F: Future + ?Sized>(mut future: Pin<&mut F>) {
poll_fn(|context| {
assert!(future.as_mut().poll(context).is_pending());
Poll::Ready(())
})
.await;
}
async fn pending_consumers(
admin: &mut TestConnection,
stream: &str,
) -> Vec<(String, String, u64, u64)> {
::redis::cmd("XPENDING")
.arg(stream)
.arg(TEST_GROUP)
.arg("-")
.arg("+")
.arg(100)
.query_async(admin)
.await
.expect("pending entry ownership")
}
#[tokio::test]
async fn redis_stream_receive_full_pool_waits_without_sending_and_cancellation_preserves_pel() {
for protocol in ["resp2", "resp3"] {
let Some((_server, runtime, mut admin)) = receive_runtime(protocol, 5_000).await else {
return;
};
let stream = "receive:blocked";
let fast_stream = "receive:nonblocking";
receive_group(&runtime, stream).await;
receive_group(&runtime, fast_stream).await;
let lanes = receive_lane_count();
let mut owners = tokio::task::JoinSet::new();
for index in 0..lanes {
spawn_owner(&mut owners, &runtime, stream, index);
}
wait_for_blocked(&mut admin, lanes).await;
let before = client_rows(&mut admin).await;
let owner_input_bytes = blocked_rows(&before)
.into_iter()
.map(|row| (row["id"].clone(), row.get("tot-net-in").cloned()))
.collect::<BTreeMap<_, _>>();
let mut waiter = Box::pin(RuntimeQueueStore::read_group(
&runtime,
stream,
TEST_GROUP,
"cancelled-waiter",
1,
Some(OWNER_BLOCK_MS),
));
assert_receive_pending(waiter.as_mut()).await;
// Complete unrelated round trips while the waiter remains polled and the owners block.
tokio::time::timeout(Duration::from_secs(5), async {
runtime.kv_set("receive-fast", "ready", None).await.unwrap();
assert_eq!(
runtime.kv_get("receive-fast").await.unwrap().as_deref(),
Some("ready")
);
let expected_id = append_receive(&runtime, fast_stream, 7).await;
let entries = RuntimeQueueStore::read_group(
&runtime,
fast_stream,
TEST_GROUP,
"nonblocking-reader",
1,
None,
)
.await
.unwrap();
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].id, expected_id);
assert_eq!(entries[0].fields, receive_fields(7));
})
.await
.expect("full blocking pool must not delay fast or nonblocking stream lanes");
for _ in 0..3 {
tokio::task::yield_now().await;
let rows = client_rows(&mut admin).await;
let blocked = blocked_rows(&rows);
assert_eq!(blocked.len(), lanes);
for row in blocked {
assert_eq!(
row.get("tot-net-in"),
owner_input_bytes[&row["id"]].as_ref()
);
assert_eq!(row["qbuf"], "0", "waiter must not be sent behind a BLOCK");
}
assert_receive_pending(waiter.as_mut()).await;
}
drop(waiter);
assert_eq!(blocked_rows(&client_rows(&mut admin).await).len(), lanes);
assert!(pending_consumers(&mut admin, stream).await.is_empty());
abort_owners(&mut owners).await;
wait_for_blocked(&mut admin, 0).await;
let expected_id = append_receive(&runtime, stream, 8).await;
let entries = RuntimeQueueStore::read_group(
&runtime,
stream,
TEST_GROUP,
"replacement-reader",
1,
Some(100),
)
.await
.expect("replacement blocking connection");
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].id, expected_id);
assert_eq!(entries[0].fields, receive_fields(8));
let pending = pending_consumers(&mut admin, stream).await;
assert_eq!(pending.len(), 1);
assert_eq!(pending[0].0, expected_id);
assert_eq!(pending[0].1, "replacement-reader");
::redis::cmd("SELECT")
.arg(0)
.query_async::<()>(&mut admin)
.await
.unwrap();
let other_database_len = ::redis::cmd("XLEN")
.arg(stream)
.query_async::<u64>(&mut admin)
.await
.unwrap();
assert_eq!(
other_database_len, 0,
"AUTH and SELECT must survive replacement connections"
);
}
}
#[tokio::test]
async fn redis_stream_receive_fast_consumer_reuses_free_lane_while_other_lanes_block() {
let Some((_server, runtime, mut admin)) = receive_runtime("resp2", 2_000).await else {
return;
};
let slow_stream = "receive:slow";
let fast_stream = "receive:ready";
receive_group(&runtime, slow_stream).await;
receive_group(&runtime, fast_stream).await;
let lanes = receive_lane_count();
let mut owners = tokio::task::JoinSet::new();
for index in 0..lanes - 1 {
spawn_owner(&mut owners, &runtime, slow_stream, index);
}
wait_for_blocked(&mut admin, lanes - 1).await;
let mut reader_connection_id = None;
for sequence in 0..lanes * 2 {
let expected_id = append_receive(&runtime, fast_stream, sequence).await;
let entries = RuntimeQueueStore::read_group(
&runtime,
fast_stream,
TEST_GROUP,
"fast-reader",
1,
Some(100),
)
.await
.expect("a free lane must remain reusable instead of rotating into a blocked lane");
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].id, expected_id);
assert_eq!(entries[0].fields, receive_fields(sequence));
let rows = client_rows(&mut admin).await;
let idle_readers = rows
.iter()
.filter(|row| {
row.get("cmd").map(String::as_str) == Some("xreadgroup")
&& row.get("flags").is_some_and(|flags| !flags.contains('b'))
})
.collect::<Vec<_>>();
assert_eq!(idle_readers.len(), 1);
let current_id = &idle_readers[0]["id"];
if let Some(previous_id) = reader_connection_id.as_ref() {
assert_eq!(
current_id, previous_id,
"successful reads must reuse the same free connection"
);
} else {
reader_connection_id = Some(current_id.clone());
}
}
assert_eq!(
blocked_rows(&client_rows(&mut admin).await).len(),
lanes - 1
);
abort_owners(&mut owners).await;
wait_for_blocked(&mut admin, 0).await;
}
#[tokio::test]
async fn redis_stream_receive_timeout_discards_inflight_connection_before_reuse() {
let Some((_server, runtime, mut admin)) = receive_runtime("resp3", 1_000).await else {
return;
};
let stream = "receive:timeout";
receive_group(&runtime, stream).await;
let initial_connection_ids = client_rows(&mut admin)
.await
.into_iter()
.map(|row| row["id"].clone())
.collect::<BTreeSet<_>>();
// XREADGROUP is a write command; pausing writes makes its network response exceed the
// normal BLOCK-plus-grace timeout while read-only CLIENT diagnostics remain available.
::redis::cmd("CLIENT")
.arg("PAUSE")
.arg(30_000)
.arg("WRITE")
.query_async::<()>(&mut admin)
.await
.expect("pause test Redis writes");
let result = tokio::time::timeout(
Duration::from_secs(10),
RuntimeQueueStore::read_group(
&runtime,
stream,
TEST_GROUP,
"timed-out-reader",
1,
Some(100),
),
)
.await
.expect("read reaches its configured command deadline");
assert!(matches!(result, Err(DataLayerError::TimedOut(_))));
tokio::time::timeout(Duration::from_secs(10), async {
loop {
let connection_ids = client_rows(&mut admin)
.await
.into_iter()
.map(|row| row["id"].clone())
.collect::<BTreeSet<_>>();
if connection_ids.len() + 1 == initial_connection_ids.len()
&& connection_ids.is_subset(&initial_connection_ids)
{
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("timed-out connection must disconnect while its old command is still paused");
::redis::cmd("CLIENT")
.arg("UNPAUSE")
.query_async::<()>(&mut admin)
.await
.unwrap();
let expected_id = append_receive(&runtime, stream, 9).await;
let entries =
RuntimeQueueStore::read_group(&runtime, stream, TEST_GROUP, "after-timeout", 1, Some(100))
.await
.expect("replacement read after timeout");
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].id, expected_id);
assert_eq!(entries[0].fields, receive_fields(9));
let pending = pending_consumers(&mut admin, stream).await;
assert_eq!(pending.len(), 1);
assert_eq!(pending[0].0, expected_id);
assert_eq!(pending[0].1, "after-timeout");
}
async fn set_pending_idle(admin: &mut TestConnection, stream: &str, ids: &[String]) {
let mut command = ::redis::cmd("XCLAIM");
command
.arg(stream)
.arg(TEST_GROUP)
.arg("initial-reader")
.arg(0);
for id in ids {
command.arg(id);
}
command.arg("IDLE").arg(120_000).arg("JUSTID");
let changed = command
.query_async::<Vec<String>>(admin)
.await
.expect("set pending idle");
assert_eq!(changed, ids);
}
#[tokio::test]
async fn redis_stream_receive_reclaim_pages_advance_past_fresh_prefix_and_deleted_entries() {
for protocol in ["resp2", "resp3"] {
let Some((_server, runtime, mut admin)) = receive_runtime(protocol, 5_000).await else {
return;
};
let stream = "receive:reclaim-pages";
receive_group(&runtime, stream).await;
let mut ids = Vec::new();
for sequence in 0..28 {
ids.push(append_receive(&runtime, stream, sequence).await);
}
let entries =
RuntimeQueueStore::read_group(&runtime, stream, TEST_GROUP, "initial-reader", 28, None)
.await
.expect("seed pending entries");
assert_eq!(entries.len(), ids.len());
set_pending_idle(&mut admin, stream, &ids[25..]).await;
assert_eq!(
RuntimeQueueStore::delete(&runtime, stream, &ids[26..27])
.await
.unwrap(),
1
);
let config = RuntimeQueueReclaimConfig {
min_idle_ms: 60_000,
count: 2,
};
let first = RuntimeQueueStore::claim_stale_page(
&runtime,
stream,
TEST_GROUP,
"reclaim-reader",
"0-0",
config,
)
.await
.expect("first reclaim page");
assert!(
first.entries.is_empty(),
"fresh prefix exceeds COUNT * 10 scan budget"
);
assert!(first.deleted_ids.is_empty());
assert_ne!(
first.next_start_id, "0-0",
"empty page must preserve continuation"
);
let mut cursor = first.next_start_id;
let mut reclaimed = BTreeMap::new();
let mut deleted = BTreeSet::new();
for _ in 0..8 {
let page = RuntimeQueueStore::claim_stale_page(
&runtime,
stream,
TEST_GROUP,
"reclaim-reader",
&cursor,
config,
)
.await
.expect("continued reclaim page");
assert!(page.entries.len() <= config.count);
for entry in page.entries {
assert!(reclaimed.insert(entry.id, entry.fields).is_none());
}
deleted.extend(page.deleted_ids);
cursor = page.next_start_id;
if cursor == "0-0" {
break;
}
}
assert_eq!(cursor, "0-0", "scan must eventually wrap");
assert_eq!(
reclaimed,
BTreeMap::from([
(ids[25].clone(), receive_fields(25)),
(ids[27].clone(), receive_fields(27))
])
);
assert_eq!(deleted, BTreeSet::from([ids[26].clone()]));
let pending = pending_consumers(&mut admin, stream).await;
assert_eq!(pending.len(), 27);
assert!(!pending.iter().any(|entry| entry.0 == ids[26]));
assert!(pending
.iter()
.filter(|entry| entry.1 == "reclaim-reader")
.all(|entry| entry.0 == ids[25] || entry.0 == ids[27]));
set_pending_idle(&mut admin, stream, &ids[..1]).await;
let restarted = RuntimeQueueStore::claim_stale_page(
&runtime,
stream,
TEST_GROUP,
"rescan-reader",
"0-0",
config,
)
.await
.expect("restart scan after cursor wraps");
assert_eq!(restarted.entries.len(), 1);
assert_eq!(restarted.entries[0].id, ids[0]);
assert_eq!(restarted.entries[0].fields, receive_fields(0));
assert!(restarted.deleted_ids.is_empty());
}
}
async fn interrupted_reclaim_preserves_pending(protocol: &str, cancel: bool) {
let command_timeout_ms = if cancel { 10_000 } else { 1_000 };
let Some((_server, runtime, mut admin)) = receive_runtime(protocol, command_timeout_ms).await
else {
return;
};
let stream = "receive:interrupted-reclaim";
receive_group(&runtime, stream).await;
let mut ids = Vec::new();
for sequence in 0..3 {
ids.push(append_receive(&runtime, stream, sequence).await);
}
assert_eq!(
RuntimeQueueStore::read_group(&runtime, stream, TEST_GROUP, "initial-reader", 3, None,)
.await
.unwrap()
.len(),
3
);
set_pending_idle(&mut admin, stream, &ids).await;
RuntimeQueueStore::delete(&runtime, stream, &ids[1..2])
.await
.unwrap();
let before = client_rows(&mut admin).await;
let initial_ids = before
.iter()
.map(|row| row["id"].clone())
.collect::<BTreeSet<_>>();
let input_bytes = before
.iter()
.filter(|row| row.get("user").map(String::as_str) == Some("stream-reader"))
.map(|row| {
(
row["id"].clone(),
row.get("tot-net-in")
.and_then(|value| value.parse::<u64>().ok()),
)
})
.collect::<BTreeMap<_, _>>();
::redis::cmd("CLIENT")
.arg("PAUSE")
.arg(30_000)
.arg("WRITE")
.query_async::<()>(&mut admin)
.await
.unwrap();
let config = RuntimeQueueReclaimConfig {
min_idle_ms: 60_000,
count: 1,
};
let claim_runtime = runtime.clone();
let claim = tokio::spawn(async move {
RuntimeQueueStore::claim_stale_page(
&claim_runtime,
stream,
TEST_GROUP,
"interrupted-reader",
"0-0",
config,
)
.await
});
if cancel {
tokio::time::timeout(Duration::from_secs(5), async {
loop {
let sent = client_rows(&mut admin).await.iter().any(|row| {
let Some(previous) = input_bytes.get(&row["id"]) else {
return false;
};
match (
previous,
row.get("tot-net-in")
.and_then(|value| value.parse::<u64>().ok()),
) {
(Some(previous), Some(current)) => current > *previous,
_ => row.get("cmd").map(String::as_str) == Some("xautoclaim"),
}
});
if sent {
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("claim reaches Redis before caller cancellation");
claim.abort();
assert!(claim.await.expect_err("cancelled reclaim").is_cancelled());
} else {
let result = tokio::time::timeout(Duration::from_secs(10), claim)
.await
.expect("reclaim command deadline")
.expect("reclaim task");
assert!(matches!(result, Err(DataLayerError::TimedOut(_))));
}
tokio::time::timeout(Duration::from_secs(10), async {
loop {
let remaining = client_rows(&mut admin)
.await
.into_iter()
.map(|row| row["id"].clone())
.collect::<BTreeSet<_>>();
if remaining.len() + 1 == initial_ids.len() && remaining.is_subset(&initial_ids) {
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("interrupted claim closes its socket before writes resume");
let pending = pending_consumers(&mut admin, stream).await;
assert_eq!(pending.len(), 3);
assert!(pending.iter().all(|entry| entry.1 == "initial-reader"));
::redis::cmd("CLIENT")
.arg("UNPAUSE")
.query_async::<()>(&mut admin)
.await
.unwrap();
let mut cursor = "0-0".to_string();
let mut recovered = BTreeMap::new();
let mut deleted = BTreeSet::new();
let mut successful_claims = 0;
for _ in 0..5 {
let page = RuntimeQueueStore::claim_stale_page(
&runtime,
stream,
TEST_GROUP,
"recovery-reader",
&cursor,
config,
)
.await
.expect("reclaim recovers after interrupted connection");
successful_claims += 1;
for entry in page.entries {
assert!(recovered.insert(entry.id, entry.fields).is_none());
}
deleted.extend(page.deleted_ids);
cursor = page.next_start_id;
if cursor == "0-0" {
break;
}
}
assert_eq!(cursor, "0-0");
assert_eq!(
recovered,
BTreeMap::from([
(ids[0].clone(), receive_fields(0)),
(ids[2].clone(), receive_fields(2))
])
);
assert_eq!(deleted, BTreeSet::from([ids[1].clone()]));
let pending = pending_consumers(&mut admin, stream).await;
assert_eq!(pending.len(), 2);
assert!(pending.iter().all(|entry| entry.1 == "recovery-reader"));
let diagnostics = runtime.redis_diagnostics().await.unwrap().unwrap();
let lane = diagnostics
.lanes
.iter()
.find(|lane| lane.lane == "blocking_stream")
.unwrap();
assert_eq!(lane.command_timeouts, u64::from(!cancel));
assert!(lane.command_count >= successful_claims);
}
#[tokio::test]
async fn redis_stream_receive_reclaim_cancellation_closes_connection_and_preserves_pending() {
interrupted_reclaim_preserves_pending("resp2", true).await;
}
#[tokio::test]
async fn redis_stream_receive_reclaim_timeout_closes_connection_and_preserves_pending() {
interrupted_reclaim_preserves_pending("resp3", false).await;
}
#[tokio::test]
async fn redis_stream_receive_reclaim_waits_for_read_lease_and_continues_after_release() {
let Some((_server, runtime, mut admin)) = receive_runtime("resp2", 1_000).await else {
return;
};
let blocked_stream = "receive:reclaim-pool-blocked";
let pending_stream = "receive:reclaim-pool-pending";
receive_group(&runtime, blocked_stream).await;
receive_group(&runtime, pending_stream).await;
let ids = vec![
append_receive(&runtime, pending_stream, 0).await,
append_receive(&runtime, pending_stream, 1).await,
];
assert_eq!(
RuntimeQueueStore::read_group(
&runtime,
pending_stream,
TEST_GROUP,
"initial-reader",
2,
None
)
.await
.unwrap()
.len(),
2
);
set_pending_idle(&mut admin, pending_stream, &ids).await;
let lanes = receive_lane_count();
let mut owners = tokio::task::JoinSet::new();
for index in 0..lanes - 1 {
spawn_owner(&mut owners, &runtime, blocked_stream, index);
}
let owner_runtime = runtime.clone();
let release_owner = owners.spawn(async move {
RuntimeQueueStore::read_group(
&owner_runtime,
blocked_stream,
TEST_GROUP,
"released-owner",
1,
Some(OWNER_BLOCK_MS),
)
.await
});
wait_for_blocked(&mut admin, lanes).await;
let config = RuntimeQueueReclaimConfig {
min_idle_ms: 60_000,
count: 1,
};
let result = tokio::time::timeout(
Duration::from_secs(10),
RuntimeQueueStore::claim_stale_page(
&runtime,
pending_stream,
TEST_GROUP,
"timed-out-waiter",
"0-0",
config,
),
)
.await
.expect("checkout uses the original command deadline");
assert!(matches!(result, Err(DataLayerError::TimedOut(_))));
assert_eq!(blocked_rows(&client_rows(&mut admin).await).len(), lanes);
assert!(pending_consumers(&mut admin, pending_stream)
.await
.iter()
.all(|entry| entry.1 == "initial-reader"));
let mut cancelled_claim = Box::pin(RuntimeQueueStore::claim_stale_page(
&runtime,
pending_stream,
TEST_GROUP,
"cancelled-waiter",
"0-0",
config,
));
assert_receive_pending(cancelled_claim.as_mut()).await;
drop(cancelled_claim);
assert_eq!(blocked_rows(&client_rows(&mut admin).await).len(), lanes);
let mut claim = Box::pin(RuntimeQueueStore::claim_stale_page(
&runtime,
pending_stream,
TEST_GROUP,
"after-release",
"0-0",
config,
));
assert_receive_pending(claim.as_mut()).await;
release_owner.abort();
assert!(owners
.join_next()
.await
.unwrap()
.expect_err("released owner cancelled")
.is_cancelled());
let first = tokio::time::timeout(Duration::from_secs(5), claim)
.await
.expect("claim receives released pool capacity")
.expect("claim succeeds after read releases lease");
assert_eq!(first.entries.len(), 1);
assert_eq!(first.entries[0].id, ids[0]);
assert_eq!(first.entries[0].fields, receive_fields(0));
assert_ne!(first.next_start_id, "0-0");
let next_id = append_receive(&runtime, pending_stream, 2).await;
let next_read = RuntimeQueueStore::read_group(
&runtime,
pending_stream,
TEST_GROUP,
"read-after-claim",
1,
Some(100),
)
.await
.expect("completed claim returns its lease for the next read");
assert_eq!(next_read.len(), 1);
assert_eq!(next_read[0].id, next_id);
let second = RuntimeQueueStore::claim_stale_page(
&runtime,
pending_stream,
TEST_GROUP,
"after-release",
&first.next_start_id,
config,
)
.await
.expect("completed read returns its lease for the next claim");
assert_eq!(second.entries.len(), 1);
assert_eq!(second.entries[0].id, ids[1]);
assert_eq!(second.entries[0].fields, receive_fields(1));
assert_eq!(
blocked_rows(&client_rows(&mut admin).await).len(),
lanes - 1
);
abort_owners(&mut owners).await;
wait_for_blocked(&mut admin, 0).await;
}
@@ -0,0 +1,175 @@
use std::sync::OnceLock;
use super::client::RedisBlockingStreamLease;
use super::runtime::USAGE_LIMIT_CHECK_AND_CONSUME_SCRIPT;
use super::{cmd, RedisConnectionLane, RedisConnectionRouter};
use crate::error::RedisResultExt;
use crate::{DataLayerError, UsageLimitInput};
const COPY_CHUNK_SIZE: i64 = 512;
const MAX_COPY_ATTEMPTS: usize = 8;
const COPY_SCRIPT: &str = include_str!("usage_copy.lua");
const COMMIT_PREFIX: &str = include_str!("usage_copy_commit.lua");
fn usage_args(command: &mut redis::Cmd, input: &UsageLimitInput<'_>) {
command.arg(input.now_unix_ms).arg(input.event_id);
for rule in input.rules {
command
.arg(rule.limit)
.arg(rule.window_seconds)
.arg(rule.retention_seconds);
}
}
pub(super) async fn check_and_consume(
connections: &RedisConnectionRouter,
keys: &[String],
input: &UsageLimitInput<'_>,
) -> Result<Vec<i64>, DataLayerError> {
static SCRIPT: OnceLock<redis::Script> = OnceLock::new();
let script = SCRIPT.get_or_init(|| redis::Script::new(USAGE_LIMIT_CHECK_AND_CONSUME_SCRIPT));
let mut invocation = script.prepare_invoke();
for key in keys {
invocation.key(key);
}
invocation.arg(input.now_unix_ms).arg(input.event_id);
for rule in input.rules {
invocation
.arg(rule.limit)
.arg(rule.window_seconds)
.arg(rule.retention_seconds);
}
let result: Vec<i64> = invocation
.invoke_async(&mut connections.connection(RedisConnectionLane::Fast))
.await
.map_redis_err()?;
if result.first() != Some(&2) {
return Ok(result);
}
let mut lease = connections.usage_cleanup_connection().await?;
for _ in 0..MAX_COPY_ATTEMPTS {
let mut temporary_keys = Vec::new();
let result = copy_and_commit(&mut lease, keys, input, &mut temporary_keys).await;
// EXEC clears WATCH even on a conflict. Errors discard the lease,
// including any pending WATCH/MULTI state.
if let Ok(Some((result, committed))) = result.as_ref() {
// Never add another fallible round trip after admission. EXEC already
// cleared WATCH; the recheck fast path instead discards its watched lease.
if *committed {
lease.recycle();
}
return Ok(result.clone());
} else if result.is_ok() {
lease.query(&cmd("UNWATCH")).await?;
if !temporary_keys.is_empty() {
lease.query(cmd("UNLINK").arg(&temporary_keys)).await?;
}
} else if !temporary_keys.is_empty() {
// A fresh connection cannot accidentally queue cleanup inside a failed MULTI.
let mut connection = connections.connection(RedisConnectionLane::Admin);
let _ = cmd("UNLINK")
.arg(&temporary_keys)
.query_async::<usize>(&mut connection)
.await;
}
result?;
tokio::task::yield_now().await;
}
lease.recycle();
Err(DataLayerError::Redis(
"usage window cleanup conflicted repeatedly; retry the request".to_string(),
))
}
async fn copy_and_commit(
lease: &mut RedisBlockingStreamLease,
keys: &[String],
input: &UsageLimitInput<'_>,
temporary_keys: &mut Vec<String>,
) -> Result<Option<(Vec<i64>, bool)>, DataLayerError> {
lease.query(cmd("WATCH").arg(keys)).await?;
let mut check = cmd("EVAL");
check
.arg(USAGE_LIMIT_CHECK_AND_CONSUME_SCRIPT)
.arg(keys.len())
.arg(keys);
usage_args(&mut check, input);
let plan: Vec<i64> =
redis::from_owned_redis_value(lease.query(&check).await?).map_redis_err()?;
if plan.first() != Some(&2) {
return Ok(Some((plan, false)));
}
if plan.len() < 4 || (plan.len() - 1) % 3 != 0 {
return Err(DataLayerError::UnexpectedValue(
"invalid usage window copy plan".to_string(),
));
}
let nonce = uuid::Uuid::new_v4();
for window in plan[1..].chunks_exact(3) {
let [index, expired, live] = [window[0], window[1], window[2]];
let key = usize::try_from(index - 1)
.ok()
.and_then(|index| keys.get(index))
.filter(|_| expired > 0 && live > 0)
.ok_or_else(|| {
DataLayerError::UnexpectedValue("invalid usage window copy range".to_string())
})?;
let temporary = format!("{key}:__usage_copy:{nonce}");
let mut offset = 0;
while offset < live {
let take = COPY_CHUNK_SIZE.min(live - offset);
let value = lease
.query(
cmd("EVAL")
.arg(COPY_SCRIPT)
.arg(2)
.arg(key)
.arg(&temporary)
.arg(expired + offset)
.arg(expired + offset + take - 1)
.arg(offset),
)
.await?;
let copied: i64 = redis::from_owned_redis_value(value).map_redis_err()?;
if offset == 0 && copied > 0 {
temporary_keys.push(temporary.clone());
}
if copied != take {
return Ok(None);
}
offset += take;
}
}
static COMMIT: OnceLock<String> = OnceLock::new();
let source =
COMMIT.get_or_init(|| format!("{COMMIT_PREFIX}\n{USAGE_LIMIT_CHECK_AND_CONSUME_SCRIPT}"));
let mut commit = cmd("EVAL");
commit
.arg(source)
.arg(keys.len() + temporary_keys.len())
.arg(keys)
.arg(&*temporary_keys)
.arg(keys.len());
usage_args(&mut commit, input);
for window in plan[1..].chunks_exact(3) {
commit.arg(window[0]).arg(window[2]);
}
lease.query(&cmd("MULTI")).await?;
lease.query(&commit).await?;
let replies: Option<Vec<redis::Value>> =
redis::from_owned_redis_value(lease.query(&cmd("EXEC")).await?).map_redis_err()?;
let Some(mut replies) = replies else {
return Ok(None);
};
let reply = replies.pop().ok_or_else(|| {
DataLayerError::UnexpectedValue("empty usage window commit response".to_string())
})?;
let result: Vec<i64> = redis::from_owned_redis_value(reply).map_redis_err()?;
if result.first() == Some(&2) {
return Ok(None);
}
Ok(Some((result, true)))
}
@@ -0,0 +1,22 @@
local source, target = KEYS[1], KEYS[2]
local offset = tonumber(ARGV[3])
if offset == 0 and redis.call('EXISTS', target) ~= 0 then
return redis.error_reply('usage copy temporary key already exists')
end
if offset > 0 and redis.call('ZCARD', target) ~= offset then return -1 end
local rows = redis.call('ZRANGE', source, ARGV[1], ARGV[2], 'WITHSCORES')
if #rows == 0 then return 0 end
if not redis.acl_check_cmd('PEXPIRE', target, 60000)
or not redis.acl_check_cmd('UNLINK', target) then
return redis.error_reply('usage copy temporary key permission denied')
end
local args = {}
for i = 1, #rows, 2 do
args[#args + 1] = rows[i + 1]
args[#args + 1] = rows[i]
end
-- TTL is established in the same command as the first allocation. A cancelled
-- caller or a process crash cannot leave a permanent scratch key behind.
redis.call('ZADD', target, unpack(args))
redis.call('PEXPIRE', target, 60000)
return #rows / 2
@@ -0,0 +1,35 @@
local rule_count = tonumber(table.remove(ARGV, 1))
local swaps = #KEYS - rule_count
local ttls = {}
-- WATCH covers every source, including modifications made by older instances.
-- Check every temporary key and every permission before replacing any source.
for i = 1, swaps do
local position = rule_count * 3 + 2 + (i - 1) * 2
local index = tonumber(ARGV[position + 1])
local live = tonumber(ARGV[position + 2])
local source, target = KEYS[index], KEYS[rule_count + i]
local ttl = redis.call('PTTL', source)
if ttl == 0 or ttl < -1 or redis.call('ZCARD', target) ~= live then return {2} end
if not redis.acl_check_cmd('UNLINK', source)
or not redis.acl_check_cmd('RENAME', target, source)
or not redis.acl_check_cmd('PEXPIRE', target, math.max(1, ttl))
or not redis.acl_check_cmd('PERSIST', target) then
return redis.error_reply('usage copy commit permission denied')
end
ttls[i] = ttl
end
for i = 1, swaps do
local position = rule_count * 3 + 2 + (i - 1) * 2
local index = tonumber(ARGV[position + 1])
local source, target = KEYS[index], KEYS[rule_count + i]
if ttls[i] > 0 then
redis.call('PEXPIRE', target, ttls[i])
else
redis.call('PERSIST', target)
end
redis.call('UNLINK', source)
redis.call('RENAME', target, source)
end
for i = #KEYS, rule_count + 1, -1 do KEYS[i] = nil end
ARGV[rule_count * 3 + 3] = 'inline'
-- The original prune/check/consume script follows in this same EXEC/EVAL.
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,25 @@
/// Maximum number of members examined by one server-side window aggregation.
pub const SCORE_WINDOW_AGGREGATION_MEMBER_LIMIT: usize = 512;
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct ScoreWindowU64Stats {
pub sum: u64,
pub positive_count: u64,
}
impl ScoreWindowU64Stats {
/// Members encode the value after their final colon. Invalid or zero values
/// do not contribute to the positive sample count.
pub fn from_members<'a>(members: impl IntoIterator<Item = &'a str>) -> Self {
let mut stats = Self::default();
for member in members {
let value = member
.rsplit_once(':')
.and_then(|(_, suffix)| suffix.parse::<u64>().ok())
.unwrap_or(0);
stats.sum = stats.sum.saturating_add(value);
stats.positive_count += u64::from(value > 0);
}
stats
}
}