Add flushing for peer connections as separate select branch

This restores the old behavior which was introduced to fix an issue
where a peer task could get stuck in a flush call (which was not a
select branch), thus preventing the task from ever exitting, causing a
buildup of tasks (which could keep open OS resources like file
descriptors, causing an eventual resource exhaustion). To do this,
introduce a Connection{Read,Write}Half, and add a `split` method to the
Connection trait. This mimics the behavior of the old code which wrapped
a connection in a framed and then split that entirely.

Signed-off-by: Lee Smet <lee.smet@hotmail.com>

Author: LeeSmet

mycelium/src/connection.rs

@@ -12,7 +12,10 @@ use std::{
 use crate::packet::{self, ControlPacket, DataPacket, Packet};
 
 use bytes::{Bytes, BytesMut};
-use futures::{SinkExt, StreamExt};
+use futures::{
+    stream::{SplitSink, SplitStream},
+    SinkExt, StreamExt,
+};
 use tokio::io::{AsyncRead, AsyncWrite};
 
 mod tracked;
@@ -42,7 +45,34 @@ const PACKET_PROCESSING_COST_IP6_QUIC: u16 = 7;
 // TODO
 const PACKET_PROCESSING_COST_IP4_QUIC: u16 = 12;
 
+pub trait ConnectionReadHalf: Send {
+    /// Receive a packet from the remote end.
+    fn receive_packet(&mut self) -> impl Future<Output = Option<io::Result<Packet>>> + Send;
+}
+
+pub trait ConnectionWriteHalf: Send {
+    /// Feeds a data packet on the connection. Depending on the connection you might need to call
+    /// [`Connection::flush`] before the packet is actually sent.
+    fn feed_data_packet(
+        &mut self,
+        packet: DataPacket,
+    ) -> impl Future<Output = io::Result<()>> + Send;
+
+    /// Feeds a control packet on the connection. Depending on the connection you might need to call
+    /// [`Connection::flush`] before the packet is actually sent.
+    fn feed_control_packet(
+        &mut self,
+        packet: ControlPacket,
+    ) -> impl Future<Output = io::Result<()>> + Send;
+
+    /// Flush the connection. This sends all buffered packets which haven't beend sent yet.
+    fn flush(&mut self) -> impl Future<Output = io::Result<()>> + Send;
+}
+
 pub trait Connection {
+    type ReadHalf: ConnectionReadHalf;
+    type WriteHalf: ConnectionWriteHalf;
+
     /// Feeds a data packet on the connection. Depending on the connection you might need to call
     /// [`Connection::flush`] before the packet is actually sent.
     fn feed_data_packet(
@@ -68,6 +98,9 @@ pub trait Connection {
 
     /// The static cost of using this connection
     fn static_link_cost(&self) -> Result<u16, io::Error>;
+
+    /// Split the connection in a read and write half which can be used independently
+    fn split(self) -> (Self::ReadHalf, Self::WriteHalf);
 }
 
 /// A wrapper about an asynchronous (non blocking) tcp stream.
@@ -93,6 +126,9 @@ impl TcpStream {
 }
 
 impl Connection for TcpStream {
+    type ReadHalf = TcpStreamReadHalf;
+    type WriteHalf = TcpStreamWriteHalf;
+
     async fn feed_data_packet(&mut self, packet: DataPacket) -> io::Result<()> {
         self.framed.feed(Packet::DataPacket(packet)).await
     }
@@ -122,6 +158,43 @@ impl Connection for TcpStream {
             SocketAddr::V6(_) => PACKET_PROCESSING_COST_IP6_TCP,
         })
     }
+
+    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
+        let (tx, rx) = self.framed.split();
+
+        (
+            TcpStreamReadHalf { framed: rx },
+            TcpStreamWriteHalf { framed: tx },
+        )
+    }
+}
+
+pub struct TcpStreamReadHalf {
+    framed: SplitStream<Framed<Tracked<tokio::net::TcpStream>, packet::Codec>>,
+}
+
+impl ConnectionReadHalf for TcpStreamReadHalf {
+    async fn receive_packet(&mut self) -> Option<io::Result<Packet>> {
+        self.framed.next().await
+    }
+}
+
+pub struct TcpStreamWriteHalf {
+    framed: SplitSink<Framed<Tracked<tokio::net::TcpStream>, packet::Codec>, packet::Packet>,
+}
+
+impl ConnectionWriteHalf for TcpStreamWriteHalf {
+    async fn feed_data_packet(&mut self, packet: DataPacket) -> io::Result<()> {
+        self.framed.feed(Packet::DataPacket(packet)).await
+    }
+
+    async fn feed_control_packet(&mut self, packet: ControlPacket) -> io::Result<()> {
+        self.framed.feed(Packet::ControlPacket(packet)).await
+    }
+
+    async fn flush(&mut self) -> io::Result<()> {
+        self.framed.flush().await
+    }
 }
 
 /// A wrapper around a quic send and quic receive stream, implementing the [`Connection`] trait.
