mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 16:37:46 +08:00
Merge branch 'fawney19:main' into main
This commit is contained in:
@@ -1879,6 +1879,8 @@ mod tests {
|
|||||||
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
|
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
|
||||||
use crate::AppState;
|
use crate::AppState;
|
||||||
|
|
||||||
|
const LOCAL_HTTP_SUCCESS_TIMEOUT_MS: u64 = 15_000;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn gateway_frontdoor_self_loop_guard_matches_loopback_public_ai_route() {
|
fn gateway_frontdoor_self_loop_guard_matches_loopback_public_ai_route() {
|
||||||
assert!(gateway_frontdoor_self_loop_guard_matches_with_port(
|
assert!(gateway_frontdoor_self_loop_guard_matches_with_port(
|
||||||
@@ -2244,7 +2246,7 @@ mod tests {
|
|||||||
transport_profile: None,
|
transport_profile: None,
|
||||||
timeouts: Some(ExecutionTimeouts {
|
timeouts: Some(ExecutionTimeouts {
|
||||||
connect_ms: Some(5_000),
|
connect_ms: Some(5_000),
|
||||||
total_ms: Some(5_000),
|
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
|
||||||
..ExecutionTimeouts::default()
|
..ExecutionTimeouts::default()
|
||||||
}),
|
}),
|
||||||
})
|
})
|
||||||
@@ -2398,7 +2400,7 @@ mod tests {
|
|||||||
format!("http://{addr}/chat"),
|
format!("http://{addr}/chat"),
|
||||||
true,
|
true,
|
||||||
ExecutionTimeouts {
|
ExecutionTimeouts {
|
||||||
first_byte_ms: Some(5_000),
|
first_byte_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
|
||||||
total_ms: Some(25),
|
total_ms: Some(25),
|
||||||
..ExecutionTimeouts::default()
|
..ExecutionTimeouts::default()
|
||||||
},
|
},
|
||||||
@@ -2674,7 +2676,7 @@ mod tests {
|
|||||||
})),
|
})),
|
||||||
}),
|
}),
|
||||||
timeouts: Some(ExecutionTimeouts {
|
timeouts: Some(ExecutionTimeouts {
|
||||||
total_ms: Some(5_000),
|
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
|
||||||
..ExecutionTimeouts::default()
|
..ExecutionTimeouts::default()
|
||||||
}),
|
}),
|
||||||
};
|
};
|
||||||
@@ -2797,7 +2799,7 @@ mod tests {
|
|||||||
transport_profile: None,
|
transport_profile: None,
|
||||||
timeouts: Some(ExecutionTimeouts {
|
timeouts: Some(ExecutionTimeouts {
|
||||||
connect_ms: Some(5_000),
|
connect_ms: Some(5_000),
|
||||||
total_ms: Some(5_000),
|
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
|
||||||
..ExecutionTimeouts::default()
|
..ExecutionTimeouts::default()
|
||||||
}),
|
}),
|
||||||
};
|
};
|
||||||
@@ -3156,7 +3158,7 @@ mod tests {
|
|||||||
transport_profile: None,
|
transport_profile: None,
|
||||||
timeouts: Some(ExecutionTimeouts {
|
timeouts: Some(ExecutionTimeouts {
|
||||||
connect_ms: Some(5_000),
|
connect_ms: Some(5_000),
|
||||||
total_ms: Some(5_000),
|
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
|
||||||
..ExecutionTimeouts::default()
|
..ExecutionTimeouts::default()
|
||||||
}),
|
}),
|
||||||
})
|
})
|
||||||
@@ -3209,7 +3211,7 @@ mod tests {
|
|||||||
transport_profile: None,
|
transport_profile: None,
|
||||||
timeouts: Some(ExecutionTimeouts {
|
timeouts: Some(ExecutionTimeouts {
|
||||||
connect_ms: Some(5_000),
|
connect_ms: Some(5_000),
|
||||||
total_ms: Some(5_000),
|
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
|
||||||
..ExecutionTimeouts::default()
|
..ExecutionTimeouts::default()
|
||||||
}),
|
}),
|
||||||
};
|
};
|
||||||
@@ -3358,7 +3360,7 @@ mod tests {
|
|||||||
transport_profile: None,
|
transport_profile: None,
|
||||||
timeouts: Some(ExecutionTimeouts {
|
timeouts: Some(ExecutionTimeouts {
|
||||||
connect_ms: Some(5_000),
|
connect_ms: Some(5_000),
|
||||||
total_ms: Some(5_000),
|
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
|
||||||
..ExecutionTimeouts::default()
|
..ExecutionTimeouts::default()
|
||||||
}),
|
}),
|
||||||
})
|
})
|
||||||
@@ -3437,7 +3439,7 @@ mod tests {
|
|||||||
transport_profile: None,
|
transport_profile: None,
|
||||||
timeouts: Some(ExecutionTimeouts {
|
timeouts: Some(ExecutionTimeouts {
|
||||||
connect_ms: Some(5_000),
|
connect_ms: Some(5_000),
|
||||||
total_ms: Some(5_000),
|
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
|
||||||
..ExecutionTimeouts::default()
|
..ExecutionTimeouts::default()
|
||||||
}),
|
}),
|
||||||
})
|
})
|
||||||
@@ -3519,7 +3521,7 @@ mod tests {
|
|||||||
}),
|
}),
|
||||||
timeouts: Some(ExecutionTimeouts {
|
timeouts: Some(ExecutionTimeouts {
|
||||||
connect_ms: Some(5_000),
|
connect_ms: Some(5_000),
|
||||||
total_ms: Some(5_000),
|
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
|
||||||
..ExecutionTimeouts::default()
|
..ExecutionTimeouts::default()
|
||||||
}),
|
}),
|
||||||
})
|
})
|
||||||
@@ -3579,7 +3581,7 @@ mod tests {
|
|||||||
transport_profile: None,
|
transport_profile: None,
|
||||||
timeouts: Some(ExecutionTimeouts {
|
timeouts: Some(ExecutionTimeouts {
|
||||||
connect_ms: Some(5_000),
|
connect_ms: Some(5_000),
|
||||||
total_ms: Some(5_000),
|
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
|
||||||
..ExecutionTimeouts::default()
|
..ExecutionTimeouts::default()
|
||||||
}),
|
}),
|
||||||
})
|
})
|
||||||
@@ -3717,7 +3719,7 @@ mod tests {
|
|||||||
transport_profile: None,
|
transport_profile: None,
|
||||||
timeouts: Some(ExecutionTimeouts {
|
timeouts: Some(ExecutionTimeouts {
|
||||||
connect_ms: Some(5_000),
|
connect_ms: Some(5_000),
|
||||||
total_ms: Some(5_000),
|
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
|
||||||
..ExecutionTimeouts::default()
|
..ExecutionTimeouts::default()
|
||||||
}),
|
}),
|
||||||
})
|
})
|
||||||
@@ -3778,7 +3780,7 @@ mod tests {
|
|||||||
transport_profile: None,
|
transport_profile: None,
|
||||||
timeouts: Some(ExecutionTimeouts {
|
timeouts: Some(ExecutionTimeouts {
|
||||||
connect_ms: Some(5_000),
|
connect_ms: Some(5_000),
|
||||||
total_ms: Some(5_000),
|
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
|
||||||
..ExecutionTimeouts::default()
|
..ExecutionTimeouts::default()
|
||||||
}),
|
}),
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -16,6 +16,6 @@ pub(crate) use self::reads::{
|
|||||||
};
|
};
|
||||||
pub(crate) use self::status::build_admin_provider_pool_status_payload;
|
pub(crate) use self::status::build_admin_provider_pool_status_payload;
|
||||||
pub(crate) use self::writes::{
|
pub(crate) use self::writes::{
|
||||||
admin_provider_pool_key_circuit_breaker_reason, record_admin_provider_pool_error,
|
admin_provider_pool_key_terminal_error_reason, record_admin_provider_pool_error,
|
||||||
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
|
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -222,7 +222,7 @@ fn first_error_text(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn admin_provider_pool_key_circuit_breaker_reason(
|
pub(crate) fn admin_provider_pool_key_terminal_error_reason(
|
||||||
status_code: u16,
|
status_code: u16,
|
||||||
error_body: Option<&str>,
|
error_body: Option<&str>,
|
||||||
) -> Option<String> {
|
) -> Option<String> {
|
||||||
@@ -475,7 +475,7 @@ pub(crate) async fn record_admin_provider_pool_error(
|
|||||||
|
|
||||||
if status_code == 400 {
|
if status_code == 400 {
|
||||||
// Bad Request is usually attributable to the caller payload, not key health.
|
// Bad Request is usually attributable to the caller payload, not key health.
|
||||||
// Account-level 400s are handled by the orchestration circuit-breaker path.
|
// Account-level 400s are handled by orchestration pool-score feedback.
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -588,7 +588,7 @@ pub(crate) async fn record_admin_provider_pool_stream_timeout(
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::{
|
use super::{
|
||||||
admin_provider_pool_key_circuit_breaker_reason, parse_google_quota_cooldown_seconds_at,
|
admin_provider_pool_key_terminal_error_reason, parse_google_quota_cooldown_seconds_at,
|
||||||
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
|
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
|
||||||
record_admin_provider_pool_success,
|
record_admin_provider_pool_success,
|
||||||
};
|
};
|
||||||
@@ -747,9 +747,9 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn circuit_reason_detects_workspace_deactivated_errors() {
|
fn terminal_error_reason_detects_workspace_deactivated_errors() {
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
admin_provider_pool_key_circuit_breaker_reason(
|
admin_provider_pool_key_terminal_error_reason(
|
||||||
402,
|
402,
|
||||||
Some(r#"{"error":{"message":"workspace has been deactivated"}}"#),
|
Some(r#"{"error":{"message":"workspace has been deactivated"}}"#),
|
||||||
)
|
)
|
||||||
@@ -757,7 +757,7 @@ mod tests {
|
|||||||
Some("workspace_deactivated_402:workspace has been deactivated")
|
Some("workspace_deactivated_402:workspace has been deactivated")
|
||||||
);
|
);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
admin_provider_pool_key_circuit_breaker_reason(
|
admin_provider_pool_key_terminal_error_reason(
|
||||||
400,
|
400,
|
||||||
Some(r#"{"error":{"message":"deactivated_workspace"}}"#),
|
Some(r#"{"error":{"message":"deactivated_workspace"}}"#),
|
||||||
)
|
)
|
||||||
@@ -767,9 +767,9 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn circuit_reason_detects_account_ban_errors() {
|
fn terminal_error_reason_detects_account_ban_errors() {
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
admin_provider_pool_key_circuit_breaker_reason(
|
admin_provider_pool_key_terminal_error_reason(
|
||||||
403,
|
403,
|
||||||
Some(r#"{"error":{"message":"AccountSuspendedException: account suspended"}}"#),
|
Some(r#"{"error":{"message":"AccountSuspendedException: account suspended"}}"#),
|
||||||
)
|
)
|
||||||
@@ -777,7 +777,7 @@ mod tests {
|
|||||||
Some("forbidden_403")
|
Some("forbidden_403")
|
||||||
);
|
);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
admin_provider_pool_key_circuit_breaker_reason(
|
admin_provider_pool_key_terminal_error_reason(
|
||||||
423,
|
423,
|
||||||
Some(r#"{"error":{"message":"account access denied"}}"#),
|
Some(r#"{"error":{"message":"account access denied"}}"#),
|
||||||
)
|
)
|
||||||
@@ -1099,7 +1099,7 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn severe_account_errors_use_circuit_breaker_instead_of_pool_cooldown() {
|
async fn severe_account_errors_do_not_use_pool_cooldown() {
|
||||||
let Some(redis) = start_managed_redis_or_skip().await else {
|
let Some(redis) = start_managed_redis_or_skip().await else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
@@ -1109,7 +1109,7 @@ mod tests {
|
|||||||
let key_ids = vec!["key-account-disabled".to_string()];
|
let key_ids = vec!["key-account-disabled".to_string()];
|
||||||
|
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
admin_provider_pool_key_circuit_breaker_reason(
|
admin_provider_pool_key_terminal_error_reason(
|
||||||
401,
|
401,
|
||||||
Some(r#"{"error":{"message":"account has been deactivated"}}"#),
|
Some(r#"{"error":{"message":"account has been deactivated"}}"#),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -13,7 +13,6 @@ use aether_data_contracts::repository::provider_catalog::{
|
|||||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||||
};
|
};
|
||||||
use aether_data_contracts::repository::usage::StoredProviderApiKeyWindowUsageSummary;
|
use aether_data_contracts::repository::usage::StoredProviderApiKeyWindowUsageSummary;
|
||||||
use aether_scheduler_core::provider_key_circuit_payload_is_active_open_at;
|
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
use std::collections::BTreeMap;
|
use std::collections::BTreeMap;
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
@@ -934,18 +933,6 @@ fn admin_pool_health_score(key: &StoredProviderCatalogKey) -> f64 {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn admin_pool_circuit_breaker_open(key: &StoredProviderCatalogKey, now_unix_secs: u64) -> bool {
|
|
||||||
key.circuit_breaker_by_format
|
|
||||||
.as_ref()
|
|
||||||
.and_then(serde_json::Value::as_object)
|
|
||||||
.map(|formats| {
|
|
||||||
formats
|
|
||||||
.values()
|
|
||||||
.any(|item| provider_key_circuit_payload_is_active_open_at(item, now_unix_secs))
|
|
||||||
})
|
|
||||||
.unwrap_or(false)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn admin_pool_scheduling_payload(
|
fn admin_pool_scheduling_payload(
|
||||||
key: &StoredProviderCatalogKey,
|
key: &StoredProviderCatalogKey,
|
||||||
cooldown_reason: Option<&str>,
|
cooldown_reason: Option<&str>,
|
||||||
@@ -1040,7 +1027,7 @@ pub(super) fn build_admin_pool_key_payload(
|
|||||||
.as_ref()
|
.as_ref()
|
||||||
.and_then(|_| runtime.cooldown_ttl_by_key.get(&key.id).copied());
|
.and_then(|_| runtime.cooldown_ttl_by_key.get(&key.id).copied());
|
||||||
let health_score = admin_pool_health_score(key);
|
let health_score = admin_pool_health_score(key);
|
||||||
let circuit_breaker_open = admin_pool_circuit_breaker_open(key, now_unix_secs);
|
let circuit_breaker_open = false;
|
||||||
let auth_semantics = provider_key_auth_semantics(key, provider_type);
|
let auth_semantics = provider_key_auth_semantics(key, provider_type);
|
||||||
let account_quota_exhausted = pool_config
|
let account_quota_exhausted = pool_config
|
||||||
.as_ref()
|
.as_ref()
|
||||||
|
|||||||
@@ -810,11 +810,11 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
|||||||
node_id: node_id.clone(),
|
node_id: node_id.clone(),
|
||||||
heartbeat_interval: payload.heartbeat_interval,
|
heartbeat_interval: payload.heartbeat_interval,
|
||||||
active_connections: payload.active_connections,
|
active_connections: payload.active_connections,
|
||||||
total_requests_delta: payload.total_requests,
|
total_requests_delta: payload.window_total_requests.or(payload.total_requests),
|
||||||
avg_latency_ms: payload.avg_latency_ms,
|
avg_latency_ms: payload.avg_latency_ms,
|
||||||
failed_requests_delta: payload.failed_requests,
|
failed_requests_delta: payload.window_failed_requests.or(payload.failed_requests),
|
||||||
dns_failures_delta: payload.dns_failures,
|
dns_failures_delta: payload.window_dns_failures.or(payload.dns_failures),
|
||||||
stream_errors_delta: payload.stream_errors,
|
stream_errors_delta: payload.window_stream_errors.or(payload.stream_errors),
|
||||||
proxy_metadata: payload.proxy_metadata,
|
proxy_metadata: payload.proxy_metadata,
|
||||||
proxy_version: payload.proxy_version,
|
proxy_version: payload.proxy_version,
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -402,10 +402,16 @@ pub(crate) fn parse_internal_tunnel_heartbeat_request(
|
|||||||
.is_some_and(|value| !(5..=600).contains(&value))
|
.is_some_and(|value| !(5..=600).contains(&value))
|
||||||
|| payload.active_connections.is_some_and(|value| value < 0)
|
|| payload.active_connections.is_some_and(|value| value < 0)
|
||||||
|| payload.total_requests.is_some_and(|value| value < 0)
|
|| payload.total_requests.is_some_and(|value| value < 0)
|
||||||
|
|| payload.window_total_requests.is_some_and(|value| value < 0)
|
||||||
|| payload.avg_latency_ms.is_some_and(|value| value < 0.0)
|
|| payload.avg_latency_ms.is_some_and(|value| value < 0.0)
|
||||||
|| payload.failed_requests.is_some_and(|value| value < 0)
|
|| payload.failed_requests.is_some_and(|value| value < 0)
|
||||||
|
|| payload
|
||||||
|
.window_failed_requests
|
||||||
|
.is_some_and(|value| value < 0)
|
||||||
|| payload.dns_failures.is_some_and(|value| value < 0)
|
|| payload.dns_failures.is_some_and(|value| value < 0)
|
||||||
|
|| payload.window_dns_failures.is_some_and(|value| value < 0)
|
||||||
|| payload.stream_errors.is_some_and(|value| value < 0)
|
|| payload.stream_errors.is_some_and(|value| value < 0)
|
||||||
|
|| payload.window_stream_errors.is_some_and(|value| value < 0)
|
||||||
|| payload
|
|| payload
|
||||||
.proxy_version
|
.proxy_version
|
||||||
.as_deref()
|
.as_deref()
|
||||||
|
|||||||
@@ -12,14 +12,22 @@ pub(crate) struct InternalTunnelHeartbeatRequest {
|
|||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub(crate) total_requests: Option<i64>,
|
pub(crate) total_requests: Option<i64>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
|
pub(crate) window_total_requests: Option<i64>,
|
||||||
|
#[serde(default)]
|
||||||
pub(crate) avg_latency_ms: Option<f64>,
|
pub(crate) avg_latency_ms: Option<f64>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub(crate) failed_requests: Option<i64>,
|
pub(crate) failed_requests: Option<i64>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
|
pub(crate) window_failed_requests: Option<i64>,
|
||||||
|
#[serde(default)]
|
||||||
pub(crate) dns_failures: Option<i64>,
|
pub(crate) dns_failures: Option<i64>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
|
pub(crate) window_dns_failures: Option<i64>,
|
||||||
|
#[serde(default)]
|
||||||
pub(crate) stream_errors: Option<i64>,
|
pub(crate) stream_errors: Option<i64>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
|
pub(crate) window_stream_errors: Option<i64>,
|
||||||
|
#[serde(default)]
|
||||||
pub(crate) proxy_metadata: Option<serde_json::Value>,
|
pub(crate) proxy_metadata: Option<serde_json::Value>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub(crate) proxy_version: Option<String>,
|
pub(crate) proxy_version: Option<String>,
|
||||||
|
|||||||
@@ -2,7 +2,7 @@ pub(crate) use super::super::admin::provider::pool::config::{
|
|||||||
admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value,
|
admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value,
|
||||||
};
|
};
|
||||||
pub(crate) use super::super::admin::provider::pool::runtime::{
|
pub(crate) use super::super::admin::provider::pool::runtime::{
|
||||||
admin_provider_pool_key_circuit_breaker_reason, read_admin_provider_pool_key_cooldown_reason,
|
admin_provider_pool_key_terminal_error_reason, read_admin_provider_pool_key_cooldown_reason,
|
||||||
read_admin_provider_pool_runtime_state, record_admin_provider_pool_error,
|
read_admin_provider_pool_runtime_state, record_admin_provider_pool_error,
|
||||||
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
|
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
|
||||||
release_admin_provider_pool_key_lease,
|
release_admin_provider_pool_key_lease,
|
||||||
|
|||||||
@@ -19,8 +19,7 @@ use tracing::warn;
|
|||||||
use super::{
|
use super::{
|
||||||
local_failover_error_message, project_local_adaptive_rate_limit,
|
local_failover_error_message, project_local_adaptive_rate_limit,
|
||||||
project_local_adaptive_success, project_local_failure_health, project_local_key_circuit_closed,
|
project_local_adaptive_success, project_local_failure_health, project_local_key_circuit_closed,
|
||||||
project_local_key_circuit_failure, project_local_key_circuit_open,
|
project_local_key_circuit_failure, project_local_success_health, LocalFailoverClassification,
|
||||||
project_local_success_health, LocalFailoverClassification,
|
|
||||||
};
|
};
|
||||||
use crate::ai_serving::extract_pool_sticky_session_token;
|
use crate::ai_serving::extract_pool_sticky_session_token;
|
||||||
use crate::client_session_affinity::{
|
use crate::client_session_affinity::{
|
||||||
@@ -29,7 +28,7 @@ use crate::client_session_affinity::{
|
|||||||
use crate::clock::current_unix_secs;
|
use crate::clock::current_unix_secs;
|
||||||
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
|
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
|
||||||
use crate::handlers::shared::provider_pool::{
|
use crate::handlers::shared::provider_pool::{
|
||||||
admin_provider_pool_key_circuit_breaker_reason, record_admin_provider_pool_error,
|
admin_provider_pool_key_terminal_error_reason, record_admin_provider_pool_error,
|
||||||
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
|
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
|
||||||
release_admin_provider_pool_key_lease, AdminProviderPoolConfig,
|
release_admin_provider_pool_key_lease, AdminProviderPoolConfig,
|
||||||
};
|
};
|
||||||
@@ -574,6 +573,7 @@ async fn record_health_failure_effect(
|
|||||||
else {
|
else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
|
let is_pool_provider = local_execution_plan_uses_pool(state, context.plan).await;
|
||||||
let observed_at_unix_secs = current_unix_secs();
|
let observed_at_unix_secs = current_unix_secs();
|
||||||
let Some(health_by_format) = project_local_failure_health(
|
let Some(health_by_format) = project_local_failure_health(
|
||||||
current_key.health_by_format.as_ref(),
|
current_key.health_by_format.as_ref(),
|
||||||
@@ -589,16 +589,24 @@ async fn record_health_failure_effect(
|
|||||||
.and_then(|value| value.get("consecutive_failures"))
|
.and_then(|value| value.get("consecutive_failures"))
|
||||||
.and_then(Value::as_u64)
|
.and_then(Value::as_u64)
|
||||||
.unwrap_or(0);
|
.unwrap_or(0);
|
||||||
let circuit_breaker_by_format = project_local_key_circuit_failure(
|
let circuit_breaker_update_owned = if is_pool_provider {
|
||||||
current_key.circuit_breaker_by_format.as_ref(),
|
None
|
||||||
api_format,
|
} else {
|
||||||
observed_at_unix_secs,
|
project_local_key_circuit_failure(
|
||||||
consecutive_failures,
|
current_key.circuit_breaker_by_format.as_ref(),
|
||||||
current_key.max_probe_interval_minutes,
|
api_format,
|
||||||
);
|
observed_at_unix_secs,
|
||||||
let circuit_breaker_update = circuit_breaker_by_format
|
consecutive_failures,
|
||||||
.as_ref()
|
current_key.max_probe_interval_minutes,
|
||||||
.or(current_key.circuit_breaker_by_format.as_ref());
|
)
|
||||||
|
};
|
||||||
|
let circuit_breaker_update = if is_pool_provider {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
circuit_breaker_update_owned
|
||||||
|
.as_ref()
|
||||||
|
.or(current_key.circuit_breaker_by_format.as_ref())
|
||||||
|
};
|
||||||
|
|
||||||
if let Err(err) = state
|
if let Err(err) = state
|
||||||
.update_provider_catalog_key_health_state(
|
.update_provider_catalog_key_health_state(
|
||||||
@@ -636,23 +644,35 @@ async fn record_health_success_effect(
|
|||||||
else {
|
else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
|
let is_pool_provider = local_execution_plan_uses_pool(state, context.plan).await;
|
||||||
let Some(health_by_format) =
|
let Some(health_by_format) =
|
||||||
project_local_success_health(current_key.health_by_format.as_ref(), api_format)
|
project_local_success_health(current_key.health_by_format.as_ref(), api_format)
|
||||||
else {
|
else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
let circuit_breaker_by_format = current_key
|
let circuit_breaker_update_owned = if is_pool_provider {
|
||||||
.circuit_breaker_by_format
|
None
|
||||||
.as_ref()
|
} else {
|
||||||
.and_then(|current| project_local_key_circuit_closed(Some(current), api_format));
|
current_key
|
||||||
|
.circuit_breaker_by_format
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|current| project_local_key_circuit_closed(Some(current), api_format))
|
||||||
|
};
|
||||||
if current_key.health_by_format.as_ref() == Some(&health_by_format)
|
if current_key.health_by_format.as_ref() == Some(&health_by_format)
|
||||||
&& circuit_breaker_by_format.as_ref() == current_key.circuit_breaker_by_format.as_ref()
|
&& ((is_pool_provider && current_key.circuit_breaker_by_format.is_none())
|
||||||
|
|| (!is_pool_provider
|
||||||
|
&& circuit_breaker_update_owned.as_ref()
|
||||||
|
== current_key.circuit_breaker_by_format.as_ref()))
|
||||||
{
|
{
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
let circuit_breaker_update = circuit_breaker_by_format
|
let circuit_breaker_update = if is_pool_provider {
|
||||||
.as_ref()
|
None
|
||||||
.or(current_key.circuit_breaker_by_format.as_ref());
|
} else {
|
||||||
|
circuit_breaker_update_owned
|
||||||
|
.as_ref()
|
||||||
|
.or(current_key.circuit_breaker_by_format.as_ref())
|
||||||
|
};
|
||||||
|
|
||||||
if let Err(err) = state
|
if let Err(err) = state
|
||||||
.update_provider_catalog_key_health_state(
|
.update_provider_catalog_key_health_state(
|
||||||
@@ -711,9 +731,9 @@ async fn record_pool_error_effect(
|
|||||||
context: LocalExecutionEffectContext<'_>,
|
context: LocalExecutionEffectContext<'_>,
|
||||||
effect: LocalPoolErrorEffect<'_>,
|
effect: LocalPoolErrorEffect<'_>,
|
||||||
) {
|
) {
|
||||||
let circuit_reason =
|
let terminal_error_reason =
|
||||||
admin_provider_pool_key_circuit_breaker_reason(effect.status_code, effect.error_body);
|
admin_provider_pool_key_terminal_error_reason(effect.status_code, effect.error_body);
|
||||||
if circuit_reason.is_none()
|
if terminal_error_reason.is_none()
|
||||||
&& !local_candidate_failure_should_record_pool_error(
|
&& !local_candidate_failure_should_record_pool_error(
|
||||||
effect.classification,
|
effect.classification,
|
||||||
effect.status_code,
|
effect.status_code,
|
||||||
@@ -726,10 +746,7 @@ async fn record_pool_error_effect(
|
|||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
|
|
||||||
if let Some(reason) = circuit_reason {
|
clear_pool_key_circuit_breaker(state, context).await;
|
||||||
open_pool_key_circuit_breaker(state, context, &reason).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
record_admin_provider_pool_error(
|
record_admin_provider_pool_error(
|
||||||
state.runtime_state.as_ref(),
|
state.runtime_state.as_ref(),
|
||||||
&context.plan.provider_id,
|
&context.plan.provider_id,
|
||||||
@@ -757,16 +774,10 @@ async fn record_pool_error_effect(
|
|||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn open_pool_key_circuit_breaker(
|
async fn clear_pool_key_circuit_breaker(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
context: LocalExecutionEffectContext<'_>,
|
context: LocalExecutionEffectContext<'_>,
|
||||||
reason: &str,
|
|
||||||
) {
|
) {
|
||||||
let api_format = context.plan.provider_api_format.trim();
|
|
||||||
if api_format.is_empty() {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
let Some(current_key) = state
|
let Some(current_key) = state
|
||||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&context.plan.key_id))
|
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&context.plan.key_id))
|
||||||
.await
|
.await
|
||||||
@@ -775,27 +786,21 @@ async fn open_pool_key_circuit_breaker(
|
|||||||
else {
|
else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
let Some(circuit_breaker_by_format) = project_local_key_circuit_open(
|
if current_key.circuit_breaker_by_format.is_none() {
|
||||||
current_key.circuit_breaker_by_format.as_ref(),
|
|
||||||
api_format,
|
|
||||||
reason,
|
|
||||||
current_unix_secs(),
|
|
||||||
current_key.max_probe_interval_minutes,
|
|
||||||
) else {
|
|
||||||
return;
|
return;
|
||||||
};
|
}
|
||||||
|
|
||||||
if let Err(err) = state
|
if let Err(err) = state
|
||||||
.update_provider_catalog_key_health_state(
|
.update_provider_catalog_key_health_state(
|
||||||
&context.plan.key_id,
|
&context.plan.key_id,
|
||||||
current_key.is_active,
|
current_key.is_active,
|
||||||
current_key.health_by_format.as_ref(),
|
current_key.health_by_format.as_ref(),
|
||||||
Some(&circuit_breaker_by_format),
|
None,
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
warn!(
|
warn!(
|
||||||
"gateway orchestration effects: failed to open pool key circuit for provider {} endpoint {} key {}: {:?}",
|
"gateway orchestration effects: failed to clear pool key circuit for provider {} endpoint {} key {}: {:?}",
|
||||||
context.plan.provider_id, context.plan.endpoint_id, context.plan.key_id, err
|
context.plan.provider_id, context.plan.endpoint_id, context.plan.key_id, err
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -984,6 +989,10 @@ fn pool_score_hard_state_for_status(
|
|||||||
status_code: u16,
|
status_code: u16,
|
||||||
error_body: Option<&str>,
|
error_body: Option<&str>,
|
||||||
) -> Option<PoolMemberHardState> {
|
) -> Option<PoolMemberHardState> {
|
||||||
|
if let Some(reason) = admin_provider_pool_key_terminal_error_reason(status_code, error_body) {
|
||||||
|
return Some(pool_score_hard_state_for_terminal_error_reason(&reason));
|
||||||
|
}
|
||||||
|
|
||||||
match status_code {
|
match status_code {
|
||||||
401 | 403 => Some(PoolMemberHardState::AuthInvalid),
|
401 | 403 => Some(PoolMemberHardState::AuthInvalid),
|
||||||
402 => Some(PoolMemberHardState::QuotaExhausted),
|
402 => Some(PoolMemberHardState::QuotaExhausted),
|
||||||
@@ -1006,6 +1015,16 @@ fn pool_score_hard_state_for_status(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn pool_score_hard_state_for_terminal_error_reason(reason: &str) -> PoolMemberHardState {
|
||||||
|
if reason.starts_with("payment_required_") {
|
||||||
|
PoolMemberHardState::QuotaExhausted
|
||||||
|
} else if reason.starts_with("forbidden_") {
|
||||||
|
PoolMemberHardState::AuthInvalid
|
||||||
|
} else {
|
||||||
|
PoolMemberHardState::Banned
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn pool_score_delta_for_status(status_code: u16) -> i32 {
|
fn pool_score_delta_for_status(status_code: u16) -> i32 {
|
||||||
match status_code {
|
match status_code {
|
||||||
401 | 403 => -2_000,
|
401 | 403 => -2_000,
|
||||||
@@ -1028,6 +1047,7 @@ mod tests {
|
|||||||
use aether_data_contracts::repository::candidates::{
|
use aether_data_contracts::repository::candidates::{
|
||||||
RequestCandidateStatus, StoredRequestCandidate,
|
RequestCandidateStatus, StoredRequestCandidate,
|
||||||
};
|
};
|
||||||
|
use aether_data_contracts::repository::pool_scores::PoolMemberHardState;
|
||||||
use aether_data_contracts::repository::provider_catalog::{
|
use aether_data_contracts::repository::provider_catalog::{
|
||||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||||
};
|
};
|
||||||
@@ -1036,9 +1056,10 @@ mod tests {
|
|||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
apply_local_execution_effect, local_candidate_failure_should_record_pool_error,
|
apply_local_execution_effect, local_candidate_failure_should_record_pool_error,
|
||||||
LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect,
|
pool_score_hard_state_for_status, LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect,
|
||||||
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
|
LocalAttemptFailureEffect, LocalExecutionEffect, LocalExecutionEffectContext,
|
||||||
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
|
LocalHealthFailureEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect,
|
||||||
|
LocalPoolErrorEffect,
|
||||||
};
|
};
|
||||||
use crate::data::{GatewayDataConfig, GatewayDataState};
|
use crate::data::{GatewayDataConfig, GatewayDataState};
|
||||||
use crate::orchestration::LocalFailoverClassification;
|
use crate::orchestration::LocalFailoverClassification;
|
||||||
@@ -1943,13 +1964,46 @@ mod tests {
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn terminal_pool_account_errors_project_pool_hard_state() {
|
||||||
|
assert_eq!(
|
||||||
|
pool_score_hard_state_for_status(
|
||||||
|
400,
|
||||||
|
Some(r#"{"error":{"message":"deactivated_workspace"}}"#),
|
||||||
|
),
|
||||||
|
Some(PoolMemberHardState::Banned)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
pool_score_hard_state_for_status(
|
||||||
|
402,
|
||||||
|
Some(r#"{"error":{"message":"payment required"}}"#),
|
||||||
|
),
|
||||||
|
Some(PoolMemberHardState::QuotaExhausted)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn pool_account_error_opens_key_circuit() {
|
async fn pool_account_error_does_not_open_key_circuit() {
|
||||||
let Some(redis) = start_managed_redis_or_skip().await else {
|
let Some(redis) = start_managed_redis_or_skip().await else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
let state = codex_state_with_redis(redis.redis_url(), "orchestration_pool_circuit");
|
let state = codex_state_with_redis(redis.redis_url(), "orchestration_pool_circuit");
|
||||||
let plan = sample_codex_plan();
|
let plan = sample_codex_plan();
|
||||||
|
let legacy_circuit = json!({
|
||||||
|
"openai:responses": {
|
||||||
|
"open": true,
|
||||||
|
"reason": "legacy"
|
||||||
|
}
|
||||||
|
});
|
||||||
|
state
|
||||||
|
.update_provider_catalog_key_health_state(
|
||||||
|
&plan.key_id,
|
||||||
|
true,
|
||||||
|
None,
|
||||||
|
Some(&legacy_circuit),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("legacy circuit should seed");
|
||||||
|
|
||||||
apply_local_execution_effect(
|
apply_local_execution_effect(
|
||||||
&state,
|
&state,
|
||||||
@@ -1973,15 +2027,7 @@ mod tests {
|
|||||||
.into_iter()
|
.into_iter()
|
||||||
.next()
|
.next()
|
||||||
.expect("stored key should exist");
|
.expect("stored key should exist");
|
||||||
let circuit = stored_key
|
assert_eq!(stored_key.circuit_breaker_by_format, None);
|
||||||
.circuit_breaker_by_format
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|value| value.get("openai:responses"))
|
|
||||||
.expect("format circuit should be stored");
|
|
||||||
assert_eq!(circuit["open"], json!(true));
|
|
||||||
assert_eq!(circuit["reason"], json!("account_deactivated_401"));
|
|
||||||
assert!(circuit["next_probe_at"].is_string());
|
|
||||||
assert!(circuit["next_probe_at_unix_secs"].as_u64().is_some());
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
@@ -2144,6 +2190,45 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn pool_health_failure_does_not_open_key_circuit_after_eight_consecutive_failures() {
|
||||||
|
let state = pool_health_state();
|
||||||
|
let plan = sample_plan();
|
||||||
|
|
||||||
|
for _ in 0..8 {
|
||||||
|
apply_local_execution_effect(
|
||||||
|
&state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan: &plan,
|
||||||
|
report_context: None,
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
|
||||||
|
status_code: 503,
|
||||||
|
classification: LocalFailoverClassification::RetryUpstreamFailure,
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
let stored_key = state
|
||||||
|
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
|
||||||
|
.await
|
||||||
|
.expect("provider catalog keys should load")
|
||||||
|
.into_iter()
|
||||||
|
.next()
|
||||||
|
.expect("stored key should exist");
|
||||||
|
assert_eq!(stored_key.circuit_breaker_by_format, None);
|
||||||
|
assert_eq!(
|
||||||
|
stored_key
|
||||||
|
.health_by_format
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|value| value.get("openai:chat"))
|
||||||
|
.and_then(|value| value.get("consecutive_failures"))
|
||||||
|
.and_then(Value::as_u64),
|
||||||
|
Some(8)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn health_success_projection_resets_key_health_for_format() {
|
async fn health_success_projection_resets_key_health_for_format() {
|
||||||
let state = health_state();
|
let state = health_state();
|
||||||
|
|||||||
@@ -35,8 +35,7 @@ pub(crate) use self::effects::{
|
|||||||
};
|
};
|
||||||
pub(crate) use self::health::{
|
pub(crate) use self::health::{
|
||||||
project_local_failure_health, project_local_key_circuit_closed,
|
project_local_failure_health, project_local_key_circuit_closed,
|
||||||
project_local_key_circuit_failure, project_local_key_circuit_open,
|
project_local_key_circuit_failure, project_local_success_health,
|
||||||
project_local_success_health,
|
|
||||||
};
|
};
|
||||||
pub(crate) use self::policy::{
|
pub(crate) use self::policy::{
|
||||||
append_local_failover_policy_to_value, local_failover_policy_from_report_context,
|
append_local_failover_policy_to_value, local_failover_policy_from_report_context,
|
||||||
|
|||||||
@@ -128,6 +128,7 @@ pub(super) fn is_candidate_selectable(
|
|||||||
.get(candidate.key_id.as_str())
|
.get(candidate.key_id.as_str())
|
||||||
.copied()
|
.copied()
|
||||||
.unwrap_or(false),
|
.unwrap_or(false),
|
||||||
|
enforce_key_circuit_breaker: !pool_group,
|
||||||
rpm_reset_at: (!pool_group)
|
rpm_reset_at: (!pool_group)
|
||||||
.then(|| {
|
.then(|| {
|
||||||
snapshot
|
snapshot
|
||||||
@@ -182,6 +183,7 @@ pub(super) fn current_candidate_runtime_skip_reason(
|
|||||||
.get(candidate.key_id.as_str())
|
.get(candidate.key_id.as_str())
|
||||||
.copied()
|
.copied()
|
||||||
.unwrap_or(false),
|
.unwrap_or(false),
|
||||||
|
enforce_key_circuit_breaker: !pool_group,
|
||||||
rpm_reset_at,
|
rpm_reset_at,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1553,6 +1553,71 @@ async fn selects_next_candidate_when_first_provider_key_circuit_is_open() {
|
|||||||
assert_eq!(selected.key_id, "key-b");
|
assert_eq!(selected.key_id, "key-b");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn pool_provider_ignores_key_circuit_open_runtime_skip() {
|
||||||
|
let mut first = sample_row();
|
||||||
|
first.provider_id = "provider-a".to_string();
|
||||||
|
first.provider_name = "openai-a".to_string();
|
||||||
|
first.endpoint_id = "endpoint-a".to_string();
|
||||||
|
first.key_id = "key-a".to_string();
|
||||||
|
first.key_name = "alpha".to_string();
|
||||||
|
first.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 1}));
|
||||||
|
|
||||||
|
let mut second = sample_row();
|
||||||
|
second.provider_id = "provider-b".to_string();
|
||||||
|
second.provider_name = "openai-b".to_string();
|
||||||
|
second.endpoint_id = "endpoint-b".to_string();
|
||||||
|
second.key_id = "key-b".to_string();
|
||||||
|
second.key_name = "beta".to_string();
|
||||||
|
second.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 2}));
|
||||||
|
|
||||||
|
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||||
|
first, second,
|
||||||
|
]));
|
||||||
|
let mut pool_provider = sample_provider("provider-a", None);
|
||||||
|
pool_provider.config = Some(serde_json::json!({"pool_advanced": {"enabled": true}}));
|
||||||
|
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||||
|
vec![pool_provider, sample_provider("provider-b", None)],
|
||||||
|
Vec::new(),
|
||||||
|
vec![
|
||||||
|
sample_key("key-a", "provider-a", Some(10)).with_health_fields(
|
||||||
|
Some(serde_json::json!({"openai:chat": {"health_score": 0.2}})),
|
||||||
|
Some(serde_json::json!({"openai:chat": {"open": true}})),
|
||||||
|
),
|
||||||
|
sample_key("key-b", "provider-b", Some(10)),
|
||||||
|
],
|
||||||
|
));
|
||||||
|
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||||
|
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![]));
|
||||||
|
let state = AppState::new()
|
||||||
|
.expect("state should build")
|
||||||
|
.with_data_state_for_tests(
|
||||||
|
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
|
||||||
|
candidates,
|
||||||
|
provider_catalog,
|
||||||
|
quotas,
|
||||||
|
request_candidates,
|
||||||
|
),
|
||||||
|
);
|
||||||
|
|
||||||
|
let (selected, skipped) = collect_selectable_candidates_with_skip_reasons(
|
||||||
|
state.data.as_ref(),
|
||||||
|
&state,
|
||||||
|
"openai:chat",
|
||||||
|
"gpt-4.1",
|
||||||
|
false,
|
||||||
|
None,
|
||||||
|
100,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("selection should succeed");
|
||||||
|
|
||||||
|
assert!(selected
|
||||||
|
.iter()
|
||||||
|
.any(|candidate| candidate.provider_id == "provider-a"));
|
||||||
|
assert!(skipped.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn exposes_runtime_skipped_candidates_with_skip_reasons() {
|
async fn exposes_runtime_skipped_candidates_with_skip_reasons() {
|
||||||
let mut first = sample_row();
|
let mut first = sample_row();
|
||||||
|
|||||||
@@ -1595,7 +1595,7 @@ async fn gateway_ignores_health_signals_in_pool_scheduling_status() {
|
|||||||
|
|
||||||
assert_eq!(keys.len(), 2);
|
assert_eq!(keys.len(), 2);
|
||||||
assert_eq!(keys[0]["key_name"], json!("circuit-open"));
|
assert_eq!(keys[0]["key_name"], json!("circuit-open"));
|
||||||
assert_eq!(keys[0]["circuit_breaker_open"], json!(true));
|
assert_eq!(keys[0]["circuit_breaker_open"], json!(false));
|
||||||
assert_eq!(keys[0]["scheduling_status"], json!("available"));
|
assert_eq!(keys[0]["scheduling_status"], json!("available"));
|
||||||
assert_eq!(keys[0]["scheduling_reason"], json!("available"));
|
assert_eq!(keys[0]["scheduling_reason"], json!("available"));
|
||||||
assert_eq!(keys[0]["scheduling_label"], json!("可用"));
|
assert_eq!(keys[0]["scheduling_label"], json!("可用"));
|
||||||
|
|||||||
@@ -1514,6 +1514,32 @@ async fn gateway_reports_proxy_node_metrics_and_filters_events_locally() {
|
|||||||
.expect("system time should be after epoch")
|
.expect("system time should be after epoch")
|
||||||
.as_secs();
|
.as_secs();
|
||||||
|
|
||||||
|
let baseline_heartbeat_response = client
|
||||||
|
.post(format!("{gateway_url}/api/internal/tunnel/heartbeat"))
|
||||||
|
.json(&json!({
|
||||||
|
"node_id": "node-1",
|
||||||
|
"heartbeat_id": 90,
|
||||||
|
"heartbeat_interval": 30,
|
||||||
|
"active_connections": 0,
|
||||||
|
"proxy_metadata": {
|
||||||
|
"tunnel_metrics": {
|
||||||
|
"connect_errors": 0,
|
||||||
|
"disconnects": 0,
|
||||||
|
"error_events_total": 0,
|
||||||
|
"ws_in_bytes": 0,
|
||||||
|
"ws_out_bytes": 0,
|
||||||
|
"ws_in_frames": 0,
|
||||||
|
"ws_out_frames": 0,
|
||||||
|
"heartbeat_rtt_last_ms": 0
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"proxy_version": "2.0.0"
|
||||||
|
}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.expect("baseline heartbeat request should succeed");
|
||||||
|
assert_eq!(baseline_heartbeat_response.status(), StatusCode::OK);
|
||||||
|
|
||||||
let heartbeat_response = client
|
let heartbeat_response = client
|
||||||
.post(format!("{gateway_url}/api/internal/tunnel/heartbeat"))
|
.post(format!("{gateway_url}/api/internal/tunnel/heartbeat"))
|
||||||
.json(&json!({
|
.json(&json!({
|
||||||
@@ -1564,18 +1590,21 @@ async fn gateway_reports_proxy_node_metrics_and_filters_events_locally() {
|
|||||||
.await
|
.await
|
||||||
.expect("metrics json should parse");
|
.expect("metrics json should parse");
|
||||||
assert_eq!(metrics_payload["step"], "1m");
|
assert_eq!(metrics_payload["step"], "1m");
|
||||||
assert_eq!(metrics_payload["summary"]["samples"], 1);
|
assert_eq!(metrics_payload["summary"]["samples"], 2);
|
||||||
assert_eq!(metrics_payload["summary"]["uptime_samples"], 1);
|
assert_eq!(metrics_payload["summary"]["uptime_samples"], 2);
|
||||||
assert_eq!(metrics_payload["summary"]["active_connections_max"], 7);
|
assert_eq!(metrics_payload["summary"]["active_connections_max"], 7);
|
||||||
assert_eq!(metrics_payload["summary"]["heartbeat_rtt_ms_avg"], 42.0);
|
assert_eq!(metrics_payload["summary"]["heartbeat_rtt_ms_sum"], 42);
|
||||||
|
assert_eq!(metrics_payload["summary"]["heartbeat_rtt_ms_avg"], 21.0);
|
||||||
assert_eq!(metrics_payload["summary"]["connect_errors_delta"], 3);
|
assert_eq!(metrics_payload["summary"]["connect_errors_delta"], 3);
|
||||||
assert_eq!(metrics_payload["summary"]["ws_out_frames_delta"], 20);
|
assert_eq!(metrics_payload["summary"]["ws_out_frames_delta"], 20);
|
||||||
let metric_items = metrics_payload["items"]
|
let metric_items = metrics_payload["items"]
|
||||||
.as_array()
|
.as_array()
|
||||||
.expect("metrics items should be array");
|
.expect("metrics items should be array");
|
||||||
assert_eq!(metric_items.len(), 1);
|
assert!(!metric_items.is_empty());
|
||||||
assert_eq!(metric_items[0]["node_id"], "node-1");
|
assert!(metric_items.iter().any(|item| item["node_id"] == "node-1"));
|
||||||
assert!(metric_items[0]["bucket_start"].is_string());
|
assert!(metric_items
|
||||||
|
.iter()
|
||||||
|
.all(|item| item["bucket_start"].is_string()));
|
||||||
|
|
||||||
let fleet_response = client
|
let fleet_response = client
|
||||||
.get(format!(
|
.get(format!(
|
||||||
@@ -1593,7 +1622,7 @@ async fn gateway_reports_proxy_node_metrics_and_filters_events_locally() {
|
|||||||
.json()
|
.json()
|
||||||
.await
|
.await
|
||||||
.expect("fleet json should parse");
|
.expect("fleet json should parse");
|
||||||
assert_eq!(fleet_payload["summary"]["samples"], 1);
|
assert_eq!(fleet_payload["summary"]["samples"], 2);
|
||||||
assert_eq!(fleet_payload["summary"]["error_events_delta"], 1);
|
assert_eq!(fleet_payload["summary"]["error_events_delta"], 1);
|
||||||
|
|
||||||
let events_response = client
|
let events_response = client
|
||||||
|
|||||||
@@ -52,11 +52,15 @@ async fn gateway_handles_internal_tunnel_heartbeat_locally_with_loopback() {
|
|||||||
"heartbeat_id": 77,
|
"heartbeat_id": 77,
|
||||||
"heartbeat_interval": 45,
|
"heartbeat_interval": 45,
|
||||||
"active_connections": 5,
|
"active_connections": 5,
|
||||||
"total_requests": 9,
|
"total_requests": 100,
|
||||||
"avg_latency_ms": 12.5,
|
"avg_latency_ms": 12.5,
|
||||||
"failed_requests": 1,
|
"failed_requests": 20,
|
||||||
"dns_failures": 2,
|
"dns_failures": 30,
|
||||||
"stream_errors": 3,
|
"stream_errors": 40,
|
||||||
|
"window_total_requests": 9,
|
||||||
|
"window_failed_requests": 1,
|
||||||
|
"window_dns_failures": 2,
|
||||||
|
"window_stream_errors": 3,
|
||||||
"proxy_metadata": {"arch": "arm64"},
|
"proxy_metadata": {"arch": "arm64"},
|
||||||
"proxy_version": "2.0.0",
|
"proxy_version": "2.0.0",
|
||||||
}))
|
}))
|
||||||
@@ -71,6 +75,15 @@ async fn gateway_handles_internal_tunnel_heartbeat_locally_with_loopback() {
|
|||||||
assert_eq!(payload["upgrade_to"], "1.2.3");
|
assert_eq!(payload["upgrade_to"], "1.2.3");
|
||||||
assert_eq!(payload["remote_config"]["allowed_ports"][0], 443);
|
assert_eq!(payload["remote_config"]["allowed_ports"][0], 443);
|
||||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||||
|
let node = repository
|
||||||
|
.find_proxy_node("node-123")
|
||||||
|
.await
|
||||||
|
.expect("node lookup should succeed")
|
||||||
|
.expect("node should exist");
|
||||||
|
assert_eq!(node.total_requests, 9);
|
||||||
|
assert_eq!(node.failed_requests, 1);
|
||||||
|
assert_eq!(node.dns_failures, 2);
|
||||||
|
assert_eq!(node.stream_errors, 3);
|
||||||
|
|
||||||
gateway_handle.abort();
|
gateway_handle.abort();
|
||||||
upstream_handle.abort();
|
upstream_handle.abort();
|
||||||
|
|||||||
@@ -72,14 +72,22 @@ struct InternalTunnelHeartbeatRequest {
|
|||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
total_requests: Option<i64>,
|
total_requests: Option<i64>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
|
window_total_requests: Option<i64>,
|
||||||
|
#[serde(default)]
|
||||||
avg_latency_ms: Option<f64>,
|
avg_latency_ms: Option<f64>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
failed_requests: Option<i64>,
|
failed_requests: Option<i64>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
|
window_failed_requests: Option<i64>,
|
||||||
|
#[serde(default)]
|
||||||
dns_failures: Option<i64>,
|
dns_failures: Option<i64>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
|
window_dns_failures: Option<i64>,
|
||||||
|
#[serde(default)]
|
||||||
stream_errors: Option<i64>,
|
stream_errors: Option<i64>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
|
window_stream_errors: Option<i64>,
|
||||||
|
#[serde(default)]
|
||||||
proxy_metadata: Option<serde_json::Value>,
|
proxy_metadata: Option<serde_json::Value>,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
proxy_version: Option<String>,
|
proxy_version: Option<String>,
|
||||||
@@ -932,11 +940,11 @@ async fn apply_embedded_tunnel_heartbeat(
|
|||||||
node_id: node_id.clone(),
|
node_id: node_id.clone(),
|
||||||
heartbeat_interval: payload.heartbeat_interval,
|
heartbeat_interval: payload.heartbeat_interval,
|
||||||
active_connections: payload.active_connections,
|
active_connections: payload.active_connections,
|
||||||
total_requests_delta: payload.total_requests,
|
total_requests_delta: payload.window_total_requests.or(payload.total_requests),
|
||||||
avg_latency_ms: payload.avg_latency_ms,
|
avg_latency_ms: payload.avg_latency_ms,
|
||||||
failed_requests_delta: payload.failed_requests,
|
failed_requests_delta: payload.window_failed_requests.or(payload.failed_requests),
|
||||||
dns_failures_delta: payload.dns_failures,
|
dns_failures_delta: payload.window_dns_failures.or(payload.dns_failures),
|
||||||
stream_errors_delta: payload.stream_errors,
|
stream_errors_delta: payload.window_stream_errors.or(payload.stream_errors),
|
||||||
proxy_metadata: payload.proxy_metadata,
|
proxy_metadata: payload.proxy_metadata,
|
||||||
proxy_version: payload.proxy_version,
|
proxy_version: payload.proxy_version,
|
||||||
};
|
};
|
||||||
@@ -1009,10 +1017,16 @@ fn parse_embedded_tunnel_heartbeat_request(
|
|||||||
.is_some_and(|value| !(5..=600).contains(&value))
|
.is_some_and(|value| !(5..=600).contains(&value))
|
||||||
|| payload.active_connections.is_some_and(|value| value < 0)
|
|| payload.active_connections.is_some_and(|value| value < 0)
|
||||||
|| payload.total_requests.is_some_and(|value| value < 0)
|
|| payload.total_requests.is_some_and(|value| value < 0)
|
||||||
|
|| payload.window_total_requests.is_some_and(|value| value < 0)
|
||||||
|| payload.avg_latency_ms.is_some_and(|value| value < 0.0)
|
|| payload.avg_latency_ms.is_some_and(|value| value < 0.0)
|
||||||
|| payload.failed_requests.is_some_and(|value| value < 0)
|
|| payload.failed_requests.is_some_and(|value| value < 0)
|
||||||
|
|| payload
|
||||||
|
.window_failed_requests
|
||||||
|
.is_some_and(|value| value < 0)
|
||||||
|| payload.dns_failures.is_some_and(|value| value < 0)
|
|| payload.dns_failures.is_some_and(|value| value < 0)
|
||||||
|
|| payload.window_dns_failures.is_some_and(|value| value < 0)
|
||||||
|| payload.stream_errors.is_some_and(|value| value < 0)
|
|| payload.stream_errors.is_some_and(|value| value < 0)
|
||||||
|
|| payload.window_stream_errors.is_some_and(|value| value < 0)
|
||||||
|| payload
|
|| payload
|
||||||
.proxy_version
|
.proxy_version
|
||||||
.as_deref()
|
.as_deref()
|
||||||
|
|||||||
@@ -1384,6 +1384,7 @@ mod tests {
|
|||||||
upstream_connect_timeout_secs: 30,
|
upstream_connect_timeout_secs: 30,
|
||||||
upstream_pool_max_idle_per_host: 4,
|
upstream_pool_max_idle_per_host: 4,
|
||||||
upstream_pool_idle_timeout_secs: 60,
|
upstream_pool_idle_timeout_secs: 60,
|
||||||
|
upstream_client_pool_capacity: crate::config::DEFAULT_UPSTREAM_CLIENT_POOL_CAPACITY,
|
||||||
upstream_tcp_keepalive_secs: 60,
|
upstream_tcp_keepalive_secs: 60,
|
||||||
upstream_tcp_nodelay: true,
|
upstream_tcp_nodelay: true,
|
||||||
upstream_proxy_url: None,
|
upstream_proxy_url: None,
|
||||||
|
|||||||
@@ -65,6 +65,7 @@ pub const DEFAULT_TUNNEL_SCALE_CHECK_INTERVAL_MS: u64 = 1_000;
|
|||||||
pub const DEFAULT_TUNNEL_SCALE_UP_THRESHOLD_PERCENT: u32 = 50;
|
pub const DEFAULT_TUNNEL_SCALE_UP_THRESHOLD_PERCENT: u32 = 50;
|
||||||
pub const DEFAULT_TUNNEL_SCALE_DOWN_THRESHOLD_PERCENT: u32 = 35;
|
pub const DEFAULT_TUNNEL_SCALE_DOWN_THRESHOLD_PERCENT: u32 = 35;
|
||||||
pub const DEFAULT_TUNNEL_SCALE_DOWN_GRACE_SECS: u64 = 15;
|
pub const DEFAULT_TUNNEL_SCALE_DOWN_GRACE_SECS: u64 = 15;
|
||||||
|
pub const DEFAULT_UPSTREAM_CLIENT_POOL_CAPACITY: usize = 256;
|
||||||
const AUTO_TUNNEL_CONNECTIONS_REDUNDANT_FLOOR: u64 = 2;
|
const AUTO_TUNNEL_CONNECTIONS_REDUNDANT_FLOOR: u64 = 2;
|
||||||
const AUTO_TUNNEL_CONNECTIONS_BASE_CAP: u64 = 4;
|
const AUTO_TUNNEL_CONNECTIONS_BASE_CAP: u64 = 4;
|
||||||
// Bias the automatic pool toward a per-device upper band without letting
|
// Bias the automatic pool toward a per-device upper band without letting
|
||||||
@@ -524,6 +525,14 @@ pub struct Config {
|
|||||||
)]
|
)]
|
||||||
pub upstream_pool_idle_timeout_secs: u64,
|
pub upstream_pool_idle_timeout_secs: u64,
|
||||||
|
|
||||||
|
/// Maximum number of keyed upstream HTTP clients retained by the tunnel.
|
||||||
|
#[arg(
|
||||||
|
long,
|
||||||
|
env = "AETHER_TUNNEL_UPSTREAM_CLIENT_POOL_CAPACITY",
|
||||||
|
default_value_t = DEFAULT_UPSTREAM_CLIENT_POOL_CAPACITY
|
||||||
|
)]
|
||||||
|
pub upstream_client_pool_capacity: usize,
|
||||||
|
|
||||||
/// Upstream TCP keepalive in seconds (0 disables)
|
/// Upstream TCP keepalive in seconds (0 disables)
|
||||||
#[arg(
|
#[arg(
|
||||||
long,
|
long,
|
||||||
@@ -813,6 +822,9 @@ impl Config {
|
|||||||
if self.upstream_connect_timeout_secs == 0 {
|
if self.upstream_connect_timeout_secs == 0 {
|
||||||
anyhow::bail!("upstream_connect_timeout_secs must be > 0");
|
anyhow::bail!("upstream_connect_timeout_secs must be > 0");
|
||||||
}
|
}
|
||||||
|
if self.upstream_client_pool_capacity == 0 {
|
||||||
|
anyhow::bail!("upstream_client_pool_capacity must be > 0");
|
||||||
|
}
|
||||||
if let Some(proxy_url) = normalized_proxy_url(&self.aether_outbound_proxy_url) {
|
if let Some(proxy_url) = normalized_proxy_url(&self.aether_outbound_proxy_url) {
|
||||||
crate::egress_proxy::UpstreamProxyConfig::parse(proxy_url)
|
crate::egress_proxy::UpstreamProxyConfig::parse(proxy_url)
|
||||||
.map_err(|err| anyhow::anyhow!("aether_outbound_proxy_url invalid: {err}"))?;
|
.map_err(|err| anyhow::anyhow!("aether_outbound_proxy_url invalid: {err}"))?;
|
||||||
@@ -1044,6 +1056,8 @@ pub struct ConfigFile {
|
|||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub upstream_pool_idle_timeout_secs: Option<u64>,
|
pub upstream_pool_idle_timeout_secs: Option<u64>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub upstream_client_pool_capacity: Option<usize>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub upstream_tcp_keepalive_secs: Option<u64>,
|
pub upstream_tcp_keepalive_secs: Option<u64>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub upstream_tcp_nodelay: Option<bool>,
|
pub upstream_tcp_nodelay: Option<bool>,
|
||||||
@@ -1227,6 +1241,10 @@ impl ConfigFile {
|
|||||||
"AETHER_TUNNEL_UPSTREAM_POOL_IDLE_TIMEOUT",
|
"AETHER_TUNNEL_UPSTREAM_POOL_IDLE_TIMEOUT",
|
||||||
self.upstream_pool_idle_timeout_secs
|
self.upstream_pool_idle_timeout_secs
|
||||||
);
|
);
|
||||||
|
set!(
|
||||||
|
"AETHER_TUNNEL_UPSTREAM_CLIENT_POOL_CAPACITY",
|
||||||
|
self.upstream_client_pool_capacity
|
||||||
|
);
|
||||||
set!(
|
set!(
|
||||||
"AETHER_TUNNEL_UPSTREAM_TCP_KEEPALIVE",
|
"AETHER_TUNNEL_UPSTREAM_TCP_KEEPALIVE",
|
||||||
self.upstream_tcp_keepalive_secs
|
self.upstream_tcp_keepalive_secs
|
||||||
|
|||||||
@@ -239,7 +239,10 @@ where
|
|||||||
|
|
||||||
// Create body channel and spawn handler
|
// Create body channel and spawn handler
|
||||||
let (body_tx, body_rx) = mpsc::channel::<Frame>(64);
|
let (body_tx, body_rx) = mpsc::channel::<Frame>(64);
|
||||||
streams.insert(frame.stream_id, body_tx);
|
let request_headers_end_stream = frame.is_end_stream();
|
||||||
|
if !request_headers_end_stream {
|
||||||
|
streams.insert(frame.stream_id, body_tx);
|
||||||
|
}
|
||||||
|
|
||||||
let state_clone = Arc::clone(&state);
|
let state_clone = Arc::clone(&state);
|
||||||
let server_clone = Arc::clone(&server);
|
let server_clone = Arc::clone(&server);
|
||||||
@@ -326,8 +329,16 @@ where
|
|||||||
// Trigger every 64 frames OR when the count exceeds max_streams.
|
// Trigger every 64 frames OR when the count exceeds max_streams.
|
||||||
frames_since_cleanup += 1;
|
frames_since_cleanup += 1;
|
||||||
if frames_since_cleanup >= 64 || handler_handles.len() > max_streams {
|
if frames_since_cleanup >= 64 || handler_handles.len() > max_streams {
|
||||||
|
let closed_streams = prune_closed_stream_senders(&mut streams);
|
||||||
|
if closed_streams > 0 {
|
||||||
|
debug!(closed_streams, "removed closed request body stream senders");
|
||||||
|
}
|
||||||
handler_handles.retain(|h| !h.is_finished());
|
handler_handles.retain(|h| !h.is_finished());
|
||||||
frames_since_cleanup = 0;
|
frames_since_cleanup = 0;
|
||||||
|
if draining && streams.is_empty() {
|
||||||
|
info!("tunnel drained after cleanup");
|
||||||
|
break None;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -397,6 +408,12 @@ fn try_send_stream_error(frame_tx: &FrameSender, stream_id: u32, message: &'stat
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn prune_closed_stream_senders(streams: &mut HashMap<u32, mpsc::Sender<Frame>>) -> usize {
|
||||||
|
let before = streams.len();
|
||||||
|
streams.retain(|_, tx| !tx.is_closed());
|
||||||
|
before.saturating_sub(streams.len())
|
||||||
|
}
|
||||||
|
|
||||||
/// Wait for all active stream handlers to finish (with a timeout).
|
/// Wait for all active stream handlers to finish (with a timeout).
|
||||||
async fn drain_handlers(handles: Vec<JoinHandle<()>>) {
|
async fn drain_handlers(handles: Vec<JoinHandle<()>>) {
|
||||||
if handles.is_empty() {
|
if handles.is_empty() {
|
||||||
@@ -470,4 +487,18 @@ mod tests {
|
|||||||
Bytes::from_static(b"tunnel request body dispatch stalled")
|
Bytes::from_static(b"tunnel request body dispatch stalled")
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn prune_closed_stream_senders_drops_streams_with_closed_receivers() {
|
||||||
|
let (closed_tx, closed_rx) = mpsc::channel::<Frame>(1);
|
||||||
|
let (open_tx, _open_rx) = mpsc::channel::<Frame>(1);
|
||||||
|
drop(closed_rx);
|
||||||
|
let mut streams = HashMap::from([(7, closed_tx), (9, open_tx)]);
|
||||||
|
|
||||||
|
let removed = prune_closed_stream_senders(&mut streams);
|
||||||
|
|
||||||
|
assert_eq!(removed, 1);
|
||||||
|
assert!(!streams.contains_key(&7));
|
||||||
|
assert!(streams.contains_key(&9));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -532,6 +532,7 @@ mod tests {
|
|||||||
upstream_connect_timeout_secs: 30,
|
upstream_connect_timeout_secs: 30,
|
||||||
upstream_pool_max_idle_per_host: 4,
|
upstream_pool_max_idle_per_host: 4,
|
||||||
upstream_pool_idle_timeout_secs: 60,
|
upstream_pool_idle_timeout_secs: 60,
|
||||||
|
upstream_client_pool_capacity: crate::config::DEFAULT_UPSTREAM_CLIENT_POOL_CAPACITY,
|
||||||
upstream_tcp_keepalive_secs: 60,
|
upstream_tcp_keepalive_secs: 60,
|
||||||
upstream_tcp_nodelay: true,
|
upstream_tcp_nodelay: true,
|
||||||
upstream_proxy_url: None,
|
upstream_proxy_url: None,
|
||||||
|
|||||||
@@ -521,6 +521,21 @@ fn prepare_request_body(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn prepare_bodyless_request_body(
|
||||||
|
body_rx: mpsc::Receiver<TunnelFrame>,
|
||||||
|
follow_redirects: bool,
|
||||||
|
) -> PreparedRequestBody {
|
||||||
|
drop(body_rx);
|
||||||
|
PreparedRequestBody {
|
||||||
|
first_request_body: Some(empty_request_body()),
|
||||||
|
replay_body: if follow_redirects {
|
||||||
|
ReplayableRequestBody::None
|
||||||
|
} else {
|
||||||
|
ReplayableRequestBody::NonReplayable
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
async fn collect_request_body_for_replay(
|
async fn collect_request_body_for_replay(
|
||||||
mut body_rx: mpsc::Receiver<TunnelFrame>,
|
mut body_rx: mpsc::Receiver<TunnelFrame>,
|
||||||
body_size: Arc<AtomicUsize>,
|
body_size: Arc<AtomicUsize>,
|
||||||
@@ -1379,17 +1394,7 @@ async fn handle_stream_inner(
|
|||||||
0,
|
0,
|
||||||
)
|
)
|
||||||
} else {
|
} else {
|
||||||
PreparedRequestBody {
|
prepare_bodyless_request_body(body_rx, follow_redirects)
|
||||||
first_request_body: Some(build_streaming_request_body(
|
|
||||||
body_rx,
|
|
||||||
Arc::clone(&request_body_size),
|
|
||||||
)),
|
|
||||||
replay_body: if follow_redirects {
|
|
||||||
ReplayableRequestBody::None
|
|
||||||
} else {
|
|
||||||
ReplayableRequestBody::NonReplayable
|
|
||||||
},
|
|
||||||
}
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut total_dns_ms = 0u64;
|
let mut total_dns_ms = 0u64;
|
||||||
@@ -1591,6 +1596,7 @@ async fn send_error(tx: &FrameSender, stream_id: u32, msg: &str) {
|
|||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
fn build_streaming_request_body(
|
fn build_streaming_request_body(
|
||||||
body_rx: mpsc::Receiver<TunnelFrame>,
|
body_rx: mpsc::Receiver<TunnelFrame>,
|
||||||
body_size: Arc<AtomicUsize>,
|
body_size: Arc<AtomicUsize>,
|
||||||
@@ -1620,6 +1626,7 @@ fn build_spooled_request_body(
|
|||||||
upstream_client::stream_request_body(body_stream)
|
upstream_client::stream_request_body(body_stream)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
fn build_prefixed_request_body(
|
fn build_prefixed_request_body(
|
||||||
prefix_chunks: Vec<Bytes>,
|
prefix_chunks: Vec<Bytes>,
|
||||||
body_rx: mpsc::Receiver<TunnelFrame>,
|
body_rx: mpsc::Receiver<TunnelFrame>,
|
||||||
@@ -1794,6 +1801,22 @@ mod tests {
|
|||||||
assert_eq!(body_size.load(Ordering::Relaxed), 0);
|
assert_eq!(body_size.load(Ordering::Relaxed), 0);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn bodyless_request_body_completes_without_waiting_for_tunnel_sender() {
|
||||||
|
let (_tx, rx) = mpsc::channel(4);
|
||||||
|
let mut prepared = prepare_bodyless_request_body(rx, true);
|
||||||
|
let mut body = prepared
|
||||||
|
.first_request_body
|
||||||
|
.take()
|
||||||
|
.expect("bodyless request should have an initial body");
|
||||||
|
|
||||||
|
let frame = tokio::time::timeout(Duration::from_millis(25), body.frame())
|
||||||
|
.await
|
||||||
|
.expect("bodyless request body should not wait for tunnel body frames");
|
||||||
|
assert!(frame.is_none());
|
||||||
|
assert!(matches!(prepared.replay_body, ReplayableRequestBody::None));
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn prepare_request_body_streams_immediately_and_replays_after_completion() {
|
async fn prepare_request_body_streams_immediately_and_replays_after_completion() {
|
||||||
let (tx, rx) = mpsc::channel(4);
|
let (tx, rx) = mpsc::channel(4);
|
||||||
@@ -2690,6 +2713,7 @@ mod tests {
|
|||||||
upstream_connect_timeout_secs: 30,
|
upstream_connect_timeout_secs: 30,
|
||||||
upstream_pool_max_idle_per_host: 4,
|
upstream_pool_max_idle_per_host: 4,
|
||||||
upstream_pool_idle_timeout_secs: 60,
|
upstream_pool_idle_timeout_secs: 60,
|
||||||
|
upstream_client_pool_capacity: crate::config::DEFAULT_UPSTREAM_CLIENT_POOL_CAPACITY,
|
||||||
upstream_tcp_keepalive_secs: 60,
|
upstream_tcp_keepalive_secs: 60,
|
||||||
upstream_tcp_nodelay: true,
|
upstream_tcp_nodelay: true,
|
||||||
upstream_proxy_url: None,
|
upstream_proxy_url: None,
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ use std::future::Future;
|
|||||||
use std::io;
|
use std::io;
|
||||||
use std::net::IpAddr;
|
use std::net::IpAddr;
|
||||||
use std::pin::Pin;
|
use std::pin::Pin;
|
||||||
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::sync::Mutex;
|
use std::sync::Mutex;
|
||||||
use std::task::{Context, Poll};
|
use std::task::{Context, Poll};
|
||||||
@@ -65,7 +66,14 @@ pub struct UpstreamClientPoolKey {
|
|||||||
pub struct UpstreamClientPool {
|
pub struct UpstreamClientPool {
|
||||||
config: Arc<Config>,
|
config: Arc<Config>,
|
||||||
dns_cache: Arc<DnsCache>,
|
dns_cache: Arc<DnsCache>,
|
||||||
clients: Arc<Mutex<HashMap<UpstreamClientPoolKey, UpstreamClient>>>,
|
clients: Arc<Mutex<HashMap<UpstreamClientPoolKey, UpstreamClientPoolEntry>>>,
|
||||||
|
access_counter: Arc<AtomicU64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct UpstreamClientPoolEntry {
|
||||||
|
client: UpstreamClient,
|
||||||
|
last_used: u64,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl UpstreamClientPool {
|
impl UpstreamClientPool {
|
||||||
@@ -74,18 +82,17 @@ impl UpstreamClientPool {
|
|||||||
config,
|
config,
|
||||||
dns_cache,
|
dns_cache,
|
||||||
clients: Arc::new(Mutex::new(HashMap::new())),
|
clients: Arc::new(Mutex::new(HashMap::new())),
|
||||||
|
access_counter: Arc::new(AtomicU64::new(0)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn get_or_build(&self, key: UpstreamClientPoolKey) -> Result<UpstreamClient, String> {
|
pub fn get_or_build(&self, key: UpstreamClientPoolKey) -> Result<UpstreamClient, String> {
|
||||||
if let Some(client) = self
|
|
||||||
.clients
|
|
||||||
.lock()
|
|
||||||
.expect("client pool lock")
|
|
||||||
.get(&key)
|
|
||||||
.cloned()
|
|
||||||
{
|
{
|
||||||
return Ok(client);
|
let mut clients = self.clients.lock().expect("client pool lock");
|
||||||
|
if let Some(entry) = clients.get_mut(&key) {
|
||||||
|
entry.last_used = self.next_access_id();
|
||||||
|
return Ok(entry.client.clone());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
validate_proxy_transport_backend(&key.backend)?;
|
validate_proxy_transport_backend(&key.backend)?;
|
||||||
@@ -97,12 +104,45 @@ impl UpstreamClientPool {
|
|||||||
Arc::clone(&self.dns_cache),
|
Arc::clone(&self.dns_cache),
|
||||||
http1_only,
|
http1_only,
|
||||||
)?;
|
)?;
|
||||||
self.clients
|
let mut clients = self.clients.lock().expect("client pool lock");
|
||||||
.lock()
|
if let Some(entry) = clients.get_mut(&key) {
|
||||||
.expect("client pool lock")
|
entry.last_used = self.next_access_id();
|
||||||
.insert(key, client.clone());
|
return Ok(entry.client.clone());
|
||||||
|
}
|
||||||
|
evict_lru_client_if_needed(
|
||||||
|
&mut clients,
|
||||||
|
self.config.upstream_client_pool_capacity.max(1),
|
||||||
|
);
|
||||||
|
clients.insert(
|
||||||
|
key,
|
||||||
|
UpstreamClientPoolEntry {
|
||||||
|
client: client.clone(),
|
||||||
|
last_used: self.next_access_id(),
|
||||||
|
},
|
||||||
|
);
|
||||||
Ok(client)
|
Ok(client)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn next_access_id(&self) -> u64 {
|
||||||
|
self.access_counter.fetch_add(1, Ordering::Relaxed)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn evict_lru_client_if_needed(
|
||||||
|
clients: &mut HashMap<UpstreamClientPoolKey, UpstreamClientPoolEntry>,
|
||||||
|
capacity: usize,
|
||||||
|
) {
|
||||||
|
if clients.len() < capacity {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let Some(oldest_key) = clients
|
||||||
|
.iter()
|
||||||
|
.min_by_key(|(_, entry)| entry.last_used)
|
||||||
|
.map(|(key, _)| key.clone())
|
||||||
|
else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
clients.remove(&oldest_key);
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn upstream_client_pool_key(
|
pub fn upstream_client_pool_key(
|
||||||
@@ -765,6 +805,44 @@ mod tests {
|
|||||||
assert!(error.contains("unsupported transport profile backend"));
|
assert!(error.contains("unsupported transport profile backend"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn upstream_client_pool_evicts_lru_clients_above_capacity() {
|
||||||
|
let config = Arc::new(
|
||||||
|
Config::try_parse_from([
|
||||||
|
"aether-tunnel",
|
||||||
|
"--aether-url",
|
||||||
|
"https://aether.example.com",
|
||||||
|
"--management-token",
|
||||||
|
"ae_test",
|
||||||
|
"--node-name",
|
||||||
|
"tunnel-test",
|
||||||
|
"--upstream-client-pool-capacity",
|
||||||
|
"2",
|
||||||
|
])
|
||||||
|
.expect("config should parse"),
|
||||||
|
);
|
||||||
|
let pool =
|
||||||
|
UpstreamClientPool::new(config, Arc::new(DnsCache::new(Duration::from_secs(60), 16)));
|
||||||
|
let key_a = test_pool_key("key-a");
|
||||||
|
let key_b = test_pool_key("key-b");
|
||||||
|
let key_c = test_pool_key("key-c");
|
||||||
|
|
||||||
|
pool.get_or_build(key_a.clone())
|
||||||
|
.expect("client A should build");
|
||||||
|
pool.get_or_build(key_b.clone())
|
||||||
|
.expect("client B should build");
|
||||||
|
pool.get_or_build(key_a.clone())
|
||||||
|
.expect("client A should be reused and become most recent");
|
||||||
|
pool.get_or_build(key_c.clone())
|
||||||
|
.expect("client C should build");
|
||||||
|
|
||||||
|
let clients = pool.clients.lock().expect("client pool lock");
|
||||||
|
assert_eq!(clients.len(), 2);
|
||||||
|
assert!(clients.contains_key(&key_a));
|
||||||
|
assert!(clients.contains_key(&key_c));
|
||||||
|
assert!(!clients.contains_key(&key_b));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn http_proxy_authorization_header_uses_basic_auth_for_http_proxy() {
|
fn http_proxy_authorization_header_uses_basic_auth_for_http_proxy() {
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
@@ -777,6 +855,16 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn test_pool_key(key_id: &str) -> UpstreamClientPoolKey {
|
||||||
|
upstream_client_pool_key(
|
||||||
|
Some("provider-1"),
|
||||||
|
Some("endpoint-1"),
|
||||||
|
Some(key_id),
|
||||||
|
None,
|
||||||
|
false,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn socks5h_target_address_uses_domain_name() {
|
async fn socks5h_target_address_uses_domain_name() {
|
||||||
let request = socks5_target_address("example.com", 443, true)
|
let request = socks5_target_address("example.com", 443, true)
|
||||||
|
|||||||
@@ -156,40 +156,6 @@ fn admin_pool_health_score(key: &StoredProviderCatalogKey) -> f64 {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn admin_pool_circuit_breaker_open(key: &StoredProviderCatalogKey) -> bool {
|
|
||||||
let now_unix_secs = Utc::now().timestamp().max(0) as u64;
|
|
||||||
key.circuit_breaker_by_format
|
|
||||||
.as_ref()
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.map(|formats| {
|
|
||||||
formats
|
|
||||||
.values()
|
|
||||||
.any(|item| admin_pool_circuit_payload_active_open_at(item, now_unix_secs))
|
|
||||||
})
|
|
||||||
.unwrap_or(false)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn admin_pool_circuit_payload_active_open_at(value: &Value, now_unix_secs: u64) -> bool {
|
|
||||||
let Some(item) = value.as_object() else {
|
|
||||||
return false;
|
|
||||||
};
|
|
||||||
if !item.get("open").and_then(Value::as_bool).unwrap_or(false) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
if let Some(next_probe_at) = item.get("next_probe_at_unix_secs").and_then(Value::as_u64) {
|
|
||||||
return now_unix_secs < next_probe_at;
|
|
||||||
}
|
|
||||||
if let Some(next_probe_at) = item
|
|
||||||
.get("next_probe_at")
|
|
||||||
.and_then(Value::as_str)
|
|
||||||
.and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok())
|
|
||||||
.and_then(|value| u64::try_from(value.timestamp()).ok())
|
|
||||||
{
|
|
||||||
return now_unix_secs < next_probe_at;
|
|
||||||
}
|
|
||||||
true
|
|
||||||
}
|
|
||||||
|
|
||||||
fn unix_secs_to_rfc3339(unix_secs: u64) -> Option<String> {
|
fn unix_secs_to_rfc3339(unix_secs: u64) -> Option<String> {
|
||||||
Utc.timestamp_opt(unix_secs as i64, 0)
|
Utc.timestamp_opt(unix_secs as i64, 0)
|
||||||
.single()
|
.single()
|
||||||
@@ -727,7 +693,7 @@ mod tests {
|
|||||||
let payload = build_admin_pool_key_payload(&key, &AdminPoolKeyPayloadContext::default());
|
let payload = build_admin_pool_key_payload(&key, &AdminPoolKeyPayloadContext::default());
|
||||||
|
|
||||||
assert_eq!(payload["health_score"], json!(0.2));
|
assert_eq!(payload["health_score"], json!(0.2));
|
||||||
assert_eq!(payload["circuit_breaker_open"], json!(true));
|
assert_eq!(payload["circuit_breaker_open"], json!(false));
|
||||||
assert_eq!(payload["scheduling_status"], json!("available"));
|
assert_eq!(payload["scheduling_status"], json!("available"));
|
||||||
assert_eq!(payload["scheduling_reason"], json!("available"));
|
assert_eq!(payload["scheduling_reason"], json!("available"));
|
||||||
assert_eq!(payload["scheduling_label"], json!("可用"));
|
assert_eq!(payload["scheduling_label"], json!("可用"));
|
||||||
@@ -739,7 +705,7 @@ pub fn build_admin_pool_key_payload(
|
|||||||
context: &AdminPoolKeyPayloadContext,
|
context: &AdminPoolKeyPayloadContext,
|
||||||
) -> Value {
|
) -> Value {
|
||||||
let health_score = admin_pool_health_score(key);
|
let health_score = admin_pool_health_score(key);
|
||||||
let circuit_breaker_open = admin_pool_circuit_breaker_open(key);
|
let circuit_breaker_open = false;
|
||||||
let (scheduling_status, scheduling_reason, scheduling_label, scheduling_reasons) =
|
let (scheduling_status, scheduling_reason, scheduling_label, scheduling_reasons) =
|
||||||
admin_pool_scheduling_payload(
|
admin_pool_scheduling_payload(
|
||||||
key,
|
key,
|
||||||
|
|||||||
@@ -1566,6 +1566,33 @@ VALUES ('node-1', 'registered', 'ok', 3)
|
|||||||
.await
|
.await
|
||||||
.expect("node should register");
|
.expect("node should register");
|
||||||
let now = super::current_unix_secs();
|
let now = super::current_unix_secs();
|
||||||
|
repository
|
||||||
|
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
||||||
|
node_id: registered.id.clone(),
|
||||||
|
heartbeat_interval: Some(30),
|
||||||
|
active_connections: Some(0),
|
||||||
|
total_requests_delta: None,
|
||||||
|
avg_latency_ms: None,
|
||||||
|
failed_requests_delta: None,
|
||||||
|
dns_failures_delta: None,
|
||||||
|
stream_errors_delta: None,
|
||||||
|
proxy_metadata: Some(json!({
|
||||||
|
"tunnel_metrics": {
|
||||||
|
"connect_errors": 0,
|
||||||
|
"disconnects": 0,
|
||||||
|
"error_events_total": 0,
|
||||||
|
"ws_in_bytes": 0,
|
||||||
|
"ws_out_bytes": 0,
|
||||||
|
"ws_in_frames": 0,
|
||||||
|
"ws_out_frames": 0,
|
||||||
|
"heartbeat_rtt_last_ms": 0
|
||||||
|
}
|
||||||
|
})),
|
||||||
|
proxy_version: Some("1.0.0".to_string()),
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.expect("baseline heartbeat should apply")
|
||||||
|
.expect("node should exist");
|
||||||
repository
|
repository
|
||||||
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
||||||
node_id: registered.id.clone(),
|
node_id: registered.id.clone(),
|
||||||
@@ -1609,13 +1636,44 @@ VALUES ('node-1', 'registered', 'ok', 3)
|
|||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
.expect("metrics should list");
|
.expect("metrics should list");
|
||||||
assert_eq!(metrics.len(), 1);
|
assert!(!metrics.is_empty());
|
||||||
assert_eq!(metrics[0].samples, 1);
|
assert_eq!(metrics.iter().map(|bucket| bucket.samples).sum::<i64>(), 2);
|
||||||
assert_eq!(metrics[0].uptime_samples, 1);
|
assert_eq!(
|
||||||
assert_eq!(metrics[0].active_connections_max, 5);
|
metrics
|
||||||
assert_eq!(metrics[0].heartbeat_rtt_ms_sum, 33);
|
.iter()
|
||||||
assert_eq!(metrics[0].connect_errors_delta, 4);
|
.map(|bucket| bucket.uptime_samples)
|
||||||
assert_eq!(metrics[0].ws_out_frames_delta, 6);
|
.sum::<i64>(),
|
||||||
|
2
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
metrics
|
||||||
|
.iter()
|
||||||
|
.map(|bucket| bucket.active_connections_max)
|
||||||
|
.max()
|
||||||
|
.unwrap_or_default(),
|
||||||
|
5
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
metrics
|
||||||
|
.iter()
|
||||||
|
.map(|bucket| bucket.heartbeat_rtt_ms_sum)
|
||||||
|
.sum::<i64>(),
|
||||||
|
33
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
metrics
|
||||||
|
.iter()
|
||||||
|
.map(|bucket| bucket.connect_errors_delta)
|
||||||
|
.sum::<i64>(),
|
||||||
|
4
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
metrics
|
||||||
|
.iter()
|
||||||
|
.map(|bucket| bucket.ws_out_frames_delta)
|
||||||
|
.sum::<i64>(),
|
||||||
|
6
|
||||||
|
);
|
||||||
|
|
||||||
let fleet = repository
|
let fleet = repository
|
||||||
.list_proxy_fleet_metrics(
|
.list_proxy_fleet_metrics(
|
||||||
@@ -1626,9 +1684,15 @@ VALUES ('node-1', 'registered', 'ok', 3)
|
|||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
.expect("fleet metrics should list");
|
.expect("fleet metrics should list");
|
||||||
assert_eq!(fleet.len(), 1);
|
assert!(!fleet.is_empty());
|
||||||
assert_eq!(fleet[0].samples, 1);
|
assert_eq!(fleet.iter().map(|bucket| bucket.samples).sum::<i64>(), 2);
|
||||||
assert_eq!(fleet[0].error_events_delta, 1);
|
assert_eq!(
|
||||||
|
fleet
|
||||||
|
.iter()
|
||||||
|
.map(|bucket| bucket.error_events_delta)
|
||||||
|
.sum::<i64>(),
|
||||||
|
1
|
||||||
|
);
|
||||||
|
|
||||||
let events = repository
|
let events = repository
|
||||||
.list_proxy_node_events_filtered(
|
.list_proxy_node_events_filtered(
|
||||||
|
|||||||
@@ -587,7 +587,8 @@ fn json_string(value: Option<&Value>) -> Option<String> {
|
|||||||
fn counter_delta_u64(previous: Option<u64>, current: u64) -> u64 {
|
fn counter_delta_u64(previous: Option<u64>, current: u64) -> u64 {
|
||||||
match previous {
|
match previous {
|
||||||
Some(previous) if current >= previous => current - previous,
|
Some(previous) if current >= previous => current - previous,
|
||||||
Some(_) | None => current,
|
Some(_) => current,
|
||||||
|
None => 0,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -962,6 +963,39 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn builds_tunnel_metrics_sample_uses_first_counter_report_as_baseline() {
|
||||||
|
let current = json!({
|
||||||
|
"tunnel_metrics": {
|
||||||
|
"connect_errors": 12,
|
||||||
|
"disconnects": 5,
|
||||||
|
"error_events_total": 7,
|
||||||
|
"ws_in_bytes": 1_500,
|
||||||
|
"ws_out_bytes": 2_500,
|
||||||
|
"ws_in_frames": 15,
|
||||||
|
"ws_out_frames": 25,
|
||||||
|
"heartbeat_rtt_last_ms": 44
|
||||||
|
},
|
||||||
|
"recent_tunnel_errors": [
|
||||||
|
{"timestamp_unix_secs": 101, "category": "newer", "message": "new"}
|
||||||
|
]
|
||||||
|
});
|
||||||
|
|
||||||
|
let sample = build_tunnel_metrics_sample(None, Some(¤t), 4, true)
|
||||||
|
.expect("sample should build");
|
||||||
|
|
||||||
|
assert_eq!(sample.samples, 1);
|
||||||
|
assert_eq!(sample.heartbeat_rtt_ms_sum, 44);
|
||||||
|
assert_eq!(sample.connect_errors_delta, 0);
|
||||||
|
assert_eq!(sample.disconnects_delta, 0);
|
||||||
|
assert_eq!(sample.error_events_delta, 0);
|
||||||
|
assert_eq!(sample.ws_in_bytes_delta, 0);
|
||||||
|
assert_eq!(sample.ws_out_bytes_delta, 0);
|
||||||
|
assert_eq!(sample.ws_in_frames_delta, 0);
|
||||||
|
assert_eq!(sample.ws_out_frames_delta, 0);
|
||||||
|
assert!(sample.recent_error_events.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn preserves_secure_tunnel_metadata_across_heartbeat_metadata_refresh() {
|
fn preserves_secure_tunnel_metadata_across_heartbeat_metadata_refresh() {
|
||||||
let previous = json!({
|
let previous = json!({
|
||||||
|
|||||||
@@ -309,6 +309,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
},
|
},
|
||||||
));
|
));
|
||||||
@@ -333,6 +334,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
}),
|
}),
|
||||||
None
|
None
|
||||||
@@ -358,6 +360,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
}),
|
}),
|
||||||
Some("provider_key_concurrency_limit_reached")
|
Some("provider_key_concurrency_limit_reached")
|
||||||
@@ -386,6 +389,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
}),
|
}),
|
||||||
Some("provider_key_concurrency_limit_reached")
|
Some("provider_key_concurrency_limit_reached")
|
||||||
@@ -415,6 +419,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
}),
|
}),
|
||||||
None
|
None
|
||||||
@@ -435,6 +440,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
}),
|
}),
|
||||||
None
|
None
|
||||||
@@ -458,6 +464,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
}),
|
}),
|
||||||
None
|
None
|
||||||
@@ -481,6 +488,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
}),
|
}),
|
||||||
Some("key_circuit_open")
|
Some("key_circuit_open")
|
||||||
@@ -516,6 +524,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
}),
|
}),
|
||||||
Some("key_rpm_exhausted")
|
Some("key_rpm_exhausted")
|
||||||
@@ -543,6 +552,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
}),
|
}),
|
||||||
Some("key_circuit_open")
|
Some("key_circuit_open")
|
||||||
@@ -557,6 +567,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
}),
|
}),
|
||||||
None
|
None
|
||||||
@@ -577,6 +588,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
},
|
},
|
||||||
));
|
));
|
||||||
@@ -590,6 +602,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: true,
|
provider_quota_blocks_requests: true,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
},
|
},
|
||||||
));
|
));
|
||||||
@@ -607,6 +620,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: true,
|
account_quota_exhausted: true,
|
||||||
oauth_invalid: false,
|
oauth_invalid: false,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
},
|
},
|
||||||
));
|
));
|
||||||
@@ -624,6 +638,7 @@ mod tests {
|
|||||||
provider_quota_blocks_requests: false,
|
provider_quota_blocks_requests: false,
|
||||||
account_quota_exhausted: false,
|
account_quota_exhausted: false,
|
||||||
oauth_invalid: true,
|
oauth_invalid: true,
|
||||||
|
enforce_key_circuit_breaker: true,
|
||||||
rpm_reset_at: None,
|
rpm_reset_at: None,
|
||||||
},
|
},
|
||||||
));
|
));
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ pub struct CandidateRuntimeSelectabilityInput<'a> {
|
|||||||
pub provider_quota_blocks_requests: bool,
|
pub provider_quota_blocks_requests: bool,
|
||||||
pub account_quota_exhausted: bool,
|
pub account_quota_exhausted: bool,
|
||||||
pub oauth_invalid: bool,
|
pub oauth_invalid: bool,
|
||||||
|
pub enforce_key_circuit_breaker: bool,
|
||||||
pub rpm_reset_at: Option<u64>,
|
pub rpm_reset_at: Option<u64>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -50,6 +51,7 @@ pub fn candidate_runtime_skip_reason_with_state(
|
|||||||
provider_quota_blocks_requests,
|
provider_quota_blocks_requests,
|
||||||
account_quota_exhausted,
|
account_quota_exhausted,
|
||||||
oauth_invalid,
|
oauth_invalid,
|
||||||
|
enforce_key_circuit_breaker,
|
||||||
rpm_reset_at,
|
rpm_reset_at,
|
||||||
} = input;
|
} = input;
|
||||||
|
|
||||||
@@ -94,11 +96,13 @@ pub fn candidate_runtime_skip_reason_with_state(
|
|||||||
}
|
}
|
||||||
|
|
||||||
if let Some(provider_key) = provider_key {
|
if let Some(provider_key) = provider_key {
|
||||||
if crate::is_provider_key_circuit_open_at(
|
if enforce_key_circuit_breaker
|
||||||
provider_key,
|
&& crate::is_provider_key_circuit_open_at(
|
||||||
candidate.endpoint_api_format.as_str(),
|
provider_key,
|
||||||
now_unix_secs,
|
candidate.endpoint_api_format.as_str(),
|
||||||
) {
|
now_unix_secs,
|
||||||
|
)
|
||||||
|
{
|
||||||
return Some("key_circuit_open");
|
return Some("key_circuit_open");
|
||||||
}
|
}
|
||||||
if crate::provider_key_health_score(provider_key, candidate.endpoint_api_format.as_str())
|
if crate::provider_key_health_score(provider_key, candidate.endpoint_api_format.as_str())
|
||||||
|
|||||||
Reference in New Issue
Block a user