mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat(tunnel/usage): proxy writer 双优先级队列、hub 连接压力感知选择、usage 请求记录级别控制及 trace 页面 proxy timing 增强
This commit is contained in:
@@ -16,16 +16,37 @@ use aether_data_contracts::repository::settlement::{StoredUsageSettlement, Usage
|
|||||||
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
|
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
|
||||||
use aether_data_contracts::repository::video_tasks::{StoredVideoTask, VideoTaskLookupKey};
|
use aether_data_contracts::repository::video_tasks::{StoredVideoTask, VideoTaskLookupKey};
|
||||||
use aether_usage_runtime::{
|
use aether_usage_runtime::{
|
||||||
UsageBillingEventEnricher, UsageEvent, UsageRecordWriter, UsageRuntimeAccess,
|
UsageBillingEventEnricher, UsageEvent, UsageRecordWriter, UsageRequestRecordLevel,
|
||||||
UsageSettlementWriter,
|
UsageRuntimeAccess, UsageSettlementWriter,
|
||||||
};
|
};
|
||||||
use aether_video_tasks_core::StoredVideoTaskReadSide;
|
use aether_video_tasks_core::StoredVideoTaskReadSide;
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
use super::GatewayDataState;
|
use super::GatewayDataState;
|
||||||
use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
|
use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
|
||||||
use crate::provider_transport::ProviderTransportSnapshotSource;
|
use crate::provider_transport::ProviderTransportSnapshotSource;
|
||||||
|
|
||||||
|
const REQUEST_RECORD_LEVEL_KEY: &str = "request_record_level";
|
||||||
|
const LEGACY_REQUEST_LOG_LEVEL_KEY: &str = "request_log_level";
|
||||||
|
|
||||||
|
fn usage_request_record_level_from_value(value: Option<&Value>) -> UsageRequestRecordLevel {
|
||||||
|
let Some(value) = value.and_then(Value::as_str).map(str::trim) else {
|
||||||
|
return UsageRequestRecordLevel::Full;
|
||||||
|
};
|
||||||
|
|
||||||
|
if value.eq_ignore_ascii_case("basic")
|
||||||
|
|| value.eq_ignore_ascii_case("base")
|
||||||
|
|| value.eq_ignore_ascii_case("headers")
|
||||||
|
|| value.eq_ignore_ascii_case("minimal")
|
||||||
|
|| value.eq_ignore_ascii_case("none")
|
||||||
|
{
|
||||||
|
UsageRequestRecordLevel::Basic
|
||||||
|
} else {
|
||||||
|
UsageRequestRecordLevel::Full
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
impl RequestAuditReader for GatewayDataState {
|
impl RequestAuditReader for GatewayDataState {
|
||||||
async fn find_request_usage_audit_by_request_id(
|
async fn find_request_usage_audit_by_request_id(
|
||||||
@@ -161,6 +182,7 @@ impl UsageBillingEventEnricher for GatewayDataState {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
impl UsageRuntimeAccess for GatewayDataState {
|
impl UsageRuntimeAccess for GatewayDataState {
|
||||||
fn has_usage_writer(&self) -> bool {
|
fn has_usage_writer(&self) -> bool {
|
||||||
GatewayDataState::has_usage_writer(self)
|
GatewayDataState::has_usage_writer(self)
|
||||||
@@ -173,6 +195,16 @@ impl UsageRuntimeAccess for GatewayDataState {
|
|||||||
fn usage_worker_runner(&self) -> Option<RedisStreamRunner> {
|
fn usage_worker_runner(&self) -> Option<RedisStreamRunner> {
|
||||||
GatewayDataState::usage_worker_runner(self)
|
GatewayDataState::usage_worker_runner(self)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn request_record_level(&self) -> Result<UsageRequestRecordLevel, DataLayerError> {
|
||||||
|
let value = GatewayDataState::find_system_config_value(self, REQUEST_RECORD_LEVEL_KEY)
|
||||||
|
.await?
|
||||||
|
.or(
|
||||||
|
GatewayDataState::find_system_config_value(self, LEGACY_REQUEST_LOG_LEVEL_KEY)
|
||||||
|
.await?,
|
||||||
|
);
|
||||||
|
Ok(usage_request_record_level_from_value(value.as_ref()))
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
@@ -188,10 +220,11 @@ impl UsageRecordWriter for GatewayDataState {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use aether_billing::enrich_usage_event_with_billing;
|
use aether_billing::enrich_usage_event_with_billing;
|
||||||
use serde_json::Value;
|
use aether_usage_runtime::UsageRuntimeAccess;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
use super::GatewayDataState;
|
use super::GatewayDataState;
|
||||||
use crate::usage::{UsageEvent, UsageEventData, UsageEventType};
|
use crate::usage::{UsageEvent, UsageEventData, UsageEventType, UsageRequestRecordLevel};
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn enriches_completed_usage_event_with_billing_snapshot() {
|
async fn enriches_completed_usage_event_with_billing_snapshot() {
|
||||||
@@ -255,4 +288,43 @@ mod tests {
|
|||||||
Some("complete")
|
Some("complete")
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn usage_runtime_access_reads_base_request_record_level_as_basic() {
|
||||||
|
let state = GatewayDataState::disabled().with_system_config_values_for_tests([(
|
||||||
|
"request_record_level".to_string(),
|
||||||
|
json!("base"),
|
||||||
|
)]);
|
||||||
|
|
||||||
|
let level = UsageRuntimeAccess::request_record_level(&state)
|
||||||
|
.await
|
||||||
|
.expect("request record level should read");
|
||||||
|
|
||||||
|
assert_eq!(level, UsageRequestRecordLevel::Basic);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn usage_runtime_access_falls_back_to_legacy_request_log_level_alias() {
|
||||||
|
let state = GatewayDataState::disabled().with_system_config_values_for_tests([(
|
||||||
|
"request_log_level".to_string(),
|
||||||
|
json!("headers"),
|
||||||
|
)]);
|
||||||
|
|
||||||
|
let level = UsageRuntimeAccess::request_record_level(&state)
|
||||||
|
.await
|
||||||
|
.expect("legacy request log level should read");
|
||||||
|
|
||||||
|
assert_eq!(level, UsageRequestRecordLevel::Basic);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn usage_runtime_access_defaults_missing_request_record_level_to_full() {
|
||||||
|
let state = GatewayDataState::disabled();
|
||||||
|
|
||||||
|
let level = UsageRuntimeAccess::request_record_level(&state)
|
||||||
|
.await
|
||||||
|
.expect("missing request record level should fall back");
|
||||||
|
|
||||||
|
assert_eq!(level, UsageRequestRecordLevel::Full);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
use super::super::test_support::{
|
use super::super::test_support::{
|
||||||
request_context, sample_candidate, sample_endpoint, sample_key, sample_provider,
|
request_context, sample_candidate, sample_endpoint, sample_key, sample_provider, sample_usage,
|
||||||
};
|
};
|
||||||
use super::local_monitoring_response;
|
use super::local_monitoring_response;
|
||||||
use crate::AppState;
|
use crate::AppState;
|
||||||
@@ -10,6 +10,7 @@ use std::sync::Arc;
|
|||||||
|
|
||||||
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
|
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
|
||||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||||
|
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn admin_monitoring_trace_request_returns_local_payload() {
|
async fn admin_monitoring_trace_request_returns_local_payload() {
|
||||||
@@ -156,6 +157,88 @@ async fn admin_monitoring_trace_request_keeps_format_conversion_disabled_candida
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn admin_monitoring_trace_request_enriches_proxy_timing_from_usage_audit() {
|
||||||
|
let mut candidate = sample_candidate(
|
||||||
|
"cand-used",
|
||||||
|
"request-1",
|
||||||
|
1,
|
||||||
|
RequestCandidateStatus::Success,
|
||||||
|
Some(101),
|
||||||
|
Some(33),
|
||||||
|
Some(200),
|
||||||
|
);
|
||||||
|
candidate.extra_data = Some(json!({
|
||||||
|
"proxy": {
|
||||||
|
"node_id": "proxy-node-1",
|
||||||
|
"node_name": "edge-1",
|
||||||
|
"source": "provider"
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
|
||||||
|
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![candidate]));
|
||||||
|
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||||
|
vec![sample_provider()],
|
||||||
|
vec![sample_endpoint()],
|
||||||
|
vec![sample_key()],
|
||||||
|
));
|
||||||
|
let mut usage = sample_usage(
|
||||||
|
"request-1",
|
||||||
|
"provider-1",
|
||||||
|
"OpenAI",
|
||||||
|
40,
|
||||||
|
0.02,
|
||||||
|
"completed",
|
||||||
|
Some(200),
|
||||||
|
100,
|
||||||
|
);
|
||||||
|
usage.candidate_id = Some("cand-used".to_string());
|
||||||
|
usage.response_headers = Some(json!({
|
||||||
|
"x-proxy-timing": "{\"connection_acquire_ms\":125,\"response_wait_ms\":475,\"ttfb_ms\":600}"
|
||||||
|
}));
|
||||||
|
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
|
||||||
|
let data_state =
|
||||||
|
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
|
||||||
|
request_candidates,
|
||||||
|
usage_repository,
|
||||||
|
)
|
||||||
|
.with_provider_catalog_reader(provider_catalog);
|
||||||
|
let state = AppState::new()
|
||||||
|
.expect("state should build")
|
||||||
|
.with_data_state_for_tests(data_state);
|
||||||
|
let context = request_context(
|
||||||
|
http::Method::GET,
|
||||||
|
"/api/admin/monitoring/trace/request-1?attempted_only=true",
|
||||||
|
);
|
||||||
|
|
||||||
|
let response = local_monitoring_response(&state, &context)
|
||||||
|
.await
|
||||||
|
.expect("handler should not error")
|
||||||
|
.expect("route should be handled locally");
|
||||||
|
|
||||||
|
assert_eq!(response.status(), http::StatusCode::OK);
|
||||||
|
let body = to_bytes(response.into_body(), usize::MAX)
|
||||||
|
.await
|
||||||
|
.expect("body should read");
|
||||||
|
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
|
||||||
|
assert_eq!(
|
||||||
|
payload["candidates"][0]["extra_data"]["first_byte_time_ms"],
|
||||||
|
json!(30)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
payload["candidates"][0]["extra_data"]["proxy"]["ttfb_ms"],
|
||||||
|
json!(600)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
payload["candidates"][0]["extra_data"]["proxy"]["timing"]["connection_acquire_ms"],
|
||||||
|
json!(125)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
payload["candidates"][0]["extra_data"]["proxy"]["timing"]["response_wait_ms"],
|
||||||
|
json!(475)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn admin_monitoring_trace_provider_stats_returns_local_payload() {
|
async fn admin_monitoring_trace_provider_stats_returns_local_payload() {
|
||||||
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
|
||||||
|
|||||||
@@ -51,9 +51,15 @@ pub(super) async fn build_admin_monitoring_trace_request_response(
|
|||||||
attempted_only,
|
attempted_only,
|
||||||
));
|
));
|
||||||
};
|
};
|
||||||
|
let usage = state
|
||||||
|
.data
|
||||||
|
.read_request_usage_audit(&request_id)
|
||||||
|
.await
|
||||||
|
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||||
|
|
||||||
Ok(build_admin_monitoring_trace_request_payload_response(
|
Ok(build_admin_monitoring_trace_request_payload_response(
|
||||||
&trace,
|
&trace,
|
||||||
|
usage.as_ref(),
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -371,6 +371,151 @@ async fn gateway_truncates_deep_request_echo_for_local_openai_chat_sync_usage()
|
|||||||
upstream_handle.abort();
|
upstream_handle.abort();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn gateway_strips_request_and_response_bodies_when_request_record_level_is_base() {
|
||||||
|
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||||
|
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||||
|
|
||||||
|
let upstream = Router::new().route(
|
||||||
|
"/api/internal/gateway/report-sync",
|
||||||
|
any(|_request: Request| async move { Json(json!({"ok": true})) }),
|
||||||
|
);
|
||||||
|
|
||||||
|
let execution_runtime = Router::new().route(
|
||||||
|
"/v1/execute/sync",
|
||||||
|
any(|_request: Request| async move {
|
||||||
|
Json(json!({
|
||||||
|
"request_id": "trace-openai-chat-local-report-sync-base-123",
|
||||||
|
"status_code": 200,
|
||||||
|
"headers": {
|
||||||
|
"content-type": "application/json"
|
||||||
|
},
|
||||||
|
"body": {
|
||||||
|
"json_body": {
|
||||||
|
"id": "chatcmpl-local-report-sync-base-123",
|
||||||
|
"object": "chat.completion",
|
||||||
|
"model": "gpt-5-upstream",
|
||||||
|
"choices": [{
|
||||||
|
"index": 0,
|
||||||
|
"message": {
|
||||||
|
"role": "assistant",
|
||||||
|
"content": "body should not be persisted"
|
||||||
|
}
|
||||||
|
}],
|
||||||
|
"usage": {
|
||||||
|
"prompt_tokens": 2,
|
||||||
|
"completion_tokens": 3,
|
||||||
|
"total_tokens": 5
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"telemetry": {
|
||||||
|
"elapsed_ms": 25
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||||
|
Some(hash_api_key("sk-client-openai-local-report-sync-base")),
|
||||||
|
sample_local_openai_auth_snapshot(
|
||||||
|
"api-key-openai-usage-local-base-1",
|
||||||
|
"user-openai-usage-local-base-1",
|
||||||
|
),
|
||||||
|
)]));
|
||||||
|
let candidate_selection_repository =
|
||||||
|
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||||
|
sample_local_openai_candidate_row(),
|
||||||
|
]));
|
||||||
|
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||||
|
vec![sample_local_openai_provider()],
|
||||||
|
vec![sample_local_openai_endpoint()],
|
||||||
|
vec![sample_local_openai_key()],
|
||||||
|
));
|
||||||
|
|
||||||
|
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||||
|
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||||
|
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
|
||||||
|
.with_data_state_for_tests(
|
||||||
|
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
|
||||||
|
auth_repository,
|
||||||
|
candidate_selection_repository,
|
||||||
|
provider_catalog_repository,
|
||||||
|
Arc::clone(&request_candidate_repository),
|
||||||
|
Arc::clone(&usage_repository),
|
||||||
|
DEVELOPMENT_ENCRYPTION_KEY,
|
||||||
|
)
|
||||||
|
.with_system_config_values_for_tests([(
|
||||||
|
"request_record_level".to_string(),
|
||||||
|
json!("base"),
|
||||||
|
)]),
|
||||||
|
)
|
||||||
|
.with_usage_runtime_for_tests(UsageRuntimeConfig {
|
||||||
|
enabled: true,
|
||||||
|
..UsageRuntimeConfig::default()
|
||||||
|
});
|
||||||
|
let gateway = build_router_with_state(gateway_state);
|
||||||
|
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||||
|
|
||||||
|
let response = reqwest::Client::new()
|
||||||
|
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||||
|
.header(http::header::CONTENT_TYPE, "application/json")
|
||||||
|
.header(
|
||||||
|
http::header::AUTHORIZATION,
|
||||||
|
"Bearer sk-client-openai-local-report-sync-base",
|
||||||
|
)
|
||||||
|
.header(
|
||||||
|
TRACE_ID_HEADER,
|
||||||
|
"trace-openai-chat-local-report-sync-base-123",
|
||||||
|
)
|
||||||
|
.body(
|
||||||
|
serde_json::to_string(&json!({
|
||||||
|
"model": "gpt-5",
|
||||||
|
"messages": [{
|
||||||
|
"role": "user",
|
||||||
|
"content": "request body should not be persisted"
|
||||||
|
}]
|
||||||
|
}))
|
||||||
|
.expect("request should encode"),
|
||||||
|
)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.expect("request should succeed");
|
||||||
|
|
||||||
|
assert_eq!(response.status(), StatusCode::OK);
|
||||||
|
let body_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||||
|
assert_eq!(body_json["model"], "gpt-5-upstream");
|
||||||
|
|
||||||
|
let stored_usage = wait_for_usage_status(
|
||||||
|
usage_repository.as_ref(),
|
||||||
|
"trace-openai-chat-local-report-sync-base-123",
|
||||||
|
"completed",
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert_eq!(stored_usage.status, "completed");
|
||||||
|
assert_eq!(stored_usage.total_tokens, 5);
|
||||||
|
assert_eq!(stored_usage.response_time_ms, Some(25));
|
||||||
|
assert!(stored_usage.request_body.is_none());
|
||||||
|
assert!(stored_usage.request_body_ref.is_none());
|
||||||
|
assert!(stored_usage.provider_request_body.is_none());
|
||||||
|
assert!(stored_usage.provider_request_body_ref.is_none());
|
||||||
|
assert!(stored_usage.response_body.is_none());
|
||||||
|
assert!(stored_usage.response_body_ref.is_none());
|
||||||
|
assert!(stored_usage.client_response_body.is_none());
|
||||||
|
assert!(stored_usage.client_response_body_ref.is_none());
|
||||||
|
|
||||||
|
let stored_candidates = request_candidate_repository
|
||||||
|
.list_by_request_id("trace-openai-chat-local-report-sync-base-123")
|
||||||
|
.await
|
||||||
|
.expect("request candidate trace should read");
|
||||||
|
assert_eq!(stored_candidates.len(), 1);
|
||||||
|
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||||
|
|
||||||
|
gateway_handle.abort();
|
||||||
|
execution_runtime_handle.abort();
|
||||||
|
upstream_handle.abort();
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exhaust_after_retryable_sync_failure(
|
async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exhaust_after_retryable_sync_failure(
|
||||||
) {
|
) {
|
||||||
|
|||||||
@@ -3,7 +3,7 @@ use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering}
|
|||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
use aether_runtime::{BoundedQueueSender, MetricKind, MetricSample, QueueSendError};
|
use aether_runtime::{BoundedQueueSender, MetricKind, MetricSample, QueueSendError, QueueSnapshot};
|
||||||
use axum::extract::ws::Message;
|
use axum::extract::ws::Message;
|
||||||
use bytes::Bytes;
|
use bytes::Bytes;
|
||||||
use dashmap::DashMap;
|
use dashmap::DashMap;
|
||||||
@@ -16,6 +16,8 @@ use super::control_plane::ControlPlaneClient;
|
|||||||
use super::protocol;
|
use super::protocol;
|
||||||
|
|
||||||
const MAX_REQUEST_BODY_FRAME_SIZE: usize = 32 * 1024;
|
const MAX_REQUEST_BODY_FRAME_SIZE: usize = 32 * 1024;
|
||||||
|
const SOFT_AVOID_QUEUE_PRESSURE_PERCENT: u64 = 50;
|
||||||
|
const SOFT_AVOID_STREAM_PRESSURE_PERCENT: u64 = 85;
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
pub enum SendStatus {
|
pub enum SendStatus {
|
||||||
@@ -75,6 +77,10 @@ impl BoundedOutbound {
|
|||||||
let _ = self.close_tx.send(true);
|
let _ = self.close_tx.send(true);
|
||||||
true
|
true
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn snapshot(&self) -> QueueSnapshot {
|
||||||
|
self.tx.snapshot()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub struct ProxyConn {
|
pub struct ProxyConn {
|
||||||
@@ -86,6 +92,7 @@ pub struct ProxyConn {
|
|||||||
pub stream_count: AtomicUsize,
|
pub stream_count: AtomicUsize,
|
||||||
pub max_streams: usize,
|
pub max_streams: usize,
|
||||||
draining: AtomicBool,
|
draining: AtomicBool,
|
||||||
|
congested_total: AtomicU64,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ProxyConn {
|
impl ProxyConn {
|
||||||
@@ -106,6 +113,7 @@ impl ProxyConn {
|
|||||||
stream_count: AtomicUsize::new(0),
|
stream_count: AtomicUsize::new(0),
|
||||||
max_streams,
|
max_streams,
|
||||||
draining: AtomicBool::new(false),
|
draining: AtomicBool::new(false),
|
||||||
|
congested_total: AtomicU64::new(0),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -180,6 +188,7 @@ impl ProxyConn {
|
|||||||
let was_closing = self.outbound.is_closing();
|
let was_closing = self.outbound.is_closing();
|
||||||
let status = self.outbound.send(msg);
|
let status = self.outbound.send(msg);
|
||||||
if status == SendStatus::Congested && !was_closing {
|
if status == SendStatus::Congested && !was_closing {
|
||||||
|
self.congested_total.fetch_add(1, Ordering::Relaxed);
|
||||||
warn!(
|
warn!(
|
||||||
conn_id = self.id,
|
conn_id = self.id,
|
||||||
node_id = %self.node_id,
|
node_id = %self.node_id,
|
||||||
@@ -190,6 +199,62 @@ impl ProxyConn {
|
|||||||
}
|
}
|
||||||
status
|
status
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn snapshot(&self) -> ProxyConnSnapshot {
|
||||||
|
let outbound = self.outbound.snapshot();
|
||||||
|
let stream_count = self.stream_count.load(Ordering::Relaxed);
|
||||||
|
let queue_pressure_percent = percent_u64(outbound.depth, outbound.capacity);
|
||||||
|
let stream_pressure_percent = percent_u64(stream_count, self.max_streams);
|
||||||
|
let soft_avoid = queue_pressure_percent >= SOFT_AVOID_QUEUE_PRESSURE_PERCENT
|
||||||
|
|| stream_pressure_percent >= SOFT_AVOID_STREAM_PRESSURE_PERCENT;
|
||||||
|
ProxyConnSnapshot {
|
||||||
|
conn_id: self.id,
|
||||||
|
available: self.is_available(),
|
||||||
|
closing: self.outbound.is_closing(),
|
||||||
|
draining: self.is_draining(),
|
||||||
|
stream_count,
|
||||||
|
max_streams: self.max_streams,
|
||||||
|
stream_pressure_percent,
|
||||||
|
outbound,
|
||||||
|
queue_pressure_percent,
|
||||||
|
soft_avoid,
|
||||||
|
congested_total: self.congested_total.load(Ordering::Relaxed),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy)]
|
||||||
|
struct ProxyConnSnapshot {
|
||||||
|
conn_id: u64,
|
||||||
|
available: bool,
|
||||||
|
closing: bool,
|
||||||
|
draining: bool,
|
||||||
|
stream_count: usize,
|
||||||
|
max_streams: usize,
|
||||||
|
stream_pressure_percent: u64,
|
||||||
|
outbound: QueueSnapshot,
|
||||||
|
queue_pressure_percent: u64,
|
||||||
|
soft_avoid: bool,
|
||||||
|
congested_total: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct ProxyConnCandidate {
|
||||||
|
conn: Arc<ProxyConn>,
|
||||||
|
snapshot: ProxyConnSnapshot,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ProxyConnCandidate {
|
||||||
|
fn rank_key(&self) -> (u8, u64, u64, usize, usize, u64) {
|
||||||
|
(
|
||||||
|
u8::from(self.snapshot.soft_avoid),
|
||||||
|
self.snapshot.queue_pressure_percent,
|
||||||
|
self.snapshot.stream_pressure_percent,
|
||||||
|
self.snapshot.outbound.depth,
|
||||||
|
self.snapshot.stream_count,
|
||||||
|
self.snapshot.conn_id,
|
||||||
|
)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
@@ -343,6 +408,9 @@ pub struct HubRouter {
|
|||||||
next_local_stream_id: AtomicU64,
|
next_local_stream_id: AtomicU64,
|
||||||
control_plane: ControlPlaneClient,
|
control_plane: ControlPlaneClient,
|
||||||
node_status_tx: mpsc::UnboundedSender<NodeStatusEvent>,
|
node_status_tx: mpsc::UnboundedSender<NodeStatusEvent>,
|
||||||
|
soft_avoid_selection_total: AtomicU64,
|
||||||
|
selection_retry_total: AtomicU64,
|
||||||
|
selection_unavailable_total: AtomicU64,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
@@ -391,6 +459,9 @@ impl HubRouter {
|
|||||||
next_local_stream_id: AtomicU64::new(1),
|
next_local_stream_id: AtomicU64::new(1),
|
||||||
control_plane,
|
control_plane,
|
||||||
node_status_tx,
|
node_status_tx,
|
||||||
|
soft_avoid_selection_total: AtomicU64::new(0),
|
||||||
|
selection_retry_total: AtomicU64::new(0),
|
||||||
|
selection_unavailable_total: AtomicU64::new(0),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -477,28 +548,28 @@ impl HubRouter {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn get_proxy_conn(&self, node_id: &str) -> Option<Arc<ProxyConn>> {
|
fn ranked_proxy_conn_candidates(&self, node_id: &str) -> Vec<ProxyConnCandidate> {
|
||||||
|
let conns = {
|
||||||
let map = self.proxy_conns.read();
|
let map = self.proxy_conns.read();
|
||||||
let conns = map.get(node_id)?;
|
map.get(node_id)
|
||||||
let result = conns
|
.map(|entries| entries.to_vec())
|
||||||
.iter()
|
.unwrap_or_default()
|
||||||
.filter(|c| c.is_available())
|
};
|
||||||
.min_by_key(|c| c.stream_count.load(Ordering::Relaxed))
|
let mut candidates = conns
|
||||||
.cloned();
|
.into_iter()
|
||||||
if result.is_none() && !conns.is_empty() {
|
.filter_map(|conn| {
|
||||||
warn!(
|
let snapshot = conn.snapshot();
|
||||||
node_id = %node_id,
|
snapshot
|
||||||
total_conns = conns.len(),
|
.available
|
||||||
closing = conns.iter().filter(|c| c.outbound.is_closing()).count(),
|
.then_some(ProxyConnCandidate { conn, snapshot })
|
||||||
draining = conns.iter().filter(|c| c.is_draining()).count(),
|
})
|
||||||
"no available proxy connection despite registered connections"
|
.collect::<Vec<_>>();
|
||||||
);
|
candidates.sort_by_key(|candidate| candidate.rank_key());
|
||||||
}
|
candidates
|
||||||
result
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn has_local_proxy(&self, node_id: &str) -> bool {
|
pub fn has_local_proxy(&self, node_id: &str) -> bool {
|
||||||
self.get_proxy_conn(node_id).is_some()
|
!self.ranked_proxy_conn_candidates(node_id).is_empty()
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn open_local_stream(
|
pub fn open_local_stream(
|
||||||
@@ -506,12 +577,49 @@ impl HubRouter {
|
|||||||
node_id: &str,
|
node_id: &str,
|
||||||
meta: &protocol::RequestMeta,
|
meta: &protocol::RequestMeta,
|
||||||
) -> Result<Arc<LocalStream>, String> {
|
) -> Result<Arc<LocalStream>, String> {
|
||||||
let proxy_conn = self
|
let candidates = self.ranked_proxy_conn_candidates(node_id);
|
||||||
.get_proxy_conn(node_id)
|
if candidates.is_empty() {
|
||||||
.ok_or_else(|| format!("no proxy connection for node {node_id}"))?;
|
self.selection_unavailable_total
|
||||||
let proxy_stream_id = proxy_conn
|
.fetch_add(1, Ordering::Relaxed);
|
||||||
.alloc_stream_id()
|
self.warn_no_available_proxy_connection(node_id);
|
||||||
.ok_or_else(|| format!("stream limit reached for node {node_id}"))?;
|
return Err(format!("no proxy connection for node {node_id}"));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut skipped_candidates = 0usize;
|
||||||
|
let mut selected_candidate = None;
|
||||||
|
let mut proxy_stream_id = None;
|
||||||
|
for candidate in candidates {
|
||||||
|
match candidate.conn.alloc_stream_id() {
|
||||||
|
Some(stream_id) => {
|
||||||
|
proxy_stream_id = Some(stream_id);
|
||||||
|
selected_candidate = Some(candidate);
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
None => skipped_candidates = skipped_candidates.saturating_add(1),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let Some(candidate) = selected_candidate else {
|
||||||
|
self.selection_unavailable_total
|
||||||
|
.fetch_add(1, Ordering::Relaxed);
|
||||||
|
return Err(format!("stream limit reached for node {node_id}"));
|
||||||
|
};
|
||||||
|
if skipped_candidates > 0 {
|
||||||
|
self.selection_retry_total.fetch_add(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
if candidate.snapshot.soft_avoid {
|
||||||
|
self.soft_avoid_selection_total
|
||||||
|
.fetch_add(1, Ordering::Relaxed);
|
||||||
|
debug!(
|
||||||
|
node_id = %node_id,
|
||||||
|
conn_id = candidate.snapshot.conn_id,
|
||||||
|
queue_pressure_percent = candidate.snapshot.queue_pressure_percent,
|
||||||
|
stream_pressure_percent = candidate.snapshot.stream_pressure_percent,
|
||||||
|
"selected high-pressure proxy connection because no lower-pressure alternative was available"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let proxy_conn = candidate.conn;
|
||||||
|
let proxy_stream_id = proxy_stream_id.expect("selected candidate should carry a stream id");
|
||||||
|
|
||||||
// Encode frames before registering the stream so that encoding failures
|
// Encode frames before registering the stream so that encoding failures
|
||||||
// (practically impossible but theoretically possible) don't leak a stream
|
// (practically impossible but theoretically possible) don't leak a stream
|
||||||
@@ -556,6 +664,7 @@ impl HubRouter {
|
|||||||
proxy_stream_id = proxy_stream_id,
|
proxy_stream_id = proxy_stream_id,
|
||||||
local_stream_id = local_stream_id,
|
local_stream_id = local_stream_id,
|
||||||
stream_count = proxy_conn.stream_count.load(Ordering::Relaxed),
|
stream_count = proxy_conn.stream_count.load(Ordering::Relaxed),
|
||||||
|
queue_depth = proxy_conn.outbound.snapshot().depth,
|
||||||
send_status = ?send_status,
|
send_status = ?send_status,
|
||||||
"open_local_stream dispatched"
|
"open_local_stream dispatched"
|
||||||
);
|
);
|
||||||
@@ -569,6 +678,28 @@ impl HubRouter {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn warn_no_available_proxy_connection(&self, node_id: &str) {
|
||||||
|
let conns = {
|
||||||
|
let map = self.proxy_conns.read();
|
||||||
|
map.get(node_id)
|
||||||
|
.map(|entries| entries.to_vec())
|
||||||
|
.unwrap_or_default()
|
||||||
|
};
|
||||||
|
if conns.is_empty() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let snapshots = conns.iter().map(|conn| conn.snapshot()).collect::<Vec<_>>();
|
||||||
|
warn!(
|
||||||
|
node_id = %node_id,
|
||||||
|
total_conns = snapshots.len(),
|
||||||
|
available = snapshots.iter().filter(|snapshot| snapshot.available).count(),
|
||||||
|
closing = snapshots.iter().filter(|snapshot| snapshot.closing).count(),
|
||||||
|
draining = snapshots.iter().filter(|snapshot| snapshot.draining).count(),
|
||||||
|
soft_avoid = snapshots.iter().filter(|snapshot| snapshot.soft_avoid).count(),
|
||||||
|
"no available proxy connection despite registered connections"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
pub fn push_local_request_body(
|
pub fn push_local_request_body(
|
||||||
&self,
|
&self,
|
||||||
local_stream_id: u64,
|
local_stream_id: u64,
|
||||||
@@ -873,15 +1004,72 @@ impl HubRouter {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn stats(&self) -> HubStats {
|
pub fn stats(&self) -> HubStats {
|
||||||
let proxy_conns = self.proxy_conns.read();
|
let proxy_conns = self
|
||||||
let total_proxy = proxy_conns.values().map(|v| v.len()).sum();
|
.proxy_conns_by_id
|
||||||
let nodes = proxy_conns.len();
|
.iter()
|
||||||
drop(proxy_conns);
|
.map(|entry| entry.value().snapshot())
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
let total_proxy = proxy_conns.len();
|
||||||
|
let nodes = self.proxy_conns.read().len();
|
||||||
|
let available_proxy_connections = proxy_conns
|
||||||
|
.iter()
|
||||||
|
.filter(|snapshot| snapshot.available)
|
||||||
|
.count();
|
||||||
|
let closing_proxy_connections = proxy_conns
|
||||||
|
.iter()
|
||||||
|
.filter(|snapshot| snapshot.closing)
|
||||||
|
.count();
|
||||||
|
let draining_proxy_connections = proxy_conns
|
||||||
|
.iter()
|
||||||
|
.filter(|snapshot| snapshot.draining)
|
||||||
|
.count();
|
||||||
|
let soft_avoid_proxy_connections = proxy_conns
|
||||||
|
.iter()
|
||||||
|
.filter(|snapshot| snapshot.available && snapshot.soft_avoid)
|
||||||
|
.count();
|
||||||
|
let outbound_queue_depth_total = proxy_conns
|
||||||
|
.iter()
|
||||||
|
.map(|snapshot| snapshot.outbound.depth)
|
||||||
|
.sum();
|
||||||
|
let outbound_queue_depth_max = proxy_conns
|
||||||
|
.iter()
|
||||||
|
.map(|snapshot| snapshot.outbound.depth)
|
||||||
|
.max()
|
||||||
|
.unwrap_or(0);
|
||||||
|
let outbound_queue_capacity_total = proxy_conns
|
||||||
|
.iter()
|
||||||
|
.map(|snapshot| snapshot.outbound.capacity)
|
||||||
|
.sum();
|
||||||
|
let outbound_queue_rejected_full_total = proxy_conns
|
||||||
|
.iter()
|
||||||
|
.map(|snapshot| snapshot.outbound.rejected_full_total)
|
||||||
|
.sum();
|
||||||
|
let outbound_queue_rejected_closed_total = proxy_conns
|
||||||
|
.iter()
|
||||||
|
.map(|snapshot| snapshot.outbound.rejected_closed_total)
|
||||||
|
.sum();
|
||||||
|
let proxy_connection_congested_total = proxy_conns
|
||||||
|
.iter()
|
||||||
|
.map(|snapshot| snapshot.congested_total)
|
||||||
|
.sum();
|
||||||
|
|
||||||
HubStats {
|
HubStats {
|
||||||
proxy_connections: total_proxy,
|
proxy_connections: total_proxy,
|
||||||
|
available_proxy_connections,
|
||||||
|
closing_proxy_connections,
|
||||||
|
draining_proxy_connections,
|
||||||
|
soft_avoid_proxy_connections,
|
||||||
nodes,
|
nodes,
|
||||||
active_streams: self.local_streams.len(),
|
active_streams: self.local_streams.len(),
|
||||||
|
outbound_queue_depth_total,
|
||||||
|
outbound_queue_depth_max,
|
||||||
|
outbound_queue_capacity_total,
|
||||||
|
outbound_queue_rejected_full_total,
|
||||||
|
outbound_queue_rejected_closed_total,
|
||||||
|
proxy_connection_congested_total,
|
||||||
|
soft_avoid_selection_total: self.soft_avoid_selection_total.load(Ordering::Relaxed),
|
||||||
|
selection_retry_total: self.selection_retry_total.load(Ordering::Relaxed),
|
||||||
|
selection_unavailable_total: self.selection_unavailable_total.load(Ordering::Relaxed),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -893,11 +1081,31 @@ fn current_unix_secs() -> u64 {
|
|||||||
.as_secs()
|
.as_secs()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn percent_u64(value: usize, total: usize) -> u64 {
|
||||||
|
if total == 0 {
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
((value as u128) * 100 / (total as u128)) as u64
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(serde::Serialize)]
|
#[derive(serde::Serialize)]
|
||||||
pub struct HubStats {
|
pub struct HubStats {
|
||||||
pub proxy_connections: usize,
|
pub proxy_connections: usize,
|
||||||
|
pub available_proxy_connections: usize,
|
||||||
|
pub closing_proxy_connections: usize,
|
||||||
|
pub draining_proxy_connections: usize,
|
||||||
|
pub soft_avoid_proxy_connections: usize,
|
||||||
pub nodes: usize,
|
pub nodes: usize,
|
||||||
pub active_streams: usize,
|
pub active_streams: usize,
|
||||||
|
pub outbound_queue_depth_total: usize,
|
||||||
|
pub outbound_queue_depth_max: usize,
|
||||||
|
pub outbound_queue_capacity_total: usize,
|
||||||
|
pub outbound_queue_rejected_full_total: u64,
|
||||||
|
pub outbound_queue_rejected_closed_total: u64,
|
||||||
|
pub proxy_connection_congested_total: u64,
|
||||||
|
pub soft_avoid_selection_total: u64,
|
||||||
|
pub selection_retry_total: u64,
|
||||||
|
pub selection_unavailable_total: u64,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl HubStats {
|
impl HubStats {
|
||||||
@@ -909,6 +1117,30 @@ impl HubStats {
|
|||||||
MetricKind::Gauge,
|
MetricKind::Gauge,
|
||||||
self.proxy_connections as u64,
|
self.proxy_connections as u64,
|
||||||
),
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_connections_available",
|
||||||
|
"Current number of proxy connections available for new work.",
|
||||||
|
MetricKind::Gauge,
|
||||||
|
self.available_proxy_connections as u64,
|
||||||
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_connections_closing",
|
||||||
|
"Current number of proxy connections marked closing.",
|
||||||
|
MetricKind::Gauge,
|
||||||
|
self.closing_proxy_connections as u64,
|
||||||
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_connections_draining",
|
||||||
|
"Current number of proxy connections marked draining.",
|
||||||
|
MetricKind::Gauge,
|
||||||
|
self.draining_proxy_connections as u64,
|
||||||
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_connections_soft_avoid",
|
||||||
|
"Current number of available proxy connections currently soft-avoided by the scheduler.",
|
||||||
|
MetricKind::Gauge,
|
||||||
|
self.soft_avoid_proxy_connections as u64,
|
||||||
|
),
|
||||||
MetricSample::new(
|
MetricSample::new(
|
||||||
"tunnel_nodes",
|
"tunnel_nodes",
|
||||||
"Current number of connected logical nodes.",
|
"Current number of connected logical nodes.",
|
||||||
@@ -921,6 +1153,60 @@ impl HubStats {
|
|||||||
MetricKind::Gauge,
|
MetricKind::Gauge,
|
||||||
self.active_streams as u64,
|
self.active_streams as u64,
|
||||||
),
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_outbound_queue_depth_total",
|
||||||
|
"Current aggregate depth across proxy outbound queues.",
|
||||||
|
MetricKind::Gauge,
|
||||||
|
self.outbound_queue_depth_total as u64,
|
||||||
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_outbound_queue_depth_max",
|
||||||
|
"Current maximum depth observed on a single proxy outbound queue.",
|
||||||
|
MetricKind::Gauge,
|
||||||
|
self.outbound_queue_depth_max as u64,
|
||||||
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_outbound_queue_capacity_total",
|
||||||
|
"Current aggregate capacity across proxy outbound queues.",
|
||||||
|
MetricKind::Gauge,
|
||||||
|
self.outbound_queue_capacity_total as u64,
|
||||||
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_outbound_queue_rejected_full_total",
|
||||||
|
"Total proxy outbound queue sends rejected because a queue was full.",
|
||||||
|
MetricKind::Counter,
|
||||||
|
self.outbound_queue_rejected_full_total,
|
||||||
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_outbound_queue_rejected_closed_total",
|
||||||
|
"Total proxy outbound queue sends rejected because a queue was closed.",
|
||||||
|
MetricKind::Counter,
|
||||||
|
self.outbound_queue_rejected_closed_total,
|
||||||
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_connection_congested_total",
|
||||||
|
"Total number of times a proxy outbound queue became congested.",
|
||||||
|
MetricKind::Counter,
|
||||||
|
self.proxy_connection_congested_total,
|
||||||
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_soft_avoid_selection_total",
|
||||||
|
"Total number of times the scheduler had to pick a high-pressure proxy connection.",
|
||||||
|
MetricKind::Counter,
|
||||||
|
self.soft_avoid_selection_total,
|
||||||
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_selection_retry_total",
|
||||||
|
"Total number of times the scheduler retried a lower-ranked proxy connection after a race on stream allocation.",
|
||||||
|
MetricKind::Counter,
|
||||||
|
self.selection_retry_total,
|
||||||
|
),
|
||||||
|
MetricSample::new(
|
||||||
|
"tunnel_proxy_selection_unavailable_total",
|
||||||
|
"Total number of relay selections that failed because no proxy connection was available.",
|
||||||
|
MetricKind::Counter,
|
||||||
|
self.selection_unavailable_total,
|
||||||
|
),
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -6,7 +6,8 @@ pub(crate) mod write;
|
|||||||
pub(crate) use aether_usage_runtime::UsageRuntime;
|
pub(crate) use aether_usage_runtime::UsageRuntime;
|
||||||
pub use aether_usage_runtime::UsageRuntimeConfig;
|
pub use aether_usage_runtime::UsageRuntimeConfig;
|
||||||
pub(crate) use aether_usage_runtime::{
|
pub(crate) use aether_usage_runtime::{
|
||||||
now_ms, UsageEvent, UsageEventData, UsageEventType, UsageQueue, USAGE_EVENT_VERSION,
|
now_ms, UsageEvent, UsageEventData, UsageEventType, UsageQueue, UsageRequestRecordLevel,
|
||||||
|
USAGE_EVENT_VERSION,
|
||||||
};
|
};
|
||||||
pub(crate) use reporting::{
|
pub(crate) use reporting::{
|
||||||
spawn_sync_report, submit_stream_report, submit_sync_report, GatewayStreamReportRequest,
|
spawn_sync_report, submit_stream_report, submit_sync_report, GatewayStreamReportRequest,
|
||||||
|
|||||||
@@ -404,10 +404,12 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn try_send_stream_error_emits_stream_error_frame() {
|
async fn try_send_stream_error_emits_stream_error_frame() {
|
||||||
let (frame_tx, mut frame_rx) = bounded_queue::<Frame>(4);
|
let (high_tx, mut high_rx) = bounded_queue::<Frame>(4);
|
||||||
|
let (normal_tx, _normal_rx) = bounded_queue::<Frame>(4);
|
||||||
|
let frame_tx = FrameSender::from_test_queues(high_tx, normal_tx);
|
||||||
try_send_stream_error(&frame_tx, 9, "proxy request body dispatch stalled");
|
try_send_stream_error(&frame_tx, 9, "proxy request body dispatch stalled");
|
||||||
|
|
||||||
let frame = frame_rx
|
let frame = high_rx
|
||||||
.recv()
|
.recv()
|
||||||
.await
|
.await
|
||||||
.expect("stream error frame should enqueue");
|
.expect("stream error frame should enqueue");
|
||||||
|
|||||||
@@ -1164,15 +1164,20 @@ fn build_prefixed_request_body(
|
|||||||
mod tests {
|
mod tests {
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::net::SocketAddr;
|
use std::net::SocketAddr;
|
||||||
|
use std::pin::Pin;
|
||||||
use std::sync::atomic::AtomicU64;
|
use std::sync::atomic::AtomicU64;
|
||||||
use std::sync::Once;
|
use std::sync::{Mutex, Once};
|
||||||
|
use std::task::{Context, Poll};
|
||||||
|
|
||||||
use aether_runtime::{bounded_queue, ConcurrencyGate, DistributedConcurrencyGate};
|
use aether_runtime::{ConcurrencyGate, DistributedConcurrencyGate};
|
||||||
use arc_swap::ArcSwap;
|
use arc_swap::ArcSwap;
|
||||||
use axum::body::Body;
|
use axum::body::Body;
|
||||||
use axum::http::{header, Response, StatusCode};
|
use axum::http::{header, Response, StatusCode};
|
||||||
use axum::routing::{get, post};
|
use axum::routing::{get, post};
|
||||||
use axum::Router;
|
use axum::Router;
|
||||||
|
use futures_util::Sink;
|
||||||
|
use tokio::task::JoinHandle;
|
||||||
|
use tokio_tungstenite::tungstenite::{Error as WebSocketError, Message};
|
||||||
|
|
||||||
use super::*;
|
use super::*;
|
||||||
use crate::config::Config;
|
use crate::config::Config;
|
||||||
@@ -1370,14 +1375,22 @@ mod tests {
|
|||||||
let state = sample_state_for_port(addr.port());
|
let state = sample_state_for_port(addr.port());
|
||||||
cache_test_host(&state, host, addr).await;
|
cache_test_host(&state, host, addr).await;
|
||||||
let server_ctx = sample_server(&state);
|
let server_ctx = sample_server(&state);
|
||||||
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
|
let (frame_tx, sent, writer_handle) = spawn_test_writer();
|
||||||
let (_body_tx, body_rx) = mpsc::channel(1);
|
let (_body_tx, body_rx) = mpsc::channel(1);
|
||||||
|
|
||||||
let mut meta = sample_request_meta();
|
let mut meta = sample_request_meta();
|
||||||
meta.url = format!("http://{host}:{}/start", addr.port());
|
meta.url = format!("http://{host}:{}/start", addr.port());
|
||||||
|
|
||||||
handle_stream(Arc::clone(&state), server_ctx, 5, meta, body_rx, frame_tx).await;
|
handle_stream(
|
||||||
let result = collect_stream_result(&mut frame_rx).await;
|
Arc::clone(&state),
|
||||||
|
server_ctx,
|
||||||
|
5,
|
||||||
|
meta,
|
||||||
|
body_rx,
|
||||||
|
frame_tx.clone(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
|
||||||
server.abort();
|
server.abort();
|
||||||
|
|
||||||
assert!(
|
assert!(
|
||||||
@@ -1423,14 +1436,22 @@ mod tests {
|
|||||||
let state = sample_state_for_port(addr.port());
|
let state = sample_state_for_port(addr.port());
|
||||||
cache_test_host(&state, host, addr).await;
|
cache_test_host(&state, host, addr).await;
|
||||||
let server_ctx = sample_server(&state);
|
let server_ctx = sample_server(&state);
|
||||||
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
|
let (frame_tx, sent, writer_handle) = spawn_test_writer();
|
||||||
let (_body_tx, body_rx) = mpsc::channel(1);
|
let (_body_tx, body_rx) = mpsc::channel(1);
|
||||||
|
|
||||||
let mut meta = sample_request_meta();
|
let mut meta = sample_request_meta();
|
||||||
meta.url = format!("http://{host}:{}/ok", addr.port());
|
meta.url = format!("http://{host}:{}/ok", addr.port());
|
||||||
|
|
||||||
handle_stream(Arc::clone(&state), server_ctx, 3, meta, body_rx, frame_tx).await;
|
handle_stream(
|
||||||
let result = collect_stream_result(&mut frame_rx).await;
|
Arc::clone(&state),
|
||||||
|
server_ctx,
|
||||||
|
3,
|
||||||
|
meta,
|
||||||
|
body_rx,
|
||||||
|
frame_tx.clone(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
|
||||||
server.abort();
|
server.abort();
|
||||||
|
|
||||||
assert!(
|
assert!(
|
||||||
@@ -1486,7 +1507,7 @@ mod tests {
|
|||||||
let state = sample_state_for_port(addr.port());
|
let state = sample_state_for_port(addr.port());
|
||||||
cache_test_host(&state, host, addr).await;
|
cache_test_host(&state, host, addr).await;
|
||||||
let server_ctx = sample_server(&state);
|
let server_ctx = sample_server(&state);
|
||||||
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
|
let (frame_tx, sent, writer_handle) = spawn_test_writer();
|
||||||
let (body_tx, body_rx) = mpsc::channel(4);
|
let (body_tx, body_rx) = mpsc::channel(4);
|
||||||
body_tx
|
body_tx
|
||||||
.send(TunnelFrame::new(
|
.send(TunnelFrame::new(
|
||||||
@@ -1504,8 +1525,16 @@ mod tests {
|
|||||||
meta.url = format!("http://{host}:{}/start", addr.port());
|
meta.url = format!("http://{host}:{}/start", addr.port());
|
||||||
meta.follow_redirects = Some(true);
|
meta.follow_redirects = Some(true);
|
||||||
|
|
||||||
handle_stream(Arc::clone(&state), server_ctx, 1, meta, body_rx, frame_tx).await;
|
handle_stream(
|
||||||
let result = collect_stream_result(&mut frame_rx).await;
|
Arc::clone(&state),
|
||||||
|
server_ctx,
|
||||||
|
1,
|
||||||
|
meta,
|
||||||
|
body_rx,
|
||||||
|
frame_tx.clone(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
|
||||||
server.abort();
|
server.abort();
|
||||||
|
|
||||||
assert!(
|
assert!(
|
||||||
@@ -1553,15 +1582,23 @@ mod tests {
|
|||||||
let state = sample_state_for_port(addr.port());
|
let state = sample_state_for_port(addr.port());
|
||||||
cache_test_host(&state, host, addr).await;
|
cache_test_host(&state, host, addr).await;
|
||||||
let server_ctx = sample_server(&state);
|
let server_ctx = sample_server(&state);
|
||||||
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
|
let (frame_tx, sent, writer_handle) = spawn_test_writer();
|
||||||
let (_body_tx, body_rx) = mpsc::channel(1);
|
let (_body_tx, body_rx) = mpsc::channel(1);
|
||||||
|
|
||||||
let mut meta = sample_request_meta();
|
let mut meta = sample_request_meta();
|
||||||
meta.url = format!("http://{host}:{}/start", addr.port());
|
meta.url = format!("http://{host}:{}/start", addr.port());
|
||||||
meta.follow_redirects = Some(false);
|
meta.follow_redirects = Some(false);
|
||||||
|
|
||||||
handle_stream(Arc::clone(&state), server_ctx, 7, meta, body_rx, frame_tx).await;
|
handle_stream(
|
||||||
let result = collect_stream_result(&mut frame_rx).await;
|
Arc::clone(&state),
|
||||||
|
server_ctx,
|
||||||
|
7,
|
||||||
|
meta,
|
||||||
|
body_rx,
|
||||||
|
frame_tx.clone(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
|
||||||
server.abort();
|
server.abort();
|
||||||
|
|
||||||
assert!(
|
assert!(
|
||||||
@@ -1618,7 +1655,7 @@ mod tests {
|
|||||||
let state = sample_state_for_budget(addr.port(), 0);
|
let state = sample_state_for_budget(addr.port(), 0);
|
||||||
cache_test_host(&state, host, addr).await;
|
cache_test_host(&state, host, addr).await;
|
||||||
let server_ctx = sample_server(&state);
|
let server_ctx = sample_server(&state);
|
||||||
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
|
let (frame_tx, sent, writer_handle) = spawn_test_writer();
|
||||||
let (body_tx, body_rx) = mpsc::channel(4);
|
let (body_tx, body_rx) = mpsc::channel(4);
|
||||||
body_tx
|
body_tx
|
||||||
.send(TunnelFrame::new(
|
.send(TunnelFrame::new(
|
||||||
@@ -1636,8 +1673,16 @@ mod tests {
|
|||||||
meta.url = format!("http://{host}:{}/start", addr.port());
|
meta.url = format!("http://{host}:{}/start", addr.port());
|
||||||
meta.follow_redirects = Some(true);
|
meta.follow_redirects = Some(true);
|
||||||
|
|
||||||
handle_stream(Arc::clone(&state), server_ctx, 11, meta, body_rx, frame_tx).await;
|
handle_stream(
|
||||||
let result = collect_stream_result(&mut frame_rx).await;
|
Arc::clone(&state),
|
||||||
|
server_ctx,
|
||||||
|
11,
|
||||||
|
meta,
|
||||||
|
body_rx,
|
||||||
|
frame_tx.clone(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
|
||||||
server.abort();
|
server.abort();
|
||||||
|
|
||||||
assert!(
|
assert!(
|
||||||
@@ -1663,7 +1708,7 @@ mod tests {
|
|||||||
let _permit = gate.try_acquire().expect("first permit");
|
let _permit = gate.try_acquire().expect("first permit");
|
||||||
let state = sample_state(Some(gate), None);
|
let state = sample_state(Some(gate), None);
|
||||||
let server = sample_server(&state);
|
let server = sample_server(&state);
|
||||||
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(4);
|
let (frame_tx, sent, writer_handle) = spawn_test_writer();
|
||||||
let (_body_tx, body_rx) = mpsc::channel(1);
|
let (_body_tx, body_rx) = mpsc::channel(1);
|
||||||
|
|
||||||
handle_stream(
|
handle_stream(
|
||||||
@@ -1672,11 +1717,15 @@ mod tests {
|
|||||||
7,
|
7,
|
||||||
sample_request_meta(),
|
sample_request_meta(),
|
||||||
body_rx,
|
body_rx,
|
||||||
frame_tx,
|
frame_tx.clone(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
let frame = frame_rx.recv().await.expect("overload frame");
|
let frame = collect_emitted_frames(frame_tx, sent, writer_handle)
|
||||||
|
.await
|
||||||
|
.into_iter()
|
||||||
|
.find(|frame| frame.msg_type == MsgType::StreamError)
|
||||||
|
.expect("overload frame");
|
||||||
assert_eq!(frame.stream_id, 7);
|
assert_eq!(frame.stream_id, 7);
|
||||||
assert_eq!(frame.msg_type, MsgType::StreamError);
|
assert_eq!(frame.msg_type, MsgType::StreamError);
|
||||||
assert_eq!(frame.payload, Bytes::from_static(b"proxy overloaded"));
|
assert_eq!(frame.payload, Bytes::from_static(b"proxy overloaded"));
|
||||||
@@ -1700,7 +1749,7 @@ mod tests {
|
|||||||
let _permit = gate.try_acquire().await.expect("first permit");
|
let _permit = gate.try_acquire().await.expect("first permit");
|
||||||
let state = sample_state(None, Some(gate));
|
let state = sample_state(None, Some(gate));
|
||||||
let server = sample_server(&state);
|
let server = sample_server(&state);
|
||||||
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(4);
|
let (frame_tx, sent, writer_handle) = spawn_test_writer();
|
||||||
let (_body_tx, body_rx) = mpsc::channel(1);
|
let (_body_tx, body_rx) = mpsc::channel(1);
|
||||||
|
|
||||||
handle_stream(
|
handle_stream(
|
||||||
@@ -1709,11 +1758,15 @@ mod tests {
|
|||||||
9,
|
9,
|
||||||
sample_request_meta(),
|
sample_request_meta(),
|
||||||
body_rx,
|
body_rx,
|
||||||
frame_tx,
|
frame_tx.clone(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
let frame = frame_rx.recv().await.expect("overload frame");
|
let frame = collect_emitted_frames(frame_tx, sent, writer_handle)
|
||||||
|
.await
|
||||||
|
.into_iter()
|
||||||
|
.find(|frame| frame.msg_type == MsgType::StreamError)
|
||||||
|
.expect("overload frame");
|
||||||
assert_eq!(frame.stream_id, 9);
|
assert_eq!(frame.stream_id, 9);
|
||||||
assert_eq!(frame.msg_type, MsgType::StreamError);
|
assert_eq!(frame.msg_type, MsgType::StreamError);
|
||||||
assert_eq!(frame.payload, Bytes::from_static(b"proxy overloaded"));
|
assert_eq!(frame.payload, Bytes::from_static(b"proxy overloaded"));
|
||||||
@@ -1881,20 +1934,85 @@ mod tests {
|
|||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Default)]
|
||||||
|
struct VecSink {
|
||||||
|
sent: Arc<Mutex<Vec<Message>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Sink<Message> for VecSink {
|
||||||
|
type Error = WebSocketError;
|
||||||
|
|
||||||
|
fn poll_ready(
|
||||||
|
self: Pin<&mut Self>,
|
||||||
|
_cx: &mut Context<'_>,
|
||||||
|
) -> Poll<Result<(), Self::Error>> {
|
||||||
|
Poll::Ready(Ok(()))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
|
||||||
|
self.sent.lock().expect("sink lock").push(item);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn poll_flush(
|
||||||
|
self: Pin<&mut Self>,
|
||||||
|
_cx: &mut Context<'_>,
|
||||||
|
) -> Poll<Result<(), Self::Error>> {
|
||||||
|
Poll::Ready(Ok(()))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn poll_close(
|
||||||
|
self: Pin<&mut Self>,
|
||||||
|
_cx: &mut Context<'_>,
|
||||||
|
) -> Poll<Result<(), Self::Error>> {
|
||||||
|
Poll::Ready(Ok(()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn spawn_test_writer() -> (FrameSender, Arc<Mutex<Vec<Message>>>, JoinHandle<()>) {
|
||||||
|
let sink = VecSink::default();
|
||||||
|
let sent = Arc::clone(&sink.sent);
|
||||||
|
let (frame_tx, handle) = crate::tunnel::writer::spawn_writer(sink, Duration::from_secs(60));
|
||||||
|
(frame_tx, sent, handle)
|
||||||
|
}
|
||||||
|
|
||||||
struct StreamResult {
|
struct StreamResult {
|
||||||
response: Option<ResponseMeta>,
|
response: Option<ResponseMeta>,
|
||||||
body: Bytes,
|
body: Bytes,
|
||||||
error: Option<String>,
|
error: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn collect_emitted_frames(
|
||||||
|
frame_tx: FrameSender,
|
||||||
|
sent: Arc<Mutex<Vec<Message>>>,
|
||||||
|
writer_handle: JoinHandle<()>,
|
||||||
|
) -> Vec<TunnelFrame> {
|
||||||
|
drop(frame_tx);
|
||||||
|
writer_handle.await.expect("writer should exit cleanly");
|
||||||
|
|
||||||
|
sent.lock()
|
||||||
|
.expect("sink lock")
|
||||||
|
.iter()
|
||||||
|
.filter_map(|message| match message {
|
||||||
|
Message::Binary(data) => {
|
||||||
|
Some(TunnelFrame::decode(data.clone().into()).expect("frame should decode"))
|
||||||
|
}
|
||||||
|
Message::Ping(_) | Message::Pong(_) | Message::Close(_) => None,
|
||||||
|
other => panic!("unexpected writer message: {other:?}"),
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
async fn collect_stream_result(
|
async fn collect_stream_result(
|
||||||
frame_rx: &mut aether_runtime::BoundedQueueReceiver<TunnelFrame>,
|
frame_tx: FrameSender,
|
||||||
|
sent: Arc<Mutex<Vec<Message>>>,
|
||||||
|
writer_handle: JoinHandle<()>,
|
||||||
) -> StreamResult {
|
) -> StreamResult {
|
||||||
let mut response = None;
|
let mut response = None;
|
||||||
let mut body = BytesMut::new();
|
let mut body = BytesMut::new();
|
||||||
let mut error = None;
|
let mut error = None;
|
||||||
|
|
||||||
while let Some(frame) = frame_rx.recv().await {
|
for frame in collect_emitted_frames(frame_tx, sent, writer_handle).await {
|
||||||
match frame.msg_type {
|
match frame.msg_type {
|
||||||
MsgType::ResponseHeaders => {
|
MsgType::ResponseHeaders => {
|
||||||
let payload = decompress_if_gzip(&frame).expect("headers payload");
|
let payload = decompress_if_gzip(&frame).expect("headers payload");
|
||||||
|
|||||||
@@ -7,7 +7,10 @@
|
|||||||
|
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use aether_runtime::{bounded_queue, BoundedQueueSender};
|
use aether_contracts::tunnel::MsgType;
|
||||||
|
#[cfg(test)]
|
||||||
|
use aether_runtime::QueueSnapshot;
|
||||||
|
use aether_runtime::{bounded_queue, BoundedQueueSender, QueueSendError};
|
||||||
use futures_util::SinkExt;
|
use futures_util::SinkExt;
|
||||||
use tokio::task::JoinHandle;
|
use tokio::task::JoinHandle;
|
||||||
use tokio_tungstenite::tungstenite::Message;
|
use tokio_tungstenite::tungstenite::Message;
|
||||||
@@ -15,8 +18,60 @@ use tracing::{debug, error, trace};
|
|||||||
|
|
||||||
use super::protocol::Frame;
|
use super::protocol::Frame;
|
||||||
|
|
||||||
|
const HIGH_PRIORITY_QUEUE_CAPACITY: usize = 64;
|
||||||
|
const NORMAL_PRIORITY_QUEUE_CAPACITY: usize = 256;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
enum FramePriority {
|
||||||
|
High,
|
||||||
|
Normal,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub struct FrameQueueSnapshots {
|
||||||
|
pub high: QueueSnapshot,
|
||||||
|
pub normal: QueueSnapshot,
|
||||||
|
}
|
||||||
|
|
||||||
/// Sender half — cloned by stream handlers and heartbeat.
|
/// Sender half — cloned by stream handlers and heartbeat.
|
||||||
pub type FrameSender = BoundedQueueSender<Frame>;
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct FrameSender {
|
||||||
|
high_tx: BoundedQueueSender<Frame>,
|
||||||
|
normal_tx: BoundedQueueSender<Frame>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl FrameSender {
|
||||||
|
pub async fn send(&self, frame: Frame) -> Result<(), QueueSendError<Frame>> {
|
||||||
|
match classify_frame_priority(&frame) {
|
||||||
|
FramePriority::High => self.high_tx.send(frame).await,
|
||||||
|
FramePriority::Normal => self.normal_tx.send(frame).await,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn try_send(&self, frame: Frame) -> Result<(), QueueSendError<Frame>> {
|
||||||
|
match classify_frame_priority(&frame) {
|
||||||
|
FramePriority::High => self.high_tx.try_send(frame),
|
||||||
|
FramePriority::Normal => self.normal_tx.try_send(frame),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
pub fn snapshots(&self) -> FrameQueueSnapshots {
|
||||||
|
FrameQueueSnapshots {
|
||||||
|
high: self.high_tx.snapshot(),
|
||||||
|
normal: self.normal_tx.snapshot(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
pub(crate) fn from_test_queues(
|
||||||
|
high_tx: BoundedQueueSender<Frame>,
|
||||||
|
normal_tx: BoundedQueueSender<Frame>,
|
||||||
|
) -> Self {
|
||||||
|
Self { high_tx, normal_tx }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Spawn the writer task. Returns the sender and a JoinHandle for cleanup.
|
/// Spawn the writer task. Returns the sender and a JoinHandle for cleanup.
|
||||||
///
|
///
|
||||||
@@ -26,33 +81,56 @@ pub fn spawn_writer<S>(mut sink: S, ping_interval: Duration) -> (FrameSender, Jo
|
|||||||
where
|
where
|
||||||
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
|
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
|
||||||
{
|
{
|
||||||
let (tx, mut rx) = bounded_queue::<Frame>(256);
|
let (high_tx, mut high_rx) = bounded_queue::<Frame>(HIGH_PRIORITY_QUEUE_CAPACITY);
|
||||||
|
let (normal_tx, mut normal_rx) = bounded_queue::<Frame>(NORMAL_PRIORITY_QUEUE_CAPACITY);
|
||||||
|
let tx = FrameSender { high_tx, normal_tx };
|
||||||
|
|
||||||
let handle = tokio::spawn(async move {
|
let handle = tokio::spawn(async move {
|
||||||
let mut ping_ticker = tokio::time::interval(ping_interval);
|
let mut ping_ticker = tokio::time::interval(ping_interval);
|
||||||
|
let mut high_open = true;
|
||||||
|
let mut normal_open = true;
|
||||||
ping_ticker.tick().await; // skip first immediate tick
|
ping_ticker.tick().await; // skip first immediate tick
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
|
if let Ok(frame) = high_rx.try_recv() {
|
||||||
|
if !write_frame(&mut sink, frame).await {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if !high_open && !normal_open {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
frame = rx.recv() => {
|
biased;
|
||||||
|
frame = high_rx.recv(), if high_open => {
|
||||||
match frame {
|
match frame {
|
||||||
Some(frame) => {
|
Some(frame) => {
|
||||||
let data = frame.encode();
|
if !write_frame(&mut sink, frame).await {
|
||||||
if let Err(e) = sink.send(Message::Binary(data.into())).await {
|
|
||||||
error!(error = %e, "failed to write frame to WebSocket");
|
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
None => break, // all senders dropped
|
None => high_open = false,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
_ = ping_ticker.tick() => {
|
_ = ping_ticker.tick(), if high_open || normal_open => {
|
||||||
if let Err(e) = sink.send(Message::Ping(vec![])).await {
|
if let Err(e) = sink.send(Message::Ping(vec![])).await {
|
||||||
error!(error = %e, "failed to send WebSocket ping");
|
error!(error = %e, "failed to send WebSocket ping");
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
trace!("sent WebSocket ping");
|
trace!("sent WebSocket ping");
|
||||||
}
|
}
|
||||||
|
frame = normal_rx.recv(), if normal_open => {
|
||||||
|
match frame {
|
||||||
|
Some(frame) => {
|
||||||
|
if !write_frame(&mut sink, frame).await {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => normal_open = false,
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
debug!("writer task exiting");
|
debug!("writer task exiting");
|
||||||
@@ -61,3 +139,134 @@ where
|
|||||||
|
|
||||||
(tx, handle)
|
(tx, handle)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn classify_frame_priority(frame: &Frame) -> FramePriority {
|
||||||
|
match frame.msg_type {
|
||||||
|
MsgType::ResponseHeaders
|
||||||
|
| MsgType::StreamError
|
||||||
|
| MsgType::Ping
|
||||||
|
| MsgType::Pong
|
||||||
|
| MsgType::GoAway
|
||||||
|
| MsgType::HeartbeatData
|
||||||
|
| MsgType::HeartbeatAck => FramePriority::High,
|
||||||
|
MsgType::RequestHeaders
|
||||||
|
| MsgType::RequestBody
|
||||||
|
| MsgType::ResponseBody
|
||||||
|
| MsgType::StreamEnd => FramePriority::Normal,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn write_frame<S>(sink: &mut S, frame: Frame) -> bool
|
||||||
|
where
|
||||||
|
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
|
||||||
|
{
|
||||||
|
let data = frame.encode();
|
||||||
|
if let Err(e) = sink.send(Message::Binary(data.into())).await {
|
||||||
|
error!(error = %e, "failed to write frame to WebSocket");
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use std::pin::Pin;
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
use std::task::{Context, Poll};
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use futures_util::Sink;
|
||||||
|
use tokio_tungstenite::tungstenite::{Error, Message};
|
||||||
|
|
||||||
|
use super::spawn_writer;
|
||||||
|
use crate::tunnel::protocol::Frame;
|
||||||
|
use aether_contracts::tunnel::MsgType;
|
||||||
|
|
||||||
|
#[derive(Clone, Default)]
|
||||||
|
struct VecSink {
|
||||||
|
sent: Arc<Mutex<Vec<Message>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Sink<Message> for VecSink {
|
||||||
|
type Error = Error;
|
||||||
|
|
||||||
|
fn poll_ready(
|
||||||
|
self: Pin<&mut Self>,
|
||||||
|
_cx: &mut Context<'_>,
|
||||||
|
) -> Poll<Result<(), Self::Error>> {
|
||||||
|
Poll::Ready(Ok(()))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
|
||||||
|
self.sent.lock().expect("sink lock").push(item);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn poll_flush(
|
||||||
|
self: Pin<&mut Self>,
|
||||||
|
_cx: &mut Context<'_>,
|
||||||
|
) -> Poll<Result<(), Self::Error>> {
|
||||||
|
Poll::Ready(Ok(()))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn poll_close(
|
||||||
|
self: Pin<&mut Self>,
|
||||||
|
_cx: &mut Context<'_>,
|
||||||
|
) -> Poll<Result<(), Self::Error>> {
|
||||||
|
Poll::Ready(Ok(()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn prioritizes_control_frames_ahead_of_buffered_body_frames() {
|
||||||
|
let sink = VecSink::default();
|
||||||
|
let sent = Arc::clone(&sink.sent);
|
||||||
|
let (sender, handle) = spawn_writer(sink, Duration::from_secs(60));
|
||||||
|
|
||||||
|
for idx in 0..8u8 {
|
||||||
|
sender
|
||||||
|
.try_send(Frame::new(
|
||||||
|
7,
|
||||||
|
MsgType::ResponseBody,
|
||||||
|
0,
|
||||||
|
bytes::Bytes::from(vec![idx; 32]),
|
||||||
|
))
|
||||||
|
.expect("frame send should succeed");
|
||||||
|
}
|
||||||
|
sender
|
||||||
|
.try_send(Frame::new(
|
||||||
|
7,
|
||||||
|
MsgType::StreamError,
|
||||||
|
0,
|
||||||
|
bytes::Bytes::from_static(b"boom"),
|
||||||
|
))
|
||||||
|
.expect("frame send should succeed");
|
||||||
|
let snapshots = sender.snapshots();
|
||||||
|
assert!(snapshots.high.enqueued_total >= 1);
|
||||||
|
assert!(snapshots.normal.enqueued_total >= 8);
|
||||||
|
|
||||||
|
tokio::time::sleep(Duration::from_millis(30)).await;
|
||||||
|
drop(sender);
|
||||||
|
handle.await.expect("writer should exit cleanly");
|
||||||
|
|
||||||
|
let sent = sent.lock().expect("sink lock");
|
||||||
|
assert!(
|
||||||
|
sent.len() >= 2,
|
||||||
|
"writer should flush both body and control frames"
|
||||||
|
);
|
||||||
|
let first = match &sent[0] {
|
||||||
|
Message::Binary(data) => {
|
||||||
|
Frame::decode(data.clone().into()).expect("frame should decode")
|
||||||
|
}
|
||||||
|
other => panic!("unexpected first message: {other:?}"),
|
||||||
|
};
|
||||||
|
let second = match &sent[1] {
|
||||||
|
Message::Binary(data) => {
|
||||||
|
Frame::decode(data.clone().into()).expect("frame should decode")
|
||||||
|
}
|
||||||
|
other => panic!("unexpected second message: {other:?}"),
|
||||||
|
};
|
||||||
|
assert_eq!(first.msg_type, MsgType::StreamError);
|
||||||
|
assert_eq!(second.msg_type, MsgType::ResponseBody);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
use aether_data_contracts::repository::{
|
use aether_data_contracts::repository::{
|
||||||
candidates::{DecisionTrace, DecisionTraceCandidate},
|
candidates::{DecisionTrace, DecisionTraceCandidate, RequestCandidateStatus},
|
||||||
provider_catalog::StoredProviderCatalogKey,
|
provider_catalog::StoredProviderCatalogKey,
|
||||||
|
usage::StoredRequestUsageAudit,
|
||||||
};
|
};
|
||||||
use axum::{
|
use axum::{
|
||||||
body::Body,
|
body::Body,
|
||||||
@@ -261,11 +262,20 @@ pub fn build_admin_monitoring_trace_provider_stats_payload_response(
|
|||||||
|
|
||||||
pub fn build_admin_monitoring_trace_request_payload_response(
|
pub fn build_admin_monitoring_trace_request_payload_response(
|
||||||
trace: &DecisionTrace,
|
trace: &DecisionTrace,
|
||||||
|
usage: Option<&StoredRequestUsageAudit>,
|
||||||
) -> Response<Body> {
|
) -> Response<Body> {
|
||||||
|
let usage_candidate_id =
|
||||||
|
usage.and_then(|item| resolve_admin_monitoring_usage_candidate_id(trace, item));
|
||||||
let candidates = trace
|
let candidates = trace
|
||||||
.candidates
|
.candidates
|
||||||
.iter()
|
.iter()
|
||||||
.map(build_admin_monitoring_trace_request_candidate_payload)
|
.map(|item| {
|
||||||
|
let matched_usage = usage_candidate_id
|
||||||
|
.as_deref()
|
||||||
|
.filter(|candidate_id| *candidate_id == item.candidate.id.as_str())
|
||||||
|
.and(usage);
|
||||||
|
build_admin_monitoring_trace_request_candidate_payload(item, matched_usage)
|
||||||
|
})
|
||||||
.collect::<Vec<_>>();
|
.collect::<Vec<_>>();
|
||||||
Json(json!({
|
Json(json!({
|
||||||
"request_id": trace.request_id,
|
"request_id": trace.request_id,
|
||||||
@@ -279,6 +289,7 @@ pub fn build_admin_monitoring_trace_request_payload_response(
|
|||||||
|
|
||||||
pub fn build_admin_monitoring_trace_request_candidate_payload(
|
pub fn build_admin_monitoring_trace_request_candidate_payload(
|
||||||
item: &DecisionTraceCandidate,
|
item: &DecisionTraceCandidate,
|
||||||
|
usage: Option<&StoredRequestUsageAudit>,
|
||||||
) -> Value {
|
) -> Value {
|
||||||
let candidate = &item.candidate;
|
let candidate = &item.candidate;
|
||||||
json!({
|
json!({
|
||||||
@@ -316,13 +327,133 @@ pub fn build_admin_monitoring_trace_request_candidate_payload(
|
|||||||
"error_message": candidate.error_message,
|
"error_message": candidate.error_message,
|
||||||
"latency_ms": candidate.latency_ms,
|
"latency_ms": candidate.latency_ms,
|
||||||
"concurrent_requests": candidate.concurrent_requests,
|
"concurrent_requests": candidate.concurrent_requests,
|
||||||
"extra_data": candidate.extra_data,
|
"extra_data": build_admin_monitoring_trace_candidate_extra_data(candidate.extra_data.as_ref(), usage),
|
||||||
"created_at": unix_ms_to_rfc3339(candidate.created_at_unix_ms),
|
"created_at": unix_ms_to_rfc3339(candidate.created_at_unix_ms),
|
||||||
"started_at": candidate.started_at_unix_ms.and_then(unix_ms_to_rfc3339),
|
"started_at": candidate.started_at_unix_ms.and_then(unix_ms_to_rfc3339),
|
||||||
"finished_at": candidate.finished_at_unix_ms.and_then(unix_ms_to_rfc3339),
|
"finished_at": candidate.finished_at_unix_ms.and_then(unix_ms_to_rfc3339),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn resolve_admin_monitoring_usage_candidate_id(
|
||||||
|
trace: &DecisionTrace,
|
||||||
|
usage: &StoredRequestUsageAudit,
|
||||||
|
) -> Option<String> {
|
||||||
|
if usage.request_id.trim() != trace.request_id {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(candidate_id) = usage
|
||||||
|
.routing_candidate_id()
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
{
|
||||||
|
return Some(candidate_id.to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
let candidate_index = usage.routing_candidate_index()?;
|
||||||
|
trace
|
||||||
|
.candidates
|
||||||
|
.iter()
|
||||||
|
.filter(|item| u64::from(item.candidate.candidate_index) == candidate_index)
|
||||||
|
.max_by_key(|item| {
|
||||||
|
(
|
||||||
|
item.candidate.retry_index,
|
||||||
|
item.candidate.finished_at_unix_ms.unwrap_or_default(),
|
||||||
|
item.candidate.started_at_unix_ms.unwrap_or_default(),
|
||||||
|
admin_monitoring_candidate_status_rank(item.candidate.status),
|
||||||
|
)
|
||||||
|
})
|
||||||
|
.map(|item| item.candidate.id.clone())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn admin_monitoring_candidate_status_rank(status: RequestCandidateStatus) -> u8 {
|
||||||
|
match status {
|
||||||
|
RequestCandidateStatus::Success => 7,
|
||||||
|
RequestCandidateStatus::Streaming => 6,
|
||||||
|
RequestCandidateStatus::Pending => 5,
|
||||||
|
RequestCandidateStatus::Failed => 4,
|
||||||
|
RequestCandidateStatus::Cancelled => 3,
|
||||||
|
RequestCandidateStatus::Skipped => 2,
|
||||||
|
RequestCandidateStatus::Unused => 1,
|
||||||
|
RequestCandidateStatus::Available => 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_admin_monitoring_trace_candidate_extra_data(
|
||||||
|
existing: Option<&Value>,
|
||||||
|
usage: Option<&StoredRequestUsageAudit>,
|
||||||
|
) -> Value {
|
||||||
|
let mut extra_data = match existing {
|
||||||
|
Some(Value::Object(object)) => Some(object.clone()),
|
||||||
|
Some(other) => return other.clone(),
|
||||||
|
None => None,
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Some(usage) = usage {
|
||||||
|
let extra_object = extra_data.get_or_insert_with(serde_json::Map::new);
|
||||||
|
if let Some(first_byte_time_ms) = usage.first_byte_time_ms {
|
||||||
|
extra_object
|
||||||
|
.entry("first_byte_time_ms".to_string())
|
||||||
|
.or_insert_with(|| json!(first_byte_time_ms));
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(proxy_value) = extra_object.get_mut("proxy") {
|
||||||
|
if let Some(proxy_object) = proxy_value.as_object_mut() {
|
||||||
|
let proxy_timing = parse_admin_monitoring_usage_proxy_timing(usage);
|
||||||
|
let proxy_ttfb_ms = proxy_timing
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|timing| timing.get("ttfb_ms"))
|
||||||
|
.and_then(Value::as_u64)
|
||||||
|
.or(usage.first_byte_time_ms);
|
||||||
|
if let Some(ttfb_ms) = proxy_ttfb_ms {
|
||||||
|
proxy_object
|
||||||
|
.entry("ttfb_ms".to_string())
|
||||||
|
.or_insert_with(|| json!(ttfb_ms));
|
||||||
|
}
|
||||||
|
if let Some(timing) = proxy_timing {
|
||||||
|
proxy_object.entry("timing".to_string()).or_insert(timing);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
match extra_data {
|
||||||
|
Some(object) => Value::Object(object),
|
||||||
|
None => Value::Null,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_admin_monitoring_usage_proxy_timing(usage: &StoredRequestUsageAudit) -> Option<Value> {
|
||||||
|
admin_monitoring_header_value(usage.response_headers.as_ref(), "x-proxy-timing")
|
||||||
|
.or_else(|| {
|
||||||
|
admin_monitoring_header_value(usage.client_response_headers.as_ref(), "x-proxy-timing")
|
||||||
|
})
|
||||||
|
.and_then(|raw| parse_admin_monitoring_proxy_timing_value(&raw))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn admin_monitoring_header_value(headers: Option<&Value>, name: &str) -> Option<String> {
|
||||||
|
headers
|
||||||
|
.and_then(Value::as_object)
|
||||||
|
.and_then(|object| {
|
||||||
|
object
|
||||||
|
.iter()
|
||||||
|
.find(|(key, _)| key.eq_ignore_ascii_case(name))
|
||||||
|
.map(|(_, value)| value)
|
||||||
|
})
|
||||||
|
.and_then(|value| match value {
|
||||||
|
Value::String(text) => Some(text.trim().to_string()),
|
||||||
|
Value::Object(object) => Some(Value::Object(object.clone()).to_string()),
|
||||||
|
_ => None,
|
||||||
|
})
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_admin_monitoring_proxy_timing_value(raw: &str) -> Option<Value> {
|
||||||
|
serde_json::from_str::<Value>(raw)
|
||||||
|
.ok()
|
||||||
|
.filter(Value::is_object)
|
||||||
|
}
|
||||||
|
|
||||||
#[allow(clippy::too_many_arguments)]
|
#[allow(clippy::too_many_arguments)]
|
||||||
pub fn build_admin_monitoring_system_status_payload_response(
|
pub fn build_admin_monitoring_system_status_payload_response(
|
||||||
timestamp: chrono::DateTime<chrono::Utc>,
|
timestamp: chrono::DateTime<chrono::Utc>,
|
||||||
|
|||||||
@@ -1201,7 +1201,7 @@ pub fn admin_system_config_default_value(key: &str) -> Option<serde_json::Value>
|
|||||||
"site_subtitle" => Some(json!("AI Gateway")),
|
"site_subtitle" => Some(json!("AI Gateway")),
|
||||||
"default_user_initial_gift_usd" => Some(json!(10.0)),
|
"default_user_initial_gift_usd" => Some(json!(10.0)),
|
||||||
"password_policy_level" => Some(json!("weak")),
|
"password_policy_level" => Some(json!("weak")),
|
||||||
REQUEST_RECORD_LEVEL_KEY => Some(json!("basic")),
|
REQUEST_RECORD_LEVEL_KEY => Some(json!("full")),
|
||||||
"max_request_body_size" => Some(json!(5_242_880)),
|
"max_request_body_size" => Some(json!(5_242_880)),
|
||||||
"max_response_body_size" => Some(json!(5_242_880)),
|
"max_response_body_size" => Some(json!(5_242_880)),
|
||||||
"sensitive_headers" => Some(json!([
|
"sensitive_headers" => Some(json!([
|
||||||
|
|||||||
@@ -30,7 +30,9 @@ pub use report_context::{
|
|||||||
build_locally_actionable_report_context_from_request_candidate,
|
build_locally_actionable_report_context_from_request_candidate,
|
||||||
build_locally_actionable_report_context_from_video_task, report_context_is_locally_actionable,
|
build_locally_actionable_report_context_from_video_task, report_context_is_locally_actionable,
|
||||||
};
|
};
|
||||||
pub use runtime::{UsageBillingEventEnricher, UsageRuntime, UsageRuntimeAccess};
|
pub use runtime::{
|
||||||
|
UsageBillingEventEnricher, UsageRequestRecordLevel, UsageRuntime, UsageRuntimeAccess,
|
||||||
|
};
|
||||||
pub use settlement::{settle_usage_if_needed, UsageSettlementWriter};
|
pub use settlement::{settle_usage_if_needed, UsageSettlementWriter};
|
||||||
pub use standardized_usage::StandardizedUsage;
|
pub use standardized_usage::StandardizedUsage;
|
||||||
pub use usage_mapper::{map_usage, map_usage_from_response, UsageMapper};
|
pub use usage_mapper::{map_usage, map_usage_from_response, UsageMapper};
|
||||||
|
|||||||
@@ -24,12 +24,24 @@ pub trait UsageBillingEventEnricher: Send + Sync {
|
|||||||
async fn enrich_usage_event(&self, event: &mut UsageEvent) -> Result<(), DataLayerError>;
|
async fn enrich_usage_event(&self, event: &mut UsageEvent) -> Result<(), DataLayerError>;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||||
|
pub enum UsageRequestRecordLevel {
|
||||||
|
Basic,
|
||||||
|
#[default]
|
||||||
|
Full,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
pub trait UsageRuntimeAccess:
|
pub trait UsageRuntimeAccess:
|
||||||
UsageRecordWriter + UsageSettlementWriter + UsageBillingEventEnricher + Send + Sync
|
UsageRecordWriter + UsageSettlementWriter + UsageBillingEventEnricher + Send + Sync
|
||||||
{
|
{
|
||||||
fn has_usage_writer(&self) -> bool;
|
fn has_usage_writer(&self) -> bool;
|
||||||
fn has_usage_worker_runner(&self) -> bool;
|
fn has_usage_worker_runner(&self) -> bool;
|
||||||
fn usage_worker_runner(&self) -> Option<RedisStreamRunner>;
|
fn usage_worker_runner(&self) -> Option<RedisStreamRunner>;
|
||||||
|
|
||||||
|
async fn request_record_level(&self) -> Result<UsageRequestRecordLevel, DataLayerError> {
|
||||||
|
Ok(UsageRequestRecordLevel::Full)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
@@ -197,6 +209,7 @@ impl UsageRuntime {
|
|||||||
spawn_on_usage_background_runtime(boxed_usage_task(async move {
|
spawn_on_usage_background_runtime(boxed_usage_task(async move {
|
||||||
match build_sync_terminal_usage_event_offthread(input).await {
|
match build_sync_terminal_usage_event_offthread(input).await {
|
||||||
Ok(mut event) => {
|
Ok(mut event) => {
|
||||||
|
apply_request_record_level_from_data(&data, &mut event).await;
|
||||||
if let Err(err) = data.enrich_usage_event(&mut event).await {
|
if let Err(err) = data.enrich_usage_event(&mut event).await {
|
||||||
warn!(
|
warn!(
|
||||||
event_name = "usage_sync_terminal_billing_enrichment_failed",
|
event_name = "usage_sync_terminal_billing_enrichment_failed",
|
||||||
@@ -244,6 +257,7 @@ impl UsageRuntime {
|
|||||||
spawn_on_usage_background_runtime(boxed_usage_task(async move {
|
spawn_on_usage_background_runtime(boxed_usage_task(async move {
|
||||||
match build_stream_terminal_usage_event_offthread(input).await {
|
match build_stream_terminal_usage_event_offthread(input).await {
|
||||||
Ok(mut event) => {
|
Ok(mut event) => {
|
||||||
|
apply_request_record_level_from_data(&data, &mut event).await;
|
||||||
if let Err(err) = data.enrich_usage_event(&mut event).await {
|
if let Err(err) = data.enrich_usage_event(&mut event).await {
|
||||||
warn!(
|
warn!(
|
||||||
event_name = "usage_stream_terminal_billing_enrichment_failed",
|
event_name = "usage_stream_terminal_billing_enrichment_failed",
|
||||||
@@ -289,6 +303,7 @@ impl UsageRuntime {
|
|||||||
if !self.is_enabled() {
|
if !self.is_enabled() {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
apply_request_record_level_from_data(data, &mut event).await;
|
||||||
if let Err(err) = data.enrich_usage_event(&mut event).await {
|
if let Err(err) = data.enrich_usage_event(&mut event).await {
|
||||||
warn!(
|
warn!(
|
||||||
event_name = "usage_terminal_billing_enrichment_failed",
|
event_name = "usage_terminal_billing_enrichment_failed",
|
||||||
@@ -431,6 +446,40 @@ fn join_error_to_data_layer(err: tokio::task::JoinError) -> DataLayerError {
|
|||||||
DataLayerError::UnexpectedValue(format!("usage builder task join failed: {err}"))
|
DataLayerError::UnexpectedValue(format!("usage builder task join failed: {err}"))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn apply_request_record_level_from_data<T>(data: &T, event: &mut UsageEvent)
|
||||||
|
where
|
||||||
|
T: UsageRuntimeAccess,
|
||||||
|
{
|
||||||
|
match data.request_record_level().await {
|
||||||
|
Ok(level) => apply_request_record_level(level, event),
|
||||||
|
Err(err) => {
|
||||||
|
warn!(
|
||||||
|
event_name = "usage_request_record_level_read_failed",
|
||||||
|
log_type = "event",
|
||||||
|
request_id = %event.request_id,
|
||||||
|
fallback = "full",
|
||||||
|
error = %err,
|
||||||
|
"usage runtime failed to read request record level; keeping full capture"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_request_record_level(level: UsageRequestRecordLevel, event: &mut UsageEvent) {
|
||||||
|
if !matches!(level, UsageRequestRecordLevel::Basic) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
event.data.request_body = None;
|
||||||
|
event.data.request_body_ref = None;
|
||||||
|
event.data.provider_request_body = None;
|
||||||
|
event.data.provider_request_body_ref = None;
|
||||||
|
event.data.response_body = None;
|
||||||
|
event.data.response_body_ref = None;
|
||||||
|
event.data.client_response_body = None;
|
||||||
|
event.data.client_response_body_ref = None;
|
||||||
|
}
|
||||||
|
|
||||||
fn boxed_usage_task<F>(task: F) -> Pin<Box<dyn Future<Output = ()> + Send>>
|
fn boxed_usage_task<F>(task: F) -> Pin<Box<dyn Future<Output = ()> + Send>>
|
||||||
where
|
where
|
||||||
F: Future<Output = ()> + Send + 'static,
|
F: Future<Output = ()> + Send + 'static,
|
||||||
@@ -444,3 +493,51 @@ fn now_unix_secs() -> u64 {
|
|||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
.as_secs()
|
.as_secs()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use serde_json::json;
|
||||||
|
|
||||||
|
use super::{apply_request_record_level, UsageRequestRecordLevel};
|
||||||
|
use crate::{UsageEvent, UsageEventData, UsageEventType};
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn basic_request_record_level_strips_body_capture_but_preserves_derived_fields() {
|
||||||
|
let mut event = UsageEvent::new(
|
||||||
|
UsageEventType::Failed,
|
||||||
|
"req-basic-1",
|
||||||
|
UsageEventData {
|
||||||
|
provider_name: "OpenAI".to_string(),
|
||||||
|
model: "gpt-5".to_string(),
|
||||||
|
total_tokens: Some(42),
|
||||||
|
error_message: Some("upstream failed".to_string()),
|
||||||
|
request_body: Some(json!({"messages":[{"role":"user","content":"hello"}]})),
|
||||||
|
request_body_ref: Some("usage://request/req-basic-1/request_body".to_string()),
|
||||||
|
provider_request_body: Some(json!({"model":"gpt-5"})),
|
||||||
|
provider_request_body_ref: Some(
|
||||||
|
"usage://request/req-basic-1/provider_request_body".to_string(),
|
||||||
|
),
|
||||||
|
response_body: Some(json!({"error":{"message":"bad gateway"}})),
|
||||||
|
response_body_ref: Some("usage://request/req-basic-1/response_body".to_string()),
|
||||||
|
client_response_body: Some(json!({"detail":"bad gateway"})),
|
||||||
|
client_response_body_ref: Some(
|
||||||
|
"usage://request/req-basic-1/client_response_body".to_string(),
|
||||||
|
),
|
||||||
|
..UsageEventData::default()
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
apply_request_record_level(UsageRequestRecordLevel::Basic, &mut event);
|
||||||
|
|
||||||
|
assert_eq!(event.data.total_tokens, Some(42));
|
||||||
|
assert_eq!(event.data.error_message.as_deref(), Some("upstream failed"));
|
||||||
|
assert!(event.data.request_body.is_none());
|
||||||
|
assert!(event.data.request_body_ref.is_none());
|
||||||
|
assert!(event.data.provider_request_body.is_none());
|
||||||
|
assert!(event.data.provider_request_body_ref.is_none());
|
||||||
|
assert!(event.data.response_body.is_none());
|
||||||
|
assert!(event.data.response_body_ref.is_none());
|
||||||
|
assert!(event.data.client_response_body.is_none());
|
||||||
|
assert!(event.data.client_response_body_ref.is_none());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user