@@ -213,6 +286,10 @@ impl AsyncWrite for QuicStream {
 }
 
 impl Connection for Quic {
+    type ReadHalf = QuicReadHalf;
+
+    type WriteHalf = QuicWriteHalf;
+
     async fn feed_data_packet(&mut self, packet: DataPacket) -> io::Result<()> {
         let mut codec = packet::Codec::new();
         let mut buffer = BytesMut::with_capacity(1500);
@@ -271,6 +348,91 @@ impl Connection for Quic {
             SocketAddr::V6(_) => PACKET_PROCESSING_COST_IP6_QUIC,
         })
     }
+
+    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
+        let Self {
+            framed,
+            con,
+            read,
+            write,
+        } = self;
+
+        let (tx, rx) = framed.split();
+
+        (
+            QuicReadHalf {
+                framed: rx,
+                con: con.clone(),
+                read,
+            },
+            QuicWriteHalf {
+                framed: tx,
+                con,
+                write,
+            },
+        )
+    }
+}
+
+pub struct QuicReadHalf {
+    framed: SplitStream<Framed<Tracked<QuicStream>, packet::Codec>>,
+    con: quinn::Connection,
+    read: Arc<AtomicU64>,
+}
+
+pub struct QuicWriteHalf {
+    framed: SplitSink<Framed<Tracked<QuicStream>, packet::Codec>, packet::Packet>,
+    con: quinn::Connection,
+    write: Arc<AtomicU64>,
+}
+
+impl ConnectionReadHalf for QuicReadHalf {
+    async fn receive_packet(&mut self) -> Option<io::Result<Packet>> {
+        tokio::select! {
+            datagram = self.con.read_datagram() => {
+                let datagram_bytes = match datagram {
+                    Ok(buffer) => buffer,
+                    Err(e) => return Some(Err(e.into())),
+                };
+                let recv_len = datagram_bytes.len();
+                self.read.fetch_add(recv_len as u64, Ordering::Relaxed);
+                let mut codec = packet::Codec::new();
+                match codec.decode(&mut datagram_bytes.into()) {
+                    Ok(Some(packet)) => Some(Ok(packet)),
+                    // Partial? packet read. We consider this to be a stream hangup
+                    // TODO: verify
+                    Ok(None) => None,
+                    Err(e) => Some(Err(e)),
+                }
+            },
+            packet = self.framed.next() => {
+                packet
+            }
+
+        }
+    }
+}
+
+impl ConnectionWriteHalf for QuicWriteHalf {
+    async fn feed_data_packet(&mut self, packet: DataPacket) -> io::Result<()> {
+        let mut codec = packet::Codec::new();
+        let mut buffer = BytesMut::with_capacity(1500);
+        codec.encode(Packet::DataPacket(packet), &mut buffer)?;
+
+        let data: Bytes = buffer.into();
+        let tx_len = data.len();
+        self.write.fetch_add(tx_len as u64, Ordering::Relaxed);
+
+        self.con.send_datagram(data).map_err(io::Error::other)
+    }
+
+    async fn feed_control_packet(&mut self, packet: ControlPacket) -> io::Result<()> {
+        self.framed.feed(Packet::ControlPacket(packet)).await
+    }
+
+    async fn flush(&mut self) -> io::Result<()> {
+        self.framed.flush().await
+    }
 }
 
 #[cfg(test)]
