mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
refactor: extract runtime state backends
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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>;
|
||||
}
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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},
|
||||
};
|
||||
"#;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(),
|
||||
|
||||
21
crates/aether-runtime-state/Cargo.toml
Normal file
21
crates/aether-runtime-state/Cargo.toml
Normal 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
|
||||
15
crates/aether-runtime-state/src/error.rs
Normal file
15
crates/aether-runtime-state/src/error.rs
Normal 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)
|
||||
}
|
||||
}
|
||||
1879
crates/aether-runtime-state/src/lib.rs
Normal file
1879
crates/aether-runtime-state/src/lib.rs
Normal file
File diff suppressed because it is too large
Load Diff
537
crates/aether-runtime-state/src/memory.rs
Normal file
537
crates/aether-runtime-state/src/memory.rs
Normal 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
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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 {
|
||||
@@ -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 {
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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 {
|
||||
@@ -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 {
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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>> {
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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>
|
||||
|
||||
Reference in New Issue
Block a user