From 6d1711fd9c3185a66f6eeb1acfdc6832d9c44203 Mon Sep 17 00:00:00 2001 From: Sajjad Pourali Date: Sun, 1 Nov 2026 20:15:58 -0700 Subject: [PATCH 1/2] Replace TCP read and write waker mutexes with AtomicWaker --- Cargo.toml | 1 + src/stream/tcp.rs | 111 +++++++++++++++++++++++++++++++++------------- 2 files changed, 81 insertions(+), 31 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 9f6d55a..be7c38a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,6 +10,7 @@ homepage = 'https://narrowlink.com' readme = "README.md" [dependencies] +atomic-waker = "1.1" ahash = { version = "0.8", default-features = false, features = ["std"] } etherparse = { version = "0.20", default-features = false, features = ["std"] } log = { version = "0.4", default-features = false } diff --git a/src/stream/tcp.rs b/src/stream/tcp.rs index 212b7a0..8838d26 100644 --- a/src/stream/tcp.rs +++ b/src/stream/tcp.rs @@ -9,6 +9,7 @@ use crate::{ }, stream::tcb::{MAX_COUNT_FOR_DUP_ACK, MAX_RETRANSMIT_COUNT, MAX_UNACK, PacketType, READ_BUFFER_SIZE, READ_CHUNK, RTO, Tcb, TcpState}, }; +use atomic_waker::AtomicWaker; use etherparse::{IpNumber, Ipv4Header, Ipv6FlowLabel, TcpHeader, TcpOptionElement}; use std::{ future::Future, @@ -162,12 +163,12 @@ pub struct IpStackTcpStream { up_packet_sender: PacketSender, tcb: TcbPtr, shutdown: std::sync::Arc>, - write_notify: std::sync::Arc>>, + write_notify: Arc, destroy_messenger: Option<::tokio::sync::oneshot::Sender<()>>, timeout: Pin>, data_tx: tokio::sync::mpsc::Sender>, data_rx: tokio::sync::mpsc::Receiver>, - read_notify: std::sync::Arc>>, + read_notify: Arc, drain_notify: Arc, arm_notify: Arc, task_handle: Option>>, @@ -231,12 +232,12 @@ impl IpStackTcpStream { up_packet_sender, tcb: std::sync::Arc::new(std::sync::Mutex::new(tcb.clone())), shutdown: std::sync::Arc::new(std::sync::Mutex::new(Shutdown::None)), - write_notify: std::sync::Arc::new(std::sync::Mutex::new(None)), + write_notify: Arc::new(AtomicWaker::new()), destroy_messenger, timeout: Box::pin(tokio::time::sleep_until(deadline)), data_tx, data_rx, - read_notify: std::sync::Arc::new(std::sync::Mutex::new(None)), + read_notify: Arc::new(AtomicWaker::new()), drain_notify: Arc::new(tokio::sync::Notify::new()), arm_notify: Arc::new(tokio::sync::Notify::new()), task_handle: None, @@ -310,6 +311,8 @@ impl AsyncRead for IpStackTcpStream { let network_tuple = self.network_tuple(); + // Register before checking state so a concurrent close cannot miss this reader. + self.read_notify.register(cx.waker()); let state = self.tcb.lock().unwrap().get_state(); if state == TcpState::Closed { // Data on hand reaches the reader before end-of-stream. @@ -321,7 +324,7 @@ impl AsyncRead for IpStackTcpStream { return Poll::Ready(Ok(())); } self.shutdown.lock().unwrap().ready(); - self.write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + self.write_notify.wake(); return Poll::Ready(Ok(())); } @@ -360,10 +363,7 @@ impl AsyncRead for IpStackTcpStream { Poll::Ready(Ok(())) } Poll::Ready(None) => Poll::Ready(Ok(())), - Poll::Pending => { - self.read_notify.lock().unwrap().replace(cx.waker().clone()); - Poll::Pending - } + Poll::Pending => Poll::Pending, } } } @@ -373,6 +373,8 @@ impl AsyncWrite for IpStackTcpStream { let nt = self.network_tuple(); self.reset_timeout(); + // Register before checking readiness so a concurrent ACK cannot miss this writer. + self.write_notify.register(cx.waker()); let mut tcb = self.tcb.lock().unwrap(); let state = tcb.get_state(); let send_window = tcb.get_send_window(); @@ -380,12 +382,11 @@ impl AsyncWrite for IpStackTcpStream { if state == TcpState::Closed { self.shutdown.lock().unwrap().ready(); - self.read_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + self.read_notify.wake(); return Poll::Ready(Err(std::io::Error::new(BrokenPipe, "TCP connection closed"))); } if send_window == 0 || is_full { - self.write_notify.lock().unwrap().replace(cx.waker().clone()); log::trace!( "{nt} {state:?}: [poll_write] current send window: {send_window}, send buffer full: {is_full}, waiting for the other side to send ACK..." ); @@ -554,8 +555,8 @@ async fn tcp_main_logic_loop( up_packet_sender: PacketSender, exit_notifier: tokio::sync::mpsc::Sender<()>, network_tuple: NetworkTuple, - write_notify: std::sync::Arc>>, - read_notify: std::sync::Arc>>, + write_notify: Arc, + read_notify: Arc, data_tx: tokio::sync::mpsc::Sender>, drain_notify: Arc, arm_notify: Arc, @@ -682,7 +683,7 @@ async fn tcp_main_logic_loop( let network_packet = tokio::select! { _ = exit_monitor.recv() => { log::debug!("{network_tuple} task exited due to exit signal"); - read_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + read_notify.wake(); break; } _ = drain_notify.notified() => { @@ -706,8 +707,8 @@ async fn tcp_main_logic_loop( log::warn!("{network_tuple} segment reached R2 transmissions, closing session"); write_packet_to_device(&up_packet_sender, network_tuple, &tcb, None, ACK | RST, None, None)?; tcb.change_state(TcpState::Closed); - write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); - read_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + write_notify.wake(); + read_notify.wake(); break; } for packet in packets { @@ -732,8 +733,8 @@ async fn tcp_main_logic_loop( let state = { tcb.lock().unwrap().get_state() }; log::debug!("{network_tuple} {state:?}: session closed unexpectedly by pipe broken, exiting task"); tcb.lock().unwrap().change_state(TcpState::Closed); - write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); - read_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + write_notify.wake(); + read_notify.wake(); break; }; @@ -790,13 +791,13 @@ async fn tcp_main_logic_loop( )?; } tcb.change_state(TcpState::Established); - write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + write_notify.wake(); } TcpState::Established => { if flags == ACK { match pkt_type { PacketType::WindowUpdate => { - write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + write_notify.wake(); } PacketType::KeepAlive => { write_packet_to_device(&up_packet_sender, network_tuple, &tcb, None, ACK, None, None)?; @@ -814,10 +815,10 @@ async fn tcp_main_logic_loop( PacketType::NewPacket => { let nt = network_tuple; receive_segment(&up_packet_sender, &mut tcb, nt, incoming_seq, payload, &data_tx, &read_notify)?; - write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + write_notify.wake(); } PacketType::Ack => { - write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + write_notify.wake(); } PacketType::Invalid => {} } @@ -855,7 +856,7 @@ async fn tcp_main_logic_loop( )); } else { // Upstream data pending, wake write_notify and wait - write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + write_notify.wake(); log::debug!("{network_tuple} {state:?}: Waiting for upstream data to complete, inflight packets: {len}",); // Spawn a timeout task to force FIN if upstream is unresponsive @@ -885,7 +886,7 @@ async fn tcp_main_logic_loop( &read_notify, )?; } - write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + write_notify.wake(); } else { // unnormal case, we do nothing here log::trace!("{network_tuple} {state:?}: {l_info}, {pkt_type:?}, unnormal case, we do nothing here"); @@ -912,7 +913,7 @@ async fn tcp_main_logic_loop( config.last_ack_max_retries, )); } else { - write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + write_notify.wake(); } } TcpState::LastAck if flags & ACK == ACK => { @@ -948,7 +949,7 @@ async fn tcp_main_logic_loop( &data_tx, &read_notify, )?; - write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + write_notify.wake(); } let new_state = tcb.get_state(); log::trace!("{network_tuple} {state:?}: Received ACK, transitioned to {new_state:?}"); @@ -985,7 +986,7 @@ async fn tcp_main_logic_loop( &data_tx, &read_notify, )?; - write_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + write_notify.wake(); } if flags & FIN == FIN && tcb.get_ack() == incoming_seq + len as u32 { tcb.increase_ack(); @@ -1038,7 +1039,7 @@ fn receive_segment( seq: SeqNum, payload: Vec, data_tx: &tokio::sync::mpsc::Sender>, - read_notify: &std::sync::Arc>>, + read_notify: &AtomicWaker, ) -> std::io::Result<()> { tcb.add_unordered_packet(seq, payload); if !extract_data_n_write_upstream(up_packet_sender, tcb, network_tuple, data_tx, read_notify)? { @@ -1054,7 +1055,7 @@ fn extract_data_n_write_upstream( tcb: &mut Tcb, network_tuple: NetworkTuple, data_tx: &tokio::sync::mpsc::Sender>, - read_notify: &std::sync::Arc>>, + read_notify: &AtomicWaker, ) -> std::io::Result { let (state, seq, ack) = (tcb.get_state(), tcb.get_seq(), tcb.get_ack()); let l_info = LocalInfo(seq, ack); @@ -1081,7 +1082,7 @@ fn extract_data_n_write_upstream( let hint = if state == TcpState::Established { "normally" } else { "still" }; log::trace!("{network_tuple} {state:?}: {l_info} {hint} receiving data, len = {}", data.len()); permit.send(data); - read_notify.lock().unwrap().take().map(|w| w.wake_by_ref()).unwrap_or(()); + read_notify.wake(); write_packet_to_device(up_packet_sender, network_tuple, tcb, None, ACK, None, None)?; return Ok(true); } @@ -1399,12 +1400,38 @@ mod tests { } assert!(advance(just_before(BACKOFF[MAX_RETRANSMIT_COUNT - 1]), &mut up_rx).await.is_none()); + + // Leave both halves waiting immediately before the final timeout closes the connection. + stream.reset_timeout(); + let (reader_woken, reader_waker) = recording_waker(); + let mut read_cx = Context::from_waker(&reader_waker); + let mut bytes = [0; 1]; + let mut read_buf = tokio::io::ReadBuf::new(&mut bytes); + assert!(Pin::new(&mut stream).poll_read(&mut read_cx, &mut read_buf).is_pending()); + // The default window cap is 16 KiB; fill the remainder behind the initial five bytes. + let remaining = vec![0; MAX_UNACK as usize - 5]; + stream.tcb.lock().unwrap().add_inflight_packet(remaining).unwrap(); + let (writer_woken, writer_waker) = recording_waker(); + let mut write_cx = Context::from_waker(&writer_waker); + assert!(Pin::new(&mut stream).poll_write(&mut write_cx, b"waiting").is_pending()); + let reset = advance(Duration::from_millis(1), &mut up_rx).await.expect("no reset"); assert!(tcp_header_of(&reset).rst); assert!(up_rx.try_recv().is_err()); assert!((tokio::time::Instant::now() - opened).as_secs() > 100); assert_eq!(stream.tcb.lock().unwrap().get_state(), TcpState::Closed); + assert!(reader_woken.load(std::sync::atomic::Ordering::SeqCst)); + assert!(writer_woken.load(std::sync::atomic::Ordering::SeqCst)); + assert!(matches!( + Pin::new(&mut stream).poll_read(&mut read_cx, &mut read_buf), + Poll::Ready(Ok(())) + )); + assert!(read_buf.filled().is_empty()); + assert!(matches!( + Pin::new(&mut stream).poll_write(&mut write_cx, b"waiting"), + Poll::Ready(Err(err)) if err.kind() == BrokenPipe + )); } /// RFC 9293 ยง3.10.7.4: a reset closes the connection and ends its retransmission. @@ -1466,7 +1493,7 @@ mod tests { async fn extract_reserves_before_consuming() { let (up_tx, _up_rx) = tokio::sync::mpsc::unbounded_channel::(); let (data_tx, mut data_rx) = tokio::sync::mpsc::channel::>(1); - let read_notify = std::sync::Arc::new(std::sync::Mutex::new(None)); + let read_notify = Arc::new(AtomicWaker::new()); let nt = NetworkTuple::new("1.1.1.1:1".parse().unwrap(), "2.2.2.2:2".parse().unwrap(), true); let mut tcb = Tcb::new( @@ -1575,6 +1602,28 @@ mod tests { assert_eq!(tcb.get_send_window(), 40_000 << 14); } + #[tokio::test] + async fn pending_reader_wakes_for_each_new_segment() { + let (mut stream, syn_ack, mut up_rx) = open(4000, &[]).await; + let peer_seq = SeqNum(syn_ack.acknowledgment_number); + feed(&stream, &syn_ack, peer_seq, 4000, &[]); + tokio::task::yield_now().await; + + for offset in 0..2 { + let (woken, waker) = recording_waker(); + let mut cx = Context::from_waker(&waker); + let mut bytes = [0; 4]; + let mut buf = tokio::io::ReadBuf::new(&mut bytes); + assert!(Pin::new(&mut stream).poll_read(&mut cx, &mut buf).is_pending()); + + feed(&stream, &syn_ack, peer_seq + offset * 4, 4000, b"data"); + up_rx.recv().await.unwrap(); + assert!(woken.load(std::sync::atomic::Ordering::SeqCst)); + assert!(matches!(Pin::new(&mut stream).poll_read(&mut cx, &mut buf), Poll::Ready(Ok(())))); + assert_eq!(buf.filled(), b"data"); + } + } + /// A closed peer window holds the writer, and the segment that reopens it wakes the writer, both /// when it completes the handshake and when it arrives on an established connection. #[tokio::test] From bb66d6aa2ded7ef2a43a4668257beeb093312a91 Mon Sep 17 00:00:00 2001 From: SajjadPourali Date: Sun, 1 Nov 2026 20:26:14 -0700 Subject: [PATCH 2/2] Replace TCP shutdown mutex with ShutdownSignal --- src/stream/tcp.rs | 112 ++++++++++++++++++++++------------------------ 1 file changed, 53 insertions(+), 59 deletions(-) diff --git a/src/stream/tcp.rs b/src/stream/tcp.rs index 3e4ba85..59c8c53 100644 --- a/src/stream/tcp.rs +++ b/src/stream/tcp.rs @@ -16,8 +16,11 @@ use std::{ io::ErrorKind::{BrokenPipe, ConnectionRefused, InvalidInput, UnexpectedEof}, net::SocketAddr, pin::Pin, - sync::Arc, - task::{Context, Poll, Waker}, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + task::{Context, Poll}, time::Duration, }; use tokio::io::{AsyncRead, AsyncWrite}; @@ -83,41 +86,27 @@ impl Default for TcpConfig { } } -#[derive(Debug)] -enum Shutdown { - None, - Pending(Waker), - Ready, +/// Remembers shutdown completion and wakes the latest task polling for it. +#[derive(Debug, Default)] +struct ShutdownSignal { + completed: AtomicBool, + waker: AtomicWaker, } -impl Shutdown { - fn pending(&mut self, w: Waker) { - *self = Shutdown::Pending(w); - } - fn ready(&mut self) { - if let Shutdown::Pending(w) = self { - w.wake_by_ref(); - } - *self = Shutdown::Ready; - } - - // Just for comparison purpose - fn fake_clone(&self) -> Shutdown { - match self { - Shutdown::None => Shutdown::None, - Shutdown::Pending(_) => Shutdown::Pending(Waker::noop().clone()), - Shutdown::Ready => Shutdown::Ready, +impl ShutdownSignal { + fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<()> { + // Register first: completion either wakes this task or is observed below. + self.waker.register(cx.waker()); + if self.completed.load(Ordering::Acquire) { + Poll::Ready(()) + } else { + Poll::Pending } } -} -impl std::fmt::Display for Shutdown { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Shutdown::None => write!(f, "None"), - Shutdown::Pending(_) => write!(f, "Pending"), - Shutdown::Ready => write!(f, "Ready"), - } + fn complete(&self) { + self.completed.store(true, Ordering::Release); + self.waker.wake(); } } @@ -162,7 +151,7 @@ pub struct IpStackTcpStream { stream_receiver: Option, up_packet_sender: PacketSender, tcb: TcbPtr, - shutdown: std::sync::Arc>, + shutdown: Arc, write_notify: Arc, destroy_messenger: Option<::tokio::sync::oneshot::Sender<()>>, timeout: Pin>, @@ -231,7 +220,7 @@ impl IpStackTcpStream { stream_receiver: Some(stream_receiver), up_packet_sender, tcb: std::sync::Arc::new(std::sync::Mutex::new(tcb.clone())), - shutdown: std::sync::Arc::new(std::sync::Mutex::new(Shutdown::None)), + shutdown: Arc::new(ShutdownSignal::default()), write_notify: Arc::new(AtomicWaker::new()), destroy_messenger, timeout: Box::pin(tokio::time::sleep_until(deadline)), @@ -325,7 +314,7 @@ impl AsyncRead for IpStackTcpStream { self.temp_read_buffer.extend_from_slice(&data[len..]); return Poll::Ready(Ok(())); } - self.shutdown.lock().unwrap().ready(); + self.shutdown.complete(); self.write_notify.wake(); return Poll::Ready(Ok(())); } @@ -343,7 +332,7 @@ impl AsyncRead for IpStackTcpStream { let state = tcb.get_state(); log::debug!("{network_tuple} {state:?}: [poll_read] {l_info}, session notified to close"); } - self.shutdown.lock().unwrap().ready(); + self.shutdown.complete(); return Poll::Ready(Err(std::io::Error::from(std::io::ErrorKind::TimedOut))); } @@ -383,7 +372,7 @@ impl AsyncWrite for IpStackTcpStream { let is_full = tcb.is_send_buffer_full(); if state == TcpState::Closed { - self.shutdown.lock().unwrap().ready(); + self.shutdown.complete(); self.read_notify.wake(); return Poll::Ready(Err(std::io::Error::new(BrokenPipe, "TCP connection closed"))); } @@ -419,36 +408,26 @@ impl AsyncWrite for IpStackTcpStream { } fn poll_shutdown(self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let shutdown = { self.shutdown.lock().unwrap().fake_clone() }; + if self.shutdown.poll_ready(cx).is_ready() { + return Poll::Ready(Ok(())); + } let (nt, state, seq, is_ready) = { let tcb = self.tcb.lock().unwrap(); let is_ready = tcb.get_inflight_packets_total_len() == 0; (self.network_tuple(), tcb.get_state(), tcb.get_seq(), is_ready) }; - log::trace!("{nt} {state:?}: [poll_shutdown] seq = {seq}, ready = {is_ready}, shutdown {shutdown}",); + log::trace!("{nt} {state:?}: [poll_shutdown] seq = {seq}, ready = {is_ready}",); #[cfg(test)] tests::run_after_state_snapshot(); if state == TcpState::Closed { return Poll::Ready(Ok(())); } - match shutdown { - Shutdown::None | Shutdown::Pending(_) => { - if is_ready && state == TcpState::Established { - let mut tcb = self.tcb.lock().unwrap(); - send_fin_n_change_state_to_fin_wait1("[poll_shutdown]", nt, &self.up_packet_sender, &mut tcb)?; - } - // The task may have completed since the snapshot above. Check Ready - // and register the latest waker under the same lock used by ready(). - let mut shutdown = self.shutdown.lock().unwrap(); - if matches!(*shutdown, Shutdown::Ready) { - Poll::Ready(Ok(())) - } else { - shutdown.pending(cx.waker().clone()); - Poll::Pending - } - } - Shutdown::Ready => Poll::Ready(Ok(())), + if is_ready && state == TcpState::Established { + let mut tcb = self.tcb.lock().unwrap(); + send_fin_n_change_state_to_fin_wait1("[poll_shutdown]", nt, &self.up_packet_sender, &mut tcb)?; } + // Completion may have happened since the initial poll or state snapshot. + self.shutdown.poll_ready(cx).map(Ok) } } @@ -542,8 +521,8 @@ impl IpStackTcpStream { } _ = destroy_messenger.map(|m| m.send(())).unwrap_or(Ok(())); log::trace!("{network_tuple} task completed, destroy messenger sent successfully"); - shutdown.lock().unwrap().ready(); - log::trace!("{network_tuple} shutdown.lock().unwrap().ready() =========="); + shutdown.complete(); + log::trace!("{network_tuple} shutdown completed"); v }); self.task_handle = Some(task_handle); @@ -1242,6 +1221,7 @@ pub(crate) fn create_raw_packet( mod tests { use super::*; use crate::stream::tcb::{MAX_COUNT_FOR_DUP_ACK, MAX_RETRANSMIT_COUNT, MAX_UNACK, READ_BUFFER_SIZE, RTO}; + use std::task::Waker; /// A waker that records whether it was woken, and the flag it sets. fn recording_waker() -> (Arc, Waker) { @@ -1255,6 +1235,20 @@ mod tests { (woken.clone(), Waker::from(Arc::new(Recording(woken)))) } + #[test] + fn shutdown_completion_is_remembered_before_any_waiter() { + let signal = ShutdownSignal::default(); + signal.complete(); + signal.complete(); + + // Completion stays ready across polls, including polls with a different task waker. + for _ in 0..2 { + let (_, waker) = recording_waker(); + let mut cx = Context::from_waker(&waker); + assert!(signal.poll_ready(&mut cx).is_ready()); + } + } + use tokio::io::AsyncWriteExt; const PEER: &str = "10.0.0.2:12345"; @@ -1516,7 +1510,7 @@ mod tests { let mut cx = Context::from_waker(&waker); assert!(matches!(Pin::new(&mut stream).poll_shutdown(&mut cx), Poll::Ready(Ok(())))); assert_eq!(stream.tcb.lock().unwrap().get_state(), TcpState::Closed); - assert!(matches!(*stream.shutdown.lock().unwrap(), Shutdown::Ready)); + assert!(stream.shutdown.poll_ready(&mut cx).is_ready()); } /// An acknowledgement from the peer, carrying no data.