refactor: extract runtime state backends

This commit is contained in:
fawney19
2026-05-08 00:18:12 +08:00
parent 6f620d92be
commit 6247ac3edc
111 changed files with 4358 additions and 3203 deletions

View File

@@ -17,7 +17,6 @@ chrono.workspace = true
chrono-tz.workspace = true
futures-util.workspace = true
flate2.workspace = true
redis.workspace = true
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true

View File

@@ -1,7 +1,7 @@
# aether-data
`aether-data` is the runtime data-access crate. It owns concrete database and
Redis clients, concrete repository implementations, migration/backfill/export
`aether-data` is the runtime data-access crate. It owns concrete SQL/database
drivers, concrete repository implementations, migration/backfill/export
workflows, and the composition layer that wires those pieces into the rest of
the application.
@@ -14,10 +14,9 @@ task crates live in `../aether-data-contracts`.
| Path | Responsibility |
|---|---|
| `src/database.rs` | Logical SQL driver selection and shared pool configuration. |
| `src/config.rs` | Data-layer config that combines SQL and Redis settings. |
| `src/config.rs` | Data-layer config for SQL drivers and repository wiring. |
| `src/maintenance.rs` | Maintenance DTOs and aggregation summaries used by backend dispatch and runtime maintenance entrypoints. |
| `src/driver/{postgres,mysql,sqlite}` | Low-level SQL driver primitives such as pools, transactions, and leases. These modules should not contain domain repository logic. |
| `src/driver/redis` | Low-level Redis clients, locks, streams, and namespaces. |
| `src/repository` | Domain repository traits/types re-exported from contracts plus concrete in-memory/Postgres/MySQL/SQLite implementations. |
| `src/backend` | Composition root. Builds concrete driver backends and exposes app-facing read/write/worker/lock/lease handles. |
| `src/backend/{maintenance,stats,wallet,system}.rs` | Backend-owned maintenance, aggregation, wallet ledger, and system config workflows that are not normal request-path repositories. |
@@ -40,9 +39,8 @@ The crate is easiest to read as five layers:
1. Contracts: DTOs, input structs, repository traits, and `DataLayerError`.
Prefer `aether-data-contracts` for anything that another crate needs to
compile against.
2. Driver primitives: `driver/postgres`, `driver/mysql`, `driver/sqlite`, and
`driver/redis` connect to
infrastructure and expose pools/runners.
2. Driver primitives: `driver/postgres`, `driver/mysql`, and `driver/sqlite`
connect to infrastructure and expose pools/runners.
3. Repository implementations: `repository/<domain>/{sql,mysql,sqlite,memory}`
translate contract types to driver-specific SQL.
4. Backend composition: `backend` chooses one SQL driver from config and wires

View File

@@ -1,58 +0,0 @@
use std::fmt;
use super::RedisBackend;
use crate::driver::redis::{RedisLockRunner, RedisLockRunnerConfig};
use crate::DataLayerError;
#[derive(Clone, Default)]
pub struct DataLockBackends {
redis: Option<RedisLockRunner>,
}
impl fmt::Debug for DataLockBackends {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataLockBackends")
.field("has_redis", &self.redis.is_some())
.finish()
}
}
impl DataLockBackends {
pub(crate) fn from_redis(redis: Option<&RedisBackend>) -> Result<Self, DataLayerError> {
Ok(Self {
redis: redis
.map(|backend| backend.lock_runner(RedisLockRunnerConfig::default()))
.transpose()?,
})
}
pub fn redis(&self) -> Option<RedisLockRunner> {
self.redis.clone()
}
pub fn has_any(&self) -> bool {
self.redis.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataLockBackends;
use crate::backend::RedisBackend;
use crate::driver::redis::RedisClientConfig;
#[test]
fn builds_redis_lock_runner_from_backend() {
let backend = RedisBackend::from_config(RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
})
.expect("redis backend should build");
let locks =
DataLockBackends::from_redis(Some(&backend)).expect("lock backends should build");
assert!(locks.has_any());
assert!(locks.redis().is_some());
}
}

View File

@@ -2,37 +2,31 @@
//!
//! `DataBackends` chooses the configured SQL driver, builds low-level pools,
//! instantiates concrete repositories, and exposes app-facing read/write,
//! lease, lock, worker, and maintenance handles. Request-path repository SQL
//! lease, transaction, and maintenance handles. Request-path repository SQL
//! belongs in `repository/*`; backend-owned maintenance SQL lives in focused
//! modules such as `stats`, `wallet`, and `system`. Pool/client primitives
//! belong in `driver/*`.
mod leases;
mod locks;
mod maintenance;
mod mysql;
mod postgres;
mod read;
mod redis;
mod sqlite;
mod stats;
mod stats_common;
mod system;
mod transactions;
mod wallet;
mod workers;
mod write;
use crate::maintenance::DatabasePoolSummary;
pub use leases::DataLeaseBackends;
pub use locks::DataLockBackends;
pub use mysql::MysqlBackend;
pub use postgres::PostgresBackend;
pub use read::DataReadRepositories;
pub use redis::RedisBackend;
pub use sqlite::SqliteBackend;
pub use transactions::DataTransactionBackends;
pub use workers::DataWorkerBackends;
pub use write::DataWriteRepositories;
use crate::database::DatabaseDriver;
@@ -51,12 +45,9 @@ pub struct DataBackends {
postgres: Option<PostgresBackend>,
mysql: Option<MysqlBackend>,
sqlite: Option<SqliteBackend>,
redis: Option<RedisBackend>,
leases: DataLeaseBackends,
locks: DataLockBackends,
read: DataReadRepositories,
transactions: DataTransactionBackends,
workers: DataWorkerBackends,
write: DataWriteRepositories,
}
@@ -111,17 +102,10 @@ impl DataBackends {
}
_ => None,
};
let redis = config
.redis
.clone()
.map(RedisBackend::from_config)
.transpose()?;
let leases = DataLeaseBackends::from_postgres(postgres.as_ref())?;
let locks = DataLockBackends::from_redis(redis.as_ref())?;
let read =
DataReadRepositories::from_backends(postgres.as_ref(), mysql.as_ref(), sqlite.as_ref());
let transactions = DataTransactionBackends::from_postgres(postgres.as_ref());
let workers = DataWorkerBackends::from_redis(redis.as_ref())?;
let write = DataWriteRepositories::from_backends(
postgres.as_ref(),
mysql.as_ref(),
@@ -133,12 +117,9 @@ impl DataBackends {
postgres,
mysql,
sqlite,
redis,
leases,
locks,
read,
transactions,
workers,
write,
})
}
@@ -165,10 +146,6 @@ impl DataBackends {
self.sqlite.as_ref()
}
pub fn redis(&self) -> Option<&RedisBackend> {
self.redis.as_ref()
}
pub fn read(&self) -> &DataReadRepositories {
&self.read
}
@@ -177,18 +154,10 @@ impl DataBackends {
&self.leases
}
pub fn locks(&self) -> &DataLockBackends {
&self.locks
}
pub fn transactions(&self) -> &DataTransactionBackends {
&self.transactions
}
pub fn workers(&self) -> &DataWorkerBackends {
&self.workers
}
pub fn write(&self) -> &DataWriteRepositories {
&self.write
}
@@ -197,12 +166,9 @@ impl DataBackends {
self.postgres.is_some()
|| self.mysql.is_some()
|| self.sqlite.is_some()
|| self.redis.is_some()
|| self.leases.has_any()
|| self.locks.has_any()
|| self.read.has_any()
|| self.transactions.has_any()
|| self.workers.has_any()
|| self.write.has_any()
}
}
@@ -224,9 +190,7 @@ mod tests {
assert!(backends.postgres().is_none());
assert!(backends.mysql().is_none());
assert!(backends.sqlite().is_none());
assert!(backends.redis().is_none());
assert!(backends.leases().postgres().is_none());
assert!(backends.locks().redis().is_none());
assert!(backends.read().auth_api_keys().is_none());
assert!(backends.read().auth_modules().is_none());
assert!(backends.read().billing().is_none());
@@ -241,7 +205,6 @@ mod tests {
assert!(backends.read().usage().is_none());
assert!(backends.read().video_tasks().is_none());
assert!(backends.transactions().postgres().is_none());
assert!(backends.workers().redis().is_none());
assert!(backends.write().settlement().is_none());
assert!(backends.write().usage().is_none());
}
@@ -260,7 +223,6 @@ mod tests {
statement_cache_capacity: 64,
require_ssl: false,
}),
redis: None,
})
.expect("postgres backend should build");
@@ -308,7 +270,6 @@ mod tests {
pool: SqlPoolConfig::default(),
}),
postgres: None,
redis: None,
})
.expect("mysql backend should build");
@@ -360,7 +321,6 @@ mod tests {
pool: SqlPoolConfig::default(),
}),
postgres: None,
redis: None,
})
.expect("sqlite backend should build");
@@ -401,34 +361,4 @@ mod tests {
assert!(backends.write().wallets().is_some());
assert!(backends.config().effective_database().is_some());
}
#[test]
fn builds_redis_backend_from_config() {
let backends = DataBackends::from_config(DataLayerConfig {
database: None,
postgres: None,
redis: Some(crate::driver::redis::RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
}),
})
.expect("redis backend should build");
assert!(backends.has_runtime_backends());
assert!(backends.postgres().is_none());
assert!(backends.mysql().is_none());
assert!(backends.sqlite().is_none());
assert!(backends.redis().is_some());
assert!(backends.leases().postgres().is_none());
assert!(backends.locks().redis().is_some());
assert!(backends.workers().redis().is_some());
assert!(backends.read().auth_api_keys().is_none());
assert!(backends.read().auth_modules().is_none());
assert!(backends.read().global_models().is_none());
assert!(backends.read().oauth_providers().is_none());
assert!(backends.transactions().postgres().is_none());
assert!(backends.write().settlement().is_none());
assert!(backends.write().usage().is_none());
assert!(backends.config().redis.is_some());
}
}

View File

@@ -1,86 +0,0 @@
use crate::driver::redis::{
RedisClient, RedisClientConfig, RedisClientFactory, RedisKeyspace, RedisKvRunner,
RedisKvRunnerConfig, RedisLockRunner, RedisLockRunnerConfig, RedisStreamRunner,
RedisStreamRunnerConfig,
};
use crate::DataLayerError;
#[derive(Debug, Clone)]
pub struct RedisBackend {
config: RedisClientConfig,
client: RedisClient,
}
impl RedisBackend {
pub fn from_config(config: RedisClientConfig) -> Result<Self, DataLayerError> {
let factory = RedisClientFactory::new(config.clone())?;
let client = factory.connect_lazy()?;
Ok(Self { config, client })
}
pub fn config(&self) -> &RedisClientConfig {
&self.config
}
pub fn client(&self) -> &RedisClient {
&self.client
}
pub fn client_clone(&self) -> RedisClient {
self.client.clone()
}
pub fn keyspace(&self) -> RedisKeyspace {
self.config.keyspace()
}
pub fn lock_runner(
&self,
config: RedisLockRunnerConfig,
) -> Result<RedisLockRunner, DataLayerError> {
RedisLockRunner::new(self.client_clone(), self.keyspace(), config)
}
pub fn stream_runner(
&self,
config: RedisStreamRunnerConfig,
) -> Result<RedisStreamRunner, DataLayerError> {
RedisStreamRunner::new(self.client_clone(), self.keyspace(), config)
}
pub fn kv_runner(&self, config: RedisKvRunnerConfig) -> Result<RedisKvRunner, DataLayerError> {
RedisKvRunner::new(self.client_clone(), self.keyspace(), config)
}
}
#[cfg(test)]
mod tests {
use super::RedisBackend;
use crate::driver::redis::{
RedisClientConfig, RedisKvRunnerConfig, RedisLockRunnerConfig, RedisStreamRunnerConfig,
};
#[test]
fn backend_retains_config_client_and_shared_runners() {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let backend = RedisBackend::from_config(config.clone()).expect("backend should build");
assert_eq!(backend.config(), &config);
assert_eq!(backend.keyspace().key("audit"), "aether:audit");
let _client_ref = backend.client();
let _client_clone = backend.client_clone();
let _lock_runner = backend
.lock_runner(RedisLockRunnerConfig::default())
.expect("lock runner should build");
let _stream_runner = backend
.stream_runner(RedisStreamRunnerConfig::default())
.expect("stream runner should build");
let _kv_runner = backend
.kv_runner(RedisKvRunnerConfig::default())
.expect("kv runner should build");
}
}

