Tokio 0.3 Upgrade (#2319)
Co-authored-by: Urhengulas <johann.hemmann@code.berlin> Co-authored-by: Eliza Weisman <eliza@buoyant.io>
This commit is contained in:
		| @@ -809,9 +809,9 @@ where | ||||
|     type Output = Result<Connection<I, S, E>, FE>; | ||||
|  | ||||
|     fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> { | ||||
|         let me = self.project(); | ||||
|         let mut me = self.project(); | ||||
|         let service = ready!(me.future.poll(cx))?; | ||||
|         let io = me.io.take().expect("polled after complete"); | ||||
|         let io = Option::take(&mut me.io).expect("polled after complete"); | ||||
|         Poll::Ready(Ok(me.protocol.serve_connection(io, service))) | ||||
|     } | ||||
| } | ||||
|   | ||||
| @@ -4,7 +4,7 @@ use std::net::{SocketAddr, TcpListener as StdTcpListener}; | ||||
| use std::time::Duration; | ||||
|  | ||||
| use tokio::net::TcpListener; | ||||
| use tokio::time::Delay; | ||||
| use tokio::time::Sleep; | ||||
|  | ||||
| use crate::common::{task, Future, Pin, Poll}; | ||||
|  | ||||
| @@ -19,7 +19,7 @@ pub struct AddrIncoming { | ||||
|     sleep_on_errors: bool, | ||||
|     tcp_keepalive_timeout: Option<Duration>, | ||||
|     tcp_nodelay: bool, | ||||
|     timeout: Option<Delay>, | ||||
|     timeout: Option<Sleep>, | ||||
| } | ||||
|  | ||||
| impl AddrIncoming { | ||||
| @@ -30,6 +30,10 @@ impl AddrIncoming { | ||||
|     } | ||||
|  | ||||
|     pub(super) fn from_std(std_listener: StdTcpListener) -> crate::Result<Self> { | ||||
|         // TcpListener::from_std doesn't set O_NONBLOCK | ||||
|         std_listener | ||||
|             .set_nonblocking(true) | ||||
|             .map_err(crate::Error::new_listen)?; | ||||
|         let listener = TcpListener::from_std(std_listener).map_err(crate::Error::new_listen)?; | ||||
|         let addr = listener.local_addr().map_err(crate::Error::new_listen)?; | ||||
|         Ok(AddrIncoming { | ||||
| @@ -98,9 +102,46 @@ impl AddrIncoming { | ||||
|             match ready!(self.listener.poll_accept(cx)) { | ||||
|                 Ok((socket, addr)) => { | ||||
|                     if let Some(dur) = self.tcp_keepalive_timeout { | ||||
|                         // Convert the Tokio `TcpStream` into a `socket2` socket | ||||
|                         // so we can call `set_keepalive`. | ||||
|                         // TODO(eliza): if Tokio's `TcpSocket` API grows a few | ||||
|                         // more methods in the future, hopefully we shouldn't | ||||
|                         // have to do the `from_raw_fd` dance any longer... | ||||
|                         #[cfg(unix)] | ||||
|                         let socket = unsafe { | ||||
|                             // Safety: `socket2`'s socket will try to close the | ||||
|                             // underlying fd when it's dropped. However, we | ||||
|                             // can't take ownership of the fd from the tokio | ||||
|                             // TcpStream, so instead we will call `into_raw_fd` | ||||
|                             // on the socket2 socket before dropping it. This | ||||
|                             // prevents it from trying to close the fd. | ||||
|                             use std::os::unix::io::{AsRawFd, FromRawFd}; | ||||
|                             socket2::Socket::from_raw_fd(socket.as_raw_fd()) | ||||
|                         }; | ||||
|                         #[cfg(windows)] | ||||
|                         let socket = unsafe { | ||||
|                             // Safety: `socket2`'s socket will try to close the | ||||
|                             // underlying SOCKET when it's dropped. However, we | ||||
|                             // can't take ownership of the SOCKET from the tokio | ||||
|                             // TcpStream, so instead we will call `into_raw_socket` | ||||
|                             // on the socket2 socket before dropping it. This | ||||
|                             // prevents it from trying to close the SOCKET. | ||||
|                             use std::os::windows::io::{AsRawSocket, FromRawSocket}; | ||||
|                             socket2::Socket::from_raw_socket(socket.as_raw_socket()) | ||||
|                         }; | ||||
|  | ||||
|                         // Actually set the TCP keepalive timeout. | ||||
|                         if let Err(e) = socket.set_keepalive(Some(dur)) { | ||||
|                             trace!("error trying to set TCP keepalive: {}", e); | ||||
|                         } | ||||
|  | ||||
|                         // Take ownershop of the fd/socket back from the socket2 | ||||
|                         // `Socket`, so that socket2 doesn't try to close it | ||||
|                         // when it's dropped. | ||||
|                         #[cfg(unix)] | ||||
|                         drop(std::os::unix::io::IntoRawFd::into_raw_fd(socket)); | ||||
|                         #[cfg(windows)] | ||||
|                         drop(std::os::windows::io::IntoRawSocket::into_raw_socket(socket)); | ||||
|                     } | ||||
|                     if let Err(e) = socket.set_nodelay(self.tcp_nodelay) { | ||||
|                         trace!("error trying to set TCP nodelay: {}", e); | ||||
| @@ -119,7 +160,7 @@ impl AddrIncoming { | ||||
|                         error!("accept error: {}", e); | ||||
|  | ||||
|                         // Sleep 1s. | ||||
|                         let mut timeout = tokio::time::delay_for(Duration::from_secs(1)); | ||||
|                         let mut timeout = tokio::time::sleep(Duration::from_secs(1)); | ||||
|  | ||||
|                         match Pin::new(&mut timeout).poll(cx) { | ||||
|                             Poll::Ready(()) => { | ||||
| @@ -181,19 +222,20 @@ impl fmt::Debug for AddrIncoming { | ||||
| } | ||||
|  | ||||
| mod addr_stream { | ||||
|     use bytes::{Buf, BufMut}; | ||||
|     use std::io; | ||||
|     use std::net::SocketAddr; | ||||
|     #[cfg(unix)] | ||||
|     use std::os::unix::io::{AsRawFd, RawFd}; | ||||
|     use tokio::io::{AsyncRead, AsyncWrite}; | ||||
|     use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; | ||||
|     use tokio::net::TcpStream; | ||||
|  | ||||
|     use crate::common::{task, Pin, Poll}; | ||||
|  | ||||
|     /// A transport returned yieled by `AddrIncoming`. | ||||
|     #[pin_project::pin_project] | ||||
|     #[derive(Debug)] | ||||
|     pub struct AddrStream { | ||||
|         #[pin] | ||||
|         inner: TcpStream, | ||||
|         pub(super) remote_addr: SocketAddr, | ||||
|     } | ||||
| @@ -231,49 +273,24 @@ mod addr_stream { | ||||
|     } | ||||
|  | ||||
|     impl AsyncRead for AddrStream { | ||||
|         unsafe fn prepare_uninitialized_buffer( | ||||
|             &self, | ||||
|             buf: &mut [std::mem::MaybeUninit<u8>], | ||||
|         ) -> bool { | ||||
|             self.inner.prepare_uninitialized_buffer(buf) | ||||
|         } | ||||
|  | ||||
|         #[inline] | ||||
|         fn poll_read( | ||||
|             mut self: Pin<&mut Self>, | ||||
|             self: Pin<&mut Self>, | ||||
|             cx: &mut task::Context<'_>, | ||||
|             buf: &mut [u8], | ||||
|         ) -> Poll<io::Result<usize>> { | ||||
|             Pin::new(&mut self.inner).poll_read(cx, buf) | ||||
|         } | ||||
|  | ||||
|         #[inline] | ||||
|         fn poll_read_buf<B: BufMut>( | ||||
|             mut self: Pin<&mut Self>, | ||||
|             cx: &mut task::Context<'_>, | ||||
|             buf: &mut B, | ||||
|         ) -> Poll<io::Result<usize>> { | ||||
|             Pin::new(&mut self.inner).poll_read_buf(cx, buf) | ||||
|             buf: &mut ReadBuf<'_>, | ||||
|         ) -> Poll<io::Result<()>> { | ||||
|             self.project().inner.poll_read(cx, buf) | ||||
|         } | ||||
|     } | ||||
|  | ||||
|     impl AsyncWrite for AddrStream { | ||||
|         #[inline] | ||||
|         fn poll_write( | ||||
|             mut self: Pin<&mut Self>, | ||||
|             self: Pin<&mut Self>, | ||||
|             cx: &mut task::Context<'_>, | ||||
|             buf: &[u8], | ||||
|         ) -> Poll<io::Result<usize>> { | ||||
|             Pin::new(&mut self.inner).poll_write(cx, buf) | ||||
|         } | ||||
|  | ||||
|         #[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) | ||||
|             self.project().inner.poll_write(cx, buf) | ||||
|         } | ||||
|  | ||||
|         #[inline] | ||||
| @@ -283,11 +300,8 @@ mod addr_stream { | ||||
|         } | ||||
|  | ||||
|         #[inline] | ||||
|         fn poll_shutdown( | ||||
|             mut self: Pin<&mut Self>, | ||||
|             cx: &mut task::Context<'_>, | ||||
|         ) -> Poll<io::Result<()>> { | ||||
|             Pin::new(&mut self.inner).poll_shutdown(cx) | ||||
|         fn poll_shutdown(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> { | ||||
|             self.project().inner.poll_shutdown(cx) | ||||
|         } | ||||
|     } | ||||
|  | ||||
|   | ||||
		Reference in New Issue
	
	Block a user