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

feat: 引入 aether-runtime/cache/data/http/testkit 基础 crate，完善并发门控与审计系统

Browse Source 
新增 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:
fawney19
2026-03-24 15:12:56 +08:00
parent eaf8475f9e
commit b5a0070023
157 changed files with 22097 additions and 448 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
crates/aether-runtime/Cargo.toml Normal file
View File

@@ -0,0 +1,19 @@
[package]
name = "aether-runtime"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Shared runtime/bootstrap helpers for Aether Rust services"
[dependencies]
async-stream.workspace = true
axum = { version = "0.8" }
futures-util.workspace = true
redis.workspace = true
thiserror.workspace = true
tokio.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
url.workspace = true
uuid.workspace = true

134
crates/aether-runtime/src/admission.rs Normal file
View File

@@ -0,0 +1,134 @@
use async_stream::stream;
use axum::body::Body;
use axum::http::Response;
use futures_util::StreamExt;
use crate::concurrency::ConcurrencyPermit;
use crate::distributed::DistributedConcurrencyPermit;
#[derive(Debug)]
pub struct AdmissionPermit {
_local: Option<ConcurrencyPermit>,
_distributed: Option<DistributedConcurrencyPermit>,
}
impl AdmissionPermit {
pub fn from_parts(
local: Option<ConcurrencyPermit>,
distributed: Option<DistributedConcurrencyPermit>,
) -> Option<Self> {
if local.is_none() && distributed.is_none() {
None
} else {
Some(Self {
_local: local,
_distributed: distributed,
})
}
}
}
impl From<ConcurrencyPermit> for AdmissionPermit {
fn from(value: ConcurrencyPermit) -> Self {
Self {
_local: Some(value),
_distributed: None,
}
}
}
pub fn maybe_hold_axum_response_permit(
response: Response<Body>,
permit: Option<AdmissionPermit>,
) -> Response<Body> {
match permit {
Some(permit) => hold_axum_response_permit(response, permit),
None => response,
}
}
pub async fn hold_admission_permit_until<T, F>(permit: Option<AdmissionPermit>, future: F) -> T
where
F: std::future::Future<Output = T>,
{
let _permit = permit;
future.await
}
fn hold_axum_response_permit(response: Response<Body>, permit: AdmissionPermit) -> Response<Body> {
let (parts, body) = response.into_parts();
let stream = stream! {
let _permit = permit;
let mut body_stream = body.into_data_stream();
while let Some(item) = body_stream.next().await {
yield item;
}
};
Response::from_parts(parts, Body::from_stream(stream))
}
#[cfg(test)]
mod tests {
use super::{hold_admission_permit_until, maybe_hold_axum_response_permit, AdmissionPermit};
use crate::{ConcurrencyGate, DistributedConcurrencyGate};
use axum::body::{to_bytes, Body};
use axum::http::Response;
#[tokio::test]
async fn holds_permit_until_response_body_is_consumed() {
let gate = ConcurrencyGate::new("test", 1);
let permit = gate.try_acquire().expect("first permit");
let response = Response::new(Body::from_stream(
async_stream::stream! { yield Ok::<_, std::convert::Infallible>(axum::body::Bytes::from_static(b"ok")); },
));
let wrapped = maybe_hold_axum_response_permit(response, Some(permit.into()));
assert_eq!(gate.snapshot().in_flight, 1);
assert!(gate.try_acquire().is_err(), "permit should still be held");
let body = to_bytes(wrapped.into_body(), usize::MAX)
.await
.expect("body should drain");
assert_eq!(body.as_ref(), b"ok");
assert_eq!(gate.snapshot().in_flight, 0);
}
#[tokio::test]
async fn holds_combined_local_and_distributed_permit_until_future_finishes() {
let local_gate = ConcurrencyGate::new("local", 1);
let distributed_gate = DistributedConcurrencyGate::new_in_memory("distributed", 1);
let local = local_gate.try_acquire().expect("local permit");
let distributed = distributed_gate
.try_acquire()
.await
.expect("distributed permit");
let task = tokio::spawn(hold_admission_permit_until(
AdmissionPermit::from_parts(Some(local), Some(distributed)),
async {
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
},
));
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
assert!(
local_gate.try_acquire().is_err(),
"local permit should still be held"
);
assert!(
distributed_gate.try_acquire().await.is_err(),
"distributed permit should still be held"
);
task.await.expect("task should complete");
assert_eq!(local_gate.snapshot().in_flight, 0);
assert_eq!(
distributed_gate
.snapshot()
.await
.expect("snapshot should build")
.in_flight,
0
);
}
}

