Create a flow control for tcp (#265)

This branch adds a FlowControl mechanism that side channels the queues.
The control handling process essentially applies back pressure by
tracking the in flight requests end to end and enforcing the
tcp_capacity on the system.

Author: atreyas

src/host.rs

@@ -1,6 +1,7 @@
 use crate::envelope::{hex, Datagram, Protocol, Segment, Syn};
 #[cfg(feature = "unstable-fs")]
 use crate::fs::{Fs, FsConfig};
+use crate::net::tcp::stream::BidiFlowControl;
 use crate::net::{SocketPair, TcpListener, UdpSocket};
 use crate::{Envelope, TRACING_TARGET};
 
@@ -11,10 +12,9 @@ use std::fmt::Display;
 use std::io;
 use std::net::{IpAddr, SocketAddr};
 use std::ops::RangeInclusive;
-#[cfg(not(feature = "unstable-fs"))]
 use std::sync::Arc;
 #[cfg(feature = "unstable-fs")]
-use std::sync::{Arc, Mutex};
+use std::sync::Mutex;
 use tokio::sync::{mpsc, Notify};
 use tokio::time::{Duration, Instant};
 
@@ -346,11 +346,11 @@ struct ServerSocket {
 }
 
 struct StreamSocket {
-    local_addr: SocketAddr,
     buf: IndexMap<u64, SequencedSegment>,
     next_send_seq: u64,
     recv_seq: u64,
     sender: mpsc::Sender<SequencedSegment>,
+    flow_control: BidiFlowControl,
     /// A simple reference counter for tracking read/write half drops. Once 0, the
     /// socket may be removed from the host.
     ref_ct: usize,
@@ -374,18 +374,19 @@ impl Display for SequencedSegment {
 }
 
 impl StreamSocket {
-    fn new(local_addr: SocketAddr, capacity: usize) -> (Self, mpsc::Receiver<SequencedSegment>) {
+    fn new(capacity: usize) -> (Self, mpsc::Receiver<SequencedSegment>, BidiFlowControl) {
         let (tx, rx) = mpsc::channel(capacity);
+        let flow_control = BidiFlowControl::new(capacity);
         let sock = Self {
-            local_addr,
             buf: IndexMap::new(),
             next_send_seq: 1,
             recv_seq: 0,
             sender: tx,
+            flow_control: flow_control.clone(),
             ref_ct: 2,
         };
 
-        (sock, rx)
+        (sock, rx, flow_control)
     }
 
     fn assign_seq(&mut self) -> u64 {
@@ -406,11 +407,17 @@ impl StreamSocket {
         while self.buf.contains_key(&(self.recv_seq + 1)) {
             self.recv_seq += 1;
 
-            let segment = self.buf.swap_remove(&self.recv_seq).unwrap();
-            self.sender.try_send(segment).map_err(|e| match e {
-                Closed(_) => Protocol::Tcp(Segment::Rst),
-                Full(_) => panic!("{} socket buffer full", self.local_addr),
-            })?;
+            match self.sender.try_reserve() {
+                Ok(permit) => {
+                    let segment = self.buf.swap_remove(&self.recv_seq).unwrap();
+                    permit.send(segment)
+                }
+                Err(Closed(())) => return Err(Protocol::Tcp(Segment::Rst)),
+                Err(Full(())) => {
+                    self.recv_seq -= 1;
+                    break;
+                }
+            }
         }
 
         Ok(())
@@ -451,14 +458,25 @@ impl Tcp {
         Ok(TcpListener::new(addr, notify))
     }
 
-    pub(crate) fn new_stream(&mut self, pair: SocketPair) -> mpsc::Receiver<SequencedSegment> {
-        let (sock, rx) = StreamSocket::new(pair.local, self.socket_capacity);
+    pub(crate) fn new_stream(
+        &mut self,
+        pair: SocketPair,
+    ) -> (mpsc::Receiver<SequencedSegment>, BidiFlowControl) {
+        let (sock, rx, bidi) = StreamSocket::new(self.socket_capacity);
 
         let exists = self.sockets.insert(pair, sock);
 
         assert!(exists.is_none(), "{pair:?} is already connected");
 
-        rx
+        (rx, bidi)
+    }
+
+    pub(crate) fn flow_control(&self, pair: SocketPair) -> BidiFlowControl {
+        self.sockets
+            .get(&pair)
+            .expect("missing stream socket")
+            .flow_control
+            .clone()
     }
 
     pub(crate) fn stream_count(&self) -> usize {

src/net/tcp/listener.rs

@@ -71,15 +71,12 @@ impl TcpListener {
             });
 
             let Some((syn, origin)) = maybe_accept else {
-                // Wait for a new incoming connection, then retry.
                 self.notify.notified().await;
                 continue;
             };
 
             tracing::trace!(target: TRACING_TARGET, src = ?origin, dst = ?self.local_addr, protocol = %"TCP SYN", "Recv");
 
-            // Send SYN-ACK -> origin. If Ok we proceed (acts as the ACK), else
-            // we retry.
             let ack = syn.ack.send(());
             tracing::trace!(target: TRACING_TARGET, src = ?self.local_addr, dst = ?origin, protocol = %"TCP SYN-ACK", "Send");
 
@@ -89,19 +86,35 @@ impl TcpListener {
         };
 
         let stream = World::current(|world| {
-            let host = world.current_host_mut();
+            let (pair, rx) = {
+                let host = world.current_host_mut();
 
-            let mut my_addr = self.local_addr;
-            if origin.ip().is_loopback() {
-                my_addr.set_ip(origin.ip());
-            }
-            if my_addr.ip().is_unspecified() {
-                my_addr.set_ip(host.addr);
-            }
+                let mut my_addr = self.local_addr;
+                if origin.ip().is_loopback() {
+                    my_addr.set_ip(origin.ip());
+                }
+                if my_addr.ip().is_unspecified() {
+                    my_addr.set_ip(host.addr);
+                }
+
+                let pair = SocketPair::new(my_addr, origin);
+                let (rx, _) = host.tcp.new_stream(pair);
+                (pair, rx)
+            };
 
-            let pair = SocketPair::new(my_addr, origin);
-            let rx = host.tcp.new_stream(pair);
-            TcpStream::new(pair, rx)
+            let client_ip = if origin.ip().is_loopback() {
+                world.current.expect("current host missing")
+            } else {
+                origin.ip()
+            };
+            let client_pair = SocketPair::new(origin, pair.local);
+            let bidi = world
+                .hosts
+                .get(&client_ip)
+                .expect("client host missing")
+                .tcp
+                .flow_control(client_pair);
+            TcpStream::new(pair, rx, bidi.invert())
         });
 
         tracing::trace!(target: TRACING_TARGET, src = ?self.local_addr, dst = ?origin, "Accepted");