2026-02-08 13:33:08 +08:00
//! Application lifecycle: initialization, task orchestration, and shutdown.
2026-04-14 22:51:02 +08:00
use std ::collections ::BTreeMap ;
use std ::sync ::atomic ::{ AtomicU64 , Ordering };
2026-02-08 13:33:08 +08:00
use std ::sync ::{ Arc , RwLock };
2026-04-14 22:51:02 +08:00
use std ::time ::{ Duration , Instant };
2026-02-08 13:33:08 +08:00
2026-04-12 16:02:38 +08:00
use aether_http ::{ jittered_delay_for_retry , HttpRetryConfig };
2026-05-19 23:49:36 +08:00
use aether_runtime ::{
init_reloadable_service_tracing , prometheus_response , wait_for_shutdown_signal , ConcurrencyGate ,
};
2026-05-08 00:18:12 +08:00
use aether_runtime_state ::{ RedisClientConfig , RuntimeSemaphoreConfig , RuntimeState };
2026-02-28 01:32:28 +08:00
use arc_swap ::ArcSwap ;
2026-05-19 23:49:36 +08:00
use axum ::extract ::State as AxumState ;
use axum ::routing ::get ;
use axum ::{ Json , Router };
2026-02-28 01:32:28 +08:00
use tokio ::sync ::{ watch , Mutex };
2026-04-12 16:02:38 +08:00
use tokio ::task ::JoinHandle ;
2026-02-25 21:59:29 +08:00
use tracing ::{ error , info , warn };
2026-02-08 13:33:08 +08:00
2026-05-22 09:47:28 +08:00
use crate ::config ::{
2026-09-04 03:45:52 +08:00
aether_url_for_log , effective_tunnel_security , validate_tunnel_encryption_key , Config ,
ServerEntry , TunnelPoolSizing ,
2026-05-22 09:47:28 +08:00
};
2026-02-08 13:33:08 +08:00
use crate ::net ;
use crate ::registration ::client ::AetherClient ;
use crate ::runtime ::{ self , DynamicConfig };
2026-05-20 01:02:01 +08:00
use crate ::state ::{ AppState , ServerContext , TunnelMetrics , TunnelRequestMetrics };
2026-03-06 17:55:14 +08:00
use crate ::upstream_client ;
2026-02-25 21:59:29 +08:00
use crate ::{ hardware , target_filter , tunnel };
2026-02-08 13:33:08 +08:00
2026-04-12 16:02:38 +08:00
type TaskHandles = Arc < Mutex < Vec < JoinHandle < () >>>> ;
2026-05-12 18:19:29 +08:00
const AUTO_STREAM_LIMIT_MIN : usize = 16 ;
2026-05-19 23:49:36 +08:00
// Keep the automatic fallback large enough for real load while still
// protecting tiny nodes from overcommitting by default.
const AUTO_STREAM_LIMIT_MAX : usize = 2048 ;
// Bias toward throughput: let a 16-core class box auto-land near 2k streams,
// then rely on FD / memory estimates and the hard cap to keep smaller hosts safe.
const AUTO_STREAM_LIMIT_PER_CPU : u64 = 128 ;
2026-05-12 18:19:29 +08:00
const AUTO_STREAM_LIMIT_MEMORY_MB_PER_STREAM : u64 = 4 ;
2026-05-19 23:49:36 +08:00
const AUTO_STREAM_LIMIT_ESTIMATED_DIVISOR : u64 = 12 ;
2026-05-12 18:19:29 +08:00
2026-04-14 22:51:02 +08:00
#[derive(Debug, Clone, Copy)]
struct TunnelPoolPolicy {
min_connections : usize ,
max_connections : usize ,
max_streams_per_tunnel : usize ,
scale_check_interval : Duration ,
scale_up_threshold_percent : u32 ,
scale_down_threshold_percent : u32 ,
scale_down_grace : Duration ,
}
impl TunnelPoolPolicy {
fn from_config ( config : & Config , sizing : TunnelPoolSizing ) -> Self {
Self {
min_connections : sizing . initial_connections . max ( 1 ) as usize ,
max_connections : sizing
. max_connections
. max ( sizing . initial_connections )
. max ( 1 ) as usize ,
max_streams_per_tunnel : config . tunnel_max_streams . unwrap_or ( 128 ). max ( 1 ) as usize ,
scale_check_interval : Duration ::from_millis ( config . tunnel_scale_check_interval_ms ),
scale_up_threshold_percent : config . tunnel_scale_up_threshold_percent ,
scale_down_threshold_percent : config . tunnel_scale_down_threshold_percent ,
scale_down_grace : Duration ::from_secs ( config . tunnel_scale_down_grace_secs ),
}
}
fn scale_up_high_water_mark ( & self ) -> u64 {
occupancy_threshold ( self . max_streams_per_tunnel , self . scale_up_threshold_percent )
}
fn scale_down_low_water_mark ( & self ) -> u64 {
occupancy_threshold (
self . max_streams_per_tunnel ,
self . scale_down_threshold_percent ,
)
}
}
struct ManagedTunnel {
slot_id : usize ,
drain_tx : watch ::Sender < bool > ,
handle : JoinHandle < () > ,
draining : bool ,
}
2026-05-19 23:49:36 +08:00
#[derive(Clone)]
struct DiagnosticsState {
state : Arc < AppState > ,
server_contexts : Arc < Mutex < Vec < Arc < ServerContext >>>> ,
}
2026-02-08 13:33:08 +08:00
/// Run the full application lifecycle after config has been parsed.
2026-02-25 21:59:29 +08:00
pub async fn run ( mut config : Config , servers : Vec < ServerEntry > ) -> anyhow ::Result < () > {
2026-02-28 01:32:28 +08:00
config . validate () ? ;
2026-09-04 03:45:52 +08:00
for ( index , server ) in servers . iter (). enumerate () {
server
. validate ()
. map_err ( | error | anyhow ::anyhow! ( "servers[{index}] invalid: {error}" )) ? ;
}
2026-02-08 13:33:08 +08:00
init_tracing ( & config );
info! (
version = env! ( "CARGO_PKG_VERSION" ),
node_name = % config . node_name ,
2026-02-25 21:59:29 +08:00
server_count = servers . len (),
2026-05-20 01:02:01 +08:00
"aether-tunnel starting (tunnel mode)"
2026-02-08 13:33:08 +08:00
);
2026-05-20 01:02:01 +08:00
if let Some ( proxy_url ) = config . effective_aether_outbound_proxy_url () {
2026-05-08 22:03:17 +08:00
if let Ok ( proxy ) = crate ::egress_proxy ::UpstreamProxyConfig ::parse ( proxy_url ) {
info! (
2026-05-20 01:02:01 +08:00
aether_outbound_proxy_url = % proxy . redacted_url (),
2026-05-08 22:03:17 +08:00
"Aether control and tunnel egress proxy configured"
);
}
}
2026-05-07 15:01:18 +08:00
if let Some ( proxy_url ) = config
. upstream_proxy_url
. as_deref ()
. map ( str ::trim )
. filter ( | value | ! value . is_empty ())
{
if let Ok ( proxy ) = crate ::egress_proxy ::UpstreamProxyConfig ::parse ( proxy_url ) {
info! (
upstream_proxy_url = % proxy . redacted_url (),
2026-09-08 17:44:59 +08:00
upstream_proxy_remote_dns = config . upstream_proxy_remote_dns ,
2026-05-07 15:01:18 +08:00
"provider upstream egress proxy configured"
);
}
}
2026-02-08 13:33:08 +08:00
2026-09-08 17:44:59 +08:00
if config . upstream_proxy_remote_dns {
warn! ( "provider hostname DNS resolution and destination IP access controls are delegated to the trusted upstream proxy" );
}
2026-02-25 21:59:29 +08:00
// Resolve public IP (best-effort for region info)
2026-02-08 13:33:08 +08:00
let public_ip = match & config . public_ip {
Some ( ip ) => ip . clone (),
2026-02-25 21:59:29 +08:00
None => net ::detect_public_ip ()
. await
. unwrap_or_else ( | _ | "0.0.0.0" . to_string ()),
2026-02-08 13:33:08 +08:00
};
// Auto-detect region if not configured
if config . node_region . is_none () {
if let Some ( region ) = net ::detect_region ( & public_ip ). await {
config . node_region = Some ( region );
}
}
2026-02-25 21:59:29 +08:00
// Collect hardware info (once at startup, sent during registration)
2026-02-08 13:33:08 +08:00
let hw_info = hardware ::collect ();
2026-05-12 18:19:29 +08:00
if config . max_in_flight_streams . is_none () {
let auto = auto_max_in_flight_streams ( & hw_info );
config . max_in_flight_streams = Some ( auto );
info! (
max_in_flight_streams = auto ,
"auto-detected max_in_flight_streams from hardware"
);
}
2026-02-25 21:59:29 +08:00
// Auto-detect tunnel_max_streams from hardware if not explicitly set
if config . tunnel_max_streams . is_none () {
2026-05-12 18:19:29 +08:00
let stream_limit = config
. max_in_flight_streams
. unwrap_or_else ( || auto_max_in_flight_streams ( & hw_info ));
let auto = stream_limit . clamp ( 1 , 1024 ) as u32 ;
2026-02-25 21:59:29 +08:00
config . tunnel_max_streams = Some ( auto );
info! (
tunnel_max_streams = auto ,
2026-05-12 18:19:29 +08:00
"auto-detected tunnel_max_streams from stream admission limit"
2026-02-25 21:59:29 +08:00
);
}
2026-04-14 22:51:02 +08:00
let tunnel_pool_sizing = config . resolve_tunnel_pool_sizing ( & hw_info ) ? ;
let tunnel_pool_policy = TunnelPoolPolicy ::from_config ( & config , tunnel_pool_sizing );
info! (
tunnel_connections_initial = tunnel_pool_policy . min_connections ,
tunnel_connections_max = tunnel_pool_policy . max_connections ,
tunnel_max_streams = tunnel_pool_policy . max_streams_per_tunnel ,
2026-06-01 01:36:49 +08:00
tunnel_profile = % config . tunnel_profile ,
tunnel_stream_initial_window_bytes = config . tunnel_stream_initial_window_bytes ,
tunnel_drain_deadline_ms = config . tunnel_drain_deadline_ms ,
2026-04-14 22:51:02 +08:00
scale_check_interval_ms = tunnel_pool_policy . scale_check_interval . as_millis (),
scale_up_threshold_percent = tunnel_pool_policy . scale_up_threshold_percent ,
scale_down_threshold_percent = tunnel_pool_policy . scale_down_threshold_percent ,
scale_down_grace_secs = tunnel_pool_policy . scale_down_grace . as_secs (),
auto_sizing = config . tunnel_connections . is_none (),
"resolved tunnel pool policy"
);
2026-02-25 21:59:29 +08:00
2026-02-11 11:30:29 +08:00
info! (
2026-02-25 21:59:29 +08:00
max_concurrency = hw_info . estimated_max_concurrency ,
"hardware info collected"
2026-02-11 11:30:29 +08:00
);
2026-02-25 21:59:29 +08:00
let dns_cache = Arc ::new ( target_filter ::DnsCache ::new (
2026-02-11 11:30:29 +08:00
Duration ::from_secs ( config . dns_cache_ttl_secs ),
config . dns_cache_capacity ,
));
2026-05-05 22:21:23 +08:00
let config = Arc ::new ( config );
// Build a profile-keyed Hyper client pool for tunnel upstream requests.
let upstream_client_pool =
upstream_client ::UpstreamClientPool ::new ( Arc ::clone ( & config ), Arc ::clone ( & dns_cache ));
2026-05-11 02:59:18 +08:00
let resource_monitor = Arc ::new ( hardware ::RuntimeResourceMonitor ::new ());
2026-02-25 21:59:29 +08:00
2026-02-28 01:32:28 +08:00
// Register with each Aether server and build per-server contexts.
// Wrapped in Arc<Mutex> so retry_failed_registrations can append later.
let server_contexts : Arc < Mutex < Vec < Arc < ServerContext >>>> = Arc ::new ( Mutex ::new ( Vec ::new ()));
let mut failed_entries : Vec < ( String , ServerEntry ) > = Vec ::new ();
2026-02-25 21:59:29 +08:00
for ( i , entry ) in servers . iter (). enumerate () {
let label = if servers . len () == 1 {
"server" . to_string ()
} else {
format! ( "server- {} " , i )
};
let node_name = entry
. node_name
. clone ()
. unwrap_or_else ( || config . node_name . clone ());
let client = Arc ::new ( AetherClient ::new (
2026-02-08 13:33:08 +08:00
& config ,
2026-02-25 21:59:29 +08:00
& entry . aether_url ,
& entry . management_token ,
));
2026-05-22 09:47:28 +08:00
let tunnel_security = effective_tunnel_security (
& entry . aether_url ,
entry . tunnel_security ,
entry . tunnel_encryption_key . as_deref (),
);
if tunnel_security == crate ::config ::TunnelSecurity ::NonTlsRequired {
2026-05-23 14:17:04 +08:00
if entry . aether_url . trim_start (). starts_with ( "http://" ) {
warn! (
server = % label ,
2026-09-04 03:45:52 +08:00
url = % aether_url_for_log ( & entry . aether_url ),
2026-05-23 14:17:04 +08:00
"secure tunnel frame encryption starts after registration; deliver install and registration credentials over HTTPS or another trusted bootstrap channel"
);
}
2026-05-22 09:47:28 +08:00
let key = entry
. tunnel_encryption_key
. as_deref ()
. map ( str ::trim )
. filter ( | value | ! value . is_empty ())
. ok_or_else ( || {
anyhow ::anyhow! ( "tunnel_encryption_key must be set for secure non-TLS tunnel" )
}) ? ;
validate_tunnel_encryption_key ( key ) ? ;
}
2026-02-25 21:59:29 +08:00
match client
2026-05-22 09:47:28 +08:00
. register ( & config , entry , & node_name , & public_ip , Some ( & hw_info ))
2026-02-25 21:59:29 +08:00
. await
{
2026-09-04 03:45:52 +08:00
Ok ( registration ) => {
info! (
server = % label ,
node_id = % registration . node_id ,
tunnel_generation = % registration . tunnel_generation ,
url = % aether_url_for_log ( & entry . aether_url ),
node_name = % node_name ,
"registered"
);
2026-04-12 16:02:38 +08:00
server_contexts . lock (). await . push ( build_server_context (
2026-09-04 03:45:52 +08:00
& config ,
& label ,
entry ,
client ,
& node_name ,
registration . node_id ,
registration . tunnel_generation ,
2026-04-12 16:02:38 +08:00
));
2026-02-25 21:59:29 +08:00
}
Err ( e ) => {
warn! (
server = % label ,
2026-09-04 03:45:52 +08:00
url = % aether_url_for_log ( & entry . aether_url ),
2026-02-25 21:59:29 +08:00
error = % e ,
2026-02-28 01:32:28 +08:00
"registration failed, will retry in background"
2026-02-25 21:59:29 +08:00
);
2026-02-28 01:32:28 +08:00
failed_entries . push (( label , entry . clone ()));
2026-02-25 21:59:29 +08:00
}
}
}
2026-02-08 13:33:08 +08:00
2026-04-12 16:02:38 +08:00
let ctx_count = server_contexts . lock (). await . len ();
if ctx_count == 0 && failed_entries . is_empty () {
anyhow ::bail! ( "no servers configured" );
}
if ctx_count == 0 {
warn! (
failed_servers = failed_entries . len (),
"no servers registered successfully at startup; continuing with background recovery"
);
2026-02-25 21:59:29 +08:00
}
2026-02-08 13:33:08 +08:00
// Build shared application state
2026-02-27 21:25:42 +08:00
let tunnel_tls_config = Arc ::new ( crate ::tunnel ::client ::build_tls_config ());
2026-03-24 15:12:56 +08:00
let mut state = AppState {
2026-05-05 22:21:23 +08:00
config ,
2026-02-11 11:30:29 +08:00
dns_cache ,
2026-05-05 22:21:23 +08:00
upstream_client_pool ,
2026-02-27 21:25:42 +08:00
tunnel_tls_config ,
2026-05-11 02:59:18 +08:00
resource_monitor ,
2026-03-24 15:12:56 +08:00
stream_gate : None ,
distributed_stream_gate : None ,
};
if let Some ( limit ) = state . config . max_in_flight_streams {
state = state
2026-05-20 01:02:01 +08:00
. with_stream_concurrency_gate ( Arc ::new ( ConcurrencyGate ::new ( "tunnel_streams" , limit )));
2026-03-24 15:12:56 +08:00
}
if let Some ( limit ) = state . config . distributed_stream_limit {
let redis_url = state
. config
. distributed_stream_redis_url
. clone ()
. expect ( "distributed stream redis url should be validated" );
2026-05-08 00:18:12 +08:00
let runtime = RuntimeState ::redis (
RedisClientConfig {
2026-03-24 15:12:56 +08:00
url : redis_url ,
key_prefix : state . config . distributed_stream_redis_key_prefix . clone (),
2026-05-08 00:18:12 +08:00
},
Some ( state . config . distributed_stream_command_timeout_ms ),
)
. await ? ;
let distributed_gate = runtime . semaphore (
2026-05-20 01:02:01 +08:00
"tunnel_streams_distributed" ,
2026-05-08 00:18:12 +08:00
limit ,
RuntimeSemaphoreConfig {
2026-03-24 15:12:56 +08:00
lease_ttl_ms : state . config . distributed_stream_lease_ttl_ms ,
renew_interval_ms : state . config . distributed_stream_renew_interval_ms ,
command_timeout_ms : Some ( state . config . distributed_stream_command_timeout_ms ),
},
) ? ;
state = state . with_distributed_stream_concurrency_gate ( Arc ::new ( distributed_gate ));
}
let state = Arc ::new ( state );
2026-02-08 13:33:08 +08:00
// Shutdown signal channel
let ( shutdown_tx , shutdown_rx ) = watch ::channel ( false );
2026-05-19 23:49:36 +08:00
let diagnostics_handle = if let Some ( bind_addr ) = state . config . diagnostics_bind {
let listener = tokio ::net ::TcpListener ::bind ( bind_addr ). await ? ;
Some ( spawn_diagnostics_server (
listener ,
DiagnosticsState {
state : Arc ::clone ( & state ),
server_contexts : Arc ::clone ( & server_contexts ),
},
shutdown_rx . clone (),
) ? )
} else {
None
};
2026-02-25 21:59:29 +08:00
info! (
2026-02-28 01:32:28 +08:00
active_servers = server_contexts . lock (). await . len (),
2026-02-25 21:59:29 +08:00
"running in tunnel mode"
);
2026-04-14 22:51:02 +08:00
// Spawn tunnel pool manager per server.
2026-04-12 16:02:38 +08:00
let tunnel_handles : TaskHandles = Arc ::new ( Mutex ::new ( Vec ::new ()));
let retry_handles : TaskHandles = Arc ::new ( Mutex ::new ( Vec ::new ()));
2026-02-28 01:32:28 +08:00
for server in server_contexts . lock (). await . iter () {
2026-04-14 22:51:02 +08:00
spawn_tunnel_pool_manager (
2026-04-12 16:02:38 +08:00
Arc ::clone ( & state ),
Arc ::clone ( server ),
2026-04-14 22:51:02 +08:00
tunnel_pool_policy ,
2026-04-12 16:02:38 +08:00
shutdown_rx . clone (),
Arc ::clone ( & tunnel_handles ),
)
. await ;
2026-02-08 13:33:08 +08:00
}
2026-02-28 01:32:28 +08:00
// Spawn background retry for failed server registrations
if ! failed_entries . is_empty () {
2026-04-12 16:02:38 +08:00
spawn_registration_recovery_tasks (
Arc ::clone ( & state ),
Arc ::clone ( & server_contexts ),
failed_entries ,
public_ip . clone (),
hw_info . clone (),
2026-04-14 22:51:02 +08:00
tunnel_pool_policy ,
2026-04-12 16:02:38 +08:00
shutdown_rx . clone (),
Arc ::clone ( & tunnel_handles ),
Arc ::clone ( & retry_handles ),
)
. await ;
2026-02-28 01:32:28 +08:00
}
2026-02-25 21:59:29 +08:00
// Wait for shutdown signal
wait_for_shutdown (). await ;
info! ( "shutdown signal received, cleaning up..." );
let _ = shutdown_tx . send ( true );
2026-05-19 23:49:36 +08:00
if let Some ( handle ) = diagnostics_handle {
let _ = handle . await ;
}
2026-04-12 16:02:38 +08:00
await_all_handles ( & retry_handles ). await ;
2026-02-25 21:59:29 +08:00
2026-02-28 01:32:28 +08:00
// Graceful unregister from all servers (including retry-registered ones)
for server in server_contexts . lock (). await . iter () {
2026-02-25 21:59:29 +08:00
let node_id = server . node_id . read (). unwrap (). clone ();
if let Err ( e ) = server . aether_client . unregister ( & node_id ). await {
error! (
server = % server . server_label ,
error = % e ,
"unregister failed during shutdown"
);
}
}
// Wait for all tunnel tasks
2026-04-12 16:02:38 +08:00
await_all_handles ( & tunnel_handles ). await ;
2026-02-08 13:33:08 +08:00
2026-05-20 01:02:01 +08:00
info! ( "aether-tunnel stopped" );
2026-02-08 13:33:08 +08:00
Ok (())
}
2026-05-19 23:49:36 +08:00
fn spawn_diagnostics_server (
listener : tokio ::net ::TcpListener ,
diagnostics_state : DiagnosticsState ,
mut shutdown : watch ::Receiver < bool > ,
) -> std ::io ::Result < JoinHandle < () >> {
let bind_addr = listener . local_addr () ? ;
let app = Router ::new ()
. route ( "/health" , get ( diagnostics_health ))
. route ( "/metrics" , get ( diagnostics_metrics ))
. route ( "/stats" , get ( diagnostics_stats ))
. with_state ( diagnostics_state );
2026-05-20 01:02:01 +08:00
info! ( bind = % bind_addr , "tunnel diagnostics server listening" );
2026-05-19 23:49:36 +08:00
Ok ( tokio ::spawn ( async move {
let graceful_shutdown = async move {
while !* shutdown . borrow () {
if shutdown . changed (). await . is_err () {
break ;
}
}
};
if let Err ( error ) = axum ::serve ( listener , app )
. with_graceful_shutdown ( graceful_shutdown )
. await
{
2026-05-20 01:02:01 +08:00
error! ( error = % error , "tunnel diagnostics server exited with error" );
2026-05-19 23:49:36 +08:00
}
}))
}
async fn diagnostics_health (
AxumState ( diagnostics ) : AxumState < DiagnosticsState > ,
) -> Json < serde_json ::Value > {
let servers = diagnostics . server_contexts . lock (). await . clone ();
let active_connections = servers
. iter ()
. map ( | server | server . active_connections . load ( Ordering ::Acquire ))
. sum ::< u64 > ();
let stream_concurrency = diagnostics
. state
. stream_concurrency_snapshot ()
. map ( concurrency_snapshot_json );
let distributed_stream_concurrency =
distributed_stream_concurrency_json ( & diagnostics . state ). await ;
Json ( serde_json ::json! ({
"status" : "ok" ,
2026-05-20 01:02:01 +08:00
"service" : "aether-tunnel" ,
2026-05-19 23:49:36 +08:00
"version" : env ! ( "CARGO_PKG_VERSION" ),
"protocol_version" : aether_contracts ::tunnel ::CURRENT_TUNNEL_PROTOCOL_VERSION ,
"server_count" : servers . len (),
"active_connections" : active_connections ,
"stream_concurrency" : stream_concurrency ,
"distributed_stream_concurrency" : distributed_stream_concurrency ,
}))
}
async fn diagnostics_metrics (
AxumState ( diagnostics ) : AxumState < DiagnosticsState > ,
) -> impl axum ::response ::IntoResponse {
let mut samples = diagnostics . state . metric_samples (). await ;
let servers = diagnostics . server_contexts . lock (). await . clone ();
for server in servers {
samples . extend ( server . metric_samples ());
}
prometheus_response ( & samples )
}
async fn diagnostics_stats (
AxumState ( diagnostics ) : AxumState < DiagnosticsState > ,
) -> Json < serde_json ::Value > {
let servers = diagnostics . server_contexts . lock (). await . clone ();
let active_connections = servers
. iter ()
. map ( | server | server . active_connections . load ( Ordering ::Acquire ))
. sum ::< u64 > ();
let server_stats = servers
. iter ()
. map ( | server | diagnostics_server_stats ( server ))
. collect ::< Vec < _ >> ();
let stream_concurrency = diagnostics
. state
. stream_concurrency_snapshot ()
. map ( concurrency_snapshot_json );
let distributed_stream_concurrency =
distributed_stream_concurrency_json ( & diagnostics . state ). await ;
Json ( serde_json ::json! ({
"status" : "ok" ,
2026-05-20 01:02:01 +08:00
"service" : "aether-tunnel" ,
2026-05-19 23:49:36 +08:00
"version" : env ! ( "CARGO_PKG_VERSION" ),
"protocol_version" : aether_contracts ::tunnel ::CURRENT_TUNNEL_PROTOCOL_VERSION ,
"capacities" : {
"max_concurrent_connections" : diagnostics . state . config . max_concurrent_connections ,
"max_in_flight_streams" : diagnostics . state . config . max_in_flight_streams ,
"distributed_stream_limit" : diagnostics . state . config . distributed_stream_limit ,
"tunnel_max_streams" : diagnostics . state . config . tunnel_max_streams ,
2026-06-01 01:36:49 +08:00
"tunnel_profile" : diagnostics . state . config . tunnel_profile . to_string (),
"tunnel_stream_initial_window_bytes" : diagnostics . state . config . tunnel_stream_initial_window_bytes ,
"tunnel_drain_deadline_ms" : diagnostics . state . config . tunnel_drain_deadline_ms ,
2026-05-19 23:49:36 +08:00
"tunnel_connections" : diagnostics . state . config . tunnel_connections ,
"tunnel_connections_max" : diagnostics . state . config . tunnel_connections_max ,
"diagnostics_bind" : diagnostics . state . config . diagnostics_bind . map ( | addr | addr . to_string ()),
},
"server_count" : servers . len (),
"active_connections" : active_connections ,
"stream_concurrency" : stream_concurrency ,
"distributed_stream_concurrency" : distributed_stream_concurrency ,
"resource_usage" : diagnostics . state . resource_monitor . snapshot (),
"servers" : server_stats ,
}))
}
fn diagnostics_server_stats ( server : & ServerContext ) -> serde_json ::Value {
let node_id = server . node_id . read (). unwrap (). clone ();
let dynamic = server . dynamic . load ();
serde_json ::json! ({
"server" : server . server_label . clone (),
"node_id" : node_id ,
"node_name" : dynamic . node_name . clone (),
"active_connections" : server . active_connections . load ( Ordering ::Acquire ),
2026-05-20 01:02:01 +08:00
"request_metrics" : server . metrics . snapshot (),
2026-05-19 23:49:36 +08:00
"tunnel_metrics" : server . tunnel_metrics . snapshot (),
"recent_tunnel_errors" : server . tunnel_metrics . recent_errors ( 16 ),
})
}
fn concurrency_snapshot_json ( snapshot : aether_runtime ::ConcurrencySnapshot ) -> serde_json ::Value {
serde_json ::json! ({
"limit" : snapshot . limit ,
"in_flight" : snapshot . in_flight ,
"available_permits" : snapshot . available_permits ,
"high_watermark" : snapshot . high_watermark ,
"rejected_total" : snapshot . rejected ,
})
}
fn runtime_semaphore_snapshot_json (
snapshot : aether_runtime_state ::RuntimeSemaphoreSnapshot ,
) -> serde_json ::Value {
serde_json ::json! ({
"limit" : snapshot . limit ,
"in_flight" : snapshot . in_flight ,
"available_permits" : snapshot . available_permits ,
"high_watermark" : snapshot . high_watermark ,
"rejected_total" : snapshot . rejected ,
})
}
async fn distributed_stream_concurrency_json ( state : & AppState ) -> Option < serde_json ::Value > {
match state . distributed_stream_concurrency_snapshot (). await {
Ok ( Some ( snapshot )) => Some ( runtime_semaphore_snapshot_json ( snapshot )),
Ok ( None ) => None ,
Err ( error ) => Some ( serde_json ::json! ({ "error" : error . to_string () })),
}
}
2026-04-12 16:19:23 +08:00
#[allow(clippy::too_many_arguments)]
2026-04-12 16:02:38 +08:00
async fn spawn_registration_recovery_tasks (
2026-02-28 01:32:28 +08:00
state : Arc < AppState > ,
server_contexts : Arc < Mutex < Vec < Arc < ServerContext >>>> ,
failed : Vec < ( String , ServerEntry ) > ,
public_ip : String ,
hw_info : crate ::hardware ::HardwareInfo ,
2026-04-14 22:51:02 +08:00
tunnel_pool_policy : TunnelPoolPolicy ,
2026-04-12 16:02:38 +08:00
shutdown : watch ::Receiver < bool > ,
tunnel_handles : TaskHandles ,
retry_handles : TaskHandles ,
2026-02-28 01:32:28 +08:00
) {
2026-04-12 16:02:38 +08:00
let mut handles = Vec ::with_capacity ( failed . len ());
for ( label , entry ) in failed {
let retry_state = Arc ::clone ( & state );
let retry_contexts = Arc ::clone ( & server_contexts );
let retry_public_ip = public_ip . clone ();
let retry_hw_info = hw_info . clone ();
let retry_shutdown = shutdown . clone ();
let retry_tunnels = Arc ::clone ( & tunnel_handles );
handles . push ( tokio ::spawn ( async move {
retry_failed_registration (
retry_state ,
retry_contexts ,
label ,
entry ,
retry_public_ip ,
retry_hw_info ,
2026-04-14 22:51:02 +08:00
tunnel_pool_policy ,
2026-04-12 16:02:38 +08:00
retry_shutdown ,
retry_tunnels ,
)
. await ;
}));
}
retry_handles . lock (). await . extend ( handles );
}
2026-02-28 01:32:28 +08:00
2026-04-12 16:02:38 +08:00
/// Background task that retries registration for a single server until either
/// registration succeeds or shutdown is requested.
2026-04-12 16:19:23 +08:00
#[allow(clippy::too_many_arguments)]
2026-04-12 16:02:38 +08:00
async fn retry_failed_registration (
state : Arc < AppState > ,
server_contexts : Arc < Mutex < Vec < Arc < ServerContext >>>> ,
label : String ,
entry : ServerEntry ,
public_ip : String ,
hw_info : crate ::hardware ::HardwareInfo ,
2026-04-14 22:51:02 +08:00
tunnel_pool_policy : TunnelPoolPolicy ,
2026-04-12 16:02:38 +08:00
mut shutdown : watch ::Receiver < bool > ,
tunnel_handles : TaskHandles ,
) {
let node_name = entry
. node_name
. clone ()
. unwrap_or_else ( || state . config . node_name . clone ());
let client = Arc ::new ( AetherClient ::new (
& state . config ,
& entry . aether_url ,
& entry . management_token ,
));
let retry_policy = registration_retry_policy ( & state . config );
let mut attempt = 0 u32 ;
2026-02-28 01:32:28 +08:00
2026-04-12 16:02:38 +08:00
loop {
if * shutdown . borrow () {
info! ( server = % label , "shutdown during registration retry" );
return ;
2026-02-28 01:32:28 +08:00
}
2026-04-12 16:02:38 +08:00
attempt = attempt . saturating_add ( 1 );
let delay = jittered_delay_for_retry ( retry_policy , attempt . saturating_sub ( 1 ));
tokio ::select! {
_ = tokio ::time ::sleep ( delay ) => {}
_ = shutdown . changed () => {
info! ( server = % label , "shutdown during registration retry" );
return ;
}
}
match client
2026-05-22 09:47:28 +08:00
. register (
& state . config ,
& entry ,
& node_name ,
& public_ip ,
Some ( & hw_info ),
)
2026-04-12 16:02:38 +08:00
. await
{
2026-09-04 03:45:52 +08:00
Ok ( registration ) => {
info! ( server = % label , node_id = % registration . node_id , tunnel_generation = % registration . tunnel_generation , attempt , "registration retry succeeded" );
2026-04-12 16:02:38 +08:00
let server = build_server_context (
& state . config ,
& label ,
& entry ,
client ,
& node_name ,
2026-09-04 03:45:52 +08:00
registration . node_id ,
registration . tunnel_generation ,
2026-04-12 16:02:38 +08:00
);
server_contexts . lock (). await . push ( Arc ::clone ( & server ));
2026-04-14 22:51:02 +08:00
spawn_tunnel_pool_manager (
2026-04-12 16:02:38 +08:00
Arc ::clone ( & state ),
server ,
2026-04-14 22:51:02 +08:00
tunnel_pool_policy ,
2026-04-12 16:02:38 +08:00
shutdown ,
tunnel_handles ,
)
. await ;
return ;
}
Err ( e ) => {
let next_delay =
jittered_delay_for_retry ( retry_policy , attempt . min ( u32 ::MAX . saturating_sub ( 1 )));
warn! (
server = % label ,
attempt ,
next_delay_ms = next_delay . as_millis (),
error = % e ,
"registration retry failed"
);
}
}
}
}
fn registration_retry_policy ( config : & Config ) -> HttpRetryConfig {
HttpRetryConfig {
max_attempts : u32 ::MAX ,
base_delay_ms : config . aether_retry_base_delay_ms ,
max_delay_ms : config . aether_retry_max_delay_ms ,
}
. normalized ()
}
2026-05-12 18:19:29 +08:00
fn auto_max_in_flight_streams ( hw_info : & crate ::hardware ::HardwareInfo ) -> usize {
let by_cpu = u64 ::from ( hw_info . cpu_cores . max ( 1 )). saturating_mul ( AUTO_STREAM_LIMIT_PER_CPU );
let by_memory = hw_info
. total_memory_mb
. max ( AUTO_STREAM_LIMIT_MEMORY_MB_PER_STREAM )
/ AUTO_STREAM_LIMIT_MEMORY_MB_PER_STREAM ;
let by_estimate = hw_info
. estimated_max_concurrency
. max ( AUTO_STREAM_LIMIT_ESTIMATED_DIVISOR )
/ AUTO_STREAM_LIMIT_ESTIMATED_DIVISOR ;
let raw = by_cpu . min ( by_memory ). min ( by_estimate ). max ( 1 );
usize ::try_from ( raw )
. unwrap_or ( AUTO_STREAM_LIMIT_MAX )
. clamp ( AUTO_STREAM_LIMIT_MIN , AUTO_STREAM_LIMIT_MAX )
}
2026-04-12 16:02:38 +08:00
fn build_server_context (
config : & Config ,
label : & str ,
entry : & ServerEntry ,
client : Arc < AetherClient > ,
node_name : & str ,
node_id : String ,
2026-09-04 03:45:52 +08:00
tunnel_generation : String ,
2026-04-12 16:02:38 +08:00
) -> Arc < ServerContext > {
let mut dynamic = DynamicConfig ::from_config ( config );
dynamic . node_name = node_name . to_string ();
Arc ::new ( ServerContext {
server_label : label . to_string (),
aether_url : entry . aether_url . clone (),
management_token : entry . management_token . clone (),
2026-05-22 09:47:28 +08:00
tunnel_security : effective_tunnel_security (
& entry . aether_url ,
entry . tunnel_security ,
entry . tunnel_encryption_key . as_deref (),
),
tunnel_encryption_key : entry . tunnel_encryption_key . clone (),
2026-04-12 16:02:38 +08:00
node_name : node_name . to_string (),
node_id : Arc ::new ( RwLock ::new ( node_id )),
2026-09-04 03:45:52 +08:00
tunnel_generation ,
2026-04-12 16:02:38 +08:00
aether_client : client ,
dynamic : Arc ::new ( ArcSwap ::from_pointee ( dynamic )),
active_connections : Arc ::new ( AtomicU64 ::new ( 0 )),
2026-05-20 01:02:01 +08:00
metrics : Arc ::new ( TunnelRequestMetrics ::new ()),
2026-05-08 22:03:17 +08:00
tunnel_metrics : Arc ::new ( TunnelMetrics ::new ()),
2026-04-12 16:02:38 +08:00
})
}
2026-04-14 22:51:02 +08:00
async fn spawn_tunnel_pool_manager (
2026-04-12 16:02:38 +08:00
state : Arc < AppState > ,
server : Arc < ServerContext > ,
2026-04-14 22:51:02 +08:00
policy : TunnelPoolPolicy ,
2026-04-12 16:02:38 +08:00
shutdown : watch ::Receiver < bool > ,
tunnel_handles : TaskHandles ,
) {
2026-04-14 22:51:02 +08:00
let handle = tokio ::spawn ( async move {
run_tunnel_pool_manager ( state , server , policy , shutdown ). await ;
});
tunnel_handles . lock (). await . push ( handle );
}
async fn run_tunnel_pool_manager (
state : Arc < AppState > ,
server : Arc < ServerContext > ,
policy : TunnelPoolPolicy ,
mut shutdown : watch ::Receiver < bool > ,
) {
let mut tunnels = BTreeMap ::< usize , ManagedTunnel > ::new ();
ensure_tunnel_capacity (
& mut tunnels ,
policy . min_connections ,
& policy ,
& state ,
& server ,
& shutdown ,
);
let mut ticker = tokio ::time ::interval ( policy . scale_check_interval );
ticker . tick (). await ;
let mut low_load_since : Option < Instant > = None ;
loop {
tokio ::select! {
_ = shutdown . changed () => {
info! ( server = % server . server_label , "tunnel pool manager shutting down" );
break ;
}
_ = ticker . tick () => {
reap_finished_tunnels ( & mut tunnels ). await ;
let available = tunnels . values (). filter ( | tunnel | ! tunnel . draining ). count ();
if available < policy . min_connections {
ensure_tunnel_capacity (
& mut tunnels ,
policy . min_connections ,
& policy ,
& state ,
& server ,
& shutdown ,
);
low_load_since = None ;
continue ;
}
let active_connections = server . active_connections . load ( Ordering ::Acquire );
let desired_connections = desired_tunnel_connections ( active_connections , & policy );
if desired_connections > available {
ensure_tunnel_capacity (
& mut tunnels ,
desired_connections ,
& policy ,
& state ,
& server ,
& shutdown ,
);
info! (
server = % server . server_label ,
active_connections ,
available_connections = available ,
target_connections = desired_connections ,
"scaled tunnel pool up"
);
low_load_since = None ;
continue ;
}
if should_scale_down ( active_connections , available , & policy ) {
match low_load_since {
Some ( since ) if since . elapsed () >= policy . scale_down_grace => {
if request_tunnel_drain ( & mut tunnels , policy . min_connections ) {
info! (
server = % server . server_label ,
active_connections ,
available_connections = available ,
"requested tunnel drain for scale-down"
);
}
low_load_since = None ;
}
None => {
low_load_since = Some ( Instant ::now ());
}
Some ( _ ) => {}
}
} else {
low_load_since = None ;
}
}
}
2026-04-12 16:02:38 +08:00
}
2026-04-14 22:51:02 +08:00
for tunnel in tunnels . values_mut () {
let _ = tunnel . drain_tx . send ( true );
tunnel . draining = true ;
}
while ! tunnels . is_empty () {
reap_finished_tunnels ( & mut tunnels ). await ;
if ! tunnels . is_empty () {
tokio ::time ::sleep ( Duration ::from_millis ( 50 )). await ;
}
}
}
fn ensure_tunnel_capacity (
tunnels : & mut BTreeMap < usize , ManagedTunnel > ,
target_connections : usize ,
policy : & TunnelPoolPolicy ,
state : & Arc < AppState > ,
server : & Arc < ServerContext > ,
shutdown : & watch ::Receiver < bool > ,
) {
let target_connections = target_connections . min ( policy . max_connections );
while tunnels . values (). filter ( | tunnel | ! tunnel . draining ). count () < target_connections {
let Some ( slot_id ) = next_available_tunnel_slot ( tunnels , policy . max_connections ) else {
break ;
};
tunnels . insert (
slot_id ,
spawn_managed_tunnel (
Arc ::clone ( state ),
Arc ::clone ( server ),
slot_id ,
shutdown . clone (),
),
);
}
}
fn spawn_managed_tunnel (
state : Arc < AppState > ,
server : Arc < ServerContext > ,
slot_id : usize ,
shutdown : watch ::Receiver < bool > ,
) -> ManagedTunnel {
let ( drain_tx , drain_rx ) = watch ::channel ( false );
let handle = tokio ::spawn ( async move {
tunnel ::run ( & state , & server , slot_id , shutdown , drain_rx ). await ;
});
ManagedTunnel {
slot_id ,
drain_tx ,
handle ,
draining : false ,
}
}
async fn reap_finished_tunnels ( tunnels : & mut BTreeMap < usize , ManagedTunnel > ) {
let finished_slots = tunnels
. iter ()
. filter_map ( | ( slot_id , tunnel ) | tunnel . handle . is_finished (). then_some ( * slot_id ))
. collect ::< Vec < _ >> ();
for slot_id in finished_slots {
if let Some ( tunnel ) = tunnels . remove ( & slot_id ) {
let _ = tunnel . handle . await ;
}
}
}
fn next_available_tunnel_slot (
tunnels : & BTreeMap < usize , ManagedTunnel > ,
max_connections : usize ,
) -> Option < usize > {
( 0 .. max_connections ). find ( | slot_id | ! tunnels . contains_key ( slot_id ))
}
fn request_tunnel_drain (
tunnels : & mut BTreeMap < usize , ManagedTunnel > ,
min_connections : usize ,
) -> bool {
let available = tunnels . values (). filter ( | tunnel | ! tunnel . draining ). count ();
if available <= min_connections {
return false ;
}
let Some (( _ , tunnel )) = tunnels
. iter_mut ()
. rev ()
. find ( | ( _ , tunnel ) | tunnel . slot_id != 0 && ! tunnel . draining )
else {
return false ;
};
if tunnel . drain_tx . send ( true ). is_ok () {
tunnel . draining = true ;
return true ;
}
false
}
fn desired_tunnel_connections ( active_connections : u64 , policy : & TunnelPoolPolicy ) -> usize {
let required = div_ceil_u64 ( active_connections . max ( 1 ), policy . scale_up_high_water_mark ());
required . clamp ( policy . min_connections as u64 , policy . max_connections as u64 ) as usize
}
fn should_scale_down (
active_connections : u64 ,
available_connections : usize ,
policy : & TunnelPoolPolicy ,
) -> bool {
if available_connections <= policy . min_connections {
return false ;
}
active_connections
<= ( available_connections as u64 )
. saturating_sub ( 1 )
. saturating_mul ( policy . scale_down_low_water_mark ())
}
fn occupancy_threshold ( max_streams_per_tunnel : usize , percent : u32 ) -> u64 {
div_ceil_u64 (
( max_streams_per_tunnel as u64 ). saturating_mul ( percent as u64 ),
100 ,
)
. max ( 1 )
}
fn div_ceil_u64 ( value : u64 , divisor : u64 ) -> u64 {
if divisor == 0 {
return value ;
}
value . saturating_add ( divisor . saturating_sub ( 1 )) / divisor
2026-04-12 16:02:38 +08:00
}
async fn await_all_handles ( handles : & TaskHandles ) {
let mut pending = handles . lock (). await . drain ( .. ). collect ::< Vec < _ >> ();
while let Some ( handle ) = pending . pop () {
let _ = handle . await ;
2026-02-28 01:32:28 +08:00
}
}
2026-02-08 13:33:08 +08:00
fn init_tracing ( config : & Config ) {
2026-04-05 20:23:16 +08:00
let reloader = init_reloadable_service_tracing (
& config . log_level ,
config
. service_runtime_config ()
2026-05-20 01:02:01 +08:00
. expect ( "tunnel service runtime config should be valid" ),
2026-04-05 20:23:16 +08:00
)
2026-05-20 01:02:01 +08:00
. expect ( "tunnel tracing should initialize" );
2026-03-24 15:12:56 +08:00
runtime ::set_log_reloader ( reloader );
2026-02-08 13:33:08 +08:00
}
async fn wait_for_shutdown () {
2026-03-24 15:12:56 +08:00
wait_for_shutdown_signal ()
. await
. expect ( "failed to install shutdown signal handler" );
2026-02-08 13:33:08 +08:00
}
2026-04-12 16:02:38 +08:00
#[cfg(test)]
mod tests {
use std ::sync ::atomic ::{ AtomicUsize , Ordering };
use std ::sync ::Once ;
use axum ::extract ::State as AxumState ;
use axum ::http ::StatusCode as AxumStatusCode ;
2026-05-19 23:49:36 +08:00
use axum ::routing ::{ get , post };
2026-04-12 16:02:38 +08:00
use axum ::Router ;
use serde_json ::json ;
2026-07-30 01:03:05 +08:00
use crate ::config ::{ TunnelLogDestinationArg , TunnelLogRotationArg };
2026-04-12 16:02:38 +08:00
use crate ::hardware ::HardwareInfo ;
2026-05-20 01:02:01 +08:00
use crate ::state ::AppState as TunnelAppState ;
2026-04-12 16:02:38 +08:00
use crate ::target_filter ::DnsCache ;
use super ::* ;
#[tokio::test]
async fn registration_recovery_survives_all_startup_failures_and_connects_later () {
ensure_rustls_provider ();
let gateway_port = reserve_local_port (). expect ( "gateway port should reserve" );
let gateway_base_url = format! ( "http://127.0.0.1: {gateway_port} " );
let state = sample_state ( sample_config ( & gateway_base_url ));
let server_contexts : Arc < Mutex < Vec < Arc < ServerContext >>>> = Arc ::new ( Mutex ::new ( Vec ::new ()));
let tunnel_handles : TaskHandles = Arc ::new ( Mutex ::new ( Vec ::new ()));
let retry_handles : TaskHandles = Arc ::new ( Mutex ::new ( Vec ::new ()));
let register_hits = Arc ::new ( AtomicUsize ::new ( 0 ));
let failed = vec! [(
"server" . to_string (),
ServerEntry {
aether_url : gateway_base_url . clone (),
management_token : "token" . to_string (),
node_name : Some ( "node-recovery" . to_string ()),
2026-05-21 16:07:27 +08:00
tunnel_security : None ,
tunnel_encryption_key : None ,
2026-04-12 16:02:38 +08:00
},
)];
let ( shutdown_tx , shutdown_rx ) = watch ::channel ( false );
2026-04-14 22:51:02 +08:00
let tunnel_pool_policy =
TunnelPoolPolicy ::from_config ( & state . config , sample_tunnel_pool_sizing ());
2026-04-12 16:02:38 +08:00
spawn_registration_recovery_tasks (
Arc ::clone ( & state ),
Arc ::clone ( & server_contexts ),
failed ,
"127.0.0.1" . to_string (),
sample_hardware_info (),
2026-04-14 22:51:02 +08:00
tunnel_pool_policy ,
2026-04-12 16:02:38 +08:00
shutdown_rx . clone (),
Arc ::clone ( & tunnel_handles ),
Arc ::clone ( & retry_handles ),
)
. await ;
tokio ::time ::sleep ( Duration ::from_millis ( 150 )). await ;
assert! (
server_contexts . lock (). await . is_empty (),
"registration should still be pending while gateway is down"
);
let gateway_handle =
start_fake_gateway_on_port_retry ( gateway_port , Arc ::clone ( & register_hits ))
. await
. expect ( "gateway should start" );
let server = wait_for_registered_server ( & server_contexts ). await ;
assert_eq! ( server . server_label , "server" );
assert_eq! ( server . node_id . read (). unwrap (). as_str (), "node-recovery" );
assert! (
register_hits . load ( Ordering ::SeqCst ) >= 1 ,
"fake control plane should observe at least one register request"
);
let _ = shutdown_tx . send ( true );
await_all_handles ( & retry_handles ). await ;
await_all_handles ( & tunnel_handles ). await ;
gateway_handle . abort ();
}
2026-05-19 23:49:36 +08:00
#[tokio::test]
async fn diagnostics_routes_report_health_metrics_and_stats () {
ensure_rustls_provider ();
let state = sample_state ( sample_config ( "https://aether.example.com" ));
let server = sample_registered_server ( & state , "server" , "node-diagnostics" );
let server_contexts = Arc ::new ( Mutex ::new ( vec! [ server ]));
let router = Router ::new ()
. route ( "/health" , get ( diagnostics_health ))
. route ( "/metrics" , get ( diagnostics_metrics ))
. route ( "/stats" , get ( diagnostics_stats ))
. with_state ( DiagnosticsState {
state : Arc ::clone ( & state ),
server_contexts ,
});
let port = reserve_local_port (). expect ( "diagnostics port should reserve" );
let handle = spawn_router_on_port ( port , router )
. await
. expect ( "diagnostics test server should start" );
let client = reqwest ::Client ::new ();
let base_url = format! ( "http://127.0.0.1: {port} " );
let health : serde_json ::Value = client
. get ( format! ( " {base_url} /health" ))
. send ()
. await
. expect ( "health request should send" )
. error_for_status ()
. expect ( "health response should be success" )
. json ()
. await
. expect ( "health response should parse" );
assert_eq! ( health [ "status" ], "ok" );
2026-05-20 01:02:01 +08:00
assert_eq! ( health [ "service" ], "aether-tunnel" );
2026-05-19 23:49:36 +08:00
assert_eq! ( health [ "server_count" ], 1 );
let metrics = client
. get ( format! ( " {base_url} /metrics" ))
. send ()
. await
. expect ( "metrics request should send" )
. error_for_status ()
. expect ( "metrics response should be success" )
. text ()
. await
. expect ( "metrics response should read" );
2026-05-20 01:02:01 +08:00
assert! ( metrics . contains ( "service_up{service= \" aether-tunnel \" } 1" ));
assert! ( metrics . contains ( "tunnel_active_connections{server= \" server \" } 0" ));
2026-05-19 23:49:36 +08:00
let stats : serde_json ::Value = client
. get ( format! ( " {base_url} /stats" ))
. send ()
. await
. expect ( "stats request should send" )
. error_for_status ()
. expect ( "stats response should be success" )
. json ()
. await
. expect ( "stats response should parse" );
assert_eq! ( stats [ "status" ], "ok" );
2026-06-01 01:36:49 +08:00
assert_eq! (
stats [ "protocol_version" ],
aether_contracts ::tunnel ::CURRENT_TUNNEL_PROTOCOL_VERSION
);
2026-05-19 23:49:36 +08:00
assert_eq! ( stats [ "servers" ][ 0 ][ "node_id" ], "node-diagnostics" );
handle . abort ();
}
2026-04-14 22:51:02 +08:00
#[test]
fn desired_tunnel_connections_expands_when_load_crosses_high_water () {
let policy = TunnelPoolPolicy {
min_connections : 1 ,
max_connections : 6 ,
max_streams_per_tunnel : 1024 ,
scale_check_interval : Duration ::from_secs ( 1 ),
scale_up_threshold_percent : 70 ,
scale_down_threshold_percent : 35 ,
scale_down_grace : Duration ::from_secs ( 15 ),
};
assert_eq! ( desired_tunnel_connections ( 1 , & policy ), 1 );
assert_eq! ( desired_tunnel_connections ( 2_000 , & policy ), 3 );
assert_eq! ( desired_tunnel_connections ( 5_000 , & policy ), 6 );
}
#[test]
fn should_scale_down_requires_load_to_fit_remaining_tunnels () {
let policy = TunnelPoolPolicy {
min_connections : 1 ,
max_connections : 6 ,
max_streams_per_tunnel : 1024 ,
scale_check_interval : Duration ::from_secs ( 1 ),
scale_up_threshold_percent : 70 ,
scale_down_threshold_percent : 35 ,
scale_down_grace : Duration ::from_secs ( 15 ),
};
assert! ( ! should_scale_down ( 800 , 3 , & policy ));
assert! ( should_scale_down ( 600 , 3 , & policy ));
assert! ( ! should_scale_down ( 200 , 1 , & policy ));
}
2026-05-12 18:19:29 +08:00
#[test]
fn auto_stream_limit_is_conservative_for_tiny_nodes () {
let hw = HardwareInfo {
cpu_cores : 1 ,
total_memory_mb : 183 ,
os_info : "test" . to_string (),
fd_limit : 65_535 ,
estimated_max_concurrency : 2_000 ,
};
assert_eq! ( auto_max_in_flight_streams ( & hw ), 45 );
}
2026-05-19 23:49:36 +08:00
#[test]
fn auto_stream_limit_scales_to_high_band_on_mid_size_nodes () {
let hw = HardwareInfo {
cpu_cores : 16 ,
total_memory_mb : 65_536 ,
os_info : "test" . to_string (),
fd_limit : 1_048_576 ,
estimated_max_concurrency : 500_000 ,
};
assert_eq! ( auto_max_in_flight_streams ( & hw ), AUTO_STREAM_LIMIT_MAX );
}
2026-05-12 18:19:29 +08:00
#[test]
fn auto_stream_limit_caps_large_nodes () {
let hw = HardwareInfo {
cpu_cores : 64 ,
total_memory_mb : 262_144 ,
os_info : "test" . to_string (),
fd_limit : 1_048_576 ,
estimated_max_concurrency : 500_000 ,
};
assert_eq! ( auto_max_in_flight_streams ( & hw ), AUTO_STREAM_LIMIT_MAX );
}
2026-04-12 16:02:38 +08:00
async fn wait_for_registered_server (
server_contexts : & Arc < Mutex < Vec < Arc < ServerContext >>>> ,
) -> Arc < ServerContext > {
let deadline = tokio ::time ::Instant ::now () + Duration ::from_secs ( 10 );
loop {
if let Some ( server ) = server_contexts . lock (). await . first (). cloned () {
return server ;
}
assert! (
tokio ::time ::Instant ::now () < deadline ,
"server context did not appear after registration recovery"
);
tokio ::time ::sleep ( Duration ::from_millis ( 25 )). await ;
}
}
async fn start_fake_gateway_on_port_retry (
port : u16 ,
register_hits : Arc < AtomicUsize > ,
) -> Result < tokio ::task ::JoinHandle < () > , std ::io ::Error > {
let mut attempts = 0 usize ;
loop {
match start_fake_gateway_on_port ( port , Arc ::clone ( & register_hits )). await {
Ok ( server ) => return Ok ( server ),
Err ( err ) => {
attempts += 1 ;
if attempts >= 20 {
return Err ( err );
}
tokio ::time ::sleep ( Duration ::from_millis ( 50 )). await ;
}
}
}
}
async fn start_fake_gateway_on_port (
port : u16 ,
register_hits : Arc < AtomicUsize > ,
) -> Result < tokio ::task ::JoinHandle < () > , std ::io ::Error > {
let router = Router ::new ()
. route ( "/api/admin/proxy-nodes/register" , post ( fake_register ))
. with_state ( register_hits );
spawn_router_on_port ( port , router ). await
}
async fn spawn_router_on_port (
port : u16 ,
app : Router ,
) -> Result < tokio ::task ::JoinHandle < () > , std ::io ::Error > {
let listener = tokio ::net ::TcpListener ::bind (( "127.0.0.1" , port )). await ? ;
Ok ( tokio ::spawn ( async move {
axum ::serve (
listener ,
app . into_make_service_with_connect_info ::< std ::net ::SocketAddr > (),
)
. await
. expect ( "gateway test server should run" );
}))
}
fn reserve_local_port () -> Result < u16 , std ::io ::Error > {
let listener = std ::net ::TcpListener ::bind ( "127.0.0.1:0" ) ? ;
let port = listener . local_addr () ? . port ();
drop ( listener );
Ok ( port )
}
async fn fake_register (
AxumState ( register_hits ) : AxumState < Arc < AtomicUsize >> ,
) -> ( AxumStatusCode , axum ::Json < serde_json ::Value > ) {
register_hits . fetch_add ( 1 , Ordering ::SeqCst );
(
AxumStatusCode ::OK ,
2026-09-04 03:45:52 +08:00
axum ::Json ( json! ({
"node_id" : "node-recovery" ,
"tunnel_generation" : "test-generation-recovery"
})),
2026-04-12 16:02:38 +08:00
)
}
fn sample_hardware_info () -> HardwareInfo {
HardwareInfo {
cpu_cores : 2 ,
total_memory_mb : 2048 ,
os_info : "test" . to_string (),
fd_limit : 1024 ,
estimated_max_concurrency : 512 ,
}
}
2026-04-14 22:51:02 +08:00
fn sample_tunnel_pool_sizing () -> TunnelPoolSizing {
TunnelPoolSizing {
initial_connections : 1 ,
max_connections : 1 ,
}
}
2026-05-20 01:02:01 +08:00
fn sample_state ( config : Config ) -> Arc < TunnelAppState > {
2026-04-12 16:02:38 +08:00
let config = Arc ::new ( config );
let dns_cache = Arc ::new ( DnsCache ::new ( Duration ::from_secs ( 60 ), 128 ));
2026-05-05 22:21:23 +08:00
let upstream_client_pool =
upstream_client ::UpstreamClientPool ::new ( Arc ::clone ( & config ), Arc ::clone ( & dns_cache ));
2026-05-20 01:02:01 +08:00
Arc ::new ( TunnelAppState {
2026-04-12 16:02:38 +08:00
config ,
dns_cache ,
2026-05-05 22:21:23 +08:00
upstream_client_pool ,
2026-04-12 16:02:38 +08:00
tunnel_tls_config : Arc ::new ( crate ::tunnel ::client ::build_tls_config ()),
2026-05-11 02:59:18 +08:00
resource_monitor : Arc ::new ( crate ::hardware ::RuntimeResourceMonitor ::new ()),
2026-04-12 16:02:38 +08:00
stream_gate : None ,
distributed_stream_gate : None ,
})
}
2026-05-19 23:49:36 +08:00
fn sample_registered_server (
2026-05-20 01:02:01 +08:00
state : & Arc < TunnelAppState > ,
2026-05-19 23:49:36 +08:00
label : & str ,
node_id : & str ,
) -> Arc < ServerContext > {
let entry = ServerEntry {
aether_url : state . config . aether_url . clone (),
management_token : state . config . management_token . clone (),
node_name : Some ( state . config . node_name . clone ()),
2026-05-21 16:07:27 +08:00
tunnel_security : Some ( state . config . tunnel_security ),
tunnel_encryption_key : state . config . tunnel_encryption_key . clone (),
2026-05-19 23:49:36 +08:00
};
let client = Arc ::new ( AetherClient ::new (
& state . config ,
& state . config . aether_url ,
& state . config . management_token ,
));
build_server_context (
& state . config ,
label ,
& entry ,
client ,
& state . config . node_name ,
node_id . to_string (),
2026-09-04 03:45:52 +08:00
"test-generation-1" . to_string (),
2026-05-19 23:49:36 +08:00
)
}
2026-04-12 16:02:38 +08:00
fn sample_config ( aether_url : & str ) -> Config {
Config {
aether_url : aether_url . to_string (),
management_token : "token" . to_string (),
public_ip : None ,
2026-05-20 01:02:01 +08:00
node_name : "tunnel-test" . to_string (),
2026-05-21 16:07:27 +08:00
tunnel_security : crate ::config ::TunnelSecurity ::Off ,
tunnel_encryption_key : None ,
2026-04-12 16:02:38 +08:00
node_region : None ,
heartbeat_interval : 1 ,
allowed_ports : vec ! [ 80 , 443 ],
2026-04-15 00:31:36 +08:00
allow_private_targets : false ,
2026-04-12 16:02:38 +08:00
aether_request_timeout_secs : 10 ,
aether_connect_timeout_secs : 2 ,
aether_pool_max_idle_per_host : 8 ,
aether_pool_idle_timeout_secs : 90 ,
aether_tcp_keepalive_secs : 60 ,
aether_tcp_nodelay : true ,
aether_http2 : true ,
2026-05-20 01:02:01 +08:00
aether_outbound_proxy_url : None ,
2026-04-12 16:02:38 +08:00
aether_retry_max_attempts : 1 ,
aether_retry_base_delay_ms : 50 ,
aether_retry_max_delay_ms : 100 ,
2026-05-19 23:49:36 +08:00
diagnostics_bind : None ,
2026-04-12 16:02:38 +08:00
max_concurrent_connections : None ,
max_in_flight_streams : None ,
distributed_stream_limit : None ,
distributed_stream_redis_url : None ,
distributed_stream_redis_key_prefix : None ,
distributed_stream_lease_ttl_ms : 30_000 ,
distributed_stream_renew_interval_ms : 10_000 ,
distributed_stream_command_timeout_ms : 1_000 ,
dns_cache_ttl_secs : 60 ,
dns_cache_capacity : 128 ,
upstream_connect_timeout_secs : 30 ,
upstream_pool_max_idle_per_host : 4 ,
upstream_pool_idle_timeout_secs : 60 ,
2026-05-28 20:34:08 +08:00
upstream_client_pool_capacity : crate ::config ::DEFAULT_UPSTREAM_CLIENT_POOL_CAPACITY ,
2026-04-12 16:02:38 +08:00
upstream_tcp_keepalive_secs : 60 ,
upstream_tcp_nodelay : true ,
2026-05-07 15:01:18 +08:00
upstream_proxy_url : None ,
2026-09-08 17:44:59 +08:00
upstream_proxy_remote_dns : false ,
2026-07-30 01:03:05 +08:00
legacy_redirect_replay_budget_bytes_ignored : None ,
2026-05-19 23:49:36 +08:00
emit_proxy_timing_header : true ,
2026-04-12 16:02:38 +08:00
log_level : "info" . to_string (),
2026-05-20 01:02:01 +08:00
log_destination : TunnelLogDestinationArg ::Stdout ,
2026-04-12 16:02:38 +08:00
log_dir : None ,
2026-05-20 01:02:01 +08:00
log_rotation : TunnelLogRotationArg ::Daily ,
2026-04-12 16:02:38 +08:00
log_retention_days : 7 ,
log_max_files : 30 ,
tunnel_reconnect_base_ms : 50 ,
tunnel_reconnect_max_ms : 250 ,
2026-04-14 22:51:02 +08:00
tunnel_ping_interval_ms : 1_000 ,
2026-04-12 16:02:38 +08:00
tunnel_max_streams : Some ( 8 ),
2026-06-01 01:36:49 +08:00
tunnel_profile : crate ::config ::TunnelProfileArg ::Lite ,
tunnel_stream_initial_window_bytes :
crate ::config ::DEFAULT_TUNNEL_STREAM_INITIAL_WINDOW_BYTES ,
tunnel_drain_deadline_ms : crate ::config ::DEFAULT_TUNNEL_DRAIN_DEADLINE_MS ,
2026-04-14 22:51:02 +08:00
tunnel_connect_timeout_ms : 2_000 ,
2026-05-20 14:42:52 +08:00
tunnel_ipv4_only : false ,
tunnel_ipv6_only : false ,
2026-04-12 16:02:38 +08:00
tunnel_tcp_keepalive_secs : 30 ,
tunnel_tcp_nodelay : true ,
2026-04-14 22:51:02 +08:00
tunnel_stale_timeout_ms : 5_000 ,
tunnel_connections : Some ( 1 ),
tunnel_connections_max : Some ( 1 ),
tunnel_scale_check_interval_ms : 1_000 ,
tunnel_scale_up_threshold_percent : 70 ,
tunnel_scale_down_threshold_percent : 35 ,
tunnel_scale_down_grace_secs : 15 ,
2026-04-12 16:02:38 +08:00
}
}
fn ensure_rustls_provider () {
static INIT : Once = Once ::new ();
INIT . call_once ( || {
let _ = rustls ::crypto ::ring ::default_provider (). install_default ();
});
}
}