View File

@@ -1,58 +0,0 @@
use std::fmt;
use super::RedisBackend;
use crate::driver::redis::{RedisStreamRunner, RedisStreamRunnerConfig};
use crate::DataLayerError;
#[derive(Clone, Default)]
pub struct DataWorkerBackends {
redis: Option<RedisStreamRunner>,
}
impl fmt::Debug for DataWorkerBackends {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataWorkerBackends")
.field("has_redis", &self.redis.is_some())
.finish()
}
}
impl DataWorkerBackends {
pub(crate) fn from_redis(redis: Option<&RedisBackend>) -> Result<Self, DataLayerError> {
Ok(Self {
redis: redis
.map(|backend| backend.stream_runner(RedisStreamRunnerConfig::default()))
.transpose()?,
})
}
pub fn redis(&self) -> Option<RedisStreamRunner> {
self.redis.clone()
}
pub fn has_any(&self) -> bool {
self.redis.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataWorkerBackends;
use crate::backend::RedisBackend;
use crate::driver::redis::RedisClientConfig;
#[test]
fn builds_redis_stream_runner_from_backend() {
let backend = RedisBackend::from_config(RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
})
.expect("redis backend should build");
let workers =
DataWorkerBackends::from_redis(Some(&backend)).expect("worker backends should build");
assert!(workers.has_any());
assert!(workers.redis().is_some());
}
}

View File

@@ -1,13 +1,11 @@
use crate::database::SqlDatabaseConfig;
use crate::driver::postgres::PostgresPoolConfig;
use crate::driver::redis::RedisClientConfig;
use crate::DataLayerError;
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct DataLayerConfig {
pub database: Option<SqlDatabaseConfig>,
pub postgres: Option<PostgresPoolConfig>,
pub redis: Option<RedisClientConfig>,
}
impl DataLayerConfig {
@@ -15,7 +13,6 @@ impl DataLayerConfig {
Self {
database: Some(database),
postgres: None,
redis: None,
}
}
@@ -23,7 +20,6 @@ impl DataLayerConfig {
Self {
database: Some(SqlDatabaseConfig::from_postgres_config(postgres)),
postgres: None,
redis: None,
}
}
@@ -42,14 +38,11 @@ impl DataLayerConfig {
if let Some(postgres) = &self.postgres {
postgres.validate()?;
}
if let Some(redis) = &self.redis {
redis.validate()?;
}
Ok(())
}
pub fn has_persistent_backends(&self) -> bool {
self.effective_database().is_some() || self.redis.is_some()
self.effective_database().is_some()
}
}
@@ -58,7 +51,6 @@ mod tests {
use super::DataLayerConfig;
use crate::database::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig};
use crate::driver::postgres::PostgresPoolConfig;
use crate::driver::redis::RedisClientConfig;
#[test]
fn validates_nested_backend_configs() {
@@ -74,10 +66,6 @@ mod tests {
statement_cache_capacity: 64,
require_ssl: false,
}),
redis: Some(RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
}),
};
assert!(config.validate().is_ok());
@@ -98,7 +86,6 @@ mod tests {
statement_cache_capacity: 64,
require_ssl: false,
}),
redis: None,
};
assert!(config.validate().is_err());
@@ -122,7 +109,6 @@ mod tests {
statement_cache_capacity: 64,
require_ssl: false,
}),
redis: None,
};
let effective = config

View File

@@ -1,10 +1,9 @@
//! Low-level data driver primitives.
//!
//! These modules own pools, transactions, leases, and Redis client helpers.
//! These modules own pools, transactions, and lease helpers.
//! Domain repository logic belongs in `repository/*`, and application-facing
//! composition belongs in `backend`.
pub mod mysql;
pub mod postgres;
pub mod redis;
pub mod sqlite;

View File

@@ -4,10 +4,6 @@ pub(crate) fn postgres_error(error: impl std::fmt::Display) -> DataLayerError {
DataLayerError::postgres(error)
}
pub(crate) fn redis_error(error: impl std::fmt::Display) -> DataLayerError {
DataLayerError::redis(error)
}
pub(crate) fn sql_error(error: impl std::fmt::Display) -> DataLayerError {
DataLayerError::sql(error)
}
@@ -22,16 +18,6 @@ impl<T> SqlxResultExt<T> for Result<T, sqlx::Error> {
}
}
pub(crate) trait RedisResultExt<T> {
fn map_redis_err(self) -> Result<T, DataLayerError>;
}
impl<T> RedisResultExt<T> for Result<T, redis::RedisError> {
fn map_redis_err(self) -> Result<T, DataLayerError> {
self.map_err(redis_error)
}
}
pub(crate) trait SqlResultExt<T> {
fn map_sql_err(self) -> Result<T, DataLayerError>;
}

View File

@@ -1,6 +1,6 @@
//! Runtime data access for Aether.
//!
//! This crate contains concrete database/Redis clients, repository
//! This crate contains concrete database clients, repository
//! implementations, migration/backfill/export workflows, and the backend
//! composition layer. Shared repository contracts that other crates compile
//! against live in `aether-data-contracts`.
@@ -17,9 +17,8 @@ pub mod maintenance;
pub mod repository;
pub use backend::{
DataBackends, DataLeaseBackends, DataLockBackends, DataReadRepositories,
DataTransactionBackends, DataWorkerBackends, DataWriteRepositories, PostgresBackend,
RedisBackend,
DataBackends, DataLeaseBackends, DataReadRepositories, DataTransactionBackends,
DataWriteRepositories, PostgresBackend,
};
pub use config::DataLayerConfig;
pub use database::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig, DEFAULT_SQLITE_DATABASE_URL};

View File

