lightway-core: io: Add a CoW abstraction for sending buffer

Depending on the specific implementation of the trait they may or may not need
an owned version of the buffer. Since we have one already in the core in some
paths we can expose that through the stack and avoid some extra allocations.

In the UDP send path we can easily and cheaply freeze the `BytesMut` buffer to
get an owned `Bytes` (with a little more overhead while `aggressive_send` is
enabled). However in the TCP path the use of `SendBuffer` makes this harder (if
not impossible) to achieve (cheaply at least).

So add `CowBytes` which can contain either a `Bytes` or a `&[u8]` slice.

Note that right now every consumer still uses the byte slice.

Author: xv-ian-c

lightway-client/src/io/outside/tcp.rs

@@ -4,7 +4,7 @@ use std::{net::SocketAddr, sync::Arc};
 use tokio::net::TcpStream;
 
 use super::OutsideIO;
-use lightway_core::{IOCallbackResult, OutsideIOSendCallback, OutsideIOSendCallbackArg};
+use lightway_core::{CowBytes, IOCallbackResult, OutsideIOSendCallback, OutsideIOSendCallbackArg};
 
 pub struct Tcp(tokio::net::TcpStream, SocketAddr);
 
@@ -58,8 +58,8 @@ impl OutsideIO for Tcp {
 }
 
 impl OutsideIOSendCallback for Tcp {
-    fn send(&self, buf: &[u8]) -> IOCallbackResult<usize> {
-        match self.0.try_write(buf) {
+    fn send(&self, buf: CowBytes) -> IOCallbackResult<usize> {
+        match self.0.try_write(buf.as_bytes()) {
             Ok(nr) => IOCallbackResult::Ok(nr),
             Err(err) if matches!(err.kind(), std::io::ErrorKind::WouldBlock) => {
                 IOCallbackResult::WouldBlock

lightway-client/src/io/outside/udp.rs

@@ -5,7 +5,7 @@ use tokio::net::UdpSocket;
 
 use super::OutsideIO;
 use lightway_app_utils::sockopt;
-use lightway_core::{IOCallbackResult, OutsideIOSendCallback, OutsideIOSendCallbackArg};
+use lightway_core::{CowBytes, IOCallbackResult, OutsideIOSendCallback, OutsideIOSendCallbackArg};
 
 pub struct Udp {
     sock: tokio::net::UdpSocket,
@@ -67,8 +67,8 @@ impl OutsideIO for Udp {
 }
 
 impl OutsideIOSendCallback for Udp {
-    fn send(&self, buf: &[u8]) -> IOCallbackResult<usize> {
-        match self.sock.try_send_to(buf, self.peer_addr) {
+    fn send(&self, buf: CowBytes) -> IOCallbackResult<usize> {
+        match self.sock.try_send_to(buf.as_bytes(), self.peer_addr) {
             Ok(nr) => IOCallbackResult::Ok(nr),
             Err(err) if matches!(err.kind(), std::io::ErrorKind::WouldBlock) => {
                 IOCallbackResult::WouldBlock

lightway-core/src/connection/io_adapter.rs

@@ -5,7 +5,8 @@ use more_asserts::*;
 use wolfssl::IOCallbackResult;
 
 use crate::{
-    plugin::PluginList, wire, ConnectionType, OutsideIOSendCallbackArg, PluginResult, Version,
+    plugin::PluginList, wire, ConnectionType, CowBytes, OutsideIOSendCallbackArg, PluginResult,
+    Version,
 };
 
 pub(crate) struct SendBuffer {
@@ -164,26 +165,28 @@ impl WolfSSLIOAdapter {
             }
         }
 
+        let b = b.freeze();
+
         // Send header + buf. If we are in aggressive mode we send it
         // a total of three times. On any send error we return
         // immediately without the remaining tries, otherwise we
         // return the result of the final attempt.
 
         if self.aggressive_send {
-            match self.io.send(&b[..]) {
+            match self.io.send(CowBytes::Owned(b.clone())) {
                 IOCallbackResult::Ok(_) => {}
                 wb @ IOCallbackResult::WouldBlock => return wb,
                 err @ IOCallbackResult::Err(_) => return err,
             }
 
-            match self.io.send(&b[..]) {
+            match self.io.send(CowBytes::Owned(b.clone())) {
                 IOCallbackResult::Ok(_) => {}
                 wb @ IOCallbackResult::WouldBlock => return wb,
                 err @ IOCallbackResult::Err(_) => return err,
             }
         }
 
-        match self.io.send(&b[..]) {
+        match self.io.send(CowBytes::Owned(b)) {
             IOCallbackResult::Ok(n) => {
                 // We've sent `n` bytes successfully out of
                 // `wire::Header::WIRE_SIZE` + `b.len()` that we
@@ -250,7 +253,7 @@ impl WolfSSLIOAdapter {
             debug_assert_le!(send_buffer.original_len(), buf.len());
         }
 
-        match self.io.send(send_buffer.as_bytes()) {
+        match self.io.send(CowBytes::Borrowed(send_buffer.as_bytes())) {
             IOCallbackResult::Ok(n) if n == send_buffer.actual_len() => {
                 // We've now sent everything we were originally
                 // asked to, so signal completion of that original
@@ -335,7 +338,8 @@ mod tests {
     }
 
     impl OutsideIOSendCallback for FakeOutsideIOSend {
-        fn send(&self, buf: &[u8]) -> IOCallbackResult<usize> {
+        fn send(&self, buf: CowBytes) -> IOCallbackResult<usize> {
+            let buf = buf.as_bytes();
             let (fakes, sent) = &mut *self.0.lock().unwrap();
             match fakes.pop_front() {
                 Some(IOCallbackResult::Ok(n)) => {

lightway-core/src/io.rs

@@ -1,6 +1,6 @@
 use std::{net::SocketAddr, sync::Arc};
 
-use bytes::BytesMut;
+use bytes::{Bytes, BytesMut};
 use wolfssl::IOCallbackResult;
 
 /// Application provided callback used to send inside data.
@@ -20,6 +20,33 @@ pub trait InsideIOSendCallback<AppState> {
 /// Convenience type to use as function arguments
 pub type InsideIOSendCallbackArg<AppState> = Arc<dyn InsideIOSendCallback<AppState> + Send + Sync>;
 
+/// A byte buffer to be sent, may be owned or borrowed.
+pub enum CowBytes<'a> {
+    /// An owned buffer
+    Owned(Bytes),
+    /// A borrowed buffer
+    Borrowed(&'a [u8]),
+}
+
+impl<'a> CowBytes<'a> {
+    /// Convert this buffer into an owned `Bytes`. Cheap if this
+    /// instance if `::Owned`, but copied if not.
+    pub fn into_owned(self) -> Bytes {
+        match self {
+            CowBytes::Owned(b) => b,
+            CowBytes::Borrowed(b) => Bytes::copy_from_slice(b),
+        }
+    }
+
+    /// Gain access to the underlying byte buffer.
+    pub fn as_bytes(&self) -> &[u8] {
+        match self {
+            CowBytes::Owned(b) => b.as_ref(),
+            CowBytes::Borrowed(b) => b,
+        }
+    }
+}
+
 /// Application provided callback used to send outside data.
 pub trait OutsideIOSendCallback {
     /// Called when Lightway wishes to send some outside data
@@ -30,7 +57,7 @@ pub trait OutsideIOSendCallback {
     /// [`IOCallbackResult::WouldBlock`].
     ///
     /// This is the same method as [`wolfssl::IOCallbacks::send`].
-    fn send(&self, buf: &[u8]) -> IOCallbackResult<usize>;
+    fn send(&self, buf: CowBytes) -> IOCallbackResult<usize>;
 
     /// Get the peer's [`SocketAddr`]
     fn peer_addr(&self) -> SocketAddr;

lightway-core/src/lib.rs

@@ -38,7 +38,8 @@ pub use context::{
     ServerAuthArg, ServerAuthHandle, ServerAuthResult, ServerContext, ServerContextBuilder,
 };
 pub use io::{
-    InsideIOSendCallback, InsideIOSendCallbackArg, OutsideIOSendCallback, OutsideIOSendCallbackArg,
+    CowBytes, InsideIOSendCallback, InsideIOSendCallbackArg, OutsideIOSendCallback,
+    OutsideIOSendCallbackArg,
 };
 pub use packet::OutsidePacket;
 pub use plugin::{

lightway-core/tests/connection.rs

@@ -106,7 +106,8 @@ impl TestSock for TestDatagramSock {
 }
 
 impl OutsideIOSendCallback for TestDatagramSock {
-    fn send(&self, buf: &[u8]) -> IOCallbackResult<usize> {
+    fn send(&self, buf: CowBytes) -> IOCallbackResult<usize> {
+        let buf = buf.as_bytes();
         match self.0.try_send(buf) {
             Ok(nr) => IOCallbackResult::Ok(nr),
             Err(err) if matches!(err.kind(), std::io::ErrorKind::WouldBlock) => {
@@ -156,8 +157,8 @@ impl TestSock for TestStreamSock {
 }
 
 impl OutsideIOSendCallback for TestStreamSock {
-    fn send(&self, buf: &[u8]) -> IOCallbackResult<usize> {
-        match self.0.try_write(buf) {
+    fn send(&self, buf: CowBytes) -> IOCallbackResult<usize> {
+        match self.0.try_write(buf.as_bytes()) {
             Ok(nr) => IOCallbackResult::Ok(nr),
             Err(err) if matches!(err.kind(), std::io::ErrorKind::WouldBlock) => {
                 IOCallbackResult::WouldBlock

lightway-server/src/io/outside/tcp.rs

@@ -4,7 +4,7 @@ use anyhow::{anyhow, Result};
 use async_trait::async_trait;
 use bytes::BytesMut;
 use lightway_core::{
-    ConnectionType, IOCallbackResult, OutsideIOSendCallback, OutsidePacket, Version,
+    ConnectionType, CowBytes, IOCallbackResult, OutsideIOSendCallback, OutsidePacket, Version,
     MAX_OUTSIDE_MTU,
 };
 use socket2::SockRef;
@@ -21,8 +21,8 @@ struct TcpStream {
 }
 
 impl OutsideIOSendCallback for TcpStream {
-    fn send(&self, buf: &[u8]) -> IOCallbackResult<usize> {
-        match self.sock.try_write(buf) {
+    fn send(&self, buf: CowBytes) -> IOCallbackResult<usize> {
+        match self.sock.try_write(buf.as_bytes()) {
             Ok(nr) => IOCallbackResult::Ok(nr),
             Err(err) if matches!(err.kind(), std::io::ErrorKind::WouldBlock) => {
                 IOCallbackResult::WouldBlock

lightway-server/src/io/outside/udp.rs

@@ -11,8 +11,8 @@ use bytes::BytesMut;
 use bytesize::ByteSize;
 use lightway_app_utils::sockopt::socket_enable_pktinfo;
 use lightway_core::{
-    ConnectionType, Header, IOCallbackResult, OutsideIOSendCallback, OutsidePacket, SessionId,
-    Version, MAX_OUTSIDE_MTU,
+    ConnectionType, CowBytes, Header, IOCallbackResult, OutsideIOSendCallback, OutsidePacket,
+    SessionId, Version, MAX_OUTSIDE_MTU,
 };
 use socket2::{MaybeUninitSlice, MsgHdr, MsgHdrMut, SockAddr, SockRef};
 use tokio::io::Interest;
@@ -87,9 +87,9 @@ struct UdpSocket {
 }
 
 impl OutsideIOSendCallback for UdpSocket {
-    fn send(&self, buf: &[u8]) -> IOCallbackResult<usize> {
+    fn send(&self, buf: CowBytes) -> IOCallbackResult<usize> {
         let peer_addr = self.peer_addr.read().unwrap();
-        send_to_socket(&self.sock, buf, &peer_addr.1, self.reply_pktinfo)
+        send_to_socket(&self.sock, buf.as_bytes(), &peer_addr.1, self.reply_pktinfo)
     }
 
     fn peer_addr(&self) -> SocketAddr {