tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-09-08 00:00:13 +02:00
Code Issues Packages Projects Releases Wiki Activity

net: move into tokio crate (#1683)

Browse Source 
A step towards collapsing Tokio sub crates into a single `tokio`
crate (#1318).

The `net` implementation is now provided by the main `tokio` crate.
Functionality can be opted out of by using the various net related
feature flags.
This commit is contained in:
Carl Lerche
2019-10-25 12:50:15 -07:00
committed by GitHub
parent 03a9378297
commit 227533d456
111 changed files with 578 additions and 868 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/Cargo.toml
+46 -15
View File
@@ -27,7 +27,7 @@ keywords = ["io", "async", "non-blocking", "futures"]
default = [
"fs",
"io",
"net",
"net-full",
"process",
"rt-full",
"signal",
@@ -35,31 +35,47 @@ default = [
"timer",
]
fs = []
io = ["tokio-io"]
fs = ["tokio-executor/blocking"]
io = ["tokio-io", "bytes", "iovec"]
macros = ["tokio-macros"]
net = ["tcp", "udp", "uds"]
net-full = ["tcp", "udp", "uds"]
net-driver = ["mio", "tokio-executor/blocking"]
rt-current-thread = [
"timer",
"tokio-net",
"tokio-executor/current-thread",
]
rt-full = [
"macros",
"num_cpus",
"net",
"net-full",
"sync",
"timer",
"tokio-executor/current-thread",
"tokio-executor/thread-pool",
]
signal = ["tokio-net/signal"]
signal = [
"lazy_static",
"libc",
"mio-uds",
"net-driver",
"signal-hook-registry"
]
sync = ["tokio-sync"]
tcp = ["io", "tokio-net/tcp"]
tcp = ["io", "net-driver"]
timer = ["crossbeam-utils", "slab"]
udp = ["io", "tokio-net/udp"]
uds = ["io", "tokio-net/uds"]
process = ["io", "tokio-net/process"]
udp = ["io", "net-driver"]
uds = ["io", "net-driver", "mio-uds", "libc"]
process = [
"crossbeam-queue",
"io",
"libc",
"mio-named-pipes",
"signal",
"winapi/consoleapi",
"winapi/minwindef",
"winapi/threadpoollegacyapiset",
"winapi/winerror",
]
[dependencies]
futures-core-preview = "=0.3.0-alpha.19"
@@ -67,23 +83,38 @@ futures-sink-preview = "=0.3.0-alpha.19"
futures-util-preview = { version = "=0.3.0-alpha.19", features = ["sink"] }
# Everything else is optional...
bytes = { version = "0.4", optional = true }
crossbeam-utils = { version = "0.6.0", optional = true }
iovec = { version = "0.1", optional = true }
lazy_static = { version = "1.0.2", optional = true }
mio = { version = "0.6.14", optional = true }
num_cpus = { version = "1.8.0", optional = true }
# Backs `DelayQueue`
slab = { version = "0.4.1", optional = true }
tokio-io = { version = "=0.2.0-alpha.6", optional = true, features = ["util"], path = "../tokio-io" }
tokio-executor = { version = "=0.2.0-alpha.6", optional = true, path = "../tokio-executor" }
tokio-macros = { version = "=0.2.0-alpha.6", optional = true, path = "../tokio-macros" }
tokio-net = { version = "=0.2.0-alpha.6", optional = true, features = ["async-traits"], path = "../tokio-net" }
tokio-sync = { version = "=0.2.0-alpha.6", optional = true, path = "../tokio-sync", features = ["async-traits"] }
[target.'cfg(unix)'.dependencies]
crossbeam-queue = { version = "0.1.2", optional = true }
mio-uds = { version = "0.6.5", optional = true }
libc = { version = "0.2.42", optional = true }
signal-hook-registry = { version = "1.1.1", optional = true }
[target.'cfg(windows)'.dependencies]
mio-named-pipes = { version = "0.1.6", optional = true }
[target.'cfg(windows)'.dependencies.winapi]
version = "0.3.8"
default-features = false
optional = true
[dev-dependencies]
tokio-test = { version = "=0.2.0-alpha.6", path = "../tokio-test" }
tokio-util = { version = "=0.2.0-alpha.6", path = "../tokio-util" }
futures-preview = "=0.3.0-alpha.19"
futures-util-preview = "=0.3.0-alpha.19"
pin-utils = "=0.1.0-alpha.4"
futures-preview = { version = "=0.3.0-alpha.19", features = ["async-await"] }
env_logger = { version = "0.6", default-features = false }
flate2 = { version = "1", features = ["tokio"] }
http = "0.1"
tokio/src/lib.rs
+14 -2
View File
@@ -26,7 +26,7 @@
//!
//! Guide level documentation is found on the [website].
//!
//! [driver]: tokio_net::driver
//! [driver]: tokio::net::driver
//! [website]: https://tokio.rs/docs/
//!
//! # Examples
@@ -82,23 +82,34 @@ macro_rules! if_runtime {
#[cfg(feature = "timer")]
pub mod clock;
#[cfg(feature = "codec")]
pub mod codec;
#[cfg(feature = "fs")]
pub mod fs;
pub mod future;
#[cfg(feature = "io")]
pub mod io;
#[cfg(any(feature = "tcp", feature = "udp", feature = "uds"))]
#[cfg(feature = "net-driver")]
pub mod net;
pub mod prelude;
#[cfg(feature = "process")]
pub mod process;
#[cfg(feature = "signal")]
pub mod signal;
pub mod stream;
#[cfg(feature = "sync")]
pub mod sync;
#[cfg(feature = "timer")]
pub mod timer;
@@ -113,6 +124,7 @@ if_runtime! {
#[cfg(feature = "macros")]
#[doc(inline)]
pub use tokio_macros::main;
#[cfg(feature = "macros")]
#[doc(inline)]
pub use tokio_macros::test;
tokio/src/net.rs
-67
View File
@@ -1,67 +0,0 @@
//! TCP/UDP/Unix bindings for `tokio`.
//!
//! This module contains the TCP/UDP/Unix networking types, similar to the standard
//! library, which can be used to implement networking protocols.
//!
//! # Organization
//!
//! * [`TcpListener`] and [`TcpStream`] provide functionality for communication over TCP
//! * [`UdpSocket`] provides functionality for communication over UDP
//! * [`UnixListener`] and [`UnixStream`] provide functionality for communication over a
//! Unix Domain Stream Socket **(available on Unix only)**
//! * [`UnixDatagram`] and [`UnixDatagramFramed`] provide functionality for communication
//! over Unix Domain Datagram Socket **(available on Unix only)**
//!
//! [`TcpListener`]: struct.TcpListener.html
//! [`TcpStream`]: struct.TcpStream.html
//! [`UdpSocket`]: struct.UdpSocket.html
//! [`UnixListener`]: struct.UnixListener.html
//! [`UnixStream`]: struct.UnixStream.html
//! [`UnixDatagram`]: struct.UnixDatagram.html
//! [`UnixDatagramFramed`]: struct.UnixDatagramFramed.html
#[cfg(feature = "tcp")]
pub mod tcp {
//! TCP bindings for `tokio`.
//!
//! Connecting to an address, via TCP, can be done using [`TcpStream`]'s
//! [`connect`] method, which returns [`ConnectFuture`]. `ConnectFuture`
//! implements a future which returns a `TcpStream`.
//!
//! To listen on an address [`TcpListener`] can be used. `TcpListener`'s
//! [`incoming`][incoming_method] method can be used to accept new connections.
//! It return the [`Incoming`] struct, which implements a stream which returns
//! `TcpStream`s.
//!
//! [`TcpStream`]: struct.TcpStream.html
//! [`connect`]: struct.TcpStream.html#method.connect
//! [`ConnectFuture`]: struct.ConnectFuture.html
//! [`TcpListener`]: struct.TcpListener.html
//! [incoming_method]: struct.TcpListener.html#method.incoming
//! [`Incoming`]: struct.Incoming.html
pub use tokio_net::tcp::{split, Incoming, TcpListener, TcpStream};
}
#[cfg(feature = "tcp")]
pub use self::tcp::{TcpListener, TcpStream};
#[cfg(feature = "udp")]
pub mod udp {
//! UDP bindings for `tokio`.
//!
//! The main struct for UDP is the [`UdpSocket`], which represents a UDP socket.
//!
//! [`UdpSocket`]: struct.UdpSocket.html
pub use tokio_net::udp::{split, UdpSocket};
}
#[cfg(feature = "udp")]
pub use self::udp::UdpSocket;
#[cfg(all(unix, feature = "uds"))]
pub mod unix {
//! Unix domain socket bindings for `tokio` (only available on unix systems).
pub use tokio_net::uds::{split, UCred, UnixDatagram, UnixListener, UnixStream};
}
#[cfg(all(unix, feature = "uds"))]
pub use self::unix::{UnixDatagram, UnixListener, UnixStream};
tokio/src/net/addr.rs
+214
View File
@@ -0,0 +1,214 @@
use tokio_executor::blocking;
use futures_util::future;
use std::io;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
/// Convert or resolve without blocking to one or more `SocketAddr` values.
///
/// Currently, this trait is only used as an argument to Tokio functions that
/// need to reference a target socket address.
///
/// This trait is sealed and is intended to be opaque. Users of Tokio should
/// only use `ToSocketAddrs` in trait bounds and __must not__ attempt to call
/// the functions directly or reference associated types. Changing these is not
/// considered a breaking change.
pub trait ToSocketAddrs: sealed::ToSocketAddrsPriv {}
type ReadyFuture<T> = future::Ready<io::Result<T>>;
// ===== impl SocketAddr =====
impl ToSocketAddrs for SocketAddr {}
impl sealed::ToSocketAddrsPriv for SocketAddr {
type Iter = std::option::IntoIter<SocketAddr>;
type Future = ReadyFuture<Self::Iter>;
fn to_socket_addrs(&self) -> Self::Future {
let iter = Some(*self).into_iter();
future::ready(Ok(iter))
}
}
// ===== impl str =====
impl ToSocketAddrs for str {}
impl sealed::ToSocketAddrsPriv for str {
type Iter = sealed::OneOrMore;
type Future = sealed::MaybeReady;
fn to_socket_addrs(&self) -> Self::Future {
use sealed::MaybeReady;
// First check if the input parses as a socket address
let res: Result<SocketAddr, _> = self.parse();
if let Ok(addr) = res {
return MaybeReady::Ready(Some(addr));
}
// Run DNS lookup on the blocking pool
let s = self.to_owned();
MaybeReady::Blocking(blocking::run(move || {
std::net::ToSocketAddrs::to_socket_addrs(&s)
}))
}
}
// ===== impl (&str, u16) =====
impl ToSocketAddrs for (&'_ str, u16) {}
impl sealed::ToSocketAddrsPriv for (&'_ str, u16) {
type Iter = sealed::OneOrMore;
type Future = sealed::MaybeReady;
fn to_socket_addrs(&self) -> Self::Future {
use sealed::MaybeReady;
use std::net::{SocketAddrV4, SocketAddrV6};
let (host, port) = *self;
// try to parse the host as a regular IP address first
if let Ok(addr) = host.parse::<Ipv4Addr>() {
let addr = SocketAddrV4::new(addr, port);
let addr = SocketAddr::V4(addr);
return MaybeReady::Ready(Some(addr));
}
if let Ok(addr) = host.parse::<Ipv6Addr>() {
let addr = SocketAddrV6::new(addr, port, 0, 0);
let addr = SocketAddr::V6(addr);
return MaybeReady::Ready(Some(addr));
}
let host = host.to_owned();
MaybeReady::Blocking(blocking::run(move || {
std::net::ToSocketAddrs::to_socket_addrs(&(&host[..], port))
}))
}
}
// ===== impl (IpAddr, u16) =====
impl ToSocketAddrs for (IpAddr, u16) {}
impl sealed::ToSocketAddrsPriv for (IpAddr, u16) {
type Iter = std::option::IntoIter<SocketAddr>;
type Future = ReadyFuture<Self::Iter>;
fn to_socket_addrs(&self) -> Self::Future {
let iter = Some(SocketAddr::from(*self)).into_iter();
future::ready(Ok(iter))
}
}
// ===== impl String =====
impl ToSocketAddrs for String {}
impl sealed::ToSocketAddrsPriv for String {
type Iter = <str as sealed::ToSocketAddrsPriv>::Iter;
type Future = <str as sealed::ToSocketAddrsPriv>::Future;
fn to_socket_addrs(&self) -> Self::Future {
(&self[..]).to_socket_addrs()
}
}
// ===== impl &'_ impl ToSocketAddrs =====
impl<T: ToSocketAddrs + ?Sized> ToSocketAddrs for &'_ T {}
impl<T> sealed::ToSocketAddrsPriv for &'_ T
where
T: sealed::ToSocketAddrsPriv + ?Sized,
{
type Iter = T::Iter;
type Future = T::Future;
fn to_socket_addrs(&self) -> Self::Future {
(**self).to_socket_addrs()
}
}
pub(crate) mod sealed {
//! The contents of this trait are intended to remain private and __not__
//! part of the `ToSocketAddrs` public API. The details will change over
//! time.
use tokio_executor::blocking::Blocking;
use futures_core::ready;
use std::future::Future;
use std::io;
use std::net::SocketAddr;
use std::option;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::vec;
#[doc(hidden)]
pub trait ToSocketAddrsPriv {
type Iter: Iterator<Item = SocketAddr> + Send + 'static;
type Future: Future<Output = io::Result<Self::Iter>> + Send + 'static;
fn to_socket_addrs(&self) -> Self::Future;
}
#[doc(hidden)]
#[derive(Debug)]
pub enum MaybeReady {
Ready(Option<SocketAddr>),
Blocking(Blocking<io::Result<vec::IntoIter<SocketAddr>>>),
}
#[doc(hidden)]
#[derive(Debug)]
pub enum OneOrMore {
One(option::IntoIter<SocketAddr>),
More(vec::IntoIter<SocketAddr>),
}
impl Future for MaybeReady {
type Output = io::Result<OneOrMore>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match *self {
MaybeReady::Ready(ref mut i) => {
let iter = OneOrMore::One(i.take().into_iter());
Poll::Ready(Ok(iter))
}
MaybeReady::Blocking(ref mut rx) => {
let res = ready!(Pin::new(rx).poll(cx)).map(OneOrMore::More);
Poll::Ready(res)
}
}
}
}
impl Iterator for OneOrMore {
type Item = SocketAddr;
fn next(&mut self) -> Option<Self::Item> {
match self {
OneOrMore::One(i) => i.next(),
OneOrMore::More(i) => i.next(),
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
match self {
OneOrMore::One(i) => i.size_hint(),
OneOrMore::More(i) => i.size_hint(),
}
}
}
}
tokio/src/net/driver/mod.rs
+133
View File
@@ -0,0 +1,133 @@
//! Event loop that drives Tokio I/O resources.
//!
//! This module contains [`Reactor`], which is the event loop that drives all
//! Tokio I/O resources. It is the reactor's job to receive events from the
//! operating system ([epoll], [kqueue], [IOCP], etc...) and forward them to
//! waiting tasks. It is the bridge between operating system and the futures
//! model.
//!
//! # Overview
//!
//! When using Tokio, all operations are asynchronous and represented by
//! futures. These futures, representing the application logic, are scheduled by
//! an executor (see [runtime model] for more details). Executors wait for
//! notifications before scheduling the future for execution time, i.e., nothing
//! happens until an event is received indicating that the task can make
//! progress.
//!
//! The reactor receives events from the operating system and notifies the
//! executor.
//!
//! Let's start with a basic example, establishing a TCP connection.
//!
//! ```
//! use tokio::net::TcpStream;
//!
//! # async fn process<T>(_t: T) {}
//!
//! # #[tokio::main]
//! # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
//! let stream = TcpStream::connect("93.184.216.34:9243").await?;
//!
//! println!("successfully connected");
//!
//! process(stream).await;
//! # Ok(())
//! # }
//! ```
//!
//! Establishing a TCP connection usually cannot be completed immediately.
//! [`TcpStream::connect`] does not block the current thread. Instead, it
//! returns a [future][connect-future] that resolves once the TCP connection has
//! been established. The connect future itself has no way of knowing when the
//! TCP connection has been established.
//!
//! Before returning the future, [`TcpStream::connect`] registers the socket
//! with a reactor. This registration process, handled by [`Registration`], is
//! what links the [`TcpStream`] with the [`Reactor`] instance. At this point,
//! the reactor starts listening for connection events from the operating system
//! for that socket.
//!
//! Once the connect future is passed to [`tokio::run`], it is spawned onto a
//! thread pool. The thread pool waits until it is notified that the connection
//! has completed.
//!
//! When the TCP connection is established, the reactor receives an event from
//! the operating system. It then notifies the thread pool, telling it that the
//! connect future can complete. At this point, the thread pool will schedule
//! the task to run on one of its worker threads. This results in the `and_then`
//! closure to get executed.
//!
//! ## Eager registration
//!
//! Notice how the snippet does not explicitly reference a reactor. When
//! [`TcpStream::connect`] is called, it registers the socket with the current
//! reactor, but no reactor is specified. This works because a reactor
//! instance is automatically made available when using the Tokio [runtime],
//! which is done using [`tokio::main`]. The Tokio runtime's executor sets a
//! thread-local variable referencing the associated [`Reactor`] instance and
//! [`Handle::current`] (used by [`Registration`]) returns the reference.
//!
//! ## Implementation
//!
//! The reactor implementation uses [`mio`] to interface with the operating
//! system's event queue. A call to [`Reactor::poll`] results in a single
//! call to [`Poll::poll`] which in turn results in a single call to the
//! operating system's selector.
//!
//! The reactor maintains state for each registered I/O resource. This tracks
//! the executor task to notify when events are provided by the operating
//! system's selector. This state is stored in a `Sync` data structure and
//! referenced by [`Registration`]. When the [`Registration`] instance is
//! dropped, this state is cleaned up. Because the state is stored in a `Sync`
//! data structure, the [`Registration`] instance is able to be moved to other
//! threads.
//!
//! By default, a runtime's default reactor runs on a background thread. This
//! ensures that application code cannot significantly impact the reactor's
//! responsiveness.
//!
//! ## Integrating with the reactor
//!
//! Tokio comes with a number of I/O resources, like TCP and UDP sockets, that
//! automatically integrate with the reactor. However, library authors or
//! applications may wish to implement their own resources that are also backed
//! by the reactor.
//!
//! There are a couple of ways to do this.
//!
//! If the custom I/O resource implements [`mio::Evented`] and implements
//! [`std::io::Read`] and / or [`std::io::Write`], then [`PollEvented`] is the
//! most suited.
//!
//! Otherwise, [`Registration`] can be used directly. This provides the lowest
//! level primitive needed for integrating with the reactor: a stream of
//! readiness events.
//!
//! [`Reactor`]: struct.Reactor.html
//! [`Registration`]: struct.Registration.html
//! [runtime model]: https://tokio.rs/docs/internals/runtime-model/
//! [epoll]: http://man7.org/linux/man-pages/man7/epoll.7.html
//! [kqueue]: https://www.freebsd.org/cgi/man.cgi?query=kqueue&sektion=2
//! [IOCP]: https://msdn.microsoft.com/en-us/library/windows/desktop/aa365198(v=vs.85).aspx
//! [`TcpStream::connect`]: ../net/struct.TcpStream.html#method.connect
//! [`connect`]: ../net/struct.TcpStream.html#method.connect
//! [connect-future]: ../net/struct.ConnectFuture.html
//! [`tokio::run`]: ../runtime/fn.run.html
//! [`TcpStream`]: ../net/struct.TcpStream.html
//! [runtime]: ../runtime
//! [`Handle::current`]: struct.Handle.html#method.current
//! [`mio`]: https://github.com/carllerche/mio
//! [`Reactor::poll`]: struct.Reactor.html#method.poll
//! [`Poll::poll`]: https://docs.rs/mio/0.6/mio/struct.Poll.html#method.poll
//! [`mio::Evented`]: https://docs.rs/mio/0.6/mio/trait.Evented.html
//! [`PollEvented`]: struct.PollEvented.html
//! [`std::io::Read`]: https://doc.rust-lang.org/std/io/trait.Read.html
//! [`std::io::Write`]: https://doc.rust-lang.org/std/io/trait.Write.html
pub(crate) mod platform;
mod reactor;
mod registration;
pub use self::reactor::{set_default, DefaultGuard, Handle, Reactor};
pub use self::registration::Registration;
tokio/src/net/driver/platform.rs
+28
View File
@@ -0,0 +1,28 @@
pub(crate) use self::sys::*;
#[cfg(unix)]
mod sys {
use mio::unix::UnixReady;
use mio::Ready;
pub(crate) fn hup() -> Ready {
UnixReady::hup().into()
}
pub(crate) fn is_hup(ready: Ready) -> bool {
UnixReady::from(ready).is_hup()
}
}
#[cfg(windows)]
mod sys {
use mio::Ready;
pub(crate) fn hup() -> Ready {
Ready::empty()
}
pub(crate) fn is_hup(_: Ready) -> bool {
false
}
}
tokio/src/net/driver/reactor.rs
+451
View File
@@ -0,0 +1,451 @@
use super::platform;
use tokio_executor::park::{Park, Unpark};
use tokio_sync::AtomicWaker;
use mio::event::Evented;
use slab::Slab;
use std::cell::RefCell;
use std::io;
use std::marker::PhantomData;
#[cfg(all(unix, not(target_os = "fuchsia")))]
use std::os::unix::io::{AsRawFd, RawFd};
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{Relaxed, SeqCst};
use std::sync::{Arc, RwLock, Weak};
use std::task::Waker;
use std::time::Duration;
use std::{fmt, usize};
/// The core reactor, or event loop.
///
/// The event loop is the main source of blocking in an application which drives
/// all other I/O events and notifications happening. Each event loop can have
/// multiple handles pointing to it, each of which can then be used to create
/// various I/O objects to interact with the event loop in interesting ways.
pub struct Reactor {
/// Reuse the `mio::Events` value across calls to poll.
events: mio::Events,
/// State shared between the reactor and the handles.
inner: Arc<Inner>,
_wakeup_registration: mio::Registration,
}
/// A reference to a reactor.
///
/// A `Handle` is used for associating I/O objects with an event loop
/// explicitly. Typically though you won't end up using a `Handle` that often
/// and will instead use the default reactor for the execution context.
#[derive(Clone)]
pub struct Handle {
inner: Weak<Inner>,
}
/// Return value from the `turn` method on `Reactor`.
///
/// Currently this value doesn't actually provide any functionality, but it may
/// in the future give insight into what happened during `turn`.
#[derive(Debug)]
pub struct Turn {
_priv: (),
}
pub(super) struct Inner {
/// The underlying system event queue.
io: mio::Poll,
/// ABA guard counter
next_aba_guard: AtomicUsize,
/// Dispatch slabs for I/O and futures events
pub(super) io_dispatch: RwLock<Slab<ScheduledIo>>,
/// Used to wake up the reactor from a call to `turn`
wakeup: mio::SetReadiness,
}
pub(super) struct ScheduledIo {
aba_guard: usize,
pub(super) readiness: AtomicUsize,
pub(super) reader: AtomicWaker,
pub(super) writer: AtomicWaker,
}
#[derive(Debug, Eq, PartialEq, Clone, Copy)]
pub(super) enum Direction {
Read,
Write,
}
thread_local! {
/// Tracks the reactor for the current execution context.
static CURRENT_REACTOR: RefCell<Option<Handle>> = RefCell::new(None)
}
const TOKEN_SHIFT: usize = 22;
// Kind of arbitrary, but this reserves some token space for later usage.
const MAX_SOURCES: usize = (1 << TOKEN_SHIFT) - 1;
const TOKEN_WAKEUP: mio::Token = mio::Token(MAX_SOURCES);
fn _assert_kinds() {
fn _assert<T: Send + Sync>() {}
_assert::<Handle>();
}
// ===== impl Reactor =====
#[derive(Debug)]
/// Guard that resets current reactor on drop.
pub struct DefaultGuard<'a> {
_lifetime: PhantomData<&'a u8>,
}
impl Drop for DefaultGuard<'_> {
fn drop(&mut self) {
CURRENT_REACTOR.with(|current| {
let mut current = current.borrow_mut();
*current = None;
});
}
}
/// Sets handle for a default reactor, returning guard that unsets it on drop.
pub fn set_default(handle: &Handle) -> DefaultGuard<'_> {
CURRENT_REACTOR.with(|current| {
let mut current = current.borrow_mut();
assert!(
current.is_none(),
"default Tokio reactor already set \
for execution context"
);
*current = Some(handle.clone());
});
DefaultGuard {
_lifetime: PhantomData,
}
}
impl Reactor {
/// Creates a new event loop, returning any error that happened during the
/// creation.
pub fn new() -> io::Result<Reactor> {
let io = mio::Poll::new()?;
let wakeup_pair = mio::Registration::new2();
io.register(
&wakeup_pair.0,
TOKEN_WAKEUP,
mio::Ready::readable(),
mio::PollOpt::level(),
)?;
Ok(Reactor {
events: mio::Events::with_capacity(1024),
_wakeup_registration: wakeup_pair.0,
inner: Arc::new(Inner {
io,
next_aba_guard: AtomicUsize::new(0),
io_dispatch: RwLock::new(Slab::with_capacity(1)),
wakeup: wakeup_pair.1,
}),
})
}
/// Returns a handle to this event loop which can be sent across threads
/// and can be used as a proxy to the event loop itself.
///
/// Handles are cloneable and clones always refer to the same event loop.
/// This handle is typically passed into functions that create I/O objects
/// to bind them to this event loop.
pub fn handle(&self) -> Handle {
Handle {
inner: Arc::downgrade(&self.inner),
}
}
/// Performs one iteration of the event loop, blocking on waiting for events
/// for at most `max_wait` (forever if `None`).
///
/// This method is the primary method of running this reactor and processing
/// I/O events that occur. This method executes one iteration of an event
/// loop, blocking at most once waiting for events to happen.
///
/// If a `max_wait` is specified then the method should block no longer than
/// the duration specified, but this shouldn't be used as a super-precise
/// timer but rather a "ballpark approximation"
///
/// # Return value
///
/// This function returns an instance of `Turn`
///
/// `Turn` as of today has no extra information with it and can be safely
/// discarded. In the future `Turn` may contain information about what
/// happened while this reactor blocked.
///
/// # Errors
///
/// This function may also return any I/O error which occurs when polling
/// for readiness of I/O objects with the OS. This is quite unlikely to
/// arise and typically mean that things have gone horribly wrong at that
/// point. Currently this is primarily only known to happen for internal
/// bugs to `tokio` itself.
pub fn turn(&mut self, max_wait: Option<Duration>) -> io::Result<Turn> {
self.poll(max_wait)?;
Ok(Turn { _priv: () })
}
/// Returns true if the reactor is currently idle.
///
/// Idle is defined as all tasks that have been spawned have completed,
/// either successfully or with an error.
pub fn is_idle(&self) -> bool {
self.inner.io_dispatch.read().unwrap().is_empty()
}
fn poll(&mut self, max_wait: Option<Duration>) -> io::Result<()> {
// Block waiting for an event to happen, peeling out how many events
// happened.
match self.inner.io.poll(&mut self.events, max_wait) {
Ok(_) => {}
Err(e) => return Err(e),
}
// Process all the events that came in, dispatching appropriately
for event in self.events.iter() {
let token = event.token();
if token == TOKEN_WAKEUP {
self.inner
.wakeup
.set_readiness(mio::Ready::empty())
.unwrap();
} else {
self.dispatch(token, event.readiness());
}
}
Ok(())
}
fn dispatch(&self, token: mio::Token, ready: mio::Ready) {
let aba_guard = token.0 & !MAX_SOURCES;
let token = token.0 & MAX_SOURCES;
let mut rd = None;
let mut wr = None;
// Create a scope to ensure that notifying the tasks stays out of the
// lock's critical section.
{
let io_dispatch = self.inner.io_dispatch.read().unwrap();
let io = match io_dispatch.get(token) {
Some(io) => io,
None => return,
};
if aba_guard != io.aba_guard {
return;
}
io.readiness.fetch_or(ready.as_usize(), Relaxed);
if ready.is_writable() || platform::is_hup(ready) {
wr = io.writer.take_waker();
}
if !(ready & (!mio::Ready::writable())).is_empty() {
rd = io.reader.take_waker();
}
}
if let Some(w) = rd {
w.wake();
}
if let Some(w) = wr {
w.wake();
}
}
}
#[cfg(all(unix, not(target_os = "fuchsia")))]
impl AsRawFd for Reactor {
fn as_raw_fd(&self) -> RawFd {
self.inner.io.as_raw_fd()
}
}
impl Park for Reactor {
type Unpark = Handle;
type Error = io::Error;
fn unpark(&self) -> Self::Unpark {
self.handle()
}
fn park(&mut self) -> io::Result<()> {
self.turn(None)?;
Ok(())
}
fn park_timeout(&mut self, duration: Duration) -> io::Result<()> {
self.turn(Some(duration))?;
Ok(())
}
}
impl fmt::Debug for Reactor {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "Reactor")
}
}
// ===== impl Handle =====
impl Handle {
/// Returns a handle to the current reactor
///
/// # Panics
///
/// This function panics if there is no current reactor set.
pub(super) fn current() -> Self {
CURRENT_REACTOR.with(|current| match *current.borrow() {
Some(ref handle) => handle.clone(),
None => panic!("no current reactor"),
})
}
/// Forces a reactor blocked in a call to `turn` to wakeup, or otherwise
/// makes the next call to `turn` return immediately.
///
/// This method is intended to be used in situations where a notification
/// needs to otherwise be sent to the main reactor. If the reactor is
/// currently blocked inside of `turn` then it will wake up and soon return
/// after this method has been called. If the reactor is not currently
/// blocked in `turn`, then the next call to `turn` will not block and
/// return immediately.
fn wakeup(&self) {
if let Some(inner) = self.inner() {
inner.wakeup.set_readiness(mio::Ready::readable()).unwrap();
}
}
pub(super) fn inner(&self) -> Option<Arc<Inner>> {
self.inner.upgrade()
}
}
impl Unpark for Handle {
fn unpark(&self) {
self.wakeup();
}
}
impl fmt::Debug for Handle {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "Handle")
}
}
// ===== impl Inner =====
impl Inner {
/// Register an I/O resource with the reactor.
///
/// The registration token is returned.
pub(super) fn add_source(&self, source: &dyn Evented) -> io::Result<usize> {
// Get an ABA guard value
let aba_guard = self.next_aba_guard.fetch_add(1 << TOKEN_SHIFT, Relaxed);
let key = {
// Block to contain the write lock
let mut io_dispatch = self.io_dispatch.write().unwrap();
if io_dispatch.len() == MAX_SOURCES {
return Err(io::Error::new(
io::ErrorKind::Other,
"reactor at max \
registered I/O resources",
));
}
io_dispatch.insert(ScheduledIo {
aba_guard,
readiness: AtomicUsize::new(0),
reader: AtomicWaker::new(),
writer: AtomicWaker::new(),
})
};
let token = aba_guard | key;
self.io.register(
source,
mio::Token(token),
mio::Ready::all(),
mio::PollOpt::edge(),
)?;
Ok(key)
}
/// Deregisters an I/O resource from the reactor.
pub(super) fn deregister_source(&self, source: &dyn Evented) -> io::Result<()> {
self.io.deregister(source)
}
pub(super) fn drop_source(&self, token: usize) {
self.io_dispatch.write().unwrap().remove(token);
}
/// Registers interest in the I/O resource associated with `token`.
pub(super) fn register(&self, token: usize, dir: Direction, w: Waker) {
let io_dispatch = self.io_dispatch.read().unwrap();
let sched = io_dispatch.get(token).unwrap();
let (waker, ready) = match dir {
Direction::Read => (&sched.reader, !mio::Ready::writable()),
Direction::Write => (&sched.writer, mio::Ready::writable()),
};
waker.register(w);
if sched.readiness.load(SeqCst) & ready.as_usize() != 0 {
waker.wake();
}
}
}
impl Drop for Inner {
fn drop(&mut self) {
// When a reactor is dropped it needs to wake up all blocked tasks as
// they'll never receive a notification, and all connected I/O objects
// will start returning errors pretty quickly.
let io = self.io_dispatch.read().unwrap();
for (_, io) in io.iter() {
io.writer.wake();
io.reader.wake();
}
}
}
impl Direction {
pub(super) fn mask(self) -> mio::Ready {
match self {
Direction::Read => {
// Everything except writable is signaled through read.
mio::Ready::all() - mio::Ready::writable()
}
Direction::Write => mio::Ready::writable() | platform::hup(),
}
}
}
tokio/src/net/driver/registration.rs
+269
View File
@@ -0,0 +1,269 @@
use super::platform;
use super::reactor::{Direction, Handle};
use mio::{self, Evented};
use std::sync::atomic::Ordering::SeqCst;
use std::task::{Context, Poll};
use std::{io, usize};
/// Associates an I/O resource with the reactor instance that drives it.
///
/// A registration represents an I/O resource registered with a Reactor such
/// that it will receive task notifications on readiness. This is the lowest
/// level API for integrating with a reactor.
///
/// The association between an I/O resource is made by calling [`new`]. Once
/// the association is established, it remains established until the
/// registration instance is dropped.
///
/// A registration instance represents two separate readiness streams. One for
/// the read readiness and one for write readiness. These streams are
/// independent and can be consumed from separate tasks.
///
/// **Note**: while `Registration` is `Sync`, the caller must ensure that there
/// are at most two tasks that use a registration instance concurrently. One
/// task for [`poll_read_ready`] and one task for [`poll_write_ready`]. While
/// violating this requirement is "safe" from a Rust memory safety point of
/// view, it will result in unexpected behavior in the form of lost
/// notifications and tasks hanging.
///
/// ## Platform-specific events
///
/// `Registration` also allows receiving platform-specific `mio::Ready` events.
/// These events are included as part of the read readiness event stream. The
/// write readiness event stream is only for `Ready::writable()` events.
///
/// [`new`]: #method.new
/// [`poll_read_ready`]: #method.poll_read_ready`]
/// [`poll_write_ready`]: #method.poll_write_ready`]
#[derive(Debug)]
pub struct Registration {
handle: Handle,
token: usize,
}
// ===== impl Registration =====
impl Registration {
/// Register the I/O resource with the default reactor.
///
/// # Return
///
/// - `Ok` if the registration happened successfully
/// - `Err` if an error was encountered during registration
pub fn new<T>(io: &T) -> io::Result<Self>
where
T: Evented,
{
let handle = Handle::current();
let token = if let Some(inner) = handle.inner() {
inner.add_source(io)?
} else {
return Err(io::Error::new(
io::ErrorKind::Other,
"failed to find event loop",
));
};
Ok(Self { handle, token })
}
/// Deregister the I/O resource from the reactor it is associated with.
///
/// This function must be called before the I/O resource associated with the
/// registration is dropped.
///
/// Note that deregistering does not guarantee that the I/O resource can be
/// registered with a different reactor. Some I/O resource types can only be
/// associated with a single reactor instance for their lifetime.
///
/// # Return
///
/// If the deregistration was successful, `Ok` is returned. Any calls to
/// `Reactor::turn` that happen after a successful call to `deregister` will
/// no longer result in notifications getting sent for this registration.
///
/// `Err` is returned if an error is encountered.
pub fn deregister<T>(&mut self, io: &T) -> io::Result<()>
where
T: Evented,
{
let inner = match self.handle.inner() {
Some(inner) => inner,
None => return Err(io::Error::new(io::ErrorKind::Other, "reactor gone")),
};
inner.deregister_source(io)
}
/// Poll for events on the I/O resource's read readiness stream.
///
/// If the I/O resource receives a new read readiness event since the last
/// call to `poll_read_ready`, it is returned. If it has not, the current
/// task is notified once a new event is received.
///
/// All events except `HUP` are [edge-triggered]. Once `HUP` is returned,
/// the function will always return `Ready(HUP)`. This should be treated as
/// the end of the readiness stream.
///
/// Ensure that [`register`] has been called first.
///
/// # Return value
///
/// There are several possible return values:
///
/// * `Poll::Ready(Ok(readiness))` means that the I/O resource has received
/// a new readiness event. The readiness value is included.
///
/// * `Poll::Pending` means that no new readiness events have been received
/// since the last call to `poll_read_ready`.
///
/// * `Poll::Ready(Err(err))` means that the registration has encountered an
/// error. This error either represents a permanent internal error **or**
/// the fact that [`register`] was not called first.
///
/// [`register`]: #method.register
/// [edge-triggered]: https://docs.rs/mio/0.6/mio/struct.Poll.html#edge-triggered-and-level-triggered
///
/// # Panics
///
/// This function will panic if called from outside of a task context.
pub fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<mio::Ready>> {
let v = self.poll_ready(Direction::Read, Some(cx))?;
match v {
Some(v) => Poll::Ready(Ok(v)),
None => Poll::Pending,
}
}
/// Consume any pending read readiness event.
///
/// This function is identical to [`poll_read_ready`] **except** that it
/// will not notify the current task when a new event is received. As such,
/// it is safe to call this function from outside of a task context.
///
/// [`poll_read_ready`]: #method.poll_read_ready
pub fn take_read_ready(&self) -> io::Result<Option<mio::Ready>> {
self.poll_ready(Direction::Read, None)
}
/// Poll for events on the I/O resource's write readiness stream.
///
/// If the I/O resource receives a new write readiness event since the last
/// call to `poll_write_ready`, it is returned. If it has not, the current
/// task is notified once a new event is received.
///
/// All events except `HUP` are [edge-triggered]. Once `HUP` is returned,
/// the function will always return `Ready(HUP)`. This should be treated as
/// the end of the readiness stream.
///
/// Ensure that [`register`] has been called first.
///
/// # Return value
///
/// There are several possible return values:
///
/// * `Poll::Ready(Ok(readiness))` means that the I/O resource has received
/// a new readiness event. The readiness value is included.
///
/// * `Poll::Pending` means that no new readiness events have been received
/// since the last call to `poll_write_ready`.
///
/// * `Poll::Ready(Err(err))` means that the registration has encountered an
/// error. This error either represents a permanent internal error **or**
/// the fact that [`register`] was not called first.
///
/// [`register`]: #method.register
/// [edge-triggered]: https://docs.rs/mio/0.6/mio/struct.Poll.html#edge-triggered-and-level-triggered
///
/// # Panics
///
/// This function will panic if called from outside of a task context.
pub fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<mio::Ready>> {
let v = self.poll_ready(Direction::Write, Some(cx))?;
match v {
Some(v) => Poll::Ready(Ok(v)),
None => Poll::Pending,
}
}
/// Consume any pending write readiness event.
///
/// This function is identical to [`poll_write_ready`] **except** that it
/// will not notify the current task when a new event is received. As such,
/// it is safe to call this function from outside of a task context.
///
/// [`poll_write_ready`]: #method.poll_write_ready
pub fn take_write_ready(&self) -> io::Result<Option<mio::Ready>> {
self.poll_ready(Direction::Write, None)
}
/// Poll for events on the I/O resource's `direction` readiness stream.
///
/// If called with a task context, notify the task when a new event is
/// received.
fn poll_ready(
&self,
direction: Direction,
cx: Option<&mut Context<'_>>,
) -> io::Result<Option<mio::Ready>> {
let inner = match self.handle.inner() {
Some(inner) => inner,
None => return Err(io::Error::new(io::ErrorKind::Other, "reactor gone")),
};
// If the task should be notified about new events, ensure that it has
// been registered
if let Some(ref cx) = cx {
inner.register(self.token, direction, cx.waker().clone())
}
let mask = direction.mask();
let mask_no_hup = (mask - platform::hup()).as_usize();
let io_dispatch = inner.io_dispatch.read().unwrap();
let sched = &io_dispatch[self.token];
// This consumes the current readiness state **except** for HUP. HUP is
// excluded because a) it is a final state and never transitions out of
// HUP and b) both the read AND the write directions need to be able to
// observe this state.
//
// If HUP were to be cleared when `direction` is `Read`, then when
// `poll_ready` is called again with a _`direction` of `Write`, the HUP
// state would not be visible.
let mut ready =
mask & mio::Ready::from_usize(sched.readiness.fetch_and(!mask_no_hup, SeqCst));
if ready.is_empty() {
if let Some(cx) = cx {
// Update the task info
match direction {
Direction::Read => sched.reader.register_by_ref(cx.waker()),
Direction::Write => sched.writer.register_by_ref(cx.waker()),
}
// Try again
ready =
mask & mio::Ready::from_usize(sched.readiness.fetch_and(!mask_no_hup, SeqCst));
}
}
if ready.is_empty() {
Ok(None)
} else {
Ok(Some(ready))
}
}
}
unsafe impl Send for Registration {}
unsafe impl Sync for Registration {}
impl Drop for Registration {
fn drop(&mut self) {
let inner = match self.handle.inner() {
Some(inner) => inner,
None => return,
};
inner.drop_source(self.token);
}
}
tokio/src/net/driver/sharded_rwlock.rs
+217
View File
@@ -0,0 +1,217 @@
//! A scalable reader-writer lock.
//!
//! This implementation makes read operations faster and more scalable due to less contention,
//! while making write operations slower. It also incurs much higher memory overhead than
//! traditional reader-writer locks.
use crossbeam_utils::CachePadded;
use lazy_static::lazy_static;
use num_cpus;
use parking_lot;
use std::cell::UnsafeCell;
use std::collections::HashMap;
use std::marker::PhantomData;
use std::mem;
use std::ops::{Deref, DerefMut};
use std::sync::Mutex;
use std::thread::{self, ThreadId};
/// A scalable read-writer lock.
///
/// This type of lock allows a number of readers or at most one writer at any point in time. The
/// write portion of this lock typically allows modification of the underlying data (exclusive
/// access) and the read portion of this lock typically allows for read-only access (shared
/// access).
///
/// This reader-writer lock differs from typical implementations in that it internally creates a
/// list of reader-writer locks called 'shards'. Shards are aligned and padded to the cache line
/// size.
///
/// Read operations lock only one shard specific to the current thread, while write operations lock
/// every shard in succession. This strategy makes concurrent read operations faster due to less
/// contention, but write operations are slower due to increased amount of locking.
pub(crate) struct RwLock<T> {
/// A list of locks protecting the internal data.
shards: Vec<CachePadded<parking_lot::RwLock<()>>>,
/// The internal data.
value: UnsafeCell<T>,
}
unsafe impl<T: Send> Send for RwLock<T> {}
unsafe impl<T: Send + Sync> Sync for RwLock<T> {}
impl<T> RwLock<T> {
/// Creates a new `RwLock` initialized with `value`.
pub(crate) fn new(value: T) -> RwLock<T> {
// The number of shards is a power of two so that the modulo operation in `read` becomes a
// simple bitwise "and".
let num_shards = num_cpus::get().next_power_of_two();
RwLock {
shards: (0..num_shards)
.map(|_| CachePadded::new(parking_lot::RwLock::new(())))
.collect(),
value: UnsafeCell::new(value),
}
}
/// Locks this `RwLock` with shared read access, blocking the current thread until it can be
/// acquired.
///
/// The calling thread will be blocked until there are no more writers which hold the lock.
/// There may be other readers currently inside the lock when this method returns. This method
/// does not provide any guarantees with respect to the ordering of whether contentious readers
/// or writers will acquire the lock first.
///
/// Returns an RAII guard which will release this thread's shared access once it is dropped.
pub(crate) fn read(&self) -> RwLockReadGuard<'_, T> {
// Take the current thread index and map it to a shard index. Thread indices will tend to
// distribute shards among threads equally, thus reducing contention due to read-locking.
let shard_index = thread_index() & (self.shards.len() - 1);
RwLockReadGuard {
parent: self,
_guard: self.shards[shard_index].read(),
_marker: PhantomData,
}
}
/// Locks this rwlock with exclusive write access, blocking the current thread until it can be
/// acquired.
///
/// This function will not return while other writers or other readers currently have access to
/// the lock.
///
/// Returns an RAII guard which will drop the write access of this rwlock when dropped.
pub(crate) fn write(&self) -> RwLockWriteGuard<'_, T> {
// Write-lock each shard in succession.
for shard in &self.shards {
// The write guard is forgotten, but the lock will be manually unlocked in `drop`.
mem::forget(shard.write());
}
RwLockWriteGuard {
parent: self,
_marker: PhantomData,
}
}
}
/// A guard used to release the shared read access of a `RwLock` when dropped.
pub(crate) struct RwLockReadGuard<'a, T> {
parent: &'a RwLock<T>,
_guard: parking_lot::RwLockReadGuard<'a, ()>,
_marker: PhantomData<parking_lot::RwLockReadGuard<'a, T>>,
}
unsafe impl<'a, T: Sync> Sync for RwLockReadGuard<'a, T> {}
impl<'a, T> Deref for RwLockReadGuard<'a, T> {
type Target = T;
fn deref(&self) -> &T {
unsafe { &*self.parent.value.get() }
}
}
/// A guard used to release the exclusive write access of a `RwLock` when dropped.
pub(crate) struct RwLockWriteGuard<'a, T> {
parent: &'a RwLock<T>,
_marker: PhantomData<parking_lot::RwLockWriteGuard<'a, T>>,
}
unsafe impl<'a, T: Sync> Sync for RwLockWriteGuard<'a, T> {}
impl<'a, T> Drop for RwLockWriteGuard<'a, T> {
fn drop(&mut self) {
// Unlock the shards in reverse order of locking.
for shard in self.parent.shards.iter().rev() {
unsafe {
shard.force_unlock_write();
}
}
}
}
impl<'a, T> Deref for RwLockWriteGuard<'a, T> {
type Target = T;
fn deref(&self) -> &T {
unsafe { &*self.parent.value.get() }
}
}
impl<'a, T> DerefMut for RwLockWriteGuard<'a, T> {
fn deref_mut(&mut self) -> &mut T {
unsafe { &mut *self.parent.value.get() }
}
}
/// Returns a `usize` that identifies the current thread.
///
/// Each thread is associated with an 'index'. Indices usually tend to be consecutive numbers
/// between 0 and the number of running threads, but there are no guarantees. During TLS teardown
/// the associated index might change.
#[inline]
pub(crate) fn thread_index() -> usize {
REGISTRATION.try_with(|reg| reg.index).unwrap_or(0)
}
/// The global registry keeping track of registered threads and indices.
struct ThreadIndices {
/// Mapping from `ThreadId` to thread index.
mapping: HashMap<ThreadId, usize>,
/// A list of free indices.
free_list: Vec<usize>,
/// The next index to allocate if the free list is empty.
next_index: usize,
}
lazy_static! {
static ref THREAD_INDICES: Mutex<ThreadIndices> = Mutex::new(ThreadIndices {
mapping: HashMap::new(),
free_list: Vec::new(),
next_index: 0,
});
}
/// A registration of a thread with an index.
///
/// When dropped, unregisters the thread and frees the reserved index.
struct Registration {
index: usize,
thread_id: ThreadId,
}
impl Drop for Registration {
fn drop(&mut self) {
let mut indices = THREAD_INDICES.lock().unwrap();
indices.mapping.remove(&self.thread_id);
indices.free_list.push(self.index);
}
}
thread_local! {
static REGISTRATION: Registration = {
let thread_id = thread::current().id();
let mut indices = THREAD_INDICES.lock().unwrap();
let index = match indices.free_list.pop() {
Some(i) => i,
None => {
let i = indices.next_index;
indices.next_index += 1;
i
}
};
indices.mapping.insert(thread_id, index);
Registration {
index,
thread_id,
}
};
}
tokio/src/net/mod.rs
+47
View File
@@ -0,0 +1,47 @@
//! TCP/UDP/Unix bindings for `tokio`.
//!
//! This module contains the TCP/UDP/Unix networking types, similar to the standard
//! library, which can be used to implement networking protocols.
//!
//! # Organization
//!
//! * [`TcpListener`] and [`TcpStream`] provide functionality for communication over TCP
//! * [`UdpSocket`] provides functionality for communication over UDP
//! * [`UnixListener`] and [`UnixStream`] provide functionality for communication over a
//! Unix Domain Stream Socket **(available on Unix only)**
//! * [`UnixDatagram`] and [`UnixDatagramFramed`] provide functionality for communication
//! over Unix Domain Datagram Socket **(available on Unix only)**
//!
//! [`TcpListener`]: struct.TcpListener.html
//! [`TcpStream`]: struct.TcpStream.html
//! [`UdpSocket`]: struct.UdpSocket.html
//! [`UnixListener`]: struct.UnixListener.html
//! [`UnixStream`]: struct.UnixStream.html
//! [`UnixDatagram`]: struct.UnixDatagram.html
//! [`UnixDatagramFramed`]: struct.UnixDatagramFramed.html
mod addr;
pub use addr::ToSocketAddrs;
pub mod driver;
pub mod util;
#[cfg(feature = "tcp")]
pub mod tcp;
#[cfg(feature = "tcp")]
pub use self::tcp::{TcpListener, TcpStream};
#[cfg(feature = "udp")]
pub mod udp;
#[cfg(feature = "udp")]
pub use self::udp::UdpSocket;
#[cfg(all(unix, feature = "uds"))]
pub mod unix;
#[cfg(all(unix, feature = "uds"))]
pub use self::unix::{UnixDatagram, UnixListener, UnixStream};
tokio/src/net/tcp/incoming.rs
+31
View File
@@ -0,0 +1,31 @@
use crate::net::tcp::TcpListener;
use crate::net::tcp::TcpStream;
use futures_core::ready;
use futures_core::stream::Stream;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
/// Stream returned by the `TcpListener::incoming` function representing the
/// stream of sockets received from a listener.
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct Incoming {
inner: TcpListener,
}
impl Incoming {
pub(crate) fn new(listener: TcpListener) -> Incoming {
Incoming { inner: listener }
}
}
impl Stream for Incoming {
type Item = io::Result<TcpStream>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let (socket, _) = ready!(self.inner.poll_accept(cx))?;
Poll::Ready(Some(Ok(socket)))
}
}
tokio/src/net/tcp/listener.rs
+358
View File
@@ -0,0 +1,358 @@
use crate::net::tcp::{Incoming, TcpStream};
use crate::net::util::PollEvented;
use crate::net::ToSocketAddrs;
use futures_core::ready;
use futures_util::future::poll_fn;
use std::convert::TryFrom;
use std::fmt;
use std::io;
use std::net::{self, SocketAddr};
use std::task::{Context, Poll};
/// An I/O object representing a TCP socket listening for incoming connections.
///
/// This object can be converted into a stream of incoming connections for
/// various forms of processing.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpListener;
///
/// use std::io;
/// # async fn process_socket<T>(_socket: T) {}
///
/// #[tokio::main]
/// async fn main() -> io::Result<()> {
/// let mut listener = TcpListener::bind("127.0.0.1:8080").await?;
///
/// loop {
/// let (socket, _) = listener.accept().await?;
/// process_socket(socket).await;
/// }
/// }
/// ```
pub struct TcpListener {
io: PollEvented<mio::net::TcpListener>,
}
impl TcpListener {
/// Creates a new TcpListener which will be bound to the specified address.
///
/// The returned listener is ready for accepting connections.
///
/// Binding with a port number of 0 will request that the OS assigns a port
/// to this listener. The port allocated can be queried via the `local_addr`
/// method.
///
/// The address type can be any implementor of `ToSocketAddrs` trait.
///
/// If `addr` yields multiple addresses, bind will be attempted with each of
/// the addresses until one succeeds and returns the listener. If none of
/// the addresses succeed in creating a listener, the error returned from
/// the last attempt (the last address) is returned.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpListener;
///
/// use std::io;
///
/// #[tokio::main]
/// async fn main() -> io::Result<()> {
/// let listener = TcpListener::bind("127.0.0.1:0").await?;
///
/// // use the listener
///
/// # let _ = listener;
/// Ok(())
/// }
/// ```
pub async fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<TcpListener> {
let addrs = addr.to_socket_addrs().await?;
let mut last_err = None;
for addr in addrs {
match TcpListener::bind_addr(addr) {
Ok(listener) => return Ok(listener),
Err(e) => last_err = Some(e),
}
}
Err(last_err.unwrap_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"could not resolve to any addresses",
)
}))
}
fn bind_addr(addr: SocketAddr) -> io::Result<TcpListener> {
let listener = mio::net::TcpListener::bind(&addr)?;
TcpListener::new(listener)
}
/// Accept a new incoming connection from this listener.
///
/// This function will yield once a new TCP connection is established. When
/// established, the corresponding [`TcpStream`] and the remote peer's
/// address will be returned.
///
/// [`TcpStream`]: ../struct.TcpStream.html
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpListener;
///
/// use std::io;
///
/// #[tokio::main]
/// async fn main() -> io::Result<()> {
/// let mut listener = TcpListener::bind("127.0.0.1:8080").await?;
///
/// match listener.accept().await {
/// Ok((_socket, addr)) => println!("new client: {:?}", addr),
/// Err(e) => println!("couldn't get client: {:?}", e),
/// }
///
/// Ok(())
/// }
/// ```
pub async fn accept(&mut self) -> io::Result<(TcpStream, SocketAddr)> {
poll_fn(|cx| self.poll_accept(cx)).await
}
pub(crate) fn poll_accept(
&mut self,
cx: &mut Context<'_>,
) -> Poll<io::Result<(TcpStream, SocketAddr)>> {
let (io, addr) = ready!(self.poll_accept_std(cx))?;
let io = mio::net::TcpStream::from_stream(io)?;
let io = TcpStream::new(io)?;
Poll::Ready(Ok((io, addr)))
}
fn poll_accept_std(
&mut self,
cx: &mut Context<'_>,
) -> Poll<io::Result<(net::TcpStream, SocketAddr)>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
match self.io.get_ref().accept_std() {
Ok(pair) => Poll::Ready(Ok(pair)),
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
Err(e) => Poll::Ready(Err(e)),
}
}
/// Create a new TCP listener from the standard library's TCP listener.
///
/// This method can be used when the `Handle::tcp_listen` method isn't
/// sufficient because perhaps some more configuration is needed in terms of
/// before the calls to `bind` and `listen`.
///
/// This API is typically paired with the `net2` crate and the `TcpBuilder`
/// type to build up and customize a listener before it's shipped off to the
/// backing event loop. This allows configuration of options like
/// `SO_REUSEPORT`, binding to multiple addresses, etc.
///
/// The `addr` argument here is one of the addresses that `listener` is
/// bound to and the listener will only be guaranteed to accept connections
/// of the same address type currently.
///
/// The platform specific behavior of this function looks like:
///
/// * On Unix, the socket is placed into nonblocking mode and connections
/// can be accepted as normal
///
/// * On Windows, the address is stored internally and all future accepts
/// will only be for the same IP version as `addr` specified. That is, if
/// `addr` is an IPv4 address then all sockets accepted will be IPv4 as
/// well (same for IPv6).
///
/// # Examples
///
/// ```rust,no_run
/// use std::error::Error;
/// use tokio::net::TcpListener;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn Error>> {
/// let std_listener = std::net::TcpListener::bind("127.0.0.1:0")?;
/// let listener = TcpListener::from_std(std_listener)?;
/// Ok(())
/// }
/// ```
pub fn from_std(listener: net::TcpListener) -> io::Result<TcpListener> {
let io = mio::net::TcpListener::from_std(listener)?;
let io = PollEvented::new(io)?;
Ok(TcpListener { io })
}
fn new(listener: mio::net::TcpListener) -> io::Result<TcpListener> {
let io = PollEvented::new(listener)?;
Ok(TcpListener { io })
}
/// Returns the local address that this listener is bound to.
///
/// This can be useful, for example, when binding to port 0 to figure out
/// which port was actually bound.
///
/// # Examples
///
/// ```rust,no_run
/// use tokio::net::TcpListener;
///
/// use std::io;
/// use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
///
/// #[tokio::main]
/// async fn main() -> io::Result<()> {
/// let listener = TcpListener::bind("127.0.0.1:8080").await?;
///
/// assert_eq!(listener.local_addr()?,
/// SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 8080)));
///
/// Ok(())
/// }
/// ```
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.io.get_ref().local_addr()
}
/// Consumes this listener, returning a stream of the sockets this listener
/// accepts.
///
/// This method returns an implementation of the `Stream` trait which
/// resolves to the sockets the are accepted on this listener.
///
/// # Errors
///
/// Note that accepting a connection can lead to various errors and not all of them are
/// necessarily fatal ‒ for example having too many open file descriptors or the other side
/// closing the connection while it waits in an accept queue. These would terminate the stream
/// if not handled in any way.
pub fn incoming(self) -> Incoming {
Incoming::new(self)
}
/// Gets the value of the `IP_TTL` option for this socket.
///
/// For more information about this option, see [`set_ttl`].
///
/// [`set_ttl`]: #method.set_ttl
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpListener;
///
/// use std::io;
///
/// #[tokio::main]
/// async fn main() -> io::Result<()> {
/// let listener = TcpListener::bind("127.0.0.1:0").await?;
///
/// listener.set_ttl(100).expect("could not set TTL");
/// assert_eq!(listener.ttl()?, 100);
///
/// Ok(())
/// }
/// ```
pub fn ttl(&self) -> io::Result<u32> {
self.io.get_ref().ttl()
}
/// Sets the value for the `IP_TTL` option on this socket.
///
/// This value sets the time-to-live field that is used in every packet sent
/// from this socket.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpListener;
///
/// use std::io;
///
/// #[tokio::main]
/// async fn main() -> io::Result<()> {
/// let listener = TcpListener::bind("127.0.0.1:0").await?;
///
/// listener.set_ttl(100).expect("could not set TTL");
///
/// Ok(())
/// }
/// ```
pub fn set_ttl(&self, ttl: u32) -> io::Result<()> {
self.io.get_ref().set_ttl(ttl)
}
}
impl TryFrom<TcpListener> for mio::net::TcpListener {
type Error = io::Error;
/// Consumes value, returning the mio I/O object.
///
/// See [`PollEvented::into_inner`] for more details about
/// resource deregistration that happens during the call.
///
/// [`PollEvented::into_inner`]: crate::util::PollEvented::into_inner
fn try_from(value: TcpListener) -> Result<Self, Self::Error> {
value.io.into_inner()
}
}
impl TryFrom<net::TcpListener> for TcpListener {
type Error = io::Error;
/// Consumes stream, returning the tokio I/O object.
///
/// This is equivalent to
/// [`TcpListener::from_std(stream)`](TcpListener::from_std).
fn try_from(stream: net::TcpListener) -> Result<Self, Self::Error> {
Self::from_std(stream)
}
}
impl fmt::Debug for TcpListener {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.io.get_ref().fmt(f)
}
}
#[cfg(unix)]
mod sys {
use super::TcpListener;
use std::os::unix::prelude::*;
impl AsRawFd for TcpListener {
fn as_raw_fd(&self) -> RawFd {
self.io.get_ref().as_raw_fd()
}
}
}
#[cfg(windows)]
mod sys {
// TODO: let's land these upstream with mio and then we can add them here.
//
// use std::os::windows::prelude::*;
// use super::{TcpListener;
//
// impl AsRawHandle for TcpListener {
// fn as_raw_handle(&self) -> RawHandle {
// self.listener.io().as_raw_handle()
// }
// }
}
tokio/src/net/tcp/mod.rs
+29
View File
@@ -0,0 +1,29 @@
//! TCP bindings for `tokio`.
//!
//! This module contains the TCP networking types, similar to the standard
//! library, which can be used to implement networking protocols.
//!
//! Connecting to an address, via TCP, can be done using [`TcpStream`]'s
//! [`connect`] method, which returns a future which returns a `TcpStream`.
//!
//! To listen on an address [`TcpListener`] can be used. `TcpListener`'s
//! [`incoming`][incoming_method] method can be used to accept new connections.
//! It return the [`Incoming`] struct, which implements a stream which returns
//! `TcpStream`s.
//!
//! [`TcpStream`]: struct.TcpStream.html
//! [`connect`]: struct.TcpStream.html#method.connect
//! [`TcpListener`]: struct.TcpListener.html
//! [incoming_method]: struct.TcpListener.html#method.incoming
//! [`Incoming`]: struct.Incoming.html
mod incoming;
pub use self::incoming::Incoming;
mod listener;
pub use self::listener::TcpListener;
pub mod split;
mod stream;
pub use self::stream::TcpStream;
tokio/src/net/tcp/split.rs
+97
View File
@@ -0,0 +1,97 @@
//! `TcpStream` split support.
//!
//! A `TcpStream` can be split into a `ReadHalf` and a
//! `WriteHalf` with the `TcpStream::split` method. `ReadHalf`
//! implements `AsyncRead` while `WriteHalf` implements `AsyncWrite`.
//!
//! Compared to the generic split of `AsyncRead + AsyncWrite`, this specialized
//! split has no associated overhead and enforces all invariants at the type
//! level.
use super::TcpStream;
use tokio_io::{AsyncRead, AsyncWrite};
use bytes::{Buf, BufMut};
use std::io;
use std::net::Shutdown;
use std::pin::Pin;
use std::task::{Context, Poll};
/// Read half of a `TcpStream`.
#[derive(Debug)]
pub struct ReadHalf<'a>(&'a TcpStream);
/// Write half of a `TcpStream`.
///
/// Note that in the `AsyncWrite` implemenation of `TcpStreamWriteHalf`,
/// `poll_shutdown` actually shuts down the TCP stream in the write direction.
#[derive(Debug)]
pub struct WriteHalf<'a>(&'a TcpStream);
pub(crate) fn split(stream: &mut TcpStream) -> (ReadHalf<'_>, WriteHalf<'_>) {
(ReadHalf(&*stream), WriteHalf(&*stream))
}
impl AsyncRead for ReadHalf<'_> {
unsafe fn prepare_uninitialized_buffer(&self, _: &mut [u8]) -> bool {
false
}
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
self.0.poll_read_priv(cx, buf)
}
fn poll_read_buf<B: BufMut>(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
self.0.poll_read_buf_priv(cx, buf)
}
}
impl AsyncWrite for WriteHalf<'_> {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.0.poll_write_priv(cx, buf)
}
#[inline]
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
// tcp flush is a no-op
Poll::Ready(Ok(()))
}
// `poll_shutdown` on a write half shutdowns the stream in the "write" direction.
fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
self.0.shutdown(Shutdown::Write).into()
}
fn poll_write_buf<B: Buf>(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
self.0.poll_write_buf_priv(cx, buf)
}
}
impl AsRef<TcpStream> for ReadHalf<'_> {
fn as_ref(&self) -> &TcpStream {
self.0
}
}
impl AsRef<TcpStream> for WriteHalf<'_> {
fn as_ref(&self) -> &TcpStream {
self.0
}
}
tokio/src/net/tcp/stream.rs
+827
View File
@@ -0,0 +1,827 @@
use crate::net::tcp::split::{split, ReadHalf, WriteHalf};
use crate::net::util::PollEvented;
use crate::net::ToSocketAddrs;
use tokio_io::{AsyncRead, AsyncWrite};
use bytes::{Buf, BufMut};
use futures_core::ready;
use futures_util::future::poll_fn;
use iovec::IoVec;
use std::convert::TryFrom;
use std::fmt;
use std::io::{self, Read, Write};
use std::net::{self, Shutdown, SocketAddr};
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
/// An I/O object representing a TCP stream connected to a remote endpoint.
///
/// A TCP stream can either be created by connecting to an endpoint, via the
/// [`connect`] method, or by [accepting] a connection from a [listener].
///
/// [`connect`]: struct.TcpStream.html#method.connect
/// [accepting]: struct.TcpListener.html#method.accept
/// [listener]: struct.TcpListener.html
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
/// use tokio::prelude::*;
/// use std::error::Error;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn Error>> {
/// // Connect to a peer
/// let mut stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// // Write some data.
/// stream.write_all(b"hello world!").await?;
///
/// Ok(())
/// }
/// ```
pub struct TcpStream {
io: PollEvented<mio::net::TcpStream>,
}
impl TcpStream {
/// Opens a TCP connection to a remote host.
///
/// `addr` is an address of the remote host. Anything which implements
/// `ToSocketAddrs` trait can be supplied for the address.
///
/// If `addr` yields multiple addresses, connect will be attempted with each
/// of the addresses until a connection is successful. If none of the
/// addresses result in a successful connection, the error returned from the
/// last connection attempt (the last address) is returned.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
/// use tokio::prelude::*;
/// use std::error::Error;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn Error>> {
/// // Connect to a peer
/// let mut stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// // Write some data.
/// stream.write_all(b"hello world!").await?;
///
/// Ok(())
/// }
/// ```
pub async fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<TcpStream> {
let addrs = addr.to_socket_addrs().await?;
let mut last_err = None;
for addr in addrs {
match TcpStream::connect_addr(addr).await {
Ok(stream) => return Ok(stream),
Err(e) => last_err = Some(e),
}
}
Err(last_err.unwrap_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"could not resolve to any addresses",
)
}))
}
/// Establish a connection to the specified `addr`.
async fn connect_addr(addr: SocketAddr) -> io::Result<TcpStream> {
let sys = mio::net::TcpStream::connect(&addr)?;
let stream = TcpStream::new(sys)?;
// Once we've connected, wait for the stream to be writable as
// that's when the actual connection has been initiated. Once we're
// writable we check for `take_socket_error` to see if the connect
// actually hit an error or not.
//
// If all that succeeded then we ship everything on up.
poll_fn(|cx| stream.io.poll_write_ready(cx)).await?;
if let Some(e) = stream.io.get_ref().take_error()? {
return Err(e);
}
Ok(stream)
}
pub(crate) fn new(connected: mio::net::TcpStream) -> io::Result<TcpStream> {
let io = PollEvented::new(connected)?;
Ok(TcpStream { io })
}
/// Create a new `TcpStream` from a `std::net::TcpStream`.
///
/// This function will convert a TCP stream created by the standard library
/// to a TCP stream ready to be used with the provided event loop handle.
///
/// # Examples
///
/// ```rust,no_run
/// use std::error::Error;
/// use tokio::net::TcpStream;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn Error>> {
/// let std_stream = std::net::TcpStream::connect("127.0.0.1:34254")?;
/// let stream = TcpStream::from_std(std_stream)?;
/// Ok(())
/// }
/// ```
pub fn from_std(stream: net::TcpStream) -> io::Result<TcpStream> {
let io = mio::net::TcpStream::from_stream(stream)?;
let io = PollEvented::new(io)?;
Ok(TcpStream { io })
}
// Connect a TcpStream asynchronously that may be built with a net2 TcpBuilder.
//
// This should be removed in favor of some in-crate TcpSocket builder API.
#[doc(hidden)]
pub async fn connect_std(stream: net::TcpStream, addr: &SocketAddr) -> io::Result<TcpStream> {
let io = mio::net::TcpStream::connect_stream(stream, addr)?;
let io = PollEvented::new(io)?;
let stream = TcpStream { io };
// Once we've connected, wait for the stream to be writable as
// that's when the actual connection has been initiated. Once we're
// writable we check for `take_socket_error` to see if the connect
// actually hit an error or not.
//
// If all that succeeded then we ship everything on up.
poll_fn(|cx| stream.io.poll_write_ready(cx)).await?;
if let Some(e) = stream.io.get_ref().take_error()? {
return Err(e);
}
Ok(stream)
}
/// Returns the local address that this stream is bound to.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// println!("{:?}", stream.local_addr()?);
/// # Ok(())
/// # }
/// ```
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.io.get_ref().local_addr()
}
/// Returns the remote address that this stream is connected to.
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// println!("{:?}", stream.peer_addr()?);
/// # Ok(())
/// # }
/// ```
pub fn peer_addr(&self) -> io::Result<SocketAddr> {
self.io.get_ref().peer_addr()
}
fn poll_peek(&mut self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
match self.io.get_ref().peek(buf) {
Ok(ret) => Poll::Ready(Ok(ret)),
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
Err(e) => Poll::Ready(Err(e)),
}
}
/// Receives data on the socket from the remote address to which it is
/// connected, without removing that data from the queue. On success,
/// returns the number of bytes peeked.
///
/// Successive calls return the same data. This is accomplished by passing
/// `MSG_PEEK` as a flag to the underlying recv system call.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
/// use tokio::prelude::*;
/// use std::error::Error;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn Error>> {
/// // Connect to a peer
/// let mut stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// let mut b1 = [0; 10];
/// let mut b2 = [0; 10];
///
/// // Peek at the data
/// let n = stream.peek(&mut b1).await?;
///
/// // Read the data
/// assert_eq!(n, stream.read(&mut b2[..n]).await?);
/// assert_eq!(&b1[..n], &b2[..n]);
///
/// Ok(())
/// }
/// ```
pub async fn peek(&mut self, buf: &mut [u8]) -> io::Result<usize> {
poll_fn(|cx| self.poll_peek(cx, buf)).await
}
/// Shuts down the read, write, or both halves of this connection.
///
/// This function will cause all pending and future I/O on the specified
/// portions to return immediately with an appropriate value (see the
/// documentation of `Shutdown`).
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
/// use std::error::Error;
/// use std::net::Shutdown;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn Error>> {
/// // Connect to a peer
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// // Shutdown the stream
/// stream.shutdown(Shutdown::Write)?;
///
/// Ok(())
/// }
/// ```
pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
self.io.get_ref().shutdown(how)
}
/// Gets the value of the `TCP_NODELAY` option on this socket.
///
/// For more information about this option, see [`set_nodelay`].
///
/// [`set_nodelay`]: #method.set_nodelay
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// println!("{:?}", stream.nodelay()?);
/// # Ok(())
/// # }
/// ```
pub fn nodelay(&self) -> io::Result<bool> {
self.io.get_ref().nodelay()
}
/// Sets the value of the `TCP_NODELAY` option on this socket.
///
/// If set, this option disables the Nagle algorithm. This means that
/// segments are always sent as soon as possible, even if there is only a
/// small amount of data. When not set, data is buffered until there is a
/// sufficient amount to send out, thereby avoiding the frequent sending of
/// small packets.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// stream.set_nodelay(true)?;
/// # Ok(())
/// # }
/// ```
pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
self.io.get_ref().set_nodelay(nodelay)
}
/// Gets the value of the `SO_RCVBUF` option on this socket.
///
/// For more information about this option, see [`set_recv_buffer_size`].
///
/// [`set_recv_buffer_size`]: #tymethod.set_recv_buffer_size
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// println!("{:?}", stream.recv_buffer_size()?);
/// # Ok(())
/// # }
/// ```
pub fn recv_buffer_size(&self) -> io::Result<usize> {
self.io.get_ref().recv_buffer_size()
}
/// Sets the value of the `SO_RCVBUF` option on this socket.
///
/// Changes the size of the operating system's receive buffer associated
/// with the socket.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// stream.set_recv_buffer_size(100)?;
/// # Ok(())
/// # }
/// ```
pub fn set_recv_buffer_size(&self, size: usize) -> io::Result<()> {
self.io.get_ref().set_recv_buffer_size(size)
}
/// Gets the value of the `SO_SNDBUF` option on this socket.
///
/// For more information about this option, see [`set_send_buffer`].
///
/// [`set_send_buffer`]: #tymethod.set_send_buffer
///
/// # Examples
///
/// Returns whether keepalive messages are enabled on this socket, and if so
/// the duration of time between them.
///
/// For more information about this option, see [`set_keepalive`].
///
/// [`set_keepalive`]: #tymethod.set_keepalive
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// println!("{:?}", stream.send_buffer_size()?);
/// # Ok(())
/// # }
/// ```
pub fn send_buffer_size(&self) -> io::Result<usize> {
self.io.get_ref().send_buffer_size()
}
/// Sets the value of the `SO_SNDBUF` option on this socket.
///
/// Changes the size of the operating system's send buffer associated with
/// the socket.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// stream.set_send_buffer_size(100)?;
/// # Ok(())
/// # }
/// ```
pub fn set_send_buffer_size(&self, size: usize) -> io::Result<()> {
self.io.get_ref().set_send_buffer_size(size)
}
/// Returns whether keepalive messages are enabled on this socket, and if so
/// the duration of time between them.
///
/// For more information about this option, see [`set_keepalive`].
///
/// [`set_keepalive`]: #tymethod.set_keepalive
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// println!("{:?}", stream.keepalive()?);
/// # Ok(())
/// # }
/// ```
pub fn keepalive(&self) -> io::Result<Option<Duration>> {
self.io.get_ref().keepalive()
}
/// Sets whether keepalive messages are enabled to be sent on this socket.
///
/// On Unix, this option will set the `SO_KEEPALIVE` as well as the
/// `TCP_KEEPALIVE` or `TCP_KEEPIDLE` option (depending on your platform).
/// On Windows, this will set the `SIO_KEEPALIVE_VALS` option.
///
/// If `None` is specified then keepalive messages are disabled, otherwise
/// the duration specified will be the time to remain idle before sending a
/// TCP keepalive probe.
///
/// Some platforms specify this value in seconds, so sub-second
/// specifications may be omitted.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// stream.set_keepalive(None)?;
/// # Ok(())
/// # }
/// ```
pub fn set_keepalive(&self, keepalive: Option<Duration>) -> io::Result<()> {
self.io.get_ref().set_keepalive(keepalive)
}
/// Gets the value of the `IP_TTL` option for this socket.
///
/// For more information about this option, see [`set_ttl`].
///
/// [`set_ttl`]: #tymethod.set_ttl
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// println!("{:?}", stream.ttl()?);
/// # Ok(())
/// # }
/// ```
pub fn ttl(&self) -> io::Result<u32> {
self.io.get_ref().ttl()
}
/// Sets the value for the `IP_TTL` option on this socket.
///
/// This value sets the time-to-live field that is used in every packet sent
/// from this socket.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// stream.set_ttl(123)?;
/// # Ok(())
/// # }
/// ```
pub fn set_ttl(&self, ttl: u32) -> io::Result<()> {
self.io.get_ref().set_ttl(ttl)
}
/// Reads the linger duration for this socket by getting the `SO_LINGER`
/// option.
///
/// For more information about this option, see [`set_linger`].
///
/// [`set_linger`]: #tymethod.set_linger
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// println!("{:?}", stream.linger()?);
/// # Ok(())
/// # }
/// ```
pub fn linger(&self) -> io::Result<Option<Duration>> {
self.io.get_ref().linger()
}
/// Sets the linger duration of this socket by setting the `SO_LINGER`
/// option.
///
/// This option controls the action taken when a stream has unsent messages
/// and the stream is closed. If `SO_LINGER` is set, the system
/// shall block the process until it can transmit the data or until the
/// time expires.
///
/// If `SO_LINGER` is not specified, and the stream is closed, the system
/// handles the call in a way that allows the process to continue as quickly
/// as possible.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// stream.set_linger(None)?;
/// # Ok(())
/// # }
/// ```
pub fn set_linger(&self, dur: Option<Duration>) -> io::Result<()> {
self.io.get_ref().set_linger(dur)
}
/// Split a `TcpStream` into a read half and a write half, which can be used
/// to read and write the stream concurrently.
///
/// See the module level documenation of [`split`](super::split) for more
/// details.
pub fn split(&mut self) -> (ReadHalf<'_>, WriteHalf<'_>) {
split(self)
}
// == Poll IO functions that takes `&self` ==
//
// They are not public because (taken from the doc of `PollEvented`):
//
// While `PollEvented` is `Sync` (if the underlying I/O type is `Sync`), the
// caller must ensure that there are at most two tasks that use a
// `PollEvented` instance concurrently. One for reading and one for writing.
// While violating this requirement is "safe" from a Rust memory model point
// of view, it will result in unexpected behavior in the form of lost
// notifications and tasks hanging.
pub(crate) fn poll_read_priv(
&self,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
match self.io.get_ref().read(buf) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
pub(crate) fn poll_read_buf_priv<B: BufMut>(
&self,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
let r = unsafe {
// The `IoVec` type can't have a 0-length size, so we create a bunch
// of dummy versions on the stack with 1 length which we'll quickly
// overwrite.
let b1: &mut [u8] = &mut [0];
let b2: &mut [u8] = &mut [0];
let b3: &mut [u8] = &mut [0];
let b4: &mut [u8] = &mut [0];
let b5: &mut [u8] = &mut [0];
let b6: &mut [u8] = &mut [0];
let b7: &mut [u8] = &mut [0];
let b8: &mut [u8] = &mut [0];
let b9: &mut [u8] = &mut [0];
let b10: &mut [u8] = &mut [0];
let b11: &mut [u8] = &mut [0];
let b12: &mut [u8] = &mut [0];
let b13: &mut [u8] = &mut [0];
let b14: &mut [u8] = &mut [0];
let b15: &mut [u8] = &mut [0];
let b16: &mut [u8] = &mut [0];
let mut bufs: [&mut IoVec; 16] = [
b1.into(),
b2.into(),
b3.into(),
b4.into(),
b5.into(),
b6.into(),
b7.into(),
b8.into(),
b9.into(),
b10.into(),
b11.into(),
b12.into(),
b13.into(),
b14.into(),
b15.into(),
b16.into(),
];
let n = buf.bytes_vec_mut(&mut bufs);
self.io.get_ref().read_bufs(&mut bufs[..n])
};
match r {
Ok(n) => {
unsafe {
buf.advance_mut(n);
}
Poll::Ready(Ok(n))
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
Err(e) => Poll::Ready(Err(e)),
}
}
pub(crate) fn poll_write_priv(
&self,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_write_ready(cx))?;
match self.io.get_ref().write(buf) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_write_ready(cx)?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
pub(crate) fn poll_write_buf_priv<B: Buf>(
&self,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_write_ready(cx))?;
let r = {
// The `IoVec` type can't have a zero-length size, so create a dummy
// version from a 1-length slice which we'll overwrite with the
// `bytes_vec` method.
static DUMMY: &[u8] = &[0];
let iovec = <&IoVec>::from(DUMMY);
let mut bufs = [iovec; 64];
let n = buf.bytes_vec(&mut bufs);
self.io.get_ref().write_bufs(&bufs[..n])
};
match r {
Ok(n) => {
buf.advance(n);
Poll::Ready(Ok(n))
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_write_ready(cx)?;
Poll::Pending
}
Err(e) => Poll::Ready(Err(e)),
}
}
}
impl TryFrom<TcpStream> for mio::net::TcpStream {
type Error = io::Error;
/// Consumes value, returning the mio I/O object.
///
/// See [`PollEvented::into_inner`] for more details about
/// resource deregistration that happens during the call.
///
/// [`PollEvented::into_inner`]: crate::util::PollEvented::into_inner
fn try_from(value: TcpStream) -> Result<Self, Self::Error> {
value.io.into_inner()
}
}
impl TryFrom<net::TcpStream> for TcpStream {
type Error = io::Error;
/// Consumes stream, returning the tokio I/O object.
///
/// This is equivalent to
/// [`TcpStream::from_std(stream)`](TcpStream::from_std).
fn try_from(stream: net::TcpStream) -> Result<Self, Self::Error> {
Self::from_std(stream)
}
}
// ===== impl Read / Write =====
impl AsyncRead for TcpStream {
unsafe fn prepare_uninitialized_buffer(&self, _: &mut [u8]) -> bool {
false
}
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
self.poll_read_priv(cx, buf)
}
fn poll_read_buf<B: BufMut>(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
self.poll_read_buf_priv(cx, buf)
}
}
impl AsyncWrite for TcpStream {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.poll_write_priv(cx, buf)
}
#[inline]
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
// tcp flush is a no-op
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
self.shutdown(std::net::Shutdown::Write)?;
Poll::Ready(Ok(()))
}
fn poll_write_buf<B: Buf>(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
self.poll_write_buf_priv(cx, buf)
}
}
impl fmt::Debug for TcpStream {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.io.get_ref().fmt(f)
}
}
#[cfg(unix)]
mod sys {
use super::TcpStream;
use std::os::unix::prelude::*;
impl AsRawFd for TcpStream {
fn as_raw_fd(&self) -> RawFd {
self.io.get_ref().as_raw_fd()
}
}
}
#[cfg(windows)]
mod sys {
// TODO: let's land these upstream with mio and then we can add them here.
//
// use std::os::windows::prelude::*;
// use super::TcpStream;
//
// impl AsRawHandle for TcpStream {
// fn as_raw_handle(&self) -> RawHandle {
// self.io.get_ref().as_raw_handle()
// }
// }
}
tokio/src/net/udp/mod.rs
+13
View File
@@ -0,0 +1,13 @@
//! UDP bindings for `tokio`.
//!
//! This module contains the UDP networking types, similar to the standard
//! library, which can be used to implement networking protocols.
//!
//! The main struct for UDP is the [`UdpSocket`], which represents a UDP socket.
//!
//! [`UdpSocket`]: struct.UdpSocket
mod socket;
pub mod split;
pub use self::socket::UdpSocket;
tokio/src/net/udp/socket.rs
+419
View File
@@ -0,0 +1,419 @@
use crate::net::udp::split::{split, UdpSocketRecvHalf, UdpSocketSendHalf};
use crate::net::util::PollEvented;
use crate::net::ToSocketAddrs;
use futures_core::ready;
use futures_util::future::poll_fn;
use std::convert::TryFrom;
use std::fmt;
use std::io;
use std::net::{self, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::task::{Context, Poll};
/// An I/O object representing a UDP socket.
pub struct UdpSocket {
io: PollEvented<mio::net::UdpSocket>,
}
impl UdpSocket {
/// This function will create a new UDP socket and attempt to bind it to
/// the `addr` provided.
pub async fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<UdpSocket> {
let addrs = addr.to_socket_addrs().await?;
let mut last_err = None;
for addr in addrs {
match UdpSocket::bind_addr(addr) {
Ok(socket) => return Ok(socket),
Err(e) => last_err = Some(e),
}
}
Err(last_err.unwrap_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"could not resolve to any addresses",
)
}))
}
fn bind_addr(addr: SocketAddr) -> io::Result<UdpSocket> {
let sys = mio::net::UdpSocket::bind(&addr)?;
UdpSocket::new(sys)
}
fn new(socket: mio::net::UdpSocket) -> io::Result<UdpSocket> {
let io = PollEvented::new(socket)?;
Ok(UdpSocket { io })
}
/// Creates a new `UdpSocket` from the previously bound socket provided.
///
/// The socket given will be registered with the event loop that `handle`
/// is associated with. This function requires that `socket` has previously
/// been bound to an address to work correctly.
///
/// This can be used in conjunction with net2's `UdpBuilder` interface to
/// configure a socket before it's handed off, such as setting options like
/// `reuse_address` or binding to multiple addresses.
pub fn from_std(socket: net::UdpSocket) -> io::Result<UdpSocket> {
let io = mio::net::UdpSocket::from_socket(socket)?;
let io = PollEvented::new(io)?;
Ok(UdpSocket { io })
}
/// Split the `UdpSocket` into a receive half and a send half. The two parts
/// can be used to receive and send datagrams concurrently, even from two
/// different tasks.
///
/// See the module level documenation of [`split`](super::split) for more
/// details.
pub fn split(self) -> (UdpSocketRecvHalf, UdpSocketSendHalf) {
split(self)
}
/// Returns the local address that this socket is bound to.
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.io.get_ref().local_addr()
}
/// Connects the UDP socket setting the default destination for send() and
/// limiting packets that are read via recv from the address specified in
/// `addr`.
pub async fn connect<A: ToSocketAddrs>(&self, addr: A) -> io::Result<()> {
let addrs = addr.to_socket_addrs().await?;
let mut last_err = None;
for addr in addrs {
match self.io.get_ref().connect(addr) {
Ok(_) => return Ok(()),
Err(e) => last_err = Some(e),
}
}
Err(last_err.unwrap_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"could not resolve to any addresses",
)
}))
}
/// Returns a future that sends data on the socket to the remote address to which it is connected.
/// On success, the future will resolve to the number of bytes written.
///
/// The [`connect`] method will connect this socket to a remote address. The future
/// will resolve to an error if the socket is not connected.
///
/// [`connect`]: #method.connect
pub async fn send(&mut self, buf: &[u8]) -> io::Result<usize> {
poll_fn(|cx| self.poll_send(cx, buf)).await
}
// Poll IO functions that takes `&self` are provided for the split API.
//
// They are not public because (taken from the doc of `PollEvented`):
//
// While `PollEvented` is `Sync` (if the underlying I/O type is `Sync`), the
// caller must ensure that there are at most two tasks that use a
// `PollEvented` instance concurrently. One for reading and one for writing.
// While violating this requirement is "safe" from a Rust memory model point
// of view, it will result in unexpected behavior in the form of lost
// notifications and tasks hanging.
#[doc(hidden)]
pub fn poll_send(&self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
ready!(self.io.poll_write_ready(cx))?;
match self.io.get_ref().send(buf) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_write_ready(cx)?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
/// Returns a future that receives a single datagram message on the socket from
/// the remote address to which it is connected. On success, the future will resolve
/// to the number of bytes read.
///
/// The function must be called with valid byte array `buf` of sufficient size to
/// hold the message bytes. If a message is too long to fit in the supplied buffer,
/// excess bytes may be discarded.
///
/// The [`connect`] method will connect this socket to a remote address. The future
/// will fail if the socket is not connected.
///
/// [`connect`]: #method.connect
pub async fn recv(&mut self, buf: &mut [u8]) -> io::Result<usize> {
poll_fn(|cx| self.poll_recv(cx, buf)).await
}
#[doc(hidden)]
pub fn poll_recv(&self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
match self.io.get_ref().recv(buf) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
/// Returns a future that sends data on the socket to the given address.
/// On success, the future will resolve to the number of bytes written.
///
/// The future will resolve to an error if the IP version of the socket does
/// not match that of `target`.
pub async fn send_to<A: ToSocketAddrs>(&mut self, buf: &[u8], target: A) -> io::Result<usize> {
let mut addrs = target.to_socket_addrs().await?;
match addrs.next() {
Some(target) => poll_fn(|cx| self.poll_send_to(cx, buf, &target)).await,
None => Err(io::Error::new(
io::ErrorKind::InvalidInput,
"no addresses to send data to",
)),
}
}
// TODO: Public or not?
#[doc(hidden)]
pub fn poll_send_to(
&self,
cx: &mut Context<'_>,
buf: &[u8],
target: &SocketAddr,
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_write_ready(cx))?;
match self.io.get_ref().send_to(buf, target) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_write_ready(cx)?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
/// Returns a future that receives a single datagram on the socket. On success,
/// the future resolves to the number of bytes read and the origin.
///
/// The function must be called with valid byte array `buf` of sufficient size
/// to hold the message bytes. If a message is too long to fit in the supplied
/// buffer, excess bytes may be discarded.
pub async fn recv_from(&mut self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
poll_fn(|cx| self.poll_recv_from(cx, buf)).await
}
#[doc(hidden)]
pub fn poll_recv_from(
&self,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<Result<(usize, SocketAddr), io::Error>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
match self.io.get_ref().recv_from(buf) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
/// Gets the value of the `SO_BROADCAST` option for this socket.
///
/// For more information about this option, see [`set_broadcast`].
///
/// [`set_broadcast`]: #method.set_broadcast
pub fn broadcast(&self) -> io::Result<bool> {
self.io.get_ref().broadcast()
}
/// Sets the value of the `SO_BROADCAST` option for this socket.
///
/// When enabled, this socket is allowed to send packets to a broadcast
/// address.
pub fn set_broadcast(&self, on: bool) -> io::Result<()> {
self.io.get_ref().set_broadcast(on)
}
/// Gets the value of the `IP_MULTICAST_LOOP` option for this socket.
///
/// For more information about this option, see [`set_multicast_loop_v4`].
///
/// [`set_multicast_loop_v4`]: #method.set_multicast_loop_v4
pub fn multicast_loop_v4(&self) -> io::Result<bool> {
self.io.get_ref().multicast_loop_v4()
}
/// Sets the value of the `IP_MULTICAST_LOOP` option for this socket.
///
/// If enabled, multicast packets will be looped back to the local socket.
///
/// # Note
///
/// This may not have any affect on IPv6 sockets.
pub fn set_multicast_loop_v4(&self, on: bool) -> io::Result<()> {
self.io.get_ref().set_multicast_loop_v4(on)
}
/// Gets the value of the `IP_MULTICAST_TTL` option for this socket.
///
/// For more information about this option, see [`set_multicast_ttl_v4`].
///
/// [`set_multicast_ttl_v4`]: #method.set_multicast_ttl_v4
pub fn multicast_ttl_v4(&self) -> io::Result<u32> {
self.io.get_ref().multicast_ttl_v4()
}
/// Sets the value of the `IP_MULTICAST_TTL` option for this socket.
///
/// Indicates the time-to-live value of outgoing multicast packets for
/// this socket. The default value is 1 which means that multicast packets
/// don't leave the local network unless explicitly requested.
///
/// # Note
///
/// This may not have any affect on IPv6 sockets.
pub fn set_multicast_ttl_v4(&self, ttl: u32) -> io::Result<()> {
self.io.get_ref().set_multicast_ttl_v4(ttl)
}
/// Gets the value of the `IPV6_MULTICAST_LOOP` option for this socket.
///
/// For more information about this option, see [`set_multicast_loop_v6`].
///
/// [`set_multicast_loop_v6`]: #method.set_multicast_loop_v6
pub fn multicast_loop_v6(&self) -> io::Result<bool> {
self.io.get_ref().multicast_loop_v6()
}
/// Sets the value of the `IPV6_MULTICAST_LOOP` option for this socket.
///
/// Controls whether this socket sees the multicast packets it sends itself.
///
/// # Note
///
/// This may not have any affect on IPv4 sockets.
pub fn set_multicast_loop_v6(&self, on: bool) -> io::Result<()> {
self.io.get_ref().set_multicast_loop_v6(on)
}
/// Gets the value of the `IP_TTL` option for this socket.
///
/// For more information about this option, see [`set_ttl`].
///
/// [`set_ttl`]: #method.set_ttl
pub fn ttl(&self) -> io::Result<u32> {
self.io.get_ref().ttl()
}
/// Sets the value for the `IP_TTL` option on this socket.
///
/// This value sets the time-to-live field that is used in every packet sent
/// from this socket.
pub fn set_ttl(&self, ttl: u32) -> io::Result<()> {
self.io.get_ref().set_ttl(ttl)
}
/// Executes an operation of the `IP_ADD_MEMBERSHIP` type.
///
/// This function specifies a new multicast group for this socket to join.
/// The address must be a valid multicast address, and `interface` is the
/// address of the local interface with which the system should join the
/// multicast group. If it's equal to `INADDR_ANY` then an appropriate
/// interface is chosen by the system.
pub fn join_multicast_v4(&self, multiaddr: Ipv4Addr, interface: Ipv4Addr) -> io::Result<()> {
self.io.get_ref().join_multicast_v4(&multiaddr, &interface)
}
/// Executes an operation of the `IPV6_ADD_MEMBERSHIP` type.
///
/// This function specifies a new multicast group for this socket to join.
/// The address must be a valid multicast address, and `interface` is the
/// index of the interface to join/leave (or 0 to indicate any interface).
pub fn join_multicast_v6(&self, multiaddr: &Ipv6Addr, interface: u32) -> io::Result<()> {
self.io.get_ref().join_multicast_v6(multiaddr, interface)
}
/// Executes an operation of the `IP_DROP_MEMBERSHIP` type.
///
/// For more information about this option, see [`join_multicast_v4`].
///
/// [`join_multicast_v4`]: #method.join_multicast_v4
pub fn leave_multicast_v4(&self, multiaddr: Ipv4Addr, interface: Ipv4Addr) -> io::Result<()> {
self.io.get_ref().leave_multicast_v4(&multiaddr, &interface)
}
/// Executes an operation of the `IPV6_DROP_MEMBERSHIP` type.
///
/// For more information about this option, see [`join_multicast_v6`].
///
/// [`join_multicast_v6`]: #method.join_multicast_v6
pub fn leave_multicast_v6(&self, multiaddr: &Ipv6Addr, interface: u32) -> io::Result<()> {
self.io.get_ref().leave_multicast_v6(multiaddr, interface)
}
}
impl TryFrom<UdpSocket> for mio::net::UdpSocket {
type Error = io::Error;
/// Consumes value, returning the mio I/O object.
///
/// See [`PollEvented::into_inner`] for more details about
/// resource deregistration that happens during the call.
///
/// [`PollEvented::into_inner`]: crate::util::PollEvented::into_inner
fn try_from(value: UdpSocket) -> Result<Self, Self::Error> {
value.io.into_inner()
}
}
impl TryFrom<net::UdpSocket> for UdpSocket {
type Error = io::Error;
/// Consumes stream, returning the tokio I/O object.
///
/// This is equivalent to
/// [`UdpSocket::from_std(stream)`](UdpSocket::from_std).
fn try_from(stream: net::UdpSocket) -> Result<Self, Self::Error> {
Self::from_std(stream)
}
}
impl fmt::Debug for UdpSocket {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.io.get_ref().fmt(f)
}
}
#[cfg(all(unix))]
mod sys {
use super::UdpSocket;
use std::os::unix::prelude::*;
impl AsRawFd for UdpSocket {
fn as_raw_fd(&self) -> RawFd {
self.io.get_ref().as_raw_fd()
}
}
}
#[cfg(windows)]
mod sys {
// TODO: let's land these upstream with mio and then we can add them here.
//
// use std::os::windows::prelude::*;
// use super::UdpSocket;
//
// impl AsRawHandle for UdpSocket {
// fn as_raw_handle(&self) -> RawHandle {
// self.io.get_ref().as_raw_handle()
// }
// }
}
tokio/src/net/udp/split.rs
+148
View File
@@ -0,0 +1,148 @@
//! [`UdpSocket`](../struct.UdpSocket.html) split support.
//!
//! The [`split`](../struct.UdpSocket.html#method.split) method splits a
//! `UdpSocket` into a receive half and a send half, which can be used to
//! receive and send datagrams concurrently, even from two different tasks.
//!
//! The halves provide access to the underlying socket, implementing
//! `AsRef<UdpSocket>`. This allows you to call `UdpSocket` methods that takes
//! `&self`, e.g., to get local address, to get and set socket options, to join
//! or leave multicast groups, etc.
//!
//! The halves can be reunited to the original socket with their `reunite`
//! methods.
use super::UdpSocket;
use futures_util::future::poll_fn;
use std::error::Error;
use std::fmt;
use std::io;
use std::net::SocketAddr;
use std::sync::Arc;
/// The send half after [`split`](super::UdpSocket::split).
///
/// Use [`send_to`](#method.send_to) or [`send`](#method.send) to send
/// datagrams.
#[derive(Debug)]
pub struct UdpSocketSendHalf(Arc<UdpSocket>);
/// The recv half after [`split`](super::UdpSocket::split).
///
/// Use [`recv_from`](#method.recv_from) or [`recv`](#method.recv) to receive
/// datagrams.
#[derive(Debug)]
pub struct UdpSocketRecvHalf(Arc<UdpSocket>);
pub(crate) fn split(socket: UdpSocket) -> (UdpSocketRecvHalf, UdpSocketSendHalf) {
let shared = Arc::new(socket);
let send = shared.clone();
let recv = shared;
(UdpSocketRecvHalf(recv), UdpSocketSendHalf(send))
}
/// Error indicating two halves were not from the same socket, and thus could
/// not be `reunite`d.
#[derive(Debug)]
pub struct ReuniteError(pub UdpSocketSendHalf, pub UdpSocketRecvHalf);
impl fmt::Display for ReuniteError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"tried to reunite halves that are not from the same socket"
)
}
}
impl Error for ReuniteError {}
fn reunite(s: UdpSocketSendHalf, r: UdpSocketRecvHalf) -> Result<UdpSocket, ReuniteError> {
if Arc::ptr_eq(&s.0, &r.0) {
drop(r);
// Only two instances of the `Arc` are ever created, one for the
// receiver and one for the sender, and those `Arc`s are never exposed
// externally. And so when we drop one here, the other one must be the
// only remaining one.
Ok(Arc::try_unwrap(s.0).expect("udp: try_unwrap failed in reunite"))
} else {
Err(ReuniteError(s, r))
}
}
impl UdpSocketRecvHalf {
/// Attempts to put the two "halves" of a `UdpSocket` back together and
/// recover the original socket. Succeeds only if the two "halves"
/// originated from the same call to `UdpSocket::split`.
pub fn reunite(self, other: UdpSocketSendHalf) -> Result<UdpSocket, ReuniteError> {
reunite(other, self)
}
/// Returns a future that receives a single datagram on the socket. On success,
/// the future resolves to the number of bytes read and the origin.
///
/// The function must be called with valid byte array `buf` of sufficient size
/// to hold the message bytes. If a message is too long to fit in the supplied
/// buffer, excess bytes may be discarded.
pub async fn recv_from(&mut self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
poll_fn(|cx| self.0.poll_recv_from(cx, buf)).await
}
/// Returns a future that receives a single datagram message on the socket from
/// the remote address to which it is connected. On success, the future will resolve
/// to the number of bytes read.
///
/// The function must be called with valid byte array `buf` of sufficient size to
/// hold the message bytes. If a message is too long to fit in the supplied buffer,
/// excess bytes may be discarded.
///
/// The [`connect`] method will connect this socket to a remote address. The future
/// will fail if the socket is not connected.
///
/// [`connect`]: super::UdpSocket::connect
pub async fn recv(&mut self, buf: &mut [u8]) -> io::Result<usize> {
poll_fn(|cx| self.0.poll_recv(cx, buf)).await
}
}
impl UdpSocketSendHalf {
/// Attempts to put the two "halves" of a `UdpSocket` back together and
/// recover the original socket. Succeeds only if the two "halves"
/// originated from the same call to `UdpSocket::split`.
pub fn reunite(self, other: UdpSocketRecvHalf) -> Result<UdpSocket, ReuniteError> {
reunite(self, other)
}
/// Returns a future that sends data on the socket to the given address.
/// On success, the future will resolve to the number of bytes written.
///
/// The future will resolve to an error if the IP version of the socket does
/// not match that of `target`.
pub async fn send_to(&mut self, buf: &[u8], target: &SocketAddr) -> io::Result<usize> {
poll_fn(|cx| self.0.poll_send_to(cx, buf, target)).await
}
/// Returns a future that sends data on the socket to the remote address to which it is connected.
/// On success, the future will resolve to the number of bytes written.
///
/// The [`connect`] method will connect this socket to a remote address. The future
/// will resolve to an error if the socket is not connected.
///
/// [`connect`]: super::UdpSocket::connect
pub async fn send(&mut self, buf: &[u8]) -> io::Result<usize> {
poll_fn(|cx| self.0.poll_send(cx, buf)).await
}
}
impl AsRef<UdpSocket> for UdpSocketSendHalf {
fn as_ref(&self) -> &UdpSocket {
&self.0
}
}
impl AsRef<UdpSocket> for UdpSocketRecvHalf {
fn as_ref(&self) -> &UdpSocket {
&self.0
}
}
tokio/src/net/unix/datagram.rs
+233
View File
@@ -0,0 +1,233 @@
use crate::net::util::PollEvented;
use futures_core::ready;
use futures_util::future::poll_fn;
use std::convert::TryFrom;
use std::fmt;
use std::io;
use std::net::Shutdown;
use std::os::unix::io::{AsRawFd, RawFd};
use std::os::unix::net::{self, SocketAddr};
use std::path::Path;
use std::task::{Context, Poll};
/// An I/O object representing a Unix datagram socket.
pub struct UnixDatagram {
io: PollEvented<mio_uds::UnixDatagram>,
}
impl UnixDatagram {
/// Creates a new `UnixDatagram` bound to the specified path.
pub fn bind<P>(path: P) -> io::Result<UnixDatagram>
where
P: AsRef<Path>,
{
let socket = mio_uds::UnixDatagram::bind(path)?;
UnixDatagram::new(socket)
}
/// Creates an unnamed pair of connected sockets.
///
/// This function will create a pair of interconnected Unix sockets for
/// communicating back and forth between one another. Each socket will
/// be associated with the default event loop's handle.
pub fn pair() -> io::Result<(UnixDatagram, UnixDatagram)> {
let (a, b) = mio_uds::UnixDatagram::pair()?;
let a = UnixDatagram::new(a)?;
let b = UnixDatagram::new(b)?;
Ok((a, b))
}
/// Consumes a `UnixDatagram` in the standard library and returns a
/// nonblocking `UnixDatagram` from this crate.
///
/// The returned datagram will be associated with the given event loop
/// specified by `handle` and is ready to perform I/O.
pub fn from_std(datagram: net::UnixDatagram) -> io::Result<UnixDatagram> {
let socket = mio_uds::UnixDatagram::from_datagram(datagram)?;
let io = PollEvented::new(socket)?;
Ok(UnixDatagram { io })
}
fn new(socket: mio_uds::UnixDatagram) -> io::Result<UnixDatagram> {
let io = PollEvented::new(socket)?;
Ok(UnixDatagram { io })
}
/// Creates a new `UnixDatagram` which is not bound to any address.
pub fn unbound() -> io::Result<UnixDatagram> {
let socket = mio_uds::UnixDatagram::unbound()?;
UnixDatagram::new(socket)
}
/// Connects the socket to the specified address.
///
/// The `send` method may be used to send data to the specified address.
/// `recv` and `recv_from` will only receive data from that address.
pub fn connect<P: AsRef<Path>>(&self, path: P) -> io::Result<()> {
self.io.get_ref().connect(path)
}
/// Sends data on the socket to the socket's peer.
pub async fn send(&mut self, buf: &[u8]) -> io::Result<usize> {
poll_fn(|cx| self.poll_send_priv(cx, buf)).await
}
// Poll IO functions that takes `&self` are provided for the split API.
//
// They are not public because (taken from the doc of `PollEvented`):
//
// While `PollEvented` is `Sync` (if the underlying I/O type is `Sync`), the
// caller must ensure that there are at most two tasks that use a
// `PollEvented` instance concurrently. One for reading and one for writing.
// While violating this requirement is "safe" from a Rust memory model point
// of view, it will result in unexpected behavior in the form of lost
// notifications and tasks hanging.
pub(crate) fn poll_send_priv(
&self,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_write_ready(cx))?;
match self.io.get_ref().send(buf) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_write_ready(cx)?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
/// Receives data from the socket.
pub async fn recv(&mut self, buf: &mut [u8]) -> io::Result<usize> {
poll_fn(|cx| self.poll_recv_priv(cx, buf)).await
}
pub(crate) fn poll_recv_priv(
&self,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
match self.io.get_ref().recv(buf) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
/// Sends data on the socket to the specified address.
pub async fn send_to<P>(&mut self, buf: &[u8], target: P) -> io::Result<usize>
where
P: AsRef<Path> + Unpin,
{
poll_fn(|cx| self.poll_send_to_priv(cx, buf, target.as_ref())).await
}
pub(crate) fn poll_send_to_priv(
&self,
cx: &mut Context<'_>,
buf: &[u8],
target: &Path,
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_write_ready(cx))?;
match self.io.get_ref().send_to(buf, target) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_write_ready(cx)?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
/// Receives data from the socket.
pub async fn recv_from(&mut self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
poll_fn(|cx| self.poll_recv_from_priv(cx, buf)).await
}
pub(crate) fn poll_recv_from_priv(
&self,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<Result<(usize, SocketAddr), io::Error>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
match self.io.get_ref().recv_from(buf) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
/// Returns the local address that this socket is bound to.
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.io.get_ref().local_addr()
}
/// Returns the address of this socket's peer.
///
/// The `connect` method will connect the socket to a peer.
pub fn peer_addr(&self) -> io::Result<SocketAddr> {
self.io.get_ref().peer_addr()
}
/// Returns the value of the `SO_ERROR` option.
pub fn take_error(&self) -> io::Result<Option<io::Error>> {
self.io.get_ref().take_error()
}
/// Shut down the read, write, or both halves of this connection.
///
/// This function will cause all pending and future I/O calls on the
/// specified portions to immediately return with an appropriate value
/// (see the documentation of `Shutdown`).
pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
self.io.get_ref().shutdown(how)
}
}
impl TryFrom<UnixDatagram> for mio_uds::UnixDatagram {
type Error = io::Error;
/// Consumes value, returning the mio I/O object.
///
/// See [`PollEvented::into_inner`] for more details about
/// resource deregistration that happens during the call.
///
/// [`PollEvented::into_inner`]: crate::util::PollEvented::into_inner
fn try_from(value: UnixDatagram) -> Result<Self, Self::Error> {
value.io.into_inner()
}
}
impl TryFrom<net::UnixDatagram> for UnixDatagram {
type Error = io::Error;
/// Consumes stream, returning the tokio I/O object.
///
/// This is equivalent to
/// [`UnixDatagram::from_std(stream)`](UnixDatagram::from_std).
fn try_from(stream: net::UnixDatagram) -> Result<Self, Self::Error> {
Self::from_std(stream)
}
}
impl fmt::Debug for UnixDatagram {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.io.get_ref().fmt(f)
}
}
impl AsRawFd for UnixDatagram {
fn as_raw_fd(&self) -> RawFd {
self.io.get_ref().as_raw_fd()
}
}
tokio/src/net/unix/incoming.rs
+31
View File
@@ -0,0 +1,31 @@
#![cfg(feature = "async-traits")]
use super::{UnixListener, UnixStream};
use futures_core::ready;
use futures_core::stream::Stream;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
/// Stream of listeners
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct Incoming {
inner: UnixListener,
}
impl Incoming {
pub(crate) fn new(listener: UnixListener) -> Incoming {
Incoming { inner: listener }
}
}
impl Stream for Incoming {
type Item = io::Result<UnixStream>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let (socket, _) = ready!(Pin::new(&mut self.inner).poll_accept(cx))?;
Poll::Ready(Some(Ok(socket)))
}
}
tokio/src/net/unix/listener.rs
+135
View File
@@ -0,0 +1,135 @@
use crate::net::unix::UnixStream;
use crate::net::util::PollEvented;
use futures_core::ready;
use futures_util::future::poll_fn;
use mio::Ready;
use mio_uds;
use std::convert::TryFrom;
use std::fmt;
use std::io;
use std::os::unix::io::{AsRawFd, RawFd};
use std::os::unix::net::{self, SocketAddr};
use std::path::Path;
use std::task::{Context, Poll};
/// A Unix socket which can accept connections from other Unix sockets.
pub struct UnixListener {
io: PollEvented<mio_uds::UnixListener>,
}
impl UnixListener {
/// Creates a new `UnixListener` bound to the specified path.
pub fn bind<P>(path: P) -> io::Result<UnixListener>
where
P: AsRef<Path>,
{
let listener = mio_uds::UnixListener::bind(path)?;
let io = PollEvented::new(listener)?;
Ok(UnixListener { io })
}
/// Consumes a `UnixListener` in the standard library and returns a
/// nonblocking `UnixListener` from this crate.
///
/// The returned listener will be associated with the given event loop
/// specified by `handle` and is ready to perform I/O.
pub fn from_std(listener: net::UnixListener) -> io::Result<UnixListener> {
let listener = mio_uds::UnixListener::from_listener(listener)?;
let io = PollEvented::new(listener)?;
Ok(UnixListener { io })
}
/// Returns the local socket address of this listener.
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.io.get_ref().local_addr()
}
/// Returns the value of the `SO_ERROR` option.
pub fn take_error(&self) -> io::Result<Option<io::Error>> {
self.io.get_ref().take_error()
}
/// Accepts a new incoming connection to this listener.
pub async fn accept(&mut self) -> io::Result<(UnixStream, SocketAddr)> {
poll_fn(|cx| self.poll_accept(cx)).await
}
pub(crate) fn poll_accept(
&mut self,
cx: &mut Context<'_>,
) -> Poll<io::Result<(UnixStream, SocketAddr)>> {
let (io, addr) = ready!(self.poll_accept_std(cx))?;
let io = mio_uds::UnixStream::from_stream(io)?;
Ok((UnixStream::new(io)?, addr)).into()
}
fn poll_accept_std(
&mut self,
cx: &mut Context<'_>,
) -> Poll<io::Result<(net::UnixStream, SocketAddr)>> {
ready!(self.io.poll_read_ready(cx, Ready::readable()))?;
match self.io.get_ref().accept_std() {
Ok(None) => {
self.io.clear_read_ready(cx, Ready::readable())?;
Poll::Pending
}
Ok(Some((sock, addr))) => Ok((sock, addr)).into(),
Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, Ready::readable())?;
Poll::Pending
}
Err(err) => Err(err).into(),
}
}
/// Consumes this listener, returning a stream of the sockets this listener
/// accepts.
///
/// This method returns an implementation of the `Stream` trait which
/// resolves to the sockets the are accepted on this listener.
#[cfg(feature = "async-traits")]
pub fn incoming(self) -> super::Incoming {
super::Incoming::new(self)
}
}
impl TryFrom<UnixListener> for mio_uds::UnixListener {
type Error = io::Error;
/// Consumes value, returning the mio I/O object.
///
/// See [`PollEvented::into_inner`] for more details about
/// resource deregistration that happens during the call.
///
/// [`PollEvented::into_inner`]: crate::util::PollEvented::into_inner
fn try_from(value: UnixListener) -> Result<Self, Self::Error> {
value.io.into_inner()
}
}
impl TryFrom<net::UnixListener> for UnixListener {
type Error = io::Error;
/// Consumes stream, returning the tokio I/O object.
///
/// This is equivalent to
/// [`UnixListener::from_std(stream)`](UnixListener::from_std).
fn try_from(stream: net::UnixListener) -> io::Result<Self> {
Self::from_std(stream)
}
}
impl fmt::Debug for UnixListener {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.io.get_ref().fmt(f)
}
}
impl AsRawFd for UnixListener {
fn as_raw_fd(&self) -> RawFd {
self.io.get_ref().as_raw_fd()
}
}
tokio/src/net/unix/mod.rs
+21
View File
@@ -0,0 +1,21 @@
//! Unix Domain Sockets for Tokio.
//!
//! This crate provides APIs for using Unix Domain Sockets with Tokio.
mod datagram;
pub use self::datagram::UnixDatagram;
mod incoming;
#[cfg(feature = "async-traits")]
pub use self::incoming::Incoming;
mod listener;
pub use self::listener::UnixListener;
pub mod split;
mod stream;
pub use self::stream::UnixStream;
mod ucred;
pub use self::ucred::UCred;
tokio/src/net/unix/split.rs
+91
View File
@@ -0,0 +1,91 @@
//! `UnixStream` split support.
//!
//! A `UnixStream` can be split into a read half and a write half with
//! `UnixStream::split`. The read half implements `AsyncRead` while the write
//! half implements `AsyncWrite`.
//!
//! Compared to the generic split of `AsyncRead + AsyncWrite`, this specialized
//! split has no associated overhead and enforces all invariants at the type
//! level.
use super::UnixStream;
use tokio_io::{AsyncRead, AsyncWrite};
use bytes::{Buf, BufMut};
use std::io;
use std::net::Shutdown;
use std::pin::Pin;
use std::task::{Context, Poll};
/// Read half of a `UnixStream`.
#[derive(Debug)]
pub struct ReadHalf<'a>(&'a UnixStream);
/// Write half of a `UnixStream`.
#[derive(Debug)]
pub struct WriteHalf<'a>(&'a UnixStream);
pub(crate) fn split(stream: &mut UnixStream) -> (ReadHalf<'_>, WriteHalf<'_>) {
(ReadHalf(stream), WriteHalf(stream))
}
impl AsyncRead for ReadHalf<'_> {
unsafe fn prepare_uninitialized_buffer(&self, _: &mut [u8]) -> bool {
false
}
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
self.0.poll_read_priv(cx, buf)
}
fn poll_read_buf<B: BufMut>(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
self.0.poll_read_buf_priv(cx, buf)
}
}
impl AsyncWrite for WriteHalf<'_> {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.0.poll_write_priv(cx, buf)
}
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
self.0.shutdown(Shutdown::Write).into()
}
fn poll_write_buf<B: Buf>(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
self.0.poll_write_buf_priv(cx, buf)
}
}
impl AsRef<UnixStream> for ReadHalf<'_> {
fn as_ref(&self) -> &UnixStream {
self.0
}
}
impl AsRef<UnixStream> for WriteHalf<'_> {
fn as_ref(&self) -> &UnixStream {
self.0
}
}
tokio/src/net/unix/stream.rs
+339
View File
@@ -0,0 +1,339 @@
use crate::net::unix::split::{split, ReadHalf, WriteHalf};
use crate::net::unix::ucred::{self, UCred};
use crate::net::util::PollEvented;
use tokio_io::{AsyncRead, AsyncWrite};
use bytes::{Buf, BufMut};
use futures_core::ready;
use futures_util::future::poll_fn;
use iovec::IoVec;
use std::convert::TryFrom;
use std::fmt;
use std::io::{self, Read, Write};
use std::net::Shutdown;
use std::os::unix::io::{AsRawFd, RawFd};
use std::os::unix::net::{self, SocketAddr};
use std::path::Path;
use std::pin::Pin;
use std::task::{Context, Poll};
/// A structure representing a connected Unix socket.
///
/// This socket can be connected directly with `UnixStream::connect` or accepted
/// from a listener with `UnixListener::incoming`. Additionally, a pair of
/// anonymous Unix sockets can be created with `UnixStream::pair`.
pub struct UnixStream {
io: PollEvented<mio_uds::UnixStream>,
}
impl UnixStream {
/// Connects to the socket named by `path`.
///
/// This function will create a new Unix socket and connect to the path
/// specified, associating the returned stream with the default event loop's
/// handle.
pub async fn connect<P>(path: P) -> io::Result<UnixStream>
where
P: AsRef<Path>,
{
let stream = mio_uds::UnixStream::connect(path)?;
let stream = UnixStream::new(stream)?;
poll_fn(|cx| stream.io.poll_write_ready(cx)).await?;
Ok(stream)
}
/// Consumes a `UnixStream` in the standard library and returns a
/// nonblocking `UnixStream` from this crate.
///
/// The returned stream will be associated with the given event loop
/// specified by `handle` and is ready to perform I/O.
pub fn from_std(stream: net::UnixStream) -> io::Result<UnixStream> {
let stream = mio_uds::UnixStream::from_stream(stream)?;
let io = PollEvented::new(stream)?;
Ok(UnixStream { io })
}
/// Creates an unnamed pair of connected sockets.
///
/// This function will create a pair of interconnected Unix sockets for
/// communicating back and forth between one another. Each socket will
/// be associated with the default event loop's handle.
pub fn pair() -> io::Result<(UnixStream, UnixStream)> {
let (a, b) = mio_uds::UnixStream::pair()?;
let a = UnixStream::new(a)?;
let b = UnixStream::new(b)?;
Ok((a, b))
}
pub(crate) fn new(stream: mio_uds::UnixStream) -> io::Result<UnixStream> {
let io = PollEvented::new(stream)?;
Ok(UnixStream { io })
}
/// Returns the socket address of the local half of this connection.
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.io.get_ref().local_addr()
}
/// Returns the socket address of the remote half of this connection.
pub fn peer_addr(&self) -> io::Result<SocketAddr> {
self.io.get_ref().peer_addr()
}
/// Returns effective credentials of the process which called `connect` or `pair`.
pub fn peer_cred(&self) -> io::Result<UCred> {
ucred::get_peer_cred(self)
}
/// Returns the value of the `SO_ERROR` option.
pub fn take_error(&self) -> io::Result<Option<io::Error>> {
self.io.get_ref().take_error()
}
/// Shuts down the read, write, or both halves of this connection.
///
/// This function will cause all pending and future I/O calls on the
/// specified portions to immediately return with an appropriate value
/// (see the documentation of `Shutdown`).
pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
self.io.get_ref().shutdown(how)
}
/// Split a `UnixStream` into a read half and a write half, which can be used
/// to read and write the stream concurrently.
///
/// See the module level documenation of [`split`](super::split) for more
/// details.
pub fn split(&mut self) -> (ReadHalf<'_>, WriteHalf<'_>) {
split(self)
}
}
impl TryFrom<UnixStream> for mio_uds::UnixStream {
type Error = io::Error;
/// Consumes value, returning the mio I/O object.
///
/// See [`PollEvented::into_inner`] for more details about
/// resource deregistration that happens during the call.
///
/// [`PollEvented::into_inner`]: crate::util::PollEvented::into_inner
fn try_from(value: UnixStream) -> Result<Self, Self::Error> {
value.io.into_inner()
}
}
impl TryFrom<net::UnixStream> for UnixStream {
type Error = io::Error;
/// Consumes stream, returning the tokio I/O object.
///
/// This is equivalent to
/// [`UnixStream::from_std(stream)`](UnixStream::from_std).
fn try_from(stream: net::UnixStream) -> io::Result<Self> {
Self::from_std(stream)
}
}
impl AsyncRead for UnixStream {
unsafe fn prepare_uninitialized_buffer(&self, _: &mut [u8]) -> bool {
false
}
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
self.poll_read_priv(cx, buf)
}
fn poll_read_buf<B: BufMut>(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
self.poll_read_buf_priv(cx, buf)
}
}
impl AsyncWrite for UnixStream {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.poll_write_priv(cx, buf)
}
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_write_buf<B: Buf>(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
self.poll_write_buf_priv(cx, buf)
}
}
impl UnixStream {
// == Poll IO functions that takes `&self` ==
//
// They are not public because (taken from the doc of `PollEvented`):
//
// While `PollEvented` is `Sync` (if the underlying I/O type is `Sync`), the
// caller must ensure that there are at most two tasks that use a
// `PollEvented` instance concurrently. One for reading and one for writing.
// While violating this requirement is "safe" from a Rust memory model point
// of view, it will result in unexpected behavior in the form of lost
// notifications and tasks hanging.
pub(crate) fn poll_read_priv(
&self,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
match self.io.get_ref().read(buf) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
pub(crate) fn poll_read_buf_priv<B: BufMut>(
&self,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
let r = unsafe {
// The `IoVec` type can't have a 0-length size, so we create a bunch
// of dummy versions on the stack with 1 length which we'll quickly
// overwrite.
let b1: &mut [u8] = &mut [0];
let b2: &mut [u8] = &mut [0];
let b3: &mut [u8] = &mut [0];
let b4: &mut [u8] = &mut [0];
let b5: &mut [u8] = &mut [0];
let b6: &mut [u8] = &mut [0];
let b7: &mut [u8] = &mut [0];
let b8: &mut [u8] = &mut [0];
let b9: &mut [u8] = &mut [0];
let b10: &mut [u8] = &mut [0];
let b11: &mut [u8] = &mut [0];
let b12: &mut [u8] = &mut [0];
let b13: &mut [u8] = &mut [0];
let b14: &mut [u8] = &mut [0];
let b15: &mut [u8] = &mut [0];
let b16: &mut [u8] = &mut [0];
let mut bufs: [&mut IoVec; 16] = [
b1.into(),
b2.into(),
b3.into(),
b4.into(),
b5.into(),
b6.into(),
b7.into(),
b8.into(),
b9.into(),
b10.into(),
b11.into(),
b12.into(),
b13.into(),
b14.into(),
b15.into(),
b16.into(),
];
let n = buf.bytes_vec_mut(&mut bufs);
self.io.get_ref().read_bufs(&mut bufs[..n])
};
match r {
Ok(n) => {
unsafe {
buf.advance_mut(n);
}
Poll::Ready(Ok(n))
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
Err(e) => Poll::Ready(Err(e)),
}
}
pub(crate) fn poll_write_priv(
&self,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_write_ready(cx))?;
match self.io.get_ref().write(buf) {
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_write_ready(cx)?;
Poll::Pending
}
x => Poll::Ready(x),
}
}
pub(crate) fn poll_write_buf_priv<B: Buf>(
&self,
cx: &mut Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_write_ready(cx))?;
let r = {
// The `IoVec` type can't have a zero-length size, so create a dummy
// version from a 1-length slice which we'll overwrite with the
// `bytes_vec` method.
static DUMMY: &[u8] = &[0];
let iovec = <&IoVec>::from(DUMMY);
let mut bufs = [iovec; 64];
let n = buf.bytes_vec(&mut bufs);
self.io.get_ref().write_bufs(&bufs[..n])
};
match r {
Ok(n) => {
buf.advance(n);
Poll::Ready(Ok(n))
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_write_ready(cx)?;
Poll::Pending
}
Err(e) => Poll::Ready(Err(e)),
}
}
}
impl fmt::Debug for UnixStream {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.io.get_ref().fmt(f)
}
}
impl AsRawFd for UnixStream {
fn as_raw_fd(&self) -> RawFd {
self.io.get_ref().as_raw_fd()
}
}
tokio/src/net/unix/ucred.rs
+154
View File
@@ -0,0 +1,154 @@
use libc::{gid_t, uid_t};
/// Credentials of a process
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub struct UCred {
/// UID (user ID) of the process
pub uid: uid_t,
/// GID (group ID) of the process
pub gid: gid_t,
}
#[cfg(any(target_os = "linux", target_os = "android"))]
pub(crate) use self::impl_linux::get_peer_cred;
#[cfg(any(
target_os = "dragonfly",
target_os = "macos",
target_os = "ios",
target_os = "freebsd",
target_os = "netbsd",
target_os = "openbsd"
))]
pub(crate) use self::impl_macos::get_peer_cred;
#[cfg(any(target_os = "solaris"))]
pub(crate) use self::impl_solaris::get_peer_cred;
#[cfg(any(target_os = "linux", target_os = "android"))]
pub(crate) mod impl_linux {
use crate::net::unix::UnixStream;
use libc::{c_void, getsockopt, socklen_t, SOL_SOCKET, SO_PEERCRED};
use std::{io, mem};
use libc::ucred;
pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> {
use std::os::unix::io::AsRawFd;
unsafe {
let raw_fd = sock.as_raw_fd();
let mut ucred = ucred {
pid: 0,
uid: 0,
gid: 0,
};
let ucred_size = mem::size_of::<ucred>();
// These paranoid checks should be optimized-out
assert!(mem::size_of::<u32>() <= mem::size_of::<usize>());
assert!(ucred_size <= u32::max_value() as usize);
let mut ucred_size = ucred_size as socklen_t;
let ret = getsockopt(
raw_fd,
SOL_SOCKET,
SO_PEERCRED,
&mut ucred as *mut ucred as *mut c_void,
&mut ucred_size,
);
if ret == 0 && ucred_size as usize == mem::size_of::<ucred>() {
Ok(super::UCred {
uid: ucred.uid,
gid: ucred.gid,
})
} else {
Err(io::Error::last_os_error())
}
}
}
}
#[cfg(any(
target_os = "dragonfly",
target_os = "macos",
target_os = "ios",
target_os = "freebsd",
target_os = "netbsd",
target_os = "openbsd"
))]
pub(crate) mod impl_macos {
use crate::net::unix::UnixStream;
use libc::getpeereid;
use std::io;
use std::mem::MaybeUninit;
use std::os::unix::io::AsRawFd;
pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> {
unsafe {
let raw_fd = sock.as_raw_fd();
let mut uid = MaybeUninit::uninit();
let mut gid = MaybeUninit::uninit();
let ret = getpeereid(raw_fd, uid.as_mut_ptr(), gid.as_mut_ptr());
if ret == 0 {
Ok(super::UCred {
uid: uid.assume_init(),
gid: gid.assume_init(),
})
} else {
Err(io::Error::last_os_error())
}
}
}
}
#[cfg(any(target_os = "solaris"))]
pub(crate) mod impl_solaris {
use std::io;
use std::os::unix::io::AsRawFd;
use std::ptr;
use UnixStream;
#[allow(non_camel_case_types)]
enum ucred_t {}
extern "C" {
fn ucred_free(cred: *mut ucred_t);
fn ucred_geteuid(cred: *const ucred_t) -> super::uid_t;
fn ucred_getegid(cred: *const ucred_t) -> super::gid_t;
fn getpeerucred(
fd: ::std::os::raw::c_int,
cred: *mut *mut ucred_t,
) -> ::std::os::raw::c_int;
}
pub(crate) fn get_peer_cred(sock: &UnixStream) -> io::Result<super::UCred> {
unsafe {
let raw_fd = sock.as_raw_fd();
let mut cred = ptr::null_mut::<*mut ucred_t>() as *mut ucred_t;
let ret = getpeerucred(raw_fd, &mut cred);
if ret == 0 {
let uid = ucred_geteuid(cred);
let gid = ucred_getegid(cred);
ucred_free(cred);
Ok(super::UCred { uid, gid })
} else {
Err(io::Error::last_os_error())
}
}
}
}
tokio/src/net/util/mod.rs
+4
View File
@@ -0,0 +1,4 @@
//! Utilities for implementing networking types.
mod poll_evented;
pub use self::poll_evented::PollEvented;
tokio/src/net/util/poll_evented.rs
+415
View File
@@ -0,0 +1,415 @@
use crate::net::driver::{platform, Registration};
use tokio_io::{AsyncRead, AsyncWrite};
use futures_core::ready;
use mio;
use mio::event::Evented;
use std::fmt;
use std::io::{self, Read, Write};
use std::marker::Unpin;
use std::pin::Pin;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::Relaxed;
use std::task::{Context, Poll};
/// Associates an I/O resource that implements the [`std::io::Read`] and/or
/// [`std::io::Write`] traits with the reactor that drives it.
///
/// `PollEvented` uses [`Registration`] internally to take a type that
/// implements [`mio::Evented`] as well as [`std::io::Read`] and or
/// [`std::io::Write`] and associate it with a reactor that will drive it.
///
/// Once the [`mio::Evented`] type is wrapped by `PollEvented`, it can be
/// used from within the future's execution model. As such, the `PollEvented`
/// type provides [`AsyncRead`] and [`AsyncWrite`] implementations using the
/// underlying I/O resource as well as readiness events provided by the reactor.
///
/// **Note**: While `PollEvented` is `Sync` (if the underlying I/O type is
/// `Sync`), the caller must ensure that there are at most two tasks that use a
/// `PollEvented` instance concurrently. One for reading and one for writing.
/// While violating this requirement is "safe" from a Rust memory model point of
/// view, it will result in unexpected behavior in the form of lost
/// notifications and tasks hanging.
///
/// ## Readiness events
///
/// Besides just providing [`AsyncRead`] and [`AsyncWrite`] implementations,
/// this type also supports access to the underlying readiness event stream.
/// While similar in function to what [`Registration`] provides, the semantics
/// are a bit different.
///
/// Two functions are provided to access the readiness events:
/// [`poll_read_ready`] and [`poll_write_ready`]. These functions return the
/// current readiness state of the `PollEvented` instance. If
/// [`poll_read_ready`] indicates read readiness, immediately calling
/// [`poll_read_ready`] again will also indicate read readiness.
///
/// When the operation is attempted and is unable to succeed due to the I/O
/// resource not being ready, the caller must call [`clear_read_ready`] or
/// [`clear_write_ready`]. This clears the readiness state until a new readiness
/// event is received.
///
/// This allows the caller to implement additional functions. For example,
/// [`TcpListener`] implements poll_accept by using [`poll_read_ready`] and
/// [`clear_read_ready`].
///
/// ```rust
/// use tokio::net::util::PollEvented;
///
/// use futures_core::ready;
/// use mio::Ready;
/// use mio::net::{TcpStream, TcpListener};
/// use std::io;
/// use std::task::{Context, Poll};
///
/// struct MyListener {
/// poll_evented: PollEvented<TcpListener>,
/// }
///
/// impl MyListener {
/// pub fn poll_accept(&mut self, cx: &mut Context<'_>) -> Poll<Result<TcpStream, io::Error>> {
/// let ready = Ready::readable();
///
/// ready!(self.poll_evented.poll_read_ready(cx, ready))?;
///
/// match self.poll_evented.get_ref().accept() {
/// Ok((socket, _)) => Poll::Ready(Ok(socket)),
/// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
/// self.poll_evented.clear_read_ready(cx, ready)?;
/// Poll::Pending
/// }
/// Err(e) => Poll::Ready(Err(e)),
/// }
/// }
/// }
/// ```
///
/// ## Platform-specific events
///
/// `PollEvented` also allows receiving platform-specific `mio::Ready` events.
/// These events are included as part of the read readiness event stream. The
/// write readiness event stream is only for `Ready::writable()` events.
///
/// [`std::io::Read`]: https://doc.rust-lang.org/std/io/trait.Read.html
/// [`std::io::Write`]: https://doc.rust-lang.org/std/io/trait.Write.html
/// [`AsyncRead`]: ../io/trait.AsyncRead.html
/// [`AsyncWrite`]: ../io/trait.AsyncWrite.html
/// [`mio::Evented`]: https://docs.rs/mio/0.6/mio/trait.Evented.html
/// [`Registration`]: struct.Registration.html
/// [`TcpListener`]: ../net/struct.TcpListener.html
/// [`clear_read_ready`]: #method.clear_read_ready
/// [`clear_write_ready`]: #method.clear_write_ready
/// [`poll_read_ready`]: #method.poll_read_ready
/// [`poll_write_ready`]: #method.poll_write_ready
pub struct PollEvented<E: Evented> {
io: Option<E>,
inner: Inner,
}
struct Inner {
registration: Registration,
/// Currently visible read readiness
read_readiness: AtomicUsize,
/// Currently visible write readiness
write_readiness: AtomicUsize,
}
// ===== impl PollEvented =====
macro_rules! poll_ready {
($me:expr, $mask:expr, $cache:ident, $take:ident, $poll:expr) => {{
// Load cached & encoded readiness.
let mut cached = $me.inner.$cache.load(Relaxed);
let mask = $mask | platform::hup();
// See if the current readiness matches any bits.
let mut ret = mio::Ready::from_usize(cached) & $mask;
if ret.is_empty() {
// Readiness does not match, consume the registration's readiness
// stream. This happens in a loop to ensure that the stream gets
// drained.
loop {
let ready = match $poll? {
Poll::Ready(v) => v,
Poll::Pending => return Poll::Pending,
};
cached |= ready.as_usize();
// Update the cache store
$me.inner.$cache.store(cached, Relaxed);
ret |= ready & mask;
if !ret.is_empty() {
return Poll::Ready(Ok(ret));
}
}
} else {
// Check what's new with the registration stream. This will not
// request to be notified
if let Some(ready) = $me.inner.registration.$take()? {
cached |= ready.as_usize();
$me.inner.$cache.store(cached, Relaxed);
}
Poll::Ready(Ok(mio::Ready::from_usize(cached)))
}
}};
}
impl<E> PollEvented<E>
where
E: Evented,
{
/// Creates a new `PollEvented` associated with the default reactor.
pub fn new(io: E) -> io::Result<Self> {
let registration = Registration::new(&io)?;
Ok(Self {
io: Some(io),
inner: Inner {
registration,
read_readiness: AtomicUsize::new(0),
write_readiness: AtomicUsize::new(0),
},
})
}
/// Returns a shared reference to the underlying I/O object this readiness
/// stream is wrapping.
pub fn get_ref(&self) -> &E {
self.io.as_ref().unwrap()
}
/// Returns a mutable reference to the underlying I/O object this readiness
/// stream is wrapping.
pub fn get_mut(&mut self) -> &mut E {
self.io.as_mut().unwrap()
}
/// Consumes self, returning the inner I/O object
///
/// This function will deregister the I/O resource from the reactor before
/// returning. If the deregistration operation fails, an error is returned.
///
/// Note that deregistering does not guarantee that the I/O resource can be
/// registered with a different reactor. Some I/O resource types can only be
/// associated with a single reactor instance for their lifetime.
pub fn into_inner(mut self) -> io::Result<E> {
let io = self.io.take().unwrap();
self.inner.registration.deregister(&io)?;
Ok(io)
}
/// Check the I/O resource's read readiness state.
///
/// The mask argument allows specifying what readiness to notify on. This
/// can be any value, including platform specific readiness, **except**
/// `writable`. HUP is always implicitly included on platforms that support
/// it.
///
/// If the resource is not ready for a read then `Poll::Pending` is returned
/// and the current task is notified once a new event is received.
///
/// The I/O resource will remain in a read-ready state until readiness is
/// cleared by calling [`clear_read_ready`].
///
/// [`clear_read_ready`]: #method.clear_read_ready
///
/// # Panics
///
/// This function panics if:
///
/// * `ready` includes writable.
/// * called from outside of a task context.
pub fn poll_read_ready(
&self,
cx: &mut Context<'_>,
mask: mio::Ready,
) -> Poll<io::Result<mio::Ready>> {
assert!(!mask.is_writable(), "cannot poll for write readiness");
poll_ready!(
self,
mask,
read_readiness,
take_read_ready,
self.inner.registration.poll_read_ready(cx)
)
}
/// Clears the I/O resource's read readiness state and registers the current
/// task to be notified once a read readiness event is received.
///
/// After calling this function, `poll_read_ready` will return
/// `Poll::Pending` until a new read readiness event has been received.
///
/// The `mask` argument specifies the readiness bits to clear. This may not
/// include `writable` or `hup`.
///
/// # Panics
///
/// This function panics if:
///
/// * `ready` includes writable or HUP
/// * called from outside of a task context.
pub fn clear_read_ready(&self, cx: &mut Context<'_>, ready: mio::Ready) -> io::Result<()> {
// Cannot clear write readiness
assert!(!ready.is_writable(), "cannot clear write readiness");
assert!(!platform::is_hup(ready), "cannot clear HUP readiness");
self.inner
.read_readiness
.fetch_and(!ready.as_usize(), Relaxed);
if self.poll_read_ready(cx, ready)?.is_ready() {
// Notify the current task
cx.waker().wake_by_ref();
}
Ok(())
}
/// Check the I/O resource's write readiness state.
///
/// This always checks for writable readiness and also checks for HUP
/// readiness on platforms that support it.
///
/// If the resource is not ready for a write then `Async::NotReady` is
/// returned and the current task is notified once a new event is received.
///
/// The I/O resource will remain in a write-ready state until readiness is
/// cleared by calling [`clear_write_ready`].
///
/// [`clear_write_ready`]: #method.clear_write_ready
///
/// # Panics
///
/// This function panics if:
///
/// * `ready` contains bits besides `writable` and `hup`.
/// * called from outside of a task context.
pub fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<mio::Ready>> {
poll_ready!(
self,
mio::Ready::writable(),
write_readiness,
take_write_ready,
self.inner.registration.poll_write_ready(cx)
)
}
/// Resets the I/O resource's write readiness state and registers the current
/// task to be notified once a write readiness event is received.
///
/// This only clears writable readiness. HUP (on platforms that support HUP)
/// cannot be cleared as it is a final state.
///
/// After calling this function, `poll_write_ready(Ready::writable())` will
/// return `NotReady` until a new write readiness event has been received.
///
/// # Panics
///
/// This function will panic if called from outside of a task context.
pub fn clear_write_ready(&self, cx: &mut Context<'_>) -> io::Result<()> {
let ready = mio::Ready::writable();
self.inner
.write_readiness
.fetch_and(!ready.as_usize(), Relaxed);
if self.poll_write_ready(cx)?.is_ready() {
// Notify the current task
cx.waker().wake_by_ref();
}
Ok(())
}
}
// ===== Read / Write impls =====
impl<E> AsyncRead for PollEvented<E>
where
E: Evented + Read + Unpin,
{
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
ready!(self.poll_read_ready(cx, mio::Ready::readable()))?;
let r = (*self).get_mut().read(buf);
if is_wouldblock(&r) {
self.clear_read_ready(cx, mio::Ready::readable())?;
return Poll::Pending;
}
Poll::Ready(r)
}
}
impl<E> AsyncWrite for PollEvented<E>
where
E: Evented + Write + Unpin,
{
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
ready!(self.poll_write_ready(cx))?;
let r = (*self).get_mut().write(buf);
if is_wouldblock(&r) {
self.clear_write_ready(cx)?;
return Poll::Pending;
}
Poll::Ready(r)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
ready!(self.poll_write_ready(cx))?;
let r = (*self).get_mut().flush();
if is_wouldblock(&r) {
self.clear_write_ready(cx)?;
return Poll::Pending;
}
Poll::Ready(r)
}
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
}
fn is_wouldblock<T>(r: &io::Result<T>) -> bool {
match *r {
Ok(_) => false,
Err(ref e) => e.kind() == io::ErrorKind::WouldBlock,
}
}
impl<E: Evented + fmt::Debug> fmt::Debug for PollEvented<E> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("PollEvented").field("io", &self.io).finish()
}
}
impl<E: Evented> Drop for PollEvented<E> {
fn drop(&mut self) {
if let Some(io) = self.io.take() {
// Ignore errors
let _ = self.inner.registration.deregister(&io);
}
}
}
tokio/src/process.rs
-2
View File
@@ -1,2 +0,0 @@
//! An implementation of asynchronous process management for Tokio.
pub use tokio_net::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command};
tokio/src/process/kill.rs
+13
View File
@@ -0,0 +1,13 @@
use std::io;
/// An interface for killing a running process.
pub(crate) trait Kill {
/// Forcefully kill the process.
fn kill(&mut self) -> io::Result<()>;
}
impl<T: Kill> Kill for &mut T {
fn kill(&mut self) -> io::Result<()> {
(**self).kill()
}
}
tokio/src/process/mod.rs
+1060
View File
File diff suppressed because it is too large Load Diff
tokio/src/process/unix/mod.rs
+225
View File
@@ -0,0 +1,225 @@
//! Unix handling of child processes
//!
//! Right now the only "fancy" thing about this is how we implement the
//! `Future` implementation on `Child` to get the exit status. Unix offers
//! no way to register a child with epoll, and the only real way to get a
//! notification when a process exits is the SIGCHLD signal.
//!
//! Signal handling in general is *super* hairy and complicated, and it's even
//! more complicated here with the fact that signals are coalesced, so we may
//! not get a SIGCHLD-per-child.
//!
//! Our best approximation here is to check *all spawned processes* for all
//! SIGCHLD signals received. To do that we create a `Signal`, implemented in
//! the `tokio-net` crate, which is a stream over signals being received.
//!
//! Later when we poll the process's exit status we simply check to see if a
//! SIGCHLD has happened since we last checked, and while that returns "yes" we
//! keep trying.
//!
//! Note that this means that this isn't really scalable, but then again
//! processes in general aren't scalable (e.g. millions) so it shouldn't be that
//! bad in theory...
mod orphan;
mod reap;
use self::orphan::{AtomicOrphanQueue, OrphanQueue, Wait};
use self::reap::Reaper;
use super::SpawnedChild;
use crate::net::util::PollEvented;
use crate::process::kill::Kill;
use crate::signal::unix::{signal, Signal, SignalKind};
use mio::event::Evented;
use mio::unix::{EventedFd, UnixReady};
use mio::{Poll as MioPoll, PollOpt, Ready, Token};
use std::fmt;
use std::future::Future;
use std::io;
use std::os::unix::io::{AsRawFd, RawFd};
use std::pin::Pin;
use std::process::{self, ExitStatus};
use std::task::Context;
use std::task::Poll;
impl Wait for process::Child {
fn id(&self) -> u32 {
self.id()
}
fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
self.try_wait()
}
}
impl Kill for process::Child {
fn kill(&mut self) -> io::Result<()> {
self.kill()
}
}
lazy_static::lazy_static! {
static ref ORPHAN_QUEUE: AtomicOrphanQueue<process::Child> = AtomicOrphanQueue::new();
}
struct GlobalOrphanQueue;
impl fmt::Debug for GlobalOrphanQueue {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
ORPHAN_QUEUE.fmt(fmt)
}
}
impl OrphanQueue<process::Child> for GlobalOrphanQueue {
fn push_orphan(&self, orphan: process::Child) {
ORPHAN_QUEUE.push_orphan(orphan)
}
fn reap_orphans(&self) {
ORPHAN_QUEUE.reap_orphans()
}
}
#[must_use = "futures do nothing unless polled"]
pub(crate) struct Child {
inner: Reaper<process::Child, GlobalOrphanQueue, Signal>,
}
impl fmt::Debug for Child {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Child")
.field("pid", &self.inner.id())
.finish()
}
}
pub(crate) fn spawn_child(cmd: &mut process::Command) -> io::Result<SpawnedChild> {
let mut child = cmd.spawn()?;
let stdin = stdio(child.stdin.take())?;
let stdout = stdio(child.stdout.take())?;
let stderr = stdio(child.stderr.take())?;
let signal = signal(SignalKind::child())?;
Ok(SpawnedChild {
child: Child {
inner: Reaper::new(child, GlobalOrphanQueue, signal),
},
stdin,
stdout,
stderr,
})
}
impl Child {
pub(crate) fn id(&self) -> u32 {
self.inner.id()
}
}
impl Kill for Child {
fn kill(&mut self) -> io::Result<()> {
self.inner.kill()
}
}
impl Future for Child {
type Output = io::Result<ExitStatus>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
Pin::new(&mut self.inner).poll(cx)
}
}
#[derive(Debug)]
pub(crate) struct Fd<T> {
inner: T,
}
impl<T> io::Read for Fd<T>
where
T: io::Read,
{
fn read(&mut self, bytes: &mut [u8]) -> io::Result<usize> {
self.inner.read(bytes)
}
}
impl<T> io::Write for Fd<T>
where
T: io::Write,
{
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
self.inner.write(bytes)
}
fn flush(&mut self) -> io::Result<()> {
self.inner.flush()
}
}
impl<T> AsRawFd for Fd<T>
where
T: AsRawFd,
{
fn as_raw_fd(&self) -> RawFd {
self.inner.as_raw_fd()
}
}
impl<T> Evented for Fd<T>
where
T: AsRawFd,
{
fn register(
&self,
poll: &MioPoll,
token: Token,
interest: Ready,
opts: PollOpt,
) -> io::Result<()> {
EventedFd(&self.as_raw_fd()).register(poll, token, interest | UnixReady::hup(), opts)
}
fn reregister(
&self,
poll: &MioPoll,
token: Token,
interest: Ready,
opts: PollOpt,
) -> io::Result<()> {
EventedFd(&self.as_raw_fd()).reregister(poll, token, interest | UnixReady::hup(), opts)
}
fn deregister(&self, poll: &MioPoll) -> io::Result<()> {
EventedFd(&self.as_raw_fd()).deregister(poll)
}
}
pub(crate) type ChildStdin = PollEvented<Fd<process::ChildStdin>>;
pub(crate) type ChildStdout = PollEvented<Fd<process::ChildStdout>>;
pub(crate) type ChildStderr = PollEvented<Fd<process::ChildStderr>>;
fn stdio<T>(option: Option<T>) -> io::Result<Option<PollEvented<Fd<T>>>>
where
T: AsRawFd,
{
let io = match option {
Some(io) => io,
None => return Ok(None),
};
// Set the fd to nonblocking before we pass it to the event loop
unsafe {
let fd = io.as_raw_fd();
let r = libc::fcntl(fd, libc::F_GETFL);
if r == -1 {
return Err(io::Error::last_os_error());
}
let r = libc::fcntl(fd, libc::F_SETFL, r | libc::O_NONBLOCK);
if r == -1 {
return Err(io::Error::last_os_error());
}
}
Ok(Some(PollEvented::new(Fd { inner: io })?))
}
tokio/src/process/unix/orphan.rs
+190
View File
@@ -0,0 +1,190 @@
use crossbeam_queue::SegQueue;
use std::io;
use std::process::ExitStatus;
/// An interface for waiting on a process to exit.
pub(crate) trait Wait {
/// Get the identifier for this process or diagnostics.
fn id(&self) -> u32;
/// Try waiting for a process to exit in a non-blocking manner.
fn try_wait(&mut self) -> io::Result<Option<ExitStatus>>;
}
impl<T: Wait> Wait for &mut T {
fn id(&self) -> u32 {
(**self).id()
}
fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
(**self).try_wait()
}
}
/// An interface for queueing up an orphaned process so that it can be reaped.
pub(crate) trait OrphanQueue<T> {
/// Add an orphan to the queue.
fn push_orphan(&self, orphan: T);
/// Attempt to reap every process in the queue, ignoring any errors and
/// enqueueing any orphans which have not yet exited.
fn reap_orphans(&self);
}
impl<T, O: OrphanQueue<T>> OrphanQueue<T> for &O {
fn push_orphan(&self, orphan: T) {
(**self).push_orphan(orphan);
}
fn reap_orphans(&self) {
(**self).reap_orphans()
}
}
/// An atomic implementation of `OrphanQueue`.
#[derive(Debug)]
pub(crate) struct AtomicOrphanQueue<T> {
queue: SegQueue<T>,
}
impl<T> AtomicOrphanQueue<T> {
pub(crate) fn new() -> Self {
Self {
queue: SegQueue::new(),
}
}
}
impl<T: Wait> OrphanQueue<T> for AtomicOrphanQueue<T> {
fn push_orphan(&self, orphan: T) {
self.queue.push(orphan)
}
fn reap_orphans(&self) {
let len = self.queue.len();
if len == 0 {
return;
}
let mut orphans = Vec::with_capacity(len);
while let Ok(mut orphan) = self.queue.pop() {
match orphan.try_wait() {
Ok(Some(_)) => {}
Err(_) => {
// TODO: bubble up error some how. Is this an internal bug?
// Shoudl we panic? Is it OK for this to be silently
// dropped?
}
// Still not done yet, we need to put it back in the queue
// when were done draining it, so that we don't get stuck
// in an infinite loop here
Ok(None) => orphans.push(orphan),
}
}
for orphan in orphans {
self.queue.push(orphan);
}
}
}
#[cfg(test)]
mod test {
use super::Wait;
use super::{AtomicOrphanQueue, OrphanQueue};
use std::cell::Cell;
use std::io;
use std::os::unix::process::ExitStatusExt;
use std::process::ExitStatus;
use std::rc::Rc;
struct MockWait {
total_waits: Rc<Cell<usize>>,
num_wait_until_status: usize,
return_err: bool,
}
impl MockWait {
fn new(num_wait_until_status: usize) -> Self {
Self {
total_waits: Rc::new(Cell::new(0)),
num_wait_until_status,
return_err: false,
}
}
fn with_err() -> Self {
Self {
total_waits: Rc::new(Cell::new(0)),
num_wait_until_status: 0,
return_err: true,
}
}
}
impl Wait for MockWait {
fn id(&self) -> u32 {
42
}
fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
let waits = self.total_waits.get();
let ret = if self.num_wait_until_status == waits {
if self.return_err {
Ok(Some(ExitStatus::from_raw(0)))
} else {
Err(io::Error::new(io::ErrorKind::Other, "mock err"))
}
} else {
Ok(None)
};
self.total_waits.set(waits + 1);
ret
}
}
#[test]
fn drain_attempts_a_single_reap_of_all_queued_orphans() {
let first_orphan = MockWait::new(0);
let second_orphan = MockWait::new(1);
let third_orphan = MockWait::new(2);
let fourth_orphan = MockWait::with_err();
let first_waits = first_orphan.total_waits.clone();
let second_waits = second_orphan.total_waits.clone();
let third_waits = third_orphan.total_waits.clone();
let fourth_waits = fourth_orphan.total_waits.clone();
let orphanage = AtomicOrphanQueue::new();
orphanage.push_orphan(first_orphan);
orphanage.push_orphan(third_orphan);
orphanage.push_orphan(second_orphan);
orphanage.push_orphan(fourth_orphan);
assert_eq!(orphanage.queue.len(), 4);
orphanage.reap_orphans();
assert_eq!(orphanage.queue.len(), 2);
assert_eq!(first_waits.get(), 1);
assert_eq!(second_waits.get(), 1);
assert_eq!(third_waits.get(), 1);
assert_eq!(fourth_waits.get(), 1);
orphanage.reap_orphans();
assert_eq!(orphanage.queue.len(), 1);
assert_eq!(first_waits.get(), 1);
assert_eq!(second_waits.get(), 2);
assert_eq!(third_waits.get(), 2);
assert_eq!(fourth_waits.get(), 1);
orphanage.reap_orphans();
assert_eq!(orphanage.queue.len(), 0);
assert_eq!(first_waits.get(), 1);
assert_eq!(second_waits.get(), 2);
assert_eq!(third_waits.get(), 3);
assert_eq!(fourth_waits.get(), 1);
orphanage.reap_orphans(); // Safe to reap when empty
}
}
tokio/src/process/unix/reap.rs
+334
View File
@@ -0,0 +1,334 @@
use super::orphan::{OrphanQueue, Wait};
use crate::process::kill::Kill;
use futures_core::stream::Stream;
use std::future::Future;
use std::io;
use std::ops::Deref;
use std::pin::Pin;
use std::process::ExitStatus;
use std::task::Context;
use std::task::Poll;
/// Orchestrates between registering interest for receiving signals when a
/// child process has exited, and attempting to poll for process completion.
#[derive(Debug)]
pub(crate) struct Reaper<W, Q, S>
where
W: Wait + Unpin,
Q: OrphanQueue<W>,
{
inner: Option<W>,
orphan_queue: Q,
signal: S,
}
impl<W, Q, S> Deref for Reaper<W, Q, S>
where
W: Wait + Unpin,
Q: OrphanQueue<W>,
{
type Target = W;
fn deref(&self) -> &Self::Target {
self.inner()
}
}
impl<W, Q, S> Reaper<W, Q, S>
where
W: Wait + Unpin,
Q: OrphanQueue<W>,
{
pub(crate) fn new(inner: W, orphan_queue: Q, signal: S) -> Self {
Self {
inner: Some(inner),
orphan_queue,
signal,
}
}
fn inner(&self) -> &W {
self.inner.as_ref().expect("inner has gone away")
}
fn inner_mut(&mut self) -> &mut W {
self.inner.as_mut().expect("inner has gone away")
}
}
impl<W, Q, S> Future for Reaper<W, Q, S>
where
W: Wait + Unpin,
Q: OrphanQueue<W> + Unpin,
S: Stream + Unpin,
{
type Output = io::Result<ExitStatus>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
loop {
// If the child hasn't exited yet, then it's our responsibility to
// ensure the current task gets notified when it might be able to
// make progress.
//
// As described in `spawn` above, we just indicate that we can
// next make progress once a SIGCHLD is received.
//
// However, we will register for a notification on the next signal
// BEFORE we poll the child. Otherwise it is possible that the child
// can exit and the signal can arrive after we last polled the child,
// but before we've registered for a notification on the next signal
// (this can cause a deadlock if there are no more spawned children
// which can generate a different signal for us). A side effect of
// pre-registering for signal notifications is that when the child
// exits, we will have already registered for an additional
// notification we don't need to consume. If another signal arrives,
// this future's task will be notified/woken up again. Since the
// futures model allows for spurious wake ups this extra wakeup
// should not cause significant issues with parent futures.
let registered_interest = Pin::new(&mut self.signal).poll_next(cx).is_pending();
self.orphan_queue.reap_orphans();
if let Some(status) = self.inner_mut().try_wait()? {
return Poll::Ready(Ok(status));
}
// If our attempt to poll for the next signal was not ready, then
// we've arranged for our task to get notified and we can bail out.
if registered_interest {
return Poll::Pending;
} else {
// Otherwise, if the signal stream delivered a signal to us, we
// won't get notified at the next signal, so we'll loop and try
// again.
continue;
}
}
}
}
impl<W, Q, S> Kill for Reaper<W, Q, S>
where
W: Kill + Wait + Unpin,
Q: OrphanQueue<W>,
{
fn kill(&mut self) -> io::Result<()> {
self.inner_mut().kill()
}
}
impl<W, Q, S> Drop for Reaper<W, Q, S>
where
W: Wait + Unpin,
Q: OrphanQueue<W>,
{
fn drop(&mut self) {
if let Ok(Some(_)) = self.inner_mut().try_wait() {
return;
}
let orphan = self.inner.take().unwrap();
self.orphan_queue.push_orphan(orphan);
}
}
#[cfg(test)]
mod test {
use super::*;
use futures_core::stream::Stream;
use futures_util::future::FutureExt;
use std::cell::{Cell, RefCell};
use std::os::unix::process::ExitStatusExt;
use std::pin::Pin;
use std::process::ExitStatus;
use std::task::Context;
use std::task::Poll;
#[derive(Debug)]
struct MockWait {
total_kills: usize,
total_waits: usize,
num_wait_until_status: usize,
status: ExitStatus,
}
impl MockWait {
fn new(status: ExitStatus, num_wait_until_status: usize) -> Self {
Self {
total_kills: 0,
total_waits: 0,
num_wait_until_status,
status,
}
}
}
impl Wait for MockWait {
fn id(&self) -> u32 {
0
}
fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
let ret = if self.num_wait_until_status == self.total_waits {
Some(self.status)
} else {
None
};
self.total_waits += 1;
Ok(ret)
}
}
impl Kill for MockWait {
fn kill(&mut self) -> io::Result<()> {
self.total_kills += 1;
Ok(())
}
}
struct MockStream {
total_polls: usize,
values: Vec<Option<()>>,
}
impl MockStream {
fn new(values: Vec<Option<()>>) -> Self {
Self {
total_polls: 0,
values,
}
}
}
impl Stream for MockStream {
type Item = io::Result<()>;
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let inner = Pin::get_mut(self);
inner.total_polls += 1;
match inner.values.remove(0) {
Some(()) => Poll::Ready(Some(Ok(()))),
None => Poll::Pending,
}
}
}
struct MockQueue<W> {
all_enqueued: RefCell<Vec<W>>,
total_reaps: Cell<usize>,
}
impl<W> MockQueue<W> {
fn new() -> Self {
Self {
all_enqueued: RefCell::new(Vec::new()),
total_reaps: Cell::new(0),
}
}
}
impl<W: Wait> OrphanQueue<W> for MockQueue<W> {
fn push_orphan(&self, orphan: W) {
self.all_enqueued.borrow_mut().push(orphan);
}
fn reap_orphans(&self) {
self.total_reaps.set(self.total_reaps.get() + 1);
}
}
#[test]
fn reaper() {
let exit = ExitStatus::from_raw(0);
let mock = MockWait::new(exit, 3);
let mut grim = Reaper::new(
mock,
MockQueue::new(),
MockStream::new(vec![None, Some(()), None, None, None]),
);
let waker = futures_util::task::noop_waker();
let mut context = Context::from_waker(&waker);
// Not yet exited, interest registered
assert!(grim.poll_unpin(&mut context).is_pending());
assert_eq!(1, grim.signal.total_polls);
assert_eq!(1, grim.total_waits);
assert_eq!(1, grim.orphan_queue.total_reaps.get());
assert!(grim.orphan_queue.all_enqueued.borrow().is_empty());
// Not yet exited, couldn't register interest the first time
// but managed to register interest the second time around
assert!(grim.poll_unpin(&mut context).is_pending());
assert_eq!(3, grim.signal.total_polls);
assert_eq!(3, grim.total_waits);
assert_eq!(3, grim.orphan_queue.total_reaps.get());
assert!(grim.orphan_queue.all_enqueued.borrow().is_empty());
// Exited
if let Poll::Ready(r) = grim.poll_unpin(&mut context) {
assert!(r.is_ok());
let exit_code = r.unwrap();
assert_eq!(exit_code, exit);
} else {
unreachable!();
}
assert_eq!(4, grim.signal.total_polls);
assert_eq!(4, grim.total_waits);
assert_eq!(4, grim.orphan_queue.total_reaps.get());
assert!(grim.orphan_queue.all_enqueued.borrow().is_empty());
}
#[test]
fn kill() {
let exit = ExitStatus::from_raw(0);
let mut grim = Reaper::new(
MockWait::new(exit, 0),
MockQueue::new(),
MockStream::new(vec![None]),
);
grim.kill().unwrap();
assert_eq!(1, grim.total_kills);
assert_eq!(0, grim.orphan_queue.total_reaps.get());
assert!(grim.orphan_queue.all_enqueued.borrow().is_empty());
}
#[test]
fn drop_reaps_if_possible() {
let exit = ExitStatus::from_raw(0);
let mut mock = MockWait::new(exit, 0);
{
let queue = MockQueue::new();
let grim = Reaper::new(&mut mock, &queue, MockStream::new(vec![]));
drop(grim);
assert_eq!(0, queue.total_reaps.get());
assert!(queue.all_enqueued.borrow().is_empty());
}
assert_eq!(1, mock.total_waits);
assert_eq!(0, mock.total_kills);
}
#[test]
fn drop_enqueues_orphan_if_wait_fails() {
let exit = ExitStatus::from_raw(0);
let mut mock = MockWait::new(exit, 2);
{
let queue = MockQueue::<&mut MockWait>::new();
let grim = Reaper::new(&mut mock, &queue, MockStream::new(vec![]));
drop(grim);
assert_eq!(0, queue.total_reaps.get());
assert_eq!(1, queue.all_enqueued.borrow().len());
}
assert_eq!(1, mock.total_waits);
assert_eq!(0, mock.total_kills);
}
}
tokio/src/process/windows.rs
+193
View File
@@ -0,0 +1,193 @@
//! Windows asynchronous process handling.
//!
//! Like with Unix we don't actually have a way of registering a process with an
//! IOCP object. As a result we similarly need another mechanism for getting a
//! signal when a process has exited. For now this is implemented with the
//! `RegisterWaitForSingleObject` function in the kernel32.dll.
//!
//! This strategy is the same that libuv takes and essentially just queues up a
//! wait for the process in a kernel32-specific thread pool. Once the object is
//! notified (e.g. the process exits) then we have a callback that basically
//! just completes a `Oneshot`.
//!
//! The `poll_exit` implementation will attempt to wait for the process in a
//! nonblocking fashion, but failing that it'll fire off a
//! `RegisterWaitForSingleObject` and then wait on the other end of the oneshot
//! from then on out.
use crate::net::util::PollEvented;
use crate::process::kill::Kill;
use crate::process::SpawnedChild;
use crate::sync::oneshot;
use futures_util::future::Fuse;
use futures_util::future::FutureExt;
use mio_named_pipes::NamedPipe;
use std::fmt;
use std::future::Future;
use std::io;
use std::os::windows::prelude::*;
use std::os::windows::process::ExitStatusExt;
use std::pin::Pin;
use std::process::{Child as StdChild, Command as StdCommand, ExitStatus};
use std::ptr;
use std::task::Context;
use std::task::Poll;
use winapi::shared::minwindef::FALSE;
use winapi::shared::winerror::WAIT_TIMEOUT;
use winapi::um::handleapi::INVALID_HANDLE_VALUE;
use winapi::um::processthreadsapi::GetExitCodeProcess;
use winapi::um::synchapi::WaitForSingleObject;
use winapi::um::threadpoollegacyapiset::UnregisterWaitEx;
use winapi::um::winbase::{RegisterWaitForSingleObject, INFINITE, WAIT_OBJECT_0};
use winapi::um::winnt::{BOOLEAN, HANDLE, PVOID, WT_EXECUTEINWAITTHREAD, WT_EXECUTEONLYONCE};
#[must_use = "futures do nothing unless polled"]
pub(crate) struct Child {
child: StdChild,
waiting: Option<Waiting>,
}
impl fmt::Debug for Child {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Child")
.field("pid", &self.id())
.field("child", &self.child)
.field("waiting", &"..")
.finish()
}
}
struct Waiting {
rx: Fuse<oneshot::Receiver<()>>,
wait_object: HANDLE,
tx: *mut Option<oneshot::Sender<()>>,
}
unsafe impl Sync for Waiting {}
unsafe impl Send for Waiting {}
pub(crate) fn spawn_child(cmd: &mut StdCommand) -> io::Result<SpawnedChild> {
let mut child = cmd.spawn()?;
let stdin = stdio(child.stdin.take());
let stdout = stdio(child.stdout.take());
let stderr = stdio(child.stderr.take());
Ok(SpawnedChild {
child: Child {
child,
waiting: None,
},
stdin,
stdout,
stderr,
})
}
impl Child {
pub(crate) fn id(&self) -> u32 {
self.child.id()
}
}
impl Kill for Child {
fn kill(&mut self) -> io::Result<()> {
self.child.kill()
}
}
impl Future for Child {
type Output = io::Result<ExitStatus>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let inner = Pin::get_mut(self);
loop {
if let Some(ref mut w) = inner.waiting {
match w.rx.poll_unpin(cx) {
Poll::Ready(Ok(())) => {}
Poll::Ready(Err(_)) => panic!("should not be canceled"),
Poll::Pending => return Poll::Pending,
}
let status = try_wait(&inner.child)?.expect("not ready yet");
return Poll::Ready(Ok(status.into()));
}
if let Some(e) = try_wait(&inner.child)? {
return Poll::Ready(Ok(e.into()));
}
let (tx, rx) = oneshot::channel();
let ptr = Box::into_raw(Box::new(Some(tx)));
let mut wait_object = ptr::null_mut();
let rc = unsafe {
RegisterWaitForSingleObject(
&mut wait_object,
inner.child.as_raw_handle(),
Some(callback),
ptr as *mut _,
INFINITE,
WT_EXECUTEINWAITTHREAD | WT_EXECUTEONLYONCE,
)
};
if rc == 0 {
let err = io::Error::last_os_error();
drop(unsafe { Box::from_raw(ptr) });
return Poll::Ready(Err(err));
}
inner.waiting = Some(Waiting {
rx: rx.fuse(),
wait_object,
tx: ptr,
});
}
}
}
impl Drop for Waiting {
fn drop(&mut self) {
unsafe {
let rc = UnregisterWaitEx(self.wait_object, INVALID_HANDLE_VALUE);
if rc == 0 {
panic!("failed to unregister: {}", io::Error::last_os_error());
}
drop(Box::from_raw(self.tx));
}
}
}
unsafe extern "system" fn callback(ptr: PVOID, _timer_fired: BOOLEAN) {
let complete = &mut *(ptr as *mut Option<oneshot::Sender<()>>);
let _ = complete.take().unwrap().send(());
}
pub(crate) fn try_wait(child: &StdChild) -> io::Result<Option<ExitStatus>> {
unsafe {
match WaitForSingleObject(child.as_raw_handle(), 0) {
WAIT_OBJECT_0 => {}
WAIT_TIMEOUT => return Ok(None),
_ => return Err(io::Error::last_os_error()),
}
let mut status = 0;
let rc = GetExitCodeProcess(child.as_raw_handle(), &mut status);
if rc == FALSE {
Err(io::Error::last_os_error())
} else {
Ok(Some(ExitStatus::from_raw(status)))
}
}
}
pub(crate) type ChildStdin = PollEvented<NamedPipe>;
pub(crate) type ChildStdout = PollEvented<NamedPipe>;
pub(crate) type ChildStderr = PollEvented<NamedPipe>;
fn stdio<T>(option: Option<T>) -> Option<PollEvented<NamedPipe>>
where
T: IntoRawHandle,
{
let io = match option {
Some(io) => io,
None => return None,
};
let pipe = unsafe { NamedPipe::from_raw_handle(io.into_raw_handle()) };
PollEvented::new(pipe).ok()
}
tokio/src/runtime/current_thread/builder.rs
+1 -1
View File
@@ -1,9 +1,9 @@
use crate::net::driver::Reactor;
use crate::runtime::current_thread::Runtime;
use crate::timer::clock::Clock;
use crate::timer::timer::Timer;
use tokio_executor::current_thread::CurrentThread;
use tokio_net::driver::Reactor;
use std::io;
tokio/src/runtime/current_thread/runtime.rs
+1 -1
View File
@@ -1,10 +1,10 @@
use crate::net::driver::{self, Reactor};
use crate::runtime::current_thread::Builder;
use crate::timer::clock::{self, Clock};
use crate::timer::timer::{self, Timer};
use tokio_executor::current_thread::Handle as ExecutorHandle;
use tokio_executor::current_thread::{self, CurrentThread};
use tokio_net::driver::{self, Reactor};
use std::error::Error;
use std::fmt;
tokio/src/runtime/mod.rs
+1 -1
View File
@@ -119,7 +119,7 @@
//! }
//! ```
//!
//! [driver]: tokio_net::driver
//! [driver]: tokio::net::driver
//! [executor]: https://tokio.rs/docs/internals/runtime-model/#executors
//! [timer]: ../timer/index.html
//! [`Runtime`]: struct.Runtime.html
tokio/src/runtime/threadpool/builder.rs
+2 -2
View File
@@ -1,9 +1,9 @@
use super::{Inner, Runtime};
use crate::net::driver::{self, Reactor};
use crate::runtime::threadpool::{Inner, Runtime};
use crate::timer::clock::{self, Clock};
use crate::timer::timer::{self, Timer};
use tokio_executor::thread_pool;
use tokio_net::driver::{self, Reactor};
use std::sync::{Arc, Mutex};
use std::{fmt, io};
tokio/src/runtime/threadpool/mod.rs
+2 -2
View File
@@ -9,10 +9,10 @@ pub use self::spawner::Spawner;
#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
pub use tokio_executor::thread_pool::JoinHandle;
use crate::net::driver;
use crate::timer::timer;
use tokio_executor::thread_pool::ThreadPool;
use tokio_net::driver;
use std::future::Future;
use std::io;
@@ -42,7 +42,7 @@ struct Inner {
pool: ThreadPool,
/// Reactor handles
reactor_handles: Vec<tokio_net::driver::Handle>,
reactor_handles: Vec<crate::net::driver::Handle>,
/// Timer handles
timer_handles: Vec<timer::Handle>,
tokio/src/signal.rs
-4
View File
@@ -1,4 +0,0 @@
//! Asynchronous signal handling for `tokio`. ctrl-C notifications are
//! supported on both unix and windows systems. For finer grained signal
//! handling support on unix systems, see `tokio_net::signal::unix::Signal`.
pub use tokio_net::signal::ctrl_c;
tokio/src/signal/ctrl_c.rs
+46
View File
@@ -0,0 +1,46 @@
#[cfg(unix)]
use super::unix::{self as os_impl, Signal as Inner};
#[cfg(windows)]
use super::windows::{self as os_impl, Event as Inner};
use futures_core::stream::Stream;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
/// Represents a stream which receives "ctrl-c" notifications sent to the process.
///
/// In general signals are handled very differently across Unix and Windows, but
/// this is somewhat cross platform in terms of how it can be handled. A ctrl-c
/// event to a console process can be represented as a stream for both Windows
/// and Unix.
///
/// Note that there are a number of caveats listening for signals, and you may
/// wish to read up on the documentation in the `unix` or `windows` module to
/// take a peek.
///
/// Notably, a notification to this process notifies *all* streams listening to
/// this event. Moreover, the notifications **are coalesced** if they aren't processed
/// quickly enough. This means that if two notifications are received back-to-back,
/// then the stream may only receive one item about the two notifications.
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct CtrlC {
inner: Inner,
}
/// Creates a new stream which receives "ctrl-c" notifications sent to the
/// process.
///
/// This function binds to the default reactor.
pub fn ctrl_c() -> io::Result<CtrlC> {
os_impl::ctrl_c().map(|inner| CtrlC { inner })
}
impl Stream for CtrlC {
type Item = ();
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Pin::new(&mut self.inner).poll_next(cx)
}
}
tokio/src/signal/mod.rs
+93
View File
@@ -0,0 +1,93 @@
//! Asynchronous signal handling for Tokio
//!
//! The primary type exported from this crate, `unix::Signal`, allows
//! listening for arbitrary signals on Unix platforms, receiving them
//! in an asynchronous fashion.
//!
//! Note that signal handling is in general a very tricky topic and should be
//! used with great care. This crate attempts to implement 'best practice' for
//! signal handling, but it should be evaluated for your own applications' needs
//! to see if it's suitable.
//!
//! The are some fundamental limitations of this crate documented on the
//! `Signal` structure as well.
//!
//! # Examples
//!
//! Print out all ctrl-C notifications received
//!
//! ```rust,no_run
//! use tokio::signal;
//!
//! use futures_util::future;
//! use futures_util::stream::StreamExt;
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
//! // Create an infinite stream of "Ctrl+C" notifications. Each item received
//! // on this stream may represent multiple ctrl-c signals.
//! let ctrl_c = signal::ctrl_c()?;
//!
//! // Process each ctrl-c as it comes in
//! let prog = ctrl_c.for_each(|_| {
//! println!("ctrl-c received!");
//! future::ready(())
//! });
//!
//! prog.await;
//!
//! Ok(())
//! }
//! ```
//!
//! Wait for SIGHUP on Unix
//!
//! ```rust,no_run
//! # #[cfg(unix)] {
//!
//! use tokio::signal::{self, unix::{signal, SignalKind}};
//!
//! use futures_util::future;
//! use futures_util::stream::StreamExt;
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
//! // Create an infinite stream of "Ctrl+C" notifications. Each item received
//! // on this stream may represent multiple ctrl-c signals.
//! let ctrl_c = signal::ctrl_c()?;
//!
//! // Process each ctrl-c as it comes in
//! let prog = ctrl_c.for_each(|_| {
//! println!("ctrl-c received!");
//! future::ready(())
//! });
//!
//! prog.await;
//!
//! // Like the previous example, this is an infinite stream of signals
//! // being received, and signals may be coalesced while pending.
//! let stream = signal(SignalKind::hangup())?;
//!
//! // Convert out stream into a future and block the program
//! let (signal, _stream) = stream.into_future().await;
//! println!("got signal {:?}", signal);
//! Ok(())
//! }
//! # }
//! ```
mod ctrl_c;
mod registry;
mod os {
#[cfg(unix)]
pub(crate) use super::unix::{OsExtraData, OsStorage};
#[cfg(windows)]
pub(crate) use super::windows::{OsExtraData, OsStorage};
}
pub mod unix;
pub mod windows;
pub use self::ctrl_c::{ctrl_c, CtrlC};
tokio/src/signal/registry.rs
+310
View File
@@ -0,0 +1,310 @@
use crate::signal::os::{OsExtraData, OsStorage};
use tokio_sync::mpsc::Sender;
use lazy_static::lazy_static;
use std::ops;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Mutex;
pub(crate) type EventId = usize;
/// State for a specific event, whether a notification is pending delivery,
/// and what listeners are registered.
#[derive(Default, Debug)]
pub(crate) struct EventInfo {
pending: AtomicBool,
recipients: Mutex<Vec<Sender<()>>>,
}
/// An interface for retrieving the `EventInfo` for a particular eventId.
pub(crate) trait Storage {
/// Get the `EventInfo` for `id` if it exists.
fn event_info(&self, id: EventId) -> Option<&EventInfo>;
/// Invoke `f` once for each defined `EventInfo` in this storage.
fn for_each<'a, F>(&'a self, f: F)
where
F: FnMut(&'a EventInfo);
}
impl Storage for Vec<EventInfo> {
fn event_info(&self, id: EventId) -> Option<&EventInfo> {
self.get(id)
}
fn for_each<'a, F>(&'a self, f: F)
where
F: FnMut(&'a EventInfo),
{
self.iter().for_each(f)
}
}
/// An interface for initializing a type. Useful for situations where we cannot
/// inject a configured instance in the constructor of another type.
pub(crate) trait Init {
fn init() -> Self;
}
/// Manages and distributes event notifications to any registered listeners.
///
/// Generic over the underlying storage to allow for domain specific
/// optimizations (e.g. eventIds may or may not be contiguous).
#[derive(Debug)]
pub(crate) struct Registry<S> {
storage: S,
}
impl<S> Registry<S> {
fn new(storage: S) -> Self {
Self { storage }
}
}
impl<S: Storage> Registry<S> {
/// Register a new listener for `event_id`.
fn register_listener(&self, event_id: EventId, listener: Sender<()>) {
self.storage
.event_info(event_id)
.unwrap_or_else(|| panic!("invalid event_id: {}", event_id))
.recipients
.lock()
.unwrap()
.push(listener);
}
/// Mark `event_id` as having been delivered, without broadcasting it to
/// any listeners.
fn record_event(&self, event_id: EventId) {
if let Some(event_info) = self.storage.event_info(event_id) {
event_info.pending.store(true, Ordering::SeqCst)
}
}
/// Broadcast all previously recorded events to their respective listeners.
///
/// Returns true if an event was delivered to at least one listener.
fn broadcast(&self) -> bool {
let mut did_notify = false;
self.storage.for_each(|event_info| {
// Any signal of this kind arrived since we checked last?
if !event_info.pending.swap(false, Ordering::SeqCst) {
return;
}
let mut recipients = event_info.recipients.lock().unwrap();
// Notify all waiters on this signal that the signal has been
// received. If we can't push a message into the queue then we don't
// worry about it as everything is coalesced anyway. If the channel
// has gone away then we can remove that slot.
for i in (0..recipients.len()).rev() {
match recipients[i].try_send(()) {
Ok(()) => did_notify = true,
Err(ref e) if e.is_closed() => {
recipients.swap_remove(i);
}
// Channel is full, ignore the error since the
// receiver has already been woken up
Err(e) => {
// Sanity check in case this error type ever gets
// additional variants we have not considered.
debug_assert!(e.is_full());
}
}
}
});
did_notify
}
}
pub(crate) struct Globals {
extra: OsExtraData,
registry: Registry<OsStorage>,
}
impl ops::Deref for Globals {
type Target = OsExtraData;
fn deref(&self) -> &Self::Target {
&self.extra
}
}
impl Globals {
/// Register a new listener for `event_id`.
pub(crate) fn register_listener(&self, event_id: EventId, listener: Sender<()>) {
self.registry.register_listener(event_id, listener);
}
/// Mark `event_id` as having been delivered, without broadcasting it to
/// any listeners.
pub(crate) fn record_event(&self, event_id: EventId) {
self.registry.record_event(event_id);
}
/// Broadcast all previously recorded events to their respective listeners.
///
/// Returns true if an event was delivered to at least one listener.
pub(crate) fn broadcast(&self) -> bool {
self.registry.broadcast()
}
#[cfg(unix)]
pub(crate) fn storage(&self) -> &OsStorage {
&self.registry.storage
}
}
pub(crate) fn globals() -> Pin<&'static Globals>
where
OsExtraData: 'static + Send + Sync + Init,
OsStorage: 'static + Send + Sync + Init,
{
lazy_static! {
static ref GLOBALS: Pin<Box<Globals>> = Box::pin(Globals {
extra: OsExtraData::init(),
registry: Registry::new(OsStorage::init()),
});
}
GLOBALS.as_ref()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::runtime::current_thread::Runtime;
use crate::sync::{mpsc, oneshot};
use futures::{future, StreamExt};
#[test]
fn smoke() {
let mut rt = Runtime::new().unwrap();
rt.block_on(async move {
let registry = Registry::new(vec![
EventInfo::default(),
EventInfo::default(),
EventInfo::default(),
]);
let (first_tx, first_rx) = mpsc::channel(3);
let (second_tx, second_rx) = mpsc::channel(3);
let (third_tx, third_rx) = mpsc::channel(3);
registry.register_listener(0, first_tx);
registry.register_listener(1, second_tx);
registry.register_listener(2, third_tx);
let (fire, wait) = oneshot::channel();
crate::spawn(async {
wait.await.expect("wait failed");
// Record some events which should get coalesced
registry.record_event(0);
registry.record_event(0);
registry.record_event(1);
registry.record_event(1);
registry.broadcast();
// Send subsequent signal
registry.record_event(0);
registry.broadcast();
drop(registry);
});
let _ = fire.send(());
let all = future::join3(
first_rx.collect::<Vec<_>>(),
second_rx.collect::<Vec<_>>(),
third_rx.collect::<Vec<_>>(),
);
let (first_results, second_results, third_results) = all.await;
assert_eq!(2, first_results.len());
assert_eq!(1, second_results.len());
assert_eq!(0, third_results.len());
});
}
#[test]
#[should_panic = "invalid event_id: 1"]
fn register_panics_on_invalid_input() {
let registry = Registry::new(vec![EventInfo::default()]);
let (tx, _) = mpsc::channel(1);
registry.register_listener(1, tx);
}
#[test]
fn record_invalid_event_does_nothing() {
let registry = Registry::new(vec![EventInfo::default()]);
registry.record_event(42);
}
#[test]
fn broadcast_cleans_up_disconnected_listeners() {
let mut rt = Runtime::new().unwrap();
rt.block_on(async {
let registry = Registry::new(vec![EventInfo::default()]);
let (first_tx, first_rx) = mpsc::channel(1);
let (second_tx, second_rx) = mpsc::channel(1);
let (third_tx, third_rx) = mpsc::channel(1);
registry.register_listener(0, first_tx);
registry.register_listener(0, second_tx);
registry.register_listener(0, third_tx);
drop(first_rx);
drop(second_rx);
let (fire, wait) = oneshot::channel();
crate::spawn(async {
wait.await.expect("wait failed");
registry.record_event(0);
registry.broadcast();
assert_eq!(1, registry.storage[0].recipients.lock().unwrap().len());
drop(registry);
});
let _ = fire.send(());
let results: Vec<()> = third_rx.collect().await;
assert_eq!(1, results.len());
});
}
#[test]
fn broadcast_returns_if_at_least_one_event_fired() {
let registry = Registry::new(vec![EventInfo::default()]);
registry.record_event(0);
assert_eq!(false, registry.broadcast());
let (first_tx, first_rx) = mpsc::channel(1);
let (second_tx, second_rx) = mpsc::channel(1);
registry.register_listener(0, first_tx);
registry.register_listener(0, second_tx);
registry.record_event(0);
assert_eq!(true, registry.broadcast());
drop(first_rx);
registry.record_event(0);
assert_eq!(false, registry.broadcast());
drop(second_rx);
}
}
tokio/src/signal/unix.rs
+426
View File
@@ -0,0 +1,426 @@
//! Unix-specific types for signal handling.
//!
//! This module is only defined on Unix platforms and contains the primary
//! `Signal` type for receiving notifications of signals.
#![cfg(unix)]
use super::registry::{globals, EventId, EventInfo, Globals, Init, Storage};
use crate::net::util::PollEvented;
use tokio_io::AsyncRead;
use tokio_sync::mpsc::{channel, Receiver};
use futures_core::stream::Stream;
use libc::c_int;
use mio_uds::UnixStream;
use std::future::Future;
use std::io::{self, Error, ErrorKind, Write};
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Once;
use std::task::{Context, Poll};
pub(crate) type OsStorage = Vec<SignalInfo>;
// Number of different unix signals
// (FreeBSD has 33)
const SIGNUM: usize = 33;
impl Init for OsStorage {
fn init() -> Self {
(0..SIGNUM).map(|_| SignalInfo::default()).collect()
}
}
impl Storage for OsStorage {
fn event_info(&self, id: EventId) -> Option<&EventInfo> {
self.get(id).map(|si| &si.event_info)
}
fn for_each<'a, F>(&'a self, f: F)
where
F: FnMut(&'a EventInfo),
{
self.iter().map(|si| &si.event_info).for_each(f)
}
}
#[derive(Debug)]
pub(crate) struct OsExtraData {
sender: UnixStream,
receiver: UnixStream,
}
impl Init for OsExtraData {
fn init() -> Self {
let (receiver, sender) = UnixStream::pair().expect("failed to create UnixStream");
Self { sender, receiver }
}
}
/// Represents the specific kind of signal to listen for.
#[derive(Debug, Clone, Copy)]
pub struct SignalKind(c_int);
impl SignalKind {
/// Allows for listening to any valid OS signal.
///
/// For example, this can be used for listening for platform-specific
/// signals.
/// ```rust,no_run
/// # use tokio::signal::unix::SignalKind;
/// # let signum = -1;
/// // let signum = libc::OS_SPECIFIC_SIGNAL;
/// let kind = SignalKind::from_raw(signum);
/// ```
pub fn from_raw(signum: c_int) -> Self {
Self(signum)
}
/// Represents the SIGALRM signal.
///
/// On Unix systems this signal is sent when a real-time timer has expired.
/// By default, the process is terminated by this signal.
pub fn alarm() -> Self {
Self(libc::SIGALRM)
}
/// Represents the SIGCHLD signal.
///
/// On Unix systems this signal is sent when the status of a child process
/// has changed. By default, this signal is ignored.
pub fn child() -> Self {
Self(libc::SIGCHLD)
}
/// Represents the SIGHUP signal.
///
/// On Unix systems this signal is sent when the terminal is disconnected.
/// By default, the process is terminated by this signal.
pub fn hangup() -> Self {
Self(libc::SIGHUP)
}
/// Represents the SIGINFO signal.
///
/// On Unix systems this signal is sent to request a status update from the
/// process. By default, this signal is ignored.
#[cfg(any(
target_os = "dragonfly",
target_os = "freebsd",
target_os = "macos",
target_os = "netbsd",
target_os = "openbsd"
))]
pub fn info() -> Self {
Self(libc::SIGINFO)
}
/// Represents the SIGINT signal.
///
/// On Unix systems this signal is sent to interrupt a program.
/// By default, the process is terminated by this signal.
pub fn interrupt() -> Self {
Self(libc::SIGINT)
}
/// Represents the SIGIO signal.
///
/// On Unix systems this signal is sent when I/O operations are possible
/// on some file descriptor. By default, this signal is ignored.
pub fn io() -> Self {
Self(libc::SIGIO)
}
/// Represents the SIGPIPE signal.
///
/// On Unix systems this signal is sent when the process attempts to write
/// to a pipe which has no reader. By default, the process is terminated by
/// this signal.
pub fn pipe() -> Self {
Self(libc::SIGPIPE)
}
/// Represents the SIGQUIT signal.
///
/// On Unix systems this signal is sent to issue a shutdown of the
/// process, after which the OS will dump the process core.
/// By default, the process is terminated by this signal.
pub fn quit() -> Self {
Self(libc::SIGQUIT)
}
/// Represents the SIGTERM signal.
///
/// On Unix systems this signal is sent to issue a shutdown of the
/// process. By default, the process is terminated by this signal.
pub fn terminate() -> Self {
Self(libc::SIGTERM)
}
/// Represents the SIGUSR1 signal.
///
/// On Unix systems this is a user defined signal.
/// By default, the process is terminated by this signal.
pub fn user_defined1() -> Self {
Self(libc::SIGUSR1)
}
/// Represents the SIGUSR2 signal.
///
/// On Unix systems this is a user defined signal.
/// By default, the process is terminated by this signal.
pub fn user_defined2() -> Self {
Self(libc::SIGUSR2)
}
/// Represents the SIGWINCH signal.
///
/// On Unix systems this signal is sent when the terminal window is resized.
/// By default, this signal is ignored.
pub fn window_change() -> Self {
Self(libc::SIGWINCH)
}
}
pub(crate) struct SignalInfo {
event_info: EventInfo,
init: Once,
initialized: AtomicBool,
}
impl Default for SignalInfo {
fn default() -> SignalInfo {
SignalInfo {
event_info: Default::default(),
init: Once::new(),
initialized: AtomicBool::new(false),
}
}
}
/// Our global signal handler for all signals registered by this module.
///
/// The purpose of this signal handler is to primarily:
///
/// 1. Flag that our specific signal was received (e.g. store an atomic flag)
/// 2. Wake up driver tasks by writing a byte to a pipe
///
/// Those two operations shoudl both be async-signal safe.
fn action(globals: Pin<&'static Globals>, signal: c_int) {
globals.record_event(signal as EventId);
// Send a wakeup, ignore any errors (anything reasonably possible is
// full pipe and then it will wake up anyway).
let mut sender = &globals.sender;
drop(sender.write(&[1]));
}
/// Enable this module to receive signal notifications for the `signal`
/// provided.
///
/// This will register the signal handler if it hasn't already been registered,
/// returning any error along the way if that fails.
fn signal_enable(signal: c_int) -> io::Result<()> {
if signal < 0 || signal_hook_registry::FORBIDDEN.contains(&signal) {
return Err(Error::new(
ErrorKind::Other,
format!("Refusing to register signal {}", signal),
));
}
let globals = globals();
let siginfo = match globals.storage().get(signal as EventId) {
Some(slot) => slot,
None => return Err(io::Error::new(io::ErrorKind::Other, "signal too large")),
};
let mut registered = Ok(());
siginfo.init.call_once(|| {
registered = unsafe {
signal_hook_registry::register(signal, move || action(globals, signal)).map(|_| ())
};
if registered.is_ok() {
siginfo.initialized.store(true, Ordering::Relaxed);
}
});
registered?;
// If the call_once failed, it won't be retried on the next attempt to register the signal. In
// such case it is not run, registered is still `Ok(())`, initialized is still false.
if siginfo.initialized.load(Ordering::Relaxed) {
Ok(())
} else {
Err(Error::new(
ErrorKind::Other,
"Failed to register signal handler",
))
}
}
#[derive(Debug)]
struct Driver {
wakeup: PollEvented<UnixStream>,
}
impl Future for Driver {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
// Drain the data from the pipe and maintain interest in getting more
self.drain(cx);
// Broadcast any signals which were received
globals().broadcast();
Poll::Pending
}
}
impl Driver {
fn new() -> io::Result<Driver> {
// NB: We give each driver a "fresh" reciever file descriptor to avoid
// the issues described in alexcrichton/tokio-process#42.
//
// In the past we would reuse the actual receiver file descriptor and
// swallow any errors around double registration of the same descriptor.
// I'm not sure if the second (failed) registration simply doesn't end up
// receiving wake up notifications, or there could be some race condition
// when consuming readiness events, but having distinct descriptors for
// distinct PollEvented instances appears to mitigate this.
//
// Unfortunately we cannot just use a single global PollEvented instance
// either, since we can't compare Handles or assume they will always
// point to the exact same reactor.
let stream = globals().receiver.try_clone()?;
let wakeup = PollEvented::new(stream)?;
Ok(Driver { wakeup })
}
/// Drain all data in the global receiver, ensuring we'll get woken up when
/// there is a write on the other end.
///
/// We do *NOT* use the existence of any read bytes as evidence a sigal was
/// received since the `pending` flags would have already been set if that
/// was the case. See #38 for more info.
fn drain(mut self: Pin<&mut Self>, cx: &mut Context<'_>) {
loop {
match Pin::new(&mut self.wakeup).poll_read(cx, &mut [0; 128]) {
Poll::Ready(Ok(0)) => panic!("EOF on self-pipe"),
Poll::Ready(Ok(_)) => {}
Poll::Ready(Err(e)) => panic!("Bad read on self-pipe: {}", e),
Poll::Pending => break,
}
}
}
}
/// An implementation of `Stream` for receiving a particular type of signal.
///
/// This structure implements the `Stream` trait and represents notifications
/// of the current process receiving a particular signal. The signal being
/// listened for is passed to `Signal::new`, and the same signal number is then
/// yielded as each element for the stream.
///
/// In general signal handling on Unix is a pretty tricky topic, and this
/// structure is no exception! There are some important limitations to keep in
/// mind when using `Signal` streams:
///
/// * Signals handling in Unix already necessitates coalescing signals
/// together sometimes. This `Signal` stream is also no exception here in
/// that it will also coalesce signals. That is, even if the signal handler
/// for this process runs multiple times, the `Signal` stream may only return
/// one signal notification. Specifically, before `poll` is called, all
/// signal notifications are coalesced into one item returned from `poll`.
/// Once `poll` has been called, however, a further signal is guaranteed to
/// be yielded as an item.
///
/// Put another way, any element pulled off the returned stream corresponds to
/// *at least one* signal, but possibly more.
///
/// * Signal handling in general is relatively inefficient. Although some
/// improvements are possible in this crate, it's recommended to not plan on
/// having millions of signal channels open.
///
/// * Currently the "driver task" to process incoming signals never exits. This
/// driver task runs in the background of the event loop provided, and
/// in general you shouldn't need to worry about it.
///
/// If you've got any questions about this feel free to open an issue on the
/// repo, though, as I'd love to chat about this! In other words, I'd love to
/// alleviate some of these limitations if possible!
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct Signal {
driver: Driver,
rx: Receiver<()>,
}
/// Creates a new stream which will receive notifications when the current
/// process receives the signal `signal`.
///
/// This function will create a new stream which binds to the default reactor.
/// The `Signal` stream is an infinite stream which will receive
/// notifications whenever a signal is received. More documentation can be
/// found on `Signal` itself, but to reiterate:
///
/// * Signals may be coalesced beyond what the kernel already does.
/// * Once a signal handler is registered with the process the underlying
/// libc signal handler is never unregistered.
///
/// A `Signal` stream can be created for a particular signal number
/// multiple times. When a signal is received then all the associated
/// channels will receive the signal notification.
///
/// # Errors
///
/// * If the lower-level C functions fail for some reason.
/// * If the previous initialization of this specific signal failed.
/// * If the signal is one of
/// [`signal_hook::FORBIDDEN`](https://docs.rs/signal-hook/*/signal_hook/fn.register.html#panics)
pub fn signal(kind: SignalKind) -> io::Result<Signal> {
let signal = kind.0;
// Turn the signal delivery on once we are ready for it
signal_enable(signal)?;
// Ensure there's a driver for our associated event loop processing
// signals.
let driver = Driver::new()?;
// One wakeup in a queue is enough, no need for us to buffer up any
// more.
let (tx, rx) = channel(1);
globals().register_listener(signal as EventId, tx);
Ok(Signal { driver, rx })
}
pub(crate) fn ctrl_c() -> io::Result<Signal> {
signal(SignalKind::interrupt())
}
impl Stream for Signal {
type Item = ();
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let _ = Pin::new(&mut self.driver).poll(cx);
self.rx.poll_recv(cx)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn signal_enable_error_on_invalid_input() {
signal_enable(-1).unwrap_err();
}
#[test]
fn signal_enable_error_on_forbidden_input() {
signal_enable(signal_hook_registry::FORBIDDEN[0]).unwrap_err();
}
}
tokio/src/signal/windows.rs
+217
View File
@@ -0,0 +1,217 @@
//! Windows-specific types for signal handling.
//!
//! This module is only defined on Windows and contains the primary `Event` type
//! for receiving notifications of events. These events are listened for via the
//! `SetConsoleCtrlHandler` function which receives events of the type
//! `CTRL_C_EVENT` and `CTRL_BREAK_EVENT`
#![cfg(windows)]
use super::registry::{globals, EventId, EventInfo, Init, Storage};
use tokio_sync::mpsc::{channel, Receiver};
use futures_core::stream::Stream;
use std::convert::TryFrom;
use std::io;
use std::pin::Pin;
use std::sync::Once;
use std::task::{Context, Poll};
use winapi::shared::minwindef::*;
use winapi::um::consoleapi::SetConsoleCtrlHandler;
use winapi::um::wincon::*;
#[derive(Debug)]
pub(crate) struct OsStorage {
ctrl_c: EventInfo,
ctrl_break: EventInfo,
}
impl Init for OsStorage {
fn init() -> Self {
Self {
ctrl_c: EventInfo::default(),
ctrl_break: EventInfo::default(),
}
}
}
impl Storage for OsStorage {
fn event_info(&self, id: EventId) -> Option<&EventInfo> {
match DWORD::try_from(id) {
Ok(CTRL_C_EVENT) => Some(&self.ctrl_c),
Ok(CTRL_BREAK_EVENT) => Some(&self.ctrl_break),
_ => None,
}
}
fn for_each<'a, F>(&'a self, mut f: F)
where
F: FnMut(&'a EventInfo),
{
f(&self.ctrl_c);
f(&self.ctrl_break);
}
}
#[derive(Debug)]
pub(crate) struct OsExtraData {}
impl Init for OsExtraData {
fn init() -> Self {
Self {}
}
}
/// Stream of events discovered via `SetConsoleCtrlHandler`.
///
/// This structure can be used to listen for events of the type `CTRL_C_EVENT`
/// and `CTRL_BREAK_EVENT`. The `Stream` trait is implemented for this struct
/// and will resolve for each notification received by the process. Note that
/// there are few limitations with this as well:
///
/// * A notification to this process notifies *all* `Event` streams for that
/// event type.
/// * Notifications to an `Event` stream **are coalesced** if they aren't
/// processed quickly enough. This means that if two notifications are
/// received back-to-back, then the stream may only receive one item about the
/// two notifications.
// FIXME: refactor and combine with unix::Signal
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub(crate) struct Event {
rx: Receiver<()>,
}
impl Event {
fn new(signum: DWORD) -> io::Result<Self> {
global_init()?;
let (tx, rx) = channel(1);
globals().register_listener(signum as EventId, tx);
Ok(Event { rx })
}
}
pub(crate) fn ctrl_c() -> io::Result<Event> {
Event::new(CTRL_C_EVENT)
}
impl Stream for Event {
type Item = ();
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.rx.poll_recv(cx)
}
}
fn global_init() -> io::Result<()> {
static INIT: Once = Once::new();
let mut init = None;
INIT.call_once(|| unsafe {
let rc = SetConsoleCtrlHandler(Some(handler), TRUE);
let ret = if rc == 0 {
Err(io::Error::last_os_error())
} else {
Ok(())
};
init = Some(ret);
});
init.unwrap_or_else(|| Ok(()))
}
unsafe extern "system" fn handler(ty: DWORD) -> BOOL {
let globals = globals();
globals.record_event(ty as EventId);
// According to https://docs.microsoft.com/en-us/windows/console/handlerroutine
// the handler routine is always invoked in a new thread, thus we don't
// have the same restrictions as in Unix signal handlers, meaning we can
// go ahead and perform the broadcast here.
if globals.broadcast() {
TRUE
} else {
// No one is listening for this notification any more
// let the OS fire the next (possibly the default) handler.
FALSE
}
}
/// Represents a stream which receives "ctrl-break" notifications sent to the process
/// via `SetConsoleCtrlHandler`.
///
/// A notification to this process notifies *all* streams listening to
/// this event. Moreover, the notifications **are coalesced** if they aren't processed
/// quickly enough. This means that if two notifications are received back-to-back,
/// then the stream may only receive one item about the two notifications.
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct CtrlBreak {
inner: Event,
}
/// Creates a new stream which receives "ctrl-break" notifications sent to the
/// process.
///
/// This function binds to the default reactor.
pub fn ctrl_break() -> io::Result<CtrlBreak> {
Event::new(CTRL_BREAK_EVENT).map(|inner| CtrlBreak { inner })
}
impl Stream for CtrlBreak {
type Item = ();
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Pin::new(&mut self.inner)
.poll_next(cx)
.map(|item| item.map(|_| ()))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::runtime::current_thread::Runtime;
use futures_util::stream::StreamExt;
#[test]
fn ctrl_c() {
let mut rt = Runtime::new().unwrap();
rt.block_on(async {
let ctrl_c = crate::signal::ctrl_c().expect("failed to create CtrlC");
// Windows doesn't have a good programmatic way of sending events
// like sending signals on Unix, so we'll stub out the actual OS
// integration and test that our handling works.
unsafe {
super::handler(CTRL_C_EVENT);
}
let _ = ctrl_c.into_future().await;
});
}
#[test]
fn ctrl_break() {
let mut rt = Runtime::new().unwrap();
rt.block_on(async {
let ctrl_break = super::ctrl_break().expect("failed to create CtrlC");
// Windows doesn't have a good programmatic way of sending events
// like sending signals on Unix, so we'll stub out the actual OS
// integration and test that our handling works.
unsafe {
super::handler(CTRL_BREAK_EVENT);
}
let _ = ctrl_break.into_future().await;
});
}
}
tokio/tests/buffered.rs
-1
View File
@@ -1,5 +1,4 @@
#![warn(rust_2018_idioms)]
#![cfg(feature = "default")]
use tokio::net::TcpListener;
use tokio::prelude::*;
tokio/tests/net_bind_resource.rs
+11
View File
@@ -0,0 +1,11 @@
use tokio::net::TcpListener;
use std::convert::TryFrom;
use std::net;
#[test]
#[should_panic]
fn no_runtime_panics_binding_net_tcp_listener() {
let listener = net::TcpListener::bind("127.0.0.1:0").expect("failed to bind listener");
let _ = TcpListener::try_from(listener);
}
tokio/tests/reactor.rs → tokio/tests/net_driver.rs
+3 -4
View File
@@ -1,8 +1,7 @@
#![warn(rust_2018_idioms)]
#![cfg(feature = "default")]
use tokio_net::driver::Reactor;
use tokio_net::tcp::TcpListener;
use tokio::net::driver::Reactor;
use tokio::net::TcpListener;
use tokio_test::{assert_ok, assert_pending};
use futures_util::task::{waker_ref, ArcWake};
@@ -66,7 +65,7 @@ fn test_drop_on_notify() {
{
let handle = reactor.handle();
let _reactor = tokio_net::driver::set_default(&handle);
let _reactor = tokio::net::driver::set_default(&handle);
let waker = waker_ref(&task);
let mut cx = Context::from_waker(&waker);
assert_pending!(task.future.lock().unwrap().as_mut().poll(&mut cx));
tokio/tests/drop-core.rs → tokio/tests/net_driver_drop.rs
+1 -2
View File
@@ -1,8 +1,7 @@
#![warn(rust_2018_idioms)]
#![cfg(feature = "default")]
use tokio::net::driver::{self, Reactor};
use tokio::net::TcpListener;
use tokio_net::driver::{self, Reactor};
use tokio_test::{assert_err, assert_pending, assert_ready, task};
#[test]
tokio/tests/process_issue_42.rs
+56
View File
@@ -0,0 +1,56 @@
#![cfg(feature = "process")]
#![cfg(unix)]
#![warn(rust_2018_idioms)]
use tokio::process::Command;
use tokio::runtime::current_thread;
use futures_util::future::FutureExt;
use futures_util::stream::FuturesOrdered;
use std::process::Stdio;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::Duration;
fn run_test() {
let finished = Arc::new(AtomicBool::new(false));
let finished_clone = finished.clone();
thread::spawn(move || {
let mut rt = current_thread::Runtime::new().expect("failed to get runtime");
let mut futures = FuturesOrdered::new();
rt.block_on(async {
for i in 0..2 {
futures.push(
Command::new("echo")
.arg(format!("I am spawned process #{}", i))
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
.unwrap()
.boxed(),
)
}
});
drop(rt);
finished_clone.store(true, Ordering::SeqCst);
});
thread::sleep(Duration::from_millis(1000));
assert!(
finished.load(Ordering::SeqCst),
"FINISHED flag not set, maybe we deadlocked?"
);
}
#[test]
fn issue_42() {
let max = 10;
for i in 0..max {
println!("running {}/{}", i, max);
run_test()
}
}
tokio/tests/process_smoke.rs
+29
View File
@@ -0,0 +1,29 @@
#![cfg(feature = "process")]
#![warn(rust_2018_idioms)]
use tokio::process::Command;
use tokio_test::assert_ok;
#[tokio::test]
async fn simple() {
let mut cmd;
if cfg!(windows) {
cmd = Command::new("cmd");
cmd.arg("/c");
} else {
cmd = Command::new("sh");
cmd.arg("-c");
}
let mut child = cmd.arg("exit 2").spawn().unwrap();
let id = child.id();
assert!(id > 0);
let status = assert_ok!((&mut child).await);
assert_eq!(status.code(), Some(2));
assert_eq!(child.id(), id);
drop(child.kill());
}
tokio/tests/signal_ctrl_c.rs
+28
View File
@@ -0,0 +1,28 @@
#![cfg(unix)]
#![warn(rust_2018_idioms)]
mod support {
pub mod signal;
}
use support::signal::send_signal;
use tokio::prelude::*;
use tokio::signal;
use tokio::sync::oneshot;
#[tokio::test]
async fn ctrl_c() {
let ctrl_c = signal::ctrl_c().expect("failed to init ctrl_c");
let (fire, wait) = oneshot::channel();
// NB: simulate a signal coming in by exercising our signal handler
// to avoid complications with sending SIGINT to the test process
tokio::spawn(async {
wait.await.expect("wait failed");
send_signal(libc::SIGINT);
});
let _ = fire.send(());
let _ = ctrl_c.into_future().await;
}
tokio/tests/signal_drop_recv.rs
+22
View File
@@ -0,0 +1,22 @@
#![cfg(unix)]
#![warn(rust_2018_idioms)]
mod support {
pub mod signal;
}
use support::signal::send_signal;
use tokio::prelude::*;
use tokio::signal::unix::{signal, SignalKind};
#[tokio::test]
async fn drop_then_get_a_signal() {
let kind = SignalKind::user_defined1();
let sig = signal(kind).expect("failed to create first signal");
drop(sig);
send_signal(libc::SIGUSR1);
let sig = signal(kind).expect("failed to create second signal");
let _ = sig.into_future().await;
}
tokio/tests/signal_drop_rt.rs
+37
View File
@@ -0,0 +1,37 @@
#![cfg(unix)]
#![warn(rust_2018_idioms)]
mod support {
pub mod signal;
}
use support::signal::send_signal;
use tokio::prelude::*;
use tokio::runtime::current_thread::Runtime;
use tokio::signal::unix::{signal, SignalKind};
#[test]
fn dropping_loops_does_not_cause_starvation() {
let kind = SignalKind::user_defined1();
let mut first_rt = Runtime::new().expect("failed to init first runtime");
let mut first_signal =
first_rt.block_on(async { signal(kind).expect("failed to register first signal") });
let mut second_rt = Runtime::new().expect("failed to init second runtime");
let mut second_signal =
second_rt.block_on(async { signal(kind).expect("failed to register second signal") });
send_signal(libc::SIGUSR1);
first_rt
.block_on(first_signal.next())
.expect("failed to await first signal");
drop(first_rt);
drop(first_signal);
send_signal(libc::SIGUSR1);
second_rt.block_on(second_signal.next());
}
tokio/tests/signal_drop_signal.rs
+26
View File
@@ -0,0 +1,26 @@
#![cfg(unix)]
#![warn(rust_2018_idioms)]
mod support {
pub mod signal;
}
use support::signal::send_signal;
use tokio::prelude::*;
use tokio::signal::unix::{signal, SignalKind};
#[tokio::test]
async fn dropping_signal_does_not_deregister_any_other_instances() {
let kind = SignalKind::user_defined1();
// Signals should not starve based on ordering
let first_duplicate_signal = signal(kind).expect("failed to register first duplicate signal");
let sig = signal(kind).expect("failed to register signal");
let second_duplicate_signal = signal(kind).expect("failed to register second duplicate signal");
drop(first_duplicate_signal);
drop(second_duplicate_signal);
send_signal(libc::SIGUSR1);
let _ = sig.into_future().await;
}
tokio/tests/signal_multi_rt.rs
+47
View File
@@ -0,0 +1,47 @@
#![cfg(unix)]
#![warn(rust_2018_idioms)]
mod support {
pub mod signal;
}
use support::signal::send_signal;
use tokio::prelude::*;
use tokio::runtime::current_thread::Runtime;
use tokio::signal::unix::{signal, SignalKind};
use std::sync::mpsc::channel;
use std::thread;
#[test]
fn multi_loop() {
// An "ordinary" (non-future) channel
let (sender, receiver) = channel();
// Run multiple times, to make sure there are no race conditions
for _ in 0..10 {
// Run multiple event loops, each one in its own thread
let threads: Vec<_> = (0..4)
.map(|_| {
let sender = sender.clone();
thread::spawn(move || {
let mut rt = Runtime::new().unwrap();
let _ = rt.block_on(async {
let signal = signal(SignalKind::hangup()).unwrap();
sender.send(()).unwrap();
signal.into_future().await
});
})
})
.collect();
// Wait for them to declare they're ready
for &_ in threads.iter() {
receiver.recv().unwrap();
}
// Send a signal
send_signal(libc::SIGHUP);
// Make sure the threads terminated correctly
for t in threads {
t.join().unwrap();
}
}
}
tokio/tests/signal_no_rt.rs
+10
View File
@@ -0,0 +1,10 @@
#![cfg(unix)]
#![warn(rust_2018_idioms)]
use tokio::signal::unix::{signal, SignalKind};
#[test]
#[should_panic]
fn no_runtime_panics_creating_signals() {
let _ = signal(SignalKind::hangup());
}
tokio/tests/signal_notify_both.rs
+23
View File
@@ -0,0 +1,23 @@
#![cfg(unix)]
#![warn(rust_2018_idioms)]
mod support {
pub mod signal;
}
use support::signal::send_signal;
use tokio::prelude::*;
use tokio::signal::unix::{signal, SignalKind};
use futures::future;
#[tokio::test]
async fn notify_both() {
let kind = SignalKind::user_defined2();
let signal1 = signal(kind).expect("failed to create signal1");
let signal2 = signal(kind).expect("failed to create signal2");
send_signal(libc::SIGUSR2);
let _ = future::join(signal1.into_future(), signal2.into_future()).await;
}
tokio/tests/signal_twice.rs
+25
View File
@@ -0,0 +1,25 @@
#![cfg(unix)]
#![warn(rust_2018_idioms)]
mod support {
pub mod signal;
}
use support::signal::send_signal;
use tokio::prelude::*;
use tokio::signal::unix::{signal, SignalKind};
#[tokio::test]
async fn twice() {
let kind = SignalKind::user_defined1();
let mut sig = signal(kind).expect("failed to get signal");
for _ in 0..2 {
send_signal(libc::SIGUSR1);
let (item, sig_next) = sig.into_future().await;
assert_eq!(item, Some(()));
sig = sig_next;
}
}
tokio/tests/signal_usr1.rs
+23
View File
@@ -0,0 +1,23 @@
#![cfg(unix)]
#![warn(rust_2018_idioms)]
mod support {
pub mod signal;
}
use support::signal::send_signal;
use tokio::prelude::*;
use tokio::signal::unix::{signal, SignalKind};
use tokio_test::assert_ok;
#[tokio::test]
async fn signal_usr1() {
let signal = assert_ok!(
signal(SignalKind::user_defined1()),
"failed to create signal"
);
send_signal(libc::SIGUSR1);
let _ = signal.into_future().await;
}
tokio/tests/support/signal.rs
+7
View File
@@ -0,0 +1,7 @@
pub fn send_signal(signal: libc::c_int) {
use libc::{getpid, kill};
unsafe {
assert_eq!(kill(getpid(), signal), 0);
}
}
tokio/tests/tcp_accept.rs
+39
View File
@@ -0,0 +1,39 @@
#![warn(rust_2018_idioms)]
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::oneshot;
use tokio_test::assert_ok;
use std::net::{IpAddr, SocketAddr};
macro_rules! test_accept {
($(($ident:ident, $target:expr),)*) => {
$(
#[tokio::test]
async fn $ident() {
let mut listener = assert_ok!(TcpListener::bind($target).await);
let addr = listener.local_addr().unwrap();
let (tx, rx) = oneshot::channel();
tokio::spawn(async move {
let (socket, _) = assert_ok!(listener.accept().await);
assert_ok!(tx.send(socket));
});
let cli = assert_ok!(TcpStream::connect(&addr).await);
let srv = assert_ok!(rx.await);
assert_eq!(cli.local_addr().unwrap(), srv.peer_addr().unwrap());
}
)*
}
}
test_accept! {
(ip_str, "127.0.0.1:0"),
(host_str, "localhost:0"),
(socket_addr, "127.0.0.1:0".parse::<SocketAddr>().unwrap()),
(str_port_tuple, ("127.0.0.1", 0)),
(ip_port_tuple, ("127.0.0.1".parse::<IpAddr>().unwrap(), 0)),
}
tokio/tests/tcp_connect.rs
+228
View File
@@ -0,0 +1,228 @@
#![warn(rust_2018_idioms)]
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::oneshot;
use tokio_test::assert_ok;
use futures::join;
#[tokio::test]
async fn connect_v4() {
let mut srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await);
let addr = assert_ok!(srv.local_addr());
assert!(addr.is_ipv4());
let (tx, rx) = oneshot::channel();
tokio::spawn(async move {
let (socket, addr) = assert_ok!(srv.accept().await);
assert_eq!(addr, assert_ok!(socket.peer_addr()));
assert_ok!(tx.send(socket));
});
let mine = assert_ok!(TcpStream::connect(&addr).await);
let theirs = assert_ok!(rx.await);
assert_eq!(
assert_ok!(mine.local_addr()),
assert_ok!(theirs.peer_addr())
);
assert_eq!(
assert_ok!(theirs.local_addr()),
assert_ok!(mine.peer_addr())
);
}
#[tokio::test]
async fn connect_v6() {
let mut srv = assert_ok!(TcpListener::bind("[::1]:0").await);
let addr = assert_ok!(srv.local_addr());
assert!(addr.is_ipv6());
let (tx, rx) = oneshot::channel();
tokio::spawn(async move {
let (socket, addr) = assert_ok!(srv.accept().await);
assert_eq!(addr, assert_ok!(socket.peer_addr()));
assert_ok!(tx.send(socket));
});
let mine = assert_ok!(TcpStream::connect(&addr).await);
let theirs = assert_ok!(rx.await);
assert_eq!(
assert_ok!(mine.local_addr()),
assert_ok!(theirs.peer_addr())
);
assert_eq!(
assert_ok!(theirs.local_addr()),
assert_ok!(mine.peer_addr())
);
}
#[tokio::test]
async fn connect_addr_ip_string() {
let mut srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await);
let addr = assert_ok!(srv.local_addr());
let addr = format!("127.0.0.1:{}", addr.port());
let server = async {
assert_ok!(srv.accept().await);
};
let client = async {
assert_ok!(TcpStream::connect(addr).await);
};
join!(server, client);
}
#[tokio::test]
async fn connect_addr_ip_str_slice() {
let mut srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await);
let addr = assert_ok!(srv.local_addr());
let addr = format!("127.0.0.1:{}", addr.port());
let server = async {
assert_ok!(srv.accept().await);
};
let client = async {
assert_ok!(TcpStream::connect(&addr[..]).await);
};
join!(server, client);
}
#[tokio::test]
async fn connect_addr_host_string() {
let mut srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await);
let addr = assert_ok!(srv.local_addr());
let addr = format!("localhost:{}", addr.port());
let server = async {
assert_ok!(srv.accept().await);
};
let client = async {
assert_ok!(TcpStream::connect(addr).await);
};
join!(server, client);
}
#[tokio::test]
async fn connect_addr_ip_port_tuple() {
let mut srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await);
let addr = assert_ok!(srv.local_addr());
let addr = (addr.ip(), addr.port());
let server = async {
assert_ok!(srv.accept().await);
};
let client = async {
assert_ok!(TcpStream::connect(&addr).await);
};
join!(server, client);
}
#[tokio::test]
async fn connect_addr_ip_str_port_tuple() {
let mut srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await);
let addr = assert_ok!(srv.local_addr());
let addr = ("127.0.0.1", addr.port());
let server = async {
assert_ok!(srv.accept().await);
};
let client = async {
assert_ok!(TcpStream::connect(&addr).await);
};
join!(server, client);
}
#[tokio::test]
async fn connect_addr_host_str_port_tuple() {
let mut srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await);
let addr = assert_ok!(srv.local_addr());
let addr = ("localhost", addr.port());
let server = async {
assert_ok!(srv.accept().await);
};
let client = async {
assert_ok!(TcpStream::connect(&addr).await);
};
join!(server, client);
}
/*
* TODO: bring this back once TCP exposes HUP again
*
#[cfg(target_os = "linux")]
mod linux {
use tokio::net::{TcpListener, TcpStream};
use tokio::prelude::*;
use tokio_test::assert_ok;
use mio::unix::UnixReady;
use futures_util::future::poll_fn;
use std::io::Write;
use std::time::Duration;
use std::{net, thread};
#[tokio::test]
fn poll_hup() {
let addr = assert_ok!("127.0.0.1:0".parse());
let mut srv = assert_ok!(TcpListener::bind(&addr));
let addr = assert_ok!(srv.local_addr());
tokio::spawn(async move {
let (mut client, _) = assert_ok!(srv.accept().await);
assert_ok!(client.set_linger(Some(Duration::from_millis(0))));
assert_ok!(client.write_all(b"hello world").await);
// TODO: Drop?
});
/*
let t = thread::spawn(move || {
let mut client = assert_ok!(srv.accept()).0;
client.set_linger(Some(Duration::from_millis(0))).unwrap();
client.write(b"hello world").unwrap();
thread::sleep(Duration::from_millis(200));
});
*/
let mut stream = assert_ok!(TcpStream::connect(&addr).await);
// Poll for HUP before reading.
future::poll_fn(|| stream.poll_read_ready(UnixReady::hup().into()))
.wait()
.unwrap();
// Same for write half
future::poll_fn(|| stream.poll_write_ready())
.wait()
.unwrap();
let mut buf = vec![0; 11];
// Read the data
future::poll_fn(|| stream.poll_read(&mut buf))
.wait()
.unwrap();
assert_eq!(b"hello world", &buf[..]);
t.join().unwrap();
}
}
*/
tokio/tests/tcp_echo.rs
+41
View File
@@ -0,0 +1,41 @@
#![warn(rust_2018_idioms)]
use tokio::net::{TcpListener, TcpStream};
use tokio::prelude::*;
use tokio::sync::oneshot;
use tokio_test::assert_ok;
#[tokio::test]
async fn echo_server() {
const ITER: usize = 1024;
let (tx, rx) = oneshot::channel();
let mut srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await);
let addr = assert_ok!(srv.local_addr());
let msg = "foo bar baz";
tokio::spawn(async move {
let mut stream = assert_ok!(TcpStream::connect(&addr).await);
for _ in 0..ITER {
// write
assert_ok!(stream.write_all(msg.as_bytes()).await);
// read
let mut buf = [0; 11];
assert_ok!(stream.read_exact(&mut buf).await);
assert_eq!(&buf[..], msg.as_bytes());
}
assert_ok!(tx.send(()));
});
let (mut stream, _) = assert_ok!(srv.accept().await);
let (mut rd, mut wr) = stream.split();
let n = assert_ok!(rd.copy(&mut wr).await);
assert_eq!(n, (ITER * msg.len()) as u64);
assert_ok!(rx.await);
}
tokio/tests/tcp_peek.rs
+28
View File
@@ -0,0 +1,28 @@
#![warn(rust_2018_idioms)]
use tokio::io::AsyncReadExt;
use tokio::net::TcpStream;
use tokio_test::assert_ok;
use std::thread;
use std::{convert::TryInto, io::Write, net};
#[tokio::test]
async fn peek() {
let listener = net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let t = thread::spawn(move || assert_ok!(listener.accept()).0);
let left = net::TcpStream::connect(&addr).unwrap();
let mut right = t.join().unwrap();
right.write(&[1, 2, 3, 4]).unwrap();
let mut left: TcpStream = left.try_into().unwrap();
let mut buf = [0u8; 16];
let n = assert_ok!(left.peek(&mut buf).await);
assert_eq!([1, 2, 3, 4], buf[..n]);
let n = assert_ok!(left.read(&mut buf).await);
assert_eq!([1, 2, 3, 4], buf[..n]);
}
tokio/tests/tcp_shutdown.rs
+28
View File
@@ -0,0 +1,28 @@
#![warn(rust_2018_idioms)]
use tokio::io::AsyncWriteExt;
use tokio::net::{TcpListener, TcpStream};
use tokio::prelude::*;
use tokio_test::assert_ok;
#[tokio::test]
async fn shutdown() {
let mut srv = assert_ok!(TcpListener::bind("127.0.0.1:0").await);
let addr = assert_ok!(srv.local_addr());
tokio::spawn(async move {
let mut stream = assert_ok!(TcpStream::connect(&addr).await);
assert_ok!(AsyncWriteExt::shutdown(&mut stream).await);
let mut buf = [0; 1];
let n = assert_ok!(stream.read(&mut buf).await);
assert_eq!(n, 0);
});
let (mut stream, _) = assert_ok!(srv.accept().await);
let (mut rd, mut wr) = stream.split();
let n = assert_ok!(rd.copy(&mut wr).await);
assert_eq!(n, 0);
}
tokio/tests/tcp_split.rs
+1
View File
@@ -0,0 +1 @@
// TODO: write tests using TcpStream::split()
tokio/tests/udp.rs
+72
View File
@@ -0,0 +1,72 @@
#![warn(rust_2018_idioms)]
use tokio::net::UdpSocket;
#[tokio::test]
async fn send_recv() -> std::io::Result<()> {
let mut sender = UdpSocket::bind("127.0.0.1:0").await?;
let mut receiver = UdpSocket::bind("127.0.0.1:0").await?;
sender.connect(receiver.local_addr()?).await?;
receiver.connect(sender.local_addr()?).await?;
let message = b"hello!";
sender.send(message).await?;
let mut recv_buf = [0u8; 32];
let len = receiver.recv(&mut recv_buf[..]).await?;
assert_eq!(&recv_buf[..len], message);
Ok(())
}
#[tokio::test]
async fn send_to_recv_from() -> std::io::Result<()> {
let mut sender = UdpSocket::bind("127.0.0.1:0").await?;
let mut receiver = UdpSocket::bind("127.0.0.1:0").await?;
let message = b"hello!";
let receiver_addr = receiver.local_addr()?;
sender.send_to(message, &receiver_addr).await?;
let mut recv_buf = [0u8; 32];
let (len, addr) = receiver.recv_from(&mut recv_buf[..]).await?;
assert_eq!(&recv_buf[..len], message);
assert_eq!(addr, sender.local_addr()?);
Ok(())
}
#[tokio::test]
async fn split() -> std::io::Result<()> {
let socket = UdpSocket::bind("127.0.0.1:0").await?;
let (mut r, mut s) = socket.split();
let msg = b"hello";
let addr = s.as_ref().local_addr()?;
tokio::spawn(async move {
s.send_to(msg, &addr).await.unwrap();
});
let mut recv_buf = [0u8; 32];
let (len, _) = r.recv_from(&mut recv_buf[..]).await?;
assert_eq!(&recv_buf[..len], msg);
Ok(())
}
#[tokio::test]
async fn reunite() -> std::io::Result<()> {
let socket = UdpSocket::bind("127.0.0.1:0").await?;
let (s, r) = socket.split();
assert!(s.reunite(r).is_ok());
Ok(())
}
#[tokio::test]
async fn reunite_error() -> std::io::Result<()> {
let socket = UdpSocket::bind("127.0.0.1:0").await?;
let socket1 = UdpSocket::bind("127.0.0.1:0").await?;
let (s, _) = socket.split();
let (_, r1) = socket1.split();
assert!(s.reunite(r1).is_err());
Ok(())
}
tokio/tests/uds_cred.rs
+29
View File
@@ -0,0 +1,29 @@
#![cfg(unix)]
#![cfg(not(target_os = "dragonfly"))]
use tokio::net::UnixStream;
use libc::getegid;
use libc::geteuid;
#[tokio::test]
#[cfg_attr(
target_os = "freebsd",
ignore = "Requires FreeBSD 12.0 or later. https://bugs.freebsd.org/bugzilla/show_bug.cgi?id=176419"
)]
#[cfg_attr(
target_os = "netbsd",
ignore = "NetBSD does not support getpeereid() for sockets created by socketpair()"
)]
async fn test_socket_pair() {
let (a, b) = UnixStream::pair().unwrap();
let cred_a = a.peer_cred().unwrap();
let cred_b = b.peer_cred().unwrap();
assert_eq!(cred_a, cred_b);
let uid = unsafe { geteuid() };
let gid = unsafe { getegid() };
assert_eq!(cred_a.uid, uid);
assert_eq!(cred_a.gid, gid);
}
tokio/tests/uds_datagram.rs
+69
View File
@@ -0,0 +1,69 @@
#![cfg(unix)]
#![warn(rust_2018_idioms)]
use tokio::net::unix::*;
use std::io;
// struct StringDatagramCodec;
// /// A codec to decode datagrams from a unix domain socket as utf-8 text messages.
// impl Encoder for StringDatagramCodec {
// type Item = String;
// type Error = io::Error;
// fn encode(&mut self, item: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
// dst.extend_from_slice(&item.into_bytes());
// Ok(())
// }
// }
// /// A codec to decode datagrams from a unix domain socket as utf-8 text messages.
// impl Decoder for StringDatagramCodec {
// type Item = String;
// type Error = io::Error;
// fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
// let decoded = str::from_utf8(buf)
// .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
// .to_string();
// Ok(Some(decoded))
// }
// }
async fn echo_server(mut socket: UnixDatagram) -> io::Result<()> {
let mut recv_buf = vec![0u8; 1024];
loop {
let (len, peer_addr) = socket.recv_from(&mut recv_buf[..]).await?;
if let Some(path) = peer_addr.as_pathname() {
socket.send_to(&recv_buf[..len], path).await?;
}
}
}
#[tokio::test]
async fn echo() -> io::Result<()> {
let dir = tempfile::tempdir().unwrap();
let server_path = dir.path().join("server.sock");
let client_path = dir.path().join("client.sock");
let server_socket = UnixDatagram::bind(server_path.clone())?;
tokio::spawn(async move {
if let Err(e) = echo_server(server_socket).await {
eprintln!("Error in echo server: {}", e);
}
});
{
let mut socket = UnixDatagram::bind(&client_path).unwrap();
socket.connect(server_path)?;
socket.send(b"ECHO").await?;
let mut recv_buf = [0u8; 16];
let len = socket.recv(&mut recv_buf[..]).await?;
assert_eq!(&recv_buf[..len], b"ECHO");
}
Ok(())
}
tokio/tests/uds_split.rs
+42
View File
@@ -0,0 +1,42 @@
#![cfg(unix)]
#![deny(warnings, rust_2018_idioms)]
use tokio::net::UnixStream;
use tokio::prelude::*;
/// Checks that `UnixStream` can be split into a read half and a write half using
/// `UnixStream::split` and `UnixStream::split_mut`.
///
/// Verifies that the implementation of `AsyncWrite::poll_shutdown` shutdowns the stream for
/// writing by reading to the end of stream on the other side of the connection.
#[tokio::test]
async fn split() -> std::io::Result<()> {
let (mut a, mut b) = UnixStream::pair()?;
let (mut a_read, mut a_write) = a.split();
let (mut b_read, mut b_write) = b.split();
let (a_response, b_response) = futures::future::try_join(
send_recv_all(&mut a_read, &mut a_write, b"A"),
send_recv_all(&mut b_read, &mut b_write, b"B"),
)
.await?;
assert_eq!(a_response, b"B");
assert_eq!(b_response, b"A");
Ok(())
}
async fn send_recv_all(
read: &mut (dyn AsyncRead + Unpin),
write: &mut (dyn AsyncWrite + Unpin),
input: &[u8],
) -> std::io::Result<Vec<u8>> {
write.write_all(input).await?;
write.shutdown().await?;
let mut output = Vec::new();
read.read_to_end(&mut output).await?;
Ok(output)
}
tokio/tests/uds_stream.rs
+34
View File
@@ -0,0 +1,34 @@
#![cfg(unix)]
#![warn(rust_2018_idioms)]
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::unix::*;
use futures::future::try_join;
#[tokio::test]
async fn accept_read_write() -> std::io::Result<()> {
let dir = tempfile::Builder::new()
.prefix("tokio-uds-tests")
.tempdir()
.unwrap();
let sock_path = dir.path().join("connect.sock");
let mut listener = UnixListener::bind(&sock_path)?;
let accept = listener.accept();
let connect = UnixStream::connect(&sock_path);
let ((mut server, _), mut client) = try_join(accept, connect).await?;
// Write to the client. TODO: Switch to write_all.
let write_len = client.write(b"hello").await?;
assert_eq!(write_len, 5);
drop(client);
// Read from the server. TODO: Switch to read_to_end.
let mut buf = [0u8; 5];
server.read_exact(&mut buf).await?;
assert_eq!(&buf, b"hello");
let len = server.read(&mut buf).await?;
assert_eq!(len, 0);
Ok(())
}