mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
feat(proxy): 实现代理节点批量升级回滚、隧道重定向跟随及远程配置管理
核心功能： - 新增代理节点批量升级回滚工作流，支持分批升级、健康探针、跳过/重试/取消等操作 - proxy 隧道流处理器支持 HTTP 重定向跟随（最多 10 跳），区分 307/308 可重播与不可重播请求体 - proxy 协议新增 follow_redirects / http1_only 字段，网关侧同步支持 - 新增代理节点远端配置变更接口（名称、允许端口、调度状态、升级目标等） - 新增代理节点注册/反注册/心跳的 Admin API，及节点过期清理维护任务 - gateway 隧道 owner-relay 支持流式代理大请求体，新增 5 MiB 默认限制 - 新增 ProxyNodeRegistrationMutation / ProxyNodeRemoteConfigMutation 数据类型 - proxy 配置新增重定向重播预算、心跳间隔等参数，TUI 安装向导同步更新 - 前端 ProxyNodes 页面新增批量升级操作面板及滚动进度展示
This commit is contained in:
fawney19
@@ -3,10 +3,18 @@ mod runtime;
|
||||
mod tests;
|
||||
|
||||
pub(crate) use runtime::{
|
||||
perform_provider_checkin_once, spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
|
||||
spawn_gemini_file_mapping_cleanup_worker, spawn_pending_cleanup_worker,
|
||||
spawn_pool_monitor_worker, spawn_provider_checkin_worker,
|
||||
cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts,
|
||||
inspect_proxy_upgrade_rollout, perform_provider_checkin_once,
|
||||
record_proxy_upgrade_traffic_success, restore_proxy_upgrade_rollout_skipped_nodes,
|
||||
retry_proxy_upgrade_rollout_node, skip_proxy_upgrade_rollout_node, spawn_audit_cleanup_worker,
|
||||
spawn_db_maintenance_worker, spawn_gemini_file_mapping_cleanup_worker,
|
||||
spawn_pending_cleanup_worker, spawn_pool_monitor_worker, spawn_provider_checkin_worker,
|
||||
spawn_proxy_node_stale_cleanup_worker, spawn_proxy_upgrade_rollout_worker,
|
||||
spawn_request_candidate_cleanup_worker, spawn_stats_aggregation_worker,
|
||||
spawn_stats_hourly_aggregation_worker, spawn_usage_cleanup_worker,
|
||||
spawn_wallet_daily_usage_aggregation_worker, ProviderCheckinRunSummary,
|
||||
spawn_wallet_daily_usage_aggregation_worker, start_proxy_upgrade_rollout,
|
||||
ProviderCheckinRunSummary, ProxyUpgradeRolloutCancelSummary,
|
||||
ProxyUpgradeRolloutConflictClearSummary, ProxyUpgradeRolloutNodeActionSummary,
|
||||
ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutSkippedRestoreSummary,
|
||||
ProxyUpgradeRolloutStatus, ProxyUpgradeRolloutTrackedNodeState,
|
||||
};
|
||||
|
||||
@@ -29,6 +29,10 @@ mod db_maintenance;
|
||||
mod pending_cleanup;
|
||||
#[path = "runtime/provider_checkin.rs"]
|
||||
mod provider_checkin;
|
||||
#[path = "runtime/proxy_node_staleness.rs"]
|
||||
mod proxy_node_staleness;
|
||||
#[path = "runtime/proxy_upgrade_rollout.rs"]
|
||||
mod proxy_upgrade_rollout;
|
||||
#[path = "runtime/request_candidate_cleanup.rs"]
|
||||
mod request_candidate_cleanup;
|
||||
#[path = "runtime/runners.rs"]
|
||||
@@ -53,6 +57,18 @@ use config::*;
|
||||
use db_maintenance::*;
|
||||
use pending_cleanup::*;
|
||||
pub(crate) use provider_checkin::{perform_provider_checkin_once, ProviderCheckinRunSummary};
|
||||
use proxy_node_staleness::*;
|
||||
use proxy_upgrade_rollout::*;
|
||||
pub(crate) use proxy_upgrade_rollout::{
|
||||
cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts,
|
||||
collect_proxy_upgrade_rollout_probes, inspect_proxy_upgrade_rollout,
|
||||
record_proxy_upgrade_traffic_success, restore_proxy_upgrade_rollout_skipped_nodes,
|
||||
retry_proxy_upgrade_rollout_node, skip_proxy_upgrade_rollout_node, start_proxy_upgrade_rollout,
|
||||
ProxyUpgradeRolloutCancelSummary, ProxyUpgradeRolloutConflictClearSummary,
|
||||
ProxyUpgradeRolloutNodeActionSummary, ProxyUpgradeRolloutPendingProbe,
|
||||
ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutSkippedRestoreSummary,
|
||||
ProxyUpgradeRolloutStatus, ProxyUpgradeRolloutSummary, ProxyUpgradeRolloutTrackedNodeState,
|
||||
};
|
||||
use request_candidate_cleanup::*;
|
||||
use runners::*;
|
||||
use schedule::*;
|
||||
@@ -71,6 +87,10 @@ pub(super) fn postgres_error(
|
||||
const AUDIT_LOG_CLEANUP_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
const GEMINI_FILE_MAPPING_CLEANUP_INTERVAL: Duration = Duration::from_secs(60 * 60);
|
||||
const PENDING_CLEANUP_INTERVAL: Duration = Duration::from_secs(5 * 60);
|
||||
const PROXY_NODE_STALE_SWEEP_INTERVAL: Duration = Duration::from_secs(30);
|
||||
const PROXY_UPGRADE_ROLLOUT_INTERVAL: Duration = Duration::from_secs(15);
|
||||
const PROXY_NODE_STALE_MIN_GRACE_SECS: u64 = 90;
|
||||
const PROXY_NODE_STALE_MISSED_HEARTBEATS: u64 = 3;
|
||||
const POOL_MONITOR_INTERVAL: Duration = Duration::from_secs(5 * 60);
|
||||
const PROVIDER_CHECKIN_CONCURRENCY: usize = 3;
|
||||
const PROVIDER_CHECKIN_DEFAULT_TIME: &str = "01:05";
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_data::repository::proxy_nodes::ProxyNodeTunnelStatusMutation;
|
||||
use aether_data_contracts::DataLayerError;
|
||||
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
use super::{PROXY_NODE_STALE_MIN_GRACE_SECS, PROXY_NODE_STALE_MISSED_HEARTBEATS};
|
||||
|
||||
fn current_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
fn stale_proxy_node_grace_secs(heartbeat_interval: i32) -> u64 {
|
||||
let interval = u64::try_from(heartbeat_interval.max(5)).unwrap_or(5);
|
||||
interval
|
||||
.saturating_mul(PROXY_NODE_STALE_MISSED_HEARTBEATS)
|
||||
.max(PROXY_NODE_STALE_MIN_GRACE_SECS)
|
||||
}
|
||||
|
||||
pub(super) async fn cleanup_stale_proxy_nodes_once(
|
||||
data: &GatewayDataState,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
if !data.has_proxy_node_reader() || !data.has_proxy_node_writer() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let now_unix_secs = current_unix_secs();
|
||||
let nodes = data.list_proxy_nodes().await?;
|
||||
let mut updated = 0usize;
|
||||
|
||||
for node in nodes {
|
||||
if node.is_manual || !node.tunnel_connected {
|
||||
continue;
|
||||
}
|
||||
|
||||
let grace_secs = stale_proxy_node_grace_secs(node.heartbeat_interval);
|
||||
let is_stale = node
|
||||
.last_heartbeat_at_unix_secs
|
||||
.map(|last_seen| last_seen.saturating_add(grace_secs) < now_unix_secs)
|
||||
.unwrap_or(true);
|
||||
if !is_stale {
|
||||
continue;
|
||||
}
|
||||
|
||||
let detail = format!(
|
||||
"[heartbeat_timeout] last_heartbeat_at={} grace_secs={}",
|
||||
node.last_heartbeat_at_unix_secs.unwrap_or(0),
|
||||
grace_secs
|
||||
);
|
||||
let mutation = ProxyNodeTunnelStatusMutation {
|
||||
node_id: node.id.clone(),
|
||||
connected: false,
|
||||
conn_count: 0,
|
||||
detail: Some(detail),
|
||||
observed_at_unix_secs: Some(now_unix_secs),
|
||||
};
|
||||
if data
|
||||
.update_proxy_node_tunnel_status(&mutation)
|
||||
.await?
|
||||
.is_some()
|
||||
{
|
||||
updated = updated.saturating_add(1);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(updated)
|
||||
}
|
||||
1070
apps/aether-gateway/src/maintenance/runtime/proxy_upgrade_rollout.rs
Normal file
1070
apps/aether-gateway/src/maintenance/runtime/proxy_upgrade_rollout.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -1,15 +1,18 @@
|
||||
use aether_data_contracts::DataLayerError;
|
||||
use tracing::info;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
use super::{
|
||||
cleanup_audit_logs_once, cleanup_expired_gemini_file_mappings_once,
|
||||
cleanup_request_candidates_once, cleanup_stale_pending_requests_once,
|
||||
perform_db_maintenance_once, perform_provider_checkin_once, perform_stats_aggregation_once,
|
||||
advance_proxy_upgrade_rollout_once, cleanup_audit_logs_once,
|
||||
cleanup_expired_gemini_file_mappings_once, cleanup_request_candidates_once,
|
||||
cleanup_stale_pending_requests_once, cleanup_stale_proxy_nodes_once,
|
||||
collect_proxy_upgrade_rollout_probes, perform_db_maintenance_once,
|
||||
perform_provider_checkin_once, perform_stats_aggregation_once,
|
||||
perform_stats_hourly_aggregation_once, perform_usage_cleanup_once,
|
||||
perform_wallet_daily_usage_aggregation_once, summarize_postgres_pool,
|
||||
perform_wallet_daily_usage_aggregation_once, record_proxy_upgrade_traffic_success,
|
||||
summarize_postgres_pool,
|
||||
};
|
||||
|
||||
pub(super) async fn run_audit_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
|
||||
@@ -42,6 +45,91 @@ pub(super) async fn run_gemini_file_mapping_cleanup_once(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) async fn run_proxy_node_stale_cleanup_once(
|
||||
data: &GatewayDataState,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let stale_marked = cleanup_stale_proxy_nodes_once(data).await?;
|
||||
if stale_marked > 0 {
|
||||
info!(
|
||||
event_name = "proxy_node_stale_cleanup_completed",
|
||||
log_type = "ops",
|
||||
worker = "proxy_node_stale_cleanup",
|
||||
stale_marked,
|
||||
"gateway marked stale proxy nodes offline"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) async fn run_proxy_upgrade_rollout_once(state: &AppState) -> Result<(), DataLayerError> {
|
||||
let mut summary = advance_proxy_upgrade_rollout_once(&state.data).await?;
|
||||
let probes = collect_proxy_upgrade_rollout_probes(&state.data).await?;
|
||||
let mut probe_recorded = false;
|
||||
for probe in probes {
|
||||
match state
|
||||
.tunnel
|
||||
.probe_node_url(&probe.node_id, &probe.url, probe.timeout_secs)
|
||||
.await
|
||||
{
|
||||
Ok(status) if (200..300).contains(&status) => {
|
||||
let _ = record_proxy_upgrade_traffic_success(&state.data, &probe.node_id).await?;
|
||||
probe_recorded = true;
|
||||
info!(
|
||||
event_name = "proxy_upgrade_rollout_probe_succeeded",
|
||||
log_type = "ops",
|
||||
worker = "proxy_upgrade_rollout",
|
||||
node_id = %probe.node_id,
|
||||
url = %probe.url,
|
||||
status,
|
||||
"gateway confirmed proxy upgrade health probe"
|
||||
);
|
||||
}
|
||||
Ok(status) => {
|
||||
warn!(
|
||||
event_name = "proxy_upgrade_rollout_probe_unhealthy",
|
||||
log_type = "ops",
|
||||
worker = "proxy_upgrade_rollout",
|
||||
node_id = %probe.node_id,
|
||||
url = %probe.url,
|
||||
status,
|
||||
"gateway proxy upgrade health probe returned non-success status"
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "proxy_upgrade_rollout_probe_failed",
|
||||
log_type = "ops",
|
||||
worker = "proxy_upgrade_rollout",
|
||||
node_id = %probe.node_id,
|
||||
url = %probe.url,
|
||||
error = %error,
|
||||
"gateway proxy upgrade health probe failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if probe_recorded {
|
||||
summary = advance_proxy_upgrade_rollout_once(&state.data).await?;
|
||||
}
|
||||
if summary.updated > 0 || !summary.pending_node_ids.is_empty() || !summary.version.is_empty() {
|
||||
info!(
|
||||
event_name = "proxy_upgrade_rollout_checked",
|
||||
log_type = "ops",
|
||||
worker = "proxy_upgrade_rollout",
|
||||
version = %summary.version,
|
||||
updated = summary.updated,
|
||||
blocked = summary.blocked,
|
||||
pending = summary.pending_node_ids.len(),
|
||||
completed = summary.completed,
|
||||
remaining = summary.remaining,
|
||||
rollout_active = summary.rollout_active,
|
||||
"gateway checked proxy upgrade rollout"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) async fn run_db_maintenance_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
|
||||
let summary = perform_db_maintenance_once(data).await?;
|
||||
if summary.attempted > 0 {
|
||||
|
||||
@@ -1,24 +1,35 @@
|
||||
use std::collections::{HashSet, VecDeque};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_data::repository::proxy_nodes::{
|
||||
InMemoryProxyNodeRepository, ProxyNodeHeartbeatMutation, ProxyNodeReadRepository,
|
||||
ProxyNodeWriteRepository, StoredProxyNode,
|
||||
};
|
||||
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::{
|
||||
cleanup_audit_logs_with, next_daily_run_after, next_db_maintenance_run_after,
|
||||
advance_proxy_upgrade_rollout_once, cleanup_audit_logs_with, 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, plan_pending_cleanup_batch,
|
||||
provider_checkin_schedule, run_db_maintenance_with, spawn_audit_cleanup_worker,
|
||||
spawn_db_maintenance_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
|
||||
spawn_provider_checkin_worker, spawn_stats_aggregation_worker,
|
||||
spawn_stats_hourly_aggregation_worker, spawn_usage_cleanup_worker,
|
||||
spawn_wallet_daily_usage_aggregation_worker, stats_aggregation_target_day,
|
||||
provider_checkin_schedule, record_proxy_upgrade_traffic_success, run_db_maintenance_with,
|
||||
run_proxy_upgrade_rollout_once, spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
|
||||
spawn_pending_cleanup_worker, spawn_pool_monitor_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_postgres_pool, usage_cleanup_settings,
|
||||
usage_cleanup_window, wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary,
|
||||
FailedPendingUsageRow, GatewayDataState, StalePendingUsageRow, UsageCleanupSettings,
|
||||
USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
|
||||
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
|
||||
FailedPendingUsageRow, GatewayDataState, ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow,
|
||||
UsageCleanupSettings, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
|
||||
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
@@ -36,11 +47,476 @@ async fn spawn_pending_cleanup_worker_skips_when_postgres_unavailable() {
|
||||
assert!(spawn_pending_cleanup_worker(Arc::new(GatewayDataState::disabled())).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn spawn_proxy_node_stale_cleanup_worker_skips_when_proxy_nodes_unavailable() {
|
||||
assert!(
|
||||
spawn_proxy_node_stale_cleanup_worker(Arc::new(GatewayDataState::disabled())).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_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() {
|
||||
assert!(spawn_pool_monitor_worker(Arc::new(GatewayDataState::disabled())).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),
|
||||
)
|
||||
}
|
||||
|
||||
#[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(),
|
||||
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(),
|
||||
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("proxy-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(),
|
||||
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(),
|
||||
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);
|
||||
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 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(),
|
||||
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 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,
|
||||
)));
|
||||
|
||||
let responder_hub = tunnel_state.hub.clone();
|
||||
let responder = tokio::spawn(async move {
|
||||
let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") {
|
||||
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 proxy_rx.recv().await.expect("body frame should arrive") {
|
||||
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_postgres_unavailable() {
|
||||
assert!(spawn_stats_aggregation_worker(Arc::new(GatewayDataState::disabled())).is_none());
|
||||
|
||||
@@ -11,13 +11,14 @@ use super::{
|
||||
duration_until_next_stats_aggregation_run, duration_until_next_stats_hourly_aggregation_run,
|
||||
maintenance_timezone, parse_hhmm_time, provider_checkin_schedule, run_audit_cleanup_once,
|
||||
run_db_maintenance_once, run_gemini_file_mapping_cleanup_once, run_pending_cleanup_once,
|
||||
run_pool_monitor_once, run_provider_checkin_once, run_request_candidate_cleanup_once,
|
||||
run_stats_aggregation_once, run_stats_hourly_aggregation_once, run_usage_cleanup_once,
|
||||
run_pool_monitor_once, run_provider_checkin_once, run_proxy_node_stale_cleanup_once,
|
||||
run_proxy_upgrade_rollout_once, run_request_candidate_cleanup_once, run_stats_aggregation_once,
|
||||
run_stats_hourly_aggregation_once, run_usage_cleanup_once,
|
||||
run_wallet_daily_usage_aggregation_once, AUDIT_LOG_CLEANUP_INTERVAL,
|
||||
GEMINI_FILE_MAPPING_CLEANUP_INTERVAL, PENDING_CLEANUP_INTERVAL, POOL_MONITOR_INTERVAL,
|
||||
PROVIDER_CHECKIN_DEFAULT_TIME, REQUEST_CANDIDATE_CLEANUP_INTERVAL, USAGE_CLEANUP_HOUR,
|
||||
USAGE_CLEANUP_MINUTE, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
|
||||
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
|
||||
PROVIDER_CHECKIN_DEFAULT_TIME, PROXY_NODE_STALE_SWEEP_INTERVAL, PROXY_UPGRADE_ROLLOUT_INTERVAL,
|
||||
REQUEST_CANDIDATE_CLEANUP_INTERVAL, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
|
||||
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
|
||||
};
|
||||
|
||||
fn log_maintenance_worker_failure(
|
||||
@@ -227,6 +228,55 @@ pub(crate) fn spawn_pending_cleanup_worker(
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_proxy_node_stale_cleanup_worker(
|
||||
data: Arc<GatewayDataState>,
|
||||
) -> Option<tokio::task::JoinHandle<()>> {
|
||||
if !data.has_proxy_node_reader() || !data.has_proxy_node_writer() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(tokio::spawn(async move {
|
||||
if let Err(err) = run_proxy_node_stale_cleanup_once(&data).await {
|
||||
log_maintenance_worker_failure("proxy_node_stale_cleanup", "startup", &err);
|
||||
}
|
||||
let mut interval = tokio::time::interval(PROXY_NODE_STALE_SWEEP_INTERVAL);
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
|
||||
interval.tick().await;
|
||||
loop {
|
||||
interval.tick().await;
|
||||
if let Err(err) = run_proxy_node_stale_cleanup_once(&data).await {
|
||||
log_maintenance_worker_failure("proxy_node_stale_cleanup", "tick", &err);
|
||||
}
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_proxy_upgrade_rollout_worker(
|
||||
state: AppState,
|
||||
) -> Option<tokio::task::JoinHandle<()>> {
|
||||
if !state.data.has_proxy_node_reader()
|
||||
|| !state.data.has_proxy_node_writer()
|
||||
|| !state.data.has_system_config_store()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(tokio::spawn(async move {
|
||||
if let Err(err) = run_proxy_upgrade_rollout_once(&state).await {
|
||||
log_maintenance_worker_failure("proxy_upgrade_rollout", "startup", &err);
|
||||
}
|
||||
let mut interval = tokio::time::interval(PROXY_UPGRADE_ROLLOUT_INTERVAL);
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
|
||||
interval.tick().await;
|
||||
loop {
|
||||
interval.tick().await;
|
||||
if let Err(err) = run_proxy_upgrade_rollout_once(&state).await {
|
||||
log_maintenance_worker_failure("proxy_upgrade_rollout", "tick", &err);
|
||||
}
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_pool_monitor_worker(
|
||||
data: Arc<GatewayDataState>,
|
||||
) -> Option<tokio::task::JoinHandle<()>> {
|
||||
|
||||
Reference in New Issue
Block a user