8
crates/aether-runtime/src/bootstrap.rs Normal file
View File

@@ -0,0 +1,8 @@
use crate::config::ServiceRuntimeConfig;
use crate::error::RuntimeBootstrapError;
pub fn init_service_runtime(config: ServiceRuntimeConfig) -> Result<(), RuntimeBootstrapError> {
crate::tracing::init_tracing(config)?;
crate::metrics::init_metrics(config);
Ok(())
}

200
crates/aether-runtime/src/concurrency.rs Normal file
View File

@@ -0,0 +1,200 @@
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use crate::metrics::{MetricKind, MetricLabel, MetricSample};
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum ConcurrencyError {
#[error("concurrency gate {gate} is saturated at {limit}")]
Saturated { gate: &'static str, limit: usize },
#[error("concurrency gate {gate} is closed")]
Closed { gate: &'static str },
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ConcurrencySnapshot {
pub limit: usize,
pub in_flight: usize,
pub available_permits: usize,
pub high_watermark: usize,
pub rejected: u64,
}
impl ConcurrencySnapshot {
pub fn to_metric_samples(&self, gate: &'static str) -> Vec<MetricSample> {
let labels = vec![MetricLabel::new("gate", gate)];
vec![
MetricSample::new(
"concurrency_in_flight",
"Current number of in-flight operations guarded by the concurrency gate.",
MetricKind::Gauge,
self.in_flight as u64,
)
.with_labels(labels.clone()),
MetricSample::new(
"concurrency_available_permits",
"Currently available permits for the concurrency gate.",
MetricKind::Gauge,
self.available_permits as u64,
)
.with_labels(labels.clone()),
MetricSample::new(
"concurrency_high_watermark",
"Highest observed in-flight count for the concurrency gate.",
MetricKind::Gauge,
self.high_watermark as u64,
)
.with_labels(labels.clone()),
MetricSample::new(
"concurrency_rejected_total",
"Number of operations rejected by the concurrency gate.",
MetricKind::Counter,
self.rejected,
)
.with_labels(labels),
]
}
}
#[derive(Debug)]
struct ConcurrencyState {
gate: &'static str,
limit: usize,
semaphore: Arc<Semaphore>,
in_flight: AtomicUsize,
high_watermark: AtomicUsize,
rejected: AtomicU64,
}
#[derive(Debug, Clone)]
pub struct ConcurrencyGate {
state: Arc<ConcurrencyState>,
}
impl ConcurrencyGate {
pub fn new(gate: &'static str, limit: usize) -> Self {
assert!(limit > 0, "concurrency gate limit must be positive");
Self {
state: Arc::new(ConcurrencyState {
gate,
limit,
semaphore: Arc::new(Semaphore::new(limit)),
in_flight: AtomicUsize::new(0),
high_watermark: AtomicUsize::new(0),
rejected: AtomicU64::new(0),
}),
}
}
pub async fn acquire(&self) -> Result<ConcurrencyPermit, ConcurrencyError> {
let permit = self
.state
.semaphore
.clone()
.acquire_owned()
.await
.map_err(|_| ConcurrencyError::Closed {
gate: self.state.gate,
})?;
Ok(ConcurrencyPermit::new(self.state.clone(), permit))
}
pub fn try_acquire(&self) -> Result<ConcurrencyPermit, ConcurrencyError> {
match self.state.semaphore.clone().try_acquire_owned() {
Ok(permit) => Ok(ConcurrencyPermit::new(self.state.clone(), permit)),
Err(tokio::sync::TryAcquireError::NoPermits) => {
self.state.rejected.fetch_add(1, Ordering::Relaxed);
Err(ConcurrencyError::Saturated {
gate: self.state.gate,
limit: self.state.limit,
})
}
Err(tokio::sync::TryAcquireError::Closed) => Err(ConcurrencyError::Closed {
gate: self.state.gate,
}),
}
}
pub fn snapshot(&self) -> ConcurrencySnapshot {
ConcurrencySnapshot {
limit: self.state.limit,
in_flight: self.state.in_flight.load(Ordering::Relaxed),
available_permits: self.state.semaphore.available_permits(),
high_watermark: self.state.high_watermark.load(Ordering::Relaxed),
rejected: self.state.rejected.load(Ordering::Relaxed),
}
}
}
#[derive(Debug)]
pub struct ConcurrencyPermit {
state: Arc<ConcurrencyState>,
_permit: OwnedSemaphorePermit,
}
impl ConcurrencyPermit {
fn new(state: Arc<ConcurrencyState>, permit: OwnedSemaphorePermit) -> Self {
let in_flight = state.in_flight.fetch_add(1, Ordering::AcqRel) + 1;
let mut observed = state.high_watermark.load(Ordering::Acquire);
while in_flight > observed {
match state.high_watermark.compare_exchange_weak(
observed,
in_flight,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => break,
Err(next) => observed = next,
}
}
Self {
state,
_permit: permit,
}
}
}
impl Drop for ConcurrencyPermit {
fn drop(&mut self) {
self.state.in_flight.fetch_sub(1, Ordering::AcqRel);
}
}
#[cfg(test)]
mod tests {
use super::{ConcurrencyError, ConcurrencyGate};
#[tokio::test]
async fn tracks_in_flight_and_high_watermark() {
let gate = ConcurrencyGate::new("test", 2);
let permit_a = gate.acquire().await.expect("permit a");
let permit_b = gate.acquire().await.expect("permit b");
let snapshot = gate.snapshot();
assert_eq!(snapshot.in_flight, 2);
assert_eq!(snapshot.high_watermark, 2);
assert_eq!(snapshot.available_permits, 0);
drop((permit_a, permit_b));
assert_eq!(gate.snapshot().in_flight, 0);
}
#[test]
fn rejects_when_saturated() {
let gate = ConcurrencyGate::new("test", 1);
let _permit = gate.try_acquire().expect("first permit");
let error = gate.try_acquire().expect_err("second permit should fail");
assert_eq!(
error,
ConcurrencyError::Saturated {
gate: "test",
limit: 1,
}
);
assert_eq!(gate.snapshot().rejected, 1);
}
}

28
crates/aether-runtime/src/config.rs Normal file
View File

@@ -0,0 +1,28 @@
use crate::observability::ServiceObservabilityConfig;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ServiceRuntimeConfig {
pub service_name: &'static str,
pub default_log_filter: &'static str,
pub observability: ServiceObservabilityConfig,
}
impl ServiceRuntimeConfig {
pub const fn new(service_name: &'static str, default_log_filter: &'static str) -> Self {
Self {
service_name,
default_log_filter,
observability: ServiceObservabilityConfig::new(crate::LogFormat::Pretty, service_name),
}
}
pub const fn with_log_format(mut self, log_format: crate::LogFormat) -> Self {
self.observability.log_format = log_format;
self
}
pub const fn with_metrics_namespace(mut self, metrics_namespace: &'static str) -> Self {
self.observability.metrics_namespace = metrics_namespace;
self
}
}

644
crates/aether-runtime/src/distributed.rs Normal file
View File

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

5
crates/aether-runtime/src/error.rs Normal file
View File

@@ -0,0 +1,5 @@
#[derive(Debug, thiserror::Error)]
pub enum RuntimeBootstrapError {
#[error("failed to initialize tracing: {0}")]
Tracing(String),
}

31
crates/aether-runtime/src/lib.rs Normal file
View File

@@ -0,0 +1,31 @@
pub mod admission;
mod bootstrap;
pub mod concurrency;
mod config;
pub mod distributed;
mod error;
pub mod metrics;
mod observability;
pub mod queue;
pub mod shutdown;
pub mod task;
mod tracing;
pub use admission::{
hold_admission_permit_until, maybe_hold_axum_response_permit, AdmissionPermit,
};
pub use bootstrap::init_service_runtime;
pub use concurrency::{ConcurrencyError, ConcurrencyGate, ConcurrencyPermit, ConcurrencySnapshot};
pub use config::ServiceRuntimeConfig;
pub use distributed::{
DistributedConcurrencyError, DistributedConcurrencyGate, DistributedConcurrencyPermit,
DistributedConcurrencySnapshot, RedisDistributedConcurrencyConfig,
};
pub use error::RuntimeBootstrapError;
pub use metrics::{prometheus_response, service_up_sample, MetricKind, MetricLabel, MetricSample};
pub use observability::ServiceObservabilityConfig;
pub use queue::{
bounded_queue, BoundedQueueReceiver, BoundedQueueSender, QueueSendError, QueueSnapshot,
};
pub use shutdown::wait_for_shutdown_signal;
pub use tracing::{init_reloadable_tracing, LogFormat, LogReloader};

186
crates/aether-runtime/src/metrics.rs Normal file
View File

@@ -0,0 +1,186 @@
use crate::config::ServiceRuntimeConfig;
use axum::body::Body;
use axum::http::header::{HeaderValue, CONTENT_TYPE};
use axum::http::Response;
static METRICS_NAMESPACE: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MetricKind {
Counter,
Gauge,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MetricLabel {
pub key: &'static str,
pub value: String,
}
impl MetricLabel {
pub fn new(key: &'static str, value: impl Into<String>) -> Self {
Self {
key,
value: value.into(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MetricSample {
pub name: &'static str,
pub help: &'static str,
pub kind: MetricKind,
pub value: u64,
pub labels: Vec<MetricLabel>,
}
impl MetricSample {
pub fn new(name: &'static str, help: &'static str, kind: MetricKind, value: u64) -> Self {
Self {
name,
help,
kind,
value,
labels: Vec::new(),
}
}
pub fn with_labels(mut self, labels: Vec<MetricLabel>) -> Self {
self.labels = labels;
self
}
}
pub fn init_metrics(config: ServiceRuntimeConfig) {
let _ = METRICS_NAMESPACE.set(config.observability.metrics_namespace);
}
pub fn metrics_namespace() -> Option<&'static str> {
METRICS_NAMESPACE.get().copied()
}
pub fn render_prometheus_text(samples: &[MetricSample]) -> String {
let mut body = String::new();
let namespace = metrics_namespace();
for sample in samples {
let metric_name = format_metric_name(namespace, sample.name);
body.push_str(&format!("# HELP {} {}\n", metric_name, sample.help));
body.push_str(&format!(
"# TYPE {} {}\n",
metric_name,
match sample.kind {
MetricKind::Counter => "counter",
MetricKind::Gauge => "gauge",
}
));
body.push_str(&metric_name);
if !sample.labels.is_empty() {
body.push('{');
for (index, label) in sample.labels.iter().enumerate() {
if index > 0 {
body.push(',');
}
body.push_str(label.key);
body.push_str("=\"");
body.push_str(&escape_prometheus_label(&label.value));
body.push('"');
}
body.push('}');
}
body.push(' ');
body.push_str(&sample.value.to_string());
body.push('\n');
}
body
}
pub fn prometheus_response(samples: &[MetricSample]) -> Response<Body> {
let mut response = Response::new(Body::from(render_prometheus_text(samples)));
response.headers_mut().insert(
CONTENT_TYPE,
HeaderValue::from_static("text/plain; version=0.0.4; charset=utf-8"),
);
response
}
pub fn service_up_sample(service: &'static str) -> MetricSample {
MetricSample::new(
"service_up",
"Whether the service process is currently up.",
MetricKind::Gauge,
1,
)
.with_labels(vec![MetricLabel::new("service", service)])
}
fn format_metric_name(namespace: Option<&str>, name: &str) -> String {
match namespace {
Some(namespace) if !namespace.is_empty() => format!("{}_{}", namespace, name),
_ => name.to_string(),
}
}
fn escape_prometheus_label(value: &str) -> String {
value
.replace('\\', "\\\\")
.replace('\n', "\\n")
.replace('"', "\\\"")
}
#[cfg(test)]
mod tests {
use super::{
prometheus_response, render_prometheus_text, service_up_sample, MetricKind, MetricLabel,
MetricSample,
};
use axum::body::to_bytes;
#[test]
fn renders_prometheus_samples_with_labels() {
let text = render_prometheus_text(&[MetricSample::new(
"queue_depth",
"Current queue depth",
MetricKind::Gauge,
3,
)
.with_labels(vec![MetricLabel::new("queue", "proxy_writer")])]);
assert!(text.contains("# HELP queue_depth Current queue depth"));
assert!(text.contains("# TYPE queue_depth gauge"));
assert!(text.contains("queue_depth{queue=\"proxy_writer\"} 3"));
}
#[test]
fn escapes_prometheus_labels() {
let text = render_prometheus_text(&[MetricSample::new(
"errors_total",
"Errors",
MetricKind::Counter,
1,
)
.with_labels(vec![MetricLabel::new("message", "bad\"line\nx")])]);
assert!(text.contains("message=\"bad\\\"line\\nx\""));
}
#[tokio::test]
async fn builds_prometheus_http_response() {
let response = prometheus_response(&[service_up_sample("gateway")]);
let content_type = response
.headers()
.get(axum::http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok());
assert_eq!(
content_type,
Some("text/plain; version=0.0.4; charset=utf-8")
);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let text = String::from_utf8(body.to_vec()).expect("body should be utf8");
assert!(text.contains("service_up{service=\"gateway\"} 1"));
}
}

16
crates/aether-runtime/src/observability.rs Normal file
View File

@@ -0,0 +1,16 @@
use crate::tracing::LogFormat;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ServiceObservabilityConfig {
pub log_format: LogFormat,
pub metrics_namespace: &'static str,
}
impl ServiceObservabilityConfig {
pub const fn new(log_format: LogFormat, metrics_namespace: &'static str) -> Self {
Self {
log_format,
metrics_namespace,
}
}
}

240
crates/aether-runtime/src/queue.rs Normal file
View File

@@ -0,0 +1,240 @@
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::sync::mpsc;
use crate::metrics::{MetricKind, MetricLabel, MetricSample};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct QueueSnapshot {
pub capacity: usize,
pub depth: usize,
pub high_watermark: usize,
pub enqueued_total: u64,
pub rejected_full_total: u64,
pub rejected_closed_total: u64,
}
impl QueueSnapshot {
pub fn to_metric_samples(&self, queue: &'static str) -> Vec<MetricSample> {
let labels = vec![MetricLabel::new("queue", queue)];
vec![
MetricSample::new(
"queue_depth",
"Current number of items buffered in the queue.",
MetricKind::Gauge,
self.depth as u64,
)
.with_labels(labels.clone()),
MetricSample::new(
"queue_high_watermark",
"Highest observed queue depth.",
MetricKind::Gauge,
self.high_watermark as u64,
)
.with_labels(labels.clone()),
MetricSample::new(
"queue_enqueued_total",
"Total number of items successfully enqueued.",
MetricKind::Counter,
self.enqueued_total,
)
.with_labels(labels.clone()),
MetricSample::new(
"queue_rejected_full_total",
"Total number of items rejected because the queue was full.",
MetricKind::Counter,
self.rejected_full_total,
)
.with_labels(labels.clone()),
MetricSample::new(
"queue_rejected_closed_total",
"Total number of items rejected because the queue was closed.",
MetricKind::Counter,
self.rejected_closed_total,
)
.with_labels(labels),
]
}
}
#[derive(Debug)]
struct QueueState {
capacity: usize,
depth: AtomicUsize,
high_watermark: AtomicUsize,
enqueued_total: AtomicU64,
rejected_full_total: AtomicU64,
rejected_closed_total: AtomicU64,
}
#[derive(Debug)]
pub enum QueueSendError<T> {
Full(T),
Closed(T),
}
#[derive(Debug, Clone)]
pub struct BoundedQueueSender<T> {
inner: mpsc::Sender<T>,
state: Arc<QueueState>,
}
#[derive(Debug)]
pub struct BoundedQueueReceiver<T> {
inner: mpsc::Receiver<T>,
state: Arc<QueueState>,
}
pub fn bounded_queue<T>(capacity: usize) -> (BoundedQueueSender<T>, BoundedQueueReceiver<T>) {
assert!(capacity > 0, "bounded queue capacity must be positive");
let (tx, rx) = mpsc::channel(capacity);
let state = Arc::new(QueueState {
capacity,
depth: AtomicUsize::new(0),
high_watermark: AtomicUsize::new(0),
enqueued_total: AtomicU64::new(0),
rejected_full_total: AtomicU64::new(0),
rejected_closed_total: AtomicU64::new(0),
});
(
BoundedQueueSender {
inner: tx,
state: state.clone(),
},
BoundedQueueReceiver { inner: rx, state },
)
}
impl<T> BoundedQueueSender<T> {
pub async fn send(&self, value: T) -> Result<(), QueueSendError<T>> {
let permit = match self.inner.reserve().await {
Ok(permit) => permit,
Err(_) => {
self.state
.rejected_closed_total
.fetch_add(1, Ordering::Relaxed);
return Err(QueueSendError::Closed(value));
}
};
self.record_enqueue();
permit.send(value);
Ok(())
}
pub fn try_send(&self, value: T) -> Result<(), QueueSendError<T>> {
let permit = match self.inner.try_reserve() {
Ok(permit) => permit,
Err(mpsc::error::TrySendError::Full(_)) => {
self.state
.rejected_full_total
.fetch_add(1, Ordering::Relaxed);
return Err(QueueSendError::Full(value));
}
Err(mpsc::error::TrySendError::Closed(_)) => {
self.state
.rejected_closed_total
.fetch_add(1, Ordering::Relaxed);
return Err(QueueSendError::Closed(value));
}
};
self.record_enqueue();
permit.send(value);
Ok(())
}
pub fn snapshot(&self) -> QueueSnapshot {
QueueSnapshot {
capacity: self.state.capacity,
depth: self.state.depth.load(Ordering::Relaxed),
high_watermark: self.state.high_watermark.load(Ordering::Relaxed),
enqueued_total: self.state.enqueued_total.load(Ordering::Relaxed),
rejected_full_total: self.state.rejected_full_total.load(Ordering::Relaxed),
rejected_closed_total: self.state.rejected_closed_total.load(Ordering::Relaxed),
}
}
pub fn capacity(&self) -> usize {
self.state.capacity
}
fn record_enqueue(&self) {
let depth = self.state.depth.fetch_add(1, Ordering::AcqRel) + 1;
self.state.enqueued_total.fetch_add(1, Ordering::Relaxed);
let mut observed = self.state.high_watermark.load(Ordering::Acquire);
while depth > observed {
match self.state.high_watermark.compare_exchange_weak(
observed,
depth,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => break,
Err(next) => observed = next,
}
}
}
}
impl<T> BoundedQueueReceiver<T> {
pub async fn recv(&mut self) -> Option<T> {
let value = self.inner.recv().await?;
self.state.depth.fetch_sub(1, Ordering::AcqRel);
Some(value)
}
pub fn try_recv(&mut self) -> Result<T, mpsc::error::TryRecvError> {
let value = self.inner.try_recv()?;
self.state.depth.fetch_sub(1, Ordering::AcqRel);
Ok(value)
}
}
#[cfg(test)]
mod tests {
use super::{bounded_queue, QueueSendError};
#[tokio::test]
async fn tracks_queue_depth_and_high_watermark() {
let (tx, mut rx) = bounded_queue::<u32>(2);
tx.send(1).await.expect("enqueue 1");
tx.send(2).await.expect("enqueue 2");
let snapshot = tx.snapshot();
assert_eq!(snapshot.depth, 2);
assert_eq!(snapshot.high_watermark, 2);
assert_eq!(snapshot.enqueued_total, 2);
assert_eq!(rx.recv().await, Some(1));
assert_eq!(tx.snapshot().depth, 1);
}
#[test]
fn counts_full_rejections() {
let (tx, _rx) = bounded_queue::<u32>(1);
tx.try_send(1).expect("first send should work");
let error = tx.try_send(2).expect_err("second send should fail");
assert!(matches!(error, QueueSendError::Full(2)));
assert_eq!(tx.snapshot().rejected_full_total, 1);
}
#[tokio::test]
async fn send_does_not_underflow_depth_when_receiver_races() {
let (tx, mut rx) = bounded_queue::<u32>(1);
let receiver = tokio::spawn(async move {
for _ in 0..256 {
assert!(rx.recv().await.is_some());
}
});
for value in 0..256 {
tx.send(value).await.expect("send should succeed");
}
receiver.await.expect("receiver task should join");
let snapshot = tx.snapshot();
assert_eq!(snapshot.depth, 0);
assert!(snapshot.high_watermark <= 1);
assert_eq!(snapshot.enqueued_total, 256);
}
}

15
crates/aether-runtime/src/shutdown.rs Normal file
View File

@@ -0,0 +1,15 @@
#[cfg(unix)]
pub async fn wait_for_shutdown_signal() -> Result<(), std::io::Error> {
use tokio::signal::unix::{signal, SignalKind};
let mut terminate = signal(SignalKind::terminate())?;
tokio::select! {
_ = tokio::signal::ctrl_c() => Ok(()),
_ = terminate.recv() => Ok(()),
}
}
#[cfg(not(unix))]
pub async fn wait_for_shutdown_signal() -> Result<(), std::io::Error> {
tokio::signal::ctrl_c().await
}

12
crates/aether-runtime/src/task.rs Normal file
View File

@@ -0,0 +1,12 @@
use std::future::Future;
pub fn spawn_named<F>(task_name: &'static str, future: F) -> tokio::task::JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
tokio::spawn(async move {
tracing::debug!(task = task_name, "spawned runtime task");
future.await
})
}

63
crates/aether-runtime/src/tracing.rs Normal file
View File

@@ -0,0 +1,63 @@
use std::sync::OnceLock;
use tracing_subscriber::prelude::*;
use tracing_subscriber::EnvFilter;
use crate::config::ServiceRuntimeConfig;
use crate::error::RuntimeBootstrapError;
static TRACING_INIT: OnceLock<Result<(), String>> = OnceLock::new();
pub type LogReloader = Box<dyn Fn(&str) + Send + Sync>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LogFormat {
Pretty,
Json,
}
pub(crate) fn init_tracing(config: ServiceRuntimeConfig) -> Result<(), RuntimeBootstrapError> {
TRACING_INIT
.get_or_init(|| {
let filter = EnvFilter::try_from_default_env()
.unwrap_or_else(|_| config.default_log_filter.into());
match config.observability.log_format {
LogFormat::Pretty => tracing_subscriber::fmt().with_env_filter(filter).try_init(),
LogFormat::Json => tracing_subscriber::fmt()
.json()
.with_env_filter(filter)
.try_init(),
}
.map_err(|err| err.to_string())
})
.clone()
.map_err(RuntimeBootstrapError::Tracing)
}
pub fn init_reloadable_tracing(
initial_filter: &str,
format: LogFormat,
) -> Result<LogReloader, RuntimeBootstrapError> {
use tracing_subscriber::reload;
let filter = EnvFilter::try_new(initial_filter).unwrap_or_else(|_| EnvFilter::new("info"));
let (filter_layer, reload_handle) = reload::Layer::new(filter);
match format {
LogFormat::Pretty => tracing_subscriber::registry()
.with(filter_layer)
.with(tracing_subscriber::fmt::layer())
.try_init(),
LogFormat::Json => tracing_subscriber::registry()
.with(filter_layer)
.with(tracing_subscriber::fmt::layer().json())
.try_init(),
}
.map_err(|err| RuntimeBootstrapError::Tracing(err.to_string()))?;
Ok(Box::new(move |level: &str| {
if let Ok(new_filter) = EnvFilter::try_new(level) {
let _ = reload_handle.modify(|filter| *filter = new_filter);
}
}))
}