mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat: 引入 aether-runtime/cache/data/http/testkit 基础 crate,完善并发门控与审计系统
新增 crate: - aether-runtime: 服务运行时基础设施(并发门控、分布式并发、指标、队列、优雅关闭、tracing) - aether-cache: 通用 TTL 缓存与命名空间抽象 - aether-data: 数据访问层(PostgreSQL/Redis 后端、repository 模式) - aether-http: HTTP 客户端封装(重试、配置) - aether-testkit: 集成测试工具集(gateway/executor/hub/proxy fixture、等待、负载测试) gateway 扩展: - 引入 audit 模块(shadow 执行审计、决策链路追踪、请求审计 bundle) - 引入 cache 模块(AuthContext 缓存、direct-plan bypass 缓存) - 引入 data 模块(auth/candidates/config/usage/video_tasks 数据访问) - 集成 ConcurrencyGate/DistributedConcurrencyGate 请求门控 - 新增本地 auth 拒绝、过载响应构建器 - 补充 control/auth_cache/video/concurrency 集成测试 aether-proxy 扩展: - AppState 集成 stream_gate / distributed_stream_gate 并发门控 - 新增 ProxyAdmissionError 及准入拒绝流程 - stream_handler 补充门控饱和/不可用场景测试 - 配置与注册客户端逻辑完善 aether-hub 扩展: - main.rs 引入运行时初始化、指标端点、健康检查 - local_relay 重构为 lib.rs 暴露公共接口
This commit is contained in:
9
crates/aether-cache/Cargo.toml
Normal file
9
crates/aether-cache/Cargo.toml
Normal file
@@ -0,0 +1,9 @@
|
||||
[package]
|
||||
name = "aether-cache"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Shared in-memory cache primitives for Aether Rust services"
|
||||
|
||||
[dependencies]
|
||||
5
crates/aether-cache/src/lib.rs
Normal file
5
crates/aether-cache/src/lib.rs
Normal file
@@ -0,0 +1,5 @@
|
||||
mod namespace;
|
||||
mod ttl_map;
|
||||
|
||||
pub use namespace::CacheKeyNamespace;
|
||||
pub use ttl_map::ExpiringMap;
|
||||
51
crates/aether-cache/src/namespace.rs
Normal file
51
crates/aether-cache/src/namespace.rs
Normal file
@@ -0,0 +1,51 @@
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct CacheKeyNamespace {
|
||||
prefix: String,
|
||||
}
|
||||
|
||||
impl CacheKeyNamespace {
|
||||
pub fn new(prefix: impl Into<String>) -> Self {
|
||||
Self {
|
||||
prefix: prefix.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn child(&self, suffix: &str) -> Self {
|
||||
if self.prefix.is_empty() {
|
||||
return Self::new(suffix);
|
||||
}
|
||||
if suffix.is_empty() {
|
||||
return self.clone();
|
||||
}
|
||||
Self::new(format!("{}:{}", self.prefix, suffix))
|
||||
}
|
||||
|
||||
pub fn key(&self, raw_key: &str) -> String {
|
||||
if self.prefix.is_empty() {
|
||||
return raw_key.to_string();
|
||||
}
|
||||
if raw_key.is_empty() {
|
||||
return self.prefix.clone();
|
||||
}
|
||||
format!("{}:{}", self.prefix, raw_key)
|
||||
}
|
||||
|
||||
pub fn prefix(&self) -> &str {
|
||||
&self.prefix
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::CacheKeyNamespace;
|
||||
|
||||
#[test]
|
||||
fn composes_scoped_keys() {
|
||||
let root = CacheKeyNamespace::new("aether");
|
||||
let child = root.child("auth");
|
||||
|
||||
assert_eq!(root.key("user-1"), "aether:user-1");
|
||||
assert_eq!(child.key("user-1"), "aether:auth:user-1");
|
||||
assert_eq!(child.prefix(), "aether:auth");
|
||||
}
|
||||
}
|
||||
175
crates/aether-cache/src/ttl_map.rs
Normal file
175
crates/aether-cache/src/ttl_map.rs
Normal file
@@ -0,0 +1,175 @@
|
||||
use std::collections::HashMap;
|
||||
use std::hash::Hash;
|
||||
use std::sync::Mutex;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct TimedEntry<V> {
|
||||
value: V,
|
||||
inserted_at: Instant,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ExpiringMap<K, V> {
|
||||
entries: Mutex<HashMap<K, TimedEntry<V>>>,
|
||||
}
|
||||
|
||||
impl<K, V> Default for ExpiringMap<K, V> {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
entries: Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<K, V> ExpiringMap<K, V>
|
||||
where
|
||||
K: Eq + Hash + Clone,
|
||||
{
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn insert(&self, key: K, value: V, ttl: Duration, max_entries: usize) {
|
||||
let Ok(mut entries) = self.entries.lock() else {
|
||||
return;
|
||||
};
|
||||
|
||||
prune_expired(&mut entries, ttl);
|
||||
while max_entries > 0 && entries.len() >= max_entries {
|
||||
let Some(oldest_key) = entries
|
||||
.iter()
|
||||
.min_by_key(|(_, entry)| entry.inserted_at)
|
||||
.map(|(key, _)| key.clone())
|
||||
else {
|
||||
break;
|
||||
};
|
||||
entries.remove(&oldest_key);
|
||||
}
|
||||
|
||||
entries.insert(
|
||||
key,
|
||||
TimedEntry {
|
||||
value,
|
||||
inserted_at: Instant::now(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
pub fn remove(&self, key: &K) -> Option<V> {
|
||||
let Ok(mut entries) = self.entries.lock() else {
|
||||
return None;
|
||||
};
|
||||
entries.remove(key).map(|entry| entry.value)
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.entries
|
||||
.lock()
|
||||
.map(|entries| entries.len())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
}
|
||||
|
||||
impl<K, V> ExpiringMap<K, V>
|
||||
where
|
||||
K: Eq + Hash + Clone,
|
||||
V: Clone,
|
||||
{
|
||||
pub fn get_fresh(&self, key: &K, ttl: Duration) -> Option<V> {
|
||||
let Ok(mut entries) = self.entries.lock() else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let Some(entry) = entries.get(key).cloned() else {
|
||||
return None;
|
||||
};
|
||||
|
||||
if entry.inserted_at.elapsed() > ttl {
|
||||
entries.remove(key);
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(entry.value)
|
||||
}
|
||||
|
||||
pub fn contains_fresh(&self, key: &K, ttl: Duration) -> bool {
|
||||
self.get_fresh(key, ttl).is_some()
|
||||
}
|
||||
}
|
||||
|
||||
fn prune_expired<K, V>(entries: &mut HashMap<K, TimedEntry<V>>, ttl: Duration)
|
||||
where
|
||||
K: Eq + Hash,
|
||||
{
|
||||
if ttl.is_zero() {
|
||||
entries.clear();
|
||||
return;
|
||||
}
|
||||
entries.retain(|_, entry| entry.inserted_at.elapsed() <= ttl);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::thread::sleep;
|
||||
|
||||
use super::ExpiringMap;
|
||||
|
||||
#[test]
|
||||
fn evicts_expired_entries_on_read() {
|
||||
let cache = ExpiringMap::new();
|
||||
cache.insert(
|
||||
"hello".to_string(),
|
||||
42_u32,
|
||||
std::time::Duration::from_millis(10),
|
||||
16,
|
||||
);
|
||||
|
||||
sleep(std::time::Duration::from_millis(20));
|
||||
|
||||
assert_eq!(
|
||||
cache.get_fresh(&"hello".to_string(), std::time::Duration::from_millis(10)),
|
||||
None
|
||||
);
|
||||
assert_eq!(cache.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evicts_oldest_entry_when_capacity_is_hit() {
|
||||
let cache = ExpiringMap::new();
|
||||
|
||||
cache.insert(
|
||||
"one".to_string(),
|
||||
1_u32,
|
||||
std::time::Duration::from_secs(60),
|
||||
2,
|
||||
);
|
||||
sleep(std::time::Duration::from_millis(2));
|
||||
cache.insert(
|
||||
"two".to_string(),
|
||||
2_u32,
|
||||
std::time::Duration::from_secs(60),
|
||||
2,
|
||||
);
|
||||
sleep(std::time::Duration::from_millis(2));
|
||||
cache.insert(
|
||||
"three".to_string(),
|
||||
3_u32,
|
||||
std::time::Duration::from_secs(60),
|
||||
2,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
cache.get_fresh(&"one".to_string(), std::time::Duration::from_secs(60)),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
cache.get_fresh(&"two".to_string(), std::time::Duration::from_secs(60)),
|
||||
Some(2)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.get_fresh(&"three".to_string(), std::time::Duration::from_secs(60)),
|
||||
Some(3)
|
||||
);
|
||||
}
|
||||
}
|
||||
20
crates/aether-data/Cargo.toml
Normal file
20
crates/aether-data/Cargo.toml
Normal file
@@ -0,0 +1,20 @@
|
||||
[package]
|
||||
name = "aether-data"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Shared data access contracts and config for Aether Rust services"
|
||||
|
||||
[dependencies]
|
||||
aether-cache.workspace = true
|
||||
async-trait.workspace = true
|
||||
futures-util.workspace = true
|
||||
redis.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sqlx.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
url.workspace = true
|
||||
uuid.workspace = true
|
||||
66
crates/aether-data/src/backends/leases.rs
Normal file
66
crates/aether-data/src/backends/leases.rs
Normal file
@@ -0,0 +1,66 @@
|
||||
use std::fmt;
|
||||
|
||||
use super::PostgresBackend;
|
||||
use crate::postgres::{PostgresLeaseRunner, PostgresLeaseRunnerConfig};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct DataLeaseBackends {
|
||||
postgres: Option<PostgresLeaseRunner>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for DataLeaseBackends {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("DataLeaseBackends")
|
||||
.field("has_postgres", &self.postgres.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl DataLeaseBackends {
|
||||
pub(crate) fn from_postgres(
|
||||
postgres: Option<&PostgresBackend>,
|
||||
) -> Result<Self, DataLayerError> {
|
||||
Ok(Self {
|
||||
postgres: postgres
|
||||
.map(|backend| backend.lease_runner(PostgresLeaseRunnerConfig::default()))
|
||||
.transpose()?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn postgres(&self) -> Option<PostgresLeaseRunner> {
|
||||
self.postgres.clone()
|
||||
}
|
||||
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.postgres.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataLeaseBackends;
|
||||
use crate::backends::PostgresBackend;
|
||||
use crate::postgres::PostgresPoolConfig;
|
||||
|
||||
#[tokio::test]
|
||||
async fn builds_postgres_lease_runner_from_backend() {
|
||||
let backend = PostgresBackend::from_config(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("postgres backend should build");
|
||||
|
||||
let leases =
|
||||
DataLeaseBackends::from_postgres(Some(&backend)).expect("lease backends should build");
|
||||
|
||||
assert!(leases.has_any());
|
||||
assert!(leases.postgres().is_some());
|
||||
}
|
||||
}
|
||||
58
crates/aether-data/src/backends/locks.rs
Normal file
58
crates/aether-data/src/backends/locks.rs
Normal file
@@ -0,0 +1,58 @@
|
||||
use std::fmt;
|
||||
|
||||
use super::RedisBackend;
|
||||
use crate::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::backends::RedisBackend;
|
||||
use crate::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());
|
||||
}
|
||||
}
|
||||
195
crates/aether-data/src/backends/mod.rs
Normal file
195
crates/aether-data/src/backends/mod.rs
Normal file
@@ -0,0 +1,195 @@
|
||||
mod leases;
|
||||
mod locks;
|
||||
mod postgres;
|
||||
mod read;
|
||||
mod redis;
|
||||
mod transactions;
|
||||
mod workers;
|
||||
mod write;
|
||||
|
||||
pub use leases::DataLeaseBackends;
|
||||
pub use locks::DataLockBackends;
|
||||
pub use postgres::PostgresBackend;
|
||||
pub use read::DataReadRepositories;
|
||||
pub use redis::RedisBackend;
|
||||
pub use transactions::DataTransactionBackends;
|
||||
pub use workers::DataWorkerBackends;
|
||||
pub use write::DataWriteRepositories;
|
||||
|
||||
use crate::{DataLayerConfig, DataLayerError};
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct DataBackends {
|
||||
config: DataLayerConfig,
|
||||
postgres: Option<PostgresBackend>,
|
||||
redis: Option<RedisBackend>,
|
||||
leases: DataLeaseBackends,
|
||||
locks: DataLockBackends,
|
||||
read: DataReadRepositories,
|
||||
transactions: DataTransactionBackends,
|
||||
workers: DataWorkerBackends,
|
||||
write: DataWriteRepositories,
|
||||
}
|
||||
|
||||
impl DataBackends {
|
||||
pub fn from_config(config: DataLayerConfig) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
|
||||
let postgres = config
|
||||
.postgres
|
||||
.clone()
|
||||
.map(PostgresBackend::from_config)
|
||||
.transpose()?;
|
||||
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_postgres(postgres.as_ref());
|
||||
let transactions = DataTransactionBackends::from_postgres(postgres.as_ref());
|
||||
let workers = DataWorkerBackends::from_redis(redis.as_ref())?;
|
||||
let write = DataWriteRepositories::from_postgres(postgres.as_ref());
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
postgres,
|
||||
redis,
|
||||
leases,
|
||||
locks,
|
||||
read,
|
||||
transactions,
|
||||
workers,
|
||||
write,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &DataLayerConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
pub fn postgres(&self) -> Option<&PostgresBackend> {
|
||||
self.postgres.as_ref()
|
||||
}
|
||||
|
||||
pub fn redis(&self) -> Option<&RedisBackend> {
|
||||
self.redis.as_ref()
|
||||
}
|
||||
|
||||
pub fn read(&self) -> &DataReadRepositories {
|
||||
&self.read
|
||||
}
|
||||
|
||||
pub fn leases(&self) -> &DataLeaseBackends {
|
||||
&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
|
||||
}
|
||||
|
||||
pub fn has_runtime_backends(&self) -> bool {
|
||||
self.postgres.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()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataBackends;
|
||||
use crate::{postgres::PostgresPoolConfig, DataLayerConfig};
|
||||
|
||||
#[test]
|
||||
fn builds_empty_backends_from_default_config() {
|
||||
let backends = DataBackends::from_config(DataLayerConfig::default())
|
||||
.expect("empty config should be accepted");
|
||||
|
||||
assert!(!backends.has_runtime_backends());
|
||||
assert!(backends.postgres().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().request_candidates().is_none());
|
||||
assert!(backends.read().provider_catalog().is_none());
|
||||
assert!(backends.read().usage().is_none());
|
||||
assert!(backends.read().video_tasks().is_none());
|
||||
assert!(backends.read().shadow_results().is_none());
|
||||
assert!(backends.transactions().postgres().is_none());
|
||||
assert!(backends.workers().redis().is_none());
|
||||
assert!(backends.write().shadow_results().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn builds_postgres_backend_from_config() {
|
||||
let backends = DataBackends::from_config(DataLayerConfig {
|
||||
postgres: Some(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
}),
|
||||
redis: None,
|
||||
})
|
||||
.expect("postgres backend should build");
|
||||
|
||||
assert!(backends.has_runtime_backends());
|
||||
assert!(backends.postgres().is_some());
|
||||
assert!(backends.leases().postgres().is_some());
|
||||
assert!(backends.read().auth_api_keys().is_some());
|
||||
assert!(backends.read().request_candidates().is_some());
|
||||
assert!(backends.read().provider_catalog().is_some());
|
||||
assert!(backends.read().usage().is_some());
|
||||
assert!(backends.read().video_tasks().is_some());
|
||||
assert!(backends.read().shadow_results().is_some());
|
||||
assert!(backends.transactions().postgres().is_some());
|
||||
assert!(backends.write().shadow_results().is_some());
|
||||
assert!(backends.config().postgres.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_redis_backend_from_config() {
|
||||
let backends = DataBackends::from_config(DataLayerConfig {
|
||||
postgres: None,
|
||||
redis: Some(crate::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.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.transactions().postgres().is_none());
|
||||
assert!(backends.write().shadow_results().is_none());
|
||||
assert!(backends.config().redis.is_some());
|
||||
}
|
||||
}
|
||||
123
crates/aether-data/src/backends/postgres.rs
Normal file
123
crates/aether-data/src/backends/postgres.rs
Normal file
@@ -0,0 +1,123 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::postgres::{
|
||||
PostgresLeaseRunner, PostgresLeaseRunnerConfig, PostgresPool, PostgresPoolConfig,
|
||||
PostgresPoolFactory, PostgresTransactionRunner,
|
||||
};
|
||||
use crate::repository::auth::{AuthApiKeyReadRepository, SqlxAuthApiKeySnapshotReadRepository};
|
||||
use crate::repository::candidates::{
|
||||
RequestCandidateReadRepository, SqlxRequestCandidateReadRepository,
|
||||
};
|
||||
use crate::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, SqlxProviderCatalogReadRepository,
|
||||
};
|
||||
use crate::repository::shadow_results::{
|
||||
ShadowResultReadRepository, ShadowResultWriteRepository, SqlxShadowResultRepository,
|
||||
};
|
||||
use crate::repository::usage::{SqlxUsageReadRepository, UsageReadRepository};
|
||||
use crate::repository::video_tasks::{SqlxVideoTaskReadRepository, VideoTaskReadRepository};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostgresBackend {
|
||||
config: PostgresPoolConfig,
|
||||
pool: PostgresPool,
|
||||
}
|
||||
|
||||
impl PostgresBackend {
|
||||
pub fn from_config(config: PostgresPoolConfig) -> Result<Self, DataLayerError> {
|
||||
let factory = PostgresPoolFactory::new(config.clone())?;
|
||||
let pool = factory.connect_lazy()?;
|
||||
|
||||
Ok(Self { config, pool })
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &PostgresPoolConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PostgresPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub fn pool_clone(&self) -> PostgresPool {
|
||||
self.pool.clone()
|
||||
}
|
||||
|
||||
pub fn auth_api_key_read_repository(&self) -> Arc<dyn AuthApiKeyReadRepository> {
|
||||
Arc::new(SqlxAuthApiKeySnapshotReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn request_candidate_read_repository(&self) -> Arc<dyn RequestCandidateReadRepository> {
|
||||
Arc::new(SqlxRequestCandidateReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_catalog_read_repository(&self) -> Arc<dyn ProviderCatalogReadRepository> {
|
||||
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn usage_read_repository(&self) -> Arc<dyn UsageReadRepository> {
|
||||
Arc::new(SqlxUsageReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn video_task_read_repository(&self) -> Arc<dyn VideoTaskReadRepository> {
|
||||
Arc::new(SqlxVideoTaskReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn transaction_runner(&self) -> PostgresTransactionRunner {
|
||||
PostgresTransactionRunner::new(self.pool_clone())
|
||||
}
|
||||
|
||||
pub fn lease_runner(
|
||||
&self,
|
||||
config: PostgresLeaseRunnerConfig,
|
||||
) -> Result<PostgresLeaseRunner, DataLayerError> {
|
||||
PostgresLeaseRunner::new(self.transaction_runner(), config)
|
||||
}
|
||||
|
||||
pub fn shadow_result_write_repository(&self) -> Arc<dyn ShadowResultWriteRepository> {
|
||||
Arc::new(SqlxShadowResultRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn shadow_result_read_repository(&self) -> Arc<dyn ShadowResultReadRepository> {
|
||||
Arc::new(SqlxShadowResultRepository::new(self.pool_clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::PostgresBackend;
|
||||
use crate::postgres::{PostgresLeaseRunnerConfig, PostgresPoolConfig};
|
||||
|
||||
#[tokio::test]
|
||||
async fn backend_retains_config_and_pool() {
|
||||
let config = PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
};
|
||||
|
||||
let backend =
|
||||
PostgresBackend::from_config(config.clone()).expect("backend should build lazily");
|
||||
|
||||
assert_eq!(backend.config(), &config);
|
||||
let _pool = backend.pool();
|
||||
let _pool_clone = backend.pool_clone();
|
||||
let _auth_api_key_reader = backend.auth_api_key_read_repository();
|
||||
let _request_candidate_reader = backend.request_candidate_read_repository();
|
||||
let _provider_catalog_reader = backend.provider_catalog_read_repository();
|
||||
let _usage_reader = backend.usage_read_repository();
|
||||
let _video_task_reader = backend.video_task_read_repository();
|
||||
let _transaction_runner = backend.transaction_runner();
|
||||
let _lease_runner = backend
|
||||
.lease_runner(PostgresLeaseRunnerConfig::default())
|
||||
.expect("lease runner should build");
|
||||
let _shadow_result_reader = backend.shadow_result_read_repository();
|
||||
let _shadow_result_writer = backend.shadow_result_write_repository();
|
||||
}
|
||||
}
|
||||
111
crates/aether-data/src/backends/read.rs
Normal file
111
crates/aether-data/src/backends/read.rs
Normal file
@@ -0,0 +1,111 @@
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::PostgresBackend;
|
||||
use crate::repository::auth::AuthApiKeyReadRepository;
|
||||
use crate::repository::candidates::RequestCandidateReadRepository;
|
||||
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
|
||||
use crate::repository::shadow_results::ShadowResultReadRepository;
|
||||
use crate::repository::usage::UsageReadRepository;
|
||||
use crate::repository::video_tasks::VideoTaskReadRepository;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct DataReadRepositories {
|
||||
auth_api_keys: Option<Arc<dyn AuthApiKeyReadRepository>>,
|
||||
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
|
||||
provider_catalog: Option<Arc<dyn ProviderCatalogReadRepository>>,
|
||||
usage: Option<Arc<dyn UsageReadRepository>>,
|
||||
video_tasks: Option<Arc<dyn VideoTaskReadRepository>>,
|
||||
shadow_results: Option<Arc<dyn ShadowResultReadRepository>>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for DataReadRepositories {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("DataReadRepositories")
|
||||
.field("has_auth_api_keys", &self.auth_api_keys.is_some())
|
||||
.field("has_request_candidates", &self.request_candidates.is_some())
|
||||
.field("has_provider_catalog", &self.provider_catalog.is_some())
|
||||
.field("has_usage", &self.usage.is_some())
|
||||
.field("has_video_tasks", &self.video_tasks.is_some())
|
||||
.field("has_shadow_results", &self.shadow_results.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl DataReadRepositories {
|
||||
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
|
||||
Self {
|
||||
auth_api_keys: postgres.map(PostgresBackend::auth_api_key_read_repository),
|
||||
request_candidates: postgres.map(PostgresBackend::request_candidate_read_repository),
|
||||
provider_catalog: postgres.map(PostgresBackend::provider_catalog_read_repository),
|
||||
usage: postgres.map(PostgresBackend::usage_read_repository),
|
||||
video_tasks: postgres.map(PostgresBackend::video_task_read_repository),
|
||||
shadow_results: postgres.map(PostgresBackend::shadow_result_read_repository),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn auth_api_keys(&self) -> Option<Arc<dyn AuthApiKeyReadRepository>> {
|
||||
self.auth_api_keys.clone()
|
||||
}
|
||||
|
||||
pub fn request_candidates(&self) -> Option<Arc<dyn RequestCandidateReadRepository>> {
|
||||
self.request_candidates.clone()
|
||||
}
|
||||
|
||||
pub fn provider_catalog(&self) -> Option<Arc<dyn ProviderCatalogReadRepository>> {
|
||||
self.provider_catalog.clone()
|
||||
}
|
||||
|
||||
pub fn usage(&self) -> Option<Arc<dyn UsageReadRepository>> {
|
||||
self.usage.clone()
|
||||
}
|
||||
|
||||
pub fn video_tasks(&self) -> Option<Arc<dyn VideoTaskReadRepository>> {
|
||||
self.video_tasks.clone()
|
||||
}
|
||||
|
||||
pub fn shadow_results(&self) -> Option<Arc<dyn ShadowResultReadRepository>> {
|
||||
self.shadow_results.clone()
|
||||
}
|
||||
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.auth_api_keys.is_some()
|
||||
|| self.request_candidates.is_some()
|
||||
|| self.provider_catalog.is_some()
|
||||
|| self.usage.is_some()
|
||||
|| self.video_tasks.is_some()
|
||||
|| self.shadow_results.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataReadRepositories;
|
||||
use crate::backends::PostgresBackend;
|
||||
use crate::postgres::PostgresPoolConfig;
|
||||
|
||||
#[tokio::test]
|
||||
async fn builds_read_repositories_from_postgres_backend() {
|
||||
let backend = PostgresBackend::from_config(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("postgres backend should build");
|
||||
|
||||
let read = DataReadRepositories::from_postgres(Some(&backend));
|
||||
|
||||
assert!(read.has_any());
|
||||
assert!(read.auth_api_keys().is_some());
|
||||
assert!(read.request_candidates().is_some());
|
||||
assert!(read.provider_catalog().is_some());
|
||||
assert!(read.usage().is_some());
|
||||
assert!(read.video_tasks().is_some());
|
||||
assert!(read.shadow_results().is_some());
|
||||
}
|
||||
}
|
||||
76
crates/aether-data/src/backends/redis.rs
Normal file
76
crates/aether-data/src/backends/redis.rs
Normal file
@@ -0,0 +1,76 @@
|
||||
use crate::redis::{
|
||||
RedisClient, RedisClientConfig, RedisClientFactory, RedisKeyspace, 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)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::RedisBackend;
|
||||
use crate::redis::{RedisClientConfig, 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");
|
||||
}
|
||||
}
|
||||
60
crates/aether-data/src/backends/transactions.rs
Normal file
60
crates/aether-data/src/backends/transactions.rs
Normal file
@@ -0,0 +1,60 @@
|
||||
use std::fmt;
|
||||
|
||||
use super::PostgresBackend;
|
||||
use crate::postgres::PostgresTransactionRunner;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct DataTransactionBackends {
|
||||
postgres: Option<PostgresTransactionRunner>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for DataTransactionBackends {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("DataTransactionBackends")
|
||||
.field("has_postgres", &self.postgres.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl DataTransactionBackends {
|
||||
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
|
||||
Self {
|
||||
postgres: postgres.map(PostgresBackend::transaction_runner),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn postgres(&self) -> Option<PostgresTransactionRunner> {
|
||||
self.postgres.clone()
|
||||
}
|
||||
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.postgres.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataTransactionBackends;
|
||||
use crate::backends::PostgresBackend;
|
||||
use crate::postgres::PostgresPoolConfig;
|
||||
|
||||
#[tokio::test]
|
||||
async fn builds_postgres_transaction_runner_from_backend() {
|
||||
let backend = PostgresBackend::from_config(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("postgres backend should build");
|
||||
|
||||
let transactions = DataTransactionBackends::from_postgres(Some(&backend));
|
||||
|
||||
assert!(transactions.has_any());
|
||||
assert!(transactions.postgres().is_some());
|
||||
}
|
||||
}
|
||||
58
crates/aether-data/src/backends/workers.rs
Normal file
58
crates/aether-data/src/backends/workers.rs
Normal file
@@ -0,0 +1,58 @@
|
||||
use std::fmt;
|
||||
|
||||
use super::RedisBackend;
|
||||
use crate::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::backends::RedisBackend;
|
||||
use crate::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());
|
||||
}
|
||||
}
|
||||
61
crates/aether-data/src/backends/write.rs
Normal file
61
crates/aether-data/src/backends/write.rs
Normal file
@@ -0,0 +1,61 @@
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::PostgresBackend;
|
||||
use crate::repository::shadow_results::ShadowResultWriteRepository;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct DataWriteRepositories {
|
||||
shadow_results: Option<Arc<dyn ShadowResultWriteRepository>>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for DataWriteRepositories {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("DataWriteRepositories")
|
||||
.field("has_shadow_results", &self.shadow_results.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl DataWriteRepositories {
|
||||
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
|
||||
Self {
|
||||
shadow_results: postgres.map(PostgresBackend::shadow_result_write_repository),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn shadow_results(&self) -> Option<Arc<dyn ShadowResultWriteRepository>> {
|
||||
self.shadow_results.clone()
|
||||
}
|
||||
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.shadow_results.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataWriteRepositories;
|
||||
use crate::backends::PostgresBackend;
|
||||
use crate::postgres::PostgresPoolConfig;
|
||||
|
||||
#[tokio::test]
|
||||
async fn builds_shadow_result_writer_from_postgres_backend() {
|
||||
let backend = PostgresBackend::from_config(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("postgres backend should build");
|
||||
|
||||
let write = DataWriteRepositories::from_postgres(Some(&backend));
|
||||
|
||||
assert!(write.has_any());
|
||||
assert!(write.shadow_results().is_some());
|
||||
}
|
||||
}
|
||||
74
crates/aether-data/src/config.rs
Normal file
74
crates/aether-data/src/config.rs
Normal file
@@ -0,0 +1,74 @@
|
||||
use crate::postgres::PostgresPoolConfig;
|
||||
use crate::redis::RedisClientConfig;
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
|
||||
pub struct DataLayerConfig {
|
||||
pub postgres: Option<PostgresPoolConfig>,
|
||||
pub redis: Option<RedisClientConfig>,
|
||||
}
|
||||
|
||||
impl DataLayerConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
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.postgres.is_some() || self.redis.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::DataLayerConfig;
|
||||
use crate::postgres::PostgresPoolConfig;
|
||||
use crate::redis::RedisClientConfig;
|
||||
|
||||
#[test]
|
||||
fn validates_nested_backend_configs() {
|
||||
let config = DataLayerConfig {
|
||||
postgres: Some(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 2,
|
||||
max_connections: 8,
|
||||
acquire_timeout_ms: 1_500,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
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());
|
||||
assert!(config.has_persistent_backends());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_nested_backend_configs() {
|
||||
let config = DataLayerConfig {
|
||||
postgres: Some(PostgresPoolConfig {
|
||||
database_url: String::new(),
|
||||
min_connections: 4,
|
||||
max_connections: 2,
|
||||
acquire_timeout_ms: 1_500,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
}),
|
||||
redis: None,
|
||||
};
|
||||
|
||||
assert!(config.validate().is_err());
|
||||
}
|
||||
}
|
||||
20
crates/aether-data/src/error.rs
Normal file
20
crates/aether-data/src/error.rs
Normal file
@@ -0,0 +1,20 @@
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum DataLayerError {
|
||||
#[error("invalid configuration: {0}")]
|
||||
InvalidConfiguration(String),
|
||||
|
||||
#[error("invalid input: {0}")]
|
||||
InvalidInput(String),
|
||||
|
||||
#[error("postgres error: {0}")]
|
||||
Postgres(#[from] sqlx::Error),
|
||||
|
||||
#[error("redis error: {0}")]
|
||||
Redis(#[from] redis::RedisError),
|
||||
|
||||
#[error("operation timed out: {0}")]
|
||||
TimedOut(String),
|
||||
|
||||
#[error("unexpected database value: {0}")]
|
||||
UnexpectedValue(String),
|
||||
}
|
||||
14
crates/aether-data/src/lib.rs
Normal file
14
crates/aether-data/src/lib.rs
Normal file
@@ -0,0 +1,14 @@
|
||||
pub mod backends;
|
||||
mod config;
|
||||
mod error;
|
||||
pub mod postgres;
|
||||
pub mod redis;
|
||||
pub mod repository;
|
||||
|
||||
pub use backends::{
|
||||
DataBackends, DataLeaseBackends, DataLockBackends, DataReadRepositories,
|
||||
DataTransactionBackends, DataWorkerBackends, DataWriteRepositories, PostgresBackend,
|
||||
RedisBackend,
|
||||
};
|
||||
pub use config::DataLayerConfig;
|
||||
pub use error::DataLayerError;
|
||||
417
crates/aether-data/src/postgres/lease.rs
Normal file
417
crates/aether-data/src/postgres/lease.rs
Normal file
@@ -0,0 +1,417 @@
|
||||
use crate::postgres::{DatabaseRecordId, PostgresTransactionOptions, PostgresTransactionRunner};
|
||||
use crate::DataLayerError;
|
||||
use futures_util::FutureExt;
|
||||
use sqlx::query_scalar;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct PostgresLeaseClaimOptions {
|
||||
pub batch_size: usize,
|
||||
pub lease_ms: u64,
|
||||
}
|
||||
|
||||
impl PostgresLeaseClaimOptions {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if self.batch_size == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"postgres lease batch_size must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.lease_ms == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"postgres lease lease_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PostgresLeaseClaimSpec {
|
||||
pub table: &'static str,
|
||||
pub id_column: &'static str,
|
||||
pub lease_owner_column: &'static str,
|
||||
pub lease_expires_at_column: &'static str,
|
||||
pub eligibility_predicate_sql: &'static str,
|
||||
pub order_by_sql: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct PostgresLeaseRunnerConfig {
|
||||
pub statement_timeout_ms: Option<u64>,
|
||||
pub lock_timeout_ms: Option<u64>,
|
||||
}
|
||||
|
||||
impl PostgresLeaseRunnerConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if matches!(self.statement_timeout_ms, Some(0)) {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"postgres lease statement_timeout_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if matches!(self.lock_timeout_ms, Some(0)) {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"postgres lease lock_timeout_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostgresLeaseRunner {
|
||||
transaction_runner: PostgresTransactionRunner,
|
||||
config: PostgresLeaseRunnerConfig,
|
||||
}
|
||||
|
||||
impl PostgresLeaseRunner {
|
||||
pub fn new(
|
||||
transaction_runner: PostgresTransactionRunner,
|
||||
config: PostgresLeaseRunnerConfig,
|
||||
) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
Ok(Self {
|
||||
transaction_runner,
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn config(&self) -> PostgresLeaseRunnerConfig {
|
||||
self.config
|
||||
}
|
||||
|
||||
pub fn transaction_runner(&self) -> &PostgresTransactionRunner {
|
||||
&self.transaction_runner
|
||||
}
|
||||
|
||||
pub async fn claim_ids(
|
||||
&self,
|
||||
spec: &PostgresLeaseClaimSpec,
|
||||
options: PostgresLeaseClaimOptions,
|
||||
owner: &str,
|
||||
) -> Result<Vec<DatabaseRecordId>, DataLayerError> {
|
||||
validate_lease_owner(owner)?;
|
||||
let sql = build_postgres_lease_claim_sql(spec, options)?;
|
||||
let owner = owner.trim().to_string();
|
||||
let lease_ms = i64::try_from(options.lease_ms).map_err(|_| {
|
||||
DataLayerError::InvalidInput("postgres lease lease_ms exceeds i64 range".to_string())
|
||||
})?;
|
||||
let tx_options = PostgresTransactionOptions {
|
||||
statement_timeout_ms: self.config.statement_timeout_ms,
|
||||
lock_timeout_ms: self.config.lock_timeout_ms,
|
||||
..PostgresTransactionOptions::read_write()
|
||||
};
|
||||
|
||||
self.transaction_runner
|
||||
.run(tx_options, |tx| {
|
||||
async move {
|
||||
let rows = query_scalar::<_, String>(&sql)
|
||||
.bind(owner)
|
||||
.bind(lease_ms)
|
||||
.fetch_all(&mut **tx)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(DatabaseRecordId).collect())
|
||||
}
|
||||
.boxed()
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn release_ids(
|
||||
&self,
|
||||
spec: &PostgresLeaseClaimSpec,
|
||||
ids: &[DatabaseRecordId],
|
||||
owner: &str,
|
||||
) -> Result<Vec<DatabaseRecordId>, DataLayerError> {
|
||||
validate_lease_owner(owner)?;
|
||||
if ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let sql = build_postgres_lease_release_sql(spec)?;
|
||||
let owner = owner.trim().to_string();
|
||||
let ids = ids.iter().map(|id| id.0.clone()).collect::<Vec<_>>();
|
||||
let tx_options = PostgresTransactionOptions {
|
||||
statement_timeout_ms: self.config.statement_timeout_ms,
|
||||
lock_timeout_ms: self.config.lock_timeout_ms,
|
||||
..PostgresTransactionOptions::read_write()
|
||||
};
|
||||
|
||||
self.transaction_runner
|
||||
.run(tx_options, |tx| {
|
||||
async move {
|
||||
let rows = query_scalar::<_, String>(&sql)
|
||||
.bind(ids)
|
||||
.bind(owner)
|
||||
.fetch_all(&mut **tx)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(DatabaseRecordId).collect())
|
||||
}
|
||||
.boxed()
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn renew_ids(
|
||||
&self,
|
||||
spec: &PostgresLeaseClaimSpec,
|
||||
ids: &[DatabaseRecordId],
|
||||
owner: &str,
|
||||
lease_ms: u64,
|
||||
) -> Result<Vec<DatabaseRecordId>, DataLayerError> {
|
||||
validate_lease_owner(owner)?;
|
||||
if lease_ms == 0 {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"postgres lease lease_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let sql = build_postgres_lease_renew_sql(spec)?;
|
||||
let owner = owner.trim().to_string();
|
||||
let lease_ms = i64::try_from(lease_ms).map_err(|_| {
|
||||
DataLayerError::InvalidInput("postgres lease lease_ms exceeds i64 range".to_string())
|
||||
})?;
|
||||
let ids = ids.iter().map(|id| id.0.clone()).collect::<Vec<_>>();
|
||||
let tx_options = PostgresTransactionOptions {
|
||||
statement_timeout_ms: self.config.statement_timeout_ms,
|
||||
lock_timeout_ms: self.config.lock_timeout_ms,
|
||||
..PostgresTransactionOptions::read_write()
|
||||
};
|
||||
|
||||
self.transaction_runner
|
||||
.run(tx_options, |tx| {
|
||||
async move {
|
||||
let rows = query_scalar::<_, String>(&sql)
|
||||
.bind(ids)
|
||||
.bind(owner)
|
||||
.bind(lease_ms)
|
||||
.fetch_all(&mut **tx)
|
||||
.await?;
|
||||
Ok(rows.into_iter().map(DatabaseRecordId).collect())
|
||||
}
|
||||
.boxed()
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_postgres_lease_claim_sql(
|
||||
spec: &PostgresLeaseClaimSpec,
|
||||
options: PostgresLeaseClaimOptions,
|
||||
) -> Result<String, DataLayerError> {
|
||||
options.validate()?;
|
||||
validate_lease_spec(spec)?;
|
||||
|
||||
Ok(format!(
|
||||
"WITH claimable AS (\
|
||||
SELECT {id_column} \
|
||||
FROM {table} \
|
||||
WHERE ({eligibility_predicate_sql}) \
|
||||
AND ({lease_expires_at_column} IS NULL OR {lease_expires_at_column} <= NOW()) \
|
||||
ORDER BY {order_by_sql} \
|
||||
FOR UPDATE SKIP LOCKED \
|
||||
LIMIT {batch_size}\
|
||||
) \
|
||||
UPDATE {table} AS target \
|
||||
SET {lease_owner_column} = $1, \
|
||||
{lease_expires_at_column} = NOW() + ($2::bigint * INTERVAL '1 millisecond') \
|
||||
FROM claimable \
|
||||
WHERE target.{id_column} = claimable.{id_column} \
|
||||
RETURNING target.{id_column}",
|
||||
id_column = spec.id_column,
|
||||
table = spec.table,
|
||||
eligibility_predicate_sql = spec.eligibility_predicate_sql,
|
||||
lease_expires_at_column = spec.lease_expires_at_column,
|
||||
order_by_sql = spec.order_by_sql,
|
||||
batch_size = options.batch_size,
|
||||
lease_owner_column = spec.lease_owner_column,
|
||||
))
|
||||
}
|
||||
|
||||
pub fn build_postgres_lease_release_sql(
|
||||
spec: &PostgresLeaseClaimSpec,
|
||||
) -> Result<String, DataLayerError> {
|
||||
validate_lease_spec(spec)?;
|
||||
|
||||
Ok(format!(
|
||||
"UPDATE {table} \
|
||||
SET {lease_owner_column} = NULL, \
|
||||
{lease_expires_at_column} = NULL \
|
||||
WHERE {id_column} = ANY($1) \
|
||||
AND {lease_owner_column} = $2 \
|
||||
RETURNING {id_column}",
|
||||
table = spec.table,
|
||||
id_column = spec.id_column,
|
||||
lease_owner_column = spec.lease_owner_column,
|
||||
lease_expires_at_column = spec.lease_expires_at_column,
|
||||
))
|
||||
}
|
||||
|
||||
pub fn build_postgres_lease_renew_sql(
|
||||
spec: &PostgresLeaseClaimSpec,
|
||||
) -> Result<String, DataLayerError> {
|
||||
validate_lease_spec(spec)?;
|
||||
|
||||
Ok(format!(
|
||||
"UPDATE {table} \
|
||||
SET {lease_expires_at_column} = NOW() + ($3::bigint * INTERVAL '1 millisecond') \
|
||||
WHERE {id_column} = ANY($1) \
|
||||
AND {lease_owner_column} = $2 \
|
||||
RETURNING {id_column}",
|
||||
table = spec.table,
|
||||
id_column = spec.id_column,
|
||||
lease_owner_column = spec.lease_owner_column,
|
||||
lease_expires_at_column = spec.lease_expires_at_column,
|
||||
))
|
||||
}
|
||||
|
||||
fn validate_lease_spec(spec: &PostgresLeaseClaimSpec) -> Result<(), DataLayerError> {
|
||||
for (field, value) in [
|
||||
("table", spec.table),
|
||||
("id_column", spec.id_column),
|
||||
("lease_owner_column", spec.lease_owner_column),
|
||||
("lease_expires_at_column", spec.lease_expires_at_column),
|
||||
("eligibility_predicate_sql", spec.eligibility_predicate_sql),
|
||||
("order_by_sql", spec.order_by_sql),
|
||||
] {
|
||||
if value.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidConfiguration(format!(
|
||||
"postgres lease {field} cannot be empty"
|
||||
)));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_lease_owner(owner: &str) -> Result<(), DataLayerError> {
|
||||
if owner.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"postgres lease owner cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
build_postgres_lease_claim_sql, build_postgres_lease_release_sql,
|
||||
build_postgres_lease_renew_sql, PostgresLeaseClaimOptions, PostgresLeaseClaimSpec,
|
||||
PostgresLeaseRunner, PostgresLeaseRunnerConfig,
|
||||
};
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory, PostgresTransactionRunner};
|
||||
|
||||
fn sample_spec() -> PostgresLeaseClaimSpec {
|
||||
PostgresLeaseClaimSpec {
|
||||
table: "video_tasks",
|
||||
id_column: "id",
|
||||
lease_owner_column: "lease_owner",
|
||||
lease_expires_at_column: "lease_expires_at",
|
||||
eligibility_predicate_sql: "status IN ('submitted', 'processing')",
|
||||
order_by_sql: "updated_at ASC",
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_skip_locked_claim_sql() {
|
||||
let sql = build_postgres_lease_claim_sql(
|
||||
&sample_spec(),
|
||||
PostgresLeaseClaimOptions {
|
||||
batch_size: 16,
|
||||
lease_ms: 15_000,
|
||||
},
|
||||
)
|
||||
.expect("claim SQL should build");
|
||||
|
||||
assert!(sql.contains("FOR UPDATE SKIP LOCKED"));
|
||||
assert!(sql.contains("LIMIT 16"));
|
||||
assert!(sql.contains("lease_owner = $1"));
|
||||
assert!(sql.contains("NOW() + ($2::bigint * INTERVAL '1 millisecond')"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_release_sql() {
|
||||
let sql =
|
||||
build_postgres_lease_release_sql(&sample_spec()).expect("release SQL should build");
|
||||
|
||||
assert!(sql.contains("id = ANY($1)"));
|
||||
assert!(sql.contains("lease_owner = $2"));
|
||||
assert!(sql.contains("lease_expires_at = NULL"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_renew_sql() {
|
||||
let sql = build_postgres_lease_renew_sql(&sample_spec()).expect("renew SQL should build");
|
||||
|
||||
assert!(sql.contains("id = ANY($1)"));
|
||||
assert!(sql.contains("lease_owner = $2"));
|
||||
assert!(sql.contains("NOW() + ($3::bigint * INTERVAL '1 millisecond')"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lease_runner_reuses_transaction_runner() {
|
||||
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("factory should build");
|
||||
let pool = factory.connect_lazy().expect("lazy pool should build");
|
||||
let transaction_runner = PostgresTransactionRunner::new(pool);
|
||||
|
||||
let lease_runner = PostgresLeaseRunner::new(
|
||||
transaction_runner.clone(),
|
||||
PostgresLeaseRunnerConfig {
|
||||
statement_timeout_ms: Some(2_000),
|
||||
lock_timeout_ms: Some(500),
|
||||
},
|
||||
)
|
||||
.expect("lease runner should build");
|
||||
|
||||
assert_eq!(
|
||||
lease_runner.config(),
|
||||
PostgresLeaseRunnerConfig {
|
||||
statement_timeout_ms: Some(2_000),
|
||||
lock_timeout_ms: Some(500),
|
||||
}
|
||||
);
|
||||
let _runner_ref = lease_runner.transaction_runner();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn release_and_renew_empty_ids_short_circuit() {
|
||||
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("factory should build");
|
||||
let pool = factory.connect_lazy().expect("lazy pool should build");
|
||||
let runner = PostgresLeaseRunner::new(
|
||||
PostgresTransactionRunner::new(pool),
|
||||
PostgresLeaseRunnerConfig::default(),
|
||||
)
|
||||
.expect("lease runner should build");
|
||||
|
||||
assert!(runner
|
||||
.release_ids(&sample_spec(), &[], "worker-1")
|
||||
.await
|
||||
.expect("empty release should succeed")
|
||||
.is_empty());
|
||||
assert!(runner
|
||||
.renew_ids(&sample_spec(), &[], "worker-1", 5_000)
|
||||
.await
|
||||
.expect("empty renew should succeed")
|
||||
.is_empty());
|
||||
}
|
||||
}
|
||||
15
crates/aether-data/src/postgres/mod.rs
Normal file
15
crates/aether-data/src/postgres/mod.rs
Normal file
@@ -0,0 +1,15 @@
|
||||
mod lease;
|
||||
mod pool;
|
||||
mod tx;
|
||||
mod types;
|
||||
|
||||
pub use lease::{
|
||||
build_postgres_lease_claim_sql, build_postgres_lease_release_sql,
|
||||
build_postgres_lease_renew_sql, PostgresLeaseClaimOptions, PostgresLeaseClaimSpec,
|
||||
PostgresLeaseRunner, PostgresLeaseRunnerConfig,
|
||||
};
|
||||
pub use pool::{PostgresPool, PostgresPoolConfig, PostgresPoolFactory};
|
||||
pub use tx::{
|
||||
PostgresTransaction, PostgresTransactionOptions, PostgresTransactionRunner, TransactionMode,
|
||||
};
|
||||
pub use types::DatabaseRecordId;
|
||||
122
crates/aether-data/src/postgres/pool.rs
Normal file
122
crates/aether-data/src/postgres/pool.rs
Normal file
@@ -0,0 +1,122 @@
|
||||
use crate::DataLayerError;
|
||||
use sqlx::postgres::{PgConnectOptions, PgPoolOptions, PgSslMode};
|
||||
use sqlx::PgPool;
|
||||
use std::str::FromStr;
|
||||
use std::time::Duration;
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
|
||||
pub struct PostgresPoolConfig {
|
||||
pub database_url: String,
|
||||
pub min_connections: u32,
|
||||
pub max_connections: u32,
|
||||
pub acquire_timeout_ms: u64,
|
||||
pub idle_timeout_ms: u64,
|
||||
pub max_lifetime_ms: u64,
|
||||
pub statement_cache_capacity: usize,
|
||||
pub require_ssl: bool,
|
||||
}
|
||||
|
||||
impl Default for PostgresPoolConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
database_url: String::new(),
|
||||
min_connections: 1,
|
||||
max_connections: 20,
|
||||
acquire_timeout_ms: 5_000,
|
||||
idle_timeout_ms: 60_000,
|
||||
max_lifetime_ms: 30 * 60_000,
|
||||
statement_cache_capacity: 100,
|
||||
require_ssl: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PostgresPoolConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if self.database_url.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"postgres database_url cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.min_connections > self.max_connections {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"postgres min_connections cannot exceed max_connections".to_string(),
|
||||
));
|
||||
}
|
||||
if self.statement_cache_capacity == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"postgres statement_cache_capacity must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn connect_options(&self) -> Result<PgConnectOptions, DataLayerError> {
|
||||
self.validate()?;
|
||||
|
||||
let ssl_mode = if self.require_ssl {
|
||||
PgSslMode::Require
|
||||
} else {
|
||||
PgSslMode::Prefer
|
||||
};
|
||||
|
||||
let options = PgConnectOptions::from_str(self.database_url.trim()).map_err(|err| {
|
||||
DataLayerError::InvalidConfiguration(format!("invalid postgres database_url: {err}"))
|
||||
})?;
|
||||
|
||||
Ok(options
|
||||
.ssl_mode(ssl_mode)
|
||||
.statement_cache_capacity(self.statement_cache_capacity))
|
||||
}
|
||||
}
|
||||
|
||||
pub type PostgresPool = PgPool;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostgresPoolFactory {
|
||||
config: PostgresPoolConfig,
|
||||
}
|
||||
|
||||
impl PostgresPoolFactory {
|
||||
pub fn new(config: PostgresPoolConfig) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
Ok(Self { config })
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &PostgresPoolConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
pub fn connect_lazy(&self) -> Result<PostgresPool, DataLayerError> {
|
||||
let options = self.config.connect_options()?;
|
||||
Ok(PgPoolOptions::new()
|
||||
.min_connections(self.config.min_connections)
|
||||
.max_connections(self.config.max_connections)
|
||||
.acquire_timeout(Duration::from_millis(self.config.acquire_timeout_ms))
|
||||
.idle_timeout(Duration::from_millis(self.config.idle_timeout_ms))
|
||||
.max_lifetime(Duration::from_millis(self.config.max_lifetime_ms))
|
||||
.connect_lazy_with(options))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
|
||||
#[tokio::test]
|
||||
async fn factory_builds_lazy_pool_from_valid_config() {
|
||||
let config = PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
};
|
||||
|
||||
let factory = PostgresPoolFactory::new(config).expect("factory should build");
|
||||
let _pool = factory.connect_lazy().expect("lazy pool should build");
|
||||
}
|
||||
}
|
||||
202
crates/aether-data/src/postgres/tx.rs
Normal file
202
crates/aether-data/src/postgres/tx.rs
Normal file
@@ -0,0 +1,202 @@
|
||||
use futures_util::future::BoxFuture;
|
||||
use sqlx::{Postgres, Transaction};
|
||||
|
||||
use crate::postgres::PostgresPool;
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum TransactionMode {
|
||||
ReadOnly,
|
||||
#[default]
|
||||
ReadWrite,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct PostgresTransactionOptions {
|
||||
pub mode: TransactionMode,
|
||||
pub statement_timeout_ms: Option<u64>,
|
||||
pub lock_timeout_ms: Option<u64>,
|
||||
}
|
||||
|
||||
impl PostgresTransactionOptions {
|
||||
pub fn read_only() -> Self {
|
||||
Self {
|
||||
mode: TransactionMode::ReadOnly,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn read_write() -> Self {
|
||||
Self {
|
||||
mode: TransactionMode::ReadWrite,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if matches!(self.statement_timeout_ms, Some(0)) {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"postgres statement_timeout_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if matches!(self.lock_timeout_ms, Some(0)) {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"postgres lock_timeout_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub type PostgresTransaction = Transaction<'static, Postgres>;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostgresTransactionRunner {
|
||||
pool: PostgresPool,
|
||||
}
|
||||
|
||||
impl PostgresTransactionRunner {
|
||||
pub fn new(pool: PostgresPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PostgresPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn begin(
|
||||
&self,
|
||||
options: PostgresTransactionOptions,
|
||||
) -> Result<PostgresTransaction, DataLayerError> {
|
||||
options.validate()?;
|
||||
|
||||
let mut tx = self.pool.begin().await?;
|
||||
for statement in build_transaction_setup_statements(options) {
|
||||
sqlx::query(statement.as_str()).execute(&mut *tx).await?;
|
||||
}
|
||||
Ok(tx)
|
||||
}
|
||||
|
||||
pub async fn run<T, F>(
|
||||
&self,
|
||||
options: PostgresTransactionOptions,
|
||||
f: F,
|
||||
) -> Result<T, DataLayerError>
|
||||
where
|
||||
F: for<'tx> FnOnce(
|
||||
&'tx mut PostgresTransaction,
|
||||
) -> BoxFuture<'tx, Result<T, DataLayerError>>,
|
||||
{
|
||||
let mut tx = self.begin(options).await?;
|
||||
match f(&mut tx).await {
|
||||
Ok(value) => {
|
||||
tx.commit().await?;
|
||||
Ok(value)
|
||||
}
|
||||
Err(err) => {
|
||||
let _ = tx.rollback().await;
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_read_only<T, F>(&self, f: F) -> Result<T, DataLayerError>
|
||||
where
|
||||
F: for<'tx> FnOnce(
|
||||
&'tx mut PostgresTransaction,
|
||||
) -> BoxFuture<'tx, Result<T, DataLayerError>>,
|
||||
{
|
||||
self.run(PostgresTransactionOptions::read_only(), f).await
|
||||
}
|
||||
|
||||
pub async fn run_read_write<T, F>(&self, f: F) -> Result<T, DataLayerError>
|
||||
where
|
||||
F: for<'tx> FnOnce(
|
||||
&'tx mut PostgresTransaction,
|
||||
) -> BoxFuture<'tx, Result<T, DataLayerError>>,
|
||||
{
|
||||
self.run(PostgresTransactionOptions::read_write(), f).await
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_transaction_setup_statements(
|
||||
options: PostgresTransactionOptions,
|
||||
) -> Vec<String> {
|
||||
let mut statements = Vec::new();
|
||||
if options.mode == TransactionMode::ReadOnly {
|
||||
statements.push("SET TRANSACTION READ ONLY".to_string());
|
||||
}
|
||||
if let Some(statement_timeout_ms) = options.statement_timeout_ms {
|
||||
statements.push(format!(
|
||||
"SET LOCAL statement_timeout = {}",
|
||||
statement_timeout_ms
|
||||
));
|
||||
}
|
||||
if let Some(lock_timeout_ms) = options.lock_timeout_ms {
|
||||
statements.push(format!("SET LOCAL lock_timeout = {}", lock_timeout_ms));
|
||||
}
|
||||
statements
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
build_transaction_setup_statements, PostgresTransactionOptions, PostgresTransactionRunner,
|
||||
TransactionMode,
|
||||
};
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
|
||||
#[test]
|
||||
fn validates_transaction_options() {
|
||||
assert!(PostgresTransactionOptions {
|
||||
statement_timeout_ms: Some(0),
|
||||
..PostgresTransactionOptions::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
assert!(PostgresTransactionOptions {
|
||||
lock_timeout_ms: Some(0),
|
||||
..PostgresTransactionOptions::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_expected_setup_statements() {
|
||||
let statements = build_transaction_setup_statements(PostgresTransactionOptions {
|
||||
mode: TransactionMode::ReadOnly,
|
||||
statement_timeout_ms: Some(1_500),
|
||||
lock_timeout_ms: Some(750),
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
statements,
|
||||
vec![
|
||||
"SET TRANSACTION READ ONLY".to_string(),
|
||||
"SET LOCAL statement_timeout = 1500".to_string(),
|
||||
"SET LOCAL lock_timeout = 750".to_string(),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn runner_reuses_lazy_pool() {
|
||||
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("factory should build");
|
||||
let pool = factory.connect_lazy().expect("lazy pool should build");
|
||||
|
||||
let runner = PostgresTransactionRunner::new(pool.clone());
|
||||
|
||||
let _pool_ref = runner.pool();
|
||||
}
|
||||
}
|
||||
2
crates/aether-data/src/postgres/types.rs
Normal file
2
crates/aether-data/src/postgres/types.rs
Normal file
@@ -0,0 +1,2 @@
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct DatabaseRecordId(pub String);
|
||||
68
crates/aether-data/src/redis/client.rs
Normal file
68
crates/aether-data/src/redis/client.rs
Normal file
@@ -0,0 +1,68 @@
|
||||
use crate::redis::RedisKeyspace;
|
||||
use crate::DataLayerError;
|
||||
|
||||
pub type RedisClient = redis::Client;
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
|
||||
pub struct RedisClientConfig {
|
||||
pub url: String,
|
||||
pub key_prefix: Option<String>,
|
||||
}
|
||||
|
||||
impl RedisClientConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
let raw = self.url.trim();
|
||||
if raw.is_empty() {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis url cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
url::Url::parse(raw).map_err(|err| {
|
||||
DataLayerError::InvalidConfiguration(format!("invalid redis url: {err}"))
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn keyspace(&self) -> RedisKeyspace {
|
||||
RedisKeyspace::new(self.key_prefix.as_deref())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RedisClientFactory {
|
||||
config: RedisClientConfig,
|
||||
}
|
||||
|
||||
impl RedisClientFactory {
|
||||
pub fn new(config: RedisClientConfig) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
Ok(Self { config })
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &RedisClientConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
pub fn connect_lazy(&self) -> Result<RedisClient, DataLayerError> {
|
||||
Ok(RedisClient::open(self.config.url.clone())?)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{RedisClientConfig, RedisClientFactory};
|
||||
|
||||
#[test]
|
||||
fn factory_builds_lazy_client_from_valid_config() {
|
||||
let config = RedisClientConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
};
|
||||
let factory = RedisClientFactory::new(config.clone()).expect("factory should build");
|
||||
|
||||
assert_eq!(factory.config(), &config);
|
||||
let _client = factory
|
||||
.connect_lazy()
|
||||
.expect("lazy redis client should build");
|
||||
}
|
||||
}
|
||||
307
crates/aether-data/src/redis/lock.rs
Normal file
307
crates/aether-data/src/redis/lock.rs
Normal file
@@ -0,0 +1,307 @@
|
||||
use std::future::Future;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::redis::{RedisClient, RedisKeyspace};
|
||||
use crate::DataLayerError;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct RedisLockKey(pub String);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RedisLockLease {
|
||||
pub key: RedisLockKey,
|
||||
pub owner: String,
|
||||
pub token: String,
|
||||
pub ttl_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct RedisLockRunnerConfig {
|
||||
pub command_timeout_ms: Option<u64>,
|
||||
pub default_ttl_ms: u64,
|
||||
}
|
||||
|
||||
impl Default for RedisLockRunnerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
command_timeout_ms: Some(1_000),
|
||||
default_ttl_ms: 15_000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisLockRunnerConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if matches!(self.command_timeout_ms, Some(0)) {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis lock command_timeout_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.default_ttl_ms == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis lock default_ttl_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RedisLockRunner {
|
||||
client: RedisClient,
|
||||
keyspace: RedisKeyspace,
|
||||
config: RedisLockRunnerConfig,
|
||||
}
|
||||
|
||||
impl RedisLockRunner {
|
||||
pub fn new(
|
||||
client: RedisClient,
|
||||
keyspace: RedisKeyspace,
|
||||
config: RedisLockRunnerConfig,
|
||||
) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
Ok(Self {
|
||||
client,
|
||||
keyspace,
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn client(&self) -> &RedisClient {
|
||||
&self.client
|
||||
}
|
||||
|
||||
pub fn keyspace(&self) -> &RedisKeyspace {
|
||||
&self.keyspace
|
||||
}
|
||||
|
||||
pub fn config(&self) -> RedisLockRunnerConfig {
|
||||
self.config
|
||||
}
|
||||
|
||||
pub async fn try_acquire(
|
||||
&self,
|
||||
key: &RedisLockKey,
|
||||
owner: &str,
|
||||
ttl_ms: Option<u64>,
|
||||
) -> Result<Option<RedisLockLease>, DataLayerError> {
|
||||
validate_owner(owner)?;
|
||||
validate_key(key)?;
|
||||
let ttl_ms = self.resolve_ttl_ms(ttl_ms)?;
|
||||
let token = format!("{owner}:{}", Uuid::new_v4());
|
||||
|
||||
self.run_with_timeout("redis lock acquire", async {
|
||||
let mut connection = self.client.get_multiplexed_async_connection().await?;
|
||||
let status = redis::cmd("SET")
|
||||
.arg(&key.0)
|
||||
.arg(&token)
|
||||
.arg("NX")
|
||||
.arg("PX")
|
||||
.arg(ttl_ms)
|
||||
.query_async::<Option<String>>(&mut connection)
|
||||
.await?;
|
||||
|
||||
Ok(status.map(|_| RedisLockLease {
|
||||
key: key.clone(),
|
||||
owner: owner.to_string(),
|
||||
token,
|
||||
ttl_ms,
|
||||
}))
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn release(&self, lease: &RedisLockLease) -> Result<bool, DataLayerError> {
|
||||
validate_lease(lease)?;
|
||||
self.run_with_timeout("redis lock release", async {
|
||||
let mut connection = self.client.get_multiplexed_async_connection().await?;
|
||||
let deleted = redis::Script::new(
|
||||
"if redis.call('get', KEYS[1]) == ARGV[1] then \
|
||||
return redis.call('del', KEYS[1]) \
|
||||
else \
|
||||
return 0 \
|
||||
end",
|
||||
)
|
||||
.key(&lease.key.0)
|
||||
.arg(&lease.token)
|
||||
.invoke_async::<i32>(&mut connection)
|
||||
.await?;
|
||||
Ok(deleted > 0)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn renew(
|
||||
&self,
|
||||
lease: &RedisLockLease,
|
||||
ttl_ms: Option<u64>,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
validate_lease(lease)?;
|
||||
let ttl_ms = self.resolve_ttl_ms(ttl_ms)?;
|
||||
|
||||
self.run_with_timeout("redis lock renew", async {
|
||||
let mut connection = self.client.get_multiplexed_async_connection().await?;
|
||||
let renewed = redis::Script::new(
|
||||
"if redis.call('get', KEYS[1]) == ARGV[1] then \
|
||||
return redis.call('pexpire', KEYS[1], ARGV[2]) \
|
||||
else \
|
||||
return 0 \
|
||||
end",
|
||||
)
|
||||
.key(&lease.key.0)
|
||||
.arg(&lease.token)
|
||||
.arg(ttl_ms)
|
||||
.invoke_async::<i32>(&mut connection)
|
||||
.await?;
|
||||
Ok(renewed > 0)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_with_timeout<T, F>(
|
||||
&self,
|
||||
operation: &'static str,
|
||||
future: F,
|
||||
) -> Result<T, DataLayerError>
|
||||
where
|
||||
F: Future<Output = Result<T, DataLayerError>>,
|
||||
{
|
||||
if let Some(timeout_ms) = self.config.command_timeout_ms {
|
||||
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
|
||||
})?
|
||||
} else {
|
||||
future.await
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_ttl_ms(&self, ttl_ms: Option<u64>) -> Result<u64, DataLayerError> {
|
||||
let ttl_ms = ttl_ms.unwrap_or(self.config.default_ttl_ms);
|
||||
if ttl_ms == 0 {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis lock ttl_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(ttl_ms)
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_owner(owner: &str) -> Result<(), DataLayerError> {
|
||||
if owner.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis lock owner cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_key(key: &RedisLockKey) -> Result<(), DataLayerError> {
|
||||
if key.0.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis lock key cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_lease(lease: &RedisLockLease) -> Result<(), DataLayerError> {
|
||||
validate_key(&lease.key)?;
|
||||
validate_owner(&lease.owner)?;
|
||||
if lease.token.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis lock token cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if lease.ttl_ms == 0 {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis lock ttl_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
|
||||
use crate::redis::{RedisClientConfig, RedisClientFactory};
|
||||
|
||||
fn sample_runner() -> RedisLockRunner {
|
||||
let client = RedisClientFactory::new(RedisClientConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
})
|
||||
.expect("factory should build")
|
||||
.connect_lazy()
|
||||
.expect("client should build");
|
||||
|
||||
RedisLockRunner::new(
|
||||
client,
|
||||
RedisClientConfig {
|
||||
url: "redis://127.0.0.1/0".to_string(),
|
||||
key_prefix: Some("aether".to_string()),
|
||||
}
|
||||
.keyspace(),
|
||||
RedisLockRunnerConfig::default(),
|
||||
)
|
||||
.expect("runner should build")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_runner_config() {
|
||||
assert!(RedisLockRunnerConfig {
|
||||
command_timeout_ms: Some(0),
|
||||
..RedisLockRunnerConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
assert!(RedisLockRunnerConfig {
|
||||
default_ttl_ms: 0,
|
||||
..RedisLockRunnerConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runner_reuses_client_and_keyspace() {
|
||||
let runner = sample_runner();
|
||||
|
||||
assert_eq!(runner.config(), RedisLockRunnerConfig::default());
|
||||
assert_eq!(runner.keyspace().lock_key("poller").0, "aether:lock:poller");
|
||||
let _client_ref = runner.client();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_invalid_owner_or_lease_before_network() {
|
||||
let runner = sample_runner();
|
||||
|
||||
assert!(runner
|
||||
.try_acquire(&RedisLockKey("aether:lock:poller".to_string()), "", None)
|
||||
.await
|
||||
.is_err());
|
||||
assert!(runner
|
||||
.release(&RedisLockLease {
|
||||
key: RedisLockKey("aether:lock:poller".to_string()),
|
||||
owner: "worker-1".to_string(),
|
||||
token: String::new(),
|
||||
ttl_ms: 1_000,
|
||||
})
|
||||
.await
|
||||
.is_err());
|
||||
assert!(runner
|
||||
.renew(
|
||||
&RedisLockLease {
|
||||
key: RedisLockKey("aether:lock:poller".to_string()),
|
||||
owner: "worker-1".to_string(),
|
||||
token: "token-1".to_string(),
|
||||
ttl_ms: 1_000,
|
||||
},
|
||||
Some(0),
|
||||
)
|
||||
.await
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
12
crates/aether-data/src/redis/mod.rs
Normal file
12
crates/aether-data/src/redis/mod.rs
Normal file
@@ -0,0 +1,12 @@
|
||||
mod client;
|
||||
mod lock;
|
||||
mod namespace;
|
||||
mod stream;
|
||||
|
||||
pub use client::{RedisClient, RedisClientConfig, RedisClientFactory};
|
||||
pub use lock::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
|
||||
pub use namespace::RedisKeyspace;
|
||||
pub use stream::{
|
||||
RedisConsumerGroup, RedisConsumerName, RedisStreamEntry, RedisStreamName,
|
||||
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner, RedisStreamRunnerConfig,
|
||||
};
|
||||
43
crates/aether-data/src/redis/namespace.rs
Normal file
43
crates/aether-data/src/redis/namespace.rs
Normal file
@@ -0,0 +1,43 @@
|
||||
use aether_cache::CacheKeyNamespace;
|
||||
|
||||
use crate::redis::{RedisLockKey, RedisStreamName};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RedisKeyspace {
|
||||
namespace: CacheKeyNamespace,
|
||||
}
|
||||
|
||||
impl RedisKeyspace {
|
||||
pub fn new(prefix: Option<&str>) -> Self {
|
||||
let normalized = prefix.unwrap_or_default().trim().trim_matches(':');
|
||||
Self {
|
||||
namespace: CacheKeyNamespace::new(normalized),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn key(&self, raw_key: &str) -> String {
|
||||
self.namespace.key(raw_key)
|
||||
}
|
||||
|
||||
pub fn lock_key(&self, raw_key: &str) -> RedisLockKey {
|
||||
RedisLockKey(self.namespace.child("lock").key(raw_key))
|
||||
}
|
||||
|
||||
pub fn stream_name(&self, raw_name: &str) -> RedisStreamName {
|
||||
RedisStreamName(self.namespace.child("stream").key(raw_name))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::RedisKeyspace;
|
||||
|
||||
#[test]
|
||||
fn composes_prefixed_lock_and_stream_names() {
|
||||
let keyspace = RedisKeyspace::new(Some("aether"));
|
||||
|
||||
assert_eq!(keyspace.key("auth:user"), "aether:auth:user");
|
||||
assert_eq!(keyspace.lock_key("poller").0, "aether:lock:poller");
|
||||
assert_eq!(keyspace.stream_name("audit").0, "aether:stream:audit");
|
||||
}
|
||||
}
|
||||
658
crates/aether-data/src/redis/stream.rs
Normal file
658
crates/aether-data/src/redis/stream.rs
Normal file
@@ -0,0 +1,658 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::future::Future;
|
||||
use std::time::Duration;
|
||||
|
||||
use redis::from_redis_value;
|
||||
use redis::streams::StreamReadReply;
|
||||
use redis::Value as RedisValue;
|
||||
|
||||
use crate::redis::{RedisClient, RedisKeyspace};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct RedisStreamName(pub String);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct RedisConsumerGroup(pub String);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
pub struct RedisConsumerName(pub String);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RedisStreamEntry {
|
||||
pub id: String,
|
||||
pub fields: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RedisStreamReclaimResult {
|
||||
pub next_start_id: String,
|
||||
pub entries: Vec<RedisStreamEntry>,
|
||||
pub deleted_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct RedisStreamReclaimConfig {
|
||||
pub min_idle_ms: u64,
|
||||
pub count: usize,
|
||||
}
|
||||
|
||||
impl Default for RedisStreamReclaimConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
min_idle_ms: 60_000,
|
||||
count: 32,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisStreamReclaimConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if self.min_idle_ms == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis stream reclaim min_idle_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.count == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis stream reclaim count must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct RedisStreamRunnerConfig {
|
||||
pub command_timeout_ms: Option<u64>,
|
||||
pub read_block_ms: Option<u64>,
|
||||
pub read_count: usize,
|
||||
}
|
||||
|
||||
impl Default for RedisStreamRunnerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
command_timeout_ms: Some(1_000),
|
||||
read_block_ms: Some(1_000),
|
||||
read_count: 32,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisStreamRunnerConfig {
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if matches!(self.command_timeout_ms, Some(0)) {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis stream command_timeout_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if matches!(self.read_block_ms, Some(0)) {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis stream read_block_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.read_count == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"redis stream read_count must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RedisStreamRunner {
|
||||
client: RedisClient,
|
||||
keyspace: RedisKeyspace,
|
||||
config: RedisStreamRunnerConfig,
|
||||
}
|
||||
|
||||
impl RedisStreamRunner {
|
||||
pub fn new(
|
||||
client: RedisClient,
|
||||
keyspace: RedisKeyspace,
|
||||
config: RedisStreamRunnerConfig,
|
||||
) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
Ok(Self {
|
||||
client,
|
||||
keyspace,
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn client(&self) -> &RedisClient {
|
||||
&self.client
|
||||
}
|
||||
|
||||
pub fn keyspace(&self) -> &RedisKeyspace {
|
||||
&self.keyspace
|
||||
}
|
||||
|
||||
pub fn config(&self) -> RedisStreamRunnerConfig {
|
||||
self.config
|
||||
}
|
||||
|
||||
pub async fn ensure_consumer_group(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
group: &RedisConsumerGroup,
|
||||
start_id: &str,
|
||||
) -> Result<(), DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
validate_group(group)?;
|
||||
validate_stream_position(start_id)?;
|
||||
|
||||
self.run_with_timeout("redis stream ensure consumer group", async {
|
||||
let mut connection = self.client.get_multiplexed_async_connection().await?;
|
||||
let result = redis::cmd("XGROUP")
|
||||
.arg("CREATE")
|
||||
.arg(&stream.0)
|
||||
.arg(&group.0)
|
||||
.arg(start_id)
|
||||
.arg("MKSTREAM")
|
||||
.query_async::<String>(&mut connection)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(_) => Ok(()),
|
||||
Err(err) if err.code() == Some("BUSYGROUP") => Ok(()),
|
||||
Err(err) => Err(DataLayerError::Redis(err)),
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn append_fields(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
fields: &BTreeMap<String, String>,
|
||||
) -> Result<String, DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
if fields.is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis stream fields cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
self.run_with_timeout("redis stream append", async {
|
||||
let mut connection = self.client.get_multiplexed_async_connection().await?;
|
||||
let mut command = redis::cmd("XADD");
|
||||
command.arg(&stream.0).arg("*");
|
||||
for (key, value) in fields {
|
||||
command.arg(key).arg(value);
|
||||
}
|
||||
Ok(command.query_async::<String>(&mut connection).await?)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn append_json(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
field: &str,
|
||||
payload: &serde_json::Value,
|
||||
) -> Result<String, DataLayerError> {
|
||||
if field.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis stream json field cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let mut fields = BTreeMap::new();
|
||||
fields.insert(
|
||||
field.to_string(),
|
||||
serde_json::to_string(payload).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"failed to serialize redis stream payload: {err}"
|
||||
))
|
||||
})?,
|
||||
);
|
||||
self.append_fields(stream, &fields).await
|
||||
}
|
||||
|
||||
pub async fn read_group(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
group: &RedisConsumerGroup,
|
||||
consumer: &RedisConsumerName,
|
||||
) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
validate_group(group)?;
|
||||
validate_consumer(consumer)?;
|
||||
|
||||
self.run_with_timeout("redis stream read group", async {
|
||||
let mut connection = self.client.get_multiplexed_async_connection().await?;
|
||||
let mut command = redis::cmd("XREADGROUP");
|
||||
command
|
||||
.arg("GROUP")
|
||||
.arg(&group.0)
|
||||
.arg(&consumer.0)
|
||||
.arg("COUNT")
|
||||
.arg(self.config.read_count);
|
||||
if let Some(block_ms) = self.config.read_block_ms {
|
||||
command.arg("BLOCK").arg(block_ms);
|
||||
}
|
||||
command.arg("STREAMS").arg(&stream.0).arg(">");
|
||||
|
||||
let reply = command
|
||||
.query_async::<StreamReadReply>(&mut connection)
|
||||
.await?;
|
||||
|
||||
Ok(reply
|
||||
.keys
|
||||
.into_iter()
|
||||
.flat_map(|key| key.ids.into_iter())
|
||||
.map(|id| RedisStreamEntry {
|
||||
id: id.id,
|
||||
fields: id
|
||||
.map
|
||||
.into_iter()
|
||||
.filter_map(|(field, value)| {
|
||||
redis::from_redis_value::<String>(&value)
|
||||
.ok()
|
||||
.map(|text| (field, text))
|
||||
})
|
||||
.collect(),
|
||||
})
|
||||
.collect())
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn ack(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
group: &RedisConsumerGroup,
|
||||
ids: &[String],
|
||||
) -> Result<usize, DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
validate_group(group)?;
|
||||
if ids.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
self.run_with_timeout("redis stream ack", async {
|
||||
let mut connection = self.client.get_multiplexed_async_connection().await?;
|
||||
let mut command = redis::cmd("XACK");
|
||||
command.arg(&stream.0).arg(&group.0);
|
||||
for id in ids {
|
||||
command.arg(id);
|
||||
}
|
||||
Ok(command.query_async::<usize>(&mut connection).await?)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn claim_stale(
|
||||
&self,
|
||||
stream: &RedisStreamName,
|
||||
group: &RedisConsumerGroup,
|
||||
consumer: &RedisConsumerName,
|
||||
start_id: &str,
|
||||
config: RedisStreamReclaimConfig,
|
||||
) -> Result<RedisStreamReclaimResult, DataLayerError> {
|
||||
validate_stream_name(stream)?;
|
||||
validate_group(group)?;
|
||||
validate_consumer(consumer)?;
|
||||
validate_stream_position(start_id)?;
|
||||
config.validate()?;
|
||||
|
||||
self.run_with_timeout("redis stream reclaim", async {
|
||||
let mut connection = self.client.get_multiplexed_async_connection().await?;
|
||||
let reply = redis::cmd("XAUTOCLAIM")
|
||||
.arg(&stream.0)
|
||||
.arg(&group.0)
|
||||
.arg(&consumer.0)
|
||||
.arg(config.min_idle_ms)
|
||||
.arg(start_id)
|
||||
.arg("COUNT")
|
||||
.arg(config.count)
|
||||
.query_async::<RedisValue>(&mut connection)
|
||||
.await?;
|
||||
|
||||
parse_reclaim_result(reply)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_with_timeout<T, F>(
|
||||
&self,
|
||||
operation: &'static str,
|
||||
future: F,
|
||||
) -> Result<T, DataLayerError>
|
||||
where
|
||||
F: Future<Output = Result<T, DataLayerError>>,
|
||||
{
|
||||
if let Some(timeout_ms) = self.config.command_timeout_ms {
|
||||
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
|
||||
})?
|
||||
} else {
|
||||
future.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_stream_name(stream: &RedisStreamName) -> Result<(), DataLayerError> {
|
||||
if stream.0.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis stream name cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_group(group: &RedisConsumerGroup) -> Result<(), DataLayerError> {
|
||||
if group.0.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis consumer group cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_consumer(consumer: &RedisConsumerName) -> Result<(), DataLayerError> {
|
||||
if consumer.0.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis consumer name cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_stream_position(position: &str) -> Result<(), DataLayerError> {
|
||||
if position.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"redis stream position cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_reclaim_result(value: RedisValue) -> Result<RedisStreamReclaimResult, DataLayerError> {
|
||||
let RedisValue::Array(parts) = value else {
|
||||
return Err(DataLayerError::UnexpectedValue(
|
||||
"redis xautoclaim returned non-array payload".to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
if parts.len() < 2 || parts.len() > 3 {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"redis xautoclaim returned {} top-level fields, expected 2 or 3",
|
||||
parts.len()
|
||||
)));
|
||||
}
|
||||
|
||||
let next_start_id = parse_string_value(&parts[0], "redis xautoclaim next_start_id")?;
|
||||
let entries = parse_reclaim_entries(&parts[1])?;
|
||||
let deleted_ids = match parts.get(2) {
|
||||
Some(value) => parse_string_array(value, "redis xautoclaim deleted_ids")?,
|
||||
None => Vec::new(),
|
||||
};
|
||||
|
||||
Ok(RedisStreamReclaimResult {
|
||||
next_start_id,
|
||||
entries,
|
||||
deleted_ids,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_reclaim_entries(value: &RedisValue) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
|
||||
match value {
|
||||
RedisValue::Array(entries) => entries.iter().map(parse_reclaim_entry).collect(),
|
||||
RedisValue::Nil => Ok(Vec::new()),
|
||||
_ => Err(DataLayerError::UnexpectedValue(
|
||||
"redis xautoclaim entries payload was not an array".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_reclaim_entry(value: &RedisValue) -> Result<RedisStreamEntry, DataLayerError> {
|
||||
let RedisValue::Array(parts) = value else {
|
||||
return Err(DataLayerError::UnexpectedValue(
|
||||
"redis xautoclaim entry was not an array".to_string(),
|
||||
));
|
||||
};
|
||||
if parts.len() != 2 {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"redis xautoclaim entry had {} fields, expected 2",
|
||||
parts.len()
|
||||
)));
|
||||
}
|
||||
|
||||
let id = parse_string_value(&parts[0], "redis xautoclaim entry id")?;
|
||||
let fields = parse_string_map(&parts[1], "redis xautoclaim entry fields")?;
|
||||
Ok(RedisStreamEntry { id, fields })
|
||||
}
|
||||
|
||||
fn parse_string_map(
|
||||
value: &RedisValue,
|
||||
context: &str,
|
||||
) -> Result<BTreeMap<String, String>, DataLayerError> {
|
||||
match value {
|
||||
RedisValue::Array(values) => {
|
||||
if values.len() % 2 != 0 {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"{context} expected an even number of field elements, got {}",
|
||||
values.len()
|
||||
)));
|
||||
}
|
||||
let mut fields = BTreeMap::new();
|
||||
for pair in values.chunks(2) {
|
||||
let key = parse_string_value(&pair[0], context)?;
|
||||
let value = parse_string_value(&pair[1], context)?;
|
||||
fields.insert(key, value);
|
||||
}
|
||||
Ok(fields)
|
||||
}
|
||||
RedisValue::Map(entries) => entries
|
||||
.iter()
|
||||
.map(|(key, value)| {
|
||||
Ok((
|
||||
parse_string_value(key, context)?,
|
||||
parse_string_value(value, context)?,
|
||||
))
|
||||
})
|
||||
.collect(),
|
||||
RedisValue::Nil => Ok(BTreeMap::new()),
|
||||
_ => Err(DataLayerError::UnexpectedValue(format!(
|
||||
"{context} expected a redis array/map payload"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_string_array(value: &RedisValue, context: &str) -> Result<Vec<String>, DataLayerError> {
|
||||
match value {
|
||||
RedisValue::Array(values) => values
|
||||
.iter()
|
||||
.map(|value| parse_string_value(value, context))
|
||||
.collect(),
|
||||
RedisValue::Nil => Ok(Vec::new()),
|
||||
_ => Err(DataLayerError::UnexpectedValue(format!(
|
||||
"{context} expected a redis array payload"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_string_value(value: &RedisValue, context: &str) -> Result<String, DataLayerError> {
|
||||
from_redis_value::<String>(value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"{context} was not a string-compatible redis value: {err}"
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{
|
||||
parse_reclaim_result, RedisConsumerGroup, RedisConsumerName, RedisStreamName,
|
||||
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner,
|
||||
RedisStreamRunnerConfig,
|
||||
};
|
||||
use crate::redis::{RedisClientConfig, RedisClientFactory};
|
||||
use redis::Value as RedisValue;
|
||||
|
||||
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")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_stream_runner_config() {
|
||||
assert!(RedisStreamRunnerConfig {
|
||||
command_timeout_ms: Some(0),
|
||||
..RedisStreamRunnerConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
assert!(RedisStreamRunnerConfig {
|
||||
read_block_ms: Some(0),
|
||||
..RedisStreamRunnerConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
assert!(RedisStreamRunnerConfig {
|
||||
read_count: 0,
|
||||
..RedisStreamRunnerConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_reclaim_config() {
|
||||
assert!(RedisStreamReclaimConfig {
|
||||
min_idle_ms: 0,
|
||||
..RedisStreamReclaimConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
assert!(RedisStreamReclaimConfig {
|
||||
count: 0,
|
||||
..RedisStreamReclaimConfig::default()
|
||||
}
|
||||
.validate()
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runner_reuses_client_and_keyspace() {
|
||||
let runner = sample_runner();
|
||||
|
||||
assert_eq!(runner.config(), RedisStreamRunnerConfig::default());
|
||||
assert_eq!(
|
||||
runner.keyspace().stream_name("audit").0,
|
||||
"aether:stream:audit"
|
||||
);
|
||||
let _client_ref = runner.client();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_invalid_inputs_before_network() {
|
||||
let runner = sample_runner();
|
||||
let stream = RedisStreamName("aether:stream:audit".to_string());
|
||||
let group = RedisConsumerGroup("audit-workers".to_string());
|
||||
let consumer = RedisConsumerName("worker-1".to_string());
|
||||
|
||||
assert!(runner
|
||||
.ensure_consumer_group(&stream, &group, "")
|
||||
.await
|
||||
.is_err());
|
||||
assert!(runner
|
||||
.append_fields(&stream, &BTreeMap::new())
|
||||
.await
|
||||
.is_err());
|
||||
assert!(runner
|
||||
.append_json(&stream, "", &serde_json::json!({"ok": true}))
|
||||
.await
|
||||
.is_err());
|
||||
assert!(runner
|
||||
.read_group(&stream, &group, &RedisConsumerName(String::new()))
|
||||
.await
|
||||
.is_err());
|
||||
assert_eq!(
|
||||
runner.ack(&stream, &group, &[]).await.expect("empty ack"),
|
||||
0
|
||||
);
|
||||
assert!(runner
|
||||
.claim_stale(
|
||||
&stream,
|
||||
&group,
|
||||
&consumer,
|
||||
"",
|
||||
RedisStreamReclaimConfig::default()
|
||||
)
|
||||
.await
|
||||
.is_err());
|
||||
let _ = consumer;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_reclaim_result_with_deleted_ids() {
|
||||
let parsed = parse_reclaim_result(RedisValue::Array(vec![
|
||||
RedisValue::BulkString(b"0-0".to_vec()),
|
||||
RedisValue::Array(vec![RedisValue::Array(vec![
|
||||
RedisValue::BulkString(b"1710000000000-0".to_vec()),
|
||||
RedisValue::Array(vec![
|
||||
RedisValue::BulkString(b"payload".to_vec()),
|
||||
RedisValue::BulkString(br#"{"ok":true}"#.to_vec()),
|
||||
RedisValue::BulkString(b"kind".to_vec()),
|
||||
RedisValue::BulkString(b"shadow".to_vec()),
|
||||
]),
|
||||
])]),
|
||||
RedisValue::Array(vec![RedisValue::BulkString(b"1709999999999-0".to_vec())]),
|
||||
]))
|
||||
.expect("reclaim result should parse");
|
||||
|
||||
assert_eq!(
|
||||
parsed,
|
||||
RedisStreamReclaimResult {
|
||||
next_start_id: "0-0".to_string(),
|
||||
entries: vec![super::RedisStreamEntry {
|
||||
id: "1710000000000-0".to_string(),
|
||||
fields: BTreeMap::from([
|
||||
("kind".to_string(), "shadow".to_string()),
|
||||
("payload".to_string(), r#"{"ok":true}"#.to_string()),
|
||||
]),
|
||||
}],
|
||||
deleted_ids: vec!["1709999999999-0".to_string()],
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_reclaim_result_without_deleted_ids() {
|
||||
let parsed = parse_reclaim_result(RedisValue::Array(vec![
|
||||
RedisValue::BulkString(b"0-0".to_vec()),
|
||||
RedisValue::Array(vec![]),
|
||||
]))
|
||||
.expect("reclaim result should parse");
|
||||
|
||||
assert_eq!(
|
||||
parsed,
|
||||
RedisStreamReclaimResult {
|
||||
next_start_id: "0-0".to_string(),
|
||||
entries: Vec::new(),
|
||||
deleted_ids: Vec::new(),
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
133
crates/aether-data/src/repository/auth/memory.rs
Normal file
133
crates/aether-data/src/repository/auth/memory.rs
Normal file
@@ -0,0 +1,133 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct MemoryAuthApiKeyIndex {
|
||||
by_api_key_id: BTreeMap<String, StoredAuthApiKeySnapshot>,
|
||||
by_key_hash: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryAuthApiKeySnapshotRepository {
|
||||
index: RwLock<MemoryAuthApiKeyIndex>,
|
||||
}
|
||||
|
||||
impl InMemoryAuthApiKeySnapshotRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = (Option<String>, StoredAuthApiKeySnapshot)>,
|
||||
{
|
||||
let mut by_api_key_id = BTreeMap::new();
|
||||
let mut by_key_hash = BTreeMap::new();
|
||||
for (key_hash, snapshot) in items {
|
||||
if let Some(key_hash) = key_hash {
|
||||
by_key_hash.insert(key_hash, snapshot.api_key_id.clone());
|
||||
}
|
||||
by_api_key_id.insert(snapshot.api_key_id.clone(), snapshot);
|
||||
}
|
||||
Self {
|
||||
index: RwLock::new(MemoryAuthApiKeyIndex {
|
||||
by_api_key_id,
|
||||
by_key_hash,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthApiKeyReadRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
async fn find_api_key_snapshot(
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
|
||||
let index = self
|
||||
.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(match key {
|
||||
AuthApiKeyLookupKey::KeyHash(key_hash) => index
|
||||
.by_key_hash
|
||||
.get(key_hash)
|
||||
.and_then(|api_key_id| index.by_api_key_id.get(api_key_id))
|
||||
.cloned(),
|
||||
AuthApiKeyLookupKey::ApiKeyId(api_key_id) => {
|
||||
index.by_api_key_id.get(api_key_id).cloned()
|
||||
}
|
||||
AuthApiKeyLookupKey::UserApiKeyIds {
|
||||
user_id,
|
||||
api_key_id,
|
||||
} => index
|
||||
.by_api_key_id
|
||||
.get(api_key_id)
|
||||
.filter(|snapshot| snapshot.user_id == user_id)
|
||||
.cloned(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryAuthApiKeySnapshotRepository;
|
||||
use crate::repository::auth::{
|
||||
AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
|
||||
fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(200),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
)
|
||||
.expect("snapshot should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_auth_snapshot_by_all_supported_keys() {
|
||||
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_snapshot("key-1", "user-1"),
|
||||
)]);
|
||||
|
||||
assert!(repository
|
||||
.find_api_key_snapshot(AuthApiKeyLookupKey::KeyHash("hash-1"))
|
||||
.await
|
||||
.expect("find by hash should succeed")
|
||||
.is_some());
|
||||
assert!(repository
|
||||
.find_api_key_snapshot(AuthApiKeyLookupKey::ApiKeyId("key-1"))
|
||||
.await
|
||||
.expect("find by api key id should succeed")
|
||||
.is_some());
|
||||
assert!(repository
|
||||
.find_api_key_snapshot(AuthApiKeyLookupKey::UserApiKeyIds {
|
||||
user_id: "user-1",
|
||||
api_key_id: "key-1",
|
||||
})
|
||||
.await
|
||||
.expect("find by user/api key ids should succeed")
|
||||
.is_some());
|
||||
}
|
||||
}
|
||||
9
crates/aether-data/src/repository/auth/mod.rs
Normal file
9
crates/aether-data/src/repository/auth/mod.rs
Normal file
@@ -0,0 +1,9 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryAuthApiKeySnapshotRepository;
|
||||
pub use sql::SqlxAuthApiKeySnapshotReadRepository;
|
||||
pub use types::{
|
||||
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
202
crates/aether-data/src/repository/auth/sql.rs
Normal file
202
crates/aether-data/src/repository/auth/sql.rs
Normal file
@@ -0,0 +1,202 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_BY_KEY_HASH_SQL: &str = r#"
|
||||
SELECT
|
||||
users.id AS user_id,
|
||||
users.username,
|
||||
users.email,
|
||||
users.role::text AS user_role,
|
||||
users.auth_source::text AS user_auth_source,
|
||||
users.is_active AS user_is_active,
|
||||
users.is_deleted AS user_is_deleted,
|
||||
users.allowed_providers AS user_allowed_providers,
|
||||
users.allowed_api_formats AS user_allowed_api_formats,
|
||||
users.allowed_models AS user_allowed_models,
|
||||
api_keys.id AS api_key_id,
|
||||
api_keys.name AS api_key_name,
|
||||
api_keys.is_active AS api_key_is_active,
|
||||
api_keys.is_locked AS api_key_is_locked,
|
||||
api_keys.is_standalone AS api_key_is_standalone,
|
||||
api_keys.rate_limit AS api_key_rate_limit,
|
||||
api_keys.concurrent_limit AS api_key_concurrent_limit,
|
||||
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
|
||||
api_keys.allowed_providers AS api_key_allowed_providers,
|
||||
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
|
||||
api_keys.allowed_models AS api_key_allowed_models
|
||||
FROM api_keys
|
||||
JOIN users ON users.id = api_keys.user_id
|
||||
WHERE api_keys.key_hash = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_BY_API_KEY_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
users.id AS user_id,
|
||||
users.username,
|
||||
users.email,
|
||||
users.role::text AS user_role,
|
||||
users.auth_source::text AS user_auth_source,
|
||||
users.is_active AS user_is_active,
|
||||
users.is_deleted AS user_is_deleted,
|
||||
users.allowed_providers AS user_allowed_providers,
|
||||
users.allowed_api_formats AS user_allowed_api_formats,
|
||||
users.allowed_models AS user_allowed_models,
|
||||
api_keys.id AS api_key_id,
|
||||
api_keys.name AS api_key_name,
|
||||
api_keys.is_active AS api_key_is_active,
|
||||
api_keys.is_locked AS api_key_is_locked,
|
||||
api_keys.is_standalone AS api_key_is_standalone,
|
||||
api_keys.rate_limit AS api_key_rate_limit,
|
||||
api_keys.concurrent_limit AS api_key_concurrent_limit,
|
||||
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
|
||||
api_keys.allowed_providers AS api_key_allowed_providers,
|
||||
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
|
||||
api_keys.allowed_models AS api_key_allowed_models
|
||||
FROM api_keys
|
||||
JOIN users ON users.id = api_keys.user_id
|
||||
WHERE api_keys.id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_BY_USER_API_KEY_IDS_SQL: &str = r#"
|
||||
SELECT
|
||||
users.id AS user_id,
|
||||
users.username,
|
||||
users.email,
|
||||
users.role::text AS user_role,
|
||||
users.auth_source::text AS user_auth_source,
|
||||
users.is_active AS user_is_active,
|
||||
users.is_deleted AS user_is_deleted,
|
||||
users.allowed_providers AS user_allowed_providers,
|
||||
users.allowed_api_formats AS user_allowed_api_formats,
|
||||
users.allowed_models AS user_allowed_models,
|
||||
api_keys.id AS api_key_id,
|
||||
api_keys.name AS api_key_name,
|
||||
api_keys.is_active AS api_key_is_active,
|
||||
api_keys.is_locked AS api_key_is_locked,
|
||||
api_keys.is_standalone AS api_key_is_standalone,
|
||||
api_keys.rate_limit AS api_key_rate_limit,
|
||||
api_keys.concurrent_limit AS api_key_concurrent_limit,
|
||||
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
|
||||
api_keys.allowed_providers AS api_key_allowed_providers,
|
||||
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
|
||||
api_keys.allowed_models AS api_key_allowed_models
|
||||
FROM api_keys
|
||||
JOIN users ON users.id = api_keys.user_id
|
||||
WHERE api_keys.id = $1 AND users.id = $2
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxAuthApiKeySnapshotReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxAuthApiKeySnapshotReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn find_api_key_snapshot(
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
|
||||
let row = match key {
|
||||
AuthApiKeyLookupKey::KeyHash(key_hash) => {
|
||||
sqlx::query(FIND_BY_KEY_HASH_SQL)
|
||||
.bind(key_hash)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?
|
||||
}
|
||||
AuthApiKeyLookupKey::ApiKeyId(api_key_id) => {
|
||||
sqlx::query(FIND_BY_API_KEY_ID_SQL)
|
||||
.bind(api_key_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?
|
||||
}
|
||||
AuthApiKeyLookupKey::UserApiKeyIds {
|
||||
user_id,
|
||||
api_key_id,
|
||||
} => {
|
||||
sqlx::query(FIND_BY_USER_API_KEY_IDS_SQL)
|
||||
.bind(api_key_id)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?
|
||||
}
|
||||
};
|
||||
|
||||
row.as_ref().map(map_auth_api_key_snapshot_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AuthApiKeyReadRepository for SqlxAuthApiKeySnapshotReadRepository {
|
||||
async fn find_api_key_snapshot(
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
|
||||
Self::find_api_key_snapshot(self, key).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_auth_api_key_snapshot_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> Result<StoredAuthApiKeySnapshot, DataLayerError> {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
row.try_get("user_id")?,
|
||||
row.try_get("username")?,
|
||||
row.try_get("email")?,
|
||||
row.try_get("user_role")?,
|
||||
row.try_get("user_auth_source")?,
|
||||
row.try_get("user_is_active")?,
|
||||
row.try_get("user_is_deleted")?,
|
||||
row.try_get("user_allowed_providers")?,
|
||||
row.try_get("user_allowed_api_formats")?,
|
||||
row.try_get("user_allowed_models")?,
|
||||
row.try_get("api_key_id")?,
|
||||
row.try_get("api_key_name")?,
|
||||
row.try_get("api_key_is_active")?,
|
||||
row.try_get("api_key_is_locked")?,
|
||||
row.try_get("api_key_is_standalone")?,
|
||||
row.try_get("api_key_rate_limit")?,
|
||||
row.try_get("api_key_concurrent_limit")?,
|
||||
row.try_get("api_key_expires_at_unix_secs")?,
|
||||
row.try_get("api_key_allowed_providers")?,
|
||||
row.try_get("api_key_allowed_api_formats")?,
|
||||
row.try_get("api_key_allowed_models")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxAuthApiKeySnapshotReadRepository;
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_constructs_from_lazy_pool() {
|
||||
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("factory should build");
|
||||
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let repository = SqlxAuthApiKeySnapshotReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
}
|
||||
}
|
||||
225
crates/aether-data/src/repository/auth/types.rs
Normal file
225
crates/aether-data/src/repository/auth/types.rs
Normal file
@@ -0,0 +1,225 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredAuthApiKeySnapshot {
|
||||
pub user_id: String,
|
||||
pub username: String,
|
||||
pub email: Option<String>,
|
||||
pub user_role: String,
|
||||
pub user_auth_source: String,
|
||||
pub user_is_active: bool,
|
||||
pub user_is_deleted: bool,
|
||||
pub user_allowed_providers: Option<Vec<String>>,
|
||||
pub user_allowed_api_formats: Option<Vec<String>>,
|
||||
pub user_allowed_models: Option<Vec<String>>,
|
||||
pub api_key_id: String,
|
||||
pub api_key_name: Option<String>,
|
||||
pub api_key_is_active: bool,
|
||||
pub api_key_is_locked: bool,
|
||||
pub api_key_is_standalone: bool,
|
||||
pub api_key_rate_limit: Option<i32>,
|
||||
pub api_key_concurrent_limit: Option<i32>,
|
||||
pub api_key_expires_at_unix_secs: Option<u64>,
|
||||
pub api_key_allowed_providers: Option<Vec<String>>,
|
||||
pub api_key_allowed_api_formats: Option<Vec<String>>,
|
||||
pub api_key_allowed_models: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
impl StoredAuthApiKeySnapshot {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
user_id: String,
|
||||
username: String,
|
||||
email: Option<String>,
|
||||
user_role: String,
|
||||
user_auth_source: String,
|
||||
user_is_active: bool,
|
||||
user_is_deleted: bool,
|
||||
user_allowed_providers: Option<serde_json::Value>,
|
||||
user_allowed_api_formats: Option<serde_json::Value>,
|
||||
user_allowed_models: Option<serde_json::Value>,
|
||||
api_key_id: String,
|
||||
api_key_name: Option<String>,
|
||||
api_key_is_active: bool,
|
||||
api_key_is_locked: bool,
|
||||
api_key_is_standalone: bool,
|
||||
api_key_rate_limit: Option<i32>,
|
||||
api_key_concurrent_limit: Option<i32>,
|
||||
api_key_expires_at_unix_secs: Option<i64>,
|
||||
api_key_allowed_providers: Option<serde_json::Value>,
|
||||
api_key_allowed_api_formats: Option<serde_json::Value>,
|
||||
api_key_allowed_models: Option<serde_json::Value>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
Ok(Self {
|
||||
user_id,
|
||||
username,
|
||||
email,
|
||||
user_role,
|
||||
user_auth_source,
|
||||
user_is_active,
|
||||
user_is_deleted,
|
||||
user_allowed_providers: parse_string_list(
|
||||
user_allowed_providers,
|
||||
"users.allowed_providers",
|
||||
)?,
|
||||
user_allowed_api_formats: parse_string_list(
|
||||
user_allowed_api_formats,
|
||||
"users.allowed_api_formats",
|
||||
)?,
|
||||
user_allowed_models: parse_string_list(user_allowed_models, "users.allowed_models")?,
|
||||
api_key_id,
|
||||
api_key_name,
|
||||
api_key_is_active,
|
||||
api_key_is_locked,
|
||||
api_key_is_standalone,
|
||||
api_key_rate_limit,
|
||||
api_key_concurrent_limit,
|
||||
api_key_expires_at_unix_secs: api_key_expires_at_unix_secs
|
||||
.map(|value| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid api_keys.expires_at_unix_secs: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?,
|
||||
api_key_allowed_providers: parse_string_list(
|
||||
api_key_allowed_providers,
|
||||
"api_keys.allowed_providers",
|
||||
)?,
|
||||
api_key_allowed_api_formats: parse_string_list(
|
||||
api_key_allowed_api_formats,
|
||||
"api_keys.allowed_api_formats",
|
||||
)?,
|
||||
api_key_allowed_models: parse_string_list(
|
||||
api_key_allowed_models,
|
||||
"api_keys.allowed_models",
|
||||
)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn is_currently_usable(&self, now_unix_secs: u64) -> bool {
|
||||
if !self.user_is_active || self.user_is_deleted {
|
||||
return false;
|
||||
}
|
||||
if !self.api_key_is_active {
|
||||
return false;
|
||||
}
|
||||
if self.api_key_is_locked && !self.api_key_is_standalone {
|
||||
return false;
|
||||
}
|
||||
if let Some(expires_at_unix_secs) = self.api_key_expires_at_unix_secs {
|
||||
if expires_at_unix_secs < now_unix_secs {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AuthApiKeyLookupKey<'a> {
|
||||
KeyHash(&'a str),
|
||||
ApiKeyId(&'a str),
|
||||
UserApiKeyIds {
|
||||
user_id: &'a str,
|
||||
api_key_id: &'a str,
|
||||
},
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait AuthApiKeyReadRepository: Send + Sync {
|
||||
async fn find_api_key_snapshot(
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait AuthRepository: AuthApiKeyReadRepository + Send + Sync {}
|
||||
|
||||
impl<T> AuthRepository for T where T: AuthApiKeyReadRepository + Send + Sync {}
|
||||
|
||||
fn parse_string_list(
|
||||
value: Option<serde_json::Value>,
|
||||
field_name: &str,
|
||||
) -> Result<Option<Vec<String>>, crate::DataLayerError> {
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
let array = value.as_array().ok_or_else(|| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("{field_name} is not a JSON array"))
|
||||
})?;
|
||||
let mut items = Vec::with_capacity(array.len());
|
||||
for item in array {
|
||||
let Some(item) = item.as_str() else {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"{field_name} contains a non-string item"
|
||||
)));
|
||||
};
|
||||
items.push(item.to_string());
|
||||
}
|
||||
Ok(Some(items))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::StoredAuthApiKeySnapshot;
|
||||
|
||||
#[test]
|
||||
fn rejects_non_array_allowed_providers() {
|
||||
assert!(StoredAuthApiKeySnapshot::new(
|
||||
"user-1".to_string(),
|
||||
"alice".to_string(),
|
||||
None,
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!({"bad": true})),
|
||||
None,
|
||||
None,
|
||||
"key-1".to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expired_non_standalone_key_is_not_usable() {
|
||||
let snapshot = StoredAuthApiKeySnapshot::new(
|
||||
"user-1".to_string(),
|
||||
"alice".to_string(),
|
||||
None,
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
"key-1".to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(100),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("snapshot should build");
|
||||
|
||||
assert!(!snapshot.is_currently_usable(101));
|
||||
}
|
||||
}
|
||||
148
crates/aether-data/src/repository/candidates/memory.rs
Normal file
148
crates/aether-data/src/repository/candidates/memory.rs
Normal file
@@ -0,0 +1,148 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{RequestCandidateReadRepository, StoredRequestCandidate};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryRequestCandidateRepository {
|
||||
by_id: RwLock<BTreeMap<String, StoredRequestCandidate>>,
|
||||
}
|
||||
|
||||
impl InMemoryRequestCandidateRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredRequestCandidate>,
|
||||
{
|
||||
let mut by_id = BTreeMap::new();
|
||||
for item in items {
|
||||
by_id.insert(item.id.clone(), item);
|
||||
}
|
||||
Self {
|
||||
by_id: RwLock::new(by_id),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RequestCandidateReadRepository for InMemoryRequestCandidateRepository {
|
||||
async fn list_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
let mut rows = self
|
||||
.by_id
|
||||
.read()
|
||||
.expect("request candidate repository lock")
|
||||
.values()
|
||||
.filter(|row| row.request_id == request_id)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
rows.sort_by(|left, right| {
|
||||
left.candidate_index
|
||||
.cmp(&right.candidate_index)
|
||||
.then(left.retry_index.cmp(&right.retry_index))
|
||||
.then(left.created_at_unix_secs.cmp(&right.created_at_unix_secs))
|
||||
});
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
async fn list_recent(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut rows = self
|
||||
.by_id
|
||||
.read()
|
||||
.expect("request candidate repository lock")
|
||||
.values()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
rows.sort_by(|left, right| right.created_at_unix_secs.cmp(&left.created_at_unix_secs));
|
||||
rows.truncate(limit);
|
||||
Ok(rows)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryRequestCandidateRepository;
|
||||
use crate::repository::candidates::{
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
|
||||
fn sample_candidate(
|
||||
id: &str,
|
||||
request_id: &str,
|
||||
created_at_unix_secs: i64,
|
||||
) -> StoredRequestCandidate {
|
||||
StoredRequestCandidate::new(
|
||||
id.to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
0,
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
RequestCandidateStatus::Success,
|
||||
None,
|
||||
false,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(10),
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
created_at_unix_secs,
|
||||
Some(created_at_unix_secs),
|
||||
Some(created_at_unix_secs + 1),
|
||||
)
|
||||
.expect("candidate should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_request_candidates_by_request_id_in_candidate_order() {
|
||||
let repository = InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_candidate("cand-2", "req-1", 200),
|
||||
sample_candidate("cand-1", "req-1", 100),
|
||||
sample_candidate("cand-3", "req-2", 300),
|
||||
]);
|
||||
|
||||
let rows = repository
|
||||
.list_by_request_id("req-1")
|
||||
.await
|
||||
.expect("list should succeed");
|
||||
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(rows[0].request_id, "req-1");
|
||||
assert_eq!(rows[1].request_id, "req-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lists_recent_request_candidates_in_descending_created_order() {
|
||||
let repository = InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_candidate("cand-1", "req-1", 100),
|
||||
sample_candidate("cand-2", "req-2", 200),
|
||||
]);
|
||||
|
||||
let rows = repository
|
||||
.list_recent(10)
|
||||
.await
|
||||
.expect("list recent should succeed");
|
||||
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(rows[0].id, "cand-2");
|
||||
assert_eq!(rows[1].id, "cand-1");
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/candidates/mod.rs
Normal file
10
crates/aether-data/src/repository/candidates/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryRequestCandidateRepository;
|
||||
pub use sql::SqlxRequestCandidateReadRepository;
|
||||
pub use types::{
|
||||
RequestCandidateReadRepository, RequestCandidateRepository, RequestCandidateStatus,
|
||||
StoredRequestCandidate,
|
||||
};
|
||||
189
crates/aether-data/src/repository/candidates/sql.rs
Normal file
189
crates/aether-data/src/repository/candidates/sql.rs
Normal file
@@ -0,0 +1,189 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const LIST_BY_REQUEST_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
status,
|
||||
skip_reason,
|
||||
is_cached,
|
||||
status_code,
|
||||
error_type,
|
||||
error_message,
|
||||
latency_ms,
|
||||
concurrent_requests,
|
||||
extra_data,
|
||||
required_capabilities,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM started_at) AS BIGINT) AS started_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finished_at) AS BIGINT) AS finished_at_unix_secs
|
||||
FROM request_candidates
|
||||
WHERE request_id = $1
|
||||
ORDER BY candidate_index ASC, retry_index ASC, created_at ASC
|
||||
"#;
|
||||
|
||||
const LIST_RECENT_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
status,
|
||||
skip_reason,
|
||||
is_cached,
|
||||
status_code,
|
||||
error_type,
|
||||
error_message,
|
||||
latency_ms,
|
||||
concurrent_requests,
|
||||
extra_data,
|
||||
required_capabilities,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM started_at) AS BIGINT) AS started_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finished_at) AS BIGINT) AS finished_at_unix_secs
|
||||
FROM request_candidates
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $1
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxRequestCandidateReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxRequestCandidateReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn list_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
let rows = sqlx::query(LIST_BY_REQUEST_ID_SQL)
|
||||
.bind(request_id)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_request_candidate_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_recent(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = sqlx::query(LIST_RECENT_SQL)
|
||||
.bind(i64::try_from(limit).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"invalid recent request candidate limit: {limit}"
|
||||
))
|
||||
})?)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_request_candidate_row).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RequestCandidateReadRepository for SqlxRequestCandidateReadRepository {
|
||||
async fn list_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
Self::list_by_request_id(self, request_id).await
|
||||
}
|
||||
|
||||
async fn list_recent(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
Self::list_recent(self, limit).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_request_candidate_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> Result<StoredRequestCandidate, DataLayerError> {
|
||||
let status =
|
||||
RequestCandidateStatus::from_database(row.try_get::<String, _>("status")?.as_str())?;
|
||||
StoredRequestCandidate::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("request_id")?,
|
||||
row.try_get("user_id")?,
|
||||
row.try_get("api_key_id")?,
|
||||
row.try_get("username")?,
|
||||
row.try_get("api_key_name")?,
|
||||
row.try_get("candidate_index")?,
|
||||
row.try_get("retry_index")?,
|
||||
row.try_get("provider_id")?,
|
||||
row.try_get("endpoint_id")?,
|
||||
row.try_get("key_id")?,
|
||||
status,
|
||||
row.try_get("skip_reason")?,
|
||||
row.try_get("is_cached")?,
|
||||
row.try_get("status_code")?,
|
||||
row.try_get("error_type")?,
|
||||
row.try_get("error_message")?,
|
||||
row.try_get("latency_ms")?,
|
||||
row.try_get("concurrent_requests")?,
|
||||
row.try_get("extra_data")?,
|
||||
row.try_get("required_capabilities")?,
|
||||
row.try_get("created_at_unix_secs")?,
|
||||
row.try_get("started_at_unix_secs")?,
|
||||
row.try_get("finished_at_unix_secs")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxRequestCandidateReadRepository;
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_constructs_from_lazy_pool() {
|
||||
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("factory should build");
|
||||
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let repository = SqlxRequestCandidateReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
}
|
||||
}
|
||||
289
crates/aether-data/src/repository/candidates/types.rs
Normal file
289
crates/aether-data/src/repository/candidates/types.rs
Normal file
@@ -0,0 +1,289 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RequestCandidateStatus {
|
||||
Available,
|
||||
Unused,
|
||||
Pending,
|
||||
Streaming,
|
||||
Success,
|
||||
Failed,
|
||||
Cancelled,
|
||||
Skipped,
|
||||
}
|
||||
|
||||
impl RequestCandidateStatus {
|
||||
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"available" => Ok(Self::Available),
|
||||
"unused" => Ok(Self::Unused),
|
||||
"pending" => Ok(Self::Pending),
|
||||
"streaming" => Ok(Self::Streaming),
|
||||
"success" => Ok(Self::Success),
|
||||
"failed" => Ok(Self::Failed),
|
||||
"cancelled" => Ok(Self::Cancelled),
|
||||
"skipped" => Ok(Self::Skipped),
|
||||
other => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"unsupported request_candidates.status: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_attempted(self, started_at_unix_secs: Option<u64>) -> bool {
|
||||
match self {
|
||||
Self::Available | Self::Unused | Self::Skipped => false,
|
||||
Self::Pending => started_at_unix_secs.is_some(),
|
||||
Self::Streaming | Self::Success | Self::Failed | Self::Cancelled => true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredRequestCandidate {
|
||||
pub id: String,
|
||||
pub request_id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub api_key_id: Option<String>,
|
||||
pub username: Option<String>,
|
||||
pub api_key_name: Option<String>,
|
||||
pub candidate_index: u32,
|
||||
pub retry_index: u32,
|
||||
pub provider_id: Option<String>,
|
||||
pub endpoint_id: Option<String>,
|
||||
pub key_id: Option<String>,
|
||||
pub status: RequestCandidateStatus,
|
||||
pub skip_reason: Option<String>,
|
||||
pub is_cached: bool,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_type: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub latency_ms: Option<u64>,
|
||||
pub concurrent_requests: Option<u32>,
|
||||
pub extra_data: Option<serde_json::Value>,
|
||||
pub required_capabilities: Option<serde_json::Value>,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub started_at_unix_secs: Option<u64>,
|
||||
pub finished_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl StoredRequestCandidate {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
request_id: String,
|
||||
user_id: Option<String>,
|
||||
api_key_id: Option<String>,
|
||||
username: Option<String>,
|
||||
api_key_name: Option<String>,
|
||||
candidate_index: i32,
|
||||
retry_index: i32,
|
||||
provider_id: Option<String>,
|
||||
endpoint_id: Option<String>,
|
||||
key_id: Option<String>,
|
||||
status: RequestCandidateStatus,
|
||||
skip_reason: Option<String>,
|
||||
is_cached: bool,
|
||||
status_code: Option<i32>,
|
||||
error_type: Option<String>,
|
||||
error_message: Option<String>,
|
||||
latency_ms: Option<i32>,
|
||||
concurrent_requests: Option<i32>,
|
||||
extra_data: Option<serde_json::Value>,
|
||||
required_capabilities: Option<serde_json::Value>,
|
||||
created_at_unix_secs: i64,
|
||||
started_at_unix_secs: Option<i64>,
|
||||
finished_at_unix_secs: Option<i64>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
let candidate_index = u32::try_from(candidate_index).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.candidate_index: {candidate_index}"
|
||||
))
|
||||
})?;
|
||||
let retry_index = u32::try_from(retry_index).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.retry_index: {retry_index}"
|
||||
))
|
||||
})?;
|
||||
let status_code = status_code
|
||||
.map(|value| {
|
||||
u16::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.status_code: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let latency_ms = latency_ms
|
||||
.map(|value| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.latency_ms: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let concurrent_requests = concurrent_requests
|
||||
.map(|value| {
|
||||
u32::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.concurrent_requests: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.created_at_unix_secs: {created_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
let started_at_unix_secs = started_at_unix_secs
|
||||
.map(|value| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.started_at_unix_secs: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let finished_at_unix_secs = finished_at_unix_secs
|
||||
.map(|value| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid request_candidates.finished_at_unix_secs: {value}"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
status,
|
||||
skip_reason,
|
||||
is_cached,
|
||||
status_code,
|
||||
error_type,
|
||||
error_message,
|
||||
latency_ms,
|
||||
concurrent_requests,
|
||||
extra_data,
|
||||
required_capabilities,
|
||||
created_at_unix_secs,
|
||||
started_at_unix_secs,
|
||||
finished_at_unix_secs,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait RequestCandidateReadRepository: Send + Sync {
|
||||
async fn list_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
|
||||
|
||||
async fn list_recent(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait RequestCandidateRepository: RequestCandidateReadRepository + Send + Sync {}
|
||||
|
||||
impl<T> RequestCandidateRepository for T where T: RequestCandidateReadRepository + Send + Sync {}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{RequestCandidateStatus, StoredRequestCandidate};
|
||||
|
||||
#[test]
|
||||
fn parses_status_from_database_text() {
|
||||
assert_eq!(
|
||||
RequestCandidateStatus::from_database("streaming").expect("status should parse"),
|
||||
RequestCandidateStatus::Streaming
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_database_status() {
|
||||
assert!(RequestCandidateStatus::from_database("mystery").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_candidate_index() {
|
||||
assert!(StoredRequestCandidate::new(
|
||||
"cand-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
-1,
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
RequestCandidateStatus::Pending,
|
||||
None,
|
||||
false,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(10),
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
100,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_created_at() {
|
||||
assert!(StoredRequestCandidate::new(
|
||||
"cand-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
0,
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
RequestCandidateStatus::Pending,
|
||||
None,
|
||||
false,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(10),
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
-1,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_without_started_at_is_not_attempted() {
|
||||
assert!(!RequestCandidateStatus::Pending.is_attempted(None));
|
||||
assert!(RequestCandidateStatus::Pending.is_attempted(Some(1)));
|
||||
}
|
||||
}
|
||||
6
crates/aether-data/src/repository/mod.rs
Normal file
6
crates/aether-data/src/repository/mod.rs
Normal file
@@ -0,0 +1,6 @@
|
||||
pub mod auth;
|
||||
pub mod candidates;
|
||||
pub mod provider_catalog;
|
||||
pub mod shadow_results;
|
||||
pub mod usage;
|
||||
pub mod video_tasks;
|
||||
157
crates/aether-data/src/repository/provider_catalog/memory.rs
Normal file
157
crates/aether-data/src/repository/provider_catalog/memory.rs
Normal file
@@ -0,0 +1,157 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct MemoryProviderCatalogIndex {
|
||||
providers: BTreeMap<String, StoredProviderCatalogProvider>,
|
||||
endpoints: BTreeMap<String, StoredProviderCatalogEndpoint>,
|
||||
keys: BTreeMap<String, StoredProviderCatalogKey>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryProviderCatalogReadRepository {
|
||||
index: RwLock<MemoryProviderCatalogIndex>,
|
||||
}
|
||||
|
||||
impl InMemoryProviderCatalogReadRepository {
|
||||
pub fn seed(
|
||||
providers: Vec<StoredProviderCatalogProvider>,
|
||||
endpoints: Vec<StoredProviderCatalogEndpoint>,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
) -> Self {
|
||||
Self {
|
||||
index: RwLock::new(MemoryProviderCatalogIndex {
|
||||
providers: providers
|
||||
.into_iter()
|
||||
.map(|provider| (provider.id.clone(), provider))
|
||||
.collect(),
|
||||
endpoints: endpoints
|
||||
.into_iter()
|
||||
.map(|endpoint| (endpoint.id.clone(), endpoint))
|
||||
.collect(),
|
||||
keys: keys.into_iter().map(|key| (key.id.clone(), key)).collect(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
|
||||
async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
let index = self.index.read().expect("provider catalog repository lock");
|
||||
Ok(provider_ids
|
||||
.iter()
|
||||
.filter_map(|id| index.providers.get(id).cloned())
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
let index = self.index.read().expect("provider catalog repository lock");
|
||||
Ok(endpoint_ids
|
||||
.iter()
|
||||
.filter_map(|id| index.endpoints.get(id).cloned())
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
let index = self.index.read().expect("provider catalog repository lock");
|
||||
Ok(key_ids
|
||||
.iter()
|
||||
.filter_map(|id| index.keys.get(id).cloned())
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryProviderCatalogReadRepository;
|
||||
use crate::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
|
||||
fn sample_provider(id: &str) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
id.to_string(),
|
||||
format!("provider-{id}"),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
}
|
||||
|
||||
fn sample_endpoint(id: &str, provider_id: &str) -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
id.to_string(),
|
||||
provider_id.to_string(),
|
||||
"openai:chat".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
}
|
||||
|
||||
fn sample_key(id: &str, provider_id: &str) -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
id.to_string(),
|
||||
provider_id.to_string(),
|
||||
"default".to_string(),
|
||||
"api_key".to_string(),
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_provider_catalog_items_by_id() {
|
||||
let repository = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1")],
|
||||
vec![sample_endpoint("endpoint-1", "provider-1")],
|
||||
vec![sample_key("key-1", "provider-1")],
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
repository
|
||||
.list_providers_by_ids(&["provider-1".to_string()])
|
||||
.await
|
||||
.expect("providers should read")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
repository
|
||||
.list_endpoints_by_ids(&["endpoint-1".to_string()])
|
||||
.await
|
||||
.expect("endpoints should read")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
repository
|
||||
.list_keys_by_ids(&["key-1".to_string()])
|
||||
.await
|
||||
.expect("keys should read")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/provider_catalog/mod.rs
Normal file
10
crates/aether-data/src/repository/provider_catalog/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryProviderCatalogReadRepository;
|
||||
pub use sql::SqlxProviderCatalogReadRepository;
|
||||
pub use types::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
209
crates/aether-data/src/repository/provider_catalog/sql.rs
Normal file
209
crates/aether-data/src/repository/provider_catalog/sql.rs
Normal file
@@ -0,0 +1,209 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
||||
|
||||
use super::types::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const LIST_PROVIDERS_BY_IDS_PREFIX: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
name,
|
||||
website,
|
||||
provider_type
|
||||
FROM providers
|
||||
WHERE id IN (
|
||||
"#;
|
||||
|
||||
const LIST_ENDPOINTS_BY_IDS_PREFIX: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
provider_id,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
is_active
|
||||
FROM provider_endpoints
|
||||
WHERE id IN (
|
||||
"#;
|
||||
|
||||
const LIST_KEYS_BY_IDS_PREFIX: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
provider_id,
|
||||
name,
|
||||
auth_type,
|
||||
capabilities,
|
||||
is_active
|
||||
FROM provider_api_keys
|
||||
WHERE id IN (
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxProviderCatalogReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxProviderCatalogReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
if provider_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_list_query(
|
||||
LIST_PROVIDERS_BY_IDS_PREFIX,
|
||||
provider_ids,
|
||||
" ORDER BY name ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_provider_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
if endpoint_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_list_query(
|
||||
LIST_ENDPOINTS_BY_IDS_PREFIX,
|
||||
endpoint_ids,
|
||||
" ORDER BY api_format ASC, id ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_endpoint_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
if key_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_list_query(
|
||||
LIST_KEYS_BY_IDS_PREFIX,
|
||||
key_ids,
|
||||
" ORDER BY name ASC, id ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
rows.iter().map(map_key_row).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderCatalogReadRepository for SqlxProviderCatalogReadRepository {
|
||||
async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
Self::list_providers_by_ids(self, provider_ids).await
|
||||
}
|
||||
|
||||
async fn list_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
Self::list_endpoints_by_ids(self, endpoint_ids).await
|
||||
}
|
||||
|
||||
async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
Self::list_keys_by_ids(self, key_ids).await
|
||||
}
|
||||
}
|
||||
|
||||
fn build_list_query<'a>(
|
||||
prefix: &'static str,
|
||||
ids: &'a [String],
|
||||
suffix: &'static str,
|
||||
) -> QueryBuilder<'a, Postgres> {
|
||||
let mut builder = QueryBuilder::<Postgres>::new(prefix);
|
||||
let mut separated = builder.separated(", ");
|
||||
for id in ids {
|
||||
separated.push_bind(id);
|
||||
}
|
||||
separated.push_unseparated(")");
|
||||
builder.push(suffix);
|
||||
builder
|
||||
}
|
||||
|
||||
fn map_provider_row(row: &PgRow) -> Result<StoredProviderCatalogProvider, DataLayerError> {
|
||||
StoredProviderCatalogProvider::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("name")?,
|
||||
row.try_get("website")?,
|
||||
row.try_get("provider_type")?,
|
||||
)
|
||||
}
|
||||
|
||||
fn map_endpoint_row(row: &PgRow) -> Result<StoredProviderCatalogEndpoint, DataLayerError> {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("provider_id")?,
|
||||
row.try_get("api_format")?,
|
||||
row.try_get("api_family")?,
|
||||
row.try_get("endpoint_kind")?,
|
||||
row.try_get("is_active")?,
|
||||
)
|
||||
}
|
||||
|
||||
fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError> {
|
||||
StoredProviderCatalogKey::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("provider_id")?,
|
||||
row.try_get("name")?,
|
||||
row.try_get("auth_type")?,
|
||||
row.try_get("capabilities")?,
|
||||
row.try_get("is_active")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxProviderCatalogReadRepository;
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_constructs_from_lazy_pool() {
|
||||
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("factory should build");
|
||||
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let repository = SqlxProviderCatalogReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
}
|
||||
}
|
||||
175
crates/aether-data/src/repository/provider_catalog/types.rs
Normal file
175
crates/aether-data/src/repository/provider_catalog/types.rs
Normal file
@@ -0,0 +1,175 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderCatalogProvider {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub website: Option<String>,
|
||||
pub provider_type: String,
|
||||
}
|
||||
|
||||
impl StoredProviderCatalogProvider {
|
||||
pub fn new(
|
||||
id: String,
|
||||
name: String,
|
||||
website: Option<String>,
|
||||
provider_type: String,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"providers.name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if provider_type.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"providers.provider_type is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
name,
|
||||
website,
|
||||
provider_type,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderCatalogEndpoint {
|
||||
pub id: String,
|
||||
pub provider_id: String,
|
||||
pub api_format: String,
|
||||
pub api_family: Option<String>,
|
||||
pub endpoint_kind: Option<String>,
|
||||
pub is_active: bool,
|
||||
}
|
||||
|
||||
impl StoredProviderCatalogEndpoint {
|
||||
pub fn new(
|
||||
id: String,
|
||||
provider_id: String,
|
||||
api_format: String,
|
||||
api_family: Option<String>,
|
||||
endpoint_kind: Option<String>,
|
||||
is_active: bool,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if api_format.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider_endpoints.api_format is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
provider_id,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
is_active,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderCatalogKey {
|
||||
pub id: String,
|
||||
pub provider_id: String,
|
||||
pub name: String,
|
||||
pub auth_type: String,
|
||||
pub capabilities: Option<serde_json::Value>,
|
||||
pub is_active: bool,
|
||||
}
|
||||
|
||||
impl StoredProviderCatalogKey {
|
||||
pub fn new(
|
||||
id: String,
|
||||
provider_id: String,
|
||||
name: String,
|
||||
auth_type: String,
|
||||
capabilities: Option<serde_json::Value>,
|
||||
is_active: bool,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider_api_keys.name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if auth_type.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider_api_keys.auth_type is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
provider_id,
|
||||
name,
|
||||
auth_type,
|
||||
capabilities,
|
||||
is_active,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ProviderCatalogReadRepository: Send + Sync {
|
||||
async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, crate::DataLayerError>;
|
||||
|
||||
async fn list_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, crate::DataLayerError>;
|
||||
|
||||
async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_provider_name() {
|
||||
assert!(StoredProviderCatalogProvider::new(
|
||||
"provider-1".to_string(),
|
||||
"".to_string(),
|
||||
None,
|
||||
"custom".to_string(),
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_endpoint_api_format() {
|
||||
assert!(StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"".to_string(),
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_key_auth_type() {
|
||||
assert!(StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"default".to_string(),
|
||||
"".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
170
crates/aether-data/src/repository/shadow_results/memory.rs
Normal file
170
crates/aether-data/src/repository/shadow_results/memory.rs
Normal file
@@ -0,0 +1,170 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
ShadowResultLookupKey, ShadowResultReadRepository, ShadowResultWriteRepository,
|
||||
StoredShadowResult, UpsertShadowResult,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryShadowResultRepository {
|
||||
results: RwLock<BTreeMap<(String, String), StoredShadowResult>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ShadowResultReadRepository for InMemoryShadowResultRepository {
|
||||
async fn find(
|
||||
&self,
|
||||
key: ShadowResultLookupKey<'_>,
|
||||
) -> Result<Option<StoredShadowResult>, DataLayerError> {
|
||||
let results = self.results.read().expect("shadow result repository lock");
|
||||
Ok(match key {
|
||||
ShadowResultLookupKey::TraceFingerprint {
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
} => results
|
||||
.get(&(trace_id.to_string(), request_fingerprint.to_string()))
|
||||
.cloned(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_recent(&self, limit: usize) -> Result<Vec<StoredShadowResult>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut results = self
|
||||
.results
|
||||
.read()
|
||||
.expect("shadow result repository lock")
|
||||
.values()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
results.sort_by(|left, right| right.updated_at_unix_secs.cmp(&left.updated_at_unix_secs));
|
||||
results.truncate(limit);
|
||||
Ok(results)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ShadowResultWriteRepository for InMemoryShadowResultRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
result: UpsertShadowResult,
|
||||
) -> Result<StoredShadowResult, DataLayerError> {
|
||||
let stored = result.into_stored();
|
||||
let mut results = self.results.write().expect("shadow result repository lock");
|
||||
results.insert(
|
||||
(stored.trace_id.clone(), stored.request_fingerprint.clone()),
|
||||
stored.clone(),
|
||||
);
|
||||
Ok(stored)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryShadowResultRepository;
|
||||
use crate::repository::shadow_results::{
|
||||
ShadowResultLookupKey, ShadowResultMatchStatus, ShadowResultReadRepository,
|
||||
ShadowResultWriteRepository, UpsertShadowResult,
|
||||
};
|
||||
|
||||
fn sample_result(
|
||||
trace_id: &str,
|
||||
request_fingerprint: &str,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> UpsertShadowResult {
|
||||
UpsertShadowResult {
|
||||
trace_id: trace_id.to_string(),
|
||||
request_fingerprint: request_fingerprint.to_string(),
|
||||
request_id: Some(format!("req-{trace_id}")),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
candidate_id: Some("cand-1".to_string()),
|
||||
rust_result_digest: Some("rust-digest".to_string()),
|
||||
python_result_digest: Some("python-digest".to_string()),
|
||||
match_status: ShadowResultMatchStatus::Match,
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
created_at_unix_secs: updated_at_unix_secs.saturating_sub(10),
|
||||
updated_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_result_by_trace_and_fingerprint() {
|
||||
let repo = InMemoryShadowResultRepository::default();
|
||||
repo.upsert(sample_result("trace-1", "fp-1", 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
assert!(repo
|
||||
.find(ShadowResultLookupKey::TraceFingerprint {
|
||||
trace_id: "trace-1",
|
||||
request_fingerprint: "fp-1",
|
||||
})
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_recent_returns_results_in_descending_update_order() {
|
||||
let repo = InMemoryShadowResultRepository::default();
|
||||
repo.upsert(sample_result("trace-1", "fp-1", 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(sample_result("trace-2", "fp-2", 200))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let recent = repo
|
||||
.list_recent(10)
|
||||
.await
|
||||
.expect("list recent should succeed");
|
||||
assert_eq!(recent.len(), 2);
|
||||
assert_eq!(recent[0].trace_id, "trace-2");
|
||||
assert_eq!(recent[1].trace_id, "trace-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_replaces_existing_shadow_result() {
|
||||
let repo = InMemoryShadowResultRepository::default();
|
||||
repo.upsert(sample_result("trace-1", "fp-1", 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(UpsertShadowResult {
|
||||
trace_id: "trace-1".to_string(),
|
||||
request_fingerprint: "fp-1".to_string(),
|
||||
request_id: Some("req-trace-1".to_string()),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
candidate_id: Some("cand-2".to_string()),
|
||||
rust_result_digest: Some("rust-digest-2".to_string()),
|
||||
python_result_digest: Some("python-digest-2".to_string()),
|
||||
match_status: ShadowResultMatchStatus::Mismatch,
|
||||
status_code: Some(502),
|
||||
error_message: Some("mismatch".to_string()),
|
||||
created_at_unix_secs: 100,
|
||||
updated_at_unix_secs: 200,
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let stored = repo
|
||||
.find(ShadowResultLookupKey::TraceFingerprint {
|
||||
trace_id: "trace-1",
|
||||
request_fingerprint: "fp-1",
|
||||
})
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.expect("stored result should exist");
|
||||
assert_eq!(stored.request_id.as_deref(), Some("req-trace-1"));
|
||||
assert_eq!(stored.candidate_id.as_deref(), Some("cand-2"));
|
||||
assert_eq!(stored.match_status, ShadowResultMatchStatus::Mismatch);
|
||||
}
|
||||
}
|
||||
12
crates/aether-data/src/repository/shadow_results/mod.rs
Normal file
12
crates/aether-data/src/repository/shadow_results/mod.rs
Normal file
@@ -0,0 +1,12 @@
|
||||
mod memory;
|
||||
mod record;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryShadowResultRepository;
|
||||
pub use record::{merge_shadow_result_sample, RecordShadowResultSample, ShadowResultSampleOrigin};
|
||||
pub use sql::SqlxShadowResultRepository;
|
||||
pub use types::{
|
||||
ShadowResultLookupKey, ShadowResultMatchStatus, ShadowResultReadRepository,
|
||||
ShadowResultRepository, ShadowResultWriteRepository, StoredShadowResult, UpsertShadowResult,
|
||||
};
|
||||
188
crates/aether-data/src/repository/shadow_results/record.rs
Normal file
188
crates/aether-data/src/repository/shadow_results/record.rs
Normal file
@@ -0,0 +1,188 @@
|
||||
use super::types::{ShadowResultMatchStatus, StoredShadowResult, UpsertShadowResult};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ShadowResultSampleOrigin {
|
||||
Rust,
|
||||
Python,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RecordShadowResultSample {
|
||||
pub trace_id: String,
|
||||
pub request_fingerprint: String,
|
||||
pub request_id: Option<String>,
|
||||
pub route_family: Option<String>,
|
||||
pub route_kind: Option<String>,
|
||||
pub candidate_id: Option<String>,
|
||||
pub origin: ShadowResultSampleOrigin,
|
||||
pub result_digest: String,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_message: Option<String>,
|
||||
pub recorded_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
pub fn merge_shadow_result_sample(
|
||||
existing: Option<&StoredShadowResult>,
|
||||
sample: RecordShadowResultSample,
|
||||
) -> UpsertShadowResult {
|
||||
let RecordShadowResultSample {
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
request_id,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
origin,
|
||||
result_digest,
|
||||
status_code,
|
||||
error_message,
|
||||
recorded_at_unix_secs,
|
||||
} = sample;
|
||||
|
||||
let (rust_result_digest, python_result_digest) = match origin {
|
||||
ShadowResultSampleOrigin::Rust => (
|
||||
Some(result_digest),
|
||||
existing.and_then(|row| row.python_result_digest.clone()),
|
||||
),
|
||||
ShadowResultSampleOrigin::Python => (
|
||||
existing.and_then(|row| row.rust_result_digest.clone()),
|
||||
Some(result_digest),
|
||||
),
|
||||
};
|
||||
|
||||
let match_status = resolve_match_status(
|
||||
rust_result_digest.as_deref(),
|
||||
python_result_digest.as_deref(),
|
||||
);
|
||||
|
||||
UpsertShadowResult {
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
request_id: request_id.or_else(|| existing.and_then(|row| row.request_id.clone())),
|
||||
route_family: route_family.or_else(|| existing.and_then(|row| row.route_family.clone())),
|
||||
route_kind: route_kind.or_else(|| existing.and_then(|row| row.route_kind.clone())),
|
||||
candidate_id: candidate_id.or_else(|| existing.and_then(|row| row.candidate_id.clone())),
|
||||
rust_result_digest,
|
||||
python_result_digest,
|
||||
match_status,
|
||||
status_code: status_code.or(existing.and_then(|row| row.status_code)),
|
||||
error_message: resolve_error_message(existing, error_message, match_status),
|
||||
created_at_unix_secs: existing
|
||||
.map(|row| row.created_at_unix_secs)
|
||||
.unwrap_or(recorded_at_unix_secs),
|
||||
updated_at_unix_secs: recorded_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_match_status(
|
||||
rust_result_digest: Option<&str>,
|
||||
python_result_digest: Option<&str>,
|
||||
) -> ShadowResultMatchStatus {
|
||||
match (rust_result_digest, python_result_digest) {
|
||||
(Some(rust_digest), Some(python_digest)) if rust_digest == python_digest => {
|
||||
ShadowResultMatchStatus::Match
|
||||
}
|
||||
(Some(_), Some(_)) => ShadowResultMatchStatus::Mismatch,
|
||||
_ => ShadowResultMatchStatus::Pending,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_error_message(
|
||||
existing: Option<&StoredShadowResult>,
|
||||
error_message: Option<String>,
|
||||
match_status: ShadowResultMatchStatus,
|
||||
) -> Option<String> {
|
||||
if match_status == ShadowResultMatchStatus::Mismatch {
|
||||
error_message
|
||||
.or_else(|| existing.and_then(|row| row.error_message.clone()))
|
||||
.or_else(|| Some("shadow result digest mismatch".to_string()))
|
||||
} else {
|
||||
error_message.or_else(|| existing.and_then(|row| row.error_message.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{merge_shadow_result_sample, RecordShadowResultSample, ShadowResultSampleOrigin};
|
||||
use crate::repository::shadow_results::{ShadowResultMatchStatus, UpsertShadowResult};
|
||||
|
||||
fn rust_sample(result_digest: &str, recorded_at_unix_secs: u64) -> RecordShadowResultSample {
|
||||
RecordShadowResultSample {
|
||||
trace_id: "trace-1".to_string(),
|
||||
request_fingerprint: "fp-1".to_string(),
|
||||
request_id: Some("req-1".to_string()),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
candidate_id: None,
|
||||
origin: ShadowResultSampleOrigin::Rust,
|
||||
result_digest: result_digest.to_string(),
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
recorded_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
fn python_sample(result_digest: &str, recorded_at_unix_secs: u64) -> RecordShadowResultSample {
|
||||
RecordShadowResultSample {
|
||||
trace_id: "trace-1".to_string(),
|
||||
request_fingerprint: "fp-1".to_string(),
|
||||
request_id: Some("req-1".to_string()),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
candidate_id: None,
|
||||
origin: ShadowResultSampleOrigin::Python,
|
||||
result_digest: result_digest.to_string(),
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
recorded_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
fn stored(upsert: UpsertShadowResult) -> crate::repository::shadow_results::StoredShadowResult {
|
||||
upsert.into_stored()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_pending_until_both_samples_exist() {
|
||||
let merged = merge_shadow_result_sample(None, rust_sample("digest-1", 100));
|
||||
|
||||
assert_eq!(merged.match_status, ShadowResultMatchStatus::Pending);
|
||||
assert_eq!(merged.request_id.as_deref(), Some("req-1"));
|
||||
assert_eq!(merged.rust_result_digest.as_deref(), Some("digest-1"));
|
||||
assert!(merged.python_result_digest.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn marks_match_when_rust_and_python_digests_are_equal() {
|
||||
let existing = stored(merge_shadow_result_sample(
|
||||
None,
|
||||
rust_sample("digest-1", 100),
|
||||
));
|
||||
let merged = merge_shadow_result_sample(Some(&existing), python_sample("digest-1", 200));
|
||||
|
||||
assert_eq!(merged.match_status, ShadowResultMatchStatus::Match);
|
||||
assert_eq!(merged.created_at_unix_secs, 100);
|
||||
assert_eq!(merged.updated_at_unix_secs, 200);
|
||||
assert_eq!(merged.request_id.as_deref(), Some("req-1"));
|
||||
assert_eq!(merged.rust_result_digest.as_deref(), Some("digest-1"));
|
||||
assert_eq!(merged.python_result_digest.as_deref(), Some("digest-1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn marks_mismatch_when_rust_and_python_digests_differ() {
|
||||
let existing = stored(merge_shadow_result_sample(
|
||||
None,
|
||||
rust_sample("digest-1", 100),
|
||||
));
|
||||
let merged = merge_shadow_result_sample(Some(&existing), python_sample("digest-2", 200));
|
||||
|
||||
assert_eq!(merged.match_status, ShadowResultMatchStatus::Mismatch);
|
||||
assert_eq!(
|
||||
merged.error_message.as_deref(),
|
||||
Some("shadow result digest mismatch")
|
||||
);
|
||||
assert_eq!(merged.request_id.as_deref(), Some("req-1"));
|
||||
assert_eq!(merged.rust_result_digest.as_deref(), Some("digest-1"));
|
||||
assert_eq!(merged.python_result_digest.as_deref(), Some("digest-2"));
|
||||
}
|
||||
}
|
||||
283
crates/aether-data/src/repository/shadow_results/sql.rs
Normal file
283
crates/aether-data/src/repository/shadow_results/sql.rs
Normal file
@@ -0,0 +1,283 @@
|
||||
use async_trait::async_trait;
|
||||
use futures_util::future::BoxFuture;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{
|
||||
ShadowResultLookupKey, ShadowResultMatchStatus, ShadowResultReadRepository,
|
||||
ShadowResultWriteRepository, StoredShadowResult, UpsertShadowResult,
|
||||
};
|
||||
use crate::postgres::PostgresTransactionRunner;
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_BY_TRACE_FINGERPRINT_SQL: &str = r#"
|
||||
SELECT
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
NULL::TEXT AS request_id,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
rust_result_digest,
|
||||
python_result_digest,
|
||||
match_status,
|
||||
status_code,
|
||||
error_message,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
FROM gateway_shadow_results
|
||||
WHERE trace_id = $1 AND request_fingerprint = $2
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const LIST_RECENT_SQL: &str = r#"
|
||||
SELECT
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
NULL::TEXT AS request_id,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
rust_result_digest,
|
||||
python_result_digest,
|
||||
match_status,
|
||||
status_code,
|
||||
error_message,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
FROM gateway_shadow_results
|
||||
ORDER BY updated_at DESC
|
||||
LIMIT $1
|
||||
"#;
|
||||
|
||||
const UPSERT_SQL: &str = r#"
|
||||
INSERT INTO gateway_shadow_results (
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
rust_result_digest,
|
||||
python_result_digest,
|
||||
match_status,
|
||||
status_code,
|
||||
error_message,
|
||||
created_at,
|
||||
updated_at
|
||||
) VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5,
|
||||
$6,
|
||||
$7,
|
||||
$8,
|
||||
$9,
|
||||
$10,
|
||||
TO_TIMESTAMP($11::double precision),
|
||||
TO_TIMESTAMP($12::double precision)
|
||||
)
|
||||
ON CONFLICT (trace_id, request_fingerprint)
|
||||
DO UPDATE SET
|
||||
route_family = EXCLUDED.route_family,
|
||||
route_kind = EXCLUDED.route_kind,
|
||||
candidate_id = EXCLUDED.candidate_id,
|
||||
rust_result_digest = EXCLUDED.rust_result_digest,
|
||||
python_result_digest = EXCLUDED.python_result_digest,
|
||||
match_status = EXCLUDED.match_status,
|
||||
status_code = EXCLUDED.status_code,
|
||||
error_message = EXCLUDED.error_message,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
RETURNING
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
NULL::TEXT AS request_id,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
rust_result_digest,
|
||||
python_result_digest,
|
||||
match_status,
|
||||
status_code,
|
||||
error_message,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxShadowResultRepository {
|
||||
pool: PgPool,
|
||||
tx_runner: PostgresTransactionRunner,
|
||||
}
|
||||
|
||||
impl SqlxShadowResultRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
let tx_runner = PostgresTransactionRunner::new(pool.clone());
|
||||
Self { pool, tx_runner }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub fn transaction_runner(&self) -> &PostgresTransactionRunner {
|
||||
&self.tx_runner
|
||||
}
|
||||
|
||||
pub async fn find(
|
||||
&self,
|
||||
key: ShadowResultLookupKey<'_>,
|
||||
) -> Result<Option<StoredShadowResult>, DataLayerError> {
|
||||
match key {
|
||||
ShadowResultLookupKey::TraceFingerprint {
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
} => {
|
||||
self.find_by_trace_fingerprint(trace_id, request_fingerprint)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn find_by_trace_fingerprint(
|
||||
&self,
|
||||
trace_id: &str,
|
||||
request_fingerprint: &str,
|
||||
) -> Result<Option<StoredShadowResult>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_TRACE_FINGERPRINT_SQL)
|
||||
.bind(trace_id)
|
||||
.bind(request_fingerprint)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_shadow_result_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn list_recent(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredShadowResult>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = sqlx::query(LIST_RECENT_SQL)
|
||||
.bind(i64::try_from(limit).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"invalid recent shadow result limit: {limit}"
|
||||
))
|
||||
})?)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_shadow_result_row).collect()
|
||||
}
|
||||
|
||||
pub async fn upsert(
|
||||
&self,
|
||||
result: UpsertShadowResult,
|
||||
) -> Result<StoredShadowResult, DataLayerError> {
|
||||
self.tx_runner
|
||||
.run_read_write(|tx| {
|
||||
Box::pin(async move {
|
||||
let row = sqlx::query(UPSERT_SQL)
|
||||
.bind(&result.trace_id)
|
||||
.bind(&result.request_fingerprint)
|
||||
.bind(&result.route_family)
|
||||
.bind(&result.route_kind)
|
||||
.bind(&result.candidate_id)
|
||||
.bind(&result.rust_result_digest)
|
||||
.bind(&result.python_result_digest)
|
||||
.bind(match_status_to_database(result.match_status))
|
||||
.bind(result.status_code.map(i32::from))
|
||||
.bind(&result.error_message)
|
||||
.bind(result.created_at_unix_secs as f64)
|
||||
.bind(result.updated_at_unix_secs as f64)
|
||||
.fetch_one(&mut **tx)
|
||||
.await?;
|
||||
map_shadow_result_row(&row)
|
||||
}) as BoxFuture<'_, Result<StoredShadowResult, DataLayerError>>
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ShadowResultReadRepository for SqlxShadowResultRepository {
|
||||
async fn find(
|
||||
&self,
|
||||
key: ShadowResultLookupKey<'_>,
|
||||
) -> Result<Option<StoredShadowResult>, DataLayerError> {
|
||||
Self::find(self, key).await
|
||||
}
|
||||
|
||||
async fn list_recent(&self, limit: usize) -> Result<Vec<StoredShadowResult>, DataLayerError> {
|
||||
Self::list_recent(self, limit).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ShadowResultWriteRepository for SqlxShadowResultRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
result: UpsertShadowResult,
|
||||
) -> Result<StoredShadowResult, DataLayerError> {
|
||||
Self::upsert(self, result).await
|
||||
}
|
||||
}
|
||||
|
||||
fn match_status_to_database(status: ShadowResultMatchStatus) -> &'static str {
|
||||
match status {
|
||||
ShadowResultMatchStatus::Pending => "pending",
|
||||
ShadowResultMatchStatus::Match => "match",
|
||||
ShadowResultMatchStatus::Mismatch => "mismatch",
|
||||
ShadowResultMatchStatus::Error => "error",
|
||||
}
|
||||
}
|
||||
|
||||
fn map_shadow_result_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> Result<StoredShadowResult, DataLayerError> {
|
||||
let match_status =
|
||||
ShadowResultMatchStatus::from_database(row.try_get::<String, _>("match_status")?.as_str())?;
|
||||
StoredShadowResult::new(
|
||||
row.try_get("trace_id")?,
|
||||
row.try_get("request_fingerprint")?,
|
||||
row.try_get("request_id")?,
|
||||
row.try_get("route_family")?,
|
||||
row.try_get("route_kind")?,
|
||||
row.try_get("candidate_id")?,
|
||||
row.try_get("rust_result_digest")?,
|
||||
row.try_get("python_result_digest")?,
|
||||
match_status,
|
||||
row.try_get("status_code")?,
|
||||
row.try_get("error_message")?,
|
||||
row.try_get("created_at_unix_secs")?,
|
||||
row.try_get("updated_at_unix_secs")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxShadowResultRepository;
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_constructs_from_lazy_pool() {
|
||||
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("factory should build");
|
||||
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let repository = SqlxShadowResultRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
let _ = repository.transaction_runner();
|
||||
}
|
||||
}
|
||||
227
crates/aether-data/src/repository/shadow_results/types.rs
Normal file
227
crates/aether-data/src/repository/shadow_results/types.rs
Normal file
@@ -0,0 +1,227 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub enum ShadowResultMatchStatus {
|
||||
Pending,
|
||||
Match,
|
||||
Mismatch,
|
||||
Error,
|
||||
}
|
||||
|
||||
impl ShadowResultMatchStatus {
|
||||
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"pending" => Ok(Self::Pending),
|
||||
"match" => Ok(Self::Match),
|
||||
"mismatch" => Ok(Self::Mismatch),
|
||||
"error" => Ok(Self::Error),
|
||||
other => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"unsupported gateway_shadow_results.match_status: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredShadowResult {
|
||||
pub trace_id: String,
|
||||
pub request_fingerprint: String,
|
||||
pub request_id: Option<String>,
|
||||
pub route_family: Option<String>,
|
||||
pub route_kind: Option<String>,
|
||||
pub candidate_id: Option<String>,
|
||||
pub rust_result_digest: Option<String>,
|
||||
pub python_result_digest: Option<String>,
|
||||
pub match_status: ShadowResultMatchStatus,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_message: Option<String>,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
impl StoredShadowResult {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
trace_id: String,
|
||||
request_fingerprint: String,
|
||||
request_id: Option<String>,
|
||||
route_family: Option<String>,
|
||||
route_kind: Option<String>,
|
||||
candidate_id: Option<String>,
|
||||
rust_result_digest: Option<String>,
|
||||
python_result_digest: Option<String>,
|
||||
match_status: ShadowResultMatchStatus,
|
||||
status_code: Option<i32>,
|
||||
error_message: Option<String>,
|
||||
created_at_unix_secs: i64,
|
||||
updated_at_unix_secs: i64,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
let status_code = status_code
|
||||
.map(|value| {
|
||||
u16::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid status_code: {value}"))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid created_at_unix_secs: {created_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
let updated_at_unix_secs = u64::try_from(updated_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid updated_at_unix_secs: {updated_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(Self {
|
||||
trace_id,
|
||||
request_fingerprint,
|
||||
request_id,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
rust_result_digest,
|
||||
python_result_digest,
|
||||
match_status,
|
||||
status_code,
|
||||
error_message,
|
||||
created_at_unix_secs,
|
||||
updated_at_unix_secs,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpsertShadowResult {
|
||||
pub trace_id: String,
|
||||
pub request_fingerprint: String,
|
||||
pub request_id: Option<String>,
|
||||
pub route_family: Option<String>,
|
||||
pub route_kind: Option<String>,
|
||||
pub candidate_id: Option<String>,
|
||||
pub rust_result_digest: Option<String>,
|
||||
pub python_result_digest: Option<String>,
|
||||
pub match_status: ShadowResultMatchStatus,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_message: Option<String>,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
impl UpsertShadowResult {
|
||||
pub fn into_stored(self) -> StoredShadowResult {
|
||||
StoredShadowResult {
|
||||
trace_id: self.trace_id,
|
||||
request_fingerprint: self.request_fingerprint,
|
||||
request_id: self.request_id,
|
||||
route_family: self.route_family,
|
||||
route_kind: self.route_kind,
|
||||
candidate_id: self.candidate_id,
|
||||
rust_result_digest: self.rust_result_digest,
|
||||
python_result_digest: self.python_result_digest,
|
||||
match_status: self.match_status,
|
||||
status_code: self.status_code,
|
||||
error_message: self.error_message,
|
||||
created_at_unix_secs: self.created_at_unix_secs,
|
||||
updated_at_unix_secs: self.updated_at_unix_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ShadowResultLookupKey<'a> {
|
||||
TraceFingerprint {
|
||||
trace_id: &'a str,
|
||||
request_fingerprint: &'a str,
|
||||
},
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ShadowResultReadRepository: Send + Sync {
|
||||
async fn find(
|
||||
&self,
|
||||
key: ShadowResultLookupKey<'_>,
|
||||
) -> Result<Option<StoredShadowResult>, crate::DataLayerError>;
|
||||
|
||||
async fn list_recent(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredShadowResult>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ShadowResultWriteRepository: Send + Sync {
|
||||
async fn upsert(
|
||||
&self,
|
||||
result: UpsertShadowResult,
|
||||
) -> Result<StoredShadowResult, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait ShadowResultRepository:
|
||||
ShadowResultReadRepository + ShadowResultWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> ShadowResultRepository for T where
|
||||
T: ShadowResultReadRepository + ShadowResultWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{ShadowResultMatchStatus, StoredShadowResult};
|
||||
|
||||
#[test]
|
||||
fn parses_match_status_from_database_text() {
|
||||
assert_eq!(
|
||||
ShadowResultMatchStatus::from_database("match").expect("status should parse"),
|
||||
ShadowResultMatchStatus::Match
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_database_status() {
|
||||
assert!(ShadowResultMatchStatus::from_database("mystery").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_numeric_fields() {
|
||||
assert!(StoredShadowResult::new(
|
||||
"trace-1".to_string(),
|
||||
"fp-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
ShadowResultMatchStatus::Pending,
|
||||
Some(-1),
|
||||
None,
|
||||
1,
|
||||
1,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_updated_at_values() {
|
||||
assert!(StoredShadowResult::new(
|
||||
"trace-1".to_string(),
|
||||
"fp-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
ShadowResultMatchStatus::Pending,
|
||||
Some(200),
|
||||
None,
|
||||
1,
|
||||
-1,
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
107
crates/aether-data/src/repository/usage/memory.rs
Normal file
107
crates/aether-data/src/repository/usage/memory.rs
Normal file
@@ -0,0 +1,107 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryUsageReadRepository {
|
||||
by_request_id: RwLock<BTreeMap<String, StoredRequestUsageAudit>>,
|
||||
}
|
||||
|
||||
impl InMemoryUsageReadRepository {
|
||||
pub fn seed<I>(items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredRequestUsageAudit>,
|
||||
{
|
||||
let mut by_request_id = BTreeMap::new();
|
||||
for item in items {
|
||||
by_request_id.insert(item.request_id.clone(), item);
|
||||
}
|
||||
Self {
|
||||
by_request_id: RwLock::new(by_request_id),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Ok(self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.get(request_id)
|
||||
.cloned())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryUsageReadRepository;
|
||||
use crate::repository::usage::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
|
||||
fn sample_usage(request_id: &str, created_at_unix_secs: i64) -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
"OpenAI".to_string(),
|
||||
"gpt-4.1".to_string(),
|
||||
Some("gpt-4.1-mini".to_string()),
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
false,
|
||||
100,
|
||||
50,
|
||||
150,
|
||||
0.12,
|
||||
0.18,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(420),
|
||||
Some(120),
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
created_at_unix_secs,
|
||||
created_at_unix_secs + 1,
|
||||
Some(created_at_unix_secs + 2),
|
||||
)
|
||||
.expect("usage should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn finds_usage_by_request_id() {
|
||||
let repository = InMemoryUsageReadRepository::seed(vec![
|
||||
sample_usage("req-1", 100),
|
||||
sample_usage("req-2", 200),
|
||||
]);
|
||||
|
||||
let usage = repository
|
||||
.find_by_request_id("req-2")
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.expect("usage should exist");
|
||||
|
||||
assert_eq!(usage.request_id, "req-2");
|
||||
assert_eq!(usage.total_tokens, 150);
|
||||
}
|
||||
}
|
||||
7
crates/aether-data/src/repository/usage/mod.rs
Normal file
7
crates/aether-data/src/repository/usage/mod.rs
Normal file
@@ -0,0 +1,7 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryUsageReadRepository;
|
||||
pub use sql::SqlxUsageReadRepository;
|
||||
pub use types::{StoredRequestUsageAudit, UsageReadRepository, UsageRepository};
|
||||
150
crates/aether-data/src/repository/usage/sql.rs
Normal file
150
crates/aether-data/src/repository/usage/sql.rs
Normal file
@@ -0,0 +1,150 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::types::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_BY_REQUEST_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
COALESCE(has_format_conversion, FALSE) AS has_format_conversion,
|
||||
COALESCE(is_stream, FALSE) AS is_stream,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
|
||||
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms,
|
||||
first_byte_time_ms,
|
||||
status,
|
||||
billing_status,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
FROM "usage"
|
||||
WHERE request_id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxUsageReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxUsageReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_REQUEST_ID_SQL)
|
||||
.bind(request_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_usage_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageReadRepository for SqlxUsageReadRepository {
|
||||
async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Self::find_by_request_id(self, request_id).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_usage_row(row: &sqlx::postgres::PgRow) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
||||
StoredRequestUsageAudit::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("request_id")?,
|
||||
row.try_get("user_id")?,
|
||||
row.try_get("api_key_id")?,
|
||||
row.try_get("username")?,
|
||||
row.try_get("api_key_name")?,
|
||||
row.try_get("provider_name")?,
|
||||
row.try_get("model")?,
|
||||
row.try_get("target_model")?,
|
||||
row.try_get("provider_id")?,
|
||||
row.try_get("provider_endpoint_id")?,
|
||||
row.try_get("provider_api_key_id")?,
|
||||
row.try_get("request_type")?,
|
||||
row.try_get("api_format")?,
|
||||
row.try_get("api_family")?,
|
||||
row.try_get("endpoint_kind")?,
|
||||
row.try_get("endpoint_api_format")?,
|
||||
row.try_get("provider_api_family")?,
|
||||
row.try_get("provider_endpoint_kind")?,
|
||||
row.try_get("has_format_conversion")?,
|
||||
row.try_get("is_stream")?,
|
||||
row.try_get("input_tokens")?,
|
||||
row.try_get("output_tokens")?,
|
||||
row.try_get("total_tokens")?,
|
||||
row.try_get("total_cost_usd")?,
|
||||
row.try_get("actual_total_cost_usd")?,
|
||||
row.try_get("status_code")?,
|
||||
row.try_get("error_message")?,
|
||||
row.try_get("error_category")?,
|
||||
row.try_get("response_time_ms")?,
|
||||
row.try_get("first_byte_time_ms")?,
|
||||
row.try_get("status")?,
|
||||
row.try_get("billing_status")?,
|
||||
row.try_get("created_at_unix_secs")?,
|
||||
row.try_get("updated_at_unix_secs")?,
|
||||
row.try_get("finalized_at_unix_secs")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxUsageReadRepository;
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_constructs_from_lazy_pool() {
|
||||
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("factory should build");
|
||||
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let repository = SqlxUsageReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
}
|
||||
}
|
||||
304
crates/aether-data/src/repository/usage/types.rs
Normal file
304
crates/aether-data/src/repository/usage/types.rs
Normal file
@@ -0,0 +1,304 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredRequestUsageAudit {
|
||||
pub id: String,
|
||||
pub request_id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub api_key_id: Option<String>,
|
||||
pub username: Option<String>,
|
||||
pub api_key_name: Option<String>,
|
||||
pub provider_name: String,
|
||||
pub model: String,
|
||||
pub target_model: Option<String>,
|
||||
pub provider_id: Option<String>,
|
||||
pub provider_endpoint_id: Option<String>,
|
||||
pub provider_api_key_id: Option<String>,
|
||||
pub request_type: Option<String>,
|
||||
pub api_format: Option<String>,
|
||||
pub api_family: Option<String>,
|
||||
pub endpoint_kind: Option<String>,
|
||||
pub endpoint_api_format: Option<String>,
|
||||
pub provider_api_family: Option<String>,
|
||||
pub provider_endpoint_kind: Option<String>,
|
||||
pub has_format_conversion: bool,
|
||||
pub is_stream: bool,
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub total_cost_usd: f64,
|
||||
pub actual_total_cost_usd: f64,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_message: Option<String>,
|
||||
pub error_category: Option<String>,
|
||||
pub response_time_ms: Option<u64>,
|
||||
pub first_byte_time_ms: Option<u64>,
|
||||
pub status: String,
|
||||
pub billing_status: String,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
pub finalized_at_unix_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl StoredRequestUsageAudit {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
id: String,
|
||||
request_id: String,
|
||||
user_id: Option<String>,
|
||||
api_key_id: Option<String>,
|
||||
username: Option<String>,
|
||||
api_key_name: Option<String>,
|
||||
provider_name: String,
|
||||
model: String,
|
||||
target_model: Option<String>,
|
||||
provider_id: Option<String>,
|
||||
provider_endpoint_id: Option<String>,
|
||||
provider_api_key_id: Option<String>,
|
||||
request_type: Option<String>,
|
||||
api_format: Option<String>,
|
||||
api_family: Option<String>,
|
||||
endpoint_kind: Option<String>,
|
||||
endpoint_api_format: Option<String>,
|
||||
provider_api_family: Option<String>,
|
||||
provider_endpoint_kind: Option<String>,
|
||||
has_format_conversion: bool,
|
||||
is_stream: bool,
|
||||
input_tokens: i32,
|
||||
output_tokens: i32,
|
||||
total_tokens: i32,
|
||||
total_cost_usd: f64,
|
||||
actual_total_cost_usd: f64,
|
||||
status_code: Option<i32>,
|
||||
error_message: Option<String>,
|
||||
error_category: Option<String>,
|
||||
response_time_ms: Option<i32>,
|
||||
first_byte_time_ms: Option<i32>,
|
||||
status: String,
|
||||
billing_status: String,
|
||||
created_at_unix_secs: i64,
|
||||
updated_at_unix_secs: i64,
|
||||
finalized_at_unix_secs: Option<i64>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if request_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.request_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if provider_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.provider_name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if model.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.model is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if status.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.status is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if billing_status.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.billing_status is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if !total_cost_usd.is_finite() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.total_cost_usd is not finite".to_string(),
|
||||
));
|
||||
}
|
||||
if !actual_total_cost_usd.is_finite() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"usage.actual_total_cost_usd is not finite".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
has_format_conversion,
|
||||
is_stream,
|
||||
input_tokens: parse_u64(input_tokens, "usage.input_tokens")?,
|
||||
output_tokens: parse_u64(output_tokens, "usage.output_tokens")?,
|
||||
total_tokens: parse_u64(total_tokens, "usage.total_tokens")?,
|
||||
total_cost_usd,
|
||||
actual_total_cost_usd,
|
||||
status_code: parse_u16(status_code, "usage.status_code")?,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms: parse_optional_u64(response_time_ms, "usage.response_time_ms")?,
|
||||
first_byte_time_ms: parse_optional_u64(first_byte_time_ms, "usage.first_byte_time_ms")?,
|
||||
status,
|
||||
billing_status,
|
||||
created_at_unix_secs: parse_timestamp(
|
||||
created_at_unix_secs,
|
||||
"usage.created_at_unix_secs",
|
||||
)?,
|
||||
updated_at_unix_secs: parse_timestamp(
|
||||
updated_at_unix_secs,
|
||||
"usage.updated_at_unix_secs",
|
||||
)?,
|
||||
finalized_at_unix_secs: finalized_at_unix_secs
|
||||
.map(|value| parse_timestamp(value, "usage.finalized_at_unix_secs"))
|
||||
.transpose()?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait UsageReadRepository: Send + Sync {
|
||||
async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait UsageRepository: UsageReadRepository + Send + Sync {}
|
||||
|
||||
impl<T> UsageRepository for T where T: UsageReadRepository + Send + Sync {}
|
||||
|
||||
fn parse_u64(value: i32, field_name: &str) -> Result<u64, crate::DataLayerError> {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_optional_u64(
|
||||
value: Option<i32>,
|
||||
field_name: &str,
|
||||
) -> Result<Option<u64>, crate::DataLayerError> {
|
||||
value
|
||||
.map(|value| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn parse_u16(value: Option<i32>, field_name: &str) -> Result<Option<u16>, crate::DataLayerError> {
|
||||
value
|
||||
.map(|value| {
|
||||
u16::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn parse_timestamp(value: i64, field_name: &str) -> Result<u64, crate::DataLayerError> {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::StoredRequestUsageAudit;
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_request_id() {
|
||||
assert!(StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
"".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
"OpenAI".to_string(),
|
||||
"gpt-4.1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
false,
|
||||
false,
|
||||
10,
|
||||
20,
|
||||
30,
|
||||
0.1,
|
||||
0.1,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(120),
|
||||
Some(80),
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
100,
|
||||
101,
|
||||
Some(102),
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_token_count() {
|
||||
assert!(StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
"OpenAI".to_string(),
|
||||
"gpt-4.1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
false,
|
||||
false,
|
||||
-1,
|
||||
20,
|
||||
30,
|
||||
0.1,
|
||||
0.1,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(120),
|
||||
Some(80),
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
100,
|
||||
101,
|
||||
Some(102),
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
232
crates/aether-data/src/repository/video_tasks/memory.rs
Normal file
232
crates/aether-data/src/repository/video_tasks/memory.rs
Normal file
@@ -0,0 +1,232 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{
|
||||
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
|
||||
VideoTaskWriteRepository,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct MemoryVideoTaskIndex {
|
||||
by_id: BTreeMap<String, StoredVideoTask>,
|
||||
short_to_id: BTreeMap<String, String>,
|
||||
user_external_to_id: BTreeMap<(String, String), String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryVideoTaskRepository {
|
||||
index: RwLock<MemoryVideoTaskIndex>,
|
||||
}
|
||||
|
||||
impl InMemoryVideoTaskRepository {
|
||||
fn store_locked(index: &mut MemoryVideoTaskIndex, task: StoredVideoTask) -> StoredVideoTask {
|
||||
if let Some(previous) = index.by_id.insert(task.id.clone(), task.clone()) {
|
||||
if let Some(short_id) = previous.short_id {
|
||||
index.short_to_id.remove(&short_id);
|
||||
}
|
||||
if let (Some(user_id), Some(external_task_id)) =
|
||||
(previous.user_id, previous.external_task_id)
|
||||
{
|
||||
index
|
||||
.user_external_to_id
|
||||
.remove(&(user_id, external_task_id));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(short_id) = &task.short_id {
|
||||
index.short_to_id.insert(short_id.clone(), task.id.clone());
|
||||
}
|
||||
if let (Some(user_id), Some(external_task_id)) = (&task.user_id, &task.external_task_id) {
|
||||
index
|
||||
.user_external_to_id
|
||||
.insert((user_id.clone(), external_task_id.clone()), task.id.clone());
|
||||
}
|
||||
|
||||
task
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl VideoTaskReadRepository for InMemoryVideoTaskRepository {
|
||||
async fn find(
|
||||
&self,
|
||||
key: VideoTaskLookupKey<'_>,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
let index = self.index.read().expect("video task repository lock");
|
||||
Ok(match key {
|
||||
VideoTaskLookupKey::Id(id) => index.by_id.get(id).cloned(),
|
||||
VideoTaskLookupKey::ShortId(short_id) => index
|
||||
.short_to_id
|
||||
.get(short_id)
|
||||
.and_then(|id| index.by_id.get(id))
|
||||
.cloned(),
|
||||
VideoTaskLookupKey::UserExternal {
|
||||
user_id,
|
||||
external_task_id,
|
||||
} => index
|
||||
.user_external_to_id
|
||||
.get(&(user_id.to_string(), external_task_id.to_string()))
|
||||
.and_then(|id| index.by_id.get(id))
|
||||
.cloned(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_active(&self, limit: usize) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut tasks = self
|
||||
.index
|
||||
.read()
|
||||
.expect("video task repository lock")
|
||||
.by_id
|
||||
.values()
|
||||
.filter(|task| task.status.is_active())
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
tasks.sort_by(|left, right| right.updated_at_unix_secs.cmp(&left.updated_at_unix_secs));
|
||||
tasks.truncate(limit);
|
||||
Ok(tasks)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl VideoTaskWriteRepository for InMemoryVideoTaskRepository {
|
||||
async fn upsert(&self, task: UpsertVideoTask) -> Result<StoredVideoTask, DataLayerError> {
|
||||
let mut index = self.index.write().expect("video task repository lock");
|
||||
Ok(Self::store_locked(&mut index, task.into_stored()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryVideoTaskRepository;
|
||||
use crate::repository::video_tasks::{
|
||||
UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository, VideoTaskStatus,
|
||||
VideoTaskWriteRepository,
|
||||
};
|
||||
|
||||
fn sample_task(
|
||||
id: &str,
|
||||
status: VideoTaskStatus,
|
||||
updated_at_unix_secs: u64,
|
||||
) -> UpsertVideoTask {
|
||||
UpsertVideoTask {
|
||||
id: id.to_string(),
|
||||
short_id: Some(format!("short-{id}")),
|
||||
user_id: Some("user-1".to_string()),
|
||||
external_task_id: Some(format!("ext-{id}")),
|
||||
provider_api_format: Some("openai:video".to_string()),
|
||||
model: Some("sora-2".to_string()),
|
||||
prompt: Some("hello".to_string()),
|
||||
size: Some("1280x720".to_string()),
|
||||
status,
|
||||
progress_percent: 0,
|
||||
created_at_unix_secs: updated_at_unix_secs.saturating_sub(10),
|
||||
updated_at_unix_secs,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_task_by_all_supported_lookup_keys() {
|
||||
let repo = InMemoryVideoTaskRepository::default();
|
||||
repo.upsert(sample_task("task-1", VideoTaskStatus::Submitted, 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
assert!(repo
|
||||
.find(VideoTaskLookupKey::Id("task-1"))
|
||||
.await
|
||||
.expect("find by id should succeed")
|
||||
.is_some());
|
||||
assert!(repo
|
||||
.find(VideoTaskLookupKey::ShortId("short-task-1"))
|
||||
.await
|
||||
.expect("find by short id should succeed")
|
||||
.is_some());
|
||||
assert!(repo
|
||||
.find(VideoTaskLookupKey::UserExternal {
|
||||
user_id: "user-1",
|
||||
external_task_id: "ext-task-1",
|
||||
})
|
||||
.await
|
||||
.expect("find by user/external should succeed")
|
||||
.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_active_only_returns_active_tasks_in_descending_update_order() {
|
||||
let repo = InMemoryVideoTaskRepository::default();
|
||||
repo.upsert(sample_task("task-1", VideoTaskStatus::Completed, 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(sample_task("task-2", VideoTaskStatus::Processing, 200))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
repo.upsert(sample_task("task-3", VideoTaskStatus::Queued, 150))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let active = repo
|
||||
.list_active(10)
|
||||
.await
|
||||
.expect("list active should succeed");
|
||||
assert_eq!(active.len(), 2);
|
||||
assert_eq!(active[0].id, "task-2");
|
||||
assert_eq!(active[1].id, "task-3");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_replaces_secondary_indexes() {
|
||||
let repo = InMemoryVideoTaskRepository::default();
|
||||
repo.upsert(sample_task("task-1", VideoTaskStatus::Submitted, 100))
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
repo.upsert(UpsertVideoTask {
|
||||
id: "task-1".to_string(),
|
||||
short_id: Some("short-task-1b".to_string()),
|
||||
user_id: Some("user-2".to_string()),
|
||||
external_task_id: Some("ext-task-1b".to_string()),
|
||||
provider_api_format: Some("gemini:video".to_string()),
|
||||
model: Some("veo-3".to_string()),
|
||||
prompt: Some("remix".to_string()),
|
||||
size: Some("720p".to_string()),
|
||||
status: VideoTaskStatus::Processing,
|
||||
progress_percent: 50,
|
||||
created_at_unix_secs: 150,
|
||||
updated_at_unix_secs: 200,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
assert!(repo
|
||||
.find(VideoTaskLookupKey::ShortId("short-task-1"))
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.is_none());
|
||||
assert!(repo
|
||||
.find(VideoTaskLookupKey::UserExternal {
|
||||
user_id: "user-1",
|
||||
external_task_id: "ext-task-1",
|
||||
})
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.is_none());
|
||||
assert!(repo
|
||||
.find(VideoTaskLookupKey::ShortId("short-task-1b"))
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.is_some());
|
||||
}
|
||||
}
|
||||
10
crates/aether-data/src/repository/video_tasks/mod.rs
Normal file
10
crates/aether-data/src/repository/video_tasks/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod memory;
|
||||
mod sql;
|
||||
mod types;
|
||||
|
||||
pub use memory::InMemoryVideoTaskRepository;
|
||||
pub use sql::SqlxVideoTaskReadRepository;
|
||||
pub use types::{
|
||||
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
|
||||
VideoTaskRepository, VideoTaskStatus, VideoTaskWriteRepository,
|
||||
};
|
||||
259
crates/aether-data/src/repository/video_tasks/sql.rs
Normal file
259
crates/aether-data/src/repository/video_tasks/sql.rs
Normal file
@@ -0,0 +1,259 @@
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::repository::video_tasks::{
|
||||
StoredVideoTask, VideoTaskLookupKey, VideoTaskReadRepository, VideoTaskStatus,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_BY_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
short_id,
|
||||
user_id,
|
||||
external_task_id,
|
||||
provider_api_format,
|
||||
model,
|
||||
prompt,
|
||||
size,
|
||||
status,
|
||||
progress_percent,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
,
|
||||
error_code,
|
||||
error_message,
|
||||
video_url
|
||||
FROM video_tasks
|
||||
WHERE id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_BY_SHORT_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
short_id,
|
||||
user_id,
|
||||
external_task_id,
|
||||
provider_api_format,
|
||||
model,
|
||||
prompt,
|
||||
size,
|
||||
status,
|
||||
progress_percent,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
,
|
||||
error_code,
|
||||
error_message,
|
||||
video_url
|
||||
FROM video_tasks
|
||||
WHERE short_id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_BY_USER_EXTERNAL_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
short_id,
|
||||
user_id,
|
||||
external_task_id,
|
||||
provider_api_format,
|
||||
model,
|
||||
prompt,
|
||||
size,
|
||||
status,
|
||||
progress_percent,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
,
|
||||
error_code,
|
||||
error_message,
|
||||
video_url
|
||||
FROM video_tasks
|
||||
WHERE user_id = $1 AND external_task_id = $2
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const LIST_ACTIVE_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
short_id,
|
||||
user_id,
|
||||
external_task_id,
|
||||
provider_api_format,
|
||||
model,
|
||||
prompt,
|
||||
size,
|
||||
status,
|
||||
progress_percent,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
|
||||
,
|
||||
error_code,
|
||||
error_message,
|
||||
video_url
|
||||
FROM video_tasks
|
||||
WHERE status = ANY($1)
|
||||
ORDER BY updated_at DESC
|
||||
LIMIT $2
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxVideoTaskReadRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl SqlxVideoTaskReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn find(
|
||||
&self,
|
||||
key: VideoTaskLookupKey<'_>,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
match key {
|
||||
VideoTaskLookupKey::Id(id) => self.find_by_id(id).await,
|
||||
VideoTaskLookupKey::ShortId(short_id) => self.find_by_short_id(short_id).await,
|
||||
VideoTaskLookupKey::UserExternal {
|
||||
user_id,
|
||||
external_task_id,
|
||||
} => self.find_by_user_external(user_id, external_task_id).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn find_by_id(&self, id: &str) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_ID_SQL)
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_video_task_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn find_by_short_id(
|
||||
&self,
|
||||
short_id: &str,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_SHORT_ID_SQL)
|
||||
.bind(short_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_video_task_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn find_by_user_external(
|
||||
&self,
|
||||
user_id: &str,
|
||||
external_task_id: &str,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_USER_EXTERNAL_SQL)
|
||||
.bind(user_id)
|
||||
.bind(external_task_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_video_task_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn list_active(&self, limit: usize) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let active_statuses = vec!["pending", "submitted", "queued", "processing"];
|
||||
let rows = sqlx::query(LIST_ACTIVE_SQL)
|
||||
.bind(active_statuses)
|
||||
.bind(i64::try_from(limit).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("invalid active task limit: {limit}"))
|
||||
})?)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
rows.iter().map(map_video_task_row).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl VideoTaskReadRepository for SqlxVideoTaskReadRepository {
|
||||
async fn find(
|
||||
&self,
|
||||
key: VideoTaskLookupKey<'_>,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
Self::find(self, key).await
|
||||
}
|
||||
|
||||
async fn list_active(&self, limit: usize) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
Self::list_active(self, limit).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_video_task_row(row: &sqlx::postgres::PgRow) -> Result<StoredVideoTask, DataLayerError> {
|
||||
let status = VideoTaskStatus::from_database(row.try_get::<String, _>("status")?.as_str())?;
|
||||
StoredVideoTask::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("short_id")?,
|
||||
row.try_get("user_id")?,
|
||||
row.try_get("external_task_id")?,
|
||||
row.try_get("provider_api_format")?,
|
||||
row.try_get("model")?,
|
||||
row.try_get("prompt")?,
|
||||
row.try_get("size")?,
|
||||
status,
|
||||
row.try_get("progress_percent")?,
|
||||
row.try_get("created_at_unix_secs")?,
|
||||
row.try_get("updated_at_unix_secs")?,
|
||||
row.try_get("error_code")?,
|
||||
row.try_get("error_message")?,
|
||||
row.try_get("video_url")?,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxVideoTaskReadRepository;
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
use crate::repository::video_tasks::{VideoTaskLookupKey, VideoTaskReadRepository};
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_constructs_from_lazy_pool() {
|
||||
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("factory should build");
|
||||
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let repository = SqlxVideoTaskReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_trait_delegates_to_sqlx_repository() {
|
||||
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("factory should build");
|
||||
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let repository = SqlxVideoTaskReadRepository::new(pool);
|
||||
let _ = VideoTaskReadRepository::find(&repository, VideoTaskLookupKey::Id("task-1")).await;
|
||||
}
|
||||
}
|
||||
277
crates/aether-data/src/repository/video_tasks/types.rs
Normal file
277
crates/aether-data/src/repository/video_tasks/types.rs
Normal file
@@ -0,0 +1,277 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub enum VideoTaskStatus {
|
||||
Pending,
|
||||
Submitted,
|
||||
Queued,
|
||||
Processing,
|
||||
Completed,
|
||||
Failed,
|
||||
Cancelled,
|
||||
Expired,
|
||||
Deleted,
|
||||
}
|
||||
|
||||
impl VideoTaskStatus {
|
||||
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"pending" => Ok(Self::Pending),
|
||||
"submitted" => Ok(Self::Submitted),
|
||||
"queued" => Ok(Self::Queued),
|
||||
"processing" => Ok(Self::Processing),
|
||||
"completed" => Ok(Self::Completed),
|
||||
"failed" => Ok(Self::Failed),
|
||||
"cancelled" => Ok(Self::Cancelled),
|
||||
"expired" => Ok(Self::Expired),
|
||||
"deleted" => Ok(Self::Deleted),
|
||||
other => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||
"unsupported video_tasks.status: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_active(self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::Pending | Self::Submitted | Self::Queued | Self::Processing
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredVideoTask {
|
||||
pub id: String,
|
||||
pub short_id: Option<String>,
|
||||
pub user_id: Option<String>,
|
||||
pub external_task_id: Option<String>,
|
||||
pub provider_api_format: Option<String>,
|
||||
pub model: Option<String>,
|
||||
pub prompt: Option<String>,
|
||||
pub size: Option<String>,
|
||||
pub status: VideoTaskStatus,
|
||||
pub progress_percent: u16,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
pub error_code: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub video_url: Option<String>,
|
||||
}
|
||||
|
||||
impl StoredVideoTask {
|
||||
pub fn new(
|
||||
id: String,
|
||||
short_id: Option<String>,
|
||||
user_id: Option<String>,
|
||||
external_task_id: Option<String>,
|
||||
provider_api_format: Option<String>,
|
||||
model: Option<String>,
|
||||
prompt: Option<String>,
|
||||
size: Option<String>,
|
||||
status: VideoTaskStatus,
|
||||
progress_percent: i32,
|
||||
created_at_unix_secs: i64,
|
||||
updated_at_unix_secs: i64,
|
||||
error_code: Option<String>,
|
||||
error_message: Option<String>,
|
||||
video_url: Option<String>,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
let progress_percent = u16::try_from(progress_percent).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid progress_percent: {progress_percent}"
|
||||
))
|
||||
})?;
|
||||
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid created_at_unix_secs: {created_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
let updated_at_unix_secs = u64::try_from(updated_at_unix_secs).map_err(|_| {
|
||||
crate::DataLayerError::UnexpectedValue(format!(
|
||||
"invalid updated_at_unix_secs: {updated_at_unix_secs}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(Self {
|
||||
id,
|
||||
short_id,
|
||||
user_id,
|
||||
external_task_id,
|
||||
provider_api_format,
|
||||
model,
|
||||
prompt,
|
||||
size,
|
||||
status,
|
||||
progress_percent,
|
||||
created_at_unix_secs,
|
||||
updated_at_unix_secs,
|
||||
error_code,
|
||||
error_message,
|
||||
video_url,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpsertVideoTask {
|
||||
pub id: String,
|
||||
pub short_id: Option<String>,
|
||||
pub user_id: Option<String>,
|
||||
pub external_task_id: Option<String>,
|
||||
pub provider_api_format: Option<String>,
|
||||
pub model: Option<String>,
|
||||
pub prompt: Option<String>,
|
||||
pub size: Option<String>,
|
||||
pub status: VideoTaskStatus,
|
||||
pub progress_percent: u16,
|
||||
pub created_at_unix_secs: u64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
pub error_code: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub video_url: Option<String>,
|
||||
}
|
||||
|
||||
impl UpsertVideoTask {
|
||||
pub fn into_stored(self) -> StoredVideoTask {
|
||||
StoredVideoTask {
|
||||
id: self.id,
|
||||
short_id: self.short_id,
|
||||
user_id: self.user_id,
|
||||
external_task_id: self.external_task_id,
|
||||
provider_api_format: self.provider_api_format,
|
||||
model: self.model,
|
||||
prompt: self.prompt,
|
||||
size: self.size,
|
||||
status: self.status,
|
||||
progress_percent: self.progress_percent,
|
||||
created_at_unix_secs: self.created_at_unix_secs,
|
||||
updated_at_unix_secs: self.updated_at_unix_secs,
|
||||
error_code: self.error_code,
|
||||
error_message: self.error_message,
|
||||
video_url: self.video_url,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum VideoTaskLookupKey<'a> {
|
||||
Id(&'a str),
|
||||
ShortId(&'a str),
|
||||
UserExternal {
|
||||
user_id: &'a str,
|
||||
external_task_id: &'a str,
|
||||
},
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait VideoTaskReadRepository: Send + Sync {
|
||||
async fn find(
|
||||
&self,
|
||||
key: VideoTaskLookupKey<'_>,
|
||||
) -> Result<Option<StoredVideoTask>, crate::DataLayerError>;
|
||||
|
||||
async fn list_active(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait VideoTaskWriteRepository: Send + Sync {
|
||||
async fn upsert(&self, task: UpsertVideoTask)
|
||||
-> Result<StoredVideoTask, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait VideoTaskRepository:
|
||||
VideoTaskReadRepository + VideoTaskWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> VideoTaskRepository for T where
|
||||
T: VideoTaskReadRepository + VideoTaskWriteRepository + Send + Sync
|
||||
{
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{StoredVideoTask, VideoTaskStatus};
|
||||
|
||||
#[test]
|
||||
fn parses_status_from_database_text() {
|
||||
assert_eq!(
|
||||
VideoTaskStatus::from_database("processing").expect("status should parse"),
|
||||
VideoTaskStatus::Processing
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_database_status() {
|
||||
assert!(VideoTaskStatus::from_database("mystery").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_numeric_fields() {
|
||||
assert!(StoredVideoTask::new(
|
||||
"task-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
VideoTaskStatus::Submitted,
|
||||
-1,
|
||||
1,
|
||||
1,
|
||||
None,
|
||||
None,
|
||||
None
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_updated_at_values() {
|
||||
assert!(StoredVideoTask::new(
|
||||
"task-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
VideoTaskStatus::Submitted,
|
||||
10,
|
||||
1,
|
||||
-1,
|
||||
None,
|
||||
None,
|
||||
None
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_negative_created_at_values() {
|
||||
assert!(StoredVideoTask::new(
|
||||
"task-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
VideoTaskStatus::Submitted,
|
||||
10,
|
||||
-1,
|
||||
1,
|
||||
None,
|
||||
None,
|
||||
None
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
}
|
||||
@@ -8,6 +8,8 @@ description = "Rust executor scaffold for Aether request execution"
|
||||
|
||||
[dependencies]
|
||||
aether-contracts.workspace = true
|
||||
aether-http.workspace = true
|
||||
aether-runtime.workspace = true
|
||||
async-stream.workspace = true
|
||||
axum = { version = "0.8" }
|
||||
base64.workspace = true
|
||||
@@ -22,6 +24,5 @@ serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
tracing.workspace = true
|
||||
webpki-roots.workspace = true
|
||||
|
||||
@@ -39,6 +39,12 @@ pub enum ExecutorServiceError {
|
||||
BodyEncode(serde_json::Error),
|
||||
#[error("failed to build HTTP client: {0}")]
|
||||
ClientBuild(reqwest::Error),
|
||||
#[error("failed to read executor request body: {0}")]
|
||||
RequestRead(String),
|
||||
#[error("executor request body is not valid JSON: {0}")]
|
||||
InvalidRequestJson(serde_json::Error),
|
||||
#[error("executor overloaded: gate {gate} saturated at {limit}")]
|
||||
Overloaded { gate: &'static str, limit: usize },
|
||||
#[error("failed to execute upstream request: {0}")]
|
||||
UpstreamRequest(reqwest::Error),
|
||||
#[error("hub relay request failed: {0}")]
|
||||
|
||||
@@ -4,6 +4,10 @@ use clap::Parser;
|
||||
use tracing::info;
|
||||
|
||||
use aether_executor::server;
|
||||
use aether_runtime::{
|
||||
init_service_runtime, DistributedConcurrencyGate, RedisDistributedConcurrencyConfig,
|
||||
ServiceRuntimeConfig,
|
||||
};
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "aether-executor", about = "Internal Rust executor for Aether")]
|
||||
@@ -20,28 +24,106 @@ struct Args {
|
||||
default_value = "/tmp/aether-executor.sock"
|
||||
)]
|
||||
unix_socket: PathBuf,
|
||||
|
||||
#[arg(long, env = "AETHER_EXECUTOR_MAX_IN_FLIGHT_REQUESTS")]
|
||||
max_in_flight_requests: Option<usize>,
|
||||
|
||||
#[arg(long, env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_LIMIT")]
|
||||
distributed_request_limit: Option<usize>,
|
||||
|
||||
#[arg(long, env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_REDIS_URL")]
|
||||
distributed_request_redis_url: Option<String>,
|
||||
|
||||
#[arg(long, env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_REDIS_KEY_PREFIX")]
|
||||
distributed_request_redis_key_prefix: Option<String>,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_LEASE_TTL_MS",
|
||||
default_value_t = 30_000
|
||||
)]
|
||||
distributed_request_lease_ttl_ms: u64,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_RENEW_INTERVAL_MS",
|
||||
default_value_t = 10_000
|
||||
)]
|
||||
distributed_request_renew_interval_ms: u64,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_COMMAND_TIMEOUT_MS",
|
||||
default_value_t = 1_000
|
||||
)]
|
||||
distributed_request_command_timeout_ms: u64,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "aether_executor=info".into()),
|
||||
)
|
||||
.init();
|
||||
init_service_runtime(ServiceRuntimeConfig::new(
|
||||
"aether-executor",
|
||||
"aether_executor=info",
|
||||
))?;
|
||||
|
||||
let args = Args::parse();
|
||||
let distributed_request_gate = match args.distributed_request_limit.filter(|limit| *limit > 0) {
|
||||
Some(limit) => {
|
||||
let redis_url = args
|
||||
.distributed_request_redis_url
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| {
|
||||
std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
"AETHER_EXECUTOR_DISTRIBUTED_REQUEST_REDIS_URL is required when distributed request limit is enabled",
|
||||
)
|
||||
})?;
|
||||
Some(DistributedConcurrencyGate::new_redis(
|
||||
"executor_requests_distributed",
|
||||
limit,
|
||||
RedisDistributedConcurrencyConfig {
|
||||
url: redis_url.to_string(),
|
||||
key_prefix: args
|
||||
.distributed_request_redis_key_prefix
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
lease_ttl_ms: args.distributed_request_lease_ttl_ms.max(1),
|
||||
renew_interval_ms: args.distributed_request_renew_interval_ms.max(1),
|
||||
command_timeout_ms: Some(args.distributed_request_command_timeout_ms.max(1)),
|
||||
},
|
||||
)?)
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
match args.transport.trim().to_ascii_lowercase().as_str() {
|
||||
"unix_socket" | "unix" | "uds" => {
|
||||
info!(socket = %args.unix_socket.display(), "aether-executor started");
|
||||
server::serve_unix(&args.unix_socket).await?;
|
||||
server::serve_unix(
|
||||
&args.unix_socket,
|
||||
args.max_in_flight_requests,
|
||||
distributed_request_gate.clone(),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
"tcp" => {
|
||||
info!(bind = %args.bind, "aether-executor started");
|
||||
server::serve_tcp(&args.bind).await?;
|
||||
info!(
|
||||
bind = %args.bind,
|
||||
max_in_flight_requests = args.max_in_flight_requests.unwrap_or_default(),
|
||||
distributed_request_limit = args.distributed_request_limit.unwrap_or_default(),
|
||||
"aether-executor started"
|
||||
);
|
||||
server::serve_tcp(
|
||||
&args.bind,
|
||||
args.max_in_flight_requests,
|
||||
distributed_request_gate,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
other => {
|
||||
return Err(format!("unsupported executor transport: {other}").into());
|
||||
|
||||
@@ -1,13 +1,20 @@
|
||||
use std::convert::Infallible;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_contracts::{
|
||||
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult,
|
||||
ExecutionTelemetry, StreamFrame, StreamFramePayload, StreamFrameType,
|
||||
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionTelemetry,
|
||||
StreamFrame, StreamFramePayload, StreamFrameType,
|
||||
};
|
||||
use aether_runtime::{
|
||||
maybe_hold_axum_response_permit, prometheus_response, service_up_sample, AdmissionPermit,
|
||||
ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot, DistributedConcurrencyError,
|
||||
DistributedConcurrencyGate, DistributedConcurrencySnapshot, MetricKind, MetricLabel,
|
||||
MetricSample,
|
||||
};
|
||||
use async_stream::stream;
|
||||
use axum::body::Body;
|
||||
use axum::extract::State;
|
||||
use axum::body::{to_bytes, Body};
|
||||
use axum::extract::{Request, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::{get, post};
|
||||
@@ -22,25 +29,132 @@ use crate::{encode_frame, ExecutorServiceError, SyncExecutor};
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct AppState {
|
||||
executor: SyncExecutor,
|
||||
request_gate: Option<Arc<ConcurrencyGate>>,
|
||||
distributed_request_gate: Option<Arc<DistributedConcurrencyGate>>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
fn with_request_concurrency_limit(limit: Option<usize>) -> Self {
|
||||
Self {
|
||||
executor: SyncExecutor::new(),
|
||||
request_gate: limit
|
||||
.filter(|limit| *limit > 0)
|
||||
.map(|limit| Arc::new(ConcurrencyGate::new("executor_requests", limit))),
|
||||
distributed_request_gate: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn with_distributed_request_gate(mut self, gate: DistributedConcurrencyGate) -> Self {
|
||||
self.distributed_request_gate = Some(Arc::new(gate));
|
||||
self
|
||||
}
|
||||
|
||||
fn request_concurrency_snapshot(&self) -> Option<ConcurrencySnapshot> {
|
||||
self.request_gate.as_ref().map(|gate| gate.snapshot())
|
||||
}
|
||||
|
||||
async fn distributed_request_concurrency_snapshot(
|
||||
&self,
|
||||
) -> Result<Option<DistributedConcurrencySnapshot>, DistributedConcurrencyError> {
|
||||
match self.distributed_request_gate.as_ref() {
|
||||
Some(gate) => gate.snapshot().await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn metric_samples(&self) -> Vec<MetricSample> {
|
||||
let mut samples = vec![service_up_sample("aether-executor")];
|
||||
if let Some(snapshot) = self.request_concurrency_snapshot() {
|
||||
samples.extend(snapshot.to_metric_samples("executor_requests"));
|
||||
}
|
||||
if let Some(gate) = self.distributed_request_gate.as_ref() {
|
||||
match gate.snapshot().await {
|
||||
Ok(snapshot) => {
|
||||
samples.extend(snapshot.to_metric_samples("executor_requests_distributed"));
|
||||
}
|
||||
Err(_) => samples.push(
|
||||
MetricSample::new(
|
||||
"concurrency_unavailable",
|
||||
"Whether the distributed concurrency gate is currently unavailable.",
|
||||
MetricKind::Gauge,
|
||||
1,
|
||||
)
|
||||
.with_labels(vec![MetricLabel::new(
|
||||
"gate",
|
||||
"executor_requests_distributed",
|
||||
)]),
|
||||
),
|
||||
}
|
||||
}
|
||||
samples
|
||||
}
|
||||
|
||||
async fn try_acquire_request_permit(
|
||||
&self,
|
||||
) -> Result<Option<AdmissionPermit>, RequestAdmissionError> {
|
||||
let local = self
|
||||
.request_gate
|
||||
.as_ref()
|
||||
.map(|gate| gate.try_acquire())
|
||||
.transpose()
|
||||
.map_err(RequestAdmissionError::Local)?;
|
||||
let distributed = match self.distributed_request_gate.as_ref() {
|
||||
Some(gate) => Some(
|
||||
gate.try_acquire()
|
||||
.await
|
||||
.map_err(RequestAdmissionError::Distributed)?,
|
||||
),
|
||||
None => None,
|
||||
};
|
||||
Ok(AdmissionPermit::from_parts(local, distributed))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_router() -> Router {
|
||||
Router::new()
|
||||
.route("/health", get(health))
|
||||
.route("/v1/execute/sync", post(execute_sync))
|
||||
.route("/v1/execute/stream", post(execute_stream))
|
||||
.with_state(AppState {
|
||||
executor: SyncExecutor::new(),
|
||||
})
|
||||
build_router_with_request_concurrency_limit(None)
|
||||
}
|
||||
|
||||
pub async fn serve_tcp(bind: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
pub fn build_router_with_request_concurrency_limit(limit: Option<usize>) -> Router {
|
||||
build_router_with_request_gates(limit, None)
|
||||
}
|
||||
|
||||
pub fn build_router_with_request_gates(
|
||||
limit: Option<usize>,
|
||||
distributed_gate: Option<DistributedConcurrencyGate>,
|
||||
) -> Router {
|
||||
let state = match distributed_gate {
|
||||
Some(gate) => {
|
||||
AppState::with_request_concurrency_limit(limit).with_distributed_request_gate(gate)
|
||||
}
|
||||
None => AppState::with_request_concurrency_limit(limit),
|
||||
};
|
||||
Router::new()
|
||||
.route("/health", get(health))
|
||||
.route("/metrics", get(metrics))
|
||||
.route("/v1/execute/sync", post(execute_sync))
|
||||
.route("/v1/execute/stream", post(execute_stream))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
pub async fn serve_tcp(
|
||||
bind: &str,
|
||||
max_in_flight_requests: Option<usize>,
|
||||
distributed_request_gate: Option<DistributedConcurrencyGate>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let listener = tokio::net::TcpListener::bind(bind).await?;
|
||||
axum::serve(listener, build_router()).await?;
|
||||
axum::serve(
|
||||
listener,
|
||||
build_router_with_request_gates(max_in_flight_requests, distributed_request_gate),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn serve_unix(socket_path: &Path) -> Result<(), Box<dyn std::error::Error>> {
|
||||
pub async fn serve_unix(
|
||||
socket_path: &Path,
|
||||
max_in_flight_requests: Option<usize>,
|
||||
distributed_request_gate: Option<DistributedConcurrencyGate>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
if let Some(parent) = socket_path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
@@ -49,30 +163,69 @@ pub async fn serve_unix(socket_path: &Path) -> Result<(), Box<dyn std::error::Er
|
||||
}
|
||||
|
||||
let listener = tokio::net::UnixListener::bind(socket_path)?;
|
||||
axum::serve(listener, build_router()).await?;
|
||||
axum::serve(
|
||||
listener,
|
||||
build_router_with_request_gates(max_in_flight_requests, distributed_request_gate),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn health() -> impl IntoResponse {
|
||||
Json(json!({"status": "ok"}))
|
||||
async fn health(State(state): State<AppState>) -> impl IntoResponse {
|
||||
let request_concurrency = state.request_concurrency_snapshot().map(|snapshot| {
|
||||
json!({
|
||||
"limit": snapshot.limit,
|
||||
"in_flight": snapshot.in_flight,
|
||||
"available_permits": snapshot.available_permits,
|
||||
"high_watermark": snapshot.high_watermark,
|
||||
"rejected": snapshot.rejected,
|
||||
})
|
||||
});
|
||||
let distributed_request_concurrency = state
|
||||
.distributed_request_concurrency_snapshot()
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.map(|snapshot| {
|
||||
json!({
|
||||
"limit": snapshot.limit,
|
||||
"in_flight": snapshot.in_flight,
|
||||
"available_permits": snapshot.available_permits,
|
||||
"high_watermark": snapshot.high_watermark,
|
||||
"rejected": snapshot.rejected,
|
||||
})
|
||||
});
|
||||
Json(json!({
|
||||
"status": "ok",
|
||||
"component": "aether-executor",
|
||||
"request_concurrency": request_concurrency,
|
||||
"distributed_request_concurrency": distributed_request_concurrency,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn metrics(State(state): State<AppState>) -> Response {
|
||||
prometheus_response(&state.metric_samples().await)
|
||||
}
|
||||
|
||||
async fn execute_sync(
|
||||
State(state): State<AppState>,
|
||||
Json(plan): Json<ExecutionPlan>,
|
||||
) -> Result<Json<ExecutionResult>, AppError> {
|
||||
state
|
||||
.executor
|
||||
.execute_sync(plan)
|
||||
.await
|
||||
.map(Json)
|
||||
.map_err(AppError)
|
||||
request: Request,
|
||||
) -> Result<Response, AppError> {
|
||||
let request_permit = acquire_request_permit(&state).await?;
|
||||
let plan = parse_request_json::<ExecutionPlan>(request).await?;
|
||||
let result = state.executor.execute_sync(plan).await.map_err(AppError)?;
|
||||
Ok(maybe_hold_axum_response_permit(
|
||||
Json(result).into_response(),
|
||||
request_permit,
|
||||
))
|
||||
}
|
||||
|
||||
async fn execute_stream(
|
||||
State(state): State<AppState>,
|
||||
Json(plan): Json<ExecutionPlan>,
|
||||
request: Request,
|
||||
) -> Result<Response, AppError> {
|
||||
let request_permit = acquire_request_permit(&state).await?;
|
||||
let plan = parse_request_json::<ExecutionPlan>(request).await?;
|
||||
let execution = state
|
||||
.executor
|
||||
.execute_stream(plan)
|
||||
@@ -149,7 +302,61 @@ async fn execute_stream(
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
axum::http::HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
Ok(response)
|
||||
Ok(maybe_hold_axum_response_permit(response, request_permit))
|
||||
}
|
||||
|
||||
async fn acquire_request_permit(state: &AppState) -> Result<Option<AdmissionPermit>, AppError> {
|
||||
match state.try_acquire_request_permit().await {
|
||||
Ok(permit) => Ok(permit),
|
||||
Err(RequestAdmissionError::Local(ConcurrencyError::Saturated { gate, limit }))
|
||||
| Err(RequestAdmissionError::Distributed(DistributedConcurrencyError::Saturated {
|
||||
gate,
|
||||
limit,
|
||||
}))
|
||||
| Err(RequestAdmissionError::Distributed(DistributedConcurrencyError::Unavailable {
|
||||
gate,
|
||||
limit,
|
||||
..
|
||||
})) => Err(AppError(ExecutorServiceError::Overloaded { gate, limit })),
|
||||
Err(RequestAdmissionError::Local(ConcurrencyError::Closed { gate })) => {
|
||||
Err(AppError(ExecutorServiceError::RequestRead(format!(
|
||||
"executor request concurrency gate {gate} is closed"
|
||||
))))
|
||||
}
|
||||
Err(RequestAdmissionError::Distributed(
|
||||
DistributedConcurrencyError::InvalidConfiguration(message),
|
||||
)) => Err(AppError(ExecutorServiceError::RequestRead(message))),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum RequestAdmissionError {
|
||||
Local(ConcurrencyError),
|
||||
Distributed(DistributedConcurrencyError),
|
||||
}
|
||||
|
||||
async fn parse_request_json<T>(request: Request) -> Result<T, AppError>
|
||||
where
|
||||
T: serde::de::DeserializeOwned,
|
||||
{
|
||||
let body = to_bytes(request.into_body(), usize::MAX)
|
||||
.await
|
||||
.map_err(|err| AppError(ExecutorServiceError::RequestRead(err.to_string())))?;
|
||||
serde_json::from_slice(&body)
|
||||
.map_err(|err| AppError(ExecutorServiceError::InvalidRequestJson(err)))
|
||||
}
|
||||
|
||||
fn build_overloaded_response(message: &str) -> Response {
|
||||
(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"type": "overloaded",
|
||||
"message": message,
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
@@ -158,6 +365,12 @@ struct AppError(ExecutorServiceError);
|
||||
impl IntoResponse for AppError {
|
||||
fn into_response(self) -> Response {
|
||||
let status_code = match self.0 {
|
||||
ExecutorServiceError::RequestRead(_) | ExecutorServiceError::InvalidRequestJson(_) => {
|
||||
StatusCode::BAD_REQUEST
|
||||
}
|
||||
ExecutorServiceError::Overloaded { .. } => {
|
||||
return build_overloaded_response(&self.0.to_string());
|
||||
}
|
||||
ExecutorServiceError::StreamUnsupported
|
||||
| ExecutorServiceError::RequestBodyRequired
|
||||
| ExecutorServiceError::BodyDecode(_)
|
||||
@@ -185,3 +398,227 @@ impl IntoResponse for AppError {
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{build_router_with_request_concurrency_limit, build_router_with_request_gates};
|
||||
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use std::convert::Infallible;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
async fn start_server(app: Router) -> (String, tokio::task::JoinHandle<()>) {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let addr = listener.local_addr().expect("local addr should resolve");
|
||||
let handle = tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.expect("server should run");
|
||||
});
|
||||
(format!("http://{addr}"), handle)
|
||||
}
|
||||
|
||||
fn stream_plan(url: String) -> ExecutionPlan {
|
||||
ExecutionPlan {
|
||||
request_id: "req-1".into(),
|
||||
candidate_id: Some("cand-1".into()),
|
||||
provider_name: Some("openai".into()),
|
||||
provider_id: "prov-1".into(),
|
||||
endpoint_id: "ep-1".into(),
|
||||
key_id: "key-1".into(),
|
||||
method: "GET".into(),
|
||||
url,
|
||||
headers: std::collections::BTreeMap::new(),
|
||||
content_type: None,
|
||||
content_encoding: None,
|
||||
body: RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
stream: true,
|
||||
client_api_format: "openai:chat".into(),
|
||||
provider_api_format: "openai:chat".into(),
|
||||
model_name: Some("gpt-4.1".into()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: Some(ExecutionTimeouts {
|
||||
connect_ms: Some(5_000),
|
||||
total_ms: Some(30_000),
|
||||
..ExecutionTimeouts::default()
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn executor_rejects_second_in_flight_stream_request_with_overload() {
|
||||
let upstream_hits = Arc::new(AtomicUsize::new(0));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/slow",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
upstream_hits.fetch_add(1, Ordering::SeqCst);
|
||||
let stream = async_stream::stream! {
|
||||
yield Ok::<_, Infallible>(Bytes::from_static(b"chunk-1"));
|
||||
futures_util::future::pending::<()>().await;
|
||||
};
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(stream))
|
||||
.expect("response should build")
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let executor = build_router_with_request_concurrency_limit(Some(1));
|
||||
let (executor_url, executor_handle) = start_server(executor).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let first_response = client
|
||||
.post(format!("{executor_url}/v1/execute/stream"))
|
||||
.json(&stream_plan(format!("{upstream_url}/slow")))
|
||||
.send()
|
||||
.await
|
||||
.expect("first request should succeed");
|
||||
|
||||
for _ in 0..50 {
|
||||
if upstream_hits.load(Ordering::SeqCst) == 1 {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let second_response = client
|
||||
.post(format!("{executor_url}/v1/execute/stream"))
|
||||
.json(&stream_plan(format!("{upstream_url}/slow")))
|
||||
.send()
|
||||
.await
|
||||
.expect("second request should complete");
|
||||
|
||||
assert_eq!(second_response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
second_response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("json body should decode")["error"]["type"],
|
||||
"overloaded"
|
||||
);
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
drop(first_response);
|
||||
executor_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn executor_rejects_second_in_flight_stream_request_with_distributed_overload() {
|
||||
let upstream_hits = Arc::new(AtomicUsize::new(0));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/slow",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
upstream_hits.fetch_add(1, Ordering::SeqCst);
|
||||
let stream = async_stream::stream! {
|
||||
yield Ok::<_, Infallible>(Bytes::from_static(b"chunk-1"));
|
||||
futures_util::future::pending::<()>().await;
|
||||
};
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(stream))
|
||||
.expect("response should build")
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let distributed_gate = aether_runtime::DistributedConcurrencyGate::new_in_memory(
|
||||
"executor_requests_distributed",
|
||||
1,
|
||||
);
|
||||
let executor_a = build_router_with_request_gates(None, Some(distributed_gate.clone()));
|
||||
let executor_b = build_router_with_request_gates(None, Some(distributed_gate));
|
||||
let (executor_a_url, executor_a_handle) = start_server(executor_a).await;
|
||||
let (executor_b_url, executor_b_handle) = start_server(executor_b).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let first_response = client
|
||||
.post(format!("{executor_a_url}/v1/execute/stream"))
|
||||
.json(&stream_plan(format!("{upstream_url}/slow")))
|
||||
.send()
|
||||
.await
|
||||
.expect("first request should succeed");
|
||||
|
||||
for _ in 0..50 {
|
||||
if upstream_hits.load(Ordering::SeqCst) == 1 {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let second_response = client
|
||||
.post(format!("{executor_b_url}/v1/execute/stream"))
|
||||
.json(&stream_plan(format!("{upstream_url}/slow")))
|
||||
.send()
|
||||
.await
|
||||
.expect("second request should complete");
|
||||
|
||||
assert_eq!(second_response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
second_response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("json body should decode")["error"]["type"],
|
||||
"overloaded"
|
||||
);
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
drop(first_response);
|
||||
executor_a_handle.abort();
|
||||
executor_b_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn executor_exposes_request_concurrency_metrics() {
|
||||
let executor = build_router_with_request_gates(
|
||||
Some(4),
|
||||
Some(aether_runtime::DistributedConcurrencyGate::new_in_memory(
|
||||
"executor_requests_distributed",
|
||||
6,
|
||||
)),
|
||||
);
|
||||
let (executor_url, executor_handle) = start_server(executor).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{executor_url}/metrics"))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("text/plain; version=0.0.4; charset=utf-8")
|
||||
);
|
||||
let body = response.text().await.expect("body should read");
|
||||
assert!(body.contains("service_up{service=\"aether-executor\"} 1"));
|
||||
assert!(body.contains("concurrency_available_permits{gate=\"executor_requests\"} 4"));
|
||||
assert!(body
|
||||
.contains("concurrency_available_permits{gate=\"executor_requests_distributed\"} 6"));
|
||||
|
||||
executor_handle.abort();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ use std::time::{Duration, Instant};
|
||||
use aether_contracts::{
|
||||
ExecutionPlan, ExecutionResult, ExecutionTelemetry, ProxySnapshot, ResponseBody,
|
||||
};
|
||||
use aether_http::{apply_http_client_config, HttpClientConfig};
|
||||
use base64::Engine as _;
|
||||
use flate2::write::GzEncoder;
|
||||
use flate2::Compression;
|
||||
@@ -256,10 +257,14 @@ fn gzip_bytes(body_bytes: &[u8]) -> Result<Vec<u8>, ExecutorServiceError> {
|
||||
fn build_relay_client(
|
||||
timeouts: Option<&aether_contracts::ExecutionTimeouts>,
|
||||
) -> Result<reqwest::Client, ExecutorServiceError> {
|
||||
let mut builder = reqwest::Client::builder();
|
||||
if let Some(connect_ms) = timeouts.and_then(|timeouts| timeouts.connect_ms) {
|
||||
builder = builder.connect_timeout(Duration::from_millis(connect_ms));
|
||||
}
|
||||
let builder = apply_http_client_config(
|
||||
reqwest::Client::builder(),
|
||||
&HttpClientConfig {
|
||||
connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms),
|
||||
use_rustls_tls: false,
|
||||
..HttpClientConfig::default()
|
||||
},
|
||||
);
|
||||
builder.build().map_err(ExecutorServiceError::ClientBuild)
|
||||
}
|
||||
|
||||
@@ -339,10 +344,13 @@ fn build_client(
|
||||
proxy: Option<&ProxySnapshot>,
|
||||
tls_profile: Option<&str>,
|
||||
) -> Result<reqwest::Client, ExecutorServiceError> {
|
||||
let mut builder = reqwest::Client::builder().use_rustls_tls();
|
||||
if let Some(connect_ms) = timeouts.and_then(|timeouts| timeouts.connect_ms) {
|
||||
builder = builder.connect_timeout(Duration::from_millis(connect_ms));
|
||||
}
|
||||
let mut builder = apply_http_client_config(
|
||||
reqwest::Client::builder(),
|
||||
&HttpClientConfig {
|
||||
connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms),
|
||||
..HttpClientConfig::default()
|
||||
},
|
||||
);
|
||||
builder = apply_tls_profile(builder, tls_profile);
|
||||
if let Some(proxy_url) = resolve_proxy_url(proxy)? {
|
||||
let proxy = reqwest::Proxy::all(&proxy_url).map_err(ExecutorServiceError::InvalidProxy)?;
|
||||
|
||||
@@ -7,7 +7,11 @@ repository.workspace = true
|
||||
description = "Rust ingress gateway for Aether phase 3a transparent proxy"
|
||||
|
||||
[dependencies]
|
||||
aether-cache.workspace = true
|
||||
aether-contracts.workspace = true
|
||||
aether-data.workspace = true
|
||||
aether-http.workspace = true
|
||||
aether-runtime.workspace = true
|
||||
async-stream.workspace = true
|
||||
axum = { version = "0.8" }
|
||||
base64.workspace = true
|
||||
@@ -21,7 +25,6 @@ serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
tokio-util.workspace = true
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
tracing.workspace = true
|
||||
url.workspace = true
|
||||
uuid.workspace = true
|
||||
|
||||
205
crates/aether-gateway/src/audit/http.rs
Normal file
205
crates/aether-gateway/src/audit/http.rs
Normal file
@@ -0,0 +1,205 @@
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use axum::extract::{Path, Query, State};
|
||||
use axum::response::IntoResponse;
|
||||
use axum::Json;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
|
||||
use crate::gateway::{AppState, GatewayError};
|
||||
|
||||
const DEFAULT_RECENT_LIMIT: usize = 20;
|
||||
const MAX_RECENT_LIMIT: usize = 200;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct ListRecentShadowResultsQuery {
|
||||
pub(crate) limit: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct ShadowResultStatusCounts {
|
||||
pub(crate) pending: usize,
|
||||
pub(crate) r#match: usize,
|
||||
pub(crate) mismatch: usize,
|
||||
pub(crate) error: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(crate) struct ListRecentShadowResultsResponse {
|
||||
pub(crate) items: Vec<aether_data::repository::shadow_results::StoredShadowResult>,
|
||||
pub(crate) limit_applied: usize,
|
||||
pub(crate) counts: ShadowResultStatusCounts,
|
||||
}
|
||||
|
||||
pub(crate) async fn list_recent_shadow_results(
|
||||
State(state): State<AppState>,
|
||||
Query(query): Query<ListRecentShadowResultsQuery>,
|
||||
) -> Result<Json<ListRecentShadowResultsResponse>, GatewayError> {
|
||||
let limit = query
|
||||
.limit
|
||||
.unwrap_or(DEFAULT_RECENT_LIMIT)
|
||||
.clamp(1, MAX_RECENT_LIMIT);
|
||||
let items = state.list_recent_shadow_results(limit).await?;
|
||||
|
||||
let mut counts = ShadowResultStatusCounts {
|
||||
pending: 0,
|
||||
r#match: 0,
|
||||
mismatch: 0,
|
||||
error: 0,
|
||||
};
|
||||
for item in &items {
|
||||
match item.match_status {
|
||||
aether_data::repository::shadow_results::ShadowResultMatchStatus::Pending => {
|
||||
counts.pending += 1
|
||||
}
|
||||
aether_data::repository::shadow_results::ShadowResultMatchStatus::Match => {
|
||||
counts.r#match += 1
|
||||
}
|
||||
aether_data::repository::shadow_results::ShadowResultMatchStatus::Mismatch => {
|
||||
counts.mismatch += 1
|
||||
}
|
||||
aether_data::repository::shadow_results::ShadowResultMatchStatus::Error => {
|
||||
counts.error += 1
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Json(ListRecentShadowResultsResponse {
|
||||
items,
|
||||
limit_applied: limit,
|
||||
counts,
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct GetRequestCandidateTraceQuery {
|
||||
pub(crate) attempted_only: Option<bool>,
|
||||
}
|
||||
|
||||
pub(crate) async fn get_request_candidate_trace(
|
||||
State(state): State<AppState>,
|
||||
Path(request_id): Path<String>,
|
||||
Query(query): Query<GetRequestCandidateTraceQuery>,
|
||||
) -> Result<Json<crate::gateway::data::RequestCandidateTrace>, axum::response::Response> {
|
||||
let attempted_only = query.attempted_only.unwrap_or(false);
|
||||
let trace = state
|
||||
.read_request_candidate_trace(&request_id, attempted_only)
|
||||
.await
|
||||
.map_err(IntoResponse::into_response)?;
|
||||
|
||||
match trace {
|
||||
Some(trace) => Ok(Json(trace)),
|
||||
None => Err((
|
||||
axum::http::StatusCode::NOT_FOUND,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": "Request not found",
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn get_decision_trace(
|
||||
State(state): State<AppState>,
|
||||
Path(request_id): Path<String>,
|
||||
Query(query): Query<GetRequestCandidateTraceQuery>,
|
||||
) -> Result<Json<crate::gateway::data::DecisionTrace>, axum::response::Response> {
|
||||
let attempted_only = query.attempted_only.unwrap_or(false);
|
||||
let trace = state
|
||||
.read_decision_trace(&request_id, attempted_only)
|
||||
.await
|
||||
.map_err(IntoResponse::into_response)?;
|
||||
|
||||
match trace {
|
||||
Some(trace) => Ok(Json(trace)),
|
||||
None => Err((
|
||||
axum::http::StatusCode::NOT_FOUND,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": "Decision trace not found",
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn get_request_usage_audit(
|
||||
State(state): State<AppState>,
|
||||
Path(request_id): Path<String>,
|
||||
) -> Result<Json<crate::gateway::data::RequestUsageAudit>, axum::response::Response> {
|
||||
let usage = state
|
||||
.read_request_usage_audit(&request_id)
|
||||
.await
|
||||
.map_err(IntoResponse::into_response)?;
|
||||
|
||||
match usage {
|
||||
Some(usage) => Ok(Json(usage)),
|
||||
None => Err((
|
||||
axum::http::StatusCode::NOT_FOUND,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": "Request usage not found",
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn get_request_audit_bundle(
|
||||
State(state): State<AppState>,
|
||||
Path(request_id): Path<String>,
|
||||
Query(query): Query<GetRequestCandidateTraceQuery>,
|
||||
) -> Result<Json<crate::gateway::data::RequestAuditBundle>, axum::response::Response> {
|
||||
let attempted_only = query.attempted_only.unwrap_or(false);
|
||||
let bundle = state
|
||||
.read_request_audit_bundle(&request_id, attempted_only, current_unix_secs())
|
||||
.await
|
||||
.map_err(IntoResponse::into_response)?;
|
||||
|
||||
match bundle {
|
||||
Some(bundle) => Ok(Json(bundle)),
|
||||
None => Err((
|
||||
axum::http::StatusCode::NOT_FOUND,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": "Request audit bundle not found",
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn get_auth_api_key_snapshot(
|
||||
State(state): State<AppState>,
|
||||
Path((user_id, api_key_id)): Path<(String, String)>,
|
||||
) -> Result<Json<crate::gateway::data::StoredGatewayAuthApiKeySnapshot>, axum::response::Response> {
|
||||
let snapshot = state
|
||||
.read_auth_api_key_snapshot(&user_id, &api_key_id, current_unix_secs())
|
||||
.await
|
||||
.map_err(IntoResponse::into_response)?;
|
||||
|
||||
match snapshot {
|
||||
Some(snapshot) => Ok(Json(snapshot)),
|
||||
None => Err((
|
||||
axum::http::StatusCode::NOT_FOUND,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": "Auth snapshot not found",
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response()),
|
||||
}
|
||||
}
|
||||
|
||||
fn current_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
10
crates/aether-gateway/src/audit/mod.rs
Normal file
10
crates/aether-gateway/src/audit/mod.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
mod http;
|
||||
mod shadow;
|
||||
|
||||
pub(crate) use http::get_auth_api_key_snapshot;
|
||||
pub(crate) use http::get_decision_trace;
|
||||
pub(crate) use http::get_request_audit_bundle;
|
||||
pub(crate) use http::get_request_candidate_trace;
|
||||
pub(crate) use http::get_request_usage_audit;
|
||||
pub(crate) use http::list_recent_shadow_results;
|
||||
pub(crate) use shadow::record_shadow_result_non_blocking;
|
||||
271
crates/aether-gateway/src/audit/shadow.rs
Normal file
271
crates/aether-gateway/src/audit/shadow.rs
Normal file
@@ -0,0 +1,271 @@
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_data::repository::shadow_results::{RecordShadowResultSample, ShadowResultSampleOrigin};
|
||||
use axum::body::Body;
|
||||
use axum::http::header::CONTENT_TYPE;
|
||||
use axum::http::Response;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::gateway::constants::{
|
||||
CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER, EXECUTION_PATH_CONTROL_EXECUTE_STREAM,
|
||||
EXECUTION_PATH_CONTROL_EXECUTE_SYNC,
|
||||
};
|
||||
use crate::gateway::{AppState, GatewayControlDecision};
|
||||
|
||||
pub(crate) fn record_shadow_result_non_blocking(
|
||||
state: AppState,
|
||||
trace_id: &str,
|
||||
method: &http::Method,
|
||||
path_and_query: &str,
|
||||
control_decision: Option<&GatewayControlDecision>,
|
||||
execution_path: &'static str,
|
||||
response: &Response<Body>,
|
||||
) {
|
||||
let Some(decision) =
|
||||
control_decision.filter(|decision| decision.route_class.as_deref() == Some("ai_public"))
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if !state.has_shadow_result_data_writer() {
|
||||
return;
|
||||
}
|
||||
|
||||
let route_family = decision.route_family.clone();
|
||||
let route_kind = decision.route_kind.clone();
|
||||
let status_code = response.status().as_u16();
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get(CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let candidate_id = response
|
||||
.headers()
|
||||
.get(CONTROL_CANDIDATE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let request_id = response
|
||||
.headers()
|
||||
.get(CONTROL_REQUEST_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
let sample = RecordShadowResultSample {
|
||||
trace_id: trace_id.to_string(),
|
||||
request_fingerprint: build_request_fingerprint(
|
||||
method,
|
||||
path_and_query,
|
||||
route_family.as_deref(),
|
||||
route_kind.as_deref(),
|
||||
),
|
||||
request_id,
|
||||
route_family,
|
||||
route_kind,
|
||||
candidate_id,
|
||||
origin: sample_origin_for_execution_path(execution_path),
|
||||
result_digest: build_result_digest(status_code, &content_type),
|
||||
status_code: Some(status_code),
|
||||
error_message: (status_code >= 400)
|
||||
.then(|| format!("gateway response status {status_code}")),
|
||||
recorded_at_unix_secs: now_unix_secs,
|
||||
};
|
||||
let trace_id = trace_id.to_string();
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(err) = state.record_shadow_result_sample(sample).await {
|
||||
warn!(trace_id = %trace_id, error = ?err, "gateway failed to record shadow result");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn build_request_fingerprint(
|
||||
method: &http::Method,
|
||||
path_and_query: &str,
|
||||
route_family: Option<&str>,
|
||||
route_kind: Option<&str>,
|
||||
) -> String {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
method.as_str().hash(&mut hasher);
|
||||
path_and_query.hash(&mut hasher);
|
||||
route_family.unwrap_or_default().hash(&mut hasher);
|
||||
route_kind.unwrap_or_default().hash(&mut hasher);
|
||||
format!("{:x}", hasher.finish())
|
||||
}
|
||||
|
||||
fn build_result_digest(status_code: u16, content_type: &str) -> String {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
status_code.hash(&mut hasher);
|
||||
content_type.hash(&mut hasher);
|
||||
format!("{:x}", hasher.finish())
|
||||
}
|
||||
|
||||
fn sample_origin_for_execution_path(execution_path: &str) -> ShadowResultSampleOrigin {
|
||||
match execution_path {
|
||||
EXECUTION_PATH_CONTROL_EXECUTE_SYNC | EXECUTION_PATH_CONTROL_EXECUTE_STREAM => {
|
||||
ShadowResultSampleOrigin::Python
|
||||
}
|
||||
_ => ShadowResultSampleOrigin::Rust,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_data::repository::shadow_results::{
|
||||
InMemoryShadowResultRepository, ShadowResultMatchStatus, ShadowResultReadRepository,
|
||||
};
|
||||
use axum::body::Body;
|
||||
use axum::http::header::CONTENT_TYPE;
|
||||
use axum::http::{Method, Response, StatusCode};
|
||||
|
||||
use super::record_shadow_result_non_blocking;
|
||||
use crate::gateway::constants::{
|
||||
CONTROL_REQUEST_ID_HEADER, EXECUTION_PATH_CONTROL_EXECUTE_SYNC,
|
||||
EXECUTION_PATH_EXECUTOR_SYNC,
|
||||
};
|
||||
use crate::gateway::{AppState, GatewayControlDecision};
|
||||
|
||||
fn sample_decision() -> GatewayControlDecision {
|
||||
GatewayControlDecision {
|
||||
public_path: "/v1/chat/completions".to_string(),
|
||||
public_query_string: Some("stream=true".to_string()),
|
||||
route_class: Some("ai_public".to_string()),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
auth_endpoint_signature: Some("openai:chat".to_string()),
|
||||
executor_candidate: true,
|
||||
auth_context: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn records_shadow_result_for_ai_public_response() {
|
||||
let repository = Arc::new(InMemoryShadowResultRepository::default());
|
||||
let state = AppState::new_with_executor(
|
||||
"http://127.0.0.1:18084",
|
||||
Some("http://127.0.0.1:18085".to_string()),
|
||||
Some("http://127.0.0.1:18086".to_string()),
|
||||
)
|
||||
.expect("app state should build")
|
||||
.with_shadow_result_data_writer_for_tests(repository.clone());
|
||||
let response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.header(CONTENT_TYPE, "application/json")
|
||||
.header(CONTROL_REQUEST_ID_HEADER, "req-shadow-123")
|
||||
.body(Body::from("{}"))
|
||||
.expect("response should build");
|
||||
|
||||
record_shadow_result_non_blocking(
|
||||
state,
|
||||
"trace-shadow-123",
|
||||
&Method::POST,
|
||||
"/v1/chat/completions?stream=true",
|
||||
Some(&sample_decision()),
|
||||
EXECUTION_PATH_EXECUTOR_SYNC,
|
||||
&response,
|
||||
);
|
||||
|
||||
for _ in 0..30 {
|
||||
if repository
|
||||
.list_recent(1)
|
||||
.await
|
||||
.map(|rows| !rows.is_empty())
|
||||
.unwrap_or(false)
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
|
||||
let stored = repository
|
||||
.list_recent(1)
|
||||
.await
|
||||
.expect("list should succeed")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("stored result should exist");
|
||||
assert_eq!(stored.trace_id, "trace-shadow-123");
|
||||
assert_eq!(stored.request_id.as_deref(), Some("req-shadow-123"));
|
||||
assert_eq!(stored.route_family.as_deref(), Some("openai"));
|
||||
assert_eq!(stored.route_kind.as_deref(), Some("chat"));
|
||||
assert_eq!(stored.match_status, ShadowResultMatchStatus::Pending);
|
||||
assert_eq!(stored.status_code, Some(200));
|
||||
assert!(stored.rust_result_digest.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn merges_rust_and_python_shadow_samples_into_match() {
|
||||
let repository = Arc::new(InMemoryShadowResultRepository::default());
|
||||
let state = AppState::new_with_executor(
|
||||
"http://127.0.0.1:18084",
|
||||
Some("http://127.0.0.1:18085".to_string()),
|
||||
Some("http://127.0.0.1:18086".to_string()),
|
||||
)
|
||||
.expect("app state should build")
|
||||
.with_shadow_result_data_repository_for_tests(repository.clone());
|
||||
let response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.header(CONTENT_TYPE, "application/json")
|
||||
.header(CONTROL_REQUEST_ID_HEADER, "req-shadow-compare-123")
|
||||
.body(Body::from("{}"))
|
||||
.expect("response should build");
|
||||
|
||||
record_shadow_result_non_blocking(
|
||||
state.clone(),
|
||||
"trace-shadow-compare-123",
|
||||
&Method::POST,
|
||||
"/v1/chat/completions?stream=true",
|
||||
Some(&sample_decision()),
|
||||
EXECUTION_PATH_EXECUTOR_SYNC,
|
||||
&response,
|
||||
);
|
||||
record_shadow_result_non_blocking(
|
||||
state,
|
||||
"trace-shadow-compare-123",
|
||||
&Method::POST,
|
||||
"/v1/chat/completions?stream=true",
|
||||
Some(&sample_decision()),
|
||||
EXECUTION_PATH_CONTROL_EXECUTE_SYNC,
|
||||
&response,
|
||||
);
|
||||
|
||||
for _ in 0..30 {
|
||||
if repository
|
||||
.list_recent(1)
|
||||
.await
|
||||
.map(|rows| {
|
||||
rows.first()
|
||||
.map(|row| row.match_status == ShadowResultMatchStatus::Match)
|
||||
.unwrap_or(false)
|
||||
})
|
||||
.unwrap_or(false)
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
|
||||
let stored = repository
|
||||
.list_recent(1)
|
||||
.await
|
||||
.expect("list should succeed")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("stored result should exist");
|
||||
assert_eq!(stored.request_id.as_deref(), Some("req-shadow-compare-123"));
|
||||
assert_eq!(stored.match_status, ShadowResultMatchStatus::Match);
|
||||
assert!(stored.rust_result_digest.is_some());
|
||||
assert!(stored.python_result_digest.is_some());
|
||||
}
|
||||
}
|
||||
31
crates/aether-gateway/src/cache/auth_context.rs
vendored
Normal file
31
crates/aether-gateway/src/cache/auth_context.rs
vendored
Normal file
@@ -0,0 +1,31 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_cache::ExpiringMap;
|
||||
|
||||
use crate::gateway::GatewayControlAuthContext;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub(crate) struct AuthContextCache {
|
||||
entries: ExpiringMap<String, GatewayControlAuthContext>,
|
||||
}
|
||||
|
||||
impl AuthContextCache {
|
||||
pub(crate) fn get_fresh(
|
||||
&self,
|
||||
cache_key: &str,
|
||||
ttl: Duration,
|
||||
) -> Option<GatewayControlAuthContext> {
|
||||
self.entries.get_fresh(&cache_key.to_string(), ttl)
|
||||
}
|
||||
|
||||
pub(crate) fn insert(
|
||||
&self,
|
||||
cache_key: String,
|
||||
auth_context: GatewayControlAuthContext,
|
||||
ttl: Duration,
|
||||
max_entries: usize,
|
||||
) {
|
||||
self.entries
|
||||
.insert(cache_key, auth_context, ttl, max_entries);
|
||||
}
|
||||
}
|
||||
18
crates/aether-gateway/src/cache/direct_plan_bypass.rs
vendored
Normal file
18
crates/aether-gateway/src/cache/direct_plan_bypass.rs
vendored
Normal file
@@ -0,0 +1,18 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_cache::ExpiringMap;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub(crate) struct DirectPlanBypassCache {
|
||||
entries: ExpiringMap<String, ()>,
|
||||
}
|
||||
|
||||
impl DirectPlanBypassCache {
|
||||
pub(crate) fn should_skip(&self, cache_key: &str, ttl: Duration) -> bool {
|
||||
self.entries.contains_fresh(&cache_key.to_string(), ttl)
|
||||
}
|
||||
|
||||
pub(crate) fn mark(&self, cache_key: String, ttl: Duration, max_entries: usize) {
|
||||
self.entries.insert(cache_key, (), ttl, max_entries);
|
||||
}
|
||||
}
|
||||
5
crates/aether-gateway/src/cache/mod.rs
vendored
Normal file
5
crates/aether-gateway/src/cache/mod.rs
vendored
Normal file
@@ -0,0 +1,5 @@
|
||||
mod auth_context;
|
||||
mod direct_plan_bypass;
|
||||
|
||||
pub(crate) use auth_context::AuthContextCache;
|
||||
pub(crate) use direct_plan_bypass::DirectPlanBypassCache;
|
||||
@@ -11,10 +11,15 @@ pub(crate) const EXECUTION_PATH_EXECUTOR_SYNC: &str = "executor_sync";
|
||||
pub(crate) const EXECUTION_PATH_EXECUTOR_STREAM: &str = "executor_stream";
|
||||
pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_SYNC: &str = "control_execute_sync";
|
||||
pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_STREAM: &str = "control_execute_stream";
|
||||
pub(crate) const EXECUTION_PATH_LOCAL_AUTH_DENIED: &str = "local_auth_denied";
|
||||
pub(crate) const EXECUTION_PATH_LOCAL_OVERLOADED: &str = "local_overloaded";
|
||||
pub(crate) const EXECUTION_PATH_DISTRIBUTED_OVERLOADED: &str = "distributed_overloaded";
|
||||
pub(crate) const CONTROL_ROUTE_CLASS_HEADER: &str = "x-aether-control-route-class";
|
||||
pub(crate) const CONTROL_ROUTE_FAMILY_HEADER: &str = "x-aether-control-route-family";
|
||||
pub(crate) const CONTROL_ROUTE_KIND_HEADER: &str = "x-aether-control-route-kind";
|
||||
pub(crate) const CONTROL_EXECUTOR_HEADER: &str = "x-aether-control-executor-candidate";
|
||||
pub(crate) const CONTROL_REQUEST_ID_HEADER: &str = "x-aether-control-request-id";
|
||||
pub(crate) const CONTROL_CANDIDATE_ID_HEADER: &str = "x-aether-control-candidate-id";
|
||||
pub(crate) const CONTROL_ENDPOINT_SIGNATURE_HEADER: &str = "x-aether-control-endpoint-signature";
|
||||
pub(crate) const CONTROL_EXECUTED_HEADER: &str = "x-aether-control-executed";
|
||||
pub(crate) const CONTROL_ACTION_HEADER: &str = "x-aether-control-action";
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::time::{Duration, Instant};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::http::{Response, StatusCode, Uri};
|
||||
@@ -10,7 +10,7 @@ use crate::gateway::constants::*;
|
||||
use crate::gateway::headers::{
|
||||
collect_control_headers, header_equals, header_value_str, header_value_u64, is_json_request,
|
||||
};
|
||||
use crate::gateway::{build_client_response, AppState, CachedAuthContextEntry, GatewayError};
|
||||
use crate::gateway::{build_client_response, AppState, GatewayError};
|
||||
|
||||
const AUTH_CONTEXT_CACHE_TTL: Duration = Duration::from_secs(60);
|
||||
const AUTH_CONTEXT_CACHE_MAX_ENTRIES: usize = 256;
|
||||
@@ -72,6 +72,15 @@ pub(crate) struct GatewayControlAuthContext {
|
||||
pub(crate) api_key_id: String,
|
||||
pub(crate) balance_remaining: Option<f64>,
|
||||
pub(crate) access_allowed: bool,
|
||||
#[serde(skip)]
|
||||
pub(crate) local_rejection: Option<GatewayLocalAuthRejection>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) enum GatewayLocalAuthRejection {
|
||||
InvalidApiKey,
|
||||
LockedApiKey,
|
||||
BalanceDenied { remaining: Option<f64> },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -116,10 +125,31 @@ pub(crate) async fn resolve_control_route(
|
||||
};
|
||||
decision.public_query_string = uri.query().map(ToOwned::to_owned);
|
||||
|
||||
if let Some(auth_context) = resolve_data_backed_auth_context(
|
||||
state,
|
||||
headers,
|
||||
decision.auth_endpoint_signature.as_deref(),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
if let Some(cache_key) = decision
|
||||
.auth_endpoint_signature
|
||||
.as_deref()
|
||||
.and_then(|signature| build_auth_context_cache_key(headers, uri, signature))
|
||||
{
|
||||
put_cached_auth_context(state, cache_key, auth_context.clone());
|
||||
}
|
||||
decision.auth_context = Some(auth_context);
|
||||
}
|
||||
|
||||
if state.executor_base_url.is_some() && decision.executor_candidate {
|
||||
return Ok(Some(decision));
|
||||
}
|
||||
|
||||
if decision.auth_context.is_some() {
|
||||
return Ok(Some(decision));
|
||||
}
|
||||
|
||||
match fetch_auth_context(
|
||||
state,
|
||||
control_base_url,
|
||||
@@ -170,6 +200,13 @@ pub(crate) async fn resolve_executor_auth_context(
|
||||
return Ok(Some(auth_context));
|
||||
}
|
||||
|
||||
if let Some(auth_context) =
|
||||
resolve_data_backed_auth_context(state, headers, Some(auth_endpoint_signature)).await?
|
||||
{
|
||||
put_cached_auth_context(state, cache_key, auth_context.clone());
|
||||
return Ok(Some(auth_context));
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
@@ -190,6 +227,19 @@ pub(crate) fn cache_executor_auth_context(
|
||||
put_cached_auth_context(state, cache_key, auth_context);
|
||||
}
|
||||
|
||||
pub(crate) fn trusted_auth_local_rejection(
|
||||
decision: Option<&GatewayControlDecision>,
|
||||
_headers: &http::HeaderMap,
|
||||
) -> Option<GatewayLocalAuthRejection> {
|
||||
let decision = decision?;
|
||||
if decision.route_class.as_deref() != Some("ai_public") {
|
||||
return None;
|
||||
}
|
||||
|
||||
let auth_context = decision.auth_context.as_ref()?;
|
||||
auth_context.local_rejection.clone()
|
||||
}
|
||||
|
||||
async fn fetch_auth_context(
|
||||
state: &AppState,
|
||||
control_base_url: &str,
|
||||
@@ -334,13 +384,9 @@ fn build_auth_context_cache_key(
|
||||
}
|
||||
|
||||
fn get_cached_auth_context(state: &AppState, cache_key: &str) -> Option<GatewayControlAuthContext> {
|
||||
let mut cache = state.auth_context_cache.lock().ok()?;
|
||||
let entry = cache.get(cache_key)?.clone();
|
||||
if entry.cached_at.elapsed() > AUTH_CONTEXT_CACHE_TTL {
|
||||
cache.remove(cache_key);
|
||||
return None;
|
||||
}
|
||||
Some(entry.auth_context)
|
||||
state
|
||||
.auth_context_cache
|
||||
.get_fresh(cache_key, AUTH_CONTEXT_CACHE_TTL)
|
||||
}
|
||||
|
||||
fn put_cached_auth_context(
|
||||
@@ -348,28 +394,100 @@ fn put_cached_auth_context(
|
||||
cache_key: String,
|
||||
auth_context: GatewayControlAuthContext,
|
||||
) {
|
||||
let Ok(mut cache) = state.auth_context_cache.lock() else {
|
||||
return;
|
||||
state.auth_context_cache.insert(
|
||||
cache_key,
|
||||
auth_context,
|
||||
AUTH_CONTEXT_CACHE_TTL,
|
||||
AUTH_CONTEXT_CACHE_MAX_ENTRIES,
|
||||
);
|
||||
}
|
||||
|
||||
async fn resolve_data_backed_auth_context(
|
||||
state: &AppState,
|
||||
headers: &http::HeaderMap,
|
||||
auth_endpoint_signature: Option<&str>,
|
||||
) -> Result<Option<GatewayControlAuthContext>, GatewayError> {
|
||||
let Some(signature) = auth_endpoint_signature
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let _ = signature;
|
||||
|
||||
let Some(user_id) =
|
||||
header_value_str(headers, TRUSTED_AUTH_USER_ID_HEADER).filter(|value| !value.is_empty())
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(api_key_id) =
|
||||
header_value_str(headers, TRUSTED_AUTH_API_KEY_ID_HEADER).filter(|value| !value.is_empty())
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
cache.retain(|_, entry| entry.cached_at.elapsed() <= AUTH_CONTEXT_CACHE_TTL);
|
||||
if cache.len() >= AUTH_CONTEXT_CACHE_MAX_ENTRIES {
|
||||
if let Some(oldest_key) = cache
|
||||
.iter()
|
||||
.min_by_key(|(_, entry)| entry.cached_at)
|
||||
.map(|(key, _)| key.clone())
|
||||
{
|
||||
cache.remove(&oldest_key);
|
||||
}
|
||||
}
|
||||
let snapshot = state
|
||||
.read_auth_api_key_snapshot(&user_id, &api_key_id, current_unix_secs())
|
||||
.await?;
|
||||
let Some(snapshot) = snapshot else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
cache.insert(
|
||||
cache_key,
|
||||
CachedAuthContextEntry {
|
||||
auth_context,
|
||||
cached_at: Instant::now(),
|
||||
},
|
||||
);
|
||||
let header_access_allowed = header_value_str(headers, TRUSTED_AUTH_ACCESS_ALLOWED_HEADER)
|
||||
.as_deref()
|
||||
.and_then(parse_bool_header);
|
||||
let invalid_api_key = !snapshot.user_is_active
|
||||
|| snapshot.user_is_deleted
|
||||
|| !snapshot.api_key_is_active
|
||||
|| snapshot
|
||||
.api_key_expires_at_unix_secs
|
||||
.is_some_and(|expires_at| expires_at < current_unix_secs());
|
||||
let locked_api_key = snapshot.api_key_is_locked && !snapshot.api_key_is_standalone;
|
||||
let access_allowed = header_access_allowed
|
||||
.map(|value| value && snapshot.currently_usable)
|
||||
.unwrap_or(snapshot.currently_usable);
|
||||
let local_rejection = if invalid_api_key {
|
||||
Some(GatewayLocalAuthRejection::InvalidApiKey)
|
||||
} else if locked_api_key {
|
||||
Some(GatewayLocalAuthRejection::LockedApiKey)
|
||||
} else if header_access_allowed.is_some_and(|value| !value) && snapshot.currently_usable {
|
||||
Some(GatewayLocalAuthRejection::BalanceDenied {
|
||||
remaining: header_value_str(headers, TRUSTED_AUTH_BALANCE_HEADER)
|
||||
.as_deref()
|
||||
.and_then(parse_f64_header),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Ok(Some(GatewayControlAuthContext {
|
||||
user_id: snapshot.user_id,
|
||||
api_key_id: snapshot.api_key_id,
|
||||
balance_remaining: header_value_str(headers, TRUSTED_AUTH_BALANCE_HEADER)
|
||||
.as_deref()
|
||||
.and_then(parse_f64_header),
|
||||
access_allowed,
|
||||
local_rejection,
|
||||
}))
|
||||
}
|
||||
|
||||
fn parse_bool_header(value: &str) -> Option<bool> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"true" | "1" | "yes" => Some(true),
|
||||
"false" | "0" | "no" => Some(false),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_f64_header(value: &str) -> Option<f64> {
|
||||
value.trim().parse::<f64>().ok()
|
||||
}
|
||||
|
||||
fn current_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
fn classify_control_route(
|
||||
|
||||
155
crates/aether-gateway/src/data/auth.rs
Normal file
155
crates/aether-gateway/src/data/auth.rs
Normal file
@@ -0,0 +1,155 @@
|
||||
use aether_data::repository::auth::{AuthApiKeyLookupKey, StoredAuthApiKeySnapshot};
|
||||
use aether_data::DataLayerError;
|
||||
|
||||
use super::state::GatewayDataState;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
|
||||
pub(crate) struct StoredGatewayAuthApiKeySnapshot {
|
||||
pub(crate) user_id: String,
|
||||
pub(crate) username: String,
|
||||
pub(crate) email: Option<String>,
|
||||
pub(crate) user_role: String,
|
||||
pub(crate) user_auth_source: String,
|
||||
pub(crate) user_is_active: bool,
|
||||
pub(crate) user_is_deleted: bool,
|
||||
pub(crate) user_allowed_providers: Option<Vec<String>>,
|
||||
pub(crate) user_allowed_api_formats: Option<Vec<String>>,
|
||||
pub(crate) user_allowed_models: Option<Vec<String>>,
|
||||
pub(crate) api_key_id: String,
|
||||
pub(crate) api_key_name: Option<String>,
|
||||
pub(crate) api_key_is_active: bool,
|
||||
pub(crate) api_key_is_locked: bool,
|
||||
pub(crate) api_key_is_standalone: bool,
|
||||
pub(crate) api_key_rate_limit: Option<i32>,
|
||||
pub(crate) api_key_concurrent_limit: Option<i32>,
|
||||
pub(crate) api_key_expires_at_unix_secs: Option<u64>,
|
||||
pub(crate) api_key_allowed_providers: Option<Vec<String>>,
|
||||
pub(crate) api_key_allowed_api_formats: Option<Vec<String>>,
|
||||
pub(crate) api_key_allowed_models: Option<Vec<String>>,
|
||||
pub(crate) currently_usable: bool,
|
||||
}
|
||||
|
||||
impl StoredGatewayAuthApiKeySnapshot {
|
||||
fn from_stored(snapshot: StoredAuthApiKeySnapshot, now_unix_secs: u64) -> Self {
|
||||
let currently_usable = snapshot.is_currently_usable(now_unix_secs);
|
||||
Self {
|
||||
user_id: snapshot.user_id,
|
||||
username: snapshot.username,
|
||||
email: snapshot.email,
|
||||
user_role: snapshot.user_role,
|
||||
user_auth_source: snapshot.user_auth_source,
|
||||
user_is_active: snapshot.user_is_active,
|
||||
user_is_deleted: snapshot.user_is_deleted,
|
||||
user_allowed_providers: snapshot.user_allowed_providers,
|
||||
user_allowed_api_formats: snapshot.user_allowed_api_formats,
|
||||
user_allowed_models: snapshot.user_allowed_models,
|
||||
api_key_id: snapshot.api_key_id,
|
||||
api_key_name: snapshot.api_key_name,
|
||||
api_key_is_active: snapshot.api_key_is_active,
|
||||
api_key_is_locked: snapshot.api_key_is_locked,
|
||||
api_key_is_standalone: snapshot.api_key_is_standalone,
|
||||
api_key_rate_limit: snapshot.api_key_rate_limit,
|
||||
api_key_concurrent_limit: snapshot.api_key_concurrent_limit,
|
||||
api_key_expires_at_unix_secs: snapshot.api_key_expires_at_unix_secs,
|
||||
api_key_allowed_providers: snapshot.api_key_allowed_providers,
|
||||
api_key_allowed_api_formats: snapshot.api_key_allowed_api_formats,
|
||||
api_key_allowed_models: snapshot.api_key_allowed_models,
|
||||
currently_usable,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn read_auth_api_key_snapshot(
|
||||
state: &GatewayDataState,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<StoredGatewayAuthApiKeySnapshot>, DataLayerError> {
|
||||
let snapshot = state
|
||||
.find_auth_api_key_snapshot(AuthApiKeyLookupKey::UserApiKeyIds {
|
||||
user_id,
|
||||
api_key_id,
|
||||
})
|
||||
.await?;
|
||||
Ok(snapshot
|
||||
.map(|snapshot| StoredGatewayAuthApiKeySnapshot::from_stored(snapshot, now_unix_secs)))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::GatewayDataState;
|
||||
use super::{read_auth_api_key_snapshot, StoredGatewayAuthApiKeySnapshot};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(200),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
)
|
||||
.expect("snapshot should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_trusted_auth_snapshot_and_derives_usability() {
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_snapshot("key-1", "user-1"),
|
||||
)]));
|
||||
let state = GatewayDataState::with_auth_api_key_reader_for_tests(repository);
|
||||
|
||||
let snapshot = read_auth_api_key_snapshot(&state, "user-1", "key-1", 150)
|
||||
.await
|
||||
.expect("read should succeed")
|
||||
.expect("snapshot should exist");
|
||||
|
||||
assert_eq!(
|
||||
snapshot,
|
||||
StoredGatewayAuthApiKeySnapshot {
|
||||
user_id: "user-1".to_string(),
|
||||
username: "alice".to_string(),
|
||||
email: Some("alice@example.com".to_string()),
|
||||
user_role: "user".to_string(),
|
||||
user_auth_source: "local".to_string(),
|
||||
user_is_active: true,
|
||||
user_is_deleted: false,
|
||||
user_allowed_providers: Some(vec!["openai".to_string()]),
|
||||
user_allowed_api_formats: Some(vec!["openai:chat".to_string()]),
|
||||
user_allowed_models: Some(vec!["gpt-4.1".to_string()]),
|
||||
api_key_id: "key-1".to_string(),
|
||||
api_key_name: Some("default".to_string()),
|
||||
api_key_is_active: true,
|
||||
api_key_is_locked: false,
|
||||
api_key_is_standalone: false,
|
||||
api_key_rate_limit: Some(60),
|
||||
api_key_concurrent_limit: Some(5),
|
||||
api_key_expires_at_unix_secs: Some(200),
|
||||
api_key_allowed_providers: Some(vec!["openai".to_string()]),
|
||||
api_key_allowed_api_formats: Some(vec!["openai:chat".to_string()]),
|
||||
api_key_allowed_models: Some(vec!["gpt-4.1".to_string()]),
|
||||
currently_usable: true,
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
213
crates/aether-gateway/src/data/candidates.rs
Normal file
213
crates/aether-gateway/src/data/candidates.rs
Normal file
@@ -0,0 +1,213 @@
|
||||
use aether_data::repository::candidates::{RequestCandidateStatus, StoredRequestCandidate};
|
||||
use aether_data::DataLayerError;
|
||||
|
||||
use super::state::GatewayDataState;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) enum RequestCandidateFinalStatus {
|
||||
Success,
|
||||
Failed,
|
||||
Cancelled,
|
||||
Streaming,
|
||||
Pending,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
|
||||
pub(crate) struct RequestCandidateTrace {
|
||||
pub(crate) request_id: String,
|
||||
pub(crate) total_candidates: usize,
|
||||
pub(crate) final_status: RequestCandidateFinalStatus,
|
||||
pub(crate) total_latency_ms: u64,
|
||||
pub(crate) candidates: Vec<StoredRequestCandidate>,
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_candidate_trace(
|
||||
state: &GatewayDataState,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
) -> Result<Option<RequestCandidateTrace>, DataLayerError> {
|
||||
let all_candidates = state
|
||||
.list_request_candidates_by_request_id(request_id)
|
||||
.await?;
|
||||
if all_candidates.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let candidates = if attempted_only {
|
||||
all_candidates
|
||||
.iter()
|
||||
.filter(|candidate| {
|
||||
candidate
|
||||
.status
|
||||
.is_attempted(candidate.started_at_unix_secs)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>()
|
||||
} else {
|
||||
all_candidates.clone()
|
||||
};
|
||||
|
||||
let total_latency_ms = candidates
|
||||
.iter()
|
||||
.filter(|candidate| {
|
||||
matches!(
|
||||
candidate.status,
|
||||
RequestCandidateStatus::Success
|
||||
| RequestCandidateStatus::Failed
|
||||
| RequestCandidateStatus::Cancelled
|
||||
) && candidate.latency_ms.is_some()
|
||||
})
|
||||
.map(|candidate| candidate.latency_ms.unwrap_or(0))
|
||||
.sum();
|
||||
let final_status_source = if attempted_only && candidates.is_empty() {
|
||||
&all_candidates
|
||||
} else {
|
||||
&candidates
|
||||
};
|
||||
|
||||
Ok(Some(RequestCandidateTrace {
|
||||
request_id: request_id.to_string(),
|
||||
total_candidates: candidates.len(),
|
||||
final_status: derive_final_status(final_status_source),
|
||||
total_latency_ms,
|
||||
candidates,
|
||||
}))
|
||||
}
|
||||
|
||||
fn derive_final_status(candidates: &[StoredRequestCandidate]) -> RequestCandidateFinalStatus {
|
||||
let has_success = candidates.iter().any(|candidate| {
|
||||
candidate.status == RequestCandidateStatus::Success
|
||||
|| matches!(candidate.status_code, Some(status_code) if (200..300).contains(&status_code))
|
||||
});
|
||||
if has_success {
|
||||
return RequestCandidateFinalStatus::Success;
|
||||
}
|
||||
|
||||
if candidates
|
||||
.iter()
|
||||
.any(|candidate| candidate.status == RequestCandidateStatus::Streaming)
|
||||
{
|
||||
return RequestCandidateFinalStatus::Streaming;
|
||||
}
|
||||
|
||||
if candidates
|
||||
.iter()
|
||||
.any(|candidate| candidate.status == RequestCandidateStatus::Pending)
|
||||
{
|
||||
return RequestCandidateFinalStatus::Pending;
|
||||
}
|
||||
|
||||
let has_cancelled = candidates
|
||||
.iter()
|
||||
.any(|candidate| candidate.status == RequestCandidateStatus::Cancelled);
|
||||
let has_failed = candidates
|
||||
.iter()
|
||||
.any(|candidate| candidate.status == RequestCandidateStatus::Failed);
|
||||
if has_cancelled && !has_failed {
|
||||
return RequestCandidateFinalStatus::Cancelled;
|
||||
}
|
||||
|
||||
RequestCandidateFinalStatus::Failed
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::GatewayDataState;
|
||||
use super::{derive_final_status, read_request_candidate_trace, RequestCandidateFinalStatus};
|
||||
use aether_data::repository::candidates::{
|
||||
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
fn sample_candidate(
|
||||
id: &str,
|
||||
request_id: &str,
|
||||
candidate_index: i32,
|
||||
status: RequestCandidateStatus,
|
||||
started_at_unix_secs: Option<i64>,
|
||||
latency_ms: Option<i32>,
|
||||
status_code: Option<i32>,
|
||||
) -> StoredRequestCandidate {
|
||||
StoredRequestCandidate::new(
|
||||
id.to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
candidate_index,
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
status,
|
||||
None,
|
||||
false,
|
||||
status_code,
|
||||
None,
|
||||
None,
|
||||
latency_ms,
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
100 + i64::from(candidate_index),
|
||||
started_at_unix_secs,
|
||||
started_at_unix_secs.map(|value| value + 1),
|
||||
)
|
||||
.expect("candidate should build")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn derive_final_status_prefers_success() {
|
||||
let candidates = vec![sample_candidate(
|
||||
"cand-1",
|
||||
"req-1",
|
||||
0,
|
||||
RequestCandidateStatus::Success,
|
||||
Some(100),
|
||||
Some(25),
|
||||
Some(200),
|
||||
)];
|
||||
|
||||
assert_eq!(
|
||||
derive_final_status(&candidates),
|
||||
RequestCandidateFinalStatus::Success
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_request_candidate_trace_filters_attempted_rows() {
|
||||
let repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_candidate(
|
||||
"cand-1",
|
||||
"req-1",
|
||||
0,
|
||||
RequestCandidateStatus::Pending,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
),
|
||||
sample_candidate(
|
||||
"cand-2",
|
||||
"req-1",
|
||||
1,
|
||||
RequestCandidateStatus::Failed,
|
||||
Some(101),
|
||||
Some(33),
|
||||
Some(502),
|
||||
),
|
||||
]));
|
||||
let state = GatewayDataState::with_request_candidate_reader_for_tests(repository);
|
||||
|
||||
let trace = read_request_candidate_trace(&state, "req-1", true)
|
||||
.await
|
||||
.expect("trace should succeed")
|
||||
.expect("trace should exist");
|
||||
|
||||
assert_eq!(trace.total_candidates, 1);
|
||||
assert_eq!(trace.candidates[0].id, "cand-2");
|
||||
assert_eq!(trace.final_status, RequestCandidateFinalStatus::Failed);
|
||||
assert_eq!(trace.total_latency_ms, 33);
|
||||
}
|
||||
}
|
||||
41
crates/aether-gateway/src/data/config.rs
Normal file
41
crates/aether-gateway/src/data/config.rs
Normal file
@@ -0,0 +1,41 @@
|
||||
use aether_data::postgres::PostgresPoolConfig;
|
||||
use aether_data::DataLayerConfig;
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct GatewayDataConfig {
|
||||
postgres: Option<PostgresPoolConfig>,
|
||||
}
|
||||
|
||||
impl GatewayDataConfig {
|
||||
pub fn disabled() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn from_postgres_config(postgres: PostgresPoolConfig) -> Self {
|
||||
Self {
|
||||
postgres: Some(postgres),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_postgres_url(database_url: impl Into<String>, require_ssl: bool) -> Self {
|
||||
let mut postgres = PostgresPoolConfig::default();
|
||||
postgres.database_url = database_url.into();
|
||||
postgres.require_ssl = require_ssl;
|
||||
Self::from_postgres_config(postgres)
|
||||
}
|
||||
|
||||
pub fn postgres(&self) -> Option<&PostgresPoolConfig> {
|
||||
self.postgres.as_ref()
|
||||
}
|
||||
|
||||
pub fn is_enabled(&self) -> bool {
|
||||
self.postgres.is_some()
|
||||
}
|
||||
|
||||
pub fn to_data_layer_config(&self) -> DataLayerConfig {
|
||||
DataLayerConfig {
|
||||
postgres: self.postgres.clone(),
|
||||
redis: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
267
crates/aether-gateway/src/data/decision_trace.rs
Normal file
267
crates/aether-gateway/src/data/decision_trace.rs
Normal file
@@ -0,0 +1,267 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use aether_data::repository::candidates::StoredRequestCandidate;
|
||||
use aether_data::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_data::DataLayerError;
|
||||
|
||||
use super::candidates::RequestCandidateFinalStatus;
|
||||
use super::state::GatewayDataState;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
|
||||
pub(crate) struct DecisionTraceCandidate {
|
||||
#[serde(flatten)]
|
||||
pub(crate) candidate: StoredRequestCandidate,
|
||||
pub(crate) provider_name: Option<String>,
|
||||
pub(crate) provider_website: Option<String>,
|
||||
pub(crate) provider_type: Option<String>,
|
||||
pub(crate) endpoint_api_format: Option<String>,
|
||||
pub(crate) endpoint_api_family: Option<String>,
|
||||
pub(crate) endpoint_kind: Option<String>,
|
||||
pub(crate) provider_key_name: Option<String>,
|
||||
pub(crate) provider_key_auth_type: Option<String>,
|
||||
pub(crate) provider_key_capabilities: Option<serde_json::Value>,
|
||||
pub(crate) provider_key_is_active: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
|
||||
pub(crate) struct DecisionTrace {
|
||||
pub(crate) request_id: String,
|
||||
pub(crate) total_candidates: usize,
|
||||
pub(crate) final_status: RequestCandidateFinalStatus,
|
||||
pub(crate) total_latency_ms: u64,
|
||||
pub(crate) candidates: Vec<DecisionTraceCandidate>,
|
||||
}
|
||||
|
||||
pub(crate) async fn read_decision_trace(
|
||||
state: &GatewayDataState,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
) -> Result<Option<DecisionTrace>, DataLayerError> {
|
||||
let Some(trace) = state
|
||||
.read_request_candidate_trace(request_id, attempted_only)
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let provider_ids = unique_ids(
|
||||
trace
|
||||
.candidates
|
||||
.iter()
|
||||
.filter_map(|item| item.provider_id.as_ref()),
|
||||
);
|
||||
let endpoint_ids = unique_ids(
|
||||
trace
|
||||
.candidates
|
||||
.iter()
|
||||
.filter_map(|item| item.endpoint_id.as_ref()),
|
||||
);
|
||||
let key_ids = unique_ids(
|
||||
trace
|
||||
.candidates
|
||||
.iter()
|
||||
.filter_map(|item| item.key_id.as_ref()),
|
||||
);
|
||||
|
||||
let provider_map = state
|
||||
.list_provider_catalog_providers_by_ids(&provider_ids)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|item| (item.id.clone(), item))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let endpoint_map = state
|
||||
.list_provider_catalog_endpoints_by_ids(&endpoint_ids)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|item| (item.id.clone(), item))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let key_map = state
|
||||
.list_provider_catalog_keys_by_ids(&key_ids)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|item| (item.id.clone(), item))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
|
||||
Ok(Some(DecisionTrace {
|
||||
request_id: trace.request_id,
|
||||
total_candidates: trace.total_candidates,
|
||||
final_status: trace.final_status,
|
||||
total_latency_ms: trace.total_latency_ms,
|
||||
candidates: trace
|
||||
.candidates
|
||||
.into_iter()
|
||||
.map(|candidate| enrich_candidate(candidate, &provider_map, &endpoint_map, &key_map))
|
||||
.collect(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn enrich_candidate(
|
||||
candidate: StoredRequestCandidate,
|
||||
provider_map: &BTreeMap<String, StoredProviderCatalogProvider>,
|
||||
endpoint_map: &BTreeMap<String, StoredProviderCatalogEndpoint>,
|
||||
key_map: &BTreeMap<String, StoredProviderCatalogKey>,
|
||||
) -> DecisionTraceCandidate {
|
||||
let provider = candidate
|
||||
.provider_id
|
||||
.as_ref()
|
||||
.and_then(|provider_id| provider_map.get(provider_id));
|
||||
let endpoint = candidate
|
||||
.endpoint_id
|
||||
.as_ref()
|
||||
.and_then(|endpoint_id| endpoint_map.get(endpoint_id));
|
||||
let provider_key = candidate
|
||||
.key_id
|
||||
.as_ref()
|
||||
.and_then(|key_id| key_map.get(key_id));
|
||||
|
||||
DecisionTraceCandidate {
|
||||
provider_name: provider.map(|item| item.name.clone()),
|
||||
provider_website: provider.and_then(|item| item.website.clone()),
|
||||
provider_type: provider.map(|item| item.provider_type.clone()),
|
||||
endpoint_api_format: endpoint.map(|item| item.api_format.clone()),
|
||||
endpoint_api_family: endpoint.and_then(|item| item.api_family.clone()),
|
||||
endpoint_kind: endpoint.and_then(|item| item.endpoint_kind.clone()),
|
||||
provider_key_name: provider_key
|
||||
.map(|item| item.name.clone())
|
||||
.or_else(|| candidate.api_key_name.clone()),
|
||||
provider_key_auth_type: provider_key.map(|item| item.auth_type.clone()),
|
||||
provider_key_capabilities: provider_key.and_then(|item| item.capabilities.clone()),
|
||||
provider_key_is_active: provider_key.map(|item| item.is_active),
|
||||
candidate,
|
||||
}
|
||||
}
|
||||
|
||||
fn unique_ids<'a>(items: impl Iterator<Item = &'a String>) -> Vec<String> {
|
||||
items
|
||||
.cloned()
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_data::repository::candidates::{
|
||||
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
|
||||
use super::{read_decision_trace, DecisionTrace, DecisionTraceCandidate};
|
||||
use crate::gateway::data::candidates::RequestCandidateFinalStatus;
|
||||
use crate::gateway::data::GatewayDataState;
|
||||
|
||||
fn sample_candidate(request_id: &str) -> StoredRequestCandidate {
|
||||
StoredRequestCandidate::new(
|
||||
"cand-1".to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
0,
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
RequestCandidateStatus::Failed,
|
||||
None,
|
||||
false,
|
||||
Some(502),
|
||||
Some("bad_gateway".to_string()),
|
||||
Some("upstream failed".to_string()),
|
||||
Some(37),
|
||||
Some(1),
|
||||
None,
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
100,
|
||||
Some(101),
|
||||
Some(102),
|
||||
)
|
||||
.expect("candidate should build")
|
||||
}
|
||||
|
||||
fn sample_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-1".to_string(),
|
||||
"OpenAI".to_string(),
|
||||
Some("https://openai.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
}
|
||||
|
||||
fn sample_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"openai:chat".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
}
|
||||
|
||||
fn sample_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"provider-key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"prod-key".to_string(),
|
||||
"api_key".to_string(),
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn enriches_request_candidate_trace_with_provider_catalog_metadata() {
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_candidate("req-1"),
|
||||
]));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider()],
|
||||
vec![sample_endpoint()],
|
||||
vec![sample_key()],
|
||||
));
|
||||
let state = GatewayDataState::with_decision_trace_readers_for_tests(
|
||||
request_candidates,
|
||||
provider_catalog,
|
||||
);
|
||||
|
||||
let trace = read_decision_trace(&state, "req-1", true)
|
||||
.await
|
||||
.expect("trace should read")
|
||||
.expect("trace should exist");
|
||||
|
||||
assert_eq!(
|
||||
trace,
|
||||
DecisionTrace {
|
||||
request_id: "req-1".to_string(),
|
||||
total_candidates: 1,
|
||||
final_status: RequestCandidateFinalStatus::Failed,
|
||||
total_latency_ms: 37,
|
||||
candidates: vec![DecisionTraceCandidate {
|
||||
candidate: sample_candidate("req-1"),
|
||||
provider_name: Some("OpenAI".to_string()),
|
||||
provider_website: Some("https://openai.com".to_string()),
|
||||
provider_type: Some("custom".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
provider_key_name: Some("prod-key".to_string()),
|
||||
provider_key_auth_type: Some("api_key".to_string()),
|
||||
provider_key_capabilities: Some(serde_json::json!({"cache_1h": true})),
|
||||
provider_key_is_active: Some(true),
|
||||
}],
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
80
crates/aether-gateway/src/data/gemini.rs
Normal file
80
crates/aether-gateway/src/data/gemini.rs
Normal file
@@ -0,0 +1,80 @@
|
||||
use aether_data::repository::video_tasks::{StoredVideoTask, VideoTaskStatus};
|
||||
use serde_json::json;
|
||||
|
||||
use crate::gateway::video_tasks::LocalVideoTaskReadResponse;
|
||||
|
||||
pub(super) fn map_gemini_video_task_to_read_response(
|
||||
task: StoredVideoTask,
|
||||
) -> LocalVideoTaskReadResponse {
|
||||
match task.status {
|
||||
VideoTaskStatus::Cancelled => LocalVideoTaskReadResponse {
|
||||
status_code: 404,
|
||||
body_json: json!({"detail": "Video task was cancelled"}),
|
||||
},
|
||||
VideoTaskStatus::Deleted => LocalVideoTaskReadResponse {
|
||||
status_code: 404,
|
||||
body_json: json!({"detail": "Video task not found"}),
|
||||
},
|
||||
VideoTaskStatus::Completed => LocalVideoTaskReadResponse {
|
||||
status_code: 200,
|
||||
body_json: build_gemini_completed_body(task),
|
||||
},
|
||||
VideoTaskStatus::Failed | VideoTaskStatus::Expired => LocalVideoTaskReadResponse {
|
||||
status_code: 200,
|
||||
body_json: build_gemini_failed_body(task),
|
||||
},
|
||||
_ => LocalVideoTaskReadResponse {
|
||||
status_code: 200,
|
||||
body_json: build_gemini_pending_body(task),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn build_gemini_completed_body(task: StoredVideoTask) -> serde_json::Value {
|
||||
let operation_name = operation_name(&task);
|
||||
let short_id = task.short_id.unwrap_or_default();
|
||||
|
||||
json!({
|
||||
"name": operation_name,
|
||||
"done": true,
|
||||
"response": {
|
||||
"generateVideoResponse": {
|
||||
"generatedSamples": [
|
||||
{
|
||||
"video": {
|
||||
"uri": format!("/v1beta/files/aev_{short_id}:download?alt=media"),
|
||||
"mimeType": "video/mp4"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn build_gemini_failed_body(task: StoredVideoTask) -> serde_json::Value {
|
||||
json!({
|
||||
"name": operation_name(&task),
|
||||
"done": true,
|
||||
"error": {
|
||||
"code": task.error_code.unwrap_or_else(|| "UNKNOWN".to_string()),
|
||||
"message": task
|
||||
.error_message
|
||||
.unwrap_or_else(|| "Video generation failed".to_string()),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn build_gemini_pending_body(task: StoredVideoTask) -> serde_json::Value {
|
||||
json!({
|
||||
"name": operation_name(&task),
|
||||
"done": false,
|
||||
"metadata": {}
|
||||
})
|
||||
}
|
||||
|
||||
fn operation_name(task: &StoredVideoTask) -> String {
|
||||
let model = task.model.clone().unwrap_or_else(|| "unknown".to_string());
|
||||
let short_id = task.short_id.clone().unwrap_or_else(|| task.id.clone());
|
||||
format!("models/{model}/operations/{short_id}")
|
||||
}
|
||||
21
crates/aether-gateway/src/data/mod.rs
Normal file
21
crates/aether-gateway/src/data/mod.rs
Normal file
@@ -0,0 +1,21 @@
|
||||
mod auth;
|
||||
mod candidates;
|
||||
mod config;
|
||||
mod decision_trace;
|
||||
mod gemini;
|
||||
mod openai;
|
||||
mod request_audit;
|
||||
mod state;
|
||||
mod usage;
|
||||
mod video_tasks;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
pub(crate) use auth::StoredGatewayAuthApiKeySnapshot;
|
||||
pub(crate) use candidates::RequestCandidateTrace;
|
||||
pub use config::GatewayDataConfig;
|
||||
pub(crate) use decision_trace::DecisionTrace;
|
||||
pub(crate) use request_audit::RequestAuditBundle;
|
||||
pub(crate) use state::GatewayDataState;
|
||||
pub(crate) use usage::RequestUsageAudit;
|
||||
69
crates/aether-gateway/src/data/openai.rs
Normal file
69
crates/aether-gateway/src/data/openai.rs
Normal file
@@ -0,0 +1,69 @@
|
||||
use aether_data::repository::video_tasks::{StoredVideoTask, VideoTaskStatus};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::gateway::video_tasks::LocalVideoTaskReadResponse;
|
||||
|
||||
pub(super) fn map_openai_video_task_to_read_response(
|
||||
task: StoredVideoTask,
|
||||
) -> LocalVideoTaskReadResponse {
|
||||
match task.status {
|
||||
VideoTaskStatus::Cancelled => LocalVideoTaskReadResponse {
|
||||
status_code: 404,
|
||||
body_json: json!({"detail": "Video task was cancelled"}),
|
||||
},
|
||||
VideoTaskStatus::Deleted => LocalVideoTaskReadResponse {
|
||||
status_code: 404,
|
||||
body_json: json!({"detail": "Video task not found"}),
|
||||
},
|
||||
status => LocalVideoTaskReadResponse {
|
||||
status_code: 200,
|
||||
body_json: build_openai_video_task_body(task, status),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn build_openai_video_task_body(task: StoredVideoTask, status: VideoTaskStatus) -> Value {
|
||||
let mut body = json!({
|
||||
"id": task.id,
|
||||
"object": "video",
|
||||
"status": map_openai_video_status(status),
|
||||
"progress": task.progress_percent,
|
||||
"created_at": task.created_at_unix_secs,
|
||||
});
|
||||
|
||||
if let Some(model) = task.model {
|
||||
body["model"] = Value::String(model);
|
||||
}
|
||||
if let Some(prompt) = task.prompt {
|
||||
body["prompt"] = Value::String(prompt);
|
||||
}
|
||||
if let Some(size) = task.size {
|
||||
body["size"] = Value::String(size);
|
||||
}
|
||||
if let Some(video_url) = task.video_url {
|
||||
body["video_url"] = Value::String(video_url);
|
||||
}
|
||||
if matches!(
|
||||
status,
|
||||
VideoTaskStatus::Failed | VideoTaskStatus::Expired | VideoTaskStatus::Cancelled
|
||||
) {
|
||||
body["error"] = json!({
|
||||
"code": task.error_code.unwrap_or_else(|| "unknown".to_string()),
|
||||
"message": task
|
||||
.error_message
|
||||
.unwrap_or_else(|| "Video generation failed".to_string()),
|
||||
});
|
||||
}
|
||||
|
||||
body
|
||||
}
|
||||
|
||||
fn map_openai_video_status(status: VideoTaskStatus) -> &'static str {
|
||||
match status {
|
||||
VideoTaskStatus::Pending | VideoTaskStatus::Submitted | VideoTaskStatus::Queued => "queued",
|
||||
VideoTaskStatus::Processing => "processing",
|
||||
VideoTaskStatus::Completed => "completed",
|
||||
VideoTaskStatus::Failed | VideoTaskStatus::Cancelled | VideoTaskStatus::Expired => "failed",
|
||||
VideoTaskStatus::Deleted => "deleted",
|
||||
}
|
||||
}
|
||||
56
crates/aether-gateway/src/data/request_audit.rs
Normal file
56
crates/aether-gateway/src/data/request_audit.rs
Normal file
@@ -0,0 +1,56 @@
|
||||
use aether_data::DataLayerError;
|
||||
|
||||
use super::auth::StoredGatewayAuthApiKeySnapshot;
|
||||
use super::decision_trace::DecisionTrace;
|
||||
use super::state::GatewayDataState;
|
||||
use super::usage::RequestUsageAudit;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
|
||||
pub(crate) struct RequestAuditBundle {
|
||||
pub(crate) request_id: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) usage: Option<RequestUsageAudit>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) decision_trace: Option<DecisionTrace>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) auth_snapshot: Option<StoredGatewayAuthApiKeySnapshot>,
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_audit_bundle(
|
||||
state: &GatewayDataState,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<RequestAuditBundle>, DataLayerError> {
|
||||
let usage = state.read_request_usage_audit(request_id).await?;
|
||||
let decision_trace = state
|
||||
.read_decision_trace(request_id, attempted_only)
|
||||
.await?;
|
||||
|
||||
let auth_snapshot = if let Some(usage) = usage.as_ref() {
|
||||
match (
|
||||
usage.usage.user_id.as_deref(),
|
||||
usage.usage.api_key_id.as_deref(),
|
||||
) {
|
||||
(Some(user_id), Some(api_key_id)) => {
|
||||
state
|
||||
.read_auth_api_key_snapshot(user_id, api_key_id, now_unix_secs)
|
||||
.await?
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
if usage.is_none() && decision_trace.is_none() && auth_snapshot.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(RequestAuditBundle {
|
||||
request_id: request_id.to_string(),
|
||||
usage,
|
||||
decision_trace,
|
||||
auth_snapshot,
|
||||
}))
|
||||
}
|
||||
454
crates/aether-gateway/src/data/state.rs
Normal file
454
crates/aether-gateway/src/data/state.rs
Normal file
@@ -0,0 +1,454 @@
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_data::repository::auth::{
|
||||
AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::candidates::{RequestCandidateReadRepository, StoredRequestCandidate};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_data::repository::shadow_results::{
|
||||
merge_shadow_result_sample, RecordShadowResultSample, ShadowResultLookupKey,
|
||||
ShadowResultReadRepository, ShadowResultWriteRepository, StoredShadowResult,
|
||||
};
|
||||
use aether_data::repository::usage::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
use aether_data::repository::video_tasks::{
|
||||
StoredVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
|
||||
};
|
||||
use aether_data::{DataBackends, DataLayerError};
|
||||
|
||||
use super::auth::{read_auth_api_key_snapshot, StoredGatewayAuthApiKeySnapshot};
|
||||
use super::candidates::{read_request_candidate_trace, RequestCandidateTrace};
|
||||
use super::config::GatewayDataConfig;
|
||||
use super::decision_trace::{read_decision_trace, DecisionTrace};
|
||||
use super::request_audit::{read_request_audit_bundle, RequestAuditBundle};
|
||||
use super::usage::{read_request_usage_audit, RequestUsageAudit};
|
||||
use super::video_tasks::read_video_task_response;
|
||||
use crate::gateway::video_tasks::LocalVideoTaskReadResponse;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub(crate) struct GatewayDataState {
|
||||
config: GatewayDataConfig,
|
||||
backends: Option<DataBackends>,
|
||||
auth_api_key_reader: Option<Arc<dyn AuthApiKeyReadRepository>>,
|
||||
request_candidate_reader: Option<Arc<dyn RequestCandidateReadRepository>>,
|
||||
provider_catalog_reader: Option<Arc<dyn ProviderCatalogReadRepository>>,
|
||||
usage_reader: Option<Arc<dyn UsageReadRepository>>,
|
||||
video_task_reader: Option<Arc<dyn VideoTaskReadRepository>>,
|
||||
shadow_result_reader: Option<Arc<dyn ShadowResultReadRepository>>,
|
||||
shadow_result_writer: Option<Arc<dyn ShadowResultWriteRepository>>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for GatewayDataState {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("GatewayDataState")
|
||||
.field("config", &self.config)
|
||||
.field("has_backends", &self.backends.is_some())
|
||||
.field(
|
||||
"has_auth_api_key_reader",
|
||||
&self.auth_api_key_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_request_candidate_reader",
|
||||
&self.request_candidate_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_provider_catalog_reader",
|
||||
&self.provider_catalog_reader.is_some(),
|
||||
)
|
||||
.field("has_usage_reader", &self.usage_reader.is_some())
|
||||
.field("has_video_task_reader", &self.video_task_reader.is_some())
|
||||
.field(
|
||||
"has_shadow_result_reader",
|
||||
&self.shadow_result_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_shadow_result_writer",
|
||||
&self.shadow_result_writer.is_some(),
|
||||
)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl GatewayDataState {
|
||||
pub(crate) fn disabled() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub(crate) fn from_config(config: GatewayDataConfig) -> Result<Self, DataLayerError> {
|
||||
if !config.is_enabled() {
|
||||
return Ok(Self {
|
||||
config,
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
request_candidate_reader: None,
|
||||
provider_catalog_reader: None,
|
||||
usage_reader: None,
|
||||
video_task_reader: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
});
|
||||
}
|
||||
|
||||
let backends = DataBackends::from_config(config.to_data_layer_config())?;
|
||||
let auth_api_key_reader = backends.read().auth_api_keys();
|
||||
let request_candidate_reader = backends.read().request_candidates();
|
||||
let provider_catalog_reader = backends.read().provider_catalog();
|
||||
let usage_reader = backends.read().usage();
|
||||
let video_task_reader = backends.read().video_tasks();
|
||||
let shadow_result_reader = backends.read().shadow_results();
|
||||
let shadow_result_writer = backends.write().shadow_results();
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
backends: Some(backends),
|
||||
auth_api_key_reader,
|
||||
request_candidate_reader,
|
||||
provider_catalog_reader,
|
||||
usage_reader,
|
||||
video_task_reader,
|
||||
shadow_result_reader,
|
||||
shadow_result_writer,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn has_backends(&self) -> bool {
|
||||
self.backends.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_auth_api_key_reader(&self) -> bool {
|
||||
self.auth_api_key_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_request_candidate_reader(&self) -> bool {
|
||||
self.request_candidate_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_provider_catalog_reader(&self) -> bool {
|
||||
self.provider_catalog_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_usage_reader(&self) -> bool {
|
||||
self.usage_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_video_task_reader(&self) -> bool {
|
||||
self.video_task_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_shadow_result_writer(&self) -> bool {
|
||||
self.shadow_result_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_shadow_result_reader(&self) -> bool {
|
||||
self.shadow_result_reader.is_some()
|
||||
}
|
||||
|
||||
pub(super) async fn find_video_task(
|
||||
&self,
|
||||
key: VideoTaskLookupKey<'_>,
|
||||
) -> Result<Option<StoredVideoTask>, DataLayerError> {
|
||||
match &self.video_task_reader {
|
||||
Some(repository) => repository.find(key).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn find_auth_api_key_snapshot(
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
|
||||
match &self.auth_api_key_reader {
|
||||
Some(repository) => repository.find_api_key_snapshot(key).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn list_request_candidates_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
|
||||
match &self.request_candidate_reader {
|
||||
Some(repository) => repository.list_by_request_id(request_id).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn list_provider_catalog_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => repository.list_providers_by_ids(provider_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn list_provider_catalog_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => repository.list_endpoints_by_ids(endpoint_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn list_provider_catalog_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
match &self.provider_catalog_reader {
|
||||
Some(repository) => repository.list_keys_by_ids(key_ids).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn find_request_usage_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
match &self.usage_reader {
|
||||
Some(repository) => repository.find_by_request_id(request_id).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_candidate_trace(
|
||||
&self,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
) -> Result<Option<RequestCandidateTrace>, DataLayerError> {
|
||||
read_request_candidate_trace(self, request_id, attempted_only).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_decision_trace(
|
||||
&self,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
) -> Result<Option<DecisionTrace>, DataLayerError> {
|
||||
read_decision_trace(self, request_id, attempted_only).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_usage_audit(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<RequestUsageAudit>, DataLayerError> {
|
||||
read_request_usage_audit(self, request_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_audit_bundle(
|
||||
&self,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<RequestAuditBundle>, DataLayerError> {
|
||||
read_request_audit_bundle(self, request_id, attempted_only, now_unix_secs).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_auth_api_key_snapshot(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<StoredGatewayAuthApiKeySnapshot>, DataLayerError> {
|
||||
read_auth_api_key_snapshot(self, user_id, api_key_id, now_unix_secs).await
|
||||
}
|
||||
|
||||
pub(crate) async fn read_video_task_response(
|
||||
&self,
|
||||
route_family: Option<&str>,
|
||||
request_path: &str,
|
||||
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
|
||||
read_video_task_response(self, route_family, request_path).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) async fn write_shadow_result(
|
||||
&self,
|
||||
result: aether_data::repository::shadow_results::UpsertShadowResult,
|
||||
) -> Result<Option<StoredShadowResult>, DataLayerError> {
|
||||
match &self.shadow_result_writer {
|
||||
Some(repository) => repository.upsert(result).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn record_shadow_result_sample(
|
||||
&self,
|
||||
sample: RecordShadowResultSample,
|
||||
) -> Result<Option<StoredShadowResult>, DataLayerError> {
|
||||
let Some(writer) = &self.shadow_result_writer else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let existing = match &self.shadow_result_reader {
|
||||
Some(reader) => {
|
||||
reader
|
||||
.find(ShadowResultLookupKey::TraceFingerprint {
|
||||
trace_id: &sample.trace_id,
|
||||
request_fingerprint: &sample.request_fingerprint,
|
||||
})
|
||||
.await?
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
let merged = merge_shadow_result_sample(existing.as_ref(), sample);
|
||||
|
||||
writer.upsert(merged).await.map(Some)
|
||||
}
|
||||
|
||||
pub(crate) async fn list_recent_shadow_results(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredShadowResult>, DataLayerError> {
|
||||
match &self.shadow_result_reader {
|
||||
Some(repository) => repository.list_recent(limit).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_video_task_reader_for_tests(
|
||||
repository: Arc<dyn VideoTaskReadRepository>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
request_candidate_reader: None,
|
||||
provider_catalog_reader: None,
|
||||
usage_reader: None,
|
||||
video_task_reader: Some(repository),
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_request_candidate_reader_for_tests(
|
||||
repository: Arc<dyn RequestCandidateReadRepository>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
request_candidate_reader: Some(repository),
|
||||
provider_catalog_reader: None,
|
||||
usage_reader: None,
|
||||
video_task_reader: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_usage_reader_for_tests(repository: Arc<dyn UsageReadRepository>) -> Self {
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
request_candidate_reader: None,
|
||||
provider_catalog_reader: None,
|
||||
usage_reader: Some(repository),
|
||||
video_task_reader: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_auth_api_key_reader_for_tests(
|
||||
repository: Arc<dyn AuthApiKeyReadRepository>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: Some(repository),
|
||||
request_candidate_reader: None,
|
||||
provider_catalog_reader: None,
|
||||
usage_reader: None,
|
||||
video_task_reader: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_decision_trace_readers_for_tests(
|
||||
request_candidate_repository: Arc<dyn RequestCandidateReadRepository>,
|
||||
provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
request_candidate_reader: Some(request_candidate_repository),
|
||||
provider_catalog_reader: Some(provider_catalog_repository),
|
||||
usage_reader: None,
|
||||
video_task_reader: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_request_audit_readers_for_tests(
|
||||
auth_api_key_repository: Arc<dyn AuthApiKeyReadRepository>,
|
||||
request_candidate_repository: Arc<dyn RequestCandidateReadRepository>,
|
||||
provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository>,
|
||||
usage_repository: Arc<dyn UsageReadRepository>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: Some(auth_api_key_repository),
|
||||
request_candidate_reader: Some(request_candidate_repository),
|
||||
provider_catalog_reader: Some(provider_catalog_repository),
|
||||
usage_reader: Some(usage_repository),
|
||||
video_task_reader: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_shadow_result_writer_for_tests(
|
||||
repository: Arc<dyn ShadowResultWriteRepository>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
request_candidate_reader: None,
|
||||
provider_catalog_reader: None,
|
||||
usage_reader: None,
|
||||
video_task_reader: None,
|
||||
shadow_result_reader: None,
|
||||
shadow_result_writer: Some(repository),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_shadow_result_repository_for_tests<T>(repository: Arc<T>) -> Self
|
||||
where
|
||||
T: aether_data::repository::shadow_results::ShadowResultRepository + 'static,
|
||||
{
|
||||
let shadow_result_reader: Arc<dyn ShadowResultReadRepository> = repository.clone();
|
||||
let shadow_result_writer: Arc<dyn ShadowResultWriteRepository> = repository;
|
||||
|
||||
Self {
|
||||
config: GatewayDataConfig::disabled(),
|
||||
backends: None,
|
||||
auth_api_key_reader: None,
|
||||
request_candidate_reader: None,
|
||||
provider_catalog_reader: None,
|
||||
usage_reader: None,
|
||||
video_task_reader: None,
|
||||
shadow_result_reader: Some(shadow_result_reader),
|
||||
shadow_result_writer: Some(shadow_result_writer),
|
||||
}
|
||||
}
|
||||
}
|
||||
683
crates/aether-gateway/src/data/tests.rs
Normal file
683
crates/aether-gateway/src/data/tests.rs
Normal file
@@ -0,0 +1,683 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::candidates::{
|
||||
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_data::repository::shadow_results::{
|
||||
InMemoryShadowResultRepository, RecordShadowResultSample, ShadowResultLookupKey,
|
||||
ShadowResultMatchStatus, ShadowResultReadRepository, ShadowResultSampleOrigin,
|
||||
UpsertShadowResult,
|
||||
};
|
||||
use aether_data::repository::usage::{InMemoryUsageReadRepository, StoredRequestUsageAudit};
|
||||
use aether_data::repository::video_tasks::{
|
||||
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskLookupKey, VideoTaskStatus,
|
||||
VideoTaskWriteRepository,
|
||||
};
|
||||
|
||||
use super::{GatewayDataConfig, GatewayDataState};
|
||||
use crate::gateway::AppState;
|
||||
|
||||
#[test]
|
||||
fn disabled_gateway_data_state_has_no_backends() {
|
||||
let state = GatewayDataState::from_config(GatewayDataConfig::disabled())
|
||||
.expect("disabled config should build");
|
||||
|
||||
assert!(!state.has_backends());
|
||||
assert!(!state.has_auth_api_key_reader());
|
||||
assert!(!state.has_request_candidate_reader());
|
||||
assert!(!state.has_provider_catalog_reader());
|
||||
assert!(!state.has_usage_reader());
|
||||
assert!(!state.has_video_task_reader());
|
||||
assert!(!state.has_shadow_result_reader());
|
||||
assert!(!state.has_shadow_result_writer());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn postgres_gateway_data_state_builds_video_task_reader() {
|
||||
let state = GatewayDataState::from_config(GatewayDataConfig::from_postgres_url(
|
||||
"postgres://localhost/aether",
|
||||
false,
|
||||
))
|
||||
.expect("postgres-backed state should build");
|
||||
|
||||
assert!(state.has_backends());
|
||||
assert!(state.has_auth_api_key_reader());
|
||||
assert!(state.has_request_candidate_reader());
|
||||
assert!(state.has_provider_catalog_reader());
|
||||
assert!(state.has_usage_reader());
|
||||
assert!(state.has_video_task_reader());
|
||||
assert!(state.has_shadow_result_reader());
|
||||
assert!(state.has_shadow_result_writer());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_state_find_uses_configured_read_repository() {
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
repository
|
||||
.upsert(UpsertVideoTask {
|
||||
id: "task-1".to_string(),
|
||||
short_id: Some("short-task-1".to_string()),
|
||||
user_id: Some("user-1".to_string()),
|
||||
external_task_id: Some("ext-task-1".to_string()),
|
||||
provider_api_format: Some("openai:video".to_string()),
|
||||
model: Some("sora-2".to_string()),
|
||||
prompt: Some("hello".to_string()),
|
||||
size: Some("1280x720".to_string()),
|
||||
status: VideoTaskStatus::Queued,
|
||||
progress_percent: 0,
|
||||
created_at_unix_secs: 100,
|
||||
updated_at_unix_secs: 100,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let state = GatewayDataState::with_video_task_reader_for_tests(repository);
|
||||
|
||||
let task = state
|
||||
.find_video_task(VideoTaskLookupKey::Id("task-1"))
|
||||
.await
|
||||
.expect("find should succeed");
|
||||
|
||||
assert_eq!(task.expect("task should exist").id, "task-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn app_state_wires_gateway_data_state_from_config() {
|
||||
let state = AppState::new_with_executor(
|
||||
"http://127.0.0.1:18084",
|
||||
Some("http://127.0.0.1:18085".to_string()),
|
||||
Some("http://127.0.0.1:18086".to_string()),
|
||||
)
|
||||
.expect("app state should build")
|
||||
.with_data_config(GatewayDataConfig::from_postgres_url(
|
||||
"postgres://localhost/aether",
|
||||
false,
|
||||
))
|
||||
.expect("data config should wire");
|
||||
|
||||
assert!(state.data.has_backends());
|
||||
assert!(state.data.has_auth_api_key_reader());
|
||||
assert!(state.data.has_request_candidate_reader());
|
||||
assert!(state.data.has_provider_catalog_reader());
|
||||
assert!(state.data.has_usage_reader());
|
||||
assert!(state.data.has_video_task_reader());
|
||||
assert!(state.data.has_shadow_result_reader());
|
||||
assert!(state.data.has_shadow_result_writer());
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(200),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_state_reads_auth_api_key_snapshot_from_reader() {
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_auth_snapshot("key-1", "user-1"),
|
||||
)]));
|
||||
let state = GatewayDataState::with_auth_api_key_reader_for_tests(repository);
|
||||
|
||||
let snapshot = state
|
||||
.read_auth_api_key_snapshot("user-1", "key-1", 150)
|
||||
.await
|
||||
.expect("read should succeed")
|
||||
.expect("snapshot should exist");
|
||||
|
||||
assert_eq!(snapshot.user_id, "user-1");
|
||||
assert_eq!(snapshot.api_key_id, "key-1");
|
||||
assert_eq!(snapshot.username, "alice");
|
||||
assert_eq!(
|
||||
snapshot.api_key_allowed_models,
|
||||
Some(vec!["gpt-4.1".to_string()])
|
||||
);
|
||||
assert!(snapshot.currently_usable);
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-1".to_string(),
|
||||
"OpenAI".to_string(),
|
||||
Some("https://openai.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"openai:chat".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"provider-key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"prod-key".to_string(),
|
||||
"api_key".to_string(),
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
}
|
||||
|
||||
fn sample_request_usage(request_id: &str) -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
"OpenAI".to_string(),
|
||||
"gpt-4.1".to_string(),
|
||||
Some("gpt-4.1-mini".to_string()),
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
false,
|
||||
120,
|
||||
40,
|
||||
160,
|
||||
0.24,
|
||||
0.36,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(450),
|
||||
Some(120),
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
100,
|
||||
101,
|
||||
Some(102),
|
||||
)
|
||||
.expect("usage should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_state_reads_decision_trace_with_provider_catalog_metadata() {
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
StoredRequestCandidate::new(
|
||||
"cand-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
0,
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
RequestCandidateStatus::Failed,
|
||||
None,
|
||||
false,
|
||||
Some(502),
|
||||
None,
|
||||
None,
|
||||
Some(37),
|
||||
Some(1),
|
||||
None,
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
100,
|
||||
Some(101),
|
||||
Some(102),
|
||||
)
|
||||
.expect("candidate should build"),
|
||||
]));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
let state = GatewayDataState::with_decision_trace_readers_for_tests(
|
||||
request_candidates,
|
||||
provider_catalog,
|
||||
);
|
||||
|
||||
let trace = state
|
||||
.read_decision_trace("req-1", true)
|
||||
.await
|
||||
.expect("trace should read")
|
||||
.expect("trace should exist");
|
||||
|
||||
assert_eq!(trace.request_id, "req-1");
|
||||
assert_eq!(trace.total_candidates, 1);
|
||||
assert_eq!(trace.candidates[0].provider_name.as_deref(), Some("OpenAI"));
|
||||
assert_eq!(
|
||||
trace.candidates[0].endpoint_api_format.as_deref(),
|
||||
Some("openai:chat")
|
||||
);
|
||||
assert_eq!(
|
||||
trace.candidates[0].provider_key_auth_type.as_deref(),
|
||||
Some("api_key")
|
||||
);
|
||||
assert_eq!(
|
||||
trace.candidates[0].provider_key_capabilities,
|
||||
Some(serde_json::json!({"cache_1h": true}))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_state_reads_request_usage_audit_from_reader() {
|
||||
let repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_request_usage("req-usage-1"),
|
||||
]));
|
||||
let state = GatewayDataState::with_usage_reader_for_tests(repository);
|
||||
|
||||
let usage = state
|
||||
.read_request_usage_audit("req-usage-1")
|
||||
.await
|
||||
.expect("read should succeed")
|
||||
.expect("usage should exist");
|
||||
|
||||
assert_eq!(usage.usage.request_id, "req-usage-1");
|
||||
assert_eq!(usage.usage.provider_name, "OpenAI");
|
||||
assert_eq!(usage.usage.total_tokens, 160);
|
||||
assert_eq!(usage.usage.total_cost_usd, 0.24);
|
||||
assert!(usage.usage.has_format_conversion);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_state_reads_request_audit_bundle_from_multiple_readers() {
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_auth_snapshot("api-key-1", "user-1"),
|
||||
)]));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
StoredRequestCandidate::new(
|
||||
"cand-1".to_string(),
|
||||
"req-usage-1".to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
0,
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
RequestCandidateStatus::Success,
|
||||
None,
|
||||
false,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(37),
|
||||
Some(1),
|
||||
None,
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
100,
|
||||
Some(101),
|
||||
Some(102),
|
||||
)
|
||||
.expect("candidate should build"),
|
||||
]));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_request_usage("req-usage-1"),
|
||||
]));
|
||||
let state = GatewayDataState::with_request_audit_readers_for_tests(
|
||||
auth_repository,
|
||||
request_candidates,
|
||||
provider_catalog,
|
||||
usage_repository,
|
||||
);
|
||||
|
||||
let bundle = state
|
||||
.read_request_audit_bundle("req-usage-1", true, 150)
|
||||
.await
|
||||
.expect("bundle should read")
|
||||
.expect("bundle should exist");
|
||||
|
||||
assert_eq!(bundle.request_id, "req-usage-1");
|
||||
assert_eq!(
|
||||
bundle
|
||||
.usage
|
||||
.as_ref()
|
||||
.and_then(|usage| usage.usage.target_model.as_deref()),
|
||||
Some("gpt-4.1-mini")
|
||||
);
|
||||
assert_eq!(
|
||||
bundle
|
||||
.decision_trace
|
||||
.as_ref()
|
||||
.and_then(|trace| trace.candidates.first())
|
||||
.and_then(|candidate| candidate.provider_name.as_deref()),
|
||||
Some("OpenAI")
|
||||
);
|
||||
assert_eq!(
|
||||
bundle
|
||||
.auth_snapshot
|
||||
.as_ref()
|
||||
.map(|snapshot| snapshot.currently_usable),
|
||||
Some(true)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn maps_openai_video_task_repository_row_into_read_response() {
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
repository
|
||||
.upsert(UpsertVideoTask {
|
||||
id: "task-1".to_string(),
|
||||
short_id: Some("short-task-1".to_string()),
|
||||
user_id: Some("user-1".to_string()),
|
||||
external_task_id: Some("ext-task-1".to_string()),
|
||||
provider_api_format: Some("openai:video".to_string()),
|
||||
model: Some("sora-2".to_string()),
|
||||
prompt: Some("hello".to_string()),
|
||||
size: Some("1280x720".to_string()),
|
||||
status: VideoTaskStatus::Processing,
|
||||
progress_percent: 45,
|
||||
created_at_unix_secs: 100,
|
||||
updated_at_unix_secs: 120,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let state = GatewayDataState::with_video_task_reader_for_tests(repository);
|
||||
let response = state
|
||||
.read_video_task_response(Some("openai"), "/v1/videos/task-1")
|
||||
.await
|
||||
.expect("read should succeed")
|
||||
.expect("read response should exist");
|
||||
|
||||
assert_eq!(response.status_code, 200);
|
||||
assert_eq!(response.body_json["id"], "task-1");
|
||||
assert_eq!(response.body_json["status"], "processing");
|
||||
assert_eq!(response.body_json["created_at"], 100);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn maps_gemini_video_task_repository_row_into_read_response() {
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
repository
|
||||
.upsert(UpsertVideoTask {
|
||||
id: "task-1".to_string(),
|
||||
short_id: Some("localshort123".to_string()),
|
||||
user_id: Some("user-1".to_string()),
|
||||
external_task_id: Some("operations/ext-task-1".to_string()),
|
||||
provider_api_format: Some("gemini:video".to_string()),
|
||||
model: Some("veo-3".to_string()),
|
||||
prompt: Some("hello".to_string()),
|
||||
size: Some("720p".to_string()),
|
||||
status: VideoTaskStatus::Completed,
|
||||
progress_percent: 100,
|
||||
created_at_unix_secs: 100,
|
||||
updated_at_unix_secs: 120,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let state = GatewayDataState::with_video_task_reader_for_tests(repository);
|
||||
let response = state
|
||||
.read_video_task_response(
|
||||
Some("gemini"),
|
||||
"/v1beta/models/veo-3/operations/localshort123",
|
||||
)
|
||||
.await
|
||||
.expect("read should succeed")
|
||||
.expect("read response should exist");
|
||||
|
||||
assert_eq!(response.status_code, 200);
|
||||
assert_eq!(
|
||||
response.body_json["name"],
|
||||
"models/veo-3/operations/localshort123"
|
||||
);
|
||||
assert_eq!(response.body_json["done"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_state_write_uses_configured_shadow_result_writer() {
|
||||
let repository = Arc::new(InMemoryShadowResultRepository::default());
|
||||
let state = GatewayDataState::with_shadow_result_writer_for_tests(repository.clone());
|
||||
|
||||
let written = state
|
||||
.write_shadow_result(UpsertShadowResult {
|
||||
trace_id: "trace-1".to_string(),
|
||||
request_fingerprint: "fp-1".to_string(),
|
||||
request_id: Some("req-1".to_string()),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
candidate_id: None,
|
||||
rust_result_digest: Some("rust-digest".to_string()),
|
||||
python_result_digest: None,
|
||||
match_status: ShadowResultMatchStatus::Pending,
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
created_at_unix_secs: 100,
|
||||
updated_at_unix_secs: 100,
|
||||
})
|
||||
.await
|
||||
.expect("write should succeed");
|
||||
|
||||
assert!(written.is_some());
|
||||
let stored = repository
|
||||
.find(ShadowResultLookupKey::TraceFingerprint {
|
||||
trace_id: "trace-1",
|
||||
request_fingerprint: "fp-1",
|
||||
})
|
||||
.await
|
||||
.expect("find should succeed");
|
||||
assert_eq!(
|
||||
stored.expect("stored result should exist").match_status,
|
||||
ShadowResultMatchStatus::Pending
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_state_records_shadow_result_samples_and_merges_match_status() {
|
||||
let repository = Arc::new(InMemoryShadowResultRepository::default());
|
||||
let state = GatewayDataState::with_shadow_result_repository_for_tests(repository);
|
||||
|
||||
let first = state
|
||||
.record_shadow_result_sample(RecordShadowResultSample {
|
||||
trace_id: "trace-1".to_string(),
|
||||
request_fingerprint: "fp-1".to_string(),
|
||||
request_id: Some("req-1".to_string()),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
candidate_id: None,
|
||||
origin: ShadowResultSampleOrigin::Rust,
|
||||
result_digest: "digest-1".to_string(),
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
recorded_at_unix_secs: 100,
|
||||
})
|
||||
.await
|
||||
.expect("first record should succeed")
|
||||
.expect("first stored result should exist");
|
||||
assert_eq!(first.match_status, ShadowResultMatchStatus::Pending);
|
||||
|
||||
let second = state
|
||||
.record_shadow_result_sample(RecordShadowResultSample {
|
||||
trace_id: "trace-1".to_string(),
|
||||
request_fingerprint: "fp-1".to_string(),
|
||||
request_id: Some("req-1".to_string()),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
candidate_id: None,
|
||||
origin: ShadowResultSampleOrigin::Python,
|
||||
result_digest: "digest-1".to_string(),
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
recorded_at_unix_secs: 200,
|
||||
})
|
||||
.await
|
||||
.expect("second record should succeed")
|
||||
.expect("second stored result should exist");
|
||||
|
||||
assert_eq!(second.match_status, ShadowResultMatchStatus::Match);
|
||||
assert_eq!(second.created_at_unix_secs, 100);
|
||||
assert_eq!(second.updated_at_unix_secs, 200);
|
||||
assert_eq!(second.request_id.as_deref(), Some("req-1"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_state_lists_recent_shadow_results_from_reader() {
|
||||
let repository = Arc::new(InMemoryShadowResultRepository::default());
|
||||
let state = GatewayDataState::with_shadow_result_repository_for_tests(repository.clone());
|
||||
|
||||
state
|
||||
.record_shadow_result_sample(RecordShadowResultSample {
|
||||
trace_id: "trace-1".to_string(),
|
||||
request_fingerprint: "fp-1".to_string(),
|
||||
request_id: Some("req-shadow-1".to_string()),
|
||||
route_family: Some("openai".to_string()),
|
||||
route_kind: Some("chat".to_string()),
|
||||
candidate_id: None,
|
||||
origin: ShadowResultSampleOrigin::Rust,
|
||||
result_digest: "digest-1".to_string(),
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
recorded_at_unix_secs: 100,
|
||||
})
|
||||
.await
|
||||
.expect("record should succeed");
|
||||
|
||||
let recent = state
|
||||
.list_recent_shadow_results(5)
|
||||
.await
|
||||
.expect("list recent should succeed");
|
||||
|
||||
assert_eq!(recent.len(), 1);
|
||||
assert_eq!(recent[0].trace_id, "trace-1");
|
||||
assert_eq!(recent[0].request_id.as_deref(), Some("req-shadow-1"));
|
||||
}
|
||||
|
||||
fn sample_request_candidate(
|
||||
id: &str,
|
||||
request_id: &str,
|
||||
candidate_index: i32,
|
||||
status: RequestCandidateStatus,
|
||||
started_at_unix_secs: Option<i64>,
|
||||
latency_ms: Option<i32>,
|
||||
status_code: Option<i32>,
|
||||
) -> StoredRequestCandidate {
|
||||
StoredRequestCandidate::new(
|
||||
id.to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
candidate_index,
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
status,
|
||||
None,
|
||||
false,
|
||||
status_code,
|
||||
None,
|
||||
None,
|
||||
latency_ms,
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
100 + i64::from(candidate_index),
|
||||
started_at_unix_secs,
|
||||
started_at_unix_secs.map(|value| value + 1),
|
||||
)
|
||||
.expect("candidate should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_state_reads_request_candidate_trace_from_reader() {
|
||||
let repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_request_candidate(
|
||||
"cand-1",
|
||||
"req-1",
|
||||
0,
|
||||
RequestCandidateStatus::Pending,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
),
|
||||
sample_request_candidate(
|
||||
"cand-2",
|
||||
"req-1",
|
||||
1,
|
||||
RequestCandidateStatus::Success,
|
||||
Some(101),
|
||||
Some(42),
|
||||
Some(200),
|
||||
),
|
||||
]));
|
||||
let state = GatewayDataState::with_request_candidate_reader_for_tests(repository);
|
||||
|
||||
let trace = state
|
||||
.read_request_candidate_trace("req-1", true)
|
||||
.await
|
||||
.expect("trace should succeed")
|
||||
.expect("trace should exist");
|
||||
|
||||
assert_eq!(trace.request_id, "req-1");
|
||||
assert_eq!(trace.total_candidates, 1);
|
||||
assert_eq!(
|
||||
trace.final_status,
|
||||
super::candidates::RequestCandidateFinalStatus::Success
|
||||
);
|
||||
assert_eq!(trace.total_latency_ms, 42);
|
||||
assert_eq!(trace.candidates[0].id, "cand-2");
|
||||
}
|
||||
20
crates/aether-gateway/src/data/usage.rs
Normal file
20
crates/aether-gateway/src/data/usage.rs
Normal file
@@ -0,0 +1,20 @@
|
||||
use aether_data::repository::usage::StoredRequestUsageAudit;
|
||||
use aether_data::DataLayerError;
|
||||
|
||||
use super::state::GatewayDataState;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
|
||||
pub(crate) struct RequestUsageAudit {
|
||||
#[serde(flatten)]
|
||||
pub(crate) usage: StoredRequestUsageAudit,
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_usage_audit(
|
||||
state: &GatewayDataState,
|
||||
request_id: &str,
|
||||
) -> Result<Option<RequestUsageAudit>, DataLayerError> {
|
||||
Ok(state
|
||||
.find_request_usage_by_request_id(request_id)
|
||||
.await?
|
||||
.map(|usage| RequestUsageAudit { usage }))
|
||||
}
|
||||
65
crates/aether-gateway/src/data/video_tasks.rs
Normal file
65
crates/aether-gateway/src/data/video_tasks.rs
Normal file
@@ -0,0 +1,65 @@
|
||||
use aether_data::repository::video_tasks::VideoTaskLookupKey;
|
||||
use aether_data::DataLayerError;
|
||||
|
||||
use super::gemini::map_gemini_video_task_to_read_response;
|
||||
use super::openai::map_openai_video_task_to_read_response;
|
||||
use super::state::GatewayDataState;
|
||||
use crate::gateway::video_tasks::{
|
||||
extract_gemini_short_id_from_path, extract_openai_task_id_from_path, LocalVideoTaskReadResponse,
|
||||
};
|
||||
|
||||
pub(super) async fn read_video_task_response(
|
||||
state: &GatewayDataState,
|
||||
route_family: Option<&str>,
|
||||
request_path: &str,
|
||||
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
|
||||
match route_family {
|
||||
Some("openai") => read_openai_video_task_response(state, request_path).await,
|
||||
Some("gemini") => read_gemini_video_task_response(state, request_path).await,
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_openai_video_task_response(
|
||||
state: &GatewayDataState,
|
||||
request_path: &str,
|
||||
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
|
||||
let Some(task_id) = extract_openai_task_id_from_path(request_path) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(task) = state
|
||||
.find_video_task(VideoTaskLookupKey::Id(task_id))
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if !matches!(task.provider_api_format.as_deref(), Some("openai:video")) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(map_openai_video_task_to_read_response(task)))
|
||||
}
|
||||
|
||||
async fn read_gemini_video_task_response(
|
||||
state: &GatewayDataState,
|
||||
request_path: &str,
|
||||
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
|
||||
let Some(short_id) = extract_gemini_short_id_from_path(request_path) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(task) = state
|
||||
.find_video_task(VideoTaskLookupKey::ShortId(short_id))
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
if !matches!(task.provider_api_format.as_deref(), Some("gemini:video")) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(map_gemini_video_task_to_read_response(task)))
|
||||
}
|
||||
@@ -2,7 +2,7 @@ use std::collections::hash_map::DefaultHasher;
|
||||
use std::collections::BTreeMap;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::io::Error as IoError;
|
||||
use std::time::{Duration, Instant};
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_contracts::{
|
||||
ExecutionPlan, ExecutionResult, ExecutionTelemetry, ExecutionTimeouts, ProxySnapshot,
|
||||
@@ -24,8 +24,8 @@ use crate::gateway::headers::{
|
||||
should_skip_upstream_passthrough_header,
|
||||
};
|
||||
use crate::gateway::{
|
||||
build_client_response, build_client_response_from_parts, cache_executor_auth_context,
|
||||
local_finalize::maybe_build_local_core_sync_finalize_response,
|
||||
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
|
||||
cache_executor_auth_context, local_finalize::maybe_build_local_core_sync_finalize_response,
|
||||
local_stream::maybe_build_local_stream_rewriter, resolve_executor_auth_context, AppState,
|
||||
GatewayControlAuthContext, GatewayControlDecision, GatewayError,
|
||||
};
|
||||
|
||||
@@ -161,34 +161,17 @@ pub(crate) fn build_direct_plan_bypass_cache_key(
|
||||
}
|
||||
|
||||
pub(crate) fn should_skip_direct_plan(state: &AppState, cache_key: &str) -> bool {
|
||||
let Ok(mut cache) = state.direct_plan_bypass_cache.lock() else {
|
||||
return false;
|
||||
};
|
||||
let Some(cached_at) = cache.get(cache_key).copied() else {
|
||||
return false;
|
||||
};
|
||||
if cached_at.elapsed() > DIRECT_PLAN_BYPASS_TTL {
|
||||
cache.remove(cache_key);
|
||||
return false;
|
||||
}
|
||||
true
|
||||
state
|
||||
.direct_plan_bypass_cache
|
||||
.should_skip(cache_key, DIRECT_PLAN_BYPASS_TTL)
|
||||
}
|
||||
|
||||
pub(crate) fn mark_direct_plan_bypass(state: &AppState, cache_key: String) {
|
||||
let Ok(mut cache) = state.direct_plan_bypass_cache.lock() else {
|
||||
return;
|
||||
};
|
||||
cache.retain(|_, cached_at| cached_at.elapsed() <= DIRECT_PLAN_BYPASS_TTL);
|
||||
if cache.len() >= DIRECT_PLAN_BYPASS_MAX_ENTRIES {
|
||||
if let Some(oldest_key) = cache
|
||||
.iter()
|
||||
.min_by_key(|(_, cached_at)| *cached_at)
|
||||
.map(|(key, _)| key.clone())
|
||||
{
|
||||
cache.remove(&oldest_key);
|
||||
}
|
||||
}
|
||||
cache.insert(cache_key, Instant::now());
|
||||
state.direct_plan_bypass_cache.mark(
|
||||
cache_key,
|
||||
DIRECT_PLAN_BYPASS_TTL,
|
||||
DIRECT_PLAN_BYPASS_MAX_ENTRIES,
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_direct_executor_stream_plan_kind(
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use base64::Engine as _;
|
||||
use futures_util::TryStreamExt;
|
||||
use tracing::debug;
|
||||
|
||||
use super::super::submission::{
|
||||
maybe_build_local_core_error_response, resolve_core_error_background_report_kind,
|
||||
@@ -26,6 +27,8 @@ pub(super) async fn execute_executor_stream(
|
||||
report_kind: Option<String>,
|
||||
report_context: Option<serde_json::Value>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let request_id = plan.request_id.as_str();
|
||||
let candidate_id = plan.candidate_id.as_deref();
|
||||
let response = match state
|
||||
.client
|
||||
.post(format!("{executor_base_url}/v1/execute/stream"))
|
||||
@@ -42,10 +45,10 @@ pub(super) async fn execute_executor_stream(
|
||||
};
|
||||
|
||||
if response.status() != http::StatusCode::OK {
|
||||
return Ok(Some(build_client_response(
|
||||
response,
|
||||
trace_id,
|
||||
Some(decision),
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
build_client_response(response, trace_id, Some(decision))?,
|
||||
Some(request_id),
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
|
||||
@@ -116,22 +119,30 @@ pub(super) async fn execute_executor_stream(
|
||||
payload,
|
||||
);
|
||||
}
|
||||
return Ok(Some(response));
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
response,
|
||||
Some(request_id),
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
let response = submit_sync_finalize(state, control_base_url, trace_id, payload).await?;
|
||||
return Ok(Some(build_client_response(
|
||||
response,
|
||||
trace_id,
|
||||
Some(decision),
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
build_client_response(response, trace_id, Some(decision))?,
|
||||
Some(request_id),
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
return Ok(Some(build_executor_error_response(
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
status_code,
|
||||
headers,
|
||||
error_body,
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
build_executor_error_response(
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
status_code,
|
||||
headers,
|
||||
error_body,
|
||||
)?,
|
||||
Some(request_id),
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
|
||||
@@ -216,15 +227,19 @@ pub(super) async fn execute_executor_stream(
|
||||
payload,
|
||||
);
|
||||
}
|
||||
return Ok(Some(response));
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
response,
|
||||
Some(request_id),
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
let response =
|
||||
submit_sync_finalize(state, control_base_url, trace_id, payload)
|
||||
.await?;
|
||||
return Ok(Some(build_client_response(
|
||||
response,
|
||||
trace_id,
|
||||
Some(decision),
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
build_client_response(response, trace_id, Some(decision))?,
|
||||
Some(request_id),
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
StreamPrefetchInspection::NeedMore => {}
|
||||
@@ -278,6 +293,7 @@ pub(super) async fn execute_executor_stream(
|
||||
let mut buffered_body = prefetched_body_for_report;
|
||||
let mut telemetry: Option<ExecutionTelemetry> = initial_telemetry;
|
||||
let reached_eof = initial_reached_eof;
|
||||
let mut downstream_dropped = false;
|
||||
|
||||
if !reached_eof {
|
||||
loop {
|
||||
@@ -334,6 +350,7 @@ pub(super) async fn execute_executor_stream(
|
||||
trace_id = %trace_id_owned,
|
||||
"gateway stream downstream dropped; stopping executor stream forwarding"
|
||||
);
|
||||
downstream_dropped = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -354,7 +371,12 @@ pub(super) async fn execute_executor_stream(
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(rewriter) = local_stream_rewriter.as_mut() {
|
||||
if downstream_dropped {
|
||||
debug!(
|
||||
trace_id = %trace_id_owned,
|
||||
"gateway skipped local stream flush after downstream disconnect"
|
||||
);
|
||||
} else if let Some(rewriter) = local_stream_rewriter.as_mut() {
|
||||
match rewriter.finish() {
|
||||
Ok(flushed_chunk) if !flushed_chunk.is_empty() => {
|
||||
buffered_body.extend_from_slice(&flushed_chunk);
|
||||
@@ -363,6 +385,7 @@ pub(super) async fn execute_executor_stream(
|
||||
trace_id = %trace_id_owned,
|
||||
"gateway stream downstream dropped while flushing local stream rewrite"
|
||||
);
|
||||
downstream_dropped = true;
|
||||
}
|
||||
}
|
||||
Ok(_) => {}
|
||||
@@ -374,6 +397,14 @@ pub(super) async fn execute_executor_stream(
|
||||
|
||||
drop(tx);
|
||||
|
||||
if downstream_dropped {
|
||||
debug!(
|
||||
trace_id = %trace_id_owned,
|
||||
"gateway skipped stream report because downstream disconnected before completion"
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(report_kind) = report_kind_owned {
|
||||
let report = GatewayStreamReportRequest {
|
||||
trace_id: trace_id_owned.clone(),
|
||||
@@ -407,6 +438,21 @@ pub(super) async fn execute_executor_stream(
|
||||
}
|
||||
};
|
||||
|
||||
headers.insert(
|
||||
CONTROL_REQUEST_ID_HEADER.to_string(),
|
||||
request_id.to_string(),
|
||||
);
|
||||
|
||||
if let Some(candidate_id) = candidate_id
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
headers.insert(
|
||||
CONTROL_CANDIDATE_ID_HEADER.to_string(),
|
||||
candidate_id.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
Ok(Some(build_client_response_from_parts(
|
||||
status_code,
|
||||
&headers,
|
||||
|
||||
@@ -251,6 +251,17 @@ async fn maybe_build_local_video_task_read_response(
|
||||
let read_response = state
|
||||
.video_tasks
|
||||
.read_response(decision.route_family.as_deref(), parts.uri.path());
|
||||
let read_response = match read_response {
|
||||
Some(read_response) => Some(read_response),
|
||||
None => {
|
||||
state
|
||||
.read_data_backed_video_task_response(
|
||||
decision.route_family.as_deref(),
|
||||
parts.uri.path(),
|
||||
)
|
||||
.await?
|
||||
}
|
||||
};
|
||||
let Some(read_response) = read_response else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
@@ -39,6 +39,8 @@ pub(super) async fn execute_executor_sync(
|
||||
report_kind: Option<String>,
|
||||
report_context: Option<serde_json::Value>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let plan_request_id = plan.request_id.as_str();
|
||||
let plan_candidate_id = plan.candidate_id.as_deref();
|
||||
let response = match state
|
||||
.client
|
||||
.post(format!("{executor_base_url}/v1/execute/sync"))
|
||||
@@ -55,10 +57,10 @@ pub(super) async fn execute_executor_sync(
|
||||
};
|
||||
|
||||
if response.status() != http::StatusCode::OK {
|
||||
return Ok(Some(build_client_response(
|
||||
response,
|
||||
trace_id,
|
||||
Some(decision),
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
build_client_response(response, trace_id, Some(decision))?,
|
||||
Some(plan_request_id),
|
||||
plan_candidate_id,
|
||||
)?));
|
||||
}
|
||||
|
||||
@@ -66,6 +68,10 @@ pub(super) async fn execute_executor_sync(
|
||||
.json()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let request_id = (!result.request_id.trim().is_empty())
|
||||
.then_some(result.request_id.as_str())
|
||||
.or(Some(plan_request_id));
|
||||
let candidate_id = result.candidate_id.as_deref().or(plan_candidate_id);
|
||||
let mut headers = result.headers.clone();
|
||||
let (body_bytes, body_json, body_base64) = decode_execution_result_body(&result, &mut headers)?;
|
||||
let has_body_bytes = body_base64.is_some();
|
||||
@@ -119,7 +125,11 @@ pub(super) async fn execute_executor_sync(
|
||||
payload,
|
||||
);
|
||||
}
|
||||
return Ok(Some(outcome.response));
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
outcome.response,
|
||||
request_id,
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
if let Some(outcome) = maybe_build_local_video_success_outcome(
|
||||
trace_id,
|
||||
@@ -145,7 +155,11 @@ pub(super) async fn execute_executor_sync(
|
||||
);
|
||||
}
|
||||
}
|
||||
return Ok(Some(outcome.response));
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
outcome.response,
|
||||
request_id,
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
if let Some(response) =
|
||||
maybe_build_local_sync_finalize_response(trace_id, decision, &payload)?
|
||||
@@ -172,7 +186,11 @@ pub(super) async fn execute_executor_sync(
|
||||
payload,
|
||||
);
|
||||
}
|
||||
return Ok(Some(response));
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
response,
|
||||
request_id,
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
if let Some(response) =
|
||||
maybe_build_local_video_error_response(trace_id, decision, &payload)?
|
||||
@@ -196,7 +214,11 @@ pub(super) async fn execute_executor_sync(
|
||||
payload,
|
||||
);
|
||||
}
|
||||
return Ok(Some(response));
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
response,
|
||||
request_id,
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
if let Some(response) = maybe_build_local_core_error_response(trace_id, decision, &payload)?
|
||||
{
|
||||
@@ -219,13 +241,17 @@ pub(super) async fn execute_executor_sync(
|
||||
payload,
|
||||
);
|
||||
}
|
||||
return Ok(Some(response));
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
response,
|
||||
request_id,
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
let response = submit_sync_finalize(state, control_base_url, trace_id, payload).await?;
|
||||
return Ok(Some(build_client_response(
|
||||
response,
|
||||
trace_id,
|
||||
Some(decision),
|
||||
return Ok(Some(attach_control_metadata_headers(
|
||||
build_client_response(response, trace_id, Some(decision))?,
|
||||
request_id,
|
||||
candidate_id,
|
||||
)?));
|
||||
}
|
||||
|
||||
@@ -249,6 +275,23 @@ pub(super) async fn execute_executor_sync(
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(request_id) = request_id.map(str::trim).filter(|value| !value.is_empty()) {
|
||||
headers.insert(
|
||||
CONTROL_REQUEST_ID_HEADER.to_string(),
|
||||
request_id.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(candidate_id) = candidate_id
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
headers.insert(
|
||||
CONTROL_CANDIDATE_ID_HEADER.to_string(),
|
||||
candidate_id.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
Ok(Some(build_client_response_from_parts(
|
||||
result.status_code,
|
||||
&headers,
|
||||
|
||||
@@ -1,7 +1,13 @@
|
||||
#[path = "audit/mod.rs"]
|
||||
mod audit;
|
||||
#[path = "cache/mod.rs"]
|
||||
mod cache;
|
||||
#[path = "constants.rs"]
|
||||
mod constants;
|
||||
#[path = "control.rs"]
|
||||
mod control;
|
||||
#[path = "data/mod.rs"]
|
||||
mod data;
|
||||
#[path = "error.rs"]
|
||||
mod error;
|
||||
#[path = "executor.rs"]
|
||||
@@ -22,42 +28,57 @@ mod response;
|
||||
mod video_tasks;
|
||||
|
||||
use aether_contracts::ExecutionResult;
|
||||
use aether_http::{build_http_client, HttpClientConfig};
|
||||
use aether_runtime::{
|
||||
prometheus_response, service_up_sample, AdmissionPermit, ConcurrencyError, ConcurrencyGate,
|
||||
ConcurrencySnapshot, DistributedConcurrencyError, DistributedConcurrencyGate,
|
||||
DistributedConcurrencySnapshot, MetricKind, MetricLabel, MetricSample,
|
||||
};
|
||||
use axum::http::header::{HeaderName, HeaderValue};
|
||||
use axum::routing::{any, get};
|
||||
use axum::Router;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::warn;
|
||||
|
||||
use cache::{AuthContextCache, DirectPlanBypassCache};
|
||||
|
||||
pub(crate) use audit::record_shadow_result_non_blocking;
|
||||
use audit::{
|
||||
get_auth_api_key_snapshot, get_decision_trace, get_request_audit_bundle,
|
||||
get_request_candidate_trace, get_request_usage_audit, list_recent_shadow_results,
|
||||
};
|
||||
pub(crate) use control::{
|
||||
cache_executor_auth_context, maybe_execute_via_control, resolve_control_route,
|
||||
resolve_executor_auth_context, GatewayControlAuthContext, GatewayControlDecision,
|
||||
resolve_executor_auth_context, trusted_auth_local_rejection, GatewayControlAuthContext,
|
||||
GatewayControlDecision, GatewayLocalAuthRejection,
|
||||
};
|
||||
pub use data::GatewayDataConfig;
|
||||
use data::GatewayDataState;
|
||||
pub(crate) use error::GatewayError;
|
||||
pub(crate) use executor::{maybe_execute_via_executor_stream, maybe_execute_via_executor_sync};
|
||||
use handlers::{health, proxy_request};
|
||||
pub(crate) use response::{build_client_response, build_client_response_from_parts};
|
||||
pub(crate) use response::{
|
||||
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
|
||||
build_local_auth_rejection_response, build_local_overloaded_response,
|
||||
};
|
||||
pub(crate) use video_tasks::VideoTaskService;
|
||||
pub use video_tasks::VideoTaskTruthSourceMode;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct CachedAuthContextEntry {
|
||||
pub(crate) auth_context: GatewayControlAuthContext,
|
||||
pub(crate) cached_at: Instant,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AppState {
|
||||
upstream_base_url: String,
|
||||
control_base_url: Option<String>,
|
||||
executor_base_url: Option<String>,
|
||||
data: Arc<GatewayDataState>,
|
||||
video_tasks: Arc<VideoTaskService>,
|
||||
video_task_poller: Option<VideoTaskPollerConfig>,
|
||||
request_gate: Option<Arc<ConcurrencyGate>>,
|
||||
distributed_request_gate: Option<Arc<DistributedConcurrencyGate>>,
|
||||
client: reqwest::Client,
|
||||
auth_context_cache: Arc<Mutex<HashMap<String, CachedAuthContextEntry>>>,
|
||||
direct_plan_bypass_cache: Arc<Mutex<HashMap<String, Instant>>>,
|
||||
auth_context_cache: Arc<AuthContextCache>,
|
||||
direct_plan_bypass_cache: Arc<DirectPlanBypassCache>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
@@ -79,11 +100,12 @@ impl AppState {
|
||||
control_base_url: Option<String>,
|
||||
executor_base_url: Option<String>,
|
||||
) -> Result<Self, reqwest::Error> {
|
||||
let client = reqwest::Client::builder()
|
||||
.http2_adaptive_window(true)
|
||||
.connect_timeout(std::time::Duration::from_secs(10))
|
||||
.timeout(std::time::Duration::from_secs(300))
|
||||
.build()?;
|
||||
let client = build_http_client(&HttpClientConfig {
|
||||
connect_timeout_ms: Some(10_000),
|
||||
request_timeout_ms: Some(300_000),
|
||||
http2_adaptive_window: true,
|
||||
..HttpClientConfig::default()
|
||||
})?;
|
||||
Ok(Self {
|
||||
upstream_base_url: normalize_upstream_base_url(upstream_base_url.into()),
|
||||
control_base_url: control_base_url
|
||||
@@ -92,16 +114,27 @@ impl AppState {
|
||||
executor_base_url: executor_base_url
|
||||
.map(normalize_upstream_base_url)
|
||||
.filter(|value| !value.is_empty()),
|
||||
data: Arc::new(GatewayDataState::disabled()),
|
||||
video_tasks: Arc::new(VideoTaskService::new(
|
||||
VideoTaskTruthSourceMode::PythonSyncReport,
|
||||
)),
|
||||
video_task_poller: None,
|
||||
request_gate: None,
|
||||
distributed_request_gate: None,
|
||||
client,
|
||||
auth_context_cache: Arc::new(Mutex::new(HashMap::new())),
|
||||
direct_plan_bypass_cache: Arc::new(Mutex::new(HashMap::new())),
|
||||
auth_context_cache: Arc::new(AuthContextCache::default()),
|
||||
direct_plan_bypass_cache: Arc::new(DirectPlanBypassCache::default()),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn with_data_config(
|
||||
mut self,
|
||||
config: GatewayDataConfig,
|
||||
) -> Result<Self, aether_data::DataLayerError> {
|
||||
self.data = Arc::new(GatewayDataState::from_config(config)?);
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
pub fn with_video_task_truth_source_mode(mut self, mode: VideoTaskTruthSourceMode) -> Self {
|
||||
self.video_tasks = Arc::new(VideoTaskService::new(mode));
|
||||
self
|
||||
@@ -115,6 +148,308 @@ impl AppState {
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_request_concurrency_limit(mut self, limit: usize) -> Self {
|
||||
self.request_gate = Some(Arc::new(ConcurrencyGate::new(
|
||||
"gateway_requests",
|
||||
limit.max(1),
|
||||
)));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_distributed_request_concurrency_gate(
|
||||
mut self,
|
||||
gate: DistributedConcurrencyGate,
|
||||
) -> Self {
|
||||
self.distributed_request_gate = Some(Arc::new(gate));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn has_data_backends(&self) -> bool {
|
||||
self.data.has_backends()
|
||||
}
|
||||
|
||||
pub(crate) fn request_concurrency_snapshot(&self) -> Option<ConcurrencySnapshot> {
|
||||
self.request_gate.as_ref().map(|gate| gate.snapshot())
|
||||
}
|
||||
|
||||
pub(crate) async fn distributed_request_concurrency_snapshot(
|
||||
&self,
|
||||
) -> Result<Option<DistributedConcurrencySnapshot>, DistributedConcurrencyError> {
|
||||
match self.distributed_request_gate.as_ref() {
|
||||
Some(gate) => gate.snapshot().await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn metric_samples(&self) -> Vec<MetricSample> {
|
||||
let mut samples = vec![service_up_sample("aether-gateway")];
|
||||
if let Some(snapshot) = self.request_concurrency_snapshot() {
|
||||
samples.extend(snapshot.to_metric_samples("gateway_requests"));
|
||||
}
|
||||
if let Some(gate) = self.distributed_request_gate.as_ref() {
|
||||
match gate.snapshot().await {
|
||||
Ok(snapshot) => {
|
||||
samples.extend(snapshot.to_metric_samples("gateway_requests_distributed"));
|
||||
}
|
||||
Err(_) => samples.push(
|
||||
MetricSample::new(
|
||||
"concurrency_unavailable",
|
||||
"Whether the distributed concurrency gate is currently unavailable.",
|
||||
MetricKind::Gauge,
|
||||
1,
|
||||
)
|
||||
.with_labels(vec![MetricLabel::new(
|
||||
"gate",
|
||||
"gateway_requests_distributed",
|
||||
)]),
|
||||
),
|
||||
}
|
||||
}
|
||||
samples
|
||||
}
|
||||
|
||||
pub(crate) async fn try_acquire_request_permit(
|
||||
&self,
|
||||
) -> Result<Option<AdmissionPermit>, RequestAdmissionError> {
|
||||
let local = self
|
||||
.request_gate
|
||||
.as_ref()
|
||||
.map(|gate| gate.try_acquire())
|
||||
.transpose()
|
||||
.map_err(RequestAdmissionError::Local)?;
|
||||
let distributed = match self.distributed_request_gate.as_ref() {
|
||||
Some(gate) => Some(
|
||||
gate.try_acquire()
|
||||
.await
|
||||
.map_err(RequestAdmissionError::Distributed)?,
|
||||
),
|
||||
None => None,
|
||||
};
|
||||
Ok(AdmissionPermit::from_parts(local, distributed))
|
||||
}
|
||||
|
||||
pub fn has_auth_api_key_data_reader(&self) -> bool {
|
||||
self.data.has_auth_api_key_reader()
|
||||
}
|
||||
|
||||
pub fn has_video_task_data_reader(&self) -> bool {
|
||||
self.data.has_video_task_reader()
|
||||
}
|
||||
|
||||
pub fn has_request_candidate_data_reader(&self) -> bool {
|
||||
self.data.has_request_candidate_reader()
|
||||
}
|
||||
|
||||
pub fn has_provider_catalog_data_reader(&self) -> bool {
|
||||
self.data.has_provider_catalog_reader()
|
||||
}
|
||||
|
||||
pub fn has_usage_data_reader(&self) -> bool {
|
||||
self.data.has_usage_reader()
|
||||
}
|
||||
|
||||
pub fn has_shadow_result_data_writer(&self) -> bool {
|
||||
self.data.has_shadow_result_writer()
|
||||
}
|
||||
|
||||
pub fn has_shadow_result_data_reader(&self) -> bool {
|
||||
self.data.has_shadow_result_reader()
|
||||
}
|
||||
|
||||
pub(crate) async fn read_data_backed_video_task_response(
|
||||
&self,
|
||||
route_family: Option<&str>,
|
||||
request_path: &str,
|
||||
) -> Result<Option<video_tasks::LocalVideoTaskReadResponse>, GatewayError> {
|
||||
self.data
|
||||
.read_video_task_response(route_family, request_path)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_candidate_trace(
|
||||
&self,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
) -> Result<Option<data::RequestCandidateTrace>, GatewayError> {
|
||||
self.data
|
||||
.read_request_candidate_trace(request_id, attempted_only)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn read_decision_trace(
|
||||
&self,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
) -> Result<Option<data::DecisionTrace>, GatewayError> {
|
||||
self.data
|
||||
.read_decision_trace(request_id, attempted_only)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_usage_audit(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<data::RequestUsageAudit>, GatewayError> {
|
||||
self.data
|
||||
.read_request_usage_audit(request_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_audit_bundle(
|
||||
&self,
|
||||
request_id: &str,
|
||||
attempted_only: bool,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<data::RequestAuditBundle>, GatewayError> {
|
||||
self.data
|
||||
.read_request_audit_bundle(request_id, attempted_only, now_unix_secs)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn read_auth_api_key_snapshot(
|
||||
&self,
|
||||
user_id: &str,
|
||||
api_key_id: &str,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<data::StoredGatewayAuthApiKeySnapshot>, GatewayError> {
|
||||
self.data
|
||||
.read_auth_api_key_snapshot(user_id, api_key_id, now_unix_secs)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn record_shadow_result_sample(
|
||||
&self,
|
||||
sample: aether_data::repository::shadow_results::RecordShadowResultSample,
|
||||
) -> Result<Option<aether_data::repository::shadow_results::StoredShadowResult>, GatewayError>
|
||||
{
|
||||
self.data
|
||||
.record_shadow_result_sample(sample)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_recent_shadow_results(
|
||||
&self,
|
||||
limit: usize,
|
||||
) -> Result<Vec<aether_data::repository::shadow_results::StoredShadowResult>, GatewayError>
|
||||
{
|
||||
self.data
|
||||
.list_recent_shadow_results(limit)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_video_task_data_reader_for_tests(
|
||||
mut self,
|
||||
repository: Arc<dyn aether_data::repository::video_tasks::VideoTaskReadRepository>,
|
||||
) -> Self {
|
||||
self.data = Arc::new(GatewayDataState::with_video_task_reader_for_tests(
|
||||
repository,
|
||||
));
|
||||
self
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_request_candidate_data_reader_for_tests(
|
||||
mut self,
|
||||
repository: Arc<dyn aether_data::repository::candidates::RequestCandidateReadRepository>,
|
||||
) -> Self {
|
||||
self.data = Arc::new(GatewayDataState::with_request_candidate_reader_for_tests(
|
||||
repository,
|
||||
));
|
||||
self
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_decision_trace_data_readers_for_tests(
|
||||
mut self,
|
||||
request_candidate_repository: Arc<
|
||||
dyn aether_data::repository::candidates::RequestCandidateReadRepository,
|
||||
>,
|
||||
provider_catalog_repository: Arc<
|
||||
dyn aether_data::repository::provider_catalog::ProviderCatalogReadRepository,
|
||||
>,
|
||||
) -> Self {
|
||||
self.data = Arc::new(GatewayDataState::with_decision_trace_readers_for_tests(
|
||||
request_candidate_repository,
|
||||
provider_catalog_repository,
|
||||
));
|
||||
self
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_request_audit_data_readers_for_tests(
|
||||
mut self,
|
||||
auth_api_key_repository: Arc<dyn aether_data::repository::auth::AuthApiKeyReadRepository>,
|
||||
request_candidate_repository: Arc<
|
||||
dyn aether_data::repository::candidates::RequestCandidateReadRepository,
|
||||
>,
|
||||
provider_catalog_repository: Arc<
|
||||
dyn aether_data::repository::provider_catalog::ProviderCatalogReadRepository,
|
||||
>,
|
||||
usage_repository: Arc<dyn aether_data::repository::usage::UsageReadRepository>,
|
||||
) -> Self {
|
||||
self.data = Arc::new(GatewayDataState::with_request_audit_readers_for_tests(
|
||||
auth_api_key_repository,
|
||||
request_candidate_repository,
|
||||
provider_catalog_repository,
|
||||
usage_repository,
|
||||
));
|
||||
self
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_auth_api_key_data_reader_for_tests(
|
||||
mut self,
|
||||
repository: Arc<dyn aether_data::repository::auth::AuthApiKeyReadRepository>,
|
||||
) -> Self {
|
||||
self.data = Arc::new(GatewayDataState::with_auth_api_key_reader_for_tests(
|
||||
repository,
|
||||
));
|
||||
self
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_usage_data_reader_for_tests(
|
||||
mut self,
|
||||
repository: Arc<dyn aether_data::repository::usage::UsageReadRepository>,
|
||||
) -> Self {
|
||||
self.data = Arc::new(GatewayDataState::with_usage_reader_for_tests(repository));
|
||||
self
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_shadow_result_data_writer_for_tests(
|
||||
mut self,
|
||||
repository: Arc<dyn aether_data::repository::shadow_results::ShadowResultWriteRepository>,
|
||||
) -> Self {
|
||||
self.data = Arc::new(GatewayDataState::with_shadow_result_writer_for_tests(
|
||||
repository,
|
||||
));
|
||||
self
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_shadow_result_data_repository_for_tests<T>(
|
||||
mut self,
|
||||
repository: Arc<T>,
|
||||
) -> Self
|
||||
where
|
||||
T: aether_data::repository::shadow_results::ShadowResultRepository + 'static,
|
||||
{
|
||||
self.data = Arc::new(GatewayDataState::with_shadow_result_repository_for_tests(
|
||||
repository,
|
||||
));
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_video_task_store_path(
|
||||
mut self,
|
||||
path: impl Into<std::path::PathBuf>,
|
||||
@@ -250,11 +585,48 @@ pub fn build_router_with_endpoints(
|
||||
pub fn build_router_with_state(state: AppState) -> Router {
|
||||
Router::new()
|
||||
.route("/_gateway/health", get(health))
|
||||
.route("/_gateway/metrics", get(metrics))
|
||||
.route(
|
||||
"/_gateway/audit/auth/users/{user_id}/api-keys/{api_key_id}",
|
||||
get(get_auth_api_key_snapshot),
|
||||
)
|
||||
.route(
|
||||
"/_gateway/audit/decision-trace/{request_id}",
|
||||
get(get_decision_trace),
|
||||
)
|
||||
.route(
|
||||
"/_gateway/audit/request-candidates/{request_id}",
|
||||
get(get_request_candidate_trace),
|
||||
)
|
||||
.route(
|
||||
"/_gateway/audit/request-audit/{request_id}",
|
||||
get(get_request_audit_bundle),
|
||||
)
|
||||
.route(
|
||||
"/_gateway/audit/request-usage/{request_id}",
|
||||
get(get_request_usage_audit),
|
||||
)
|
||||
.route(
|
||||
"/_gateway/audit/shadow-results/recent",
|
||||
get(list_recent_shadow_results),
|
||||
)
|
||||
.route("/", any(proxy_request))
|
||||
.route("/{*path}", any(proxy_request))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
async fn metrics(
|
||||
axum::extract::State(state): axum::extract::State<AppState>,
|
||||
) -> impl axum::response::IntoResponse {
|
||||
prometheus_response(&state.metric_samples().await)
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) enum RequestAdmissionError {
|
||||
Local(ConcurrencyError),
|
||||
Distributed(DistributedConcurrencyError),
|
||||
}
|
||||
|
||||
pub async fn serve_tcp(
|
||||
bind: &str,
|
||||
upstream_base_url: &str,
|
||||
|
||||
@@ -1,4 +1,43 @@
|
||||
use super::*;
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
|
||||
fn sample_currently_usable_auth_snapshot(
|
||||
api_key_id: &str,
|
||||
user_id: &str,
|
||||
) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_locked_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
let mut snapshot = sample_currently_usable_auth_snapshot(api_key_id, user_id);
|
||||
snapshot.api_key_is_locked = true;
|
||||
snapshot
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_reuses_cached_auth_context_for_direct_executor_plans() {
|
||||
@@ -395,3 +434,368 @@ async fn gateway_reuses_cached_auth_context_when_falling_back_to_control_execute
|
||||
executor_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_uses_data_backed_trusted_auth_context_for_direct_executor_plans() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenPlanSyncRequest {
|
||||
auth_context_present: bool,
|
||||
auth_context_user_id: String,
|
||||
auth_context_balance_remaining: String,
|
||||
auth_context_access_allowed: bool,
|
||||
}
|
||||
|
||||
let seen_plan = Arc::new(Mutex::new(None::<SeenPlanSyncRequest>));
|
||||
let seen_plan_clone = Arc::clone(&seen_plan);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |request: Request| {
|
||||
let seen_plan_inner = Arc::clone(&seen_plan_clone);
|
||||
async move {
|
||||
let raw_body = to_bytes(request.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&raw_body).expect("plan payload should parse");
|
||||
*seen_plan_inner.lock().expect("mutex should lock") =
|
||||
Some(SeenPlanSyncRequest {
|
||||
auth_context_present: payload
|
||||
.get("auth_context")
|
||||
.is_some_and(|value| !value.is_null()),
|
||||
auth_context_user_id: payload
|
||||
.get("auth_context")
|
||||
.and_then(|value| value.get("user_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
auth_context_balance_remaining: payload
|
||||
.get("auth_context")
|
||||
.and_then(|value| value.get("balance_remaining"))
|
||||
.and_then(|value| value.as_f64())
|
||||
.map(|value| value.to_string())
|
||||
.unwrap_or_default(),
|
||||
auth_context_access_allowed: payload
|
||||
.get("auth_context")
|
||||
.and_then(|value| value.get("access_allowed"))
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(false),
|
||||
});
|
||||
Json(json!({
|
||||
"action": "executor_sync",
|
||||
"plan_kind": "openai_chat_sync",
|
||||
"plan": {
|
||||
"request_id": "req-openai-chat-trusted-123",
|
||||
"provider_name": "openai",
|
||||
"provider_id": "provider-openai-chat-trusted-123",
|
||||
"endpoint_id": "endpoint-openai-chat-trusted-123",
|
||||
"key_id": "key-openai-chat-trusted-123",
|
||||
"method": "POST",
|
||||
"url": "https://api.openai.example/v1/chat/completions",
|
||||
"headers": {
|
||||
"authorization": "Bearer upstream-key",
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"model": "gpt-5",
|
||||
"messages": []
|
||||
}
|
||||
},
|
||||
"stream": false,
|
||||
"client_api_format": "openai:chat",
|
||||
"provider_api_format": "openai:chat",
|
||||
"model_name": "gpt-5"
|
||||
},
|
||||
"report_kind": "openai_chat_sync_success",
|
||||
"report_context": {
|
||||
"user_id": "user-chat-trusted-123",
|
||||
"api_key_id": "key-chat-trusted-123"
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(|_request: Request| async move { Json(json!({"ok": true})) }),
|
||||
);
|
||||
|
||||
let executor = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"request_id": "req-openai-chat-trusted-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "chatcmpl-trusted-123",
|
||||
"object": "chat.completion",
|
||||
"model": "gpt-5",
|
||||
"choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": 1,
|
||||
"completion_tokens": 2,
|
||||
"total_tokens": 3
|
||||
}
|
||||
}
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_currently_usable_auth_snapshot("key-chat-trusted-123", "user-chat-trusted-123"),
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (executor_url, executor_handle) = start_server(executor).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new_with_executor(
|
||||
upstream_url.clone(),
|
||||
Some(upstream_url.clone()),
|
||||
Some(executor_url.clone()),
|
||||
)
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, "trace-openai-chat-trusted-1")
|
||||
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-chat-trusted-123")
|
||||
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-chat-trusted-123")
|
||||
.header(TRUSTED_AUTH_BALANCE_HEADER, "7.5")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
let seen_plan_request = seen_plan
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("plan-sync should be captured");
|
||||
assert!(seen_plan_request.auth_context_present);
|
||||
assert_eq!(
|
||||
seen_plan_request.auth_context_user_id,
|
||||
"user-chat-trusted-123"
|
||||
);
|
||||
assert_eq!(seen_plan_request.auth_context_balance_remaining, "7.5");
|
||||
assert!(seen_plan_request.auth_context_access_allowed);
|
||||
|
||||
gateway_handle.abort();
|
||||
executor_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_explicit_trusted_balance_failure_before_direct_executor_plan() {
|
||||
let seen_plan = Arc::new(Mutex::new(0usize));
|
||||
let seen_plan_clone = Arc::clone(&seen_plan);
|
||||
let seen_executor = Arc::new(Mutex::new(0usize));
|
||||
let seen_executor_clone = Arc::clone(&seen_executor);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |_request: Request| {
|
||||
let seen_plan_inner = Arc::clone(&seen_plan_clone);
|
||||
async move {
|
||||
*seen_plan_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({
|
||||
"action": "executor_sync",
|
||||
"plan_kind": "openai_chat_sync"
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(|_request: Request| async move { Json(json!({"ok": true})) }),
|
||||
);
|
||||
|
||||
let executor = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |_request: Request| {
|
||||
let seen_executor_inner = Arc::clone(&seen_executor_clone);
|
||||
async move {
|
||||
*seen_executor_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({
|
||||
"request_id": "req-openai-chat-trusted-denied-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "chatcmpl-trusted-denied-123",
|
||||
"object": "chat.completion",
|
||||
"model": "gpt-5",
|
||||
"choices": []
|
||||
}
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_currently_usable_auth_snapshot("key-chat-trusted-123", "user-chat-trusted-123"),
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (executor_url, executor_handle) = start_server(executor).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new_with_executor(
|
||||
upstream_url.clone(),
|
||||
Some(upstream_url.clone()),
|
||||
Some(executor_url.clone()),
|
||||
)
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, "trace-openai-chat-trusted-denied-1")
|
||||
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-chat-trusted-123")
|
||||
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-chat-trusted-123")
|
||||
.header(TRUSTED_AUTH_BALANCE_HEADER, "0")
|
||||
.header(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, "false")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "balance_exceeded");
|
||||
assert_eq!(payload["error"]["details"]["remaining"], 0.0);
|
||||
|
||||
assert_eq!(*seen_plan.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*seen_executor.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
executor_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_locked_trusted_snapshot_before_direct_executor_plan() {
|
||||
let seen_plan = Arc::new(Mutex::new(0usize));
|
||||
let seen_plan_clone = Arc::clone(&seen_plan);
|
||||
let seen_executor = Arc::new(Mutex::new(0usize));
|
||||
let seen_executor_clone = Arc::clone(&seen_executor);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |_request: Request| {
|
||||
let seen_plan_inner = Arc::clone(&seen_plan_clone);
|
||||
async move {
|
||||
*seen_plan_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({
|
||||
"action": "executor_sync",
|
||||
"plan_kind": "openai_chat_sync"
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(|_request: Request| async move { Json(json!({"ok": true})) }),
|
||||
);
|
||||
|
||||
let executor = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |_request: Request| {
|
||||
let seen_executor_inner = Arc::clone(&seen_executor_clone);
|
||||
async move {
|
||||
*seen_executor_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({
|
||||
"request_id": "req-openai-chat-trusted-locked-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "chatcmpl-trusted-locked-123",
|
||||
"object": "chat.completion",
|
||||
"model": "gpt-5",
|
||||
"choices": []
|
||||
}
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_locked_auth_snapshot("key-chat-trusted-123", "user-chat-trusted-123"),
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (executor_url, executor_handle) = start_server(executor).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new_with_executor(
|
||||
upstream_url.clone(),
|
||||
Some(upstream_url.clone()),
|
||||
Some(executor_url.clone()),
|
||||
)
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, "trace-openai-chat-trusted-locked-1")
|
||||
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-chat-trusted-123")
|
||||
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-chat-trusted-123")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::FORBIDDEN);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"该密钥已被管理员锁定,请联系管理员"
|
||||
);
|
||||
|
||||
assert_eq!(*seen_plan.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*seen_executor.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
executor_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
840
crates/aether-gateway/src/gateway/tests/audit.rs
Normal file
840
crates/aether-gateway/src/gateway/tests/audit.rs
Normal file
@@ -0,0 +1,840 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::candidates::{
|
||||
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_data::repository::shadow_results::{
|
||||
InMemoryShadowResultRepository, ShadowResultMatchStatus, ShadowResultReadRepository,
|
||||
};
|
||||
use aether_data::repository::usage::{InMemoryUsageReadRepository, StoredRequestUsageAudit};
|
||||
use serde_json::Value;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_records_shadow_result_for_ai_public_proxy_response() {
|
||||
let repository = Arc::new(InMemoryShadowResultRepository::default());
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "chat",
|
||||
"auth_endpoint_signature": "openai:chat",
|
||||
"executor_candidate": false,
|
||||
"public_path": "/v1/chat/completions"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(|_request: Request| async move {
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from("{\"id\":\"chatcmpl-shadow\"}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway_state = AppState::new(upstream_url.clone(), Some(upstream_url))
|
||||
.expect("gateway state should build")
|
||||
.with_shadow_result_data_writer_for_tests(repository.clone());
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.body("{\"model\":\"gpt-4.1\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(CONTROL_ROUTE_CLASS_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("ai_public")
|
||||
);
|
||||
|
||||
let response_trace_id = response
|
||||
.headers()
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.expect("trace id should exist")
|
||||
.to_string();
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
"{\"id\":\"chatcmpl-shadow\"}"
|
||||
);
|
||||
|
||||
for _ in 0..50 {
|
||||
if repository
|
||||
.list_recent(1)
|
||||
.await
|
||||
.map(|rows| !rows.is_empty())
|
||||
.unwrap_or(false)
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
|
||||
let stored = repository
|
||||
.list_recent(1)
|
||||
.await
|
||||
.expect("list should succeed")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("stored result should exist");
|
||||
assert_eq!(stored.trace_id, response_trace_id);
|
||||
assert!(stored.request_id.is_none());
|
||||
assert_eq!(stored.route_family.as_deref(), Some("openai"));
|
||||
assert_eq!(stored.route_kind.as_deref(), Some("chat"));
|
||||
assert_eq!(stored.match_status, ShadowResultMatchStatus::Pending);
|
||||
assert_eq!(stored.status_code, Some(200));
|
||||
assert!(stored.rust_result_digest.is_some());
|
||||
assert!(stored.python_result_digest.is_none());
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_records_candidate_id_in_shadow_result_for_direct_executor_response() {
|
||||
let repository = Arc::new(InMemoryShadowResultRepository::default());
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "executor_sync",
|
||||
"plan_kind": "openai_chat_sync",
|
||||
"plan": {
|
||||
"request_id": "req-shadow-direct-123",
|
||||
"candidate_id": "cand-shadow-direct-123",
|
||||
"provider_name": "openai",
|
||||
"provider_id": "provider-shadow-direct-123",
|
||||
"endpoint_id": "endpoint-shadow-direct-123",
|
||||
"key_id": "key-shadow-direct-123",
|
||||
"method": "POST",
|
||||
"url": "https://api.openai.example/v1/chat/completions",
|
||||
"headers": {
|
||||
"authorization": "Bearer upstream-key",
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"model": "gpt-5",
|
||||
"messages": []
|
||||
}
|
||||
},
|
||||
"stream": false,
|
||||
"client_api_format": "openai:chat",
|
||||
"provider_api_format": "openai:chat",
|
||||
"model_name": "gpt-5"
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let executor = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"request_id": "req-shadow-direct-123",
|
||||
"candidate_id": "cand-shadow-direct-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "chatcmpl-shadow-direct-123",
|
||||
"object": "chat.completion",
|
||||
"model": "gpt-5",
|
||||
"choices": []
|
||||
}
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (executor_url, executor_handle) = start_server(executor).await;
|
||||
let gateway_state =
|
||||
AppState::new_with_executor(upstream_url.clone(), Some(upstream_url), Some(executor_url))
|
||||
.expect("gateway state should build")
|
||||
.with_shadow_result_data_repository_for_tests(repository.clone());
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(CONTROL_CANDIDATE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("cand-shadow-direct-123")
|
||||
);
|
||||
|
||||
for _ in 0..50 {
|
||||
if repository
|
||||
.list_recent(1)
|
||||
.await
|
||||
.map(|rows| !rows.is_empty())
|
||||
.unwrap_or(false)
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
|
||||
let stored = repository
|
||||
.list_recent(1)
|
||||
.await
|
||||
.expect("list should succeed")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("stored result should exist");
|
||||
assert_eq!(stored.request_id.as_deref(), Some("req-shadow-direct-123"));
|
||||
assert_eq!(
|
||||
stored.candidate_id.as_deref(),
|
||||
Some("cand-shadow-direct-123")
|
||||
);
|
||||
assert_eq!(stored.route_family.as_deref(), Some("openai"));
|
||||
assert_eq!(stored.route_kind.as_deref(), Some("chat"));
|
||||
assert_eq!(stored.match_status, ShadowResultMatchStatus::Pending);
|
||||
|
||||
gateway_handle.abort();
|
||||
executor_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_exposes_request_id_header_for_direct_executor_response() {
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_auth_snapshot("api-key-1", "user-1"),
|
||||
)]));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_request_candidate(
|
||||
"cand-1",
|
||||
"req-direct-audit-123",
|
||||
0,
|
||||
RequestCandidateStatus::Success,
|
||||
Some(101),
|
||||
Some(37),
|
||||
Some(200),
|
||||
),
|
||||
]));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_request_usage("req-direct-audit-123"),
|
||||
]));
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "executor_sync",
|
||||
"plan_kind": "openai_chat_sync",
|
||||
"plan": {
|
||||
"request_id": "req-direct-audit-123",
|
||||
"candidate_id": "cand-direct-audit-123",
|
||||
"provider_name": "openai",
|
||||
"provider_id": "provider-direct-audit-123",
|
||||
"endpoint_id": "endpoint-direct-audit-123",
|
||||
"key_id": "key-direct-audit-123",
|
||||
"method": "POST",
|
||||
"url": "https://api.openai.example/v1/chat/completions",
|
||||
"headers": {
|
||||
"authorization": "Bearer upstream-key",
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"model": "gpt-5",
|
||||
"messages": []
|
||||
}
|
||||
},
|
||||
"stream": false,
|
||||
"client_api_format": "openai:chat",
|
||||
"provider_api_format": "openai:chat",
|
||||
"model_name": "gpt-5"
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let executor = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"request_id": "req-direct-audit-123",
|
||||
"candidate_id": "cand-direct-audit-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "chatcmpl-direct-audit-123",
|
||||
"object": "chat.completion",
|
||||
"model": "gpt-5",
|
||||
"choices": []
|
||||
}
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (executor_url, executor_handle) = start_server(executor).await;
|
||||
let gateway_state =
|
||||
AppState::new_with_executor(upstream_url.clone(), Some(upstream_url), Some(executor_url))
|
||||
.expect("gateway state should build")
|
||||
.with_request_audit_data_readers_for_tests(
|
||||
auth_repository,
|
||||
request_candidates,
|
||||
provider_catalog,
|
||||
usage_repository,
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let request_id = response
|
||||
.headers()
|
||||
.get(CONTROL_REQUEST_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.expect("request id header should exist")
|
||||
.to_string();
|
||||
assert_eq!(request_id, "req-direct-audit-123");
|
||||
|
||||
let audit_response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/_gateway/audit/request-audit/{request_id}?attempted_only=true"
|
||||
))
|
||||
.send()
|
||||
.await
|
||||
.expect("request audit should succeed");
|
||||
|
||||
assert_eq!(audit_response.status(), StatusCode::OK);
|
||||
let payload: Value = audit_response.json().await.expect("payload should parse");
|
||||
assert_eq!(payload["request_id"], "req-direct-audit-123");
|
||||
assert_eq!(payload["usage"]["provider_name"], "OpenAI");
|
||||
assert_eq!(payload["decision_trace"]["total_candidates"], 1);
|
||||
assert_eq!(payload["auth_snapshot"]["api_key_id"], "api-key-1");
|
||||
|
||||
gateway_handle.abort();
|
||||
executor_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_exposes_recent_shadow_results_via_internal_audit_endpoint() {
|
||||
let repository = Arc::new(InMemoryShadowResultRepository::default());
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "chat",
|
||||
"auth_endpoint_signature": "openai:chat",
|
||||
"executor_candidate": false,
|
||||
"public_path": "/v1/chat/completions"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(|_request: Request| async move {
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from("{\"id\":\"chatcmpl-shadow-read\"}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway_state = AppState::new(upstream_url.clone(), Some(upstream_url))
|
||||
.expect("gateway state should build")
|
||||
.with_shadow_result_data_repository_for_tests(repository.clone());
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let write_response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.body("{\"model\":\"gpt-4.1\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(write_response.status(), StatusCode::OK);
|
||||
|
||||
for _ in 0..50 {
|
||||
if repository
|
||||
.list_recent(1)
|
||||
.await
|
||||
.map(|rows| !rows.is_empty())
|
||||
.unwrap_or(false)
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/_gateway/audit/shadow-results/recent?limit=5"
|
||||
))
|
||||
.send()
|
||||
.await
|
||||
.expect("audit request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: Value = response.json().await.expect("payload should parse");
|
||||
assert_eq!(payload["limit_applied"], 5);
|
||||
assert_eq!(payload["counts"]["pending"], 1);
|
||||
assert_eq!(payload["counts"]["match"], 0);
|
||||
assert_eq!(
|
||||
payload["items"].as_array().map(|items| items.len()),
|
||||
Some(1)
|
||||
);
|
||||
assert!(payload["items"][0]["request_id"].is_null());
|
||||
assert_eq!(payload["items"][0]["route_family"], "openai");
|
||||
assert_eq!(payload["items"][0]["route_kind"], "chat");
|
||||
assert_eq!(payload["items"][0]["match_status"], "Pending");
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
fn sample_request_candidate(
|
||||
id: &str,
|
||||
request_id: &str,
|
||||
candidate_index: i32,
|
||||
status: RequestCandidateStatus,
|
||||
started_at_unix_secs: Option<i64>,
|
||||
latency_ms: Option<i32>,
|
||||
status_code: Option<i32>,
|
||||
) -> StoredRequestCandidate {
|
||||
StoredRequestCandidate::new(
|
||||
id.to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
candidate_index,
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
status,
|
||||
None,
|
||||
false,
|
||||
status_code,
|
||||
None,
|
||||
None,
|
||||
latency_ms,
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
100 + i64::from(candidate_index),
|
||||
started_at_unix_secs,
|
||||
started_at_unix_secs.map(|value| value + 1),
|
||||
)
|
||||
.expect("candidate should build")
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-4.1"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-1".to_string(),
|
||||
"OpenAI".to_string(),
|
||||
Some("https://openai.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"openai:chat".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"provider-key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"prod-key".to_string(),
|
||||
"api_key".to_string(),
|
||||
Some(serde_json::json!({"cache_1h": true})),
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
}
|
||||
|
||||
fn sample_request_usage(request_id: &str) -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
"OpenAI".to_string(),
|
||||
"gpt-4.1".to_string(),
|
||||
Some("gpt-4.1-mini".to_string()),
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
false,
|
||||
120,
|
||||
40,
|
||||
160,
|
||||
0.24,
|
||||
0.36,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(450),
|
||||
Some(120),
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
100,
|
||||
101,
|
||||
Some(102),
|
||||
)
|
||||
.expect("usage should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_exposes_request_usage_via_internal_audit_endpoint() {
|
||||
let repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_request_usage("req-usage-2"),
|
||||
]));
|
||||
let gateway_state = AppState::new("http://127.0.0.1:18091", None)
|
||||
.expect("gateway state should build")
|
||||
.with_usage_data_reader_for_tests(repository);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/_gateway/audit/request-usage/req-usage-2"
|
||||
))
|
||||
.send()
|
||||
.await
|
||||
.expect("audit request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: Value = response.json().await.expect("payload should parse");
|
||||
assert_eq!(payload["request_id"], "req-usage-2");
|
||||
assert_eq!(payload["provider_name"], "OpenAI");
|
||||
assert_eq!(payload["api_format"], "openai:chat");
|
||||
assert_eq!(payload["total_tokens"], 160);
|
||||
assert_eq!(payload["total_cost_usd"], 0.24);
|
||||
assert_eq!(payload["status"], "completed");
|
||||
assert_eq!(payload["billing_status"], "settled");
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_exposes_request_audit_bundle_via_internal_audit_endpoint() {
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_auth_snapshot("api-key-1", "user-1"),
|
||||
)]));
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_request_candidate(
|
||||
"cand-1",
|
||||
"req-audit-1",
|
||||
0,
|
||||
RequestCandidateStatus::Success,
|
||||
Some(101),
|
||||
Some(37),
|
||||
Some(200),
|
||||
),
|
||||
]));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_request_usage("req-audit-1"),
|
||||
]));
|
||||
let gateway_state = AppState::new("http://127.0.0.1:18092", None)
|
||||
.expect("gateway state should build")
|
||||
.with_request_audit_data_readers_for_tests(
|
||||
auth_repository,
|
||||
request_candidates,
|
||||
provider_catalog,
|
||||
usage_repository,
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/_gateway/audit/request-audit/req-audit-1?attempted_only=true"
|
||||
))
|
||||
.send()
|
||||
.await
|
||||
.expect("audit request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: Value = response.json().await.expect("payload should parse");
|
||||
assert_eq!(payload["request_id"], "req-audit-1");
|
||||
assert_eq!(payload["usage"]["provider_name"], "OpenAI");
|
||||
assert_eq!(payload["usage"]["total_tokens"], 160);
|
||||
assert_eq!(payload["decision_trace"]["total_candidates"], 1);
|
||||
assert_eq!(
|
||||
payload["decision_trace"]["candidates"][0]["provider_key_name"],
|
||||
"prod-key"
|
||||
);
|
||||
assert_eq!(payload["auth_snapshot"]["api_key_id"], "api-key-1");
|
||||
assert_eq!(payload["auth_snapshot"]["currently_usable"], true);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_exposes_request_candidate_trace_via_internal_audit_endpoint() {
|
||||
let repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_request_candidate(
|
||||
"cand-1",
|
||||
"req-trace-1",
|
||||
0,
|
||||
RequestCandidateStatus::Pending,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
),
|
||||
sample_request_candidate(
|
||||
"cand-2",
|
||||
"req-trace-1",
|
||||
1,
|
||||
RequestCandidateStatus::Failed,
|
||||
Some(101),
|
||||
Some(37),
|
||||
Some(502),
|
||||
),
|
||||
]));
|
||||
|
||||
let gateway_state = AppState::new("http://127.0.0.1:19081", None)
|
||||
.expect("gateway state should build")
|
||||
.with_request_candidate_data_reader_for_tests(repository);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/_gateway/audit/request-candidates/req-trace-1?attempted_only=true"
|
||||
))
|
||||
.send()
|
||||
.await
|
||||
.expect("audit request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: Value = response.json().await.expect("payload should parse");
|
||||
assert_eq!(payload["request_id"], "req-trace-1");
|
||||
assert_eq!(payload["total_candidates"], 1);
|
||||
assert_eq!(payload["final_status"], "failed");
|
||||
assert_eq!(payload["total_latency_ms"], 37);
|
||||
assert_eq!(
|
||||
payload["candidates"].as_array().map(|items| items.len()),
|
||||
Some(1)
|
||||
);
|
||||
assert_eq!(payload["candidates"][0]["id"], "cand-2");
|
||||
assert_eq!(payload["candidates"][0]["status"], "failed");
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_exposes_decision_trace_via_internal_audit_endpoint() {
|
||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||
sample_request_candidate(
|
||||
"cand-1",
|
||||
"req-trace-2",
|
||||
0,
|
||||
RequestCandidateStatus::Failed,
|
||||
Some(101),
|
||||
Some(37),
|
||||
Some(502),
|
||||
),
|
||||
]));
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let gateway_state = AppState::new("http://127.0.0.1:19083", None)
|
||||
.expect("gateway state should build")
|
||||
.with_decision_trace_data_readers_for_tests(request_candidates, provider_catalog);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/_gateway/audit/decision-trace/req-trace-2?attempted_only=true"
|
||||
))
|
||||
.send()
|
||||
.await
|
||||
.expect("audit request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: Value = response.json().await.expect("payload should parse");
|
||||
assert_eq!(payload["request_id"], "req-trace-2");
|
||||
assert_eq!(payload["total_candidates"], 1);
|
||||
assert_eq!(payload["candidates"][0]["provider_name"], "OpenAI");
|
||||
assert_eq!(
|
||||
payload["candidates"][0]["provider_website"],
|
||||
"https://openai.com"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["candidates"][0]["endpoint_api_format"],
|
||||
"openai:chat"
|
||||
);
|
||||
assert_eq!(payload["candidates"][0]["provider_key_name"], "prod-key");
|
||||
assert_eq!(
|
||||
payload["candidates"][0]["provider_key_auth_type"],
|
||||
"api_key"
|
||||
);
|
||||
assert_eq!(
|
||||
payload["candidates"][0]["provider_key_capabilities"]["cache_1h"],
|
||||
true
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_exposes_auth_api_key_snapshot_via_internal_audit_endpoint() {
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_auth_snapshot("key-1", "user-1"),
|
||||
)]));
|
||||
|
||||
let gateway_state = AppState::new("http://127.0.0.1:19082", None)
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/_gateway/audit/auth/users/user-1/api-keys/key-1"
|
||||
))
|
||||
.send()
|
||||
.await
|
||||
.expect("audit request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: Value = response.json().await.expect("payload should parse");
|
||||
assert_eq!(payload["user_id"], "user-1");
|
||||
assert_eq!(payload["api_key_id"], "key-1");
|
||||
assert_eq!(payload["username"], "alice");
|
||||
assert_eq!(payload["user_role"], "user");
|
||||
assert_eq!(payload["api_key_name"], "default");
|
||||
assert_eq!(payload["currently_usable"], true);
|
||||
assert_eq!(payload["user_allowed_providers"][0], "openai");
|
||||
assert_eq!(payload["api_key_allowed_api_formats"][0], "openai:chat");
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
187
crates/aether-gateway/src/gateway/tests/concurrency.rs
Normal file
187
crates/aether-gateway/src/gateway/tests/concurrency.rs
Normal file
@@ -0,0 +1,187 @@
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_second_in_flight_stream_request_with_distributed_overload() {
|
||||
let upstream_hits = Arc::new(AtomicUsize::new(0));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/{*path}",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
upstream_hits.fetch_add(1, Ordering::SeqCst);
|
||||
let stream = async_stream::stream! {
|
||||
yield Ok::<_, Infallible>(Bytes::from_static(b"chunk-1"));
|
||||
futures_util::future::pending::<()>().await;
|
||||
};
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(stream))
|
||||
.expect("response should build")
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let distributed_gate = aether_runtime::DistributedConcurrencyGate::new_in_memory(
|
||||
"gateway_requests_distributed",
|
||||
1,
|
||||
);
|
||||
let gateway_a = build_router_with_state(
|
||||
AppState::new(upstream_url.clone(), None)
|
||||
.expect("gateway state should build")
|
||||
.with_distributed_request_concurrency_gate(distributed_gate.clone()),
|
||||
);
|
||||
let gateway_b = build_router_with_state(
|
||||
AppState::new(upstream_url, None)
|
||||
.expect("gateway state should build")
|
||||
.with_distributed_request_concurrency_gate(distributed_gate),
|
||||
);
|
||||
let (gateway_a_url, gateway_a_handle) = start_server(gateway_a).await;
|
||||
let (gateway_b_url, gateway_b_handle) = start_server(gateway_b).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let first_response = client
|
||||
.get(format!("{gateway_a_url}/v1/messages"))
|
||||
.send()
|
||||
.await
|
||||
.expect("first request should succeed");
|
||||
|
||||
wait_until(500, || upstream_hits.load(Ordering::SeqCst) == 1).await;
|
||||
|
||||
let second_response = client
|
||||
.get(format!("{gateway_b_url}/v1/messages"))
|
||||
.send()
|
||||
.await
|
||||
.expect("second request should complete");
|
||||
|
||||
assert_eq!(second_response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
second_response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_DISTRIBUTED_OVERLOADED)
|
||||
);
|
||||
assert_eq!(
|
||||
second_response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("json body should decode")["error"]["details"]["gate"],
|
||||
"gateway_requests_distributed"
|
||||
);
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
drop(first_response);
|
||||
gateway_a_handle.abort();
|
||||
gateway_b_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_second_in_flight_stream_request_with_local_overload() {
|
||||
let upstream_hits = Arc::new(AtomicUsize::new(0));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/{*path}",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
upstream_hits.fetch_add(1, Ordering::SeqCst);
|
||||
let stream = async_stream::stream! {
|
||||
yield Ok::<_, Infallible>(Bytes::from_static(b"chunk-1"));
|
||||
futures_util::future::pending::<()>().await;
|
||||
};
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(stream))
|
||||
.expect("response should build")
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new(upstream_url, None)
|
||||
.expect("gateway state should build")
|
||||
.with_request_concurrency_limit(1),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let first_response = client
|
||||
.get(format!("{gateway_url}/v1/messages"))
|
||||
.send()
|
||||
.await
|
||||
.expect("first request should succeed");
|
||||
|
||||
wait_until(500, || upstream_hits.load(Ordering::SeqCst) == 1).await;
|
||||
|
||||
let second_response = client
|
||||
.get(format!("{gateway_url}/v1/messages"))
|
||||
.send()
|
||||
.await
|
||||
.expect("second request should complete");
|
||||
|
||||
assert_eq!(second_response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
second_response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_OVERLOADED)
|
||||
);
|
||||
assert_eq!(
|
||||
second_response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("json body should decode")["error"]["type"],
|
||||
"overloaded"
|
||||
);
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
drop(first_response);
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_exposes_request_concurrency_metrics() {
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new("http://127.0.0.1:1", None)
|
||||
.expect("gateway state should build")
|
||||
.with_request_concurrency_limit(3)
|
||||
.with_distributed_request_concurrency_gate(
|
||||
aether_runtime::DistributedConcurrencyGate::new_in_memory(
|
||||
"gateway_requests_distributed",
|
||||
5,
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/_gateway/metrics"))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("text/plain; version=0.0.4; charset=utf-8")
|
||||
);
|
||||
let body = response.text().await.expect("body should read");
|
||||
assert!(body.contains("service_up{service=\"aether-gateway\"} 1"));
|
||||
assert!(body.contains("concurrency_in_flight{gate=\"gateway_requests\"} 0"));
|
||||
assert!(body.contains("concurrency_available_permits{gate=\"gateway_requests\"} 3"));
|
||||
assert!(body.contains("concurrency_in_flight{gate=\"gateway_requests_distributed\"} 0"));
|
||||
assert!(body.contains("concurrency_available_permits{gate=\"gateway_requests_distributed\"} 5"));
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
@@ -1,4 +1,43 @@
|
||||
use super::*;
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
|
||||
fn sample_currently_usable_auth_snapshot(
|
||||
api_key_id: &str,
|
||||
user_id: &str,
|
||||
) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_expired_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
let mut snapshot = sample_currently_usable_auth_snapshot(api_key_id, user_id);
|
||||
snapshot.api_key_expires_at_unix_secs = Some(1);
|
||||
snapshot
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_consults_control_api_for_ai_routes_and_propagates_decision_headers() {
|
||||
@@ -205,3 +244,284 @@ async fn gateway_consults_control_api_for_ai_routes_and_propagates_decision_head
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_uses_data_backed_trusted_auth_context_without_calling_control_auth_endpoint() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenPublicRequest {
|
||||
trusted_user_id: String,
|
||||
trusted_api_key_id: String,
|
||||
trusted_balance_remaining: String,
|
||||
trusted_access_allowed: String,
|
||||
}
|
||||
|
||||
let auth_context_hits = Arc::new(Mutex::new(0usize));
|
||||
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
|
||||
let seen_public = Arc::new(Mutex::new(None::<SeenPublicRequest>));
|
||||
let seen_public_clone = Arc::clone(&seen_public);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/auth-context",
|
||||
any(move |_request: Request| {
|
||||
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
|
||||
async move {
|
||||
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({
|
||||
"auth_context": {
|
||||
"user_id": "user-from-control",
|
||||
"api_key_id": "key-from-control",
|
||||
"balance_remaining": 99.0,
|
||||
"access_allowed": true
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |request: Request| {
|
||||
let seen_public_inner = Arc::clone(&seen_public_clone);
|
||||
async move {
|
||||
*seen_public_inner.lock().expect("mutex should lock") =
|
||||
Some(SeenPublicRequest {
|
||||
trusted_user_id: request
|
||||
.headers()
|
||||
.get(TRUSTED_AUTH_USER_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
trusted_api_key_id: request
|
||||
.headers()
|
||||
.get(TRUSTED_AUTH_API_KEY_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
trusted_balance_remaining: request
|
||||
.headers()
|
||||
.get(TRUSTED_AUTH_BALANCE_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
trusted_access_allowed: request
|
||||
.headers()
|
||||
.get(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
(
|
||||
StatusCode::OK,
|
||||
[(GATEWAY_HEADER, "python-upstream")],
|
||||
Body::from("proxied"),
|
||||
)
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_currently_usable_auth_snapshot("key-123", "user-123"),
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new(upstream_url.clone(), Some(upstream_url))
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(TRACE_ID_HEADER, "trace-control-data-auth-1")
|
||||
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-123")
|
||||
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-123")
|
||||
.header(TRUSTED_AUTH_BALANCE_HEADER, "42.5")
|
||||
.body("{\"hello\":\"world\"}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let seen_public_request = seen_public
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("public request should be captured");
|
||||
assert_eq!(seen_public_request.trusted_user_id, "user-123");
|
||||
assert_eq!(seen_public_request.trusted_api_key_id, "key-123");
|
||||
assert_eq!(seen_public_request.trusted_balance_remaining, "42.5");
|
||||
assert_eq!(seen_public_request.trusted_access_allowed, "true");
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_explicit_trusted_balance_failure_without_hitting_control_or_upstream(
|
||||
) {
|
||||
let auth_context_hits = Arc::new(Mutex::new(0usize));
|
||||
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/auth-context",
|
||||
any(move |_request: Request| {
|
||||
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
|
||||
async move {
|
||||
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({
|
||||
"auth_context": {
|
||||
"user_id": "user-from-control",
|
||||
"api_key_id": "key-from-control",
|
||||
"balance_remaining": 99.0,
|
||||
"access_allowed": true
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_currently_usable_auth_snapshot("key-123", "user-123"),
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new(upstream_url.clone(), Some(upstream_url))
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, "trace-control-balance-denied-1")
|
||||
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-123")
|
||||
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-123")
|
||||
.header(TRUSTED_AUTH_BALANCE_HEADER, "0")
|
||||
.header(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, "false")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(CONTROL_ROUTE_CLASS_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("ai_public")
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "balance_exceeded");
|
||||
assert_eq!(payload["error"]["message"], "余额不足(剩余: $0.00)");
|
||||
assert_eq!(payload["error"]["details"]["balance_type"], "USD");
|
||||
assert_eq!(payload["error"]["details"]["remaining"], 0.0);
|
||||
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_invalid_trusted_snapshot_without_hitting_control_or_upstream() {
|
||||
let auth_context_hits = Arc::new(Mutex::new(0usize));
|
||||
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/auth-context",
|
||||
any(move |_request: Request| {
|
||||
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
|
||||
async move {
|
||||
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({
|
||||
"auth_context": {
|
||||
"user_id": "user-from-control",
|
||||
"api_key_id": "key-from-control",
|
||||
"access_allowed": true
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_expired_auth_snapshot("key-123", "user-123"),
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new(upstream_url.clone(), Some(upstream_url))
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, "trace-control-invalid-trusted-1")
|
||||
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-123")
|
||||
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-123")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(payload["error"]["message"], "无效的API密钥");
|
||||
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
@@ -10,6 +10,8 @@ pub(super) use http::StatusCode;
|
||||
pub(super) use serde_json::json;
|
||||
|
||||
mod ai_execute;
|
||||
mod audit;
|
||||
mod concurrency;
|
||||
mod control;
|
||||
mod files;
|
||||
mod proxy;
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
use super::*;
|
||||
use aether_data::repository::video_tasks::{
|
||||
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskWriteRepository,
|
||||
};
|
||||
|
||||
mod error;
|
||||
mod gemini_sync_create;
|
||||
@@ -7,6 +10,186 @@ mod openai_sync_create;
|
||||
mod openai_sync_task;
|
||||
mod stream;
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_reads_openai_video_task_via_data_read_side_without_hitting_public_route() {
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "video",
|
||||
"auth_endpoint_signature": "openai:video",
|
||||
"executor_candidate": true,
|
||||
"auth_context": {
|
||||
"user_id": "user-video-db-123",
|
||||
"api_key_id": "key-video-db-123",
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/v1/videos/task-db-123"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/videos/task-db-123",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
repository
|
||||
.upsert(UpsertVideoTask {
|
||||
id: "task-db-123".to_string(),
|
||||
short_id: Some("short-task-db-123".to_string()),
|
||||
user_id: Some("user-video-db-123".to_string()),
|
||||
external_task_id: Some("ext-video-db-123".to_string()),
|
||||
provider_api_format: Some("openai:video".to_string()),
|
||||
model: Some("sora-2".to_string()),
|
||||
prompt: Some("hello from db".to_string()),
|
||||
size: Some("1280x720".to_string()),
|
||||
status: aether_data::repository::video_tasks::VideoTaskStatus::Processing,
|
||||
progress_percent: 45,
|
||||
created_at_unix_secs: 123,
|
||||
updated_at_unix_secs: 124,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new_with_executor(
|
||||
upstream_url.clone(),
|
||||
Some(upstream_url.clone()),
|
||||
Some(upstream_url.clone()),
|
||||
)
|
||||
.expect("gateway state should build")
|
||||
.with_video_task_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/v1/videos/task-db-123"))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(body["id"], "task-db-123");
|
||||
assert_eq!(body["status"], "processing");
|
||||
assert_eq!(body["progress"], 45);
|
||||
assert_eq!(body["created_at"], 123);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_reads_gemini_video_task_via_data_read_side_without_hitting_public_route() {
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "gemini",
|
||||
"route_kind": "video",
|
||||
"auth_endpoint_signature": "gemini:video",
|
||||
"executor_candidate": true,
|
||||
"auth_context": {
|
||||
"user_id": "user-video-db-123",
|
||||
"api_key_id": "key-video-db-123",
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/v1beta/models/veo-3/operations/localshort123"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1beta/models/veo-3/operations/localshort123",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
repository
|
||||
.upsert(UpsertVideoTask {
|
||||
id: "task-db-456".to_string(),
|
||||
short_id: Some("localshort123".to_string()),
|
||||
user_id: Some("user-video-db-123".to_string()),
|
||||
external_task_id: Some("operations/ext-video-db-123".to_string()),
|
||||
provider_api_format: Some("gemini:video".to_string()),
|
||||
model: Some("veo-3".to_string()),
|
||||
prompt: Some("hello from gemini db".to_string()),
|
||||
size: Some("720p".to_string()),
|
||||
status: aether_data::repository::video_tasks::VideoTaskStatus::Completed,
|
||||
progress_percent: 100,
|
||||
created_at_unix_secs: 223,
|
||||
updated_at_unix_secs: 224,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new_with_executor(
|
||||
upstream_url.clone(),
|
||||
Some(upstream_url.clone()),
|
||||
Some(upstream_url.clone()),
|
||||
)
|
||||
.expect("gateway state should build")
|
||||
.with_video_task_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/v1beta/models/veo-3/operations/localshort123"
|
||||
))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(body["name"], "models/veo-3/operations/localshort123");
|
||||
assert_eq!(body["done"], true);
|
||||
assert_eq!(
|
||||
body["response"]["generateVideoResponse"]["generatedSamples"][0]["video"]["uri"],
|
||||
"/v1beta/files/aev_localshort123:download?alt=media"
|
||||
);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_video_get_route_via_control_sync_endpoint() {
|
||||
#[derive(Debug, Clone)]
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use std::time::Instant;
|
||||
|
||||
use aether_runtime::{maybe_hold_axum_response_permit, AdmissionPermit};
|
||||
use axum::body::{to_bytes, Body};
|
||||
use axum::extract::{ConnectInfo, Request, State};
|
||||
use axum::http::header::{HeaderName, HeaderValue};
|
||||
@@ -15,16 +16,42 @@ use crate::gateway::headers::{
|
||||
extract_or_generate_trace_id, header_value_str, is_json_request, should_skip_request_header,
|
||||
};
|
||||
use crate::gateway::{
|
||||
build_client_response, maybe_execute_via_control, maybe_execute_via_executor_stream,
|
||||
maybe_execute_via_executor_sync, resolve_control_route, AppState, GatewayControlDecision,
|
||||
GatewayError,
|
||||
build_client_response, build_local_auth_rejection_response, build_local_overloaded_response,
|
||||
maybe_execute_via_control, maybe_execute_via_executor_stream, maybe_execute_via_executor_sync,
|
||||
record_shadow_result_non_blocking, resolve_control_route, trusted_auth_local_rejection,
|
||||
AppState, GatewayControlDecision, GatewayError,
|
||||
};
|
||||
|
||||
pub(crate) async fn health(State(state): State<AppState>) -> impl IntoResponse {
|
||||
let request_concurrency = state.request_concurrency_snapshot().map(|snapshot| {
|
||||
json!({
|
||||
"limit": snapshot.limit,
|
||||
"in_flight": snapshot.in_flight,
|
||||
"available_permits": snapshot.available_permits,
|
||||
"high_watermark": snapshot.high_watermark,
|
||||
"rejected": snapshot.rejected,
|
||||
})
|
||||
});
|
||||
let distributed_request_concurrency = state
|
||||
.distributed_request_concurrency_snapshot()
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.map(|snapshot| {
|
||||
json!({
|
||||
"limit": snapshot.limit,
|
||||
"in_flight": snapshot.in_flight,
|
||||
"available_permits": snapshot.available_permits,
|
||||
"high_watermark": snapshot.high_watermark,
|
||||
"rejected": snapshot.rejected,
|
||||
})
|
||||
});
|
||||
Json(json!({
|
||||
"status": "ok",
|
||||
"component": "aether-gateway",
|
||||
"control_api_enabled": state.control_base_url.is_some(),
|
||||
"request_concurrency": request_concurrency,
|
||||
"distributed_request_concurrency": distributed_request_concurrency,
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -34,6 +61,66 @@ pub(crate) async fn proxy_request(
|
||||
request: Request,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let started_at = Instant::now();
|
||||
let mut request_permit = match state.try_acquire_request_permit().await {
|
||||
Ok(permit) => permit,
|
||||
Err(crate::gateway::RequestAdmissionError::Local(
|
||||
aether_runtime::ConcurrencyError::Saturated { gate, limit },
|
||||
)) => {
|
||||
let trace_id = extract_or_generate_trace_id(request.headers());
|
||||
let response = build_local_overloaded_response(&trace_id, None, gate, limit)?;
|
||||
return Ok(finalize_gateway_response(
|
||||
&state,
|
||||
response,
|
||||
&trace_id,
|
||||
&remote_addr,
|
||||
request.method(),
|
||||
request
|
||||
.uri()
|
||||
.path_and_query()
|
||||
.map(|value| value.as_str())
|
||||
.unwrap_or("/"),
|
||||
None,
|
||||
EXECUTION_PATH_LOCAL_OVERLOADED,
|
||||
&started_at,
|
||||
None,
|
||||
));
|
||||
}
|
||||
Err(crate::gateway::RequestAdmissionError::Local(
|
||||
aether_runtime::ConcurrencyError::Closed { gate },
|
||||
)) => {
|
||||
return Err(GatewayError::Internal(format!(
|
||||
"gateway request concurrency gate {gate} is closed"
|
||||
)));
|
||||
}
|
||||
Err(crate::gateway::RequestAdmissionError::Distributed(
|
||||
aether_runtime::DistributedConcurrencyError::Saturated { gate, limit },
|
||||
))
|
||||
| Err(crate::gateway::RequestAdmissionError::Distributed(
|
||||
aether_runtime::DistributedConcurrencyError::Unavailable { gate, limit, .. },
|
||||
)) => {
|
||||
let trace_id = extract_or_generate_trace_id(request.headers());
|
||||
let response = build_local_overloaded_response(&trace_id, None, gate, limit)?;
|
||||
return Ok(finalize_gateway_response(
|
||||
&state,
|
||||
response,
|
||||
&trace_id,
|
||||
&remote_addr,
|
||||
request.method(),
|
||||
request
|
||||
.uri()
|
||||
.path_and_query()
|
||||
.map(|value| value.as_str())
|
||||
.unwrap_or("/"),
|
||||
None,
|
||||
EXECUTION_PATH_DISTRIBUTED_OVERLOADED,
|
||||
&started_at,
|
||||
None,
|
||||
));
|
||||
}
|
||||
Err(crate::gateway::RequestAdmissionError::Distributed(
|
||||
aether_runtime::DistributedConcurrencyError::InvalidConfiguration(message),
|
||||
)) => return Err(GatewayError::Internal(message)),
|
||||
};
|
||||
let (parts, body) = request.into_parts();
|
||||
let method = parts.method.clone();
|
||||
let path_and_query = parts
|
||||
@@ -46,6 +133,23 @@ pub(crate) async fn proxy_request(
|
||||
let trace_id = extract_or_generate_trace_id(&parts.headers);
|
||||
let control_decision =
|
||||
resolve_control_route(&state, &method, &parts.uri, &parts.headers, &trace_id).await?;
|
||||
if let Some(rejection) = trusted_auth_local_rejection(control_decision.as_ref(), &parts.headers)
|
||||
{
|
||||
let response =
|
||||
build_local_auth_rejection_response(&trace_id, control_decision.as_ref(), &rejection)?;
|
||||
return Ok(finalize_gateway_response(
|
||||
&state,
|
||||
response,
|
||||
&trace_id,
|
||||
&remote_addr,
|
||||
&method,
|
||||
path_and_query,
|
||||
control_decision.as_ref(),
|
||||
EXECUTION_PATH_LOCAL_AUTH_DENIED,
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
}
|
||||
let upstream_path_and_query = control_decision
|
||||
.as_ref()
|
||||
.map(|decision| decision.proxy_path_and_query())
|
||||
@@ -152,6 +256,7 @@ pub(crate) async fn proxy_request(
|
||||
.await?
|
||||
{
|
||||
return Ok(finalize_gateway_response(
|
||||
&state,
|
||||
executor_response,
|
||||
&trace_id,
|
||||
&remote_addr,
|
||||
@@ -160,6 +265,7 @@ pub(crate) async fn proxy_request(
|
||||
control_decision.as_ref(),
|
||||
EXECUTION_PATH_EXECUTOR_STREAM,
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -173,6 +279,7 @@ pub(crate) async fn proxy_request(
|
||||
.await?
|
||||
{
|
||||
return Ok(finalize_gateway_response(
|
||||
&state,
|
||||
executor_response,
|
||||
&trace_id,
|
||||
&remote_addr,
|
||||
@@ -181,6 +288,7 @@ pub(crate) async fn proxy_request(
|
||||
control_decision.as_ref(),
|
||||
EXECUTION_PATH_EXECUTOR_SYNC,
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
}
|
||||
if parts.method != http::Method::POST {
|
||||
@@ -194,6 +302,7 @@ pub(crate) async fn proxy_request(
|
||||
.await?
|
||||
{
|
||||
return Ok(finalize_gateway_response(
|
||||
&state,
|
||||
executor_response,
|
||||
&trace_id,
|
||||
&remote_addr,
|
||||
@@ -202,6 +311,7 @@ pub(crate) async fn proxy_request(
|
||||
control_decision.as_ref(),
|
||||
EXECUTION_PATH_EXECUTOR_STREAM,
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -216,6 +326,7 @@ pub(crate) async fn proxy_request(
|
||||
.await?
|
||||
{
|
||||
return Ok(finalize_gateway_response(
|
||||
&state,
|
||||
control_response,
|
||||
&trace_id,
|
||||
&remote_addr,
|
||||
@@ -228,6 +339,7 @@ pub(crate) async fn proxy_request(
|
||||
EXECUTION_PATH_CONTROL_EXECUTE_SYNC
|
||||
},
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -263,6 +375,7 @@ pub(crate) async fn proxy_request(
|
||||
|
||||
let response = build_client_response(upstream_response, &trace_id, control_decision.as_ref())?;
|
||||
Ok(finalize_gateway_response(
|
||||
&state,
|
||||
response,
|
||||
&trace_id,
|
||||
&remote_addr,
|
||||
@@ -275,6 +388,7 @@ pub(crate) async fn proxy_request(
|
||||
EXECUTION_PATH_PUBLIC_PROXY_PASSTHROUGH
|
||||
},
|
||||
&started_at,
|
||||
request_permit.take(),
|
||||
))
|
||||
}
|
||||
|
||||
@@ -292,6 +406,7 @@ fn request_wants_stream(parts: &http::request::Parts, body: &axum::body::Bytes)
|
||||
}
|
||||
|
||||
fn finalize_gateway_response(
|
||||
state: &AppState,
|
||||
mut response: Response<Body>,
|
||||
trace_id: &str,
|
||||
remote_addr: &std::net::SocketAddr,
|
||||
@@ -300,6 +415,7 @@ fn finalize_gateway_response(
|
||||
control_decision: Option<&GatewayControlDecision>,
|
||||
execution_path: &'static str,
|
||||
started_at: &Instant,
|
||||
request_permit: Option<AdmissionPermit>,
|
||||
) -> Response<Body> {
|
||||
response.headers_mut().insert(
|
||||
HeaderName::from_static(EXECUTION_PATH_HEADER),
|
||||
@@ -321,5 +437,15 @@ fn finalize_gateway_response(
|
||||
"gateway completed request"
|
||||
);
|
||||
|
||||
response
|
||||
record_shadow_result_non_blocking(
|
||||
state.clone(),
|
||||
trace_id,
|
||||
method,
|
||||
path_and_query,
|
||||
control_decision,
|
||||
execution_path,
|
||||
&response,
|
||||
);
|
||||
|
||||
maybe_hold_axum_response_permit(response, request_permit)
|
||||
}
|
||||
|
||||
@@ -2,5 +2,5 @@ mod gateway;
|
||||
|
||||
pub use gateway::{
|
||||
build_router, build_router_with_control, build_router_with_endpoints, build_router_with_state,
|
||||
serve_tcp, serve_tcp_with_endpoints, AppState, VideoTaskTruthSourceMode,
|
||||
serve_tcp, serve_tcp_with_endpoints, AppState, GatewayDataConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
|
||||
@@ -1,7 +1,14 @@
|
||||
use clap::{Parser, ValueEnum};
|
||||
use clap::{Args as ClapArgs, Parser, ValueEnum};
|
||||
use tracing::info;
|
||||
|
||||
use aether_gateway::{build_router_with_state, AppState, VideoTaskTruthSourceMode};
|
||||
use aether_data::postgres::PostgresPoolConfig;
|
||||
use aether_gateway::{
|
||||
build_router_with_state, AppState, GatewayDataConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
use aether_runtime::{
|
||||
init_service_runtime, DistributedConcurrencyGate, RedisDistributedConcurrencyConfig,
|
||||
ServiceRuntimeConfig,
|
||||
};
|
||||
|
||||
#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
|
||||
enum VideoTaskTruthSourceArg {
|
||||
@@ -20,6 +27,85 @@ impl From<VideoTaskTruthSourceArg> for VideoTaskTruthSourceMode {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(ClapArgs, Debug, Clone)]
|
||||
struct GatewayDataArgs {
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_POSTGRES_URL")]
|
||||
postgres_url: Option<String>,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_MIN_CONNECTIONS",
|
||||
default_value_t = 1
|
||||
)]
|
||||
postgres_min_connections: u32,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS",
|
||||
default_value_t = 20
|
||||
)]
|
||||
postgres_max_connections: u32,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_ACQUIRE_TIMEOUT_MS",
|
||||
default_value_t = 5_000
|
||||
)]
|
||||
postgres_acquire_timeout_ms: u64,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_IDLE_TIMEOUT_MS",
|
||||
default_value_t = 60_000
|
||||
)]
|
||||
postgres_idle_timeout_ms: u64,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_MAX_LIFETIME_MS",
|
||||
default_value_t = 1_800_000
|
||||
)]
|
||||
postgres_max_lifetime_ms: u64,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_STATEMENT_CACHE_CAPACITY",
|
||||
default_value_t = 100
|
||||
)]
|
||||
postgres_statement_cache_capacity: usize,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_REQUIRE_SSL",
|
||||
default_value_t = false
|
||||
)]
|
||||
postgres_require_ssl: bool,
|
||||
}
|
||||
|
||||
impl GatewayDataArgs {
|
||||
fn to_config(&self) -> GatewayDataConfig {
|
||||
let Some(database_url) = self
|
||||
.postgres_url
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return GatewayDataConfig::disabled();
|
||||
};
|
||||
|
||||
GatewayDataConfig::from_postgres_config(PostgresPoolConfig {
|
||||
database_url: database_url.to_string(),
|
||||
min_connections: self.postgres_min_connections,
|
||||
max_connections: self.postgres_max_connections,
|
||||
acquire_timeout_ms: self.postgres_acquire_timeout_ms,
|
||||
idle_timeout_ms: self.postgres_idle_timeout_ms,
|
||||
max_lifetime_ms: self.postgres_max_lifetime_ms,
|
||||
statement_cache_capacity: self.postgres_statement_cache_capacity,
|
||||
require_ssl: self.postgres_require_ssl,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(
|
||||
name = "aether-gateway",
|
||||
@@ -66,16 +152,50 @@ struct Args {
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_VIDEO_TASK_STORE_PATH")]
|
||||
video_task_store_path: Option<String>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_MAX_IN_FLIGHT_REQUESTS")]
|
||||
max_in_flight_requests: Option<usize>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DISTRIBUTED_REQUEST_LIMIT")]
|
||||
distributed_request_limit: Option<usize>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DISTRIBUTED_REQUEST_REDIS_URL")]
|
||||
distributed_request_redis_url: Option<String>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DISTRIBUTED_REQUEST_REDIS_KEY_PREFIX")]
|
||||
distributed_request_redis_key_prefix: Option<String>,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DISTRIBUTED_REQUEST_LEASE_TTL_MS",
|
||||
default_value_t = 30_000
|
||||
)]
|
||||
distributed_request_lease_ttl_ms: u64,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DISTRIBUTED_REQUEST_RENEW_INTERVAL_MS",
|
||||
default_value_t = 10_000
|
||||
)]
|
||||
distributed_request_renew_interval_ms: u64,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DISTRIBUTED_REQUEST_COMMAND_TIMEOUT_MS",
|
||||
default_value_t = 1_000
|
||||
)]
|
||||
distributed_request_command_timeout_ms: u64,
|
||||
|
||||
#[command(flatten)]
|
||||
data: GatewayDataArgs,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "aether_gateway=info".into()),
|
||||
)
|
||||
.init();
|
||||
init_service_runtime(ServiceRuntimeConfig::new(
|
||||
"aether-gateway",
|
||||
"aether_gateway=info",
|
||||
))?;
|
||||
|
||||
let args = Args::parse();
|
||||
let control_url = args
|
||||
@@ -97,13 +217,24 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
video_task_poller_interval_ms = args.video_task_poller_interval_ms,
|
||||
video_task_poller_batch_size = args.video_task_poller_batch_size,
|
||||
video_task_store_path = args.video_task_store_path.as_deref().unwrap_or("-"),
|
||||
max_in_flight_requests = args.max_in_flight_requests.unwrap_or_default(),
|
||||
distributed_request_limit = args.distributed_request_limit.unwrap_or_default(),
|
||||
distributed_request_redis_url = args
|
||||
.distributed_request_redis_url
|
||||
.as_deref()
|
||||
.unwrap_or("-"),
|
||||
data_postgres_url = args.data.postgres_url.as_deref().unwrap_or("-"),
|
||||
data_postgres_require_ssl = args.data.postgres_require_ssl,
|
||||
"aether-gateway started"
|
||||
);
|
||||
|
||||
let data_config = args.data.to_config();
|
||||
let mut state = AppState::new_with_executor(
|
||||
args.upstream,
|
||||
control_url.map(ToOwned::to_owned),
|
||||
executor_url.map(ToOwned::to_owned),
|
||||
)?
|
||||
.with_data_config(data_config)?
|
||||
.with_video_task_truth_source_mode(args.video_task_truth_source_mode.into());
|
||||
if matches!(
|
||||
args.video_task_truth_source_mode,
|
||||
@@ -122,6 +253,44 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
{
|
||||
state = state.with_video_task_store_path(path)?;
|
||||
}
|
||||
if let Some(limit) = args.max_in_flight_requests.filter(|limit| *limit > 0) {
|
||||
state = state.with_request_concurrency_limit(limit);
|
||||
}
|
||||
if let Some(limit) = args.distributed_request_limit.filter(|limit| *limit > 0) {
|
||||
let redis_url = args
|
||||
.distributed_request_redis_url
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| {
|
||||
std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
"AETHER_GATEWAY_DISTRIBUTED_REQUEST_REDIS_URL is required when distributed request limit is enabled",
|
||||
)
|
||||
})?;
|
||||
state =
|
||||
state.with_distributed_request_concurrency_gate(DistributedConcurrencyGate::new_redis(
|
||||
"gateway_requests_distributed",
|
||||
limit,
|
||||
RedisDistributedConcurrencyConfig {
|
||||
url: redis_url.to_string(),
|
||||
key_prefix: args
|
||||
.distributed_request_redis_key_prefix
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
lease_ttl_ms: args.distributed_request_lease_ttl_ms.max(1),
|
||||
renew_interval_ms: args.distributed_request_renew_interval_ms.max(1),
|
||||
command_timeout_ms: Some(args.distributed_request_command_timeout_ms.max(1)),
|
||||
},
|
||||
)?);
|
||||
}
|
||||
info!(
|
||||
has_data_backends = state.has_data_backends(),
|
||||
has_video_task_data_reader = state.has_video_task_data_reader(),
|
||||
"aether-gateway data layer configured"
|
||||
);
|
||||
let background_tasks = state.spawn_background_tasks();
|
||||
let listener = tokio::net::TcpListener::bind(&args.bind).await?;
|
||||
let router = build_router_with_state(state);
|
||||
|
||||
@@ -3,10 +3,14 @@ use std::collections::BTreeMap;
|
||||
use axum::body::Body;
|
||||
use axum::http::header::{HeaderName, HeaderValue};
|
||||
use axum::http::Response;
|
||||
use axum::http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::gateway::constants::*;
|
||||
use crate::gateway::headers::should_skip_response_header;
|
||||
use crate::gateway::{insert_header_if_missing, GatewayControlDecision, GatewayError};
|
||||
use crate::gateway::{
|
||||
insert_header_if_missing, GatewayControlDecision, GatewayError, GatewayLocalAuthRejection,
|
||||
};
|
||||
|
||||
pub(crate) fn build_client_response(
|
||||
upstream_response: reqwest::Response,
|
||||
@@ -90,3 +94,143 @@ pub(crate) fn build_client_response_from_parts(
|
||||
}
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
pub(crate) fn insert_candidate_id_header_if_present(
|
||||
headers: &mut http::HeaderMap,
|
||||
candidate_id: Option<&str>,
|
||||
) -> Result<(), GatewayError> {
|
||||
let Some(candidate_id) = candidate_id
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return Ok(());
|
||||
};
|
||||
insert_header_if_missing(headers, CONTROL_CANDIDATE_ID_HEADER, candidate_id)
|
||||
}
|
||||
|
||||
pub(crate) fn insert_request_id_header_if_present(
|
||||
headers: &mut http::HeaderMap,
|
||||
request_id: Option<&str>,
|
||||
) -> Result<(), GatewayError> {
|
||||
let Some(request_id) = request_id.map(str::trim).filter(|value| !value.is_empty()) else {
|
||||
return Ok(());
|
||||
};
|
||||
insert_header_if_missing(headers, CONTROL_REQUEST_ID_HEADER, request_id)
|
||||
}
|
||||
|
||||
pub(crate) fn attach_control_metadata_headers(
|
||||
mut response: Response<Body>,
|
||||
request_id: Option<&str>,
|
||||
candidate_id: Option<&str>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
insert_request_id_header_if_present(response.headers_mut(), request_id)?;
|
||||
insert_candidate_id_header_if_present(response.headers_mut(), candidate_id)?;
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_balance_denied_response(
|
||||
trace_id: &str,
|
||||
control_decision: Option<&GatewayControlDecision>,
|
||||
balance_remaining: Option<f64>,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let message = match balance_remaining {
|
||||
Some(remaining) => format!("余额不足(剩余: ${remaining:.2})"),
|
||||
None => "余额不足".to_string(),
|
||||
};
|
||||
let payload = json!({
|
||||
"error": {
|
||||
"type": "balance_exceeded",
|
||||
"message": message,
|
||||
"details": {
|
||||
"balance_type": "USD",
|
||||
"remaining": balance_remaining,
|
||||
}
|
||||
}
|
||||
});
|
||||
let body =
|
||||
serde_json::to_vec(&payload).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let headers = BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||||
build_client_response_from_parts(
|
||||
StatusCode::TOO_MANY_REQUESTS.as_u16(),
|
||||
&headers,
|
||||
Body::from(body),
|
||||
trace_id,
|
||||
control_decision,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_http_error_response(
|
||||
trace_id: &str,
|
||||
control_decision: Option<&GatewayControlDecision>,
|
||||
status_code: StatusCode,
|
||||
message: &str,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let payload = json!({
|
||||
"error": {
|
||||
"type": "http_error",
|
||||
"message": message,
|
||||
}
|
||||
});
|
||||
let body =
|
||||
serde_json::to_vec(&payload).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let headers = BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||||
build_client_response_from_parts(
|
||||
status_code.as_u16(),
|
||||
&headers,
|
||||
Body::from(body),
|
||||
trace_id,
|
||||
control_decision,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_auth_rejection_response(
|
||||
trace_id: &str,
|
||||
control_decision: Option<&GatewayControlDecision>,
|
||||
rejection: &GatewayLocalAuthRejection,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
match rejection {
|
||||
GatewayLocalAuthRejection::InvalidApiKey => build_local_http_error_response(
|
||||
trace_id,
|
||||
control_decision,
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"无效的API密钥",
|
||||
),
|
||||
GatewayLocalAuthRejection::LockedApiKey => build_local_http_error_response(
|
||||
trace_id,
|
||||
control_decision,
|
||||
StatusCode::FORBIDDEN,
|
||||
"该密钥已被管理员锁定,请联系管理员",
|
||||
),
|
||||
GatewayLocalAuthRejection::BalanceDenied { remaining } => {
|
||||
build_local_balance_denied_response(trace_id, control_decision, *remaining)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_overloaded_response(
|
||||
trace_id: &str,
|
||||
control_decision: Option<&GatewayControlDecision>,
|
||||
gate: &str,
|
||||
limit: usize,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let payload = json!({
|
||||
"error": {
|
||||
"type": "overloaded",
|
||||
"message": "服务繁忙,请稍后重试",
|
||||
"details": {
|
||||
"gate": gate,
|
||||
"limit": limit,
|
||||
}
|
||||
}
|
||||
});
|
||||
let body =
|
||||
serde_json::to_vec(&payload).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let headers = BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||||
build_client_response_from_parts(
|
||||
StatusCode::SERVICE_UNAVAILABLE.as_u16(),
|
||||
&headers,
|
||||
Body::from(body),
|
||||
trace_id,
|
||||
control_decision,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -2682,6 +2682,7 @@ mod tests {
|
||||
api_key_id: "key-123".to_string(),
|
||||
balance_remaining: None,
|
||||
access_allowed: true,
|
||||
local_rejection: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
11
crates/aether-http/Cargo.toml
Normal file
11
crates/aether-http/Cargo.toml
Normal file
@@ -0,0 +1,11 @@
|
||||
[package]
|
||||
name = "aether-http"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Shared HTTP client config and retry helpers for Aether Rust services"
|
||||
|
||||
[dependencies]
|
||||
reqwest.workspace = true
|
||||
serde.workspace = true
|
||||
75
crates/aether-http/src/client.rs
Normal file
75
crates/aether-http/src/client.rs
Normal file
@@ -0,0 +1,75 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use reqwest::header::HeaderMap;
|
||||
|
||||
use crate::HttpClientConfig;
|
||||
|
||||
pub fn apply_http_client_config(
|
||||
mut builder: reqwest::ClientBuilder,
|
||||
config: &HttpClientConfig,
|
||||
) -> reqwest::ClientBuilder {
|
||||
if config.use_rustls_tls {
|
||||
builder = builder.use_rustls_tls();
|
||||
}
|
||||
if let Some(timeout_ms) = config.connect_timeout_ms {
|
||||
builder = builder.connect_timeout(Duration::from_millis(timeout_ms));
|
||||
}
|
||||
if let Some(timeout_ms) = config.request_timeout_ms {
|
||||
builder = builder.timeout(Duration::from_millis(timeout_ms));
|
||||
}
|
||||
if let Some(timeout_ms) = config.pool_idle_timeout_ms {
|
||||
builder = builder.pool_idle_timeout(Duration::from_millis(timeout_ms));
|
||||
}
|
||||
if let Some(max_idle) = config.pool_max_idle_per_host {
|
||||
builder = builder.pool_max_idle_per_host(max_idle);
|
||||
}
|
||||
|
||||
builder = builder.tcp_keepalive(config.tcp_keepalive_ms.map(Duration::from_millis));
|
||||
builder = builder.tcp_nodelay(config.tcp_nodelay);
|
||||
|
||||
if config.http2_adaptive_window {
|
||||
builder = builder.http2_adaptive_window(true);
|
||||
}
|
||||
if let Some(user_agent) = &config.user_agent {
|
||||
builder = builder.user_agent(user_agent.clone());
|
||||
}
|
||||
|
||||
builder
|
||||
}
|
||||
|
||||
pub fn build_http_client(config: &HttpClientConfig) -> Result<reqwest::Client, reqwest::Error> {
|
||||
build_http_client_with_headers(config, HeaderMap::new())
|
||||
}
|
||||
|
||||
pub fn build_http_client_with_headers(
|
||||
config: &HttpClientConfig,
|
||||
default_headers: HeaderMap,
|
||||
) -> Result<reqwest::Client, reqwest::Error> {
|
||||
let mut builder = apply_http_client_config(reqwest::Client::builder(), config);
|
||||
if !default_headers.is_empty() {
|
||||
builder = builder.default_headers(default_headers);
|
||||
}
|
||||
builder.build()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use reqwest::header::{HeaderMap, HeaderValue};
|
||||
|
||||
use super::build_http_client_with_headers;
|
||||
use crate::HttpClientConfig;
|
||||
|
||||
#[test]
|
||||
fn builds_client_with_default_headers() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("x-test", HeaderValue::from_static("ok"));
|
||||
let config = HttpClientConfig {
|
||||
connect_timeout_ms: Some(100),
|
||||
request_timeout_ms: Some(500),
|
||||
..HttpClientConfig::default()
|
||||
};
|
||||
|
||||
let client = build_http_client_with_headers(&config, headers);
|
||||
assert!(client.is_ok());
|
||||
}
|
||||
}
|
||||
109
crates/aether-http/src/config.rs
Normal file
109
crates/aether-http/src/config.rs
Normal file
@@ -0,0 +1,109 @@
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
|
||||
pub struct HttpClientConfig {
|
||||
pub connect_timeout_ms: Option<u64>,
|
||||
pub request_timeout_ms: Option<u64>,
|
||||
pub pool_idle_timeout_ms: Option<u64>,
|
||||
pub pool_max_idle_per_host: Option<usize>,
|
||||
pub tcp_keepalive_ms: Option<u64>,
|
||||
pub tcp_nodelay: bool,
|
||||
pub http2_adaptive_window: bool,
|
||||
pub use_rustls_tls: bool,
|
||||
pub user_agent: Option<String>,
|
||||
}
|
||||
|
||||
impl Default for HttpClientConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
connect_timeout_ms: None,
|
||||
request_timeout_ms: None,
|
||||
pool_idle_timeout_ms: None,
|
||||
pool_max_idle_per_host: None,
|
||||
tcp_keepalive_ms: None,
|
||||
tcp_nodelay: true,
|
||||
http2_adaptive_window: false,
|
||||
use_rustls_tls: true,
|
||||
user_agent: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
|
||||
pub struct HttpRetryConfig {
|
||||
pub max_attempts: u32,
|
||||
pub base_delay_ms: u64,
|
||||
pub max_delay_ms: u64,
|
||||
}
|
||||
|
||||
impl Default for HttpRetryConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_attempts: 3,
|
||||
base_delay_ms: 200,
|
||||
max_delay_ms: 2_000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpRetryConfig {
|
||||
pub fn normalized(self) -> Self {
|
||||
let max_attempts = self.max_attempts.max(1);
|
||||
let base_delay_ms = self.base_delay_ms.max(1);
|
||||
let max_delay_ms = self.max_delay_ms.max(base_delay_ms);
|
||||
Self {
|
||||
max_attempts,
|
||||
base_delay_ms,
|
||||
max_delay_ms,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn delay_for_retry(self, retry_index: u32) -> std::time::Duration {
|
||||
let config = self.normalized();
|
||||
let factor = 2_u64.saturating_pow(retry_index.min(20));
|
||||
let delay_ms = config
|
||||
.base_delay_ms
|
||||
.saturating_mul(factor)
|
||||
.min(config.max_delay_ms);
|
||||
std::time::Duration::from_millis(delay_ms)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::HttpRetryConfig;
|
||||
|
||||
#[test]
|
||||
fn normalizes_retry_bounds() {
|
||||
let config = HttpRetryConfig {
|
||||
max_attempts: 0,
|
||||
base_delay_ms: 0,
|
||||
max_delay_ms: 5,
|
||||
}
|
||||
.normalized();
|
||||
|
||||
assert_eq!(config.max_attempts, 1);
|
||||
assert_eq!(config.base_delay_ms, 1);
|
||||
assert_eq!(config.max_delay_ms, 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn caps_exponential_retry_delay() {
|
||||
let config = HttpRetryConfig {
|
||||
max_attempts: 3,
|
||||
base_delay_ms: 100,
|
||||
max_delay_ms: 250,
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
config.delay_for_retry(0),
|
||||
std::time::Duration::from_millis(100)
|
||||
);
|
||||
assert_eq!(
|
||||
config.delay_for_retry(1),
|
||||
std::time::Duration::from_millis(200)
|
||||
);
|
||||
assert_eq!(
|
||||
config.delay_for_retry(2),
|
||||
std::time::Duration::from_millis(250)
|
||||
);
|
||||
}
|
||||
}
|
||||
7
crates/aether-http/src/lib.rs
Normal file
7
crates/aether-http/src/lib.rs
Normal file
@@ -0,0 +1,7 @@
|
||||
mod client;
|
||||
mod config;
|
||||
mod retry;
|
||||
|
||||
pub use client::{apply_http_client_config, build_http_client, build_http_client_with_headers};
|
||||
pub use config::{HttpClientConfig, HttpRetryConfig};
|
||||
pub use retry::jittered_delay_for_retry;
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user