async: upgrade tokio to 1.0

upgrade tokio from 0.2 to 1.0

Signed-off-by: Tim Zhang <tim@hyper.sh>

Author: Tim-Zhang

Cargo.toml

@@ -20,9 +20,9 @@ byteorder = "1.3.2"
 thiserror = "1.0"
 
 async-trait = { version = "0.1.31", optional = true }
-tokio = { version = "0.2", features = ["rt-threaded", "sync", "uds", "stream", "macros", "io-util", "time"], optional = true }
+tokio = { version = "1", features = ["rt", "sync", "io-util", "macros", "time"], optional = true }
 futures = { version = "0.3", optional = true }
-tokio-vsock = { version = "0.2.1", optional = true }
+tokio-vsock = { version = "0.3", optional = true }
 
 [build-dependencies]
 protobuf-codegen-pure = "2.14.0"

src/asynchronous/client.rs

@@ -18,10 +18,10 @@ use crate::asynchronous::stream::{receive, to_req_buf};
 use crate::r#async::utils;
 use tokio::{
     self,
-    io::split,
-    prelude::*,
+    io::{split, AsyncWriteExt},
     sync::mpsc::{channel, Receiver, Sender},
     sync::Notify,
+    time::timeout,
 };
 
 type RequestSender = Sender<(Vec<u8>, Sender<Result<Vec<u8>>>)>;
