//! Application lifecycle: initialization, task orchestration, and shutdown. use std::collections::BTreeMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, RwLock}; use std::time::{Duration, Instant}; use aether_http::{jittered_delay_for_retry, HttpRetryConfig}; use aether_runtime::{ init_reloadable_service_tracing, prometheus_response, wait_for_shutdown_signal, ConcurrencyGate, }; use aether_runtime_state::{RedisClientConfig, RuntimeSemaphoreConfig, RuntimeState}; use arc_swap::ArcSwap; use axum::extract::State as AxumState; use axum::routing::get; use axum::{Json, Router}; use tokio::sync::{watch, Mutex}; use tokio::task::JoinHandle; use tracing::{error, info, warn}; use crate::config::{ aether_url_for_log, effective_tunnel_security, validate_tunnel_encryption_key, Config, ServerEntry, TunnelPoolSizing, }; use crate::net; use crate::registration::client::AetherClient; use crate::runtime::{self, DynamicConfig}; use crate::state::{AppState, ServerContext, TunnelMetrics, TunnelRequestMetrics}; use crate::upstream_client; use crate::{hardware, target_filter, tunnel}; type TaskHandles = Arc>>>; const AUTO_STREAM_LIMIT_MIN: usize = 16; // 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; const AUTO_STREAM_LIMIT_MEMORY_MB_PER_STREAM: u64 = 4; const AUTO_STREAM_LIMIT_ESTIMATED_DIVISOR: u64 = 12; #[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, handle: JoinHandle<()>, draining: bool, } #[derive(Clone)] struct DiagnosticsState { state: Arc, server_contexts: Arc>>>, } /// Run the full application lifecycle after config has been parsed. pub async fn run(mut config: Config, servers: Vec) -> anyhow::Result<()> { config.validate()?; for (index, server) in servers.iter().enumerate() { server .validate() .map_err(|error| anyhow::anyhow!("servers[{index}] invalid: {error}"))?; } init_tracing(&config); info!( version = env!("CARGO_PKG_VERSION"), node_name = %config.node_name, server_count = servers.len(), "aether-tunnel starting (tunnel mode)" ); if let Some(proxy_url) = config.effective_aether_outbound_proxy_url() { if let Ok(proxy) = crate::egress_proxy::UpstreamProxyConfig::parse(proxy_url) { info!( aether_outbound_proxy_url = %proxy.redacted_url(), "Aether control and tunnel egress proxy configured" ); } } 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(), upstream_proxy_remote_dns = config.upstream_proxy_remote_dns, "provider upstream egress proxy configured" ); } } if config.upstream_proxy_remote_dns { warn!("provider hostname DNS resolution and destination IP access controls are delegated to the trusted upstream proxy"); } // Resolve public IP (best-effort for region info) let public_ip = match &config.public_ip { Some(ip) => ip.clone(), None => net::detect_public_ip() .await .unwrap_or_else(|_| "0.0.0.0".to_string()), }; // 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); } } // Collect hardware info (once at startup, sent during registration) let hw_info = hardware::collect(); 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" ); } // Auto-detect tunnel_max_streams from hardware if not explicitly set if config.tunnel_max_streams.is_none() { 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; config.tunnel_max_streams = Some(auto); info!( tunnel_max_streams = auto, "auto-detected tunnel_max_streams from stream admission limit" ); } 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, 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, 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" ); info!( max_concurrency = hw_info.estimated_max_concurrency, "hardware info collected" ); let dns_cache = Arc::new(target_filter::DnsCache::new( Duration::from_secs(config.dns_cache_ttl_secs), config.dns_cache_capacity, )); 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)); let resource_monitor = Arc::new(hardware::RuntimeResourceMonitor::new()); // Register with each Aether server and build per-server contexts. // Wrapped in Arc so retry_failed_registrations can append later. let server_contexts: Arc>>> = Arc::new(Mutex::new(Vec::new())); let mut failed_entries: Vec<(String, ServerEntry)> = Vec::new(); 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( &config, &entry.aether_url, &entry.management_token, )); 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 { if entry.aether_url.trim_start().starts_with("http://") { warn!( server = %label, url = %aether_url_for_log(&entry.aether_url), "secure tunnel frame encryption starts after registration; deliver install and registration credentials over HTTPS or another trusted bootstrap channel" ); } 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)?; } match client .register(&config, entry, &node_name, &public_ip, Some(&hw_info)) .await { 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" ); server_contexts.lock().await.push(build_server_context( &config, &label, entry, client, &node_name, registration.node_id, registration.tunnel_generation, )); } Err(e) => { warn!( server = %label, url = %aether_url_for_log(&entry.aether_url), error = %e, "registration failed, will retry in background" ); failed_entries.push((label, entry.clone())); } } } 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" ); } // Build shared application state let tunnel_tls_config = Arc::new(crate::tunnel::client::build_tls_config()); let mut state = AppState { config, dns_cache, upstream_client_pool, tunnel_tls_config, resource_monitor, stream_gate: None, distributed_stream_gate: None, }; if let Some(limit) = state.config.max_in_flight_streams { state = state .with_stream_concurrency_gate(Arc::new(ConcurrencyGate::new("tunnel_streams", limit))); } 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"); let runtime = RuntimeState::redis( RedisClientConfig { url: redis_url, key_prefix: state.config.distributed_stream_redis_key_prefix.clone(), }, Some(state.config.distributed_stream_command_timeout_ms), ) .await?; let distributed_gate = runtime.semaphore( "tunnel_streams_distributed", limit, RuntimeSemaphoreConfig { 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); // Shutdown signal channel let (shutdown_tx, shutdown_rx) = watch::channel(false); 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 }; info!( active_servers = server_contexts.lock().await.len(), "running in tunnel mode" ); // Spawn tunnel pool manager per server. let tunnel_handles: TaskHandles = Arc::new(Mutex::new(Vec::new())); let retry_handles: TaskHandles = Arc::new(Mutex::new(Vec::new())); for server in server_contexts.lock().await.iter() { spawn_tunnel_pool_manager( Arc::clone(&state), Arc::clone(server), tunnel_pool_policy, shutdown_rx.clone(), Arc::clone(&tunnel_handles), ) .await; } // Spawn background retry for failed server registrations if !failed_entries.is_empty() { spawn_registration_recovery_tasks( Arc::clone(&state), Arc::clone(&server_contexts), failed_entries, public_ip.clone(), hw_info.clone(), tunnel_pool_policy, shutdown_rx.clone(), Arc::clone(&tunnel_handles), Arc::clone(&retry_handles), ) .await; } // Wait for shutdown signal wait_for_shutdown().await; info!("shutdown signal received, cleaning up..."); let _ = shutdown_tx.send(true); if let Some(handle) = diagnostics_handle { let _ = handle.await; } await_all_handles(&retry_handles).await; // Graceful unregister from all servers (including retry-registered ones) for server in server_contexts.lock().await.iter() { 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 await_all_handles(&tunnel_handles).await; info!("aether-tunnel stopped"); Ok(()) } fn spawn_diagnostics_server( listener: tokio::net::TcpListener, diagnostics_state: DiagnosticsState, mut shutdown: watch::Receiver, ) -> std::io::Result> { 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); info!(bind = %bind_addr, "tunnel diagnostics server listening"); 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 { error!(error = %error, "tunnel diagnostics server exited with error"); } })) } async fn diagnostics_health( AxumState(diagnostics): AxumState, ) -> Json { let servers = diagnostics.server_contexts.lock().await.clone(); let active_connections = servers .iter() .map(|server| server.active_connections.load(Ordering::Acquire)) .sum::(); 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", "service": "aether-tunnel", "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, ) -> 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, ) -> Json { let servers = diagnostics.server_contexts.lock().await.clone(); let active_connections = servers .iter() .map(|server| server.active_connections.load(Ordering::Acquire)) .sum::(); let server_stats = servers .iter() .map(|server| diagnostics_server_stats(server)) .collect::>(); 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", "service": "aether-tunnel", "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, "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, "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), "request_metrics": server.metrics.snapshot(), "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 { 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() })), } } #[allow(clippy::too_many_arguments)] async fn spawn_registration_recovery_tasks( state: Arc, server_contexts: Arc>>>, failed: Vec<(String, ServerEntry)>, public_ip: String, hw_info: crate::hardware::HardwareInfo, tunnel_pool_policy: TunnelPoolPolicy, shutdown: watch::Receiver, tunnel_handles: TaskHandles, retry_handles: TaskHandles, ) { 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, tunnel_pool_policy, retry_shutdown, retry_tunnels, ) .await; })); } retry_handles.lock().await.extend(handles); } /// Background task that retries registration for a single server until either /// registration succeeds or shutdown is requested. #[allow(clippy::too_many_arguments)] async fn retry_failed_registration( state: Arc, server_contexts: Arc>>>, label: String, entry: ServerEntry, public_ip: String, hw_info: crate::hardware::HardwareInfo, tunnel_pool_policy: TunnelPoolPolicy, mut shutdown: watch::Receiver, 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 = 0u32; loop { if *shutdown.borrow() { info!(server = %label, "shutdown during registration retry"); return; } 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 .register( &state.config, &entry, &node_name, &public_ip, Some(&hw_info), ) .await { Ok(registration) => { info!(server = %label, node_id = %registration.node_id, tunnel_generation = %registration.tunnel_generation, attempt, "registration retry succeeded"); let server = build_server_context( &state.config, &label, &entry, client, &node_name, registration.node_id, registration.tunnel_generation, ); server_contexts.lock().await.push(Arc::clone(&server)); spawn_tunnel_pool_manager( Arc::clone(&state), server, tunnel_pool_policy, 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() } 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) } fn build_server_context( config: &Config, label: &str, entry: &ServerEntry, client: Arc, node_name: &str, node_id: String, tunnel_generation: String, ) -> Arc { 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(), 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(), node_name: node_name.to_string(), node_id: Arc::new(RwLock::new(node_id)), tunnel_generation, aether_client: client, dynamic: Arc::new(ArcSwap::from_pointee(dynamic)), active_connections: Arc::new(AtomicU64::new(0)), metrics: Arc::new(TunnelRequestMetrics::new()), tunnel_metrics: Arc::new(TunnelMetrics::new()), }) } async fn spawn_tunnel_pool_manager( state: Arc, server: Arc, policy: TunnelPoolPolicy, shutdown: watch::Receiver, tunnel_handles: TaskHandles, ) { 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, server: Arc, policy: TunnelPoolPolicy, mut shutdown: watch::Receiver, ) { let mut tunnels = BTreeMap::::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 = 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; } } } } 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, target_connections: usize, policy: &TunnelPoolPolicy, state: &Arc, server: &Arc, shutdown: &watch::Receiver, ) { 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, server: Arc, slot_id: usize, shutdown: watch::Receiver, ) -> 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) { let finished_slots = tunnels .iter() .filter_map(|(slot_id, tunnel)| tunnel.handle.is_finished().then_some(*slot_id)) .collect::>(); 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, max_connections: usize, ) -> Option { (0..max_connections).find(|slot_id| !tunnels.contains_key(slot_id)) } fn request_tunnel_drain( tunnels: &mut BTreeMap, 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 } async fn await_all_handles(handles: &TaskHandles) { let mut pending = handles.lock().await.drain(..).collect::>(); while let Some(handle) = pending.pop() { let _ = handle.await; } } fn init_tracing(config: &Config) { let reloader = init_reloadable_service_tracing( &config.log_level, config .service_runtime_config() .expect("tunnel service runtime config should be valid"), ) .expect("tunnel tracing should initialize"); runtime::set_log_reloader(reloader); } async fn wait_for_shutdown() { wait_for_shutdown_signal() .await .expect("failed to install shutdown signal handler"); } #[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; use axum::routing::{get, post}; use axum::Router; use serde_json::json; use crate::config::{TunnelLogDestinationArg, TunnelLogRotationArg}; use crate::hardware::HardwareInfo; use crate::state::AppState as TunnelAppState; 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>>> = 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()), tunnel_security: None, tunnel_encryption_key: None, }, )]; let (shutdown_tx, shutdown_rx) = watch::channel(false); let tunnel_pool_policy = TunnelPoolPolicy::from_config(&state.config, sample_tunnel_pool_sizing()); spawn_registration_recovery_tasks( Arc::clone(&state), Arc::clone(&server_contexts), failed, "127.0.0.1".to_string(), sample_hardware_info(), tunnel_pool_policy, 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(®ister_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(); } #[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"); assert_eq!(health["service"], "aether-tunnel"); 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"); assert!(metrics.contains("service_up{service=\"aether-tunnel\"} 1")); assert!(metrics.contains("tunnel_active_connections{server=\"server\"} 0")); 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"); assert_eq!( stats["protocol_version"], aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION ); assert_eq!(stats["servers"][0]["node_id"], "node-diagnostics"); handle.abort(); } #[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)); } #[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); } #[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); } #[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); } async fn wait_for_registered_server( server_contexts: &Arc>>>, ) -> Arc { 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, ) -> Result, std::io::Error> { let mut attempts = 0usize; loop { match start_fake_gateway_on_port(port, Arc::clone(®ister_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, ) -> Result, 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, 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::(), ) .await .expect("gateway test server should run"); })) } fn reserve_local_port() -> Result { 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>, ) -> (AxumStatusCode, axum::Json) { register_hits.fetch_add(1, Ordering::SeqCst); ( AxumStatusCode::OK, axum::Json(json!({ "node_id": "node-recovery", "tunnel_generation": "test-generation-recovery" })), ) } 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, } } fn sample_tunnel_pool_sizing() -> TunnelPoolSizing { TunnelPoolSizing { initial_connections: 1, max_connections: 1, } } fn sample_state(config: Config) -> Arc { let config = Arc::new(config); let dns_cache = Arc::new(DnsCache::new(Duration::from_secs(60), 128)); let upstream_client_pool = upstream_client::UpstreamClientPool::new(Arc::clone(&config), Arc::clone(&dns_cache)); Arc::new(TunnelAppState { config, dns_cache, upstream_client_pool, tunnel_tls_config: Arc::new(crate::tunnel::client::build_tls_config()), resource_monitor: Arc::new(crate::hardware::RuntimeResourceMonitor::new()), stream_gate: None, distributed_stream_gate: None, }) } fn sample_registered_server( state: &Arc, label: &str, node_id: &str, ) -> Arc { 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()), tunnel_security: Some(state.config.tunnel_security), tunnel_encryption_key: state.config.tunnel_encryption_key.clone(), }; 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(), "test-generation-1".to_string(), ) } fn sample_config(aether_url: &str) -> Config { Config { aether_url: aether_url.to_string(), management_token: "token".to_string(), public_ip: None, node_name: "tunnel-test".to_string(), tunnel_security: crate::config::TunnelSecurity::Off, tunnel_encryption_key: None, node_region: None, heartbeat_interval: 1, allowed_ports: vec![80, 443], allow_private_targets: false, 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, aether_outbound_proxy_url: None, aether_retry_max_attempts: 1, aether_retry_base_delay_ms: 50, aether_retry_max_delay_ms: 100, diagnostics_bind: None, 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, upstream_client_pool_capacity: crate::config::DEFAULT_UPSTREAM_CLIENT_POOL_CAPACITY, upstream_tcp_keepalive_secs: 60, upstream_tcp_nodelay: true, upstream_proxy_url: None, upstream_proxy_remote_dns: false, legacy_redirect_replay_budget_bytes_ignored: None, emit_proxy_timing_header: true, log_level: "info".to_string(), log_destination: TunnelLogDestinationArg::Stdout, log_dir: None, log_rotation: TunnelLogRotationArg::Daily, log_retention_days: 7, log_max_files: 30, tunnel_reconnect_base_ms: 50, tunnel_reconnect_max_ms: 250, tunnel_ping_interval_ms: 1_000, tunnel_max_streams: Some(8), 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, tunnel_connect_timeout_ms: 2_000, tunnel_ipv4_only: false, tunnel_ipv6_only: false, tunnel_tcp_keepalive_secs: 30, tunnel_tcp_nodelay: true, 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, } } fn ensure_rustls_provider() { static INIT: Once = Once::new(); INIT.call_once(|| { let _ = rustls::crypto::ring::default_provider().install_default(); }); } }