mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
fix(runtime-state): govern redis connections
This commit is contained in:
@@ -59,6 +59,11 @@ pub(super) async fn build_admin_monitoring_redis_cache_categories_response(
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let mut categories = Vec::with_capacity(ADMIN_MONITORING_REDIS_CACHE_CATEGORIES.len());
|
||||
let mut total_keys = 0usize;
|
||||
let diagnostics = state
|
||||
.runtime_state()
|
||||
.redis_diagnostics()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(format!("redis diagnostics failed: {err}")))?;
|
||||
|
||||
for (key, name, pattern, description) in ADMIN_MONITORING_REDIS_CACHE_CATEGORIES {
|
||||
let count = list_admin_monitoring_namespaced_keys(state, pattern)
|
||||
@@ -81,6 +86,7 @@ pub(super) async fn build_admin_monitoring_redis_cache_categories_response(
|
||||
"backend": state.runtime_state().backend_kind().as_str(),
|
||||
"categories": categories,
|
||||
"total_keys": total_keys,
|
||||
"diagnostics": diagnostics,
|
||||
}
|
||||
}))
|
||||
.into_response())
|
||||
|
||||
@@ -1194,6 +1194,7 @@ async fn admin_monitoring_redis_keys_returns_local_payload_without_redis() {
|
||||
assert_eq!(payload["data"]["available"], json!(true));
|
||||
assert_eq!(payload["data"]["backend"], json!("memory"));
|
||||
assert_eq!(payload["data"]["total_keys"], json!(0));
|
||||
assert_eq!(payload["data"]["diagnostics"], serde_json::Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -497,6 +497,7 @@ impl GatewayUsageArgs {
|
||||
fn to_config(&self) -> UsageRuntimeConfig {
|
||||
UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
queue_terminal_events: true,
|
||||
stream_key: self.queue_stream_key.trim().to_string(),
|
||||
consumer_group: self.queue_group.trim().to_string(),
|
||||
dlq_stream_key: self.queue_dlq_stream_key.trim().to_string(),
|
||||
|
||||
@@ -150,11 +150,20 @@ impl FrontdoorUserRpmLimiter {
|
||||
scope_key: &str,
|
||||
bucket: u64,
|
||||
) -> Result<u32, GatewayError> {
|
||||
if !state.runtime_state.is_memory() {
|
||||
let raw = state.runtime_state.kv_get(scope_key).await.map_err(|err| {
|
||||
GatewayError::Internal(format!("frontdoor user rpm runtime read failed: {err}"))
|
||||
})?;
|
||||
return Ok(raw.and_then(|value| value.parse::<u32>().ok()).unwrap_or(0));
|
||||
match state.runtime_state.kv_get(scope_key).await {
|
||||
Ok(raw) => return Ok(raw.and_then(|value| value.parse::<u32>().ok()).unwrap_or(0)),
|
||||
Err(err) if !self.config.allow_local_fallback() => {
|
||||
return Err(GatewayError::Internal(format!(
|
||||
"frontdoor user rpm runtime read failed: {err}"
|
||||
)));
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
error = ?err,
|
||||
scope_key = %scope_key,
|
||||
"frontdoor user rpm runtime count read failed; using local fallback"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let counts = self.memory_counts.lock().await;
|
||||
@@ -187,24 +196,22 @@ impl FrontdoorUserRpmLimiter {
|
||||
return Ok(FrontdoorUserRpmOutcome::Allowed);
|
||||
}
|
||||
|
||||
if !state.runtime_state.is_memory() {
|
||||
match self.check_and_consume_runtime(state, &plan).await {
|
||||
Ok(outcome) => return Ok(outcome),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
error = ?err,
|
||||
user_rpm_key = %plan.user_rpm_key,
|
||||
key_rpm_key = %plan.key_rpm_key,
|
||||
"frontdoor user rpm redis check failed"
|
||||
);
|
||||
if self.config.fail_open() {
|
||||
return Ok(FrontdoorUserRpmOutcome::NotApplicable);
|
||||
}
|
||||
if !self.config.allow_local_fallback() {
|
||||
return Err(GatewayError::Internal(
|
||||
"frontdoor user rpm runtime backend is unavailable and local fallback is disabled for the current deployment mode".to_string(),
|
||||
));
|
||||
}
|
||||
match self.check_and_consume_runtime(state, &plan).await {
|
||||
Ok(outcome) => return Ok(outcome),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
error = ?err,
|
||||
user_rpm_key = %plan.user_rpm_key,
|
||||
key_rpm_key = %plan.key_rpm_key,
|
||||
"frontdoor user rpm runtime check failed"
|
||||
);
|
||||
if self.config.fail_open() {
|
||||
return Ok(FrontdoorUserRpmOutcome::NotApplicable);
|
||||
}
|
||||
if !self.config.allow_local_fallback() {
|
||||
return Err(GatewayError::Internal(
|
||||
"frontdoor user rpm runtime backend is unavailable and local fallback is disabled for the current deployment mode".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -605,7 +612,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn limiter_rejects_missing_shared_runtime_when_local_fallback_disabled() {
|
||||
async fn limiter_uses_runtime_state_when_local_fallback_disabled() {
|
||||
let limiter = FrontdoorUserRpmLimiter::new(
|
||||
FrontdoorUserRpmConfig::new(60, 120, false).with_local_fallback(false),
|
||||
);
|
||||
@@ -626,15 +633,22 @@ mod tests {
|
||||
});
|
||||
let state = AppState::new().expect("state should build for tests");
|
||||
|
||||
let err = limiter
|
||||
let first = limiter
|
||||
.check_and_consume(&state, Some(&decision))
|
||||
.await
|
||||
.expect_err("missing shared runtime should fail in strict mode");
|
||||
match err {
|
||||
crate::GatewayError::Internal(message) => {
|
||||
assert!(message.contains("requires shared runtime state"));
|
||||
.expect("runtime check should succeed");
|
||||
assert_eq!(first, FrontdoorUserRpmOutcome::Allowed);
|
||||
|
||||
let second = limiter
|
||||
.check_and_consume(&state, Some(&decision))
|
||||
.await
|
||||
.expect("runtime check should succeed");
|
||||
match second {
|
||||
FrontdoorUserRpmOutcome::Rejected(rejection) => {
|
||||
assert_eq!(rejection.scope, "user");
|
||||
assert_eq!(rejection.limit, 1);
|
||||
}
|
||||
other => panic!("expected internal error, got {other:?}"),
|
||||
other => panic!("expected rejection, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -89,24 +89,17 @@ impl AppState {
|
||||
fn usage_worker_queue_for(
|
||||
runtime_state: &Arc<RuntimeState>,
|
||||
) -> Option<Arc<dyn RuntimeQueueStore>> {
|
||||
if runtime_state.is_redis() {
|
||||
let queue: Arc<dyn RuntimeQueueStore> = runtime_state.clone();
|
||||
Some(queue)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
let queue: Arc<dyn RuntimeQueueStore> = runtime_state.clone();
|
||||
Some(queue)
|
||||
}
|
||||
|
||||
fn spawn_scheduler_affinity_redis_write(
|
||||
fn spawn_scheduler_affinity_runtime_write(
|
||||
&self,
|
||||
cache_key: &str,
|
||||
target: &SchedulerAffinityTarget,
|
||||
ttl: Duration,
|
||||
epoch: u64,
|
||||
) {
|
||||
if self.runtime_state.is_memory() {
|
||||
return;
|
||||
}
|
||||
let Ok(handle) = tokio::runtime::Handle::try_current() else {
|
||||
return;
|
||||
};
|
||||
@@ -1088,7 +1081,7 @@ impl AppState {
|
||||
if self.scheduler_affinity_epoch() != epoch {
|
||||
return false;
|
||||
}
|
||||
self.spawn_scheduler_affinity_redis_write(cache_key, &target, ttl, epoch);
|
||||
self.spawn_scheduler_affinity_runtime_write(cache_key, &target, ttl, epoch);
|
||||
self.scheduler_affinity_cache.insert_for_epoch(
|
||||
cache_key.to_string(),
|
||||
target,
|
||||
|
||||
@@ -104,9 +104,16 @@ fn runtime_state_owns_redis_runtime_boundaries() {
|
||||
let mut violations = Vec::new();
|
||||
for root in [
|
||||
"apps/aether-gateway/src",
|
||||
"apps/aether-tunnel/src",
|
||||
"crates/aether-admin/src",
|
||||
"crates/aether-billing/src",
|
||||
"crates/aether-model-fetch/src",
|
||||
"crates/aether-provider-pool/src",
|
||||
"crates/aether-runtime/src",
|
||||
"crates/aether-task-runtime/src",
|
||||
"crates/aether-usage-runtime/src",
|
||||
"crates/aether-provider-transport/src",
|
||||
"crates/aether-wallet/src",
|
||||
] {
|
||||
for path in collect_workspace_rust_files(root) {
|
||||
if path
|
||||
@@ -132,6 +139,33 @@ fn runtime_state_owns_redis_runtime_boundaries() {
|
||||
violations.join("\n")
|
||||
);
|
||||
|
||||
let mut dependency_violations = Vec::new();
|
||||
for manifest in [
|
||||
"apps/aether-gateway/Cargo.toml",
|
||||
"apps/aether-tunnel/Cargo.toml",
|
||||
"crates/aether-admin/Cargo.toml",
|
||||
"crates/aether-billing/Cargo.toml",
|
||||
"crates/aether-model-fetch/Cargo.toml",
|
||||
"crates/aether-provider-pool/Cargo.toml",
|
||||
"crates/aether-provider-transport/Cargo.toml",
|
||||
"crates/aether-runtime/Cargo.toml",
|
||||
"crates/aether-task-runtime/Cargo.toml",
|
||||
"crates/aether-usage-runtime/Cargo.toml",
|
||||
"crates/aether-wallet/Cargo.toml",
|
||||
] {
|
||||
let cargo = read_workspace_file(manifest);
|
||||
for forbidden in ["redis.workspace", "redis ="] {
|
||||
if cargo.contains(forbidden) {
|
||||
dependency_violations.push(format!("{manifest} -> {forbidden}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
dependency_violations.is_empty(),
|
||||
"business/runtime crates must not depend on redis directly:\n{}",
|
||||
dependency_violations.join("\n")
|
||||
);
|
||||
|
||||
let mut runtime_state_violations = Vec::new();
|
||||
for path in collect_workspace_rust_files("crates/aether-runtime-state/src") {
|
||||
if path
|
||||
@@ -160,6 +194,22 @@ fn runtime_state_owns_redis_runtime_boundaries() {
|
||||
"only crates/aether-runtime-state/src/redis may depend on the redis crate directly:\n{}",
|
||||
runtime_state_violations.join("\n")
|
||||
);
|
||||
|
||||
let mut runtime_connection_violations = Vec::new();
|
||||
for path in collect_workspace_rust_files("crates/aether-runtime-state/src") {
|
||||
if path.ends_with("crates/aether-runtime-state/src/redis/client.rs") {
|
||||
continue;
|
||||
}
|
||||
let source = production_workspace_source(&path);
|
||||
if source.contains("get_multiplexed_async_connection") {
|
||||
runtime_connection_violations.push(path.display().to_string());
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
runtime_connection_violations.is_empty(),
|
||||
"runtime Redis connections must be initialized only by redis/client.rs:\n{}",
|
||||
runtime_connection_violations.join("\n")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user