mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 09:57:47 +08:00
Consolidate subscription usage policy enforcement, privacy-safe persistence, and gateway security hardening into one reviewable change. Includes bounded HTTP and execution envelopes, header and protocol guards, DNS and relay validation, authentication and secret projection hardening, secure backup/install paths, and regression coverage.
1482 lines
53 KiB
Rust
1482 lines
53 KiB
Rust
use std::collections::{HashSet, VecDeque};
|
|
use std::sync::{Arc, Mutex};
|
|
use std::time::Duration;
|
|
|
|
use aether_contracts::tunnel_security::TUNNEL_SECURITY_NON_TLS_REQUIRED;
|
|
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
|
|
use aether_data::repository::proxy_nodes::{
|
|
bucket_start_unix_secs, InMemoryProxyNodeRepository, ProxyNodeHeartbeatMutation,
|
|
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeWriteRepository, StoredProxyNode,
|
|
};
|
|
use aether_data::repository::settlement::{
|
|
InMemorySettlementRepository, ReserveUsagePolicyCostInput, ReserveUsagePolicyRequestInput,
|
|
SettlementWriteRepository, UsagePolicyCostWindow, UsagePolicyRequestWindow,
|
|
};
|
|
use aether_runtime::bounded_queue;
|
|
use axum::extract::ws::Message;
|
|
use chrono::{DateTime, Utc};
|
|
use chrono_tz::Tz;
|
|
use serde_json::json;
|
|
use tokio::sync::watch;
|
|
|
|
use super::{
|
|
advance_proxy_upgrade_rollout_once, cleanup_audit_logs_with, cleanup_proxy_node_metrics_at,
|
|
cleanup_proxy_node_metrics_once, cleanup_stale_proxy_nodes_once, inspect_proxy_upgrade_rollout,
|
|
next_daily_run_after, next_db_maintenance_run_after, next_stats_aggregation_run_after,
|
|
next_stats_hourly_aggregation_run_after, pending_cleanup_batch_size,
|
|
pending_cleanup_timeout_minutes, perform_automatic_usage_cleanup_once,
|
|
perform_manual_usage_cleanup_once, plan_pending_cleanup_batch, provider_checkin_schedule,
|
|
proxy_node_metrics_cleanup_settings, record_proxy_upgrade_traffic_success,
|
|
run_db_maintenance_with, run_proxy_upgrade_rollout_once, spawn_account_self_check_worker,
|
|
spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
|
|
spawn_fixed_provider_reconciliation_task, spawn_oauth_token_refresh_worker,
|
|
spawn_pending_cleanup_worker, spawn_pool_monitor_worker, spawn_pool_quota_probe_worker,
|
|
spawn_provider_checkin_worker, spawn_proxy_node_stale_cleanup_worker,
|
|
spawn_proxy_upgrade_rollout_worker, spawn_stats_aggregation_worker,
|
|
spawn_stats_hourly_aggregation_worker, spawn_usage_cleanup_worker,
|
|
spawn_wallet_daily_usage_aggregation_worker, start_proxy_upgrade_rollout,
|
|
stats_aggregation_target_day, stats_hourly_aggregation_target_hour, summarize_database_pool,
|
|
usage_cleanup_settings, usage_cleanup_window, usage_cleanup_window_for_mode,
|
|
usage_cleanup_window_with_override, wallet_daily_usage_aggregation_target, AppState,
|
|
DbMaintenanceRunSummary, FailedPendingUsageRow, GatewayDataState, ManualUsageCleanupMode,
|
|
ProxyNodeMetricsCleanupSettings, ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow,
|
|
UsageCleanupSettings, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
|
|
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
|
|
};
|
|
|
|
#[tokio::test]
|
|
async fn spawn_audit_cleanup_worker_skips_when_postgres_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
assert!(spawn_audit_cleanup_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_db_maintenance_worker_skips_when_database_maintenance_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
assert!(spawn_db_maintenance_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_pending_cleanup_worker_skips_when_usage_writer_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
assert!(spawn_pending_cleanup_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_proxy_node_stale_cleanup_worker_skips_when_proxy_nodes_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
assert!(spawn_proxy_node_stale_cleanup_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_proxy_upgrade_rollout_worker_skips_when_proxy_nodes_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
assert!(spawn_proxy_upgrade_rollout_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_oauth_token_refresh_worker_skips_when_provider_catalog_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
assert!(spawn_oauth_token_refresh_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_fixed_provider_reconciliation_task_skips_when_provider_catalog_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
assert!(spawn_fixed_provider_reconciliation_task(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_proxy_upgrade_rollout_worker_skips_when_system_config_unavailable() {
|
|
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![]));
|
|
let data = GatewayDataState::with_proxy_node_repository_for_tests(repository);
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(data);
|
|
assert!(spawn_proxy_upgrade_rollout_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_pool_monitor_worker_skips_when_postgres_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
assert!(spawn_pool_monitor_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_pool_quota_probe_worker_skips_when_provider_catalog_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
|
|
assert!(spawn_pool_quota_probe_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_account_self_check_worker_skips_when_provider_catalog_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
|
|
assert!(spawn_account_self_check_worker(state).is_none());
|
|
}
|
|
|
|
fn sample_connected_proxy_node(
|
|
node_id: &str,
|
|
heartbeat_interval: i32,
|
|
last_heartbeat_at_unix_secs: u64,
|
|
) -> StoredProxyNode {
|
|
StoredProxyNode::new(
|
|
node_id.to_string(),
|
|
format!("proxy-{node_id}"),
|
|
"127.0.0.1".to_string(),
|
|
0,
|
|
false,
|
|
"online".to_string(),
|
|
heartbeat_interval,
|
|
3,
|
|
10,
|
|
0,
|
|
0,
|
|
0,
|
|
true,
|
|
true,
|
|
0,
|
|
)
|
|
.expect("node should build")
|
|
.with_runtime_fields(
|
|
Some("test".to_string()),
|
|
None,
|
|
Some(last_heartbeat_at_unix_secs),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
Some(last_heartbeat_at_unix_secs),
|
|
None,
|
|
Some(last_heartbeat_at_unix_secs),
|
|
Some(last_heartbeat_at_unix_secs),
|
|
)
|
|
}
|
|
|
|
const ACTIVE_PROBE_TUNNEL_TEST_PSK: &str = "BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=";
|
|
const ACTIVE_PROBE_TUNNEL_TEST_GENERATION: &str = "active-probe-test-generation-1";
|
|
|
|
fn active_probe_tunnel_metadata(version: &str) -> serde_json::Value {
|
|
json!({
|
|
"version": version,
|
|
"tunnel_security": {
|
|
"mode": TUNNEL_SECURITY_NON_TLS_REQUIRED,
|
|
"encryption_key": ACTIVE_PROBE_TUNNEL_TEST_PSK,
|
|
}
|
|
})
|
|
}
|
|
|
|
async fn recv_tunnel_test_frame(
|
|
proxy_rx: &mut aether_runtime::BoundedQueueReceiver<Message>,
|
|
description: &str,
|
|
) -> Message {
|
|
tokio::time::timeout(Duration::from_secs(5), proxy_rx.recv())
|
|
.await
|
|
.unwrap_or_else(|_| panic!("timed out waiting for {description}"))
|
|
.unwrap_or_else(|| panic!("proxy channel closed before {description}"))
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn stale_proxy_node_cleanup_marks_timed_out_tunnel_offline() {
|
|
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
|
|
sample_connected_proxy_node("node-stale", 30, 1),
|
|
]));
|
|
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository));
|
|
|
|
let updated = cleanup_stale_proxy_nodes_once(&data)
|
|
.await
|
|
.expect("cleanup should succeed");
|
|
|
|
assert_eq!(updated, 1);
|
|
let node = repository
|
|
.find_proxy_node("node-stale")
|
|
.await
|
|
.expect("lookup should succeed")
|
|
.expect("node should exist");
|
|
assert_eq!(node.status, "offline");
|
|
assert_eq!(node.tunnel_connected, false);
|
|
assert_eq!(node.active_connections, 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn proxy_upgrade_rollout_advances_next_wave_after_version_health_confirmation() {
|
|
let mut alpha = sample_connected_proxy_node("node-alpha", 30, 1_800_000_000);
|
|
alpha.name = "alpha".to_string();
|
|
alpha.proxy_metadata = Some(json!({"version": "1.0.0"}));
|
|
alpha.remote_config = None;
|
|
alpha.config_version = 0;
|
|
|
|
let mut zeta = sample_connected_proxy_node("node-zeta", 30, 1_800_000_000);
|
|
zeta.name = "zeta".to_string();
|
|
zeta.proxy_metadata = Some(json!({"version": "1.0.0"}));
|
|
zeta.remote_config = None;
|
|
zeta.config_version = 0;
|
|
|
|
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![zeta, alpha]));
|
|
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository))
|
|
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
|
|
|
|
let first_wave = start_proxy_upgrade_rollout(&data, "2.0.0".to_string(), 1, 0, None)
|
|
.await
|
|
.expect("rollout should start");
|
|
assert_eq!(first_wave.updated, 1);
|
|
assert_eq!(first_wave.node_ids, vec!["node-alpha".to_string()]);
|
|
assert!(first_wave.rollout_active);
|
|
|
|
let alpha_after_first = repository
|
|
.find_proxy_node("node-alpha")
|
|
.await
|
|
.expect("lookup should succeed")
|
|
.expect("alpha should exist");
|
|
assert_eq!(
|
|
alpha_after_first
|
|
.remote_config
|
|
.as_ref()
|
|
.and_then(|value| value.get("upgrade_to")),
|
|
Some(&json!("2.0.0"))
|
|
);
|
|
|
|
repository
|
|
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
|
node_id: "node-alpha".to_string(),
|
|
expected_tunnel_generation: None,
|
|
heartbeat_interval: None,
|
|
active_connections: Some(2),
|
|
total_requests_delta: Some(1),
|
|
avg_latency_ms: Some(2.0),
|
|
failed_requests_delta: Some(0),
|
|
dns_failures_delta: Some(0),
|
|
stream_errors_delta: Some(0),
|
|
proxy_metadata: Some(json!({"version": "2.0.0"})),
|
|
proxy_version: Some("2.0.0".to_string()),
|
|
})
|
|
.await
|
|
.expect("heartbeat should succeed");
|
|
|
|
let observed = advance_proxy_upgrade_rollout_once(&data)
|
|
.await
|
|
.expect("rollout should observe confirmed version");
|
|
assert_eq!(observed.updated, 0);
|
|
assert!(observed.blocked);
|
|
assert_eq!(observed.pending_node_ids, vec!["node-alpha".to_string()]);
|
|
|
|
assert!(!record_proxy_upgrade_traffic_success(&data, "node-zeta")
|
|
.await
|
|
.expect("untracked node traffic should be ignored"));
|
|
assert!(record_proxy_upgrade_traffic_success(&data, "node-alpha")
|
|
.await
|
|
.expect("traffic confirmation should be recorded"));
|
|
|
|
let second_wave = advance_proxy_upgrade_rollout_once(&data)
|
|
.await
|
|
.expect("rollout should advance after a healthy observation cycle");
|
|
assert_eq!(second_wave.updated, 1);
|
|
assert_eq!(second_wave.node_ids, vec!["node-zeta".to_string()]);
|
|
assert!(second_wave.rollout_active);
|
|
|
|
repository
|
|
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
|
node_id: "node-zeta".to_string(),
|
|
expected_tunnel_generation: None,
|
|
heartbeat_interval: None,
|
|
active_connections: Some(2),
|
|
total_requests_delta: Some(1),
|
|
avg_latency_ms: Some(2.0),
|
|
failed_requests_delta: Some(0),
|
|
dns_failures_delta: Some(0),
|
|
stream_errors_delta: Some(0),
|
|
proxy_metadata: Some(json!({"version": "2.0.0"})),
|
|
proxy_version: Some("tunnel-v2.0.0".to_string()),
|
|
})
|
|
.await
|
|
.expect("heartbeat should succeed");
|
|
|
|
let zeta_observed = advance_proxy_upgrade_rollout_once(&data)
|
|
.await
|
|
.expect("rollout should observe second wave");
|
|
assert_eq!(zeta_observed.updated, 0);
|
|
assert!(zeta_observed.blocked);
|
|
assert_eq!(
|
|
zeta_observed.pending_node_ids,
|
|
vec!["node-zeta".to_string()]
|
|
);
|
|
|
|
assert!(record_proxy_upgrade_traffic_success(&data, "node-zeta")
|
|
.await
|
|
.expect("traffic confirmation should be recorded"));
|
|
|
|
let finished = advance_proxy_upgrade_rollout_once(&data)
|
|
.await
|
|
.expect("rollout should finish after the second healthy observation cycle");
|
|
assert!(!finished.rollout_active);
|
|
assert_eq!(finished.completed, 2);
|
|
assert_eq!(finished.remaining, 0);
|
|
assert!(data
|
|
.list_system_config_entries()
|
|
.await
|
|
.expect("system config list should succeed")
|
|
.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn proxy_upgrade_rollout_excludes_draining_nodes_from_online_eligible_pool() {
|
|
let mut alpha = sample_connected_proxy_node("node-alpha", 30, 1_800_000_000);
|
|
alpha.name = "alpha".to_string();
|
|
alpha.proxy_metadata = Some(json!({"version": "1.0.0"}));
|
|
alpha.remote_config = Some(json!({"scheduling_state": "draining"}));
|
|
alpha.config_version = 1;
|
|
|
|
let mut zeta = sample_connected_proxy_node("node-zeta", 30, 1_800_000_000);
|
|
zeta.name = "zeta".to_string();
|
|
zeta.proxy_metadata = Some(json!({"version": "1.0.0"}));
|
|
zeta.remote_config = None;
|
|
zeta.config_version = 0;
|
|
|
|
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![zeta, alpha]));
|
|
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository))
|
|
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
|
|
|
|
let rollout = start_proxy_upgrade_rollout(&data, "2.0.0".to_string(), 2, 0, None)
|
|
.await
|
|
.expect("rollout should start");
|
|
assert_eq!(rollout.updated, 1);
|
|
assert_eq!(rollout.node_ids, vec!["node-zeta".to_string()]);
|
|
|
|
let alpha_after = repository
|
|
.find_proxy_node("node-alpha")
|
|
.await
|
|
.expect("lookup should succeed")
|
|
.expect("alpha should exist");
|
|
assert_eq!(
|
|
alpha_after
|
|
.remote_config
|
|
.as_ref()
|
|
.and_then(|value| value.get("upgrade_to")),
|
|
None
|
|
);
|
|
|
|
let rollout_status = inspect_proxy_upgrade_rollout(&data)
|
|
.await
|
|
.expect("inspect should succeed")
|
|
.expect("rollout should exist");
|
|
assert_eq!(rollout_status.online_eligible_total, 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn proxy_upgrade_rollout_blocks_next_wave_after_post_upgrade_transport_errors() {
|
|
let mut alpha = sample_connected_proxy_node("node-alpha", 30, 1_800_000_000);
|
|
alpha.name = "alpha".to_string();
|
|
alpha.proxy_metadata = Some(json!({"version": "1.0.0"}));
|
|
alpha.remote_config = None;
|
|
alpha.config_version = 0;
|
|
|
|
let mut zeta = sample_connected_proxy_node("node-zeta", 30, 1_800_000_000);
|
|
zeta.name = "zeta".to_string();
|
|
zeta.proxy_metadata = Some(json!({"version": "1.0.0"}));
|
|
zeta.remote_config = None;
|
|
zeta.config_version = 0;
|
|
|
|
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![zeta, alpha]));
|
|
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository))
|
|
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
|
|
|
|
let first_wave = start_proxy_upgrade_rollout(&data, "2.0.0".to_string(), 1, 0, None)
|
|
.await
|
|
.expect("rollout should start");
|
|
assert_eq!(first_wave.updated, 1);
|
|
assert_eq!(first_wave.node_ids, vec!["node-alpha".to_string()]);
|
|
|
|
repository
|
|
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
|
node_id: "node-alpha".to_string(),
|
|
expected_tunnel_generation: None,
|
|
heartbeat_interval: None,
|
|
active_connections: Some(2),
|
|
total_requests_delta: Some(1),
|
|
avg_latency_ms: Some(2.0),
|
|
failed_requests_delta: Some(0),
|
|
dns_failures_delta: Some(0),
|
|
stream_errors_delta: Some(0),
|
|
proxy_metadata: Some(json!({"version": "2.0.0"})),
|
|
proxy_version: Some("2.0.0".to_string()),
|
|
})
|
|
.await
|
|
.expect("heartbeat should succeed");
|
|
|
|
let observed = advance_proxy_upgrade_rollout_once(&data)
|
|
.await
|
|
.expect("rollout should observe the first upgraded node");
|
|
assert!(observed.blocked);
|
|
assert_eq!(observed.pending_node_ids, vec!["node-alpha".to_string()]);
|
|
|
|
assert!(record_proxy_upgrade_traffic_success(&data, "node-alpha")
|
|
.await
|
|
.expect("traffic confirmation should be recorded"));
|
|
|
|
tokio::time::sleep(Duration::from_millis(5)).await;
|
|
|
|
repository
|
|
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
|
node_id: "node-alpha".to_string(),
|
|
expected_tunnel_generation: None,
|
|
heartbeat_interval: None,
|
|
active_connections: Some(2),
|
|
total_requests_delta: Some(1),
|
|
avg_latency_ms: Some(2.0),
|
|
failed_requests_delta: Some(1),
|
|
dns_failures_delta: Some(0),
|
|
stream_errors_delta: Some(0),
|
|
proxy_metadata: Some(json!({"version": "2.0.0"})),
|
|
proxy_version: Some("2.0.0".to_string()),
|
|
})
|
|
.await
|
|
.expect("heartbeat should succeed");
|
|
|
|
let blocked = advance_proxy_upgrade_rollout_once(&data)
|
|
.await
|
|
.expect("rollout should stay blocked after post-upgrade transport errors");
|
|
assert_eq!(blocked.updated, 0);
|
|
assert!(blocked.blocked);
|
|
assert_eq!(blocked.pending_node_ids, vec!["node-alpha".to_string()]);
|
|
|
|
let zeta_after = repository
|
|
.find_proxy_node("node-zeta")
|
|
.await
|
|
.expect("lookup should succeed")
|
|
.expect("zeta should exist");
|
|
assert!(zeta_after.remote_config.is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn proxy_upgrade_rollout_active_probe_advances_next_wave_after_version_confirmation() {
|
|
let mut alpha = sample_connected_proxy_node("node-alpha", 30, 1_800_000_000)
|
|
.with_tunnel_generation(ACTIVE_PROBE_TUNNEL_TEST_GENERATION.to_string());
|
|
alpha.name = "alpha".to_string();
|
|
alpha.proxy_metadata = Some(active_probe_tunnel_metadata("1.0.0"));
|
|
alpha.remote_config = None;
|
|
alpha.config_version = 0;
|
|
|
|
let mut zeta = sample_connected_proxy_node("node-zeta", 30, 1_800_000_000);
|
|
zeta.name = "zeta".to_string();
|
|
zeta.proxy_metadata = Some(json!({"version": "1.0.0"}));
|
|
zeta.remote_config = None;
|
|
zeta.config_version = 0;
|
|
|
|
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![zeta, alpha]));
|
|
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository))
|
|
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new())
|
|
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(data.clone());
|
|
|
|
let first_wave = start_proxy_upgrade_rollout(
|
|
&data,
|
|
"2.0.0".to_string(),
|
|
1,
|
|
0,
|
|
Some(ProxyUpgradeRolloutProbeConfig {
|
|
url: "https://probe.example/health".to_string(),
|
|
timeout_secs: 5,
|
|
}),
|
|
)
|
|
.await
|
|
.expect("rollout should start");
|
|
assert_eq!(first_wave.node_ids, vec!["node-alpha".to_string()]);
|
|
|
|
repository
|
|
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
|
node_id: "node-alpha".to_string(),
|
|
expected_tunnel_generation: None,
|
|
heartbeat_interval: None,
|
|
active_connections: Some(2),
|
|
total_requests_delta: Some(1),
|
|
avg_latency_ms: Some(2.0),
|
|
failed_requests_delta: Some(0),
|
|
dns_failures_delta: Some(0),
|
|
stream_errors_delta: Some(0),
|
|
proxy_metadata: Some(active_probe_tunnel_metadata("2.0.0")),
|
|
proxy_version: Some("2.0.0".to_string()),
|
|
})
|
|
.await
|
|
.expect("heartbeat should succeed");
|
|
|
|
let tunnel_state = state.tunnel.app_state();
|
|
let (proxy_tx, mut proxy_rx) = bounded_queue(8);
|
|
let (proxy_close_tx, _) = watch::channel(false);
|
|
tunnel_state.hub.register_proxy(Arc::new(
|
|
crate::tunnel::TunnelProxyConn::new(
|
|
700,
|
|
"node-alpha".to_string(),
|
|
"Node Alpha".to_string(),
|
|
proxy_tx,
|
|
proxy_close_tx,
|
|
16,
|
|
2,
|
|
)
|
|
.with_tunnel_generation(ACTIVE_PROBE_TUNNEL_TEST_GENERATION.to_string())
|
|
.with_authenticated_key(ACTIVE_PROBE_TUNNEL_TEST_PSK.to_string()),
|
|
));
|
|
|
|
let responder_hub = tunnel_state.hub.clone();
|
|
let responder = tokio::spawn(async move {
|
|
let request_headers =
|
|
match recv_tunnel_test_frame(&mut proxy_rx, "probe headers frame").await {
|
|
Message::Binary(data) => data,
|
|
other => panic!("unexpected message: {other:?}"),
|
|
};
|
|
let request_header = crate::tunnel::tunnel_protocol::FrameHeader::parse(&request_headers)
|
|
.expect("probe request headers should parse");
|
|
assert_eq!(
|
|
request_header.msg_type,
|
|
crate::tunnel::tunnel_protocol::REQUEST_HEADERS
|
|
);
|
|
|
|
let request_body = match recv_tunnel_test_frame(&mut proxy_rx, "probe body frame").await {
|
|
Message::Binary(data) => data,
|
|
other => panic!("unexpected message: {other:?}"),
|
|
};
|
|
let request_body_header = crate::tunnel::tunnel_protocol::FrameHeader::parse(&request_body)
|
|
.expect("probe request body should parse");
|
|
assert_eq!(
|
|
request_body_header.msg_type,
|
|
crate::tunnel::tunnel_protocol::REQUEST_BODY
|
|
);
|
|
|
|
let response_meta = crate::tunnel::tunnel_protocol::ResponseMeta {
|
|
status: 204,
|
|
headers: vec![],
|
|
};
|
|
let response_payload =
|
|
serde_json::to_vec(&response_meta).expect("response meta should serialize");
|
|
let mut response_headers_frame = crate::tunnel::tunnel_protocol::encode_frame(
|
|
request_header.stream_id,
|
|
crate::tunnel::tunnel_protocol::RESPONSE_HEADERS,
|
|
0,
|
|
&response_payload,
|
|
);
|
|
responder_hub
|
|
.handle_proxy_frame(700, &mut response_headers_frame)
|
|
.await;
|
|
|
|
let mut response_end_frame = crate::tunnel::tunnel_protocol::encode_frame(
|
|
request_header.stream_id,
|
|
crate::tunnel::tunnel_protocol::STREAM_END,
|
|
0,
|
|
&[],
|
|
);
|
|
responder_hub
|
|
.handle_proxy_frame(700, &mut response_end_frame)
|
|
.await;
|
|
});
|
|
|
|
run_proxy_upgrade_rollout_once(&state)
|
|
.await
|
|
.expect("rollout worker should succeed");
|
|
responder.await.expect("probe responder should complete");
|
|
|
|
let zeta_after = repository
|
|
.find_proxy_node("node-zeta")
|
|
.await
|
|
.expect("lookup should succeed")
|
|
.expect("zeta should exist");
|
|
assert_eq!(
|
|
zeta_after
|
|
.remote_config
|
|
.as_ref()
|
|
.and_then(|value| value.get("upgrade_to")),
|
|
Some(&json!("2.0.0"))
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_stats_aggregation_worker_skips_when_stats_daily_backend_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
assert!(spawn_stats_aggregation_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_stats_hourly_aggregation_worker_skips_when_stats_hourly_backend_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
assert!(spawn_stats_hourly_aggregation_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_usage_cleanup_worker_skips_when_usage_writer_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
assert!(spawn_usage_cleanup_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_usage_cleanup_worker_starts_for_settlement_writer_only() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(
|
|
GatewayDataState::disabled().with_settlement_writer_for_tests(Arc::new(
|
|
InMemorySettlementRepository::default(),
|
|
)),
|
|
);
|
|
let handle = spawn_usage_cleanup_worker(state)
|
|
.expect("settlement ledger cleanup should have a worker without a usage writer");
|
|
handle.abort();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn automatic_usage_cleanup_collects_ledgers_when_usage_cleanup_is_disabled() {
|
|
let repository = Arc::new(InMemorySettlementRepository::default());
|
|
repository
|
|
.reserve_usage_policy_request(ReserveUsagePolicyRequestInput {
|
|
request_id: "expired-request".to_string(),
|
|
subject_id: "user-1".to_string(),
|
|
event_token: "expired-event".to_string(),
|
|
admitted_at_unix_secs: 1,
|
|
retain_until_unix_secs: 2,
|
|
windows: vec![UsagePolicyRequestWindow {
|
|
starts_at_unix_secs: 0,
|
|
ends_at_unix_secs: 2,
|
|
limit_requests: 10,
|
|
}],
|
|
})
|
|
.await
|
|
.expect("request admission should be seeded");
|
|
repository
|
|
.reserve_usage_policy_cost(ReserveUsagePolicyCostInput {
|
|
request_id: "expired-cost".to_string(),
|
|
subject_id: "user-1".to_string(),
|
|
reservation_token: "expired-reservation".to_string(),
|
|
admitted_at_unix_secs: 1,
|
|
reserved_cost_units: 1,
|
|
reservation_expires_at_unix_secs: 2,
|
|
retain_until_unix_secs: 2,
|
|
windows: vec![UsagePolicyCostWindow {
|
|
window_id: "expired-window".to_string(),
|
|
starts_at_unix_secs: 0,
|
|
ends_at_unix_secs: 2,
|
|
limit_cost_units: 10,
|
|
}],
|
|
})
|
|
.await
|
|
.expect("cost reservation should be seeded");
|
|
let data = GatewayDataState::disabled()
|
|
.with_settlement_writer_for_tests(repository)
|
|
.with_system_config_values_for_tests([
|
|
("enable_auto_cleanup".to_string(), json!(false)),
|
|
("cleanup_batch_size".to_string(), json!(1)),
|
|
]);
|
|
|
|
let summary = perform_automatic_usage_cleanup_once(&data)
|
|
.await
|
|
.expect("automatic ledger cleanup should succeed");
|
|
|
|
assert_eq!(summary.cost_reservations_deleted, 1);
|
|
assert_eq!(summary.request_admissions_deleted, 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn manual_usage_cleanup_does_not_collect_policy_ledgers() {
|
|
let repository = Arc::new(InMemorySettlementRepository::default());
|
|
repository
|
|
.reserve_usage_policy_request(ReserveUsagePolicyRequestInput {
|
|
request_id: "manual-request".to_string(),
|
|
subject_id: "user-1".to_string(),
|
|
event_token: "manual-event".to_string(),
|
|
admitted_at_unix_secs: 1,
|
|
retain_until_unix_secs: 2,
|
|
windows: vec![UsagePolicyRequestWindow {
|
|
starts_at_unix_secs: 0,
|
|
ends_at_unix_secs: 2,
|
|
limit_requests: 10,
|
|
}],
|
|
})
|
|
.await
|
|
.expect("request admission should be seeded");
|
|
let data = GatewayDataState::disabled().with_settlement_writer_for_tests(repository.clone());
|
|
|
|
let summary =
|
|
perform_manual_usage_cleanup_once(&data, super::ManualUsageCleanupOptions::policy())
|
|
.await
|
|
.expect("manual usage cleanup should succeed");
|
|
|
|
assert_eq!(summary.request_admissions_deleted, 0);
|
|
assert_eq!(
|
|
repository
|
|
.cleanup_usage_policy_request_admissions(u64::MAX, 10)
|
|
.await
|
|
.expect("direct cleanup should find the retained admission"),
|
|
1
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_wallet_daily_usage_aggregation_worker_skips_when_wallet_daily_usage_backend_unavailable(
|
|
) {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
assert!(spawn_wallet_daily_usage_aggregation_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_provider_checkin_worker_skips_when_provider_catalog_unavailable() {
|
|
let state = AppState::new()
|
|
.expect("gateway state should build")
|
|
.with_data_state_for_tests(GatewayDataState::disabled());
|
|
|
|
assert!(spawn_provider_checkin_worker(state).is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn cleanup_audit_logs_respects_auto_cleanup_toggle() {
|
|
let data = GatewayDataState::disabled()
|
|
.with_system_config_values_for_tests([("enable_auto_cleanup".to_string(), json!(false))]);
|
|
|
|
let deleted = cleanup_audit_logs_with(&data, |_cutoff_time, _delete_limit| async move {
|
|
panic!("audit cleanup should not run when auto cleanup is disabled");
|
|
#[allow(unreachable_code)]
|
|
Ok(0)
|
|
})
|
|
.await
|
|
.expect("audit cleanup should short-circuit");
|
|
|
|
assert_eq!(deleted, 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn cleanup_audit_logs_uses_retention_and_batch_settings() {
|
|
let data = GatewayDataState::disabled().with_system_config_values_for_tests([
|
|
("enable_auto_cleanup".to_string(), json!(true)),
|
|
("audit_log_retention_days".to_string(), json!(21)),
|
|
("cleanup_batch_size".to_string(), json!(2)),
|
|
]);
|
|
let observed_limits = Arc::new(Mutex::new(Vec::new()));
|
|
let observed_cutoffs = Arc::new(Mutex::new(Vec::new()));
|
|
let batch_results = Arc::new(Mutex::new(VecDeque::from([2usize, 1usize])));
|
|
let started_at = Utc::now();
|
|
|
|
let deleted = cleanup_audit_logs_with(&data, {
|
|
let observed_limits = Arc::clone(&observed_limits);
|
|
let observed_cutoffs = Arc::clone(&observed_cutoffs);
|
|
let batch_results = Arc::clone(&batch_results);
|
|
move |cutoff_time, delete_limit| {
|
|
observed_limits
|
|
.lock()
|
|
.expect("observed limits lock")
|
|
.push(delete_limit);
|
|
observed_cutoffs
|
|
.lock()
|
|
.expect("observed cutoffs lock")
|
|
.push(cutoff_time);
|
|
let next = batch_results
|
|
.lock()
|
|
.expect("batch results lock")
|
|
.pop_front()
|
|
.unwrap_or_default();
|
|
async move { Ok(next) }
|
|
}
|
|
})
|
|
.await
|
|
.expect("audit cleanup should succeed");
|
|
let finished_at = Utc::now();
|
|
|
|
assert_eq!(deleted, 3);
|
|
assert_eq!(
|
|
*observed_limits.lock().expect("observed limits lock"),
|
|
vec![2, 2]
|
|
);
|
|
let observed_cutoffs = observed_cutoffs.lock().expect("observed cutoffs lock");
|
|
assert_eq!(observed_cutoffs.len(), 2);
|
|
let earliest_expected = started_at - chrono::Duration::days(21);
|
|
let latest_expected = finished_at - chrono::Duration::days(21);
|
|
for cutoff_time in observed_cutoffs.iter() {
|
|
assert!(*cutoff_time >= earliest_expected);
|
|
assert!(*cutoff_time <= latest_expected);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn pending_cleanup_settings_use_timeout_and_cap_batch_size() {
|
|
let data = GatewayDataState::disabled().with_system_config_values_for_tests([
|
|
("pending_request_timeout_minutes".to_string(), json!(25)),
|
|
("cleanup_batch_size".to_string(), json!(500)),
|
|
]);
|
|
|
|
let timeout_minutes = pending_cleanup_timeout_minutes(&data)
|
|
.await
|
|
.expect("timeout should resolve");
|
|
let batch_size = pending_cleanup_batch_size(&data)
|
|
.await
|
|
.expect("batch size should resolve");
|
|
|
|
assert_eq!(timeout_minutes, 25);
|
|
assert_eq!(batch_size, 200);
|
|
}
|
|
|
|
#[test]
|
|
fn pending_cleanup_plan_recovers_completed_requests_and_voids_failed_pending_billing() {
|
|
let plan = plan_pending_cleanup_batch(
|
|
vec![
|
|
StalePendingUsageRow {
|
|
id: "usage-1".to_string(),
|
|
request_id: "req-1".to_string(),
|
|
status: "streaming".to_string(),
|
|
billing_status: "pending".to_string(),
|
|
},
|
|
StalePendingUsageRow {
|
|
id: "usage-2".to_string(),
|
|
request_id: "req-2".to_string(),
|
|
status: "pending".to_string(),
|
|
billing_status: "pending".to_string(),
|
|
},
|
|
StalePendingUsageRow {
|
|
id: "usage-3".to_string(),
|
|
request_id: "req-3".to_string(),
|
|
status: "streaming".to_string(),
|
|
billing_status: "settled".to_string(),
|
|
},
|
|
],
|
|
&HashSet::from(["req-1".to_string()]),
|
|
10,
|
|
);
|
|
|
|
assert_eq!(plan.recovered_usage_ids, vec!["usage-1".to_string()]);
|
|
assert_eq!(plan.recovered_request_ids, vec!["req-1".to_string()]);
|
|
assert_eq!(
|
|
plan.failed_request_ids,
|
|
vec!["req-2".to_string(), "req-3".to_string()]
|
|
);
|
|
assert_eq!(
|
|
plan.failed_usage_rows,
|
|
vec![
|
|
FailedPendingUsageRow {
|
|
id: "usage-2".to_string(),
|
|
error_message: "请求超时: 状态 'pending' 超过 10 分钟未完成".to_string(),
|
|
should_void_billing: true,
|
|
},
|
|
FailedPendingUsageRow {
|
|
id: "usage-3".to_string(),
|
|
error_message: "请求超时: 状态 'streaming' 超过 10 分钟未完成".to_string(),
|
|
should_void_billing: false,
|
|
},
|
|
]
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn usage_cleanup_settings_resolve_batch_and_delete_toggle() {
|
|
let data = GatewayDataState::disabled().with_system_config_values_for_tests([
|
|
("detail_log_retention_days".to_string(), json!(7)),
|
|
("compressed_log_retention_days".to_string(), json!(30)),
|
|
("header_retention_days".to_string(), json!(90)),
|
|
("log_retention_days".to_string(), json!(365)),
|
|
("cleanup_batch_size".to_string(), json!(0)),
|
|
("auto_delete_expired_keys".to_string(), json!(true)),
|
|
]);
|
|
|
|
let settings = usage_cleanup_settings(&data)
|
|
.await
|
|
.expect("usage cleanup settings should resolve");
|
|
|
|
assert_eq!(
|
|
settings,
|
|
UsageCleanupSettings {
|
|
detail_retention_days: 7,
|
|
compressed_retention_days: 30,
|
|
header_retention_days: 90,
|
|
log_retention_days: 365,
|
|
batch_size: 1,
|
|
auto_delete_expired_keys: true,
|
|
}
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn proxy_node_metrics_cleanup_settings_use_dedicated_retention_and_batch_limits() {
|
|
let data = GatewayDataState::disabled().with_system_config_values_for_tests([
|
|
("cleanup_batch_size".to_string(), json!(250)),
|
|
("proxy_node_metrics_1m_retention_days".to_string(), json!(0)),
|
|
("proxy_node_metrics_1h_retention_days".to_string(), json!(7)),
|
|
(
|
|
"proxy_node_metrics_cleanup_batch_size".to_string(),
|
|
json!(100_000),
|
|
),
|
|
]);
|
|
|
|
let settings = proxy_node_metrics_cleanup_settings(&data)
|
|
.await
|
|
.expect("proxy metrics cleanup settings should resolve");
|
|
|
|
assert_eq!(
|
|
settings,
|
|
ProxyNodeMetricsCleanupSettings {
|
|
retain_1m_days: 1,
|
|
retain_1h_days: 7,
|
|
batch_size: 50_000,
|
|
}
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn proxy_node_metrics_cleanup_settings_fallback_to_global_batch_size() {
|
|
let data = GatewayDataState::disabled().with_system_config_values_for_tests([
|
|
("cleanup_batch_size".to_string(), json!(250)),
|
|
(
|
|
"proxy_node_metrics_cleanup_batch_size".to_string(),
|
|
json!(0),
|
|
),
|
|
]);
|
|
|
|
let settings = proxy_node_metrics_cleanup_settings(&data)
|
|
.await
|
|
.expect("proxy metrics cleanup settings should resolve");
|
|
|
|
assert_eq!(
|
|
settings,
|
|
ProxyNodeMetricsCleanupSettings {
|
|
retain_1m_days: 30,
|
|
retain_1h_days: 180,
|
|
batch_size: 250,
|
|
}
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn proxy_node_metrics_cleanup_deletes_expired_buckets_in_batches() {
|
|
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
|
|
sample_connected_proxy_node("node-metrics-1", 30, 1),
|
|
sample_connected_proxy_node("node-metrics-2", 30, 1),
|
|
sample_connected_proxy_node("node-metrics-3", 30, 1),
|
|
]));
|
|
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository))
|
|
.with_system_config_values_for_tests([
|
|
("enable_auto_cleanup".to_string(), json!(true)),
|
|
("proxy_node_metrics_1m_retention_days".to_string(), json!(1)),
|
|
("proxy_node_metrics_1h_retention_days".to_string(), json!(1)),
|
|
(
|
|
"proxy_node_metrics_cleanup_batch_size".to_string(),
|
|
json!(1),
|
|
),
|
|
]);
|
|
|
|
for (idx, node_id) in ["node-metrics-1", "node-metrics-2", "node-metrics-3"]
|
|
.into_iter()
|
|
.enumerate()
|
|
{
|
|
repository
|
|
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
|
node_id: node_id.to_string(),
|
|
expected_tunnel_generation: None,
|
|
heartbeat_interval: Some(30),
|
|
active_connections: Some(i32::try_from(idx + 1).unwrap()),
|
|
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": idx + 1,
|
|
"disconnects": 0,
|
|
"error_events_total": 0,
|
|
"ws_in_bytes": idx + 1,
|
|
"ws_out_bytes": idx + 1,
|
|
"ws_in_frames": idx + 1,
|
|
"ws_out_frames": idx + 1,
|
|
"heartbeat_rtt_last_ms": 10
|
|
}
|
|
})),
|
|
proxy_version: Some("1.0.0".to_string()),
|
|
})
|
|
.await
|
|
.expect("heartbeat should write metrics");
|
|
}
|
|
|
|
let now = chrono::Utc::now().timestamp().max(0) as u64;
|
|
let old_bucket = bucket_start_unix_secs(now, ProxyNodeMetricsStep::OneMinute);
|
|
let cleanup = data
|
|
.cleanup_proxy_node_metrics(
|
|
old_bucket.saturating_add(60),
|
|
old_bucket.saturating_add(3_600),
|
|
1,
|
|
)
|
|
.await
|
|
.expect("direct cleanup should delete a limited batch");
|
|
assert_eq!(cleanup.deleted_1m_rows, 1);
|
|
assert_eq!(cleanup.deleted_1h_rows, 1);
|
|
|
|
let cleanup = cleanup_proxy_node_metrics_at(&data, now.saturating_add(2 * 86_400))
|
|
.await
|
|
.expect("runtime cleanup should loop over batches");
|
|
assert_eq!(cleanup.deleted_1m_rows, 2);
|
|
assert_eq!(cleanup.deleted_1h_rows, 2);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn proxy_node_metrics_cleanup_respects_auto_cleanup_toggle() {
|
|
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
|
|
sample_connected_proxy_node("node-metrics-disabled", 30, 1),
|
|
]));
|
|
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository))
|
|
.with_system_config_values_for_tests([
|
|
("enable_auto_cleanup".to_string(), json!(false)),
|
|
("proxy_node_metrics_1m_retention_days".to_string(), json!(1)),
|
|
("proxy_node_metrics_1h_retention_days".to_string(), json!(1)),
|
|
(
|
|
"proxy_node_metrics_cleanup_batch_size".to_string(),
|
|
json!(1),
|
|
),
|
|
]);
|
|
|
|
repository
|
|
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
|
|
node_id: "node-metrics-disabled".to_string(),
|
|
expected_tunnel_generation: None,
|
|
heartbeat_interval: Some(30),
|
|
active_connections: Some(1),
|
|
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": 1,
|
|
"disconnects": 0,
|
|
"error_events_total": 0,
|
|
"ws_in_bytes": 1,
|
|
"ws_out_bytes": 1,
|
|
"ws_in_frames": 1,
|
|
"ws_out_frames": 1,
|
|
"heartbeat_rtt_last_ms": 10
|
|
}
|
|
})),
|
|
proxy_version: Some("1.0.0".to_string()),
|
|
})
|
|
.await
|
|
.expect("heartbeat should write metrics");
|
|
|
|
let cleanup = cleanup_proxy_node_metrics_once(&data)
|
|
.await
|
|
.expect("runtime cleanup should short-circuit");
|
|
assert_eq!(cleanup.deleted_1m_rows, 0);
|
|
assert_eq!(cleanup.deleted_1h_rows, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn usage_cleanup_window_uses_non_overlapping_ranges() {
|
|
let now_utc = "2026-03-18T03:00:00Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
let window = usage_cleanup_window(
|
|
now_utc,
|
|
UsageCleanupSettings {
|
|
detail_retention_days: 7,
|
|
compressed_retention_days: 30,
|
|
header_retention_days: 90,
|
|
log_retention_days: 365,
|
|
batch_size: 123,
|
|
auto_delete_expired_keys: false,
|
|
},
|
|
);
|
|
|
|
assert_eq!(
|
|
window.detail_cutoff.to_rfc3339(),
|
|
"2026-03-11T03:00:00+00:00"
|
|
);
|
|
assert_eq!(
|
|
window.compressed_cutoff.to_rfc3339(),
|
|
"2026-02-16T03:00:00+00:00"
|
|
);
|
|
assert_eq!(
|
|
window.header_cutoff.to_rfc3339(),
|
|
"2025-12-18T03:00:00+00:00"
|
|
);
|
|
assert_eq!(window.log_cutoff.to_rfc3339(), "2025-03-18T03:00:00+00:00");
|
|
assert!(window.detail_cutoff > window.compressed_cutoff);
|
|
assert!(window.compressed_cutoff > window.log_cutoff);
|
|
}
|
|
|
|
#[test]
|
|
fn usage_cleanup_window_with_override_is_always_non_aggressive() {
|
|
let now_utc = "2026-03-18T03:00:00Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
let settings = UsageCleanupSettings {
|
|
detail_retention_days: 7,
|
|
compressed_retention_days: 30,
|
|
header_retention_days: 90,
|
|
log_retention_days: 365,
|
|
batch_size: 123,
|
|
auto_delete_expired_keys: false,
|
|
};
|
|
let policy = usage_cleanup_window(now_utc, settings);
|
|
|
|
let override_duration = chrono::Duration::days(180);
|
|
let clamped = usage_cleanup_window_with_override(now_utc, settings, Some(override_duration));
|
|
|
|
assert_eq!(clamped.detail_cutoff, policy.detail_cutoff);
|
|
assert_eq!(clamped.compressed_cutoff, policy.compressed_cutoff);
|
|
assert_eq!(clamped.header_cutoff, policy.header_cutoff);
|
|
assert_eq!(clamped.log_cutoff, now_utc - override_duration);
|
|
assert!(clamped.log_cutoff > policy.log_cutoff);
|
|
|
|
let far_override = chrono::Duration::days(5);
|
|
let far = usage_cleanup_window_with_override(now_utc, settings, Some(far_override));
|
|
assert_eq!(far.detail_cutoff, now_utc - far_override);
|
|
assert_eq!(far.compressed_cutoff, now_utc - far_override);
|
|
assert_eq!(far.header_cutoff, now_utc - far_override);
|
|
assert_eq!(far.log_cutoff, now_utc - far_override);
|
|
assert!(far.log_cutoff > policy.log_cutoff);
|
|
|
|
let passthrough = usage_cleanup_window_with_override(now_utc, settings, None);
|
|
assert_eq!(passthrough, policy);
|
|
}
|
|
|
|
#[test]
|
|
fn usage_cleanup_before_now_window_uses_current_timestamp_only() {
|
|
let now_utc = "2026-03-18T03:00:00Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
let settings = UsageCleanupSettings {
|
|
detail_retention_days: 7,
|
|
compressed_retention_days: 30,
|
|
header_retention_days: 90,
|
|
log_retention_days: 365,
|
|
batch_size: 123,
|
|
auto_delete_expired_keys: false,
|
|
};
|
|
|
|
let window =
|
|
usage_cleanup_window_for_mode(now_utc, settings, ManualUsageCleanupMode::BeforeNow, None);
|
|
|
|
assert_eq!(window.detail_cutoff, now_utc);
|
|
assert_eq!(window.compressed_cutoff, now_utc);
|
|
assert_eq!(window.header_cutoff, now_utc);
|
|
assert_eq!(window.log_cutoff, now_utc);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn summarize_database_pool_uses_busy_connections_for_usage_rate() {
|
|
let data = GatewayDataState::from_config(crate::data::GatewayDataConfig::from_postgres_config(
|
|
aether_data::driver::postgres::PostgresPoolConfig {
|
|
database_url: "postgres://localhost/aether".to_string(),
|
|
min_connections: 1,
|
|
max_connections: 8,
|
|
acquire_timeout_ms: 1_000,
|
|
idle_timeout_ms: 5_000,
|
|
max_lifetime_ms: 30_000,
|
|
statement_cache_capacity: 64,
|
|
require_ssl: false,
|
|
},
|
|
))
|
|
.expect("gateway data state should build");
|
|
|
|
let summary = summarize_database_pool(&data).expect("pool summary should exist");
|
|
|
|
assert_eq!(summary.driver, aether_data::DatabaseDriver::Postgres);
|
|
assert_eq!(summary.checked_out, 0);
|
|
assert_eq!(summary.pool_size, 0);
|
|
assert_eq!(summary.idle, 0);
|
|
assert_eq!(summary.max_connections, 8);
|
|
assert_eq!(summary.usage_rate, 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn stats_aggregation_target_uses_previous_utc_day() {
|
|
let now_utc = "2026-04-05T10:20:30Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let target = stats_aggregation_target_day(now_utc);
|
|
|
|
assert_eq!(target.to_rfc3339(), "2026-04-04T00:00:00+00:00");
|
|
}
|
|
|
|
#[test]
|
|
fn next_stats_aggregation_run_aligns_to_same_day_when_before_slot() {
|
|
let now_utc = "2026-04-05T00:04:59Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let next = next_stats_aggregation_run_after(now_utc);
|
|
|
|
assert_eq!(next.to_rfc3339(), "2026-04-05T00:05:00+00:00");
|
|
}
|
|
|
|
#[test]
|
|
fn next_stats_aggregation_run_rolls_to_next_day_after_slot() {
|
|
let now_utc = "2026-04-05T00:05:00Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let next = next_stats_aggregation_run_after(now_utc);
|
|
|
|
assert_eq!(next.to_rfc3339(), "2026-04-06T00:05:00+00:00");
|
|
}
|
|
|
|
#[test]
|
|
fn stats_hourly_aggregation_target_uses_previous_utc_hour() {
|
|
let now_utc = "2026-04-05T10:20:30Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let target = stats_hourly_aggregation_target_hour(now_utc);
|
|
|
|
assert_eq!(target.to_rfc3339(), "2026-04-05T09:00:00+00:00");
|
|
}
|
|
|
|
#[test]
|
|
fn next_stats_hourly_aggregation_run_aligns_to_same_hour_when_before_slot() {
|
|
let now_utc = "2026-04-05T10:04:59Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let next = next_stats_hourly_aggregation_run_after(now_utc);
|
|
|
|
assert_eq!(next.to_rfc3339(), "2026-04-05T10:05:00+00:00");
|
|
}
|
|
|
|
#[test]
|
|
fn next_stats_hourly_aggregation_run_rolls_to_next_hour_after_slot() {
|
|
let now_utc = "2026-04-05T10:05:00Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let next = next_stats_hourly_aggregation_run_after(now_utc);
|
|
|
|
assert_eq!(next.to_rfc3339(), "2026-04-05T11:05:00+00:00");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn db_maintenance_respects_enable_toggle() {
|
|
let data = GatewayDataState::disabled()
|
|
.with_system_config_values_for_tests([("enable_db_maintenance".to_string(), json!(false))]);
|
|
|
|
let summary = run_db_maintenance_with(&data, |_table_name| async move {
|
|
panic!("db maintenance should not run when disabled");
|
|
#[allow(unreachable_code)]
|
|
Ok(())
|
|
})
|
|
.await
|
|
.expect("db maintenance should short-circuit");
|
|
|
|
assert_eq!(
|
|
summary,
|
|
DbMaintenanceRunSummary {
|
|
attempted: 0,
|
|
succeeded: 0,
|
|
}
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn db_maintenance_continues_across_table_failures() {
|
|
let data = GatewayDataState::disabled()
|
|
.with_system_config_values_for_tests([("enable_db_maintenance".to_string(), json!(true))]);
|
|
let seen_tables = Arc::new(Mutex::new(Vec::new()));
|
|
|
|
let summary = run_db_maintenance_with(&data, {
|
|
let seen_tables = Arc::clone(&seen_tables);
|
|
move |table_name| {
|
|
seen_tables
|
|
.lock()
|
|
.expect("seen tables lock")
|
|
.push(table_name.to_string());
|
|
async move {
|
|
if table_name == "request_candidates" {
|
|
Err(aether_data::DataLayerError::InvalidInput(
|
|
"boom".to_string(),
|
|
))
|
|
} else {
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
})
|
|
.await
|
|
.expect("db maintenance should continue after failures");
|
|
|
|
assert_eq!(
|
|
summary,
|
|
DbMaintenanceRunSummary {
|
|
attempted: 3,
|
|
succeeded: 2,
|
|
}
|
|
);
|
|
assert_eq!(
|
|
*seen_tables.lock().expect("seen tables lock"),
|
|
vec![
|
|
"usage".to_string(),
|
|
"request_candidates".to_string(),
|
|
"audit_logs".to_string(),
|
|
]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn next_db_maintenance_run_aligns_to_same_week_when_before_slot() {
|
|
let timezone: Tz = "Asia/Shanghai".parse().expect("timezone should parse");
|
|
let now_utc = "2026-04-03T20:59:00Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let next = next_db_maintenance_run_after(now_utc, timezone);
|
|
|
|
assert_eq!(next.to_rfc3339(), "2026-04-04T21:00:00+00:00");
|
|
}
|
|
|
|
#[test]
|
|
fn next_db_maintenance_run_rolls_to_next_week_after_slot() {
|
|
let timezone: Tz = "Asia/Shanghai".parse().expect("timezone should parse");
|
|
let now_utc = "2026-04-04T21:00:01Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let next = next_db_maintenance_run_after(now_utc, timezone);
|
|
|
|
assert_eq!(next.to_rfc3339(), "2026-04-11T21:00:00+00:00");
|
|
}
|
|
|
|
#[test]
|
|
fn wallet_daily_usage_aggregation_target_uses_previous_local_day_window() {
|
|
let timezone: Tz = "Asia/Shanghai".parse().expect("timezone should parse");
|
|
let now_utc = "2026-03-31T16:15:00Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let target = wallet_daily_usage_aggregation_target(now_utc, timezone);
|
|
|
|
assert_eq!(target.billing_date.to_string(), "2026-03-31");
|
|
assert_eq!(target.billing_timezone, "Asia/Shanghai");
|
|
assert_eq!(
|
|
target.window_start_utc.to_rfc3339(),
|
|
"2026-03-30T16:00:00+00:00"
|
|
);
|
|
assert_eq!(
|
|
target.window_end_utc.to_rfc3339(),
|
|
"2026-03-31T16:00:00+00:00"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn provider_checkin_schedule_uses_default_for_invalid_value() {
|
|
let data = GatewayDataState::disabled().with_system_config_values_for_tests([(
|
|
"provider_checkin_time".to_string(),
|
|
json!("25:99"),
|
|
)]);
|
|
|
|
let schedule = provider_checkin_schedule(&data)
|
|
.await
|
|
.expect("provider checkin schedule should resolve");
|
|
|
|
assert_eq!(schedule, (1, 5));
|
|
}
|
|
|
|
#[test]
|
|
fn next_provider_checkin_run_aligns_to_same_day_when_before_slot() {
|
|
let timezone: Tz = "Asia/Shanghai".parse().expect("timezone should parse");
|
|
let now_utc = "2026-03-31T16:59:00Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let next = next_daily_run_after(now_utc, timezone, 1, 5);
|
|
|
|
assert_eq!(next.to_rfc3339(), "2026-03-31T17:05:00+00:00");
|
|
}
|
|
|
|
#[test]
|
|
fn next_provider_checkin_run_rolls_to_next_day_after_slot() {
|
|
let timezone: Tz = "Asia/Shanghai".parse().expect("timezone should parse");
|
|
let now_utc = "2026-03-31T17:05:01Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let next = next_daily_run_after(now_utc, timezone, 1, 5);
|
|
|
|
assert_eq!(next.to_rfc3339(), "2026-04-01T17:05:00+00:00");
|
|
}
|
|
|
|
#[test]
|
|
fn next_usage_cleanup_run_aligns_to_same_day_when_before_slot() {
|
|
let timezone: Tz = "Asia/Shanghai".parse().expect("timezone should parse");
|
|
let now_utc = "2026-03-17T18:59:00Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let next = next_daily_run_after(now_utc, timezone, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE);
|
|
|
|
assert_eq!(next.to_rfc3339(), "2026-03-17T19:00:00+00:00");
|
|
}
|
|
|
|
#[test]
|
|
fn next_usage_cleanup_run_rolls_to_next_day_after_slot() {
|
|
let timezone: Tz = "Asia/Shanghai".parse().expect("timezone should parse");
|
|
let now_utc = "2026-03-17T19:00:01Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let next = next_daily_run_after(now_utc, timezone, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE);
|
|
|
|
assert_eq!(next.to_rfc3339(), "2026-03-18T19:00:00+00:00");
|
|
}
|
|
|
|
#[test]
|
|
fn next_wallet_daily_usage_aggregation_run_aligns_to_same_day_when_before_slot() {
|
|
let timezone: Tz = "Asia/Shanghai".parse().expect("timezone should parse");
|
|
let now_utc = "2026-03-31T16:09:00Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let next = next_daily_run_after(
|
|
now_utc,
|
|
timezone,
|
|
WALLET_DAILY_USAGE_AGGREGATION_HOUR,
|
|
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
|
|
);
|
|
|
|
assert_eq!(next.to_rfc3339(), "2026-03-31T16:10:00+00:00");
|
|
}
|
|
|
|
#[test]
|
|
fn next_wallet_daily_usage_aggregation_run_rolls_to_next_day_after_slot() {
|
|
let timezone: Tz = "Asia/Shanghai".parse().expect("timezone should parse");
|
|
let now_utc = "2026-03-31T16:10:01Z"
|
|
.parse::<DateTime<Utc>>()
|
|
.expect("timestamp should parse");
|
|
|
|
let next = next_daily_run_after(
|
|
now_utc,
|
|
timezone,
|
|
WALLET_DAILY_USAGE_AGGREGATION_HOUR,
|
|
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
|
|
);
|
|
|
|
assert_eq!(next.to_rfc3339(), "2026-04-01T16:10:00+00:00");
|
|
}
|