feat(s2n-quic-dc): add channel recv buffer impl

Author: camshaft

dc/s2n-quic-dc/src/stream/client/tokio.rs

@@ -161,5 +161,6 @@ where
 #[inline]
 fn recv_buffer() -> recv::shared::RecvBuffer {
     // TODO replace this with a parameter once everything is in place
-    recv::buffer::Local::new(msg::recv::Message::new(9000), None)
+    let recv_buffer = recv::buffer::Local::new(msg::recv::Message::new(9000), None);
+    recv::buffer::Either::A(recv_buffer)
 }

dc/s2n-quic-dc/src/stream/recv/buffer.rs

@@ -8,9 +8,11 @@ use crate::{
 use core::task::{Context, Poll};
 use std::io;
 
+pub mod channel;
 mod dispatch;
 mod local;
 
+pub use channel::Channel;
 pub use dispatch::Dispatch;
 pub use local::Local;
 
@@ -36,7 +38,6 @@ pub trait Buffer {
         R: Dispatch;
 }
 
-#[allow(dead_code)] // TODO remove this once we start using the channel buffer
 pub enum Either<A, B> {
     A(A),
     B(B),

dc/s2n-quic-dc/src/stream/recv/buffer/channel.rs

@@ -0,0 +1,99 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use super::Dispatch;
+use crate::{
+    event,
+    socket::recv::descriptor::Filled,
+    stream::{recv, socket::Socket, TransportFeatures},
+    sync::mpsc::Receiver,
+};
+use core::task::{Context, Poll};
+use s2n_quic_core::ensure;
+use std::{collections::VecDeque, io};
+
+pub struct Channel {
+    pending: VecDeque<Filled>,
+    receiver: Receiver<Filled>,
+}
+
+impl super::Buffer for Channel {
+    #[inline]
+    fn is_empty(&self) -> bool {
+        self.pending.is_empty()
+    }
+
+    #[inline]
+    fn poll_fill<S, Pub>(
+        &mut self,
+        cx: &mut Context,
+        socket: &S,
+        publisher: &mut Pub,
+    ) -> Poll<io::Result<usize>>
+    where
+        S: ?Sized + Socket,
+        Pub: event::ConnectionPublisher,
+    {
+        // check if we have any pending packets
+        debug_assert!(self.pending.is_empty(), "pending packets should be empty");
+        ensure!(self.pending.is_empty(), Ok(1).into());
+
+        let capacity = u16::MAX as usize;
+
+        // the socket isn't actually used since we're relying on another task to fill the `receiver` channel
+        let _ = socket;
+
+        let result = self
+            .receiver
+            .poll_swap(cx, &mut self.pending)
+            .map_err(|_err| io::Error::from(io::ErrorKind::BrokenPipe));
+
+        match result {
+            Poll::Ready(Ok(())) => {
+                let committed_len = self
+                    .pending
+                    .iter()
+                    .map(|segment| segment.len() as usize)
+                    .sum::<usize>();
+                publisher.on_stream_read_socket_flushed(event::builder::StreamReadSocketFlushed {
+                    capacity,
+                    committed_len,
+                });
+                Ok(committed_len).into()
+            }
+            Poll::Ready(Err(error)) => {
+                let errno = error.raw_os_error();
+                publisher.on_stream_read_socket_errored(event::builder::StreamReadSocketErrored {
+                    capacity,
+                    errno,
+                });
+                Err(error).into()
+            }
+            Poll::Pending => {
+                publisher.on_stream_read_socket_blocked(event::builder::StreamReadSocketBlocked {
+                    capacity,
+                });
+                Poll::Pending
+            }
+        }
+    }
+
+    #[inline]
+    fn process<R>(&mut self, features: TransportFeatures, router: &mut R) -> Result<(), recv::Error>
+    where
+        R: Dispatch,
+    {
+        debug_assert!(
+            !features.is_stream(),
+            "only datagram oriented transport is supported"
+        );
+
+        for mut segment in self.pending.drain(..) {
+            let remote_addr = segment.remote_address().get();
+            let ecn = segment.ecn();
+            router.on_datagram_segment(&remote_addr, ecn, segment.payload_mut())?;
+        }
+
+        Ok(())
+    }
+}

dc/s2n-quic-dc/src/stream/recv/shared.rs

@@ -28,7 +28,7 @@ use std::{
     },
 };
 
-pub type RecvBuffer = recv::buffer::Local;
+pub type RecvBuffer = recv::buffer::Either<recv::buffer::Local, recv::buffer::Channel>;
 
 /// Who will send ACKs?
 #[derive(Clone, Copy, Debug, Default)]

dc/s2n-quic-dc/src/stream/server/tokio/tcp/worker.rs

@@ -317,6 +317,7 @@ impl WorkerState {
             // TCP doesn't use the route key so just pick 0
             let queue_id = VarInt::ZERO;
             let recv_buffer = recv::buffer::Local::new(recv_buffer.take(), None);
+            let recv_buffer = recv::buffer::Either::A(recv_buffer);
 
             let stream_builder = match endpoint::accept_stream(
                 now,

dc/s2n-quic-dc/src/stream/server/tokio/udp.rs

@@ -114,6 +114,7 @@ where
         // TODO allocate a queue for this stream
         let queue_id = VarInt::ZERO;
         let recv_buffer = recv::buffer::Local::new(self.recv_buffer.take(), Some(handshake));
+        let recv_buffer = recv::buffer::Either::A(recv_buffer);
 
         let stream = match endpoint::accept_stream(
             now,

dc/s2n-quic-dc/src/stream/socket/handle.rs

@@ -1,7 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use super::{Protocol, TransportFeatures};
+use super::{recv::descriptor, Protocol, TransportFeatures};
 use crate::msg::{self, addr::Addr, cmsg};
 use core::task::{Context, Poll};
 use s2n_quic_core::inet::ExplicitCongestionNotification;
@@ -31,6 +31,23 @@ pub trait Socket: 'static + Send + Sync {
     /// Returns the amount of buffered data on the socket
     fn poll_peek_len(&self, cx: &mut Context) -> Poll<io::Result<usize>>;
 
+    #[inline]
+    fn poll_recv_desc(
+        &self,
+        cx: &mut Context,
+        desc: descriptor::Unfilled,
+    ) -> Poll<Result<descriptor::Segments, (descriptor::Unfilled, io::Error)>> {
+        let segments = desc.recv_with(|addr, cmsg, buffer| {
+            match self.poll_recv(cx, addr, cmsg, &mut [buffer]) {
+                Poll::Pending => Err(io::ErrorKind::WouldBlock.into()),
+                Poll::Ready(Ok(len)) => Ok(len),
+                Poll::Ready(Err(err)) => Err(err),
+            }
+        })?;
+
+        Ok(segments).into()
+    }
+
     #[inline]
     fn poll_recv_buffer(
         &self,