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

Add usage queue worker autoscaling

Browse Source
This commit is contained in:
elky
2026-06-26 14:02:57 +08:00
parent 6c5e70ccb1
commit 7e9424008f
21 changed files with 1801 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-runtime-state/src/lib.rs
+69 -2
View File
@@ -62,6 +62,7 @@ pub struct RuntimeStateConfig {
pub redis: Option<RedisClientConfig>,
pub memory: MemoryRuntimeStateConfig,
pub command_timeout_ms: Option<u64>,
pub blocking_stream_lanes: Option<usize>,
}
impl Default for RuntimeStateConfig {
@@ -71,6 +72,7 @@ impl Default for RuntimeStateConfig {
redis: None,
memory: MemoryRuntimeStateConfig::default(),
command_timeout_ms: Some(DEFAULT_COMMAND_TIMEOUT_MS),
blocking_stream_lanes: None,
}
}
}
@@ -138,6 +140,11 @@ impl RuntimeStateConfig {
"runtime state command_timeout_ms must be positive".to_string(),
));
}
if matches!(self.blocking_stream_lanes, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"runtime state blocking_stream_lanes must be positive".to_string(),
));
}
Ok(())
}
}
@@ -194,7 +201,12 @@ impl RuntimeState {
let redis = config.redis.clone().ok_or_else(|| {
DataLayerError::InvalidConfiguration("runtime redis config missing".to_string())
})?;
Self::redis(redis, config.command_timeout_ms).await
Self::redis_with_blocking_stream_lanes(
redis,
config.command_timeout_ms,
config.blocking_stream_lanes,
)
.await
}
RuntimeStateBackendMode::Auto => unreachable!("auto resolved above"),
}
@@ -211,10 +223,20 @@ impl RuntimeState {
pub async fn redis(
config: RedisClientConfig,
command_timeout_ms: Option<u64>,
) -> Result<Self, DataLayerError> {
Self::redis_with_blocking_stream_lanes(config, command_timeout_ms, None).await
}
pub async fn redis_with_blocking_stream_lanes(
config: RedisClientConfig,
command_timeout_ms: Option<u64>,
blocking_stream_lanes: Option<usize>,
) -> Result<Self, DataLayerError> {
let factory = redis::RedisClientFactory::new(config)?;
let keyspace = factory.config().keyspace();
let connections = factory.connect_router(command_timeout_ms).await?;
let connections = factory
.connect_router_with_blocking_stream_lanes(command_timeout_ms, blocking_stream_lanes)
.await?;
let runtime = redis::RedisRuntimeRunner::new(
connections.clone(),
keyspace.clone(),
@@ -1597,6 +1619,51 @@ mod tests {
let _ = blocking.await.expect("blocking task join");
}
#[tokio::test]
async fn redis_concurrent_blocking_stream_reads_do_not_share_single_connection() {
let Some(redis) = TestRedisServer::start().await else {
return;
};
let runtime = RuntimeState::redis(
RedisClientConfig {
url: redis.redis_url.clone(),
key_prefix: Some(format!("aether-block-pool-test-{}", std::process::id())),
},
Some(1_000),
)
.await
.expect("runtime should connect");
RuntimeQueueStore::ensure_consumer_group(&runtime, "blocking-stream", "workers", "0-0")
.await
.expect("consumer group");
let mut handles = Vec::new();
for index in 0..4 {
let blocking_runtime = runtime.clone();
handles.push(tokio::spawn(async move {
let consumer = format!("consumer-{index}");
RuntimeQueueStore::read_group(
&blocking_runtime,
"blocking-stream",
"workers",
&consumer,
1,
Some(600),
)
.await
}));
}
for handle in handles {
let result = handle.await.expect("blocking task join");
assert!(
!matches!(result, Err(DataLayerError::TimedOut(_))),
"concurrent blocking stream reads should not queue behind one connection"
);
assert!(result.expect("blocking read should succeed").is_empty());
}
}
#[tokio::test]
async fn redis_connection_manager_recovers_after_restart() {
let Some(mut redis) = TestRedisServer::start().await else {
crates/aether-runtime-state/src/redis/client.rs
+109 -13
View File
@@ -1,7 +1,7 @@
use crate::error::RedisResultExt;
use crate::redis::RedisKeyspace;
use crate::DataLayerError;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tracing::info;
@@ -9,6 +9,10 @@ use tracing::info;
pub(crate) type RedisClient = redis::Client;
pub(crate) type RedisManagedConnection = redis::aio::ConnectionManager;
const DEFAULT_BLOCKING_STREAM_LANES_FALLBACK: usize = 4;
const DEFAULT_BLOCKING_STREAM_LANES_CAP: usize = 16;
const MAX_BLOCKING_STREAM_LANES_CAP: usize = 64;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct RedisClientConfig {
pub url: String,
@@ -57,7 +61,21 @@ impl RedisClientFactory {
&self,
command_timeout_ms: Option<u64>,
) -> Result<RedisConnectionRouter, DataLayerError> {
RedisConnectionRouter::connect(self.connect_lazy()?, command_timeout_ms).await
self.connect_router_with_blocking_stream_lanes(command_timeout_ms, None)
.await
}
pub(crate) async fn connect_router_with_blocking_stream_lanes(
&self,
command_timeout_ms: Option<u64>,
blocking_stream_lanes: Option<usize>,
) -> Result<RedisConnectionRouter, DataLayerError> {
RedisConnectionRouter::connect(
self.connect_lazy()?,
command_timeout_ms,
blocking_stream_lanes,
)
.await
}
}
@@ -84,7 +102,8 @@ impl RedisConnectionLane {
pub(crate) struct RedisConnectionRouter {
fast: RedisManagedConnection,
stream: RedisManagedConnection,
blocking_stream: RedisManagedConnection,
blocking_stream: Arc<Vec<RedisManagedConnection>>,
blocking_stream_next: Arc<AtomicUsize>,
admin: RedisManagedConnection,
metrics: Arc<RedisConnectionMetrics>,
}
@@ -93,6 +112,7 @@ impl std::fmt::Debug for RedisConnectionRouter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RedisConnectionRouter")
.field("lanes", &["fast", "stream", "blocking_stream", "admin"])
.field("blocking_stream_lanes", &self.blocking_stream.len())
.finish()
}
}
@@ -101,6 +121,7 @@ impl RedisConnectionRouter {
pub(crate) async fn connect(
client: RedisClient,
command_timeout_ms: Option<u64>,
blocking_stream_lanes: Option<usize>,
) -> Result<Self, DataLayerError> {
let fast = connect_lane(
&client,
@@ -116,13 +137,9 @@ impl RedisConnectionRouter {
command_timeout_ms,
)
.await?;
let blocking_stream = connect_lane(
&client,
connection_manager_config(command_timeout_ms),
RedisConnectionLane::BlockingStream,
command_timeout_ms,
)
.await?;
let blocking_stream =
connect_blocking_stream_lanes(&client, command_timeout_ms, blocking_stream_lanes)
.await?;
let admin = connect_lane(
&client,
connection_manager_config(command_timeout_ms),
@@ -130,14 +147,17 @@ impl RedisConnectionRouter {
command_timeout_ms,
)
.await?;
let blocking_stream_lanes = blocking_stream.len();
info!(
redis_lanes = "fast,stream,blocking_stream,admin",
redis_blocking_stream_lanes = blocking_stream_lanes,
"runtime redis connection lanes initialized"
);
Ok(Self {
fast,
stream,
blocking_stream,
blocking_stream: Arc::new(blocking_stream),
blocking_stream_next: Arc::new(AtomicUsize::new(0)),
admin,
metrics: Arc::new(RedisConnectionMetrics::default()),
})
@@ -147,7 +167,11 @@ impl RedisConnectionRouter {
match lane {
RedisConnectionLane::Fast => self.fast.clone(),
RedisConnectionLane::Stream => self.stream.clone(),
RedisConnectionLane::BlockingStream => self.blocking_stream.clone(),
RedisConnectionLane::BlockingStream => {
let index = self.blocking_stream_next.fetch_add(1, Ordering::Relaxed)
% self.blocking_stream.len();
self.blocking_stream[index].clone()
}
RedisConnectionLane::Admin => self.admin.clone(),
}
}
@@ -228,6 +252,51 @@ fn connection_manager_config(
config
}
async fn connect_blocking_stream_lanes(
client: &RedisClient,
command_timeout_ms: Option<u64>,
requested_lanes: Option<usize>,
) -> Result<Vec<RedisManagedConnection>, 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?,
);
}
Ok(lanes)
}
fn blocking_stream_lane_count(requested_lanes: Option<usize>) -> Result<usize, DataLayerError> {
if matches!(requested_lanes, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"runtime redis blocking_stream_lanes must be positive".to_string(),
));
}
let default_lanes = default_blocking_stream_lane_count();
Ok(requested_lanes
.map(|lanes| lanes.max(default_lanes))
.unwrap_or(default_lanes)
.clamp(1, MAX_BLOCKING_STREAM_LANES_CAP))
}
fn default_blocking_stream_lane_count() -> usize {
std::thread::available_parallelism()
.map(|value| value.get())
.unwrap_or(DEFAULT_BLOCKING_STREAM_LANES_FALLBACK)
.clamp(
DEFAULT_BLOCKING_STREAM_LANES_FALLBACK,
DEFAULT_BLOCKING_STREAM_LANES_CAP,
)
}
async fn connect_lane(
client: &RedisClient,
config: redis::aio::ConnectionManagerConfig,
@@ -259,7 +328,10 @@ async fn connect_lane(
#[cfg(test)]
mod tests {
use super::{RedisClientConfig, RedisClientFactory};
use super::{
blocking_stream_lane_count, default_blocking_stream_lane_count, RedisClientConfig,
RedisClientFactory, MAX_BLOCKING_STREAM_LANES_CAP,
};
#[test]
fn factory_builds_lazy_client_from_valid_config() {
@@ -274,4 +346,28 @@ mod tests {
.connect_lazy()
.expect("lazy redis client should build");
}
#[test]
fn blocking_stream_lane_count_uses_requested_as_floor() {
let default_lanes = default_blocking_stream_lane_count();
assert_eq!(
blocking_stream_lane_count(None).expect("default lanes"),
default_lanes
);
assert_eq!(
blocking_stream_lane_count(Some(1)).expect("requested below default"),
default_lanes
);
assert_eq!(
blocking_stream_lane_count(Some(default_lanes + 1)).expect("requested above default"),
default_lanes + 1
);
assert_eq!(
blocking_stream_lane_count(Some(MAX_BLOCKING_STREAM_LANES_CAP + 1))
.expect("requested above cap"),
MAX_BLOCKING_STREAM_LANES_CAP
);
assert!(blocking_stream_lane_count(Some(0)).is_err());
}
}
crates/aether-runtime-state/src/redis/stream.rs
+38 -21
View File
@@ -289,27 +289,11 @@ impl RedisStreamRunner {
command.arg("STREAMS").arg(&stream.0).arg(">");
let reply = command
.query_async::<StreamReadReply>(&mut connection)
.query_async::<RedisValue>(&mut connection)
.await
.map_redis_err()?;
Ok(reply
.keys
.into_iter()
.flat_map(|key| key.ids.into_iter())
.map(|id| RedisStreamEntry {
id: id.id,
fields: id
.map
.into_iter()
.filter_map(|(field, value)| {
redis::from_redis_value::<String>(&value)
.ok()
.map(|text| (field, text))
})
.collect(),
})
.collect())
parse_stream_read_entries(reply)
})
.await
}
@@ -455,6 +439,31 @@ fn validate_stream_position(position: &str) -> Result<(), DataLayerError> {
Ok(())
}
fn parse_stream_read_entries(value: RedisValue) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
if matches!(value, RedisValue::Nil) {
return Ok(Vec::new());
}
let reply = from_redis_value::<StreamReadReply>(&value).map_err(redis_error)?;
Ok(reply
.keys
.into_iter()
.flat_map(|key| key.ids.into_iter())
.map(|id| RedisStreamEntry {
id: id.id,
fields: id
.map
.into_iter()
.filter_map(|(field, value)| {
redis::from_redis_value::<String>(&value)
.ok()
.map(|text| (field, text))
})
.collect(),
})
.collect())
}
fn parse_reclaim_result(value: RedisValue) -> Result<RedisStreamReclaimResult, DataLayerError> {
let RedisValue::Array(parts) = value else {
return Err(DataLayerError::UnexpectedValue(
@@ -573,9 +582,9 @@ mod tests {
use std::collections::BTreeMap;
use super::{
parse_reclaim_result, validate_consumer, validate_group, validate_stream_name,
validate_stream_position, RedisConsumerName, RedisStreamName, RedisStreamReclaimConfig,
RedisStreamReclaimResult, RedisStreamRunnerConfig,
parse_reclaim_result, parse_stream_read_entries, validate_consumer, validate_group,
validate_stream_name, validate_stream_position, RedisConsumerName, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunnerConfig,
};
use redis::Value as RedisValue;
@@ -639,6 +648,14 @@ mod tests {
assert!(validate_stream_position("").is_err());
}
#[test]
fn parses_empty_blocking_read_as_no_entries() {
let parsed = parse_stream_read_entries(RedisValue::Nil)
.expect("nil stream read reply should be empty");
assert!(parsed.is_empty());
}
#[test]
fn parses_reclaim_result_with_deleted_ids() {
let parsed = parse_reclaim_result(RedisValue::Array(vec![