Tokio 0.3 Upgrade (#2319)

Co-authored-by: Urhengulas <johann.hemmann@code.berlin>
Co-authored-by: Eliza Weisman <eliza@buoyant.io>

src/common/drain.rs

@@ -1,20 +1,13 @@
 use std::mem;
 
 use pin_project::pin_project;
+use tokio::stream::Stream;
 use tokio::sync::{mpsc, watch};
 
 use super::{task, Future, Never, Pin, Poll};
 
-// Sentinel value signaling that the watch is still open
-#[derive(Clone, Copy)]
-enum Action {
-    Open,
-    // Closed isn't sent via the `Action` type, but rather once
-    // the watch::Sender is dropped.
-}
-
 pub fn channel() -> (Signal, Watch) {
-    let (tx, rx) = watch::channel(Action::Open);
+    let (tx, rx) = watch::channel(());
     let (drained_tx, drained_rx) = mpsc::channel(1);
     (
         Signal {
@@ -27,17 +20,19 @@ pub fn channel() -> (Signal, Watch) {
 
 pub struct Signal {
     drained_rx: mpsc::Receiver<Never>,
-    _tx: watch::Sender<Action>,
+    _tx: watch::Sender<()>,
 }
 
+#[pin_project::pin_project]
 pub struct Draining {
+    #[pin]
     drained_rx: mpsc::Receiver<Never>,
 }
 
 #[derive(Clone)]
 pub struct Watch {
     drained_tx: mpsc::Sender<Never>,
-    rx: watch::Receiver<Action>,
+    rx: watch::Receiver<()>,
 }
 
 #[allow(missing_debug_implementations)]
@@ -46,7 +41,8 @@ pub struct Watching<F, FN> {
     #[pin]
     future: F,
     state: State<FN>,
-    watch: Watch,
+    watch: Pin<Box<dyn Future<Output = ()> + Send + Sync>>,
+    _drained_tx: mpsc::Sender<Never>,
 }
 
 enum State<F> {
@@ -66,8 +62,8 @@ impl Signal {
 impl Future for Draining {
     type Output = ();
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
-        match ready!(self.drained_rx.poll_recv(cx)) {
+    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
+        match ready!(self.project().drained_rx.poll_next(cx)) {
             Some(never) => match never {},
             None => Poll::Ready(()),
         }
@@ -80,10 +76,14 @@ impl Watch {
         F: Future,
         FN: FnOnce(Pin<&mut F>),
     {
+        let Self { drained_tx, mut rx } = self;
         Watching {
             future,
             state: State::Watch(on_drain),
-            watch: self,
+            watch: Box::pin(async move {
+                let _ = rx.changed().await;
+            }),
+            _drained_tx: drained_tx,
         }
     }
 }
@@ -100,12 +100,12 @@ where
         loop {
             match mem::replace(me.state, State::Draining) {
                 State::Watch(on_drain) => {
-                    match me.watch.rx.poll_recv_ref(cx) {
-                        Poll::Ready(None) => {
+                    match Pin::new(&mut me.watch).poll(cx) {
+                        Poll::Ready(()) => {
                             // Drain has been triggered!
                             on_drain(me.future.as_mut());
                         }
-                        Poll::Ready(Some(_ /*State::Open*/)) | Poll::Pending => {
+                        Poll::Pending => {
                             *me.state = State::Watch(on_drain);
                             return me.future.poll(cx);
                         }

src/common/io/rewind.rs

@@ -2,7 +2,7 @@ use std::marker::Unpin;
 use std::{cmp, io};
 
 use bytes::{Buf, Bytes};
-use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 
 use crate::common::{task, Pin, Poll};
 
@@ -37,36 +37,33 @@ impl<T> Rewind<T> {
         (self.inner, self.pre.unwrap_or_else(Bytes::new))
     }
 
-    pub(crate) fn get_mut(&mut self) -> &mut T {
-        &mut self.inner
-    }
+    // pub(crate) fn get_mut(&mut self) -> &mut T {
+    //     &mut self.inner
+    // }
 }
 
 impl<T> AsyncRead for Rewind<T>
 where
     T: AsyncRead + Unpin,
 {
-    #[inline]
-    unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [std::mem::MaybeUninit<u8>]) -> bool {
-        self.inner.prepare_uninitialized_buffer(buf)
-    }
-
     fn poll_read(
         mut self: Pin<&mut Self>,
         cx: &mut task::Context<'_>,
-        buf: &mut [u8],
-    ) -> Poll<io::Result<usize>> {
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
         if let Some(mut prefix) = self.pre.take() {
             // If there are no remaining bytes, let the bytes get dropped.
             if !prefix.is_empty() {
-                let copy_len = cmp::min(prefix.len(), buf.len());
-                prefix.copy_to_slice(&mut buf[..copy_len]);
+                let copy_len = cmp::min(prefix.len(), buf.remaining());
+                // TODO: There should be a way to do following two lines cleaner...
+                buf.put_slice(&prefix[..copy_len]);
+                prefix.advance(copy_len);
                 // Put back whats left
                 if !prefix.is_empty() {
                     self.pre = Some(prefix);
                 }
 
-                return Poll::Ready(Ok(copy_len));
+                return Poll::Ready(Ok(()));
             }
         }
         Pin::new(&mut self.inner).poll_read(cx, buf)
@@ -92,15 +89,6 @@ where
     fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
         Pin::new(&mut self.inner).poll_shutdown(cx)
     }
-
-    #[inline]
-    fn poll_write_buf<B: Buf>(
-        mut self: Pin<&mut Self>,
-        cx: &mut task::Context<'_>,
-        buf: &mut B,
-    ) -> Poll<io::Result<usize>> {
-        Pin::new(&mut self.inner).poll_write_buf(cx, buf)
-    }
 }
 
 #[cfg(test)]