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:
Sean McArthur
2020-11-05 17:17:21 -08:00
committed by GitHub
parent cc7d3058e8
commit 1b9af22fa0
24 changed files with 467 additions and 472 deletions

View File

@@ -12,8 +12,8 @@ use std::time::Duration;
use futures_util::future::Either;
use http::uri::{Scheme, Uri};
use pin_project::pin_project;
use tokio::net::TcpStream;
use tokio::time::Delay;
use tokio::net::{TcpSocket, TcpStream};
use tokio::time::Sleep;
use super::dns::{self, resolve, GaiResolver, Resolve};
use super::{Connected, Connection};
@@ -331,34 +331,9 @@ where
dns::IpAddrs::new(addrs)
};
let c = ConnectingTcp::new(
config.local_address_ipv4,
config.local_address_ipv6,
addrs,
config.connect_timeout,
config.happy_eyeballs_timeout,
config.reuse_address,
);
let c = ConnectingTcp::new(addrs, config);
let sock = c
.connect()
.await
.map_err(ConnectError::m("tcp connect error"))?;
if let Some(dur) = config.keep_alive_timeout {
sock.set_keepalive(Some(dur))
.map_err(ConnectError::m("tcp set_keepalive error"))?;
}
if let Some(size) = config.send_buffer_size {
sock.set_send_buffer_size(size)
.map_err(ConnectError::m("tcp set_send_buffer_size error"))?;
}
if let Some(size) = config.recv_buffer_size {
sock.set_recv_buffer_size(size)
.map_err(ConnectError::m("tcp set_recv_buffer_size error"))?;
}
let sock = c.connect().await?;
sock.set_nodelay(config.nodelay)
.map_err(ConnectError::m("tcp set_nodelay error"))?;
@@ -475,60 +450,45 @@ impl StdError for ConnectError {
}
}
struct ConnectingTcp {
local_addr_ipv4: Option<Ipv4Addr>,
local_addr_ipv6: Option<Ipv6Addr>,
struct ConnectingTcp<'a> {
preferred: ConnectingTcpRemote,
fallback: Option<ConnectingTcpFallback>,
reuse_address: bool,
config: &'a Config,
}
impl ConnectingTcp {
fn new(
local_addr_ipv4: Option<Ipv4Addr>,
local_addr_ipv6: Option<Ipv6Addr>,
remote_addrs: dns::IpAddrs,
connect_timeout: Option<Duration>,
fallback_timeout: Option<Duration>,
reuse_address: bool,
) -> ConnectingTcp {
if let Some(fallback_timeout) = fallback_timeout {
let (preferred_addrs, fallback_addrs) =
remote_addrs.split_by_preference(local_addr_ipv4, local_addr_ipv6);
impl<'a> ConnectingTcp<'a> {
fn new(remote_addrs: dns::IpAddrs, config: &'a Config) -> Self {
if let Some(fallback_timeout) = config.happy_eyeballs_timeout {
let (preferred_addrs, fallback_addrs) = remote_addrs
.split_by_preference(config.local_address_ipv4, config.local_address_ipv6);
if fallback_addrs.is_empty() {
return ConnectingTcp {
local_addr_ipv4,
local_addr_ipv6,
preferred: ConnectingTcpRemote::new(preferred_addrs, connect_timeout),
preferred: ConnectingTcpRemote::new(preferred_addrs, config.connect_timeout),
fallback: None,
reuse_address,
config,
};
}
ConnectingTcp {
local_addr_ipv4,
local_addr_ipv6,
preferred: ConnectingTcpRemote::new(preferred_addrs, connect_timeout),
preferred: ConnectingTcpRemote::new(preferred_addrs, config.connect_timeout),
fallback: Some(ConnectingTcpFallback {
delay: tokio::time::delay_for(fallback_timeout),
remote: ConnectingTcpRemote::new(fallback_addrs, connect_timeout),
delay: tokio::time::sleep(fallback_timeout),
remote: ConnectingTcpRemote::new(fallback_addrs, config.connect_timeout),
}),
reuse_address,
config,
}
} else {
ConnectingTcp {
local_addr_ipv4,
local_addr_ipv6,
preferred: ConnectingTcpRemote::new(remote_addrs, connect_timeout),
preferred: ConnectingTcpRemote::new(remote_addrs, config.connect_timeout),
fallback: None,
reuse_address,
config,
}
}
}
}
struct ConnectingTcpFallback {
delay: Delay,
delay: Sleep,
remote: ConnectingTcpRemote,
}
@@ -549,24 +509,11 @@ impl ConnectingTcpRemote {
}
impl ConnectingTcpRemote {
async fn connect(
&mut self,
local_addr_ipv4: &Option<Ipv4Addr>,
local_addr_ipv6: &Option<Ipv6Addr>,
reuse_address: bool,
) -> io::Result<TcpStream> {
async fn connect(&mut self, config: &Config) -> Result<TcpStream, ConnectError> {
let mut err = None;
for addr in &mut self.addrs {
debug!("connecting to {}", addr);
match connect(
&addr,
local_addr_ipv4,
local_addr_ipv6,
reuse_address,
self.connect_timeout,
)?
.await
{
match connect(&addr, config, self.connect_timeout)?.await {
Ok(tcp) => {
debug!("connected to {}", addr);
return Ok(tcp);
@@ -580,9 +527,9 @@ impl ConnectingTcpRemote {
match err {
Some(e) => Err(e),
None => Err(std::io::Error::new(
std::io::ErrorKind::NotConnected,
"Network unreachable",
None => Err(ConnectError::new(
"tcp connect error",
std::io::Error::new(std::io::ErrorKind::NotConnected, "Network unreachable"),
)),
}
}
@@ -618,30 +565,79 @@ fn bind_local_address(
fn connect(
addr: &SocketAddr,
local_addr_ipv4: &Option<Ipv4Addr>,
local_addr_ipv6: &Option<Ipv6Addr>,
reuse_address: bool,
config: &Config,
connect_timeout: Option<Duration>,
) -> io::Result<impl Future<Output = io::Result<TcpStream>>> {
) -> Result<impl Future<Output = Result<TcpStream, ConnectError>>, ConnectError> {
// TODO(eliza): if Tokio's `TcpSocket` gains support for setting the
// keepalive timeout and send/recv buffer size, it would be nice to use that
// instead of socket2, and avoid the unsafe `into_raw_fd`/`from_raw_fd`
// dance...
use socket2::{Domain, Protocol, Socket, Type};
let domain = match *addr {
SocketAddr::V4(_) => Domain::ipv4(),
SocketAddr::V6(_) => Domain::ipv6(),
};
let socket = Socket::new(domain, Type::stream(), Some(Protocol::tcp()))?;
let socket = Socket::new(domain, Type::stream(), Some(Protocol::tcp()))
.map_err(ConnectError::m("tcp open error"))?;
if reuse_address {
socket.set_reuse_address(true)?;
if config.reuse_address {
socket
.set_reuse_address(true)
.map_err(ConnectError::m("tcp set_reuse_address error"))?;
}
bind_local_address(&socket, addr, local_addr_ipv4, local_addr_ipv6)?;
// When constructing a Tokio `TcpSocket` from a raw fd/socket, the user is
// responsible for ensuring O_NONBLOCK is set.
socket
.set_nonblocking(true)
.map_err(ConnectError::m("tcp set_nonblocking error"))?;
let addr = *addr;
bind_local_address(
&socket,
addr,
&config.local_address_ipv4,
&config.local_address_ipv6,
)
.map_err(ConnectError::m("tcp bind local error"))?;
let std_tcp = socket.into_tcp_stream();
if let Some(dur) = config.keep_alive_timeout {
socket
.set_keepalive(Some(dur))
.map_err(ConnectError::m("tcp set_keepalive error"))?;
}
if let Some(size) = config.send_buffer_size {
socket
.set_send_buffer_size(size)
.map_err(ConnectError::m("tcp set_send_buffer_size error"))?;
}
if let Some(size) = config.recv_buffer_size {
socket
.set_recv_buffer_size(size)
.map_err(ConnectError::m("tcp set_recv_buffer_size error"))?;
}
#[cfg(unix)]
let socket = unsafe {
// Safety: `from_raw_fd` is only safe to call if ownership of the raw
// file descriptor is transferred. Since we call `into_raw_fd` on the
// socket2 socket, it gives up ownership of the fd and will not close
// it, so this is safe.
use std::os::unix::io::{FromRawFd, IntoRawFd};
TcpSocket::from_raw_fd(socket.into_raw_fd())
};
#[cfg(windows)]
let socket = unsafe {
// Safety: `from_raw_socket` is only safe to call if ownership of the raw
// Windows SOCKET is transferred. Since we call `into_raw_socket` on the
// socket2 socket, it gives up ownership of the SOCKET and will not close
// it, so this is safe.
use std::os::windows::io::{FromRawSocket, IntoRawSocket};
TcpSocket::from_raw_socket(socket.into_raw_socket())
};
let connect = socket.connect(*addr);
Ok(async move {
let connect = TcpStream::connect_std(std_tcp, &addr);
match connect_timeout {
Some(dur) => match tokio::time::timeout(dur, connect).await {
Ok(Ok(s)) => Ok(s),
@@ -650,33 +646,19 @@ fn connect(
},
None => connect.await,
}
.map_err(ConnectError::m("tcp connect error"))
})
}
impl ConnectingTcp {
async fn connect(mut self) -> io::Result<TcpStream> {
let Self {
ref local_addr_ipv4,
ref local_addr_ipv6,
reuse_address,
..
} = self;
impl ConnectingTcp<'_> {
async fn connect(mut self) -> Result<TcpStream, ConnectError> {
match self.fallback {
None => {
self.preferred
.connect(local_addr_ipv4, local_addr_ipv6, reuse_address)
.await
}
None => self.preferred.connect(self.config).await,
Some(mut fallback) => {
let preferred_fut =
self.preferred
.connect(local_addr_ipv4, local_addr_ipv6, reuse_address);
let preferred_fut = self.preferred.connect(self.config);
futures_util::pin_mut!(preferred_fut);
let fallback_fut =
fallback
.remote
.connect(local_addr_ipv4, local_addr_ipv6, reuse_address);
let fallback_fut = fallback.remote.connect(self.config);
futures_util::pin_mut!(fallback_fut);
let (result, future) =
@@ -711,7 +693,7 @@ mod tests {
use ::http::Uri;
use super::super::sealed::{Connect, ConnectSvc};
use super::HttpConnector;
use super::{Config, ConnectError, HttpConnector};
async fn connect<C>(
connector: C,
@@ -773,6 +755,7 @@ mod tests {
#[tokio::test]
async fn local_address() {
use std::net::{IpAddr, TcpListener};
let _ = pretty_env_logger::try_init();
let (bind_ip_v4, bind_ip_v6) = get_local_ips();
let server4 = TcpListener::bind("127.0.0.1:0").unwrap();
@@ -818,10 +801,8 @@ mod tests {
let server4 = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = server4.local_addr().unwrap();
let _server6 = TcpListener::bind(&format!("[::1]:{}", addr.port())).unwrap();
let mut rt = tokio::runtime::Builder::new()
.enable_io()
.enable_time()
.basic_scheduler()
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
@@ -925,16 +906,21 @@ mod tests {
.iter()
.map(|host| (host.clone(), addr.port()).into())
.collect();
let connecting_tcp = ConnectingTcp::new(
None,
None,
dns::IpAddrs::new(addrs),
None,
Some(fallback_timeout),
false,
);
let cfg = Config {
local_address_ipv4: None,
local_address_ipv6: None,
connect_timeout: None,
keep_alive_timeout: None,
happy_eyeballs_timeout: Some(fallback_timeout),
nodelay: false,
reuse_address: false,
enforce_http: false,
send_buffer_size: None,
recv_buffer_size: None,
};
let connecting_tcp = ConnectingTcp::new(dns::IpAddrs::new(addrs), &cfg);
let start = Instant::now();
Ok::<_, io::Error>((start, connecting_tcp.connect().await?))
Ok::<_, ConnectError>((start, ConnectingTcp::connect(connecting_tcp).await?))
})
.unwrap();
let res = if stream.peer_addr().unwrap().is_ipv4() {

View File

@@ -1,4 +1,5 @@
use futures_util::future;
use tokio::stream::Stream;
use tokio::sync::{mpsc, oneshot};
use crate::common::{task, Future, Pin, Poll};
@@ -131,22 +132,25 @@ impl<T, U> Clone for UnboundedSender<T, U> {
}
}
#[pin_project::pin_project(PinnedDrop)]
pub struct Receiver<T, U> {
#[pin]
inner: mpsc::UnboundedReceiver<Envelope<T, U>>,
taker: want::Taker,
}
impl<T, U> Receiver<T, U> {
pub(crate) fn poll_next(
&mut self,
self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
) -> Poll<Option<(T, Callback<T, U>)>> {
match self.inner.poll_recv(cx) {
let this = self.project();
match this.inner.poll_next(cx) {
Poll::Ready(item) => {
Poll::Ready(item.map(|mut env| env.0.take().expect("envelope not dropped")))
}
Poll::Pending => {
self.taker.want();
this.taker.want();
Poll::Pending
}
}
@@ -165,11 +169,12 @@ impl<T, U> Receiver<T, U> {
}
}
impl<T, U> Drop for Receiver<T, U> {
fn drop(&mut self) {
#[pin_project::pinned_drop]
impl<T, U> PinnedDrop for Receiver<T, U> {
fn drop(mut self: Pin<&mut Self>) {
// Notify the giver about the closure first, before dropping
// the mpsc::Receiver.
self.taker.cancel();
self.as_mut().taker.cancel();
}
}
@@ -262,7 +267,7 @@ mod tests {
impl<T, U> Future for Receiver<T, U> {
type Output = Option<(T, Callback<T, U>)>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.poll_next(cx)
}
}
@@ -344,9 +349,8 @@ mod tests {
fn giver_queue_throughput(b: &mut test::Bencher) {
use crate::{Body, Request, Response};
let mut rt = tokio::runtime::Builder::new()
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.basic_scheduler()
.build()
.unwrap();
let (mut tx, mut rx) = channel::<Request<Body>, Response<Body>>();
@@ -368,9 +372,8 @@ mod tests {
#[cfg(feature = "nightly")]
#[bench]
fn giver_queue_not_ready(b: &mut test::Bencher) {
let mut rt = tokio::runtime::Builder::new()
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.basic_scheduler()
.build()
.unwrap();
let (_tx, mut rx) = channel::<i32, ()>();

View File

@@ -706,12 +706,15 @@ impl Expiration {
}
#[cfg(feature = "runtime")]
#[pin_project::pin_project]
struct IdleTask<T> {
#[pin]
interval: Interval,
pool: WeakOpt<Mutex<PoolInner<T>>>,
// This allows the IdleTask to be notified as soon as the entire
// Pool is fully dropped, and shutdown. This channel is never sent on,
// but Err(Canceled) will be received when the Pool is dropped.
#[pin]
pool_drop_notifier: oneshot::Receiver<crate::common::Never>,
}
@@ -719,9 +722,11 @@ struct IdleTask<T> {
impl<T: Poolable + 'static> Future for IdleTask<T> {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
use tokio::stream::Stream;
let mut this = self.project();
loop {
match Pin::new(&mut self.pool_drop_notifier).poll(cx) {
match this.pool_drop_notifier.as_mut().poll(cx) {
Poll::Ready(Ok(n)) => match n {},
Poll::Pending => (),
Poll::Ready(Err(_canceled)) => {
@@ -730,9 +735,9 @@ impl<T: Poolable + 'static> Future for IdleTask<T> {
}
}
ready!(self.interval.poll_tick(cx));
ready!(this.interval.as_mut().poll_next(cx));
if let Some(inner) = self.pool.upgrade() {
if let Some(inner) = this.pool.upgrade() {
if let Ok(mut inner) = inner.lock() {
trace!("idle interval checking for expired");
inner.clear_expired();
@@ -850,7 +855,7 @@ mod tests {
let pooled = pool.pooled(c(key.clone()), Uniq(41));
drop(pooled);
tokio::time::delay_for(pool.locked().timeout.unwrap()).await;
tokio::time::sleep(pool.locked().timeout.unwrap()).await;
let mut checkout = pool.checkout(key);
let poll_once = PollOnce(&mut checkout);
let is_not_ready = poll_once.await.is_none();
@@ -871,7 +876,7 @@ mod tests {
pool.locked().idle.get(&key).map(|entries| entries.len()),
Some(3)
);
tokio::time::delay_for(pool.locked().timeout.unwrap()).await;
tokio::time::sleep(pool.locked().timeout.unwrap()).await;
let mut checkout = pool.checkout(key.clone());
let poll_once = PollOnce(&mut checkout);