@@ -54,7 +54,7 @@ impl Client {
         tokio::spawn(async move {
             let mut stream_id: u32 = 1;
 
-            while let Some((body, mut resp_tx)) = rx.recv().await {
+            while let Some((body, resp_tx)) = rx.recv().await {
                 let current_stream_id = stream_id;
                 stream_id += 2;
 
@@ -84,7 +84,7 @@ impl Client {
 
             // rx.recv will abort when client.req_tx and client is dropped.
             // notify the response-receiver to quit at this time.
-            notify.notify();
+            notify.notify_one();
         });
 
         // Response receiver
@@ -99,7 +99,7 @@ impl Client {
                         match res {
                             Ok((header, body)) => {
                                 tokio::spawn(async move {
-                                    let mut resp_tx2;
+                                    let resp_tx2;
                                     {
                                         let mut map = req_map.lock().unwrap();
                                         let resp_tx = match map.get(&header.stream_id) {
@@ -163,14 +163,13 @@ impl Client {
                 .await
                 .ok_or_else(|| Error::Others("Recive packet from recver error".to_string()))?
         } else {
-            let timeout = tokio::time::delay_for(Duration::from_nanos(req.timeout_nano as u64));
-
-            tokio::select! {
-                result = rx.recv() => {
-                    result.ok_or_else(|| Error::Others("Recive packet from recver error".to_string()))?
-                }
-                _ = timeout => {
-                    return Err(Error::Others("Recive packet from recver error: timeout".to_string()));
+            match timeout(Duration::from_nanos(req.timeout_nano as u64), rx.recv()).await {
+                Ok(result) => result
+                    .ok_or_else(|| Error::Others("Recive packet from recver error".to_string()))?,
+                Err(_) => {
+                    return Err(Error::Others(
+                        "Recive packet from recver error: timeout".to_string(),
+                    ))
                 }
             }
         };

src/asynchronous/mod.rs

@@ -10,6 +10,7 @@ pub mod server;
 pub mod stream;
 #[macro_use]
 pub mod utils;
+mod unix_incoming;
 
 #[doc(inline)]
 pub use crate::r#async::client::Client;

src/asynchronous/server.rs

@@ -10,21 +10,21 @@ use std::os::unix::io::RawFd;
 use std::sync::Arc;
 
 use crate::asynchronous::stream::{receive, respond, respond_with_status};
+use crate::asynchronous::unix_incoming::UnixIncoming;
 use crate::common::{self, Domain, MESSAGE_TYPE_REQUEST};
 use crate::error::{get_status, Error, Result};
 use crate::r#async::{MethodHandler, TtrpcContext};
 use crate::ttrpc::{Code, Request};
 use crate::MessageHeader;
+use futures::stream::Stream;
 use futures::StreamExt as _;
 use std::marker::Unpin;
 use std::os::unix::io::{AsRawFd, FromRawFd};
 use std::os::unix::net::UnixListener as SysUnixListener;
 use tokio::{
     self,
-    io::split,
+    io::{split, AsyncRead, AsyncWrite, AsyncWriteExt},
     net::UnixListener,
-    prelude::*,
-    stream::Stream,
     sync::mpsc::{channel, Receiver, Sender},
     sync::watch,
 };
@@ -116,9 +116,15 @@ impl Server {
                 unsafe {
                     sys_unix_listener = SysUnixListener::from_raw_fd(listenfd);
                 }
-                let unix_listener = UnixListener::from_std(sys_unix_listener).unwrap();
+                sys_unix_listener
+                    .set_nonblocking(true)
+                    .map_err(err_to_others_err!(e, "set_nonblocking error "))?;
+                let unix_listener = UnixListener::from_std(sys_unix_listener)
+                    .map_err(err_to_others_err!(e, "from_std error "))?;
 
-                self.do_start(listenfd, unix_listener).await
+                let incoming = UnixIncoming::new(unix_listener);
+
+                self.do_start(listenfd, incoming).await
             }
             Some(Domain::Vsock) => {
                 let incoming;
@@ -201,10 +207,10 @@ impl Server {
                                                         }
                                                     }
                                                 }
-                                                v = close_conn_rx.recv() => {
+                                                v = close_conn_rx.changed() => {
                                                     // 0 is the init value of this watch, not a valid signal
-                                                    // is_none means the tx was dropped.
-                                                    if v.is_none() || v.unwrap() != 0 {
+                                                    // is_err means the tx was dropped.
+                                                    if v.is_err() || *close_conn_rx.borrow() != 0 {
                                                         info!("Stop accepting new connections.");
                                                         break;
                                                     }
@@ -227,7 +233,7 @@ impl Server {
                         }
                     }
                     fd_tx = stop_listen_rx.recv() => {
-                        if let Some(mut fd_tx) = fd_tx {
+                        if let Some(fd_tx) = fd_tx {
                             // dup fd to keep the listener open
                             // or the listener will be closed when the incoming was dropped.
                             let dup_fd = unistd::dup(incoming.as_raw_fd()).unwrap();
@@ -254,7 +260,7 @@ impl Server {
 
     pub async fn disconnect(&mut self) {
         if let Some(tx) = self.disconnect_tx.take() {
-            tx.broadcast(1).ok();
+            tx.send(1).ok();
         }
 
         if let Some(mut rx) = self.all_conn_done_rx.take() {
@@ -263,7 +269,7 @@ impl Server {
     }
 
     pub async fn stop_listen(&mut self) {
-        if let Some(mut tx) = self.stop_listen_tx.take() {
+        if let Some(tx) = self.stop_listen_tx.take() {
             let (fd_tx, mut fd_rx) = channel(1);
             tx.send(fd_tx).await.unwrap();
 

src/asynchronous/stream.rs

@@ -120,7 +120,7 @@ fn get_response_body(res: &Response) -> Result<Vec<u8>> {
 }
 
 pub async fn respond(
-    mut tx: tokio::sync::mpsc::Sender<Vec<u8>>,
+    tx: tokio::sync::mpsc::Sender<Vec<u8>>,
     stream_id: u32,
     body: Vec<u8>,
 ) -> Result<()> {
@@ -132,7 +132,7 @@ pub async fn respond(
 }
 
 pub async fn respond_with_status(
-    mut tx: tokio::sync::mpsc::Sender<Vec<u8>>,
+    tx: tokio::sync::mpsc::Sender<Vec<u8>>,
     stream_id: u32,
     status: Status,
 ) -> Result<()> {

src/asynchronous/unix_incoming.rs

@@ -0,0 +1,43 @@
+// Copyright (c) 2021 Ant Group
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+//! Because Tokio has removed UnixIncoming since version 0.3,
+//! we define the UnixIncoming and implement the Stream for UnixIncoming.
+
+use std::io;
+use std::os::unix::io::{AsRawFd, RawFd};
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use futures::{ready, Stream};
+use tokio::net::{UnixListener, UnixStream};
+
+/// Stream of listeners
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct UnixIncoming {
+    inner: UnixListener,
+}
+
+impl UnixIncoming {
+    pub fn new(listener: UnixListener) -> Self {
+        Self { inner: listener }
+    }
+}
+
+impl Stream for UnixIncoming {
+    type Item = io::Result<UnixStream>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let (socket, _) = ready!(self.inner.poll_accept(cx))?;
+        Poll::Ready(Some(Ok(socket)))
+    }
+}
+
+impl AsRawFd for UnixIncoming {
+    fn as_raw_fd(&self) -> RawFd {
+        self.inner.as_raw_fd()
+    }
+}

src/asynchronous/utils.rs

@@ -121,5 +121,8 @@ pub fn new_unix_stream_from_raw_fd(fd: RawFd) -> UnixStream {
     unsafe {
         std_stream = std::os::unix::net::UnixStream::from_raw_fd(fd);
     }
+    // Notice: There is a big change between tokio 1.0 and 0.2
+    // we must set nonblocking by ourselves in tokio 1.0
+    std_stream.set_nonblocking(true).unwrap();
     UnixStream::from_std(std_stream).unwrap()
 }