mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 17:37:47 +08:00
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:
@@ -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
@@ -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};
|
||||
|
||||
@@ -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}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user