@@ -291,6 +453,9 @@ impl DuplexStream {
 
 #[cfg(test)]
 impl Connection for DuplexStream {
+    type ReadHalf = DuplexStreamReadHalf;
+    type WriteHalf = DuplexStreamWriteHalf;
+
     async fn feed_data_packet(&mut self, packet: DataPacket) -> io::Result<()> {
         self.framed.feed(Packet::DataPacket(packet)).await
     }
@@ -314,4 +479,45 @@ impl Connection for DuplexStream {
     fn static_link_cost(&self) -> Result<u16, io::Error> {
         Ok(1)
     }
+
+    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
+        let (tx, rx) = self.framed.split();
+
+        (
+            DuplexStreamReadHalf { framed: rx },
+            DuplexStreamWriteHalf { framed: tx },
+        )
+    }
+}
+
+#[cfg(test)]
+pub struct DuplexStreamReadHalf {
+    framed: SplitStream<Framed<tokio::io::DuplexStream, packet::Codec>>,
+}
+
+#[cfg(test)]
+pub struct DuplexStreamWriteHalf {
+    framed: SplitSink<Framed<tokio::io::DuplexStream, packet::Codec>, packet::Packet>,
+}
+
+#[cfg(test)]
+impl ConnectionReadHalf for DuplexStreamReadHalf {
+    async fn receive_packet(&mut self) -> Option<io::Result<Packet>> {
+        self.framed.next().await
+    }
+}
+
+#[cfg(test)]
+impl ConnectionWriteHalf for DuplexStreamWriteHalf {
+    async fn feed_data_packet(&mut self, packet: DataPacket) -> io::Result<()> {
+        self.framed.feed(Packet::DataPacket(packet)).await
+    }
+
+    async fn feed_control_packet(&mut self, packet: ControlPacket) -> io::Result<()> {
+        self.framed.feed(Packet::ControlPacket(packet)).await
+    }
+
+    async fn flush(&mut self) -> io::Result<()> {
+        self.framed.flush().await
+    }
 }

mycelium/src/connection/tls.rs

@@ -4,7 +4,10 @@ use std::{
     sync::{atomic::AtomicU64, Arc},
 };
 
-use futures::{SinkExt, StreamExt};
+use futures::{
+    stream::{SplitSink, SplitStream},
+    SinkExt, StreamExt,
+};
 use tokio::net::TcpStream;
 use tokio_util::codec::Framed;
 
@@ -36,6 +39,9 @@ impl TlsStream {
 }
 
 impl super::Connection for TlsStream {
+    type ReadHalf = TlsStreamReadHalf;
+    type WriteHalf = TlsStreamWriteHalf;
+
     async fn feed_data_packet(&mut self, packet: crate::packet::DataPacket) -> io::Result<()> {
         self.framed.feed(Packet::DataPacket(packet)).await
     }
@@ -68,4 +74,47 @@ impl super::Connection for TlsStream {
             SocketAddr::V6(_) => super::PACKET_PROCESSING_COST_IP6_TCP,
         })
     }
+
+    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
+        let (tx, rx) = self.framed.split();
+
+        (
+            TlsStreamReadHalf { framed: rx },
+            TlsStreamWriteHalf { framed: tx },
+        )
+    }
+}
+
+pub struct TlsStreamReadHalf {
+    framed: SplitStream<Framed<Tracked<tokio_openssl::SslStream<TcpStream>>, packet::Codec>>,
+}
+
+pub struct TlsStreamWriteHalf {
+    framed: SplitSink<
+        Framed<Tracked<tokio_openssl::SslStream<TcpStream>>, packet::Codec>,
+        packet::Packet,
+    >,
+}
+
+impl super::ConnectionReadHalf for TlsStreamReadHalf {
+    async fn receive_packet(&mut self) -> Option<io::Result<crate::packet::Packet>> {
+        self.framed.next().await
+    }
+}
+
+impl super::ConnectionWriteHalf for TlsStreamWriteHalf {
+    async fn feed_data_packet(&mut self, packet: crate::packet::DataPacket) -> io::Result<()> {
+        self.framed.feed(Packet::DataPacket(packet)).await
+    }
+
+    async fn feed_control_packet(
+        &mut self,
+        packet: crate::packet::ControlPacket,
+    ) -> io::Result<()> {
+        self.framed.feed(Packet::ControlPacket(packet)).await
+    }
+
+    async fn flush(&mut self) -> io::Result<()> {
+        self.framed.flush().await
+    }
 }