diff --git a/ostp-server/src/dispatcher.rs b/ostp-server/src/dispatcher.rs index 20e144a..22a03ca 100644 --- a/ostp-server/src/dispatcher.rs +++ b/ostp-server/src/dispatcher.rs @@ -246,6 +246,30 @@ impl Dispatcher { self.peer_machines.len() } + /// Per-session download-direction congestion headroom, in packets: + /// `(session_id, available)` where `available = clamped cwnd - in_flight`. + /// + /// Consumed by the relay's per-target-connection reader tasks (see + /// `relay::handle_relay_message`'s Connect handler) to throttle how fast + /// they pull bytes from the upstream target and forward them to the + /// client's OSTP session. Without this, a fast target (e.g. a CDN) gets + /// read and forwarded as fast as the target can serve, completely + /// ignoring the client-facing session's real congestion window - on a + /// lossy/jittery client path that self-inflicts a loss burst, which + /// wrecks the RTT/RTO estimate and can stall the session hard enough to + /// trip the client's keepalive reconnect. Same clamp(16, 16384) the + /// client uses for its own analogous uplink gate, for symmetry. + pub fn snapshot_backpressure(&self) -> Vec<(u32, i64)> { + self.peer_machines + .iter() + .map(|(&sid, ps)| { + let cwnd = (ps.machine.cwnd_packets() as i64).clamp(16, 16384); + let in_flight = ps.machine.in_flight_count() as i64; + (sid, cwnd - in_flight) + }) + .collect() + } + pub fn on_datagram(&mut self, peer: SocketAddr, packet: Bytes) -> Result { if packet.len() < 4 { return Ok(DispatchOutcome::Unauthorized); diff --git a/ostp-server/src/lib.rs b/ostp-server/src/lib.rs index 70b7960..93930ce 100644 --- a/ostp-server/src/lib.rs +++ b/ostp-server/src/lib.rs @@ -1,7 +1,9 @@ use anyhow::Result; use bytes::Bytes; +use portable_atomic::AtomicI64; use std::collections::HashMap; use std::net::IpAddr; +use std::sync::{Arc, RwLock}; use dispatcher::{DispatchOutcome, Dispatcher}; use ostp_core::relay::RelayMessage; @@ -10,6 +12,12 @@ use tokio::net::UdpSocket; use tokio::sync::mpsc; use tokio::time::{interval, Duration, Instant}; +/// Shared per-session download-direction congestion headroom (packets), +/// published by `handle_tick` from `Dispatcher::snapshot_backpressure` and +/// read lock-free by relay reader tasks. See that method's doc comment for +/// why this exists. +pub(crate) type SessionBackpressure = Arc>>>; + mod dispatcher; pub mod outbound; pub mod api; @@ -467,6 +475,7 @@ async fn run_server_loop( let mut last_empty_app_log = Instant::now() - Duration::from_secs(10); let mut peer_last_seen: HashMap = HashMap::new(); let mut peer_available: HashMap = HashMap::new(); + let session_backpressure: SessionBackpressure = Arc::new(RwLock::new(HashMap::new())); loop { tokio::select! { @@ -489,7 +498,8 @@ async fn run_server_loop( packet, peer, &mut dispatcher, &tcp_map, &socket, &mut remotes, &ui_event_tx, stream_tx.clone(), udp_reply_tx.clone(), connect_tx.clone(), router.clone(), - &mut peer_last_seen, &mut peer_available, &mut last_empty_app_log + &mut peer_last_seen, &mut peer_available, &mut last_empty_app_log, + &session_backpressure ).await { tracing::error!("handle_udp_packet error: {}", e); } @@ -533,7 +543,7 @@ async fn run_server_loop( _ = retransmit_tick.tick() => { if let Err(e) = handle_tick( &mut dispatcher, &tcp_map, &socket, &mut remotes, &ui_event_tx, - &mut peer_last_seen, &mut peer_available + &mut peer_last_seen, &mut peer_available, &session_backpressure ).await { tracing::error!("handle_tick error: {}", e); } @@ -559,6 +569,7 @@ async fn handle_udp_packet( peer_last_seen: &mut HashMap, peer_available: &mut HashMap, last_empty_app_log: &mut Instant, + session_backpressure: &SessionBackpressure, ) -> Result<()> { let size = packet.len(); match dispatcher.on_datagram(peer, packet.clone()) { @@ -621,6 +632,7 @@ async fn handle_udp_packet( connect_tx.clone(), router.clone(), tcp_map, + session_backpressure, ).await?; } } @@ -639,6 +651,7 @@ async fn handle_tick( ui_event_tx: &mpsc::UnboundedSender, peer_last_seen: &mut HashMap, peer_available: &mut HashMap, + session_backpressure: &SessionBackpressure, ) -> Result<()> { let now = Instant::now(); let peer_timeout = Duration::from_secs(45); @@ -649,6 +662,22 @@ async fn handle_tick( let _ = ui_event_tx.send(UiEvent::Log(format!("Client {peer_ip} disconnected (timeout)"))); } } + // Publish each active session's current download-direction headroom so + // relay reader tasks (running on other tasks, no access to `dispatcher`) + // can throttle without touching a lock on every read. New sessions get an + // entry created here on their first tick after the handshake; entries for + // sessions that no longer exist are pruned below alongside dropped_sessions. + { + let snapshot = dispatcher.snapshot_backpressure(); + let mut map = session_backpressure.write().unwrap_or_else(|e| e.into_inner()); + for (sid, available) in snapshot { + match map.get(&sid) { + Some(slot) => slot.store(available, std::sync::atomic::Ordering::Relaxed), + None => { map.insert(sid, Arc::new(AtomicI64::new(available))); } + } + } + } + let (frames, dropped_sessions) = dispatcher.on_tick(); for (frame, peer_addr) in frames { let mut sent_tcp = false; @@ -663,6 +692,12 @@ async fn handle_tick( let _ = socket.send_to(&frame, peer_addr).await?; } } + if !dropped_sessions.is_empty() { + let mut map = session_backpressure.write().unwrap_or_else(|e| e.into_inner()); + for sid in &dropped_sessions { + map.remove(sid); + } + } for sid in dropped_sessions { let _ = ui_event_tx.send(UiEvent::Log(format!("Session {sid} expired, releasing resources"))); let mut streams_to_cancel = Vec::new(); diff --git a/ostp-server/src/relay.rs b/ostp-server/src/relay.rs index 7cb86d1..2a93f61 100644 --- a/ostp-server/src/relay.rs +++ b/ostp-server/src/relay.rs @@ -1,6 +1,8 @@ use anyhow::Result; use bytes::Bytes; +use portable_atomic::AtomicI64; use std::collections::HashMap; +use std::sync::Arc; use ostp_core::relay::RelayMessage; use tokio::io::AsyncReadExt; @@ -8,7 +10,19 @@ use tokio::net::UdpSocket; use tokio::sync::mpsc; use crate::dispatcher::Dispatcher; -use crate::{RemoteState, UiEvent}; +use crate::{RemoteState, SessionBackpressure, UiEvent}; + +/// How long a target-connection reader task waits before rechecking the +/// client session's congestion headroom while throttled. Short enough that +/// a freed-up window (checked every server tick, 10ms) is noticed promptly; +/// long enough not to spin. +const BACKPRESSURE_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(5); +/// Upper bound on total time a single read is throttled before proceeding +/// anyway. Congestion state is a hint, not a hard guarantee - if the +/// session's headroom never frees up (e.g. a stuck/buggy state), a stream +/// must not be stalled forever; better to occasionally overshoot the window +/// than deadlock a connection. +const BACKPRESSURE_MAX_WAIT: std::time::Duration = std::time::Duration::from_secs(2); fn clean_ipv6_mapped_v4(addr: std::net::SocketAddr) -> std::net::SocketAddr { match addr { @@ -38,6 +52,7 @@ pub async fn handle_relay_message( connect_tx: mpsc::UnboundedSender<(u32, u16, String, Result<(tokio::net::tcp::OwnedWriteHalf, mpsc::Sender<()>), String>)>, router: std::sync::Arc, tcp_map: &std::sync::Arc>>>, + session_backpressure: &SessionBackpressure, ) -> Result<()> { match RelayMessage::decode(&payload)? { RelayMessage::Connect(target) => { @@ -53,15 +68,41 @@ pub async fn handle_relay_message( let connect_tx_clone = connect_tx.clone(); let stream_tx_clone = stream_tx.clone(); let router_clone = router.clone(); + let backpressure_clone = session_backpressure.clone(); tokio::spawn(async move { let stream_res = router_clone.route_tcp(&target_clone).await; match stream_res { Ok(stream) => { let (mut reader, writer) = stream.into_split(); let (cancel_tx, mut cancel_rx) = mpsc::channel::<()>(1); + // Get-or-create this session's headroom handle. A brand + // new session may not have its first tick's snapshot + // yet (up to 10ms), so default it open (matches a fresh + // congestion window) rather than stalling the very + // first read while nothing has been published. + let headroom: Arc = { + let mut map = backpressure_clone.write().unwrap_or_else(|e| e.into_inner()); + map.entry(session_id).or_insert_with(|| Arc::new(AtomicI64::new(32))).clone() + }; tokio::spawn(async move { let mut buf = [0_u8; 4096]; loop { + // Throttle to the client-facing OSTP session's + // real congestion window instead of reading from + // the target as fast as it'll send. Without this, + // a fast target blasts a lossy/jittery client + // path far beyond what it can sustain, which + // self-inflicts a loss burst, wrecks the RTT/RTO + // estimate, and can stall the session hard + // enough to trip the client's keepalive + // reconnect. See Dispatcher::snapshot_backpressure. + let mut waited = std::time::Duration::ZERO; + while headroom.load(std::sync::atomic::Ordering::Relaxed) <= 0 + && waited < BACKPRESSURE_MAX_WAIT + { + tokio::time::sleep(BACKPRESSURE_POLL_INTERVAL).await; + waited += BACKPRESSURE_POLL_INTERVAL; + } tokio::select! { _ = cancel_rx.recv() => break, read_res = reader.read(&mut buf) => {