@@ -211,7 +211,7 @@ fn relative_path(path: &Path, workspace_root: &Path) -> PathBuf {
fn grouped_import_scanner_allows_nested_new_paths() {
let source = r#"
use aether_data::{
driver::{postgres::PostgresPool, redis::RedisKvRunner},
driver::{postgres::PostgresPool, mysql::MySqlPool},
lifecycle::{backfill::PendingBackfillInfo, migrate::PendingMigrationInfo},
};
"#;

View File

@@ -10,9 +10,9 @@ description = "Provider transport core extracted from aether-gateway"
aether-ai-formats.workspace = true
aether-contracts.workspace = true
aether-crypto.workspace = true
aether-data.workspace = true
aether-data-contracts.workspace = true
aether-oauth.workspace = true
aether-runtime-state.workspace = true
aether-video-tasks-core.workspace = true
async-trait.workspace = true
http.workspace = true

View File

@@ -2,12 +2,12 @@ use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use aether_data::driver::redis::{RedisLockKey, RedisLockRunner};
use aether_oauth::core::OAuthError;
use aether_oauth::network::{
OAuthHttpExecutor, OAuthHttpRequest, OAuthHttpResponse, OAuthNetworkContext,
};
use aether_oauth::provider::ProviderOAuthTransportContext;
use aether_runtime_state::RuntimeState;
use async_trait::async_trait;
use serde_json::Value;
use thiserror::Error;
@@ -374,7 +374,7 @@ impl LocalOAuthRefreshCoordinator {
&self,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
distributed_lock: Option<&RedisLockRunner>,
distributed_lock: Option<&RuntimeState>,
distributed_owner: Option<&str>,
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
self.resolve_with_result_mode(
@@ -391,7 +391,7 @@ impl LocalOAuthRefreshCoordinator {
&self,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
distributed_lock: Option<&RedisLockRunner>,
distributed_lock: Option<&RuntimeState>,
distributed_owner: Option<&str>,
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
self.resolve_with_result_mode(
@@ -408,7 +408,7 @@ impl LocalOAuthRefreshCoordinator {
&self,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
distributed_lock: Option<&RedisLockRunner>,
distributed_lock: Option<&RuntimeState>,
distributed_owner: Option<&str>,
force_refresh: bool,
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
@@ -465,12 +465,11 @@ impl LocalOAuthRefreshCoordinator {
let distributed_lease = match (distributed_lock, distributed_owner) {
(Some(lock), Some(owner)) if !owner.trim().is_empty() => {
let lock_key = RedisLockKey(format!("provider_oauth_refresh_lock:{key_id}"));
match lock
.try_acquire(
&lock_key,
.lock_try_acquire(
&format!("provider_oauth_refresh_lock:{key_id}"),
owner,
Some(Self::DISTRIBUTED_REFRESH_LOCK_TTL_MS),
std::time::Duration::from_millis(Self::DISTRIBUTED_REFRESH_LOCK_TTL_MS),
)
.await
{
@@ -497,7 +496,7 @@ impl LocalOAuthRefreshCoordinator {
let refresh_entry = cached_entry.as_ref();
let refresh_result = adapter.refresh(executor, transport, refresh_entry).await;
if let (Some(lock), Some(lease)) = (distributed_lock, distributed_lease.as_ref()) {
if let Err(err) = lock.release(lease).await {
if let Err(err) = lock.lock_release(lease).await {
tracing::warn!(
key_id = %key_id,
provider_type = adapter.provider_type(),

View File

@@ -0,0 +1,21 @@
[package]
name = "aether-runtime-state"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Runtime state backends for Aether services"
[dependencies]
aether-cache.workspace = true
aether-data-contracts.workspace = true
aether-runtime.workspace = true
async-trait.workspace = true
redis.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio.workspace = true
tracing.workspace = true
url.workspace = true
uuid.workspace = true

View File

@@ -0,0 +1,15 @@
pub use aether_data_contracts::DataLayerError;
pub(crate) fn redis_error(error: impl std::fmt::Display) -> DataLayerError {
DataLayerError::redis(error)
}
pub(crate) trait RedisResultExt<T> {
fn map_redis_err(self) -> Result<T, DataLayerError>;
}
impl<T> RedisResultExt<T> for Result<T, redis::RedisError> {
fn map_redis_err(self) -> Result<T, DataLayerError> {
self.map_err(redis_error)
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,537 @@
use std::collections::{BTreeMap, BTreeSet, HashMap, VecDeque};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use tokio::sync::Mutex;
use crate::{RuntimeQueueEntry, RuntimeQueueReclaimConfig};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MemoryRuntimeStateConfig {
pub max_kv_entries: usize,
}
impl Default for MemoryRuntimeStateConfig {
fn default() -> Self {
Self {
max_kv_entries: 10_000,
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct MemoryKvEntry {
pub(crate) value: String,
pub(crate) inserted_at: Instant,
pub(crate) expires_at: Option<Instant>,
}
impl MemoryKvEntry {
fn is_expired(&self, now: Instant) -> bool {
self.expires_at.is_some_and(|expires_at| now >= expires_at)
}
}
#[derive(Debug, Default)]
pub(crate) struct MemoryRuntimeBackend {
config: MemoryRuntimeStateConfig,
kv: Mutex<HashMap<String, MemoryKvEntry>>,
counters: Mutex<HashMap<String, MemoryCounterEntry>>,
sets: Mutex<HashMap<String, BTreeSet<String>>>,
scores: Mutex<HashMap<String, BTreeMap<String, f64>>>,
queues: Mutex<HashMap<String, VecDeque<RuntimeQueueEntry>>>,
queue_seq: AtomicU64,
locks: Mutex<HashMap<String, MemoryLockEntry>>,
semaphores: Mutex<HashMap<String, BTreeMap<String, u64>>>,
}
#[derive(Debug, Clone)]
struct MemoryCounterEntry {
value: u32,
bucket: u64,
expires_at: Instant,
}
#[derive(Debug, Clone)]
pub(crate) struct MemoryLockEntry {
pub(crate) token: String,
#[allow(dead_code)]
pub(crate) owner: String,
pub(crate) expires_at: Instant,
}
impl MemoryRuntimeBackend {
pub(crate) fn new(config: MemoryRuntimeStateConfig) -> Self {
Self {
config,
..Self::default()
}
}
pub(crate) async fn kv_set(&self, key: &str, value: String, ttl: Option<Duration>) {
let mut kv = self.kv.lock().await;
let now = Instant::now();
if ttl.is_some_and(|ttl| ttl.is_zero()) {
kv.remove(key);
return;
}
prune_kv(&mut kv, now);
while kv.len() >= self.config.max_kv_entries.max(1) {
let Some(oldest_key) = kv
.iter()
.min_by_key(|(_, entry)| entry.inserted_at)
.map(|(key, _)| key.clone())
else {
break;
};
kv.remove(&oldest_key);
}
kv.insert(
key.to_string(),
MemoryKvEntry {
value,
inserted_at: now,
expires_at: ttl.map(|ttl| now + ttl),
},
);
}
pub(crate) fn kv_set_nowait(&self, key: &str, value: String, ttl: Option<Duration>) -> bool {
let Ok(mut kv) = self.kv.try_lock() else {
return false;
};
let now = Instant::now();
if ttl.is_some_and(|ttl| ttl.is_zero()) {
kv.remove(key);
return true;
}
prune_kv(&mut kv, now);
while kv.len() >= self.config.max_kv_entries.max(1) {
let Some(oldest_key) = kv
.iter()
.min_by_key(|(_, entry)| entry.inserted_at)
.map(|(key, _)| key.clone())
else {
break;
};
kv.remove(&oldest_key);
}
kv.insert(
key.to_string(),
MemoryKvEntry {
value,
inserted_at: now,
expires_at: ttl.map(|ttl| now + ttl),
},
);
true
}
pub(crate) async fn kv_get(&self, key: &str) -> Option<String> {
let mut kv = self.kv.lock().await;
get_fresh_locked(&mut kv, key, Instant::now())
}
pub(crate) async fn kv_take(&self, key: &str) -> Option<String> {
let mut kv = self.kv.lock().await;
let now = Instant::now();
let entry = kv.remove(key)?;
if entry.is_expired(now) {
return None;
}
Some(entry.value)
}
pub(crate) async fn kv_delete(&self, key: &str) -> bool {
self.kv.lock().await.remove(key).is_some()
}
pub(crate) async fn kv_delete_many(&self, keys: &[String]) -> usize {
let mut kv = self.kv.lock().await;
keys.iter().filter(|key| kv.remove(*key).is_some()).count()
}
pub(crate) async fn kv_exists(&self, key: &str) -> bool {
self.kv_get(key).await.is_some()
}
pub(crate) async fn kv_ttl_seconds(&self, key: &str) -> Option<i64> {
let mut kv = self.kv.lock().await;
let now = Instant::now();
let entry = kv.get(key).cloned()?;
if entry.is_expired(now) {
kv.remove(key);
return None;
}
Some(
entry
.expires_at
.map(|expires_at| {
expires_at
.saturating_duration_since(now)
.as_secs()
.try_into()
.unwrap_or(i64::MAX)
})
.unwrap_or(-1),
)
}
pub(crate) async fn kv_scan(&self, pattern: &str) -> Vec<String> {
let mut kv = self.kv.lock().await;
prune_kv(&mut kv, Instant::now());
let mut keys = kv
.keys()
.filter(|key| key_matches_pattern(key, pattern))
.cloned()
.collect::<Vec<_>>();
keys.sort();
keys
}
pub(crate) async fn check_and_consume_rate_limit(
&self,
user_key: &str,
key_key: &str,
bucket: u64,
user_limit: u32,
key_limit: u32,
ttl: Duration,
) -> Result<crate::RateLimitCheck, crate::DataLayerError> {
let mut counters = self.counters.lock().await;
let now = Instant::now();
counters.retain(|_, entry| entry.expires_at > now && entry.bucket >= bucket);
if user_limit > 0 {
let user_count = counters
.get(user_key)
.filter(|entry| entry.bucket == bucket)
.map(|entry| entry.value)
.unwrap_or_default();
if user_count >= user_limit {
return Ok(crate::RateLimitCheck::Rejected {
scope: crate::RateLimitScope::User,
limit: user_limit,
});
}
}
if key_limit > 0 {
let key_count = counters
.get(key_key)
.filter(|entry| entry.bucket == bucket)
.map(|entry| entry.value)
.unwrap_or_default();
if key_count >= key_limit {
return Ok(crate::RateLimitCheck::Rejected {
scope: crate::RateLimitScope::Key,
limit: key_limit,
});
}
}
let mut remaining = None::<u32>;
let expires_at = now + ttl;
if user_limit > 0 {
let next = counters
.entry(user_key.to_string())
.and_modify(|entry| {
entry.bucket = bucket;
entry.value = entry.value.saturating_add(1);
entry.expires_at = expires_at;
})
.or_insert(MemoryCounterEntry {
value: 1,
bucket,
expires_at,
})
.value;
remaining = Some(user_limit.saturating_sub(next));
}
if key_limit > 0 {
let next = counters
.entry(key_key.to_string())
.and_modify(|entry| {
entry.bucket = bucket;
entry.value = entry.value.saturating_add(1);
entry.expires_at = expires_at;
})
.or_insert(MemoryCounterEntry {
value: 1,
bucket,
expires_at,
})
.value;
let key_remaining = key_limit.saturating_sub(next);
remaining = Some(remaining.map_or(key_remaining, |value| value.min(key_remaining)));
}
Ok(crate::RateLimitCheck::Allowed {
remaining: remaining.unwrap_or(0),
})
}
pub(crate) async fn set_add(&self, key: &str, member: &str) -> bool {
self.sets
.lock()
.await
.entry(key.to_string())
.or_default()
.insert(member.to_string())
}
pub(crate) fn set_add_nowait(&self, key: &str, member: &str) -> bool {
let Ok(mut sets) = self.sets.try_lock() else {
return false;
};
sets.entry(key.to_string())
.or_default()
.insert(member.to_string())
}
pub(crate) async fn set_remove(&self, key: &str, member: &str) -> bool {
self.sets
.lock()
.await
.get_mut(key)
.is_some_and(|set| set.remove(member))
}
pub(crate) async fn set_members(&self, key: &str) -> Vec<String> {
self.sets
.lock()
.await
.get(key)
.map(|set| set.iter().cloned().collect())
.unwrap_or_default()
}
pub(crate) async fn set_len(&self, key: &str) -> usize {
self.sets.lock().await.get(key).map_or(0, BTreeSet::len)
}
pub(crate) async fn score_set(&self, key: &str, member: &str, score: f64) {
self.scores
.lock()
.await
.entry(key.to_string())
.or_default()
.insert(member.to_string(), score);
}
pub(crate) async fn score_many(&self, key: &str, members: &[String]) -> Vec<Option<f64>> {
let scores = self.scores.lock().await;
members
.iter()
.map(|member| scores.get(key).and_then(|set| set.get(member)).copied())
.collect()
}
pub(crate) async fn score_range_by_min(&self, key: &str, min_score: f64) -> Vec<String> {
let scores = self.scores.lock().await;
scores
.get(key)
.map(|set| {
set.iter()
.filter(|(_, score)| **score >= min_score)
.map(|(member, _)| member.clone())
.collect()
})
.unwrap_or_default()
}
pub(crate) async fn score_remove_by_score(&self, key: &str, max_score: f64) -> usize {
let mut scores = self.scores.lock().await;
let Some(set) = scores.get_mut(key) else {
return 0;
};
let before = set.len();
set.retain(|_, score| *score > max_score);
before.saturating_sub(set.len())
}
pub(crate) async fn score_len(&self, key: &str) -> usize {
self.scores.lock().await.get(key).map_or(0, BTreeMap::len)
}
pub(crate) async fn queue_append(
&self,
stream: &str,
fields: BTreeMap<String, String>,
maxlen: Option<usize>,
) -> String {
let id = format!(
"{}-0",
self.queue_seq
.fetch_add(1, Ordering::Relaxed)
.saturating_add(1)
);
let mut queues = self.queues.lock().await;
let queue = queues.entry(stream.to_string()).or_default();
queue.push_back(RuntimeQueueEntry {
id: id.clone(),
fields,
});
if let Some(maxlen) = maxlen.filter(|value| *value > 0) {
while queue.len() > maxlen {
queue.pop_front();
}
}
id
}
pub(crate) async fn queue_read(&self, stream: &str, count: usize) -> Vec<RuntimeQueueEntry> {
let mut queues = self.queues.lock().await;
let Some(queue) = queues.get_mut(stream) else {
return Vec::new();
};
let mut entries = Vec::new();
for _ in 0..count.max(1) {
let Some(entry) = queue.pop_front() else {
break;
};
entries.push(entry);
}
entries
}
pub(crate) async fn queue_claim_stale(
&self,
_stream: &str,
_config: RuntimeQueueReclaimConfig,
) -> Vec<RuntimeQueueEntry> {
Vec::new()
}
pub(crate) async fn queue_delete(&self, _stream: &str, _ids: &[String]) -> usize {
0
}
pub(crate) async fn lock_try_acquire(
&self,
key: &str,
owner: &str,
token: String,
ttl: Duration,
) -> bool {
let mut locks = self.locks.lock().await;
let now = Instant::now();
locks.retain(|_, entry| entry.expires_at > now);
if locks.contains_key(key) {
return false;
}
locks.insert(
key.to_string(),
MemoryLockEntry {
token,
owner: owner.to_string(),
expires_at: now + ttl,
},
);
true
}
pub(crate) async fn lock_release(&self, key: &str, token: &str) -> bool {
let mut locks = self.locks.lock().await;
if locks.get(key).is_some_and(|entry| entry.token == token) {
locks.remove(key);
return true;
}
false
}
pub(crate) async fn lock_renew(&self, key: &str, token: &str, ttl: Duration) -> bool {
let mut locks = self.locks.lock().await;
if let Some(entry) = locks.get_mut(key) {
if entry.token == token {
entry.expires_at = Instant::now() + ttl;
return true;
}
}
false
}
pub(crate) async fn semaphore_try_acquire(
&self,
key: &str,
token: String,
limit: usize,
ttl_ms: u64,
) -> Result<usize, usize> {
let now_ms = unix_time_ms();
let expires_at = now_ms.saturating_add(ttl_ms);
let mut semaphores = self.semaphores.lock().await;
let holders = semaphores.entry(key.to_string()).or_default();
holders.retain(|_, expires| *expires > now_ms);
let count = holders.len();
if count >= limit {
return Err(count);
}
holders.insert(token, expires_at);
Ok(holders.len())
}
pub(crate) async fn semaphore_renew(&self, key: &str, token: &str, ttl_ms: u64) -> bool {
let now_ms = unix_time_ms();
let mut semaphores = self.semaphores.lock().await;
let Some(holders) = semaphores.get_mut(key) else {
return false;
};
holders.retain(|_, expires| *expires > now_ms);
if let Some(expires) = holders.get_mut(token) {
*expires = now_ms.saturating_add(ttl_ms);
return true;
}
false
}
pub(crate) async fn semaphore_release(&self, key: &str, token: &str) {
let mut semaphores = self.semaphores.lock().await;
if let Some(holders) = semaphores.get_mut(key) {
holders.remove(token);
if holders.is_empty() {
semaphores.remove(key);
}
}
}
pub(crate) async fn semaphore_live_count(&self, key: &str) -> usize {
let now_ms = unix_time_ms();
let mut semaphores = self.semaphores.lock().await;
let Some(holders) = semaphores.get_mut(key) else {
return 0;
};
holders.retain(|_, expires| *expires > now_ms);
holders.len()
}
}
fn get_fresh_locked(
kv: &mut HashMap<String, MemoryKvEntry>,
key: &str,
now: Instant,
) -> Option<String> {
let entry = kv.get(key).cloned()?;
if entry.is_expired(now) {
kv.remove(key);
return None;
}
Some(entry.value)
}
fn prune_kv(kv: &mut HashMap<String, MemoryKvEntry>, now: Instant) {
kv.retain(|_, entry| !entry.is_expired(now));
}
pub(crate) fn key_matches_pattern(key: &str, pattern: &str) -> bool {
match pattern.strip_suffix('*') {
Some(prefix) => key.starts_with(prefix),
None => key == pattern,
}
}
fn unix_time_ms() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}

View File

@@ -1,5 +1,5 @@
use crate::driver::redis::RedisKeyspace;
use crate::error::RedisResultExt;
use crate::redis::RedisKeyspace;
use crate::DataLayerError;
pub type RedisClient = redis::Client;

View File

@@ -1,8 +1,8 @@
use std::future::Future;
use std::time::Duration;
use crate::driver::redis::{RedisClient, RedisKeyspace};
use crate::error::RedisResultExt;
use crate::redis::{RedisClient, RedisKeyspace};
use crate::DataLayerError;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -96,6 +96,59 @@ impl RedisKvRunner {
.await
}
pub async fn get(&self, key: &str) -> Result<Option<String>, DataLayerError> {
let namespaced_key = self.keyspace.key(key);
self.run_with_timeout("redis kv get", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
redis::cmd("GET")
.arg(&namespaced_key)
.query_async(&mut connection)
.await
.map_redis_err()
})
.await
}
pub async fn getdel(&self, key: &str) -> Result<Option<String>, DataLayerError> {
let namespaced_key = self.keyspace.key(key);
self.run_with_timeout("redis kv getdel", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
redis::cmd("GETDEL")
.arg(&namespaced_key)
.query_async(&mut connection)
.await
.map_redis_err()
})
.await
}
pub async fn exists(&self, key: &str) -> Result<bool, DataLayerError> {
let namespaced_key = self.keyspace.key(key);
let exists = self
.run_with_timeout("redis kv exists", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_redis_err()?;
redis::cmd("EXISTS")
.arg(&namespaced_key)
.query_async::<i64>(&mut connection)
.await
.map_redis_err()
})
.await?;
Ok(exists > 0)
}
pub async fn del(&self, key: &str) -> Result<i64, DataLayerError> {
let namespaced_key = self.keyspace.key(key);
self.run_with_timeout("redis kv del", async {
@@ -136,7 +189,7 @@ impl RedisKvRunner {
#[cfg(test)]
mod tests {
use super::{RedisKvRunner, RedisKvRunnerConfig};
use crate::driver::redis::{RedisClientConfig, RedisClientFactory, RedisKeyspace};
use crate::redis::{RedisClientConfig, RedisClientFactory, RedisKeyspace};
fn build_runner() -> RedisKvRunner {
let config = RedisClientConfig {

View File

@@ -1,8 +1,8 @@
use std::future::Future;
use std::time::Duration;
use crate::driver::redis::{RedisClient, RedisKeyspace};
use crate::error::RedisResultExt;
use crate::redis::{RedisClient, RedisKeyspace};
use crate::DataLayerError;
use uuid::Uuid;
@@ -242,7 +242,7 @@ fn validate_lease(lease: &RedisLockLease) -> Result<(), DataLayerError> {
#[cfg(test)]
mod tests {
use super::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
use crate::driver::redis::{RedisClientConfig, RedisClientFactory};
use crate::redis::{RedisClientConfig, RedisClientFactory};
fn sample_runner() -> RedisLockRunner {
let client = RedisClientFactory::new(RedisClientConfig {

View File

@@ -12,3 +12,14 @@ pub use stream::{
RedisConsumerGroup, RedisConsumerName, RedisStreamEntry, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner, RedisStreamRunnerConfig,
};
pub(crate) type RedisCmd = redis::Cmd;
pub(crate) type RedisScript = redis::Script;
pub(crate) fn cmd(name: &str) -> RedisCmd {
redis::cmd(name)
}
pub(crate) fn script(source: &str) -> RedisScript {
redis::Script::new(source)
}

View File

@@ -1,6 +1,6 @@
use aether_cache::CacheKeyNamespace;
use crate::driver::redis::{RedisLockKey, RedisStreamName};
use crate::redis::{RedisLockKey, RedisStreamName};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisKeyspace {

View File

@@ -6,8 +6,8 @@ use redis::from_redis_value;
use redis::streams::StreamReadReply;
use redis::Value as RedisValue;
use crate::driver::redis::{RedisClient, RedisKeyspace};
use crate::error::{redis_error, RedisResultExt};
use crate::redis::{RedisClient, RedisKeyspace};
use crate::DataLayerError;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -575,7 +575,7 @@ mod tests {
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner,
RedisStreamRunnerConfig,
};
use crate::driver::redis::{RedisClientConfig, RedisClientFactory};
use crate::redis::{RedisClientConfig, RedisClientFactory};
use redis::Value as RedisValue;
fn sample_runner() -> RedisStreamRunner {

View File

@@ -11,12 +11,10 @@ async-stream.workspace = true
axum = { version = "0.8" }
chrono.workspace = true
futures-util.workspace = true
redis.workspace = true
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
tokio.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
url.workspace = true
uuid.workspace = true

View File

@@ -4,25 +4,32 @@ use axum::http::Response;
use futures_util::StreamExt;
use crate::concurrency::ConcurrencyPermit;
use crate::distributed::DistributedConcurrencyPermit;
#[derive(Debug)]
pub struct AdmissionPermit {
_local: Option<ConcurrencyPermit>,
_distributed: Option<DistributedConcurrencyPermit>,
_distributed: Option<Box<dyn Send + Sync>>,
}
impl std::fmt::Debug for AdmissionPermit {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AdmissionPermit")
.field("has_local", &self._local.is_some())
.field("has_distributed", &self._distributed.is_some())
.finish()
}
}
impl AdmissionPermit {
pub fn from_parts(
pub fn from_parts<D: Send + Sync + 'static>(
local: Option<ConcurrencyPermit>,
distributed: Option<DistributedConcurrencyPermit>,
distributed: Option<D>,
) -> Option<Self> {
if local.is_none() && distributed.is_none() {
None
} else {
Some(Self {
_local: local,
_distributed: distributed,
_distributed: distributed.map(|permit| Box::new(permit) as Box<dyn Send + Sync>),
})
}
}
@@ -70,7 +77,7 @@ fn hold_axum_response_permit(response: Response<Body>, permit: AdmissionPermit)
#[cfg(test)]
mod tests {
use super::{hold_admission_permit_until, maybe_hold_axum_response_permit, AdmissionPermit};
use crate::{ConcurrencyGate, DistributedConcurrencyGate};
use crate::ConcurrencyGate;
use axum::body::{to_bytes, Body};
use axum::http::Response;
@@ -96,12 +103,9 @@ mod tests {
#[tokio::test]
async fn holds_combined_local_and_distributed_permit_until_future_finishes() {
let local_gate = ConcurrencyGate::new("local", 1);
let distributed_gate = DistributedConcurrencyGate::new_in_memory("distributed", 1);
let local = local_gate.try_acquire().expect("local permit");
let distributed = distributed_gate
.try_acquire()
.await
.expect("distributed permit");
let distributed_gate = ConcurrencyGate::new("distributed", 1);
let distributed = distributed_gate.try_acquire().expect("distributed permit");
let task = tokio::spawn(hold_admission_permit_until(
AdmissionPermit::from_parts(Some(local), Some(distributed)),
@@ -116,19 +120,12 @@ mod tests {
"local permit should still be held"
);
assert!(
distributed_gate.try_acquire().await.is_err(),
distributed_gate.try_acquire().is_err(),
"distributed permit should still be held"
);
task.await.expect("task should complete");
assert_eq!(local_gate.snapshot().in_flight, 0);
assert_eq!(
distributed_gate
.snapshot()
.await
.expect("snapshot should build")
.in_flight,
0
);
assert_eq!(distributed_gate.snapshot().in_flight, 0);
}
}

View File

@@ -1,10 +1,4 @@
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::task::JoinHandle;
use tracing::warn;
use uuid::Uuid;
use crate::concurrency::{ConcurrencyGate, ConcurrencyPermit};
use crate::metrics::{MetricKind, MetricLabel, MetricSample};
@@ -68,80 +62,11 @@ impl DistributedConcurrencySnapshot {
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisDistributedConcurrencyConfig {
pub url: String,
pub key_prefix: Option<String>,
pub lease_ttl_ms: u64,
pub renew_interval_ms: u64,
pub command_timeout_ms: Option<u64>,
}
impl Default for RedisDistributedConcurrencyConfig {
fn default() -> Self {
Self {
url: String::new(),
key_prefix: None,
lease_ttl_ms: 30_000,
renew_interval_ms: 10_000,
command_timeout_ms: Some(1_000),
}
}
}
impl RedisDistributedConcurrencyConfig {
fn validate(&self) -> Result<(), DistributedConcurrencyError> {
let raw = self.url.trim();
if raw.is_empty() {
return Err(DistributedConcurrencyError::InvalidConfiguration(
"distributed concurrency redis url cannot be empty".to_string(),
));
}
url::Url::parse(raw).map_err(|err| {
DistributedConcurrencyError::InvalidConfiguration(format!(
"invalid distributed concurrency redis url: {err}"
))
})?;
if self.lease_ttl_ms == 0 {
return Err(DistributedConcurrencyError::InvalidConfiguration(
"distributed concurrency lease_ttl_ms must be positive".to_string(),
));
}
if self.renew_interval_ms == 0 {
return Err(DistributedConcurrencyError::InvalidConfiguration(
"distributed concurrency renew_interval_ms must be positive".to_string(),
));
}
if self.renew_interval_ms >= self.lease_ttl_ms {
return Err(DistributedConcurrencyError::InvalidConfiguration(
"distributed concurrency renew_interval_ms must be smaller than lease_ttl_ms"
.to_string(),
));
}
if matches!(self.command_timeout_ms, Some(0)) {
return Err(DistributedConcurrencyError::InvalidConfiguration(
"distributed concurrency command_timeout_ms must be positive".to_string(),
));
}
Ok(())
}
fn semaphore_key(&self, gate: &'static str) -> String {
prefixed_key(self.key_prefix.as_deref(), &format!("admission:{gate}"))
}
}
#[derive(Debug)]
enum DistributedConcurrencyBackend {
InMemory(Arc<ConcurrencyGate>),
Redis(Arc<RedisDistributedState>),
}
#[derive(Debug)]
struct DistributedConcurrencyState {
gate: &'static str,
limit: usize,
backend: DistributedConcurrencyBackend,
gate_impl: Arc<ConcurrencyGate>,
}
#[derive(Debug, Clone)]
@@ -159,49 +84,11 @@ impl DistributedConcurrencyGate {
state: Arc::new(DistributedConcurrencyState {
gate,
limit,
backend: DistributedConcurrencyBackend::InMemory(Arc::new(ConcurrencyGate::new(
gate, limit,
))),
gate_impl: Arc::new(ConcurrencyGate::new(gate, limit)),
}),
}
}
pub fn new_redis(
gate: &'static str,
limit: usize,
config: RedisDistributedConcurrencyConfig,
) -> Result<Self, DistributedConcurrencyError> {
if limit == 0 {
return Err(DistributedConcurrencyError::InvalidConfiguration(
"distributed concurrency gate limit must be positive".to_string(),
));
}
config.validate()?;
let client = redis::Client::open(config.url.clone()).map_err(|err| {
DistributedConcurrencyError::InvalidConfiguration(format!(
"failed to build distributed concurrency redis client: {err}"
))
})?;
Ok(Self {
state: Arc::new(DistributedConcurrencyState {
gate,
limit,
backend: DistributedConcurrencyBackend::Redis(Arc::new(RedisDistributedState {
gate,
limit,
client,
key: config.semaphore_key(gate),
lease_ttl_ms: config.lease_ttl_ms,
renew_interval_ms: config.renew_interval_ms,
command_timeout_ms: config.command_timeout_ms,
high_watermark: AtomicUsize::new(0),
rejected: AtomicU64::new(0),
})),
}),
})
}
pub fn gate(&self) -> &'static str {
self.state.gate
}
@@ -213,364 +100,46 @@ impl DistributedConcurrencyGate {
pub async fn try_acquire(
&self,
) -> Result<DistributedConcurrencyPermit, DistributedConcurrencyError> {
match &self.state.backend {
DistributedConcurrencyBackend::InMemory(gate) => gate
.try_acquire()
.map(DistributedConcurrencyPermit::from_in_memory)
.map_err(|err| match err {
crate::ConcurrencyError::Saturated { gate, limit } => {
DistributedConcurrencyError::Saturated { gate, limit }
self.state
.gate_impl
.try_acquire()
.map(|permit| DistributedConcurrencyPermit { _permit: permit })
.map_err(|err| match err {
crate::ConcurrencyError::Saturated { gate, limit } => {
DistributedConcurrencyError::Saturated { gate, limit }
}
crate::ConcurrencyError::Closed { gate } => {
DistributedConcurrencyError::Unavailable {
gate,
limit: self.state.limit,
message: "in-memory distributed concurrency gate is closed".to_string(),
}
crate::ConcurrencyError::Closed { gate } => {
DistributedConcurrencyError::Unavailable {
gate,
limit: self.state.limit,
message: "in-memory distributed concurrency gate is closed".to_string(),
}
}
}),
DistributedConcurrencyBackend::Redis(state) => state.try_acquire().await,
}
}
})
}
pub async fn snapshot(
&self,
) -> Result<DistributedConcurrencySnapshot, DistributedConcurrencyError> {
match &self.state.backend {
DistributedConcurrencyBackend::InMemory(gate) => {
let snapshot = gate.snapshot();
Ok(DistributedConcurrencySnapshot {
limit: snapshot.limit,
in_flight: snapshot.in_flight,
available_permits: snapshot.available_permits,
high_watermark: snapshot.high_watermark,
rejected: snapshot.rejected,
})
}
DistributedConcurrencyBackend::Redis(state) => state.snapshot().await,
}
let snapshot = self.state.gate_impl.snapshot();
Ok(DistributedConcurrencySnapshot {
limit: snapshot.limit,
in_flight: snapshot.in_flight,
available_permits: snapshot.available_permits,
high_watermark: snapshot.high_watermark,
rejected: snapshot.rejected,
})
}
}
#[derive(Debug)]
pub struct DistributedConcurrencyPermit {
inner: DistributedConcurrencyPermitInner,
}
#[derive(Debug)]
enum DistributedConcurrencyPermitInner {
InMemory(ConcurrencyPermit),
Redis {
state: Arc<RedisDistributedState>,
token: String,
renew_task: JoinHandle<()>,
},
}
impl DistributedConcurrencyPermit {
fn from_in_memory(permit: ConcurrencyPermit) -> Self {
Self {
inner: DistributedConcurrencyPermitInner::InMemory(permit),
}
}
fn from_redis(
state: Arc<RedisDistributedState>,
token: String,
renew_task: JoinHandle<()>,
) -> Self {
Self {
inner: DistributedConcurrencyPermitInner::Redis {
state,
token,
renew_task,
},
}
}
}
impl Drop for DistributedConcurrencyPermit {
fn drop(&mut self) {
match &mut self.inner {
DistributedConcurrencyPermitInner::InMemory(_permit) => {}
DistributedConcurrencyPermitInner::Redis {
state,
token,
renew_task,
} => {
renew_task.abort();
let state = Arc::clone(state);
let token = token.clone();
tokio::spawn(async move {
if let Err(err) = state.release(&token).await {
warn!(
gate = state.gate,
error = %err,
"failed to release distributed concurrency permit"
);
}
});
}
}
}
}
#[derive(Debug)]
struct RedisDistributedState {
gate: &'static str,
limit: usize,
client: redis::Client,
key: String,
lease_ttl_ms: u64,
renew_interval_ms: u64,
command_timeout_ms: Option<u64>,
high_watermark: AtomicUsize,
rejected: AtomicU64,
}
impl RedisDistributedState {
async fn try_acquire(
self: &Arc<Self>,
) -> Result<DistributedConcurrencyPermit, DistributedConcurrencyError> {
let token = format!("{}:{}", self.gate, Uuid::new_v4());
let now_ms = unix_time_ms();
let expires_at_ms = now_ms.saturating_add(self.lease_ttl_ms);
let key = self.key.clone();
let result: (i64, i64) = self
.run_with_timeout("acquire", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_err(|err| self.unavailable(format!("connect failed: {err}")))?;
redis::Script::new(
"redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]); \
local count = redis.call('ZCARD', KEYS[1]); \
if count >= tonumber(ARGV[3]) then \
redis.call('PEXPIRE', KEYS[1], ARGV[5]); \
return {0, count}; \
end; \
redis.call('ZADD', KEYS[1], ARGV[2], ARGV[4]); \
count = redis.call('ZCARD', KEYS[1]); \
redis.call('PEXPIRE', KEYS[1], ARGV[5]); \
return {1, count};",
)
.key(&key)
.arg(now_ms as i64)
.arg(expires_at_ms as i64)
.arg(self.limit as i64)
.arg(&token)
.arg(self.lease_ttl_ms as i64)
.invoke_async::<(i64, i64)>(&mut connection)
.await
.map_err(|err| self.unavailable(format!("acquire failed: {err}")))
})
.await?;
let acquired = result.0 > 0;
let in_flight = result.1.max(0) as usize;
self.observe_in_flight(in_flight);
if !acquired {
self.rejected.fetch_add(1, Ordering::Relaxed);
return Err(DistributedConcurrencyError::Saturated {
gate: self.gate,
limit: self.limit,
});
}
let renew_state = Arc::clone(self);
let renew_token = token.clone();
let renew_task = tokio::spawn(async move {
let interval = Duration::from_millis(renew_state.renew_interval_ms);
loop {
tokio::time::sleep(interval).await;
if let Err(err) = renew_state.renew(&renew_token).await {
warn!(
gate = renew_state.gate,
error = %err,
"failed to renew distributed concurrency permit"
);
break;
}
}
});
Ok(DistributedConcurrencyPermit::from_redis(
Arc::clone(self),
token,
renew_task,
))
}
async fn snapshot(
&self,
) -> Result<DistributedConcurrencySnapshot, DistributedConcurrencyError> {
let in_flight = self.live_count().await?;
Ok(DistributedConcurrencySnapshot {
limit: self.limit,
in_flight,
available_permits: self.limit.saturating_sub(in_flight),
high_watermark: self.high_watermark.load(Ordering::Relaxed),
rejected: self.rejected.load(Ordering::Relaxed),
})
}
async fn renew(&self, token: &str) -> Result<(), DistributedConcurrencyError> {
let now_ms = unix_time_ms();
let expires_at_ms = now_ms.saturating_add(self.lease_ttl_ms);
let key = self.key.clone();
let renewed = self
.run_with_timeout("renew", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_err(|err| self.unavailable(format!("connect failed: {err}")))?;
redis::Script::new(
"redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]); \
local score = redis.call('ZSCORE', KEYS[1], ARGV[2]); \
if not score then \
return 0; \
end; \
redis.call('ZADD', KEYS[1], 'XX', ARGV[3], ARGV[2]); \
redis.call('PEXPIRE', KEYS[1], ARGV[4]); \
return 1;",
)
.key(&key)
.arg(now_ms as i64)
.arg(token)
.arg(expires_at_ms as i64)
.arg(self.lease_ttl_ms as i64)
.invoke_async::<i64>(&mut connection)
.await
.map_err(|err| self.unavailable(format!("renew failed: {err}")))
})
.await?;
if renewed == 0 {
return Err(self.unavailable("lease token expired".to_string()));
}
Ok(())
}
async fn release(&self, token: &str) -> Result<(), DistributedConcurrencyError> {
let key = self.key.clone();
self.run_with_timeout("release", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_err(|err| self.unavailable(format!("connect failed: {err}")))?;
redis::Script::new(
"local removed = redis.call('ZREM', KEYS[1], ARGV[1]); \
if removed > 0 and redis.call('ZCARD', KEYS[1]) == 0 then \
redis.call('DEL', KEYS[1]); \
end; \
return removed;",
)
.key(&key)
.arg(token)
.invoke_async::<i64>(&mut connection)
.await
.map_err(|err| self.unavailable(format!("release failed: {err}")))?;
Ok(())
})
.await
}
async fn live_count(&self) -> Result<usize, DistributedConcurrencyError> {
let now_ms = unix_time_ms();
let key = self.key.clone();
let count = self
.run_with_timeout("snapshot", async {
let mut connection = self
.client
.get_multiplexed_async_connection()
.await
.map_err(|err| self.unavailable(format!("connect failed: {err}")))?;
redis::Script::new(
"redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]); \
return redis.call('ZCARD', KEYS[1]);",
)
.key(&key)
.arg(now_ms as i64)
.invoke_async::<i64>(&mut connection)
.await
.map_err(|err| self.unavailable(format!("snapshot failed: {err}")))
})
.await?
.max(0) as usize;
self.observe_in_flight(count);
Ok(count)
}
async fn run_with_timeout<T, F>(
&self,
operation: &'static str,
future: F,
) -> Result<T, DistributedConcurrencyError>
where
F: std::future::Future<Output = Result<T, DistributedConcurrencyError>>,
{
if let Some(timeout_ms) = self.command_timeout_ms {
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
.await
.map_err(|_| {
self.unavailable(format!(
"{operation} exceeded {timeout_ms}ms command timeout"
))
})?
} else {
future.await
}
}
fn unavailable(&self, message: String) -> DistributedConcurrencyError {
DistributedConcurrencyError::Unavailable {
gate: self.gate,
limit: self.limit,
message,
}
}
fn observe_in_flight(&self, in_flight: usize) {
let mut observed = self.high_watermark.load(Ordering::Acquire);
while in_flight > observed {
match self.high_watermark.compare_exchange_weak(
observed,
in_flight,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => break,
Err(next) => observed = next,
}
}
}
}
fn prefixed_key(prefix: Option<&str>, raw_key: &str) -> String {
let prefix = prefix.unwrap_or_default().trim().trim_matches(':');
if prefix.is_empty() {
raw_key.trim_matches(':').to_string()
} else {
format!("{prefix}:{}", raw_key.trim_matches(':'))
}
}
fn unix_time_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
_permit: ConcurrencyPermit,
}
#[cfg(test)]
mod tests {
use super::{
DistributedConcurrencyError, DistributedConcurrencyGate, RedisDistributedConcurrencyConfig,
};
use super::{DistributedConcurrencyError, DistributedConcurrencyGate};
#[tokio::test]
async fn shared_in_memory_gate_rejects_second_acquire() {
@@ -599,46 +168,4 @@ mod tests {
let snapshot = gate.snapshot().await.expect("snapshot should build");
assert_eq!(snapshot.in_flight, 0);
}
#[test]
fn rejects_invalid_redis_config() {
let error = DistributedConcurrencyGate::new_redis(
"shared",
1,
RedisDistributedConcurrencyConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
lease_ttl_ms: 10_000,
renew_interval_ms: 10_000,
command_timeout_ms: Some(1_000),
},
)
.expect_err("equal renew interval should fail");
assert_eq!(
error,
DistributedConcurrencyError::InvalidConfiguration(
"distributed concurrency renew_interval_ms must be smaller than lease_ttl_ms"
.to_string()
)
);
}
#[test]
fn builds_redis_gate_without_touching_network() {
let gate = DistributedConcurrencyGate::new_redis(
"gateway_requests_distributed",
2,
RedisDistributedConcurrencyConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
lease_ttl_ms: 15_000,
renew_interval_ms: 5_000,
command_timeout_ms: Some(1_000),
},
)
.expect("redis gate should build");
assert_eq!(gate.gate(), "gateway_requests_distributed");
assert_eq!(gate.limit(), 2);
}
}

View File

@@ -20,7 +20,7 @@ pub use concurrency::{ConcurrencyError, ConcurrencyGate, ConcurrencyPermit, Conc
pub use config::ServiceRuntimeConfig;
pub use distributed::{
DistributedConcurrencyError, DistributedConcurrencyGate, DistributedConcurrencyPermit,
DistributedConcurrencySnapshot, RedisDistributedConcurrencyConfig,
DistributedConcurrencySnapshot,
};
pub use error::RuntimeBootstrapError;
pub use metrics::{prometheus_response, service_up_sample, MetricKind, MetricLabel, MetricSample};

View File

@@ -13,6 +13,7 @@ aether-contracts.workspace = true
aether-gateway.workspace = true
aether-http.workspace = true
aether-runtime.workspace = true
aether-runtime-state.workspace = true
axum.workspace = true
bytes.workspace = true
http.workspace = true

View File

@@ -6,12 +6,12 @@ use aether_data::driver::postgres::{
DatabaseRecordId, PostgresLeaseClaimOptions, PostgresLeaseClaimSpec, PostgresLeaseRunnerConfig,
PostgresPoolConfig,
};
use aether_data::driver::redis::{
RedisClientConfig, RedisConsumerGroup, RedisConsumerName, RedisLockLease, RedisLockRunner,
RedisLockRunnerConfig, RedisStreamName, RedisStreamReclaimConfig, RedisStreamRunner,
RedisStreamRunnerConfig,
use aether_data::PostgresBackend;
use aether_runtime_state::{
RedisClientConfig, RedisClientFactory, RedisConsumerGroup, RedisConsumerName, RedisKeyspace,
RedisLockLease, RedisLockRunner, RedisLockRunnerConfig, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamRunner, RedisStreamRunnerConfig,
};
use aether_data::{PostgresBackend, RedisBackend};
use aether_testkit::{init_test_runtime_for, ManagedPostgresServer, ManagedRedisServer};
use futures_util::stream::{self, StreamExt};
use serde::Serialize;
@@ -189,10 +189,12 @@ async fn run_suite(
})
.expect("postgres url should resolve");
let redis_backend = RedisBackend::from_config(RedisClientConfig {
let redis_factory = RedisClientFactory::new(RedisClientConfig {
url: redis_url.clone(),
key_prefix: Some(format!("aether-dependency-pressure-{}", std::process::id())),
})?;
let redis_client = redis_factory.connect_lazy()?;
let redis_keyspace = redis_factory.config().keyspace();
let postgres_backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: postgres_url.clone(),
min_connections: 1,
@@ -206,22 +208,30 @@ async fn run_suite(
bootstrap_postgres_lease_table(postgres_backend.pool_clone(), config).await?;
let lock_runner = redis_backend.lock_runner(RedisLockRunnerConfig {
command_timeout_ms: Some(config.timeout.as_millis() as u64),
default_ttl_ms: 5_000,
})?;
let stream_runner = redis_backend.stream_runner(RedisStreamRunnerConfig {
command_timeout_ms: Some(config.timeout.as_millis() as u64),
read_block_ms: Some(10),
read_count: 64,
})?;
let lock_runner = RedisLockRunner::new(
redis_client.clone(),
redis_keyspace.clone(),
RedisLockRunnerConfig {
command_timeout_ms: Some(config.timeout.as_millis() as u64),
default_ttl_ms: 5_000,
},
)?;
let stream_runner = RedisStreamRunner::new(
redis_client.clone(),
redis_keyspace.clone(),
RedisStreamRunnerConfig {
command_timeout_ms: Some(config.timeout.as_millis() as u64),
read_block_ms: Some(10),
read_count: 64,
},
)?;
let lease_runner = postgres_backend.lease_runner(PostgresLeaseRunnerConfig {
statement_timeout_ms: Some(config.timeout.as_millis() as u64),
lock_timeout_ms: Some(1_000),
})?;
let redis_lock = benchmark_redis_lock(&redis_backend, &lock_runner, config).await?;
let redis_stream = benchmark_redis_stream(&redis_backend, &stream_runner, config).await?;
let redis_lock = benchmark_redis_lock(&redis_keyspace, &lock_runner, config).await?;
let redis_stream = benchmark_redis_stream(&redis_keyspace, &stream_runner, config).await?;
let postgres_lease = benchmark_postgres_lease(&lease_runner, config).await?;
Ok(DependencyPressureBaselineReport {
@@ -263,7 +273,7 @@ async fn bootstrap_postgres_lease_table(
}
async fn benchmark_redis_lock(
backend: &RedisBackend,
keyspace: &RedisKeyspace,
runner: &RedisLockRunner,
config: &DependencyPressureBaselineConfig,
) -> Result<RedisLockPressureReport, Box<dyn std::error::Error>> {
@@ -271,7 +281,6 @@ async fn benchmark_redis_lock(
let renew = Arc::new(SummaryCollector::default());
let release = Arc::new(SummaryCollector::default());
let next = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let keyspace = backend.keyspace();
stream::iter(0..config.redis_lock_concurrency)
.for_each_concurrent(config.redis_lock_concurrency, |_| {
@@ -338,11 +347,11 @@ async fn record_redis_lock_follow_up(
}
async fn benchmark_redis_stream(
backend: &RedisBackend,
keyspace: &RedisKeyspace,
runner: &RedisStreamRunner,
config: &DependencyPressureBaselineConfig,
) -> Result<RedisStreamPressureReport, Box<dyn std::error::Error>> {
let stream = backend.keyspace().stream_name("dependency-pressure");
let stream = keyspace.stream_name("dependency-pressure");
let group = RedisConsumerGroup("dependency-group".to_string());
let consumer_a = RedisConsumerName("consumer-a".to_string());
let consumer_b = RedisConsumerName("consumer-b".to_string());

View File

@@ -7,8 +7,10 @@ use aether_data::driver::postgres::{
PostgresLeaseClaimOptions, PostgresLeaseClaimSpec, PostgresLeaseRunnerConfig,
PostgresPoolConfig, PostgresTransactionOptions,
};
use aether_data::driver::redis::{RedisClientConfig, RedisLockRunnerConfig};
use aether_data::{DataLayerError, PostgresBackend, RedisBackend};
use aether_data::{DataLayerError, PostgresBackend};
use aether_runtime_state::{
RedisClientConfig, RedisClientFactory, RedisLockRunner, RedisLockRunnerConfig,
};
use aether_testkit::{
init_test_runtime_for, reserve_local_port, ManagedPostgresServer, ManagedRedisServer,
TunnelHarness, TunnelHarnessConfig,
@@ -235,15 +237,20 @@ async fn benchmark_redis_restart_recovery(
config: &FailureRecoveryBaselineConfig,
) -> Result<RecoverySummary, Box<dyn std::error::Error>> {
let redis_url = redis_server.lock().await.redis_url().to_string();
let backend = RedisBackend::from_config(RedisClientConfig {
let factory = RedisClientFactory::new(RedisClientConfig {
url: redis_url,
key_prefix: Some(format!("aether-failure-recovery-{}", std::process::id())),
})?;
let runner = backend.lock_runner(RedisLockRunnerConfig {
command_timeout_ms: Some(250),
default_ttl_ms: 1_000,
})?;
let keyspace = backend.keyspace();
let client = factory.connect_lazy()?;
let keyspace = factory.config().keyspace();
let runner = RedisLockRunner::new(
client,
keyspace.clone(),
RedisLockRunnerConfig {
command_timeout_ms: Some(250),
default_ttl_ms: 1_000,
},
)?;
let collector = Arc::new(RecoveryCollector::default());
let next_attempt = Arc::new(AtomicUsize::new(0));
let phase = Arc::new(AtomicUsize::new(0));

View File

@@ -4,7 +4,9 @@ use std::sync::Arc;
use std::time::{Duration, Instant};
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_runtime::{DistributedConcurrencyGate, RedisDistributedConcurrencyConfig};
use aether_runtime_state::{
RedisClientConfig, RuntimeSemaphore, RuntimeSemaphoreConfig, RuntimeState,
};
use aether_testkit::{
init_test_runtime_for, run_multi_url_http_load_probe, ExecutionRuntimeHarness,
ExecutionRuntimeHarnessConfig, GatewayHarness, GatewayHarnessConfig, HttpLoadProbeConfig,
@@ -165,13 +167,15 @@ async fn start_gateway_pair(
config.request_limit,
redis_url,
"gateway-a",
)?;
)
.await?;
let gate_b = distributed_request_gate(
"gateway_requests_distributed",
config.request_limit,
redis_url,
"gateway-a",
)?;
"gateway-b",
)
.await?;
let gateway_a = GatewayHarness::start(GatewayHarnessConfig {
upstream_base_url: upstream_base_url.to_string(),
data_config: None,
@@ -209,13 +213,15 @@ async fn start_execution_runtime_pair(
config.request_limit,
redis_url,
"execution-runtime-a",
)?;
)
.await?;
let gate_b = distributed_request_gate(
"execution_runtime_requests_distributed",
config.request_limit,
redis_url,
"execution-runtime-a",
)?;
"execution-runtime-b",
)
.await?;
let runtime_a = ExecutionRuntimeHarness::start(ExecutionRuntimeHarnessConfig {
max_in_flight_requests: None,
distributed_request_gate: Some(gate_a),
@@ -245,13 +251,15 @@ async fn start_tunnel_pair(
config.tunnel_request_limit,
redis_url,
"tunnel-a",
)?;
)
.await?;
let gate_b = distributed_request_gate(
"tunnel_requests_distributed",
config.tunnel_request_limit,
redis_url,
"tunnel-a",
)?;
"tunnel-b",
)
.await?;
let tunnel_a = TunnelHarness::start(TunnelHarnessConfig {
distributed_request_gate: Some(gate_a),
..TunnelHarnessConfig::default()
@@ -279,21 +287,27 @@ async fn start_tunnel_pair(
))
}
fn distributed_request_gate(
async fn distributed_request_gate(
name: &'static str,
limit: usize,
redis_url: &str,
key_scope: &str,
) -> Result<DistributedConcurrencyGate, Box<dyn std::error::Error>> {
Ok(DistributedConcurrencyGate::new_redis(
name,
limit,
RedisDistributedConcurrencyConfig {
) -> Result<RuntimeSemaphore, Box<dyn std::error::Error>> {
let runtime = RuntimeState::redis(
RedisClientConfig {
url: redis_url.to_string(),
key_prefix: Some(format!(
"aether-baseline-{}-{name}-{key_scope}",
std::process::id()
)),
},
Some(1_000),
)
.await?;
Ok(runtime.semaphore(
name,
limit,
RuntimeSemaphoreConfig {
lease_ttl_ms: 30_000,
renew_interval_ms: 10_000,
command_timeout_ms: Some(1_000),

View File

@@ -1,11 +1,10 @@
use std::path::PathBuf;
use std::time::{Duration, Instant};
use aether_data::driver::redis::{
RedisClientConfig, RedisConsumerGroup, RedisConsumerName, RedisStreamReclaimConfig,
RedisStreamRunnerConfig,
use aether_runtime_state::{
RedisClientConfig, RedisClientFactory, RedisConsumerGroup, RedisConsumerName, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamRunner, RedisStreamRunnerConfig,
};
use aether_data::RedisBackend;
use aether_testkit::{init_test_runtime_for, ManagedRedisServer};
use serde::Serialize;
@@ -87,19 +86,25 @@ async fn run_suite(
})
.expect("redis url should be resolved");
let backend = RedisBackend::from_config(RedisClientConfig {
let factory = RedisClientFactory::new(RedisClientConfig {
url: redis_url.clone(),
key_prefix: Some(format!("aether-baseline-{}", std::process::id())),
})?;
let stream = backend.keyspace().stream_name("worker-baseline");
let client = factory.connect_lazy()?;
let keyspace = factory.config().keyspace();
let stream = keyspace.stream_name("worker-baseline");
let group = RedisConsumerGroup("worker-group".to_string());
let consumer_a = RedisConsumerName("consumer-a".to_string());
let consumer_b = RedisConsumerName("consumer-b".to_string());
let runner = backend.stream_runner(RedisStreamRunnerConfig {
command_timeout_ms: Some(2_000),
read_block_ms: Some(10),
read_count: 64,
})?;
let runner = RedisStreamRunner::new(
client,
keyspace,
RedisStreamRunnerConfig {
command_timeout_ms: Some(2_000),
read_block_ms: Some(10),
read_count: 64,
},
)?;
runner
.ensure_consumer_group(&stream, &group, "0-0")
.await
@@ -124,8 +129,8 @@ async fn run_suite(
}
async fn benchmark_append(
runner: &aether_data::driver::redis::RedisStreamRunner,
stream: &aether_data::driver::redis::RedisStreamName,
runner: &RedisStreamRunner,
stream: &RedisStreamName,
config: &RedisWorkerBaselineConfig,
) -> Result<OperationSummary, Box<dyn std::error::Error>> {
let next = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
@@ -182,8 +187,8 @@ async fn benchmark_append(
}
async fn benchmark_read_group(
runner: &aether_data::driver::redis::RedisStreamRunner,
stream: &aether_data::driver::redis::RedisStreamName,
runner: &RedisStreamRunner,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
consumer: &RedisConsumerName,
config: &RedisWorkerBaselineConfig,
@@ -211,8 +216,8 @@ async fn benchmark_read_group(
}
async fn benchmark_reclaim(
runner: &aether_data::driver::redis::RedisStreamRunner,
stream: &aether_data::driver::redis::RedisStreamName,
runner: &RedisStreamRunner,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
consumer_a: &RedisConsumerName,
consumer_b: &RedisConsumerName,
@@ -283,8 +288,8 @@ async fn benchmark_reclaim(
}
async fn benchmark_ack(
runner: &aether_data::driver::redis::RedisStreamRunner,
stream: &aether_data::driver::redis::RedisStreamName,
runner: &RedisStreamRunner,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
ids: &[String],
) -> Result<OperationSummary, Box<dyn std::error::Error>> {

View File

@@ -1,12 +1,12 @@
use aether_gateway::build_execution_runtime_router_with_request_gates;
use aether_runtime::DistributedConcurrencyGate;
use aether_runtime_state::RuntimeSemaphore;
use crate::server::SpawnedServer;
#[derive(Debug, Clone, Default)]
pub struct ExecutionRuntimeHarnessConfig {
pub max_in_flight_requests: Option<usize>,
pub distributed_request_gate: Option<DistributedConcurrencyGate>,
pub distributed_request_gate: Option<RuntimeSemaphore>,
}
#[derive(Debug)]

View File

@@ -1,5 +1,5 @@
use aether_gateway::{build_router_with_state, AppState, GatewayDataConfig};
use aether_runtime::DistributedConcurrencyGate;
use aether_runtime_state::RuntimeSemaphore;
use crate::server::SpawnedServer;
@@ -8,7 +8,7 @@ pub struct GatewayHarnessConfig {
pub upstream_base_url: String,
pub data_config: Option<GatewayDataConfig>,
pub max_in_flight_requests: Option<usize>,
pub distributed_request_gate: Option<DistributedConcurrencyGate>,
pub distributed_request_gate: Option<RuntimeSemaphore>,
pub tunnel_instance_id: Option<String>,
pub tunnel_relay_base_url: Option<String>,
}

View File

@@ -4,7 +4,7 @@ use aether_gateway::{
build_tunnel_runtime_router_with_state, TunnelConnConfig, TunnelControlPlaneClient,
TunnelRuntimeState,
};
use aether_runtime::DistributedConcurrencyGate;
use aether_runtime_state::RuntimeSemaphore;
use crate::server::SpawnedServer;
@@ -14,7 +14,7 @@ pub struct TunnelHarnessConfig {
pub ping_interval: Duration,
pub outbound_queue_capacity: usize,
pub max_in_flight_requests: Option<usize>,
pub distributed_request_gate: Option<DistributedConcurrencyGate>,
pub distributed_request_gate: Option<RuntimeSemaphore>,
}
impl Default for TunnelHarnessConfig {

View File

@@ -11,6 +11,7 @@ aether-ai-formats.workspace = true
aether-contracts.workspace = true
aether-data.workspace = true
aether-data-contracts.workspace = true
aether-runtime-state.workspace = true
async-trait.workspace = true
base64.workspace = true
serde.workspace = true

View File

@@ -1,39 +1,33 @@
use std::sync::Arc;
use serde_json::json;
use aether_data::driver::redis::{
RedisConsumerGroup, RedisConsumerName, RedisStreamEntry, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamRunner, RedisStreamRunnerConfig,
};
use aether_data_contracts::DataLayerError;
use aether_runtime_state::{RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueStore};
use super::config::UsageRuntimeConfig;
use super::event::UsageEvent;
#[derive(Debug, Clone)]
#[derive(Clone)]
pub struct UsageQueue {
runner: RedisStreamRunner,
runner: Arc<dyn RuntimeQueueStore>,
config: UsageRuntimeConfig,
stream: RedisStreamName,
group: RedisConsumerGroup,
dlq_stream: RedisStreamName,
stream: String,
group: String,
dlq_stream: String,
}
impl UsageQueue {
pub fn new(
runner: RedisStreamRunner,
runner: Arc<dyn RuntimeQueueStore>,
config: UsageRuntimeConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
let tuned_runner = RedisStreamRunner::new(
runner.client().clone(),
runner.keyspace().clone(),
usage_stream_runner_config(&config),
)?;
Ok(Self {
runner: tuned_runner,
stream: RedisStreamName(config.stream_key.clone()),
group: RedisConsumerGroup(config.consumer_group.clone()),
dlq_stream: RedisStreamName(config.dlq_stream_key.clone()),
runner,
stream: config.stream_key.clone(),
group: config.consumer_group.clone(),
dlq_stream: config.dlq_stream_key.clone(),
config,
})
}
@@ -53,32 +47,36 @@ impl UsageQueue {
pub async fn read_group(
&self,
consumer: &RedisConsumerName,
) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
consumer: &str,
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
self.runner
.read_group(&self.stream, &self.group, consumer)
.read_group(
&self.stream,
&self.group,
consumer,
self.config.consumer_batch_size.max(1),
Some(self.config.consumer_block_ms.max(1)),
)
.await
}
pub async fn claim_stale(
&self,
consumer: &RedisConsumerName,
consumer: &str,
start_id: &str,
) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
Ok(self
.runner
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
self.runner
.claim_stale(
&self.stream,
&self.group,
consumer,
start_id,
RedisStreamReclaimConfig {
RuntimeQueueReclaimConfig {
min_idle_ms: self.config.reclaim_idle_ms,
count: self.config.reclaim_count,
},
)
.await?
.entries)
.await
}
pub async fn ack_and_delete(&self, ids: &[String]) -> Result<(), DataLayerError> {
@@ -89,27 +87,37 @@ impl UsageQueue {
pub async fn push_dead_letter(
&self,
entry: &RedisStreamEntry,
entry: &RuntimeQueueEntry,
error: &str,
) -> Result<String, DataLayerError> {
let fields = std::collections::BTreeMap::from([(
"payload".to_string(),
serde_json::to_string(&json!({
"entry_id": entry.id,
"fields": entry.fields,
"error": error,
}))
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?,
)]);
self.runner
.append_json(
&self.dlq_stream,
"payload",
&json!({
"entry_id": entry.id,
"fields": entry.fields,
"error": error,
}),
)
.append_fields_with_maxlen(&self.dlq_stream, &fields, None)
.await
}
}
fn usage_stream_runner_config(config: &UsageRuntimeConfig) -> RedisStreamRunnerConfig {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg(test)]
struct UsageQueueRuntimeSettings {
command_timeout_ms: Option<u64>,
read_block_ms: Option<u64>,
read_count: usize,
}
#[cfg(test)]
fn usage_queue_runtime_settings(config: &UsageRuntimeConfig) -> UsageQueueRuntimeSettings {
let read_block_ms = config.consumer_block_ms.max(1);
let command_timeout_ms = read_block_ms.saturating_add(2_000).max(5_000);
RedisStreamRunnerConfig {
UsageQueueRuntimeSettings {
command_timeout_ms: Some(command_timeout_ms),
read_block_ms: Some(read_block_ms),
read_count: config.consumer_batch_size.max(1),
@@ -118,27 +126,10 @@ fn usage_stream_runner_config(config: &UsageRuntimeConfig) -> RedisStreamRunnerC
#[cfg(test)]
mod tests {
use super::{usage_stream_runner_config, UsageQueue};
use super::{usage_queue_runtime_settings, UsageQueue, UsageQueueRuntimeSettings};
use crate::UsageRuntimeConfig;
use aether_data::driver::redis::{RedisClientConfig, RedisClientFactory, RedisStreamRunner};
fn sample_runner() -> RedisStreamRunner {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let client = RedisClientFactory::new(config.clone())
.expect("factory should build")
.connect_lazy()
.expect("client should build");
RedisStreamRunner::new(
client,
config.keyspace(),
aether_data::driver::redis::RedisStreamRunnerConfig::default(),
)
.expect("runner should build")
}
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState};
use std::sync::Arc;
#[test]
fn usage_queue_applies_runtime_block_and_batch_settings() {
@@ -148,15 +139,19 @@ mod tests {
consumer_batch_size: 123,
..UsageRuntimeConfig::default()
};
let queue = UsageQueue::new(sample_runner(), config)
.expect("usage queue should build from runtime config");
let queue = UsageQueue::new(
Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())),
config,
)
.expect("usage queue should build from runtime config");
assert_eq!(
queue.runner.config(),
usage_stream_runner_config(&queue.config)
usage_queue_runtime_settings(&queue.config),
UsageQueueRuntimeSettings {
command_timeout_ms: Some(5_000),
read_block_ms: Some(750),
read_count: 123,
}
);
assert_eq!(queue.runner.config().read_block_ms, Some(750));
assert_eq!(queue.runner.config().read_count, 123);
assert_eq!(queue.runner.config().command_timeout_ms, Some(5_000));
}
}

View File

@@ -3,9 +3,9 @@ use std::pin::Pin;
use std::sync::Arc;
use aether_contracts::ExecutionTelemetry;
use aether_data::driver::redis::RedisStreamRunner;
use aether_data_contracts::repository::usage::UpsertUsageRecord;
use aether_data_contracts::DataLayerError;
use aether_runtime_state::RuntimeQueueStore;
use async_trait::async_trait;
use tracing::warn;
@@ -61,8 +61,8 @@ pub trait UsageRuntimeAccess:
+ Sync
{
fn has_usage_writer(&self) -> bool;
fn has_usage_worker_runner(&self) -> bool;
fn usage_worker_runner(&self) -> Option<RedisStreamRunner>;
fn has_usage_worker_queue(&self) -> bool;
fn usage_worker_queue(&self) -> Option<Arc<dyn RuntimeQueueStore>>;
async fn body_capture_policy(&self) -> Result<UsageBodyCapturePolicy, DataLayerError> {
Ok(UsageBodyCapturePolicy::default())
@@ -104,7 +104,7 @@ impl UsageRuntime {
where
T: UsageRuntimeAccess,
{
self.is_enabled() && data.has_usage_writer() && data.has_usage_worker_runner()
self.is_enabled() && data.has_usage_writer() && data.has_usage_worker_queue()
}
pub fn spawn_worker<T>(&self, data: Arc<T>) -> Option<tokio::task::JoinHandle<()>>
@@ -114,7 +114,7 @@ impl UsageRuntime {
if !self.can_spawn_worker(data.as_ref()) {
return None;
}
let runner = data.usage_worker_runner()?;
let runner = data.usage_worker_queue()?;
let worker = build_usage_queue_worker(runner, data, self.config.clone()).ok()?;
Some(worker.spawn())
}
@@ -352,7 +352,7 @@ impl UsageRuntime {
where
T: UsageRuntimeAccess,
{
if let Some(runner) = data.usage_worker_runner() {
if let Some(runner) = data.usage_worker_queue() {
match UsageQueue::new(runner, self.config.clone()) {
Ok(queue) => match queue.enqueue(&event).await {
Ok(_) => return,
@@ -543,16 +543,14 @@ fn now_unix_secs() -> u64 {
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use std::sync::{Arc, Mutex};
use aether_data::driver::redis::{
RedisClientConfig, RedisClientFactory, RedisStreamRunner, RedisStreamRunnerConfig,
};
use aether_data_contracts::repository::settlement::{
StoredUsageSettlement, UsageSettlementInput,
};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::DataLayerError;
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeQueueStore, RuntimeState};
use async_trait::async_trait;
use serde_json::json;
@@ -571,27 +569,9 @@ mod tests {
records: Mutex<Vec<UpsertUsageRecord>>,
}
struct RedisConfiguredUsageStore {
struct QueueConfiguredUsageStore {
inner: NoRedisUsageStore,
runner: RedisStreamRunner,
}
fn sample_runner() -> RedisStreamRunner {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let client = RedisClientFactory::new(config.clone())
.expect("factory should build")
.connect_lazy()
.expect("client should build");
RedisStreamRunner::new(
client,
config.keyspace(),
RedisStreamRunnerConfig::default(),
)
.expect("runner should build")
queue: Arc<dyn RuntimeQueueStore>,
}
#[async_trait]
@@ -644,17 +624,17 @@ mod tests {
true
}
fn has_usage_worker_runner(&self) -> bool {
fn has_usage_worker_queue(&self) -> bool {
false
}
fn usage_worker_runner(&self) -> Option<RedisStreamRunner> {
fn usage_worker_queue(&self) -> Option<Arc<dyn RuntimeQueueStore>> {
None
}
}
#[async_trait]
impl UsageRecordWriter for RedisConfiguredUsageStore {
impl UsageRecordWriter for QueueConfiguredUsageStore {
async fn upsert_usage_record(
&self,
record: UpsertUsageRecord,
@@ -664,7 +644,7 @@ mod tests {
}
#[async_trait]
impl UsageSettlementWriter for RedisConfiguredUsageStore {
impl UsageSettlementWriter for QueueConfiguredUsageStore {
fn has_usage_settlement_writer(&self) -> bool {
false
}
@@ -678,14 +658,14 @@ mod tests {
}
#[async_trait]
impl UsageBillingEventEnricher for RedisConfiguredUsageStore {
impl UsageBillingEventEnricher for QueueConfiguredUsageStore {
async fn enrich_usage_event(&self, _event: &mut UsageEvent) -> Result<(), DataLayerError> {
Ok(())
}
}
#[async_trait]
impl ManualProxyNodeCounter for RedisConfiguredUsageStore {
impl ManualProxyNodeCounter for QueueConfiguredUsageStore {
async fn increment_manual_proxy_node_requests(
&self,
_node_id: &str,
@@ -697,17 +677,17 @@ mod tests {
}
}
impl UsageRuntimeAccess for RedisConfiguredUsageStore {
impl UsageRuntimeAccess for QueueConfiguredUsageStore {
fn has_usage_writer(&self) -> bool {
true
}
fn has_usage_worker_runner(&self) -> bool {
fn has_usage_worker_queue(&self) -> bool {
true
}
fn usage_worker_runner(&self) -> Option<RedisStreamRunner> {
Some(self.runner.clone())
fn usage_worker_queue(&self) -> Option<Arc<dyn RuntimeQueueStore>> {
Some(Arc::clone(&self.queue))
}
}
@@ -750,10 +730,9 @@ mod tests {
..UsageRuntimeConfig::default()
})
.expect("usage runtime should build");
let runner = sample_runner();
let store = RedisConfiguredUsageStore {
let store = QueueConfiguredUsageStore {
inner: NoRedisUsageStore::default(),
runner,
queue: Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())),
};
let event = UsageEvent::new(
UsageEventType::Failed,

View File

@@ -1,9 +1,9 @@
use std::sync::Arc;
use std::time::Duration;
use aether_data::driver::redis::{RedisConsumerName, RedisStreamEntry, RedisStreamRunner};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::DataLayerError;
use aether_runtime_state::{RuntimeQueueEntry, RuntimeQueueStore};
use async_trait::async_trait;
use tracing::warn;
@@ -60,18 +60,18 @@ where
pub struct UsageQueueWorker {
queue: UsageQueue,
recorder: Arc<dyn UsageEventRecorder>,
consumer: RedisConsumerName,
consumer: String,
config: UsageRuntimeConfig,
}
impl UsageQueueWorker {
pub fn new(
runner: RedisStreamRunner,
runner: Arc<dyn RuntimeQueueStore>,
recorder: Arc<dyn UsageEventRecorder>,
config: UsageRuntimeConfig,
) -> Result<Self, DataLayerError> {
let queue = UsageQueue::new(runner, config.clone())?;
let consumer = RedisConsumerName(consumer_name());
let consumer = consumer_name();
Ok(Self {
queue,
recorder,
@@ -89,7 +89,7 @@ impl UsageQueueWorker {
warn!(
event_name = "usage_worker_consumer_group_failed",
log_type = "ops",
worker_consumer = %self.consumer.0,
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
error = %err,
"usage worker failed to ensure consumer group"
@@ -111,7 +111,7 @@ impl UsageQueueWorker {
warn!(
event_name = "usage_worker_reclaim_process_failed",
log_type = "ops",
worker_consumer = %self.consumer.0,
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
error = %err,
"usage worker failed while reclaiming stale entries"
@@ -121,7 +121,7 @@ impl UsageQueueWorker {
Err(err) => warn!(
event_name = "usage_worker_reclaim_failed",
log_type = "ops",
worker_consumer = %self.consumer.0,
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
error = %err,
"usage worker failed to reclaim stale entries"
@@ -135,7 +135,7 @@ impl UsageQueueWorker {
warn!(
event_name = "usage_worker_process_failed",
log_type = "ops",
worker_consumer = %self.consumer.0,
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
error = %err,
"usage worker failed to process queue entries"
@@ -147,7 +147,7 @@ impl UsageQueueWorker {
warn!(
event_name = "usage_worker_read_failed",
log_type = "ops",
worker_consumer = %self.consumer.0,
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
error = %err,
"usage worker failed to read queue"
@@ -160,7 +160,7 @@ impl UsageQueueWorker {
}
}
async fn process_entries(&self, entries: Vec<RedisStreamEntry>) -> Result<(), DataLayerError> {
async fn process_entries(&self, entries: Vec<RuntimeQueueEntry>) -> Result<(), DataLayerError> {
if entries.is_empty() {
return Ok(());
}
@@ -189,7 +189,7 @@ impl UsageQueueWorker {
Ok(())
}
async fn process_entry(&self, entry: &RedisStreamEntry) -> Result<bool, DataLayerError> {
async fn process_entry(&self, entry: &RuntimeQueueEntry) -> Result<bool, DataLayerError> {
let event = match UsageEvent::from_stream_fields(&entry.fields) {
Ok(event) => event,
Err(err) => {
@@ -204,7 +204,7 @@ impl UsageQueueWorker {
}
pub fn build_usage_queue_worker<T>(
runner: RedisStreamRunner,
runner: Arc<dyn RuntimeQueueStore>,
data: Arc<T>,
config: UsageRuntimeConfig,
) -> Result<UsageQueueWorker, DataLayerError>