mirror of https://github.com/ospab/ostp.git
fix(server): UoT — set TCP_NODELAY + tear down half-open connections
Two UoT (UDP-over-TCP) correctness issues: - The accepted UoT stream never had TCP_NODELAY set (the client sets it on its end, the server didn't). Nagle's algorithm then batched server->client writes and interacted with the client's delayed ACKs, adding tens-to- hundreds of ms of stall per burst — throttling the download direction badly for streaming/video. Every TCP-tunnel proxy disables Nagle; now the server matches the client. - handle_tcp_connection join!ed the reader and writer tasks, so a half-open connection (client's read side gone, no outbound data pending) parked the writer on rx.recv() forever, leaking the task and a stale tcp_map entry. Rewrote it with select! so either half closing cancels the other and the tcp_map entry is always removed. Added duplex-stream tests covering inbound reassembly across segment boundaries, outbound framing, and teardown-on-close.
This commit is contained in:
parent
7e3ada8d4d
commit
5fcc0ba7f4
|
|
@ -410,6 +410,15 @@ async fn run_server_loop(
|
|||
|
||||
loop {
|
||||
if let Ok((stream, peer_addr)) = listener.accept().await {
|
||||
// Disable Nagle's algorithm on the UoT carrier. Without
|
||||
// this the server→client (download) direction batches
|
||||
// small writes and interacts with the client's delayed
|
||||
// ACKs, adding tens-to-hundreds of ms of stall per burst
|
||||
// — which throttles throughput badly for streaming/video.
|
||||
// The client already sets nodelay on its end; the server
|
||||
// must match. (Every TCP-tunnel proxy sets TCP_NODELAY.)
|
||||
let _ = stream.set_nodelay(true);
|
||||
|
||||
// Rate limit check
|
||||
let peer_ip = peer_addr.ip();
|
||||
let allowed = {
|
||||
|
|
|
|||
|
|
@ -23,36 +23,115 @@ where
|
|||
tcp_map.write().await.insert(peer_addr, tx);
|
||||
}
|
||||
|
||||
// Process streams
|
||||
let (mut read_half, mut write_half) = tokio::io::split(stream);
|
||||
|
||||
// Spawn writer task
|
||||
let peer_clone = peer_addr;
|
||||
let tcp_map_clone = tcp_map.clone();
|
||||
let writer_task = tokio::spawn(async move {
|
||||
// Writer: length-prefix (u16 BE) each outbound datagram. OSTP datagrams are
|
||||
// always well under 64 KiB (MTU-bounded), so the u16 prefix never truncates.
|
||||
let writer = async move {
|
||||
while let Some(packet) = rx.recv().await {
|
||||
let mut out = BytesMut::with_capacity(2 + packet.len());
|
||||
out.put_u16(packet.len() as u16);
|
||||
out.put_slice(&packet);
|
||||
if write_half.write_all(&out).await.is_err() { break; }
|
||||
}
|
||||
let _ = tcp_map_clone.write().await.remove(&peer_clone);
|
||||
});
|
||||
};
|
||||
|
||||
// Spawn reader task
|
||||
let reader_task = tokio::spawn(async move {
|
||||
// Reader: reassemble length-prefixed frames off the TCP stream (read_exact
|
||||
// handles TCP segmentation) and forward each to the UDP dispatch path.
|
||||
let reader = async move {
|
||||
let mut len_buf = [0u8; 2];
|
||||
loop {
|
||||
if read_half.read_exact(&mut len_buf).await.is_err() { break; }
|
||||
let len = u16::from_be_bytes(len_buf) as usize;
|
||||
if len > 65536 { break; }
|
||||
let mut data = vec![0u8; len];
|
||||
if read_half.read_exact(&mut data).await.is_err() { break; }
|
||||
if udp_tx.send((Bytes::from(data), peer_clone)).await.is_err() { return; }
|
||||
if udp_tx.send((Bytes::from(data), peer_addr)).await.is_err() { break; }
|
||||
}
|
||||
});
|
||||
};
|
||||
|
||||
let _ = tokio::join!(writer_task, reader_task);
|
||||
// Either half completing means the connection is dead in that direction, so
|
||||
// tear the whole thing down: `select!` drops (cancels) the other half. The
|
||||
// old code `join!`ed both, so a half-open connection (client's read side
|
||||
// gone but no outbound data pending) left the writer parked on rx.recv()
|
||||
// forever, leaking the task and its stale tcp_map entry.
|
||||
tokio::select! {
|
||||
_ = writer => {},
|
||||
_ = reader => {},
|
||||
}
|
||||
|
||||
tcp_map.write().await.remove(&peer_addr);
|
||||
tracing::debug!("UoT client disconnected: {}", peer_addr);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn peer() -> SocketAddr {
|
||||
"127.0.0.1:40000".parse().unwrap()
|
||||
}
|
||||
|
||||
/// Inbound framing: length-prefixed frames written by the client are
|
||||
/// reassembled (even when split across reads) and forwarded to udp_tx.
|
||||
#[tokio::test]
|
||||
async fn reader_reassembles_framed_datagrams() {
|
||||
let (mut client, server) = tokio::io::duplex(64 * 1024);
|
||||
let tcp_map: Arc<RwLock<HashMap<SocketAddr, mpsc::Sender<Bytes>>>> =
|
||||
Arc::new(RwLock::new(HashMap::new()));
|
||||
let (udp_tx, mut udp_rx) = mpsc::channel(16);
|
||||
let handle = tokio::spawn(handle_tcp_connection(server, peer(), tcp_map.clone(), udp_tx));
|
||||
|
||||
// Two datagrams; write the second one byte-at-a-time to exercise the
|
||||
// read_exact reassembly across TCP segment boundaries.
|
||||
let d1 = b"hello".to_vec();
|
||||
client.write_all(&(d1.len() as u16).to_be_bytes()).await.unwrap();
|
||||
client.write_all(&d1).await.unwrap();
|
||||
|
||||
let d2 = vec![0xAB_u8; 1400];
|
||||
let framed2: Vec<u8> = (d2.len() as u16).to_be_bytes().iter().chain(d2.iter()).copied().collect();
|
||||
for b in &framed2 {
|
||||
client.write_all(&[*b]).await.unwrap();
|
||||
}
|
||||
|
||||
let (got1, _) = udp_rx.recv().await.unwrap();
|
||||
assert_eq!(got1.as_ref(), d1.as_slice());
|
||||
let (got2, from) = udp_rx.recv().await.unwrap();
|
||||
assert_eq!(got2.as_ref(), d2.as_slice());
|
||||
assert_eq!(from, peer());
|
||||
|
||||
drop(client);
|
||||
let _ = handle.await;
|
||||
}
|
||||
|
||||
/// Outbound framing + teardown: a datagram sent via the tcp_map sender is
|
||||
/// written to the wire with its u16 length prefix, and once the client
|
||||
/// hangs up the connection is torn down and its tcp_map entry removed.
|
||||
#[tokio::test]
|
||||
async fn writer_frames_outbound_and_cleans_up_on_close() {
|
||||
let (mut client, server) = tokio::io::duplex(64 * 1024);
|
||||
let tcp_map: Arc<RwLock<HashMap<SocketAddr, mpsc::Sender<Bytes>>>> =
|
||||
Arc::new(RwLock::new(HashMap::new()));
|
||||
let (udp_tx, _udp_rx) = mpsc::channel(16);
|
||||
let handle = tokio::spawn(handle_tcp_connection(server, peer(), tcp_map.clone(), udp_tx));
|
||||
|
||||
// Wait for registration, then push an outbound datagram.
|
||||
let tx = loop {
|
||||
if let Some(tx) = tcp_map.read().await.get(&peer()).cloned() { break tx; }
|
||||
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
|
||||
};
|
||||
tx.send(Bytes::from_static(b"world!")).await.unwrap();
|
||||
|
||||
let mut len_buf = [0u8; 2];
|
||||
client.read_exact(&mut len_buf).await.unwrap();
|
||||
assert_eq!(u16::from_be_bytes(len_buf) as usize, 6);
|
||||
let mut body = [0u8; 6];
|
||||
client.read_exact(&mut body).await.unwrap();
|
||||
assert_eq!(&body, b"world!");
|
||||
|
||||
// Client hangs up -> reader hits EOF -> select! tears down -> entry gone.
|
||||
drop(client);
|
||||
let _ = handle.await;
|
||||
assert!(!tcp_map.read().await.contains_key(&peer()), "tcp_map entry must be removed on close");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue