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

runtime: fix memory leak/growth when creating many runtimes (#3564)

Browse Source
This commit is contained in:
Ivan Petkov
2021-03-16 19:31:46 +01:00
committed by GitHub
parent e6103d6661
commit e4f76688a0
18 changed files with 437 additions and 285 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/ci.yml
+11 -2
View File
@@ -82,14 +82,23 @@ jobs:
sudo apt-get install -y valgrind
# Compile tests
- name: cargo build
- name: cargo build test-mem
run: cargo build --features rt-net --bin test-mem
working-directory: tests-integration
# Run with valgrind
- name: Run valgrind
- name: Run valgrind test-mem
run: valgrind --leak-check=full --show-leak-kinds=all ./target/debug/test-mem
# Compile tests
- name: cargo build test-process-signal
run: cargo build --features rt-process-signal --bin test-process-signal
working-directory: tests-integration
# Run with valgrind
- name: Run valgrind test-process-signal
run: valgrind --leak-check=full --show-leak-kinds=all ./target/debug/test-process-signal
test-unstable:
name: test tokio full --unstable
runs-on: ${{ matrix.os }}
tests-integration/Cargo.toml
+6
View File
@@ -12,9 +12,15 @@ name = "test-cat"
name = "test-mem"
required-features = ["rt-net"]
[[bin]]
name = "test-process-signal"
required-features = ["rt-process-signal"]
[features]
# For mem check
rt-net = ["tokio/rt", "tokio/rt-multi-thread", "tokio/net"]
# For test-process-signal
rt-process-signal = ["rt", "tokio/process", "tokio/signal"]
full = [
"macros",
tests-integration/src/bin/test-process-signal.rs
+11
View File
@@ -0,0 +1,11 @@
// https://github.com/tokio-rs/tokio/issues/3550
fn main() {
for _ in 0..1000 {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
drop(rt);
}
}
tokio/src/io/driver/interest.rs
+1 -1
View File
@@ -1,4 +1,4 @@
#![cfg_attr(not(feature = "net"), allow(unreachable_pub))]
#![cfg_attr(not(feature = "net"), allow(dead_code, unreachable_pub))]
use crate::io::driver::Ready;
tokio/src/io/driver/registration.rs
+2
View File
@@ -1,3 +1,5 @@
#![cfg_attr(not(feature = "net"), allow(dead_code))]
use crate::io::driver::{Direction, Handle, Interest, ReadyEvent, ScheduledIo};
use crate::util::slab;
tokio/src/io/poll_evented.rs
+5
View File
@@ -121,6 +121,11 @@ impl<E: Source> PollEvented<E> {
}
/// Returns a reference to the registration
#[cfg(any(
feature = "net",
all(unix, feature = "process"),
all(unix, feature = "signal"),
))]
pub(crate) fn registration(&self) -> &Registration {
&self.registration
}
tokio/src/process/unix/driver.rs
+18 -42
View File
@@ -6,8 +6,8 @@ use crate::park::Park;
use crate::process::unix::orphan::ReapOrphanQueue;
use crate::process::unix::GlobalOrphanQueue;
use crate::signal::unix::driver::Driver as SignalDriver;
use crate::signal::unix::{signal_with_handle, InternalStream, Signal, SignalKind};
use crate::sync::mpsc::error::TryRecvError;
use crate::signal::unix::{signal_with_handle, SignalKind};
use crate::sync::watch;
use std::io;
use std::time::Duration;
@@ -16,7 +16,7 @@ use std::time::Duration;
#[derive(Debug)]
pub(crate) struct Driver {
park: SignalDriver,
inner: CoreDriver<Signal, GlobalOrphanQueue>,
inner: CoreDriver<watch::Receiver<()>, GlobalOrphanQueue>,
}
#[derive(Debug)]
@@ -25,27 +25,25 @@ struct CoreDriver<S, Q> {
orphan_queue: Q,
}
trait HasChanged {
fn has_changed(&mut self) -> bool;
}
impl<T> HasChanged for watch::Receiver<T> {
fn has_changed(&mut self) -> bool {
self.try_has_changed().and_then(Result::ok).is_some()
}
}
// ===== impl CoreDriver =====
impl<S, Q> CoreDriver<S, Q>
where
S: InternalStream,
S: HasChanged,
Q: ReapOrphanQueue,
{
fn got_signal(&mut self) -> bool {
match self.sigchild.try_recv() {
Ok(()) => true,
Err(TryRecvError::Empty) => false,
Err(TryRecvError::Closed) => panic!("signal was deregistered"),
}
}
fn process(&mut self) {
if self.got_signal() {
// Drain all notifications which may have been buffered
// so we can try to reap all orphans in one batch
while self.got_signal() {}
if self.sigchild.has_changed() {
self.orphan_queue.reap_orphans();
}
}
@@ -97,8 +95,6 @@ impl Park for Driver {
mod test {
use super::*;
use crate::process::unix::orphan::test::MockQueue;
use crate::sync::mpsc::error::TryRecvError;
use std::task::{Context, Poll};
struct MockStream {
total_try_recv: usize,
@@ -114,17 +110,10 @@ mod test {
}
}
impl InternalStream for MockStream {
fn poll_recv(&mut self, _cx: &mut Context<'_>) -> Poll<Option<()>> {
unimplemented!();
}
fn try_recv(&mut self) -> Result<(), TryRecvError> {
impl HasChanged for MockStream {
fn has_changed(&mut self) -> bool {
self.total_try_recv += 1;
match self.values.remove(0) {
Some(()) => Ok(()),
None => Err(TryRecvError::Empty),
}
self.values.remove(0).is_some()
}
}
@@ -140,17 +129,4 @@ mod test {
assert_eq!(1, driver.sigchild.total_try_recv);
assert_eq!(0, driver.orphan_queue.total_reaps.get());
}
#[test]
fn coalesce_signals_before_reaping() {
let mut driver = CoreDriver {
sigchild: MockStream::new(vec![Some(()), Some(()), None]),
orphan_queue: MockQueue::<()>::new(),
};
driver.process();
assert_eq!(3, driver.sigchild.total_try_recv);
assert_eq!(1, driver.orphan_queue.total_reaps.get());
}
}
tokio/src/process/unix/reap.rs
+6 -11
View File
@@ -15,7 +15,7 @@ use std::task::Poll;
#[derive(Debug)]
pub(crate) struct Reaper<W, Q, S>
where
W: Wait + Unpin,
W: Wait,
Q: OrphanQueue<W>,
{
inner: Option<W>,
@@ -25,7 +25,7 @@ where
impl<W, Q, S> Deref for Reaper<W, Q, S>
where
W: Wait + Unpin,
W: Wait,
Q: OrphanQueue<W>,
{
type Target = W;
@@ -37,7 +37,7 @@ where
impl<W, Q, S> Reaper<W, Q, S>
where
W: Wait + Unpin,
W: Wait,
Q: OrphanQueue<W>,
{
pub(crate) fn new(inner: W, orphan_queue: Q, signal: S) -> Self {
@@ -61,7 +61,7 @@ impl<W, Q, S> Future for Reaper<W, Q, S>
where
W: Wait + Unpin,
Q: OrphanQueue<W> + Unpin,
S: InternalStream,
S: InternalStream + Unpin,
{
type Output = io::Result<ExitStatus>;
@@ -106,7 +106,7 @@ where
impl<W, Q, S> Kill for Reaper<W, Q, S>
where
W: Kill + Wait + Unpin,
W: Kill + Wait,
Q: OrphanQueue<W>,
{
fn kill(&mut self) -> io::Result<()> {
@@ -116,7 +116,7 @@ where
impl<W, Q, S> Drop for Reaper<W, Q, S>
where
W: Wait + Unpin,
W: Wait,
Q: OrphanQueue<W>,
{
fn drop(&mut self) {
@@ -134,7 +134,6 @@ mod test {
use super::*;
use crate::process::unix::orphan::test::MockQueue;
use crate::sync::mpsc::error::TryRecvError;
use futures::future::FutureExt;
use std::os::unix::process::ExitStatusExt;
use std::process::ExitStatus;
@@ -206,10 +205,6 @@ mod test {
None => Poll::Pending,
}
}
fn try_recv(&mut self) -> Result<(), TryRecvError> {
unimplemented!();
}
}
#[test]
tokio/src/signal/mod.rs
+40
View File
@@ -42,6 +42,8 @@
//! }
//! # }
//! ```
use crate::sync::watch::Receiver;
use std::task::{Context, Poll};
mod ctrl_c;
pub use ctrl_c::ctrl_c;
@@ -58,3 +60,41 @@ mod os {
pub mod unix;
pub mod windows;
mod reusable_box;
use self::reusable_box::ReusableBoxFuture;
#[derive(Debug)]
struct RxFuture {
inner: ReusableBoxFuture<Receiver<()>>,
}
async fn make_future(mut rx: Receiver<()>) -> Receiver<()> {
match rx.changed().await {
Ok(()) => rx,
Err(_) => panic!("signal sender went away"),
}
}
impl RxFuture {
fn new(rx: Receiver<()>) -> Self {
Self {
inner: ReusableBoxFuture::new(make_future(rx)),
}
}
async fn recv(&mut self) -> Option<()> {
use crate::future::poll_fn;
poll_fn(|cx| self.poll_recv(cx)).await
}
fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
match self.inner.poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(rx) => {
self.inner.set(make_future(rx));
Poll::Ready(Some(()))
}
}
}
}
tokio/src/signal/registry.rs
+42 -89
View File
@@ -2,22 +2,32 @@
use crate::signal::os::{OsExtraData, OsStorage};
use crate::sync::mpsc::Sender;
use crate::sync::watch;
use once_cell::sync::Lazy;
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)]
#[derive(Debug)]
pub(crate) struct EventInfo {
pending: AtomicBool,
recipients: Mutex<Vec<Sender<()>>>,
tx: watch::Sender<()>,
}
impl Default for EventInfo {
fn default() -> Self {
let (tx, _rx) = watch::channel(());
Self {
pending: AtomicBool::new(false),
tx,
}
}
}
/// An interface for retrieving the `EventInfo` for a particular eventId.
@@ -67,14 +77,12 @@ impl<S> Registry<S> {
impl<S: Storage> Registry<S> {
/// Registers a new listener for `event_id`.
fn register_listener(&self, event_id: EventId, listener: Sender<()>) {
fn register_listener(&self, event_id: EventId) -> watch::Receiver<()> {
self.storage
.event_info(event_id)
.unwrap_or_else(|| panic!("invalid event_id: {}", event_id))
.recipients
.lock()
.unwrap()
.push(listener);
.tx
.subscribe()
}
/// Marks `event_id` as having been delivered, without broadcasting it to
@@ -89,8 +97,6 @@ impl<S: Storage> Registry<S> {
///
/// Returns `true` if an event was delivered to at least one listener.
fn broadcast(&self) -> bool {
use crate::sync::mpsc::error::TrySendError;
let mut did_notify = false;
self.storage.for_each(|event_info| {
// Any signal of this kind arrived since we checked last?
@@ -98,23 +104,9 @@ impl<S: Storage> Registry<S> {
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(TrySendError::Closed(..)) => {
recipients.swap_remove(i);
}
// Channel is full, ignore the error since the
// receiver has already been woken up
Err(_) => {}
}
// Ignore errors if there are no listeners
if event_info.tx.send(()).is_ok() {
did_notify = true;
}
});
@@ -137,8 +129,8 @@ impl ops::Deref for Globals {
impl Globals {
/// Registers a new listener for `event_id`.
pub(crate) fn register_listener(&self, event_id: EventId, listener: Sender<()>) {
self.registry.register_listener(event_id, listener);
pub(crate) fn register_listener(&self, event_id: EventId) -> watch::Receiver<()> {
self.registry.register_listener(event_id)
}
/// Marks `event_id` as having been delivered, without broadcasting it to
@@ -179,7 +171,7 @@ where
mod tests {
use super::*;
use crate::runtime::{self, Runtime};
use crate::sync::{mpsc, oneshot};
use crate::sync::{oneshot, watch};
use futures::future;
@@ -193,13 +185,9 @@ mod tests {
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 first = registry.register_listener(0);
let second = registry.register_listener(1);
let third = registry.register_listener(2);
let (fire, wait) = oneshot::channel();
@@ -213,6 +201,9 @@ mod tests {
registry.record_event(1);
registry.broadcast();
// Yield so the previous broadcast can get received
crate::time::sleep(std::time::Duration::from_millis(10)).await;
// Send subsequent signal
registry.record_event(0);
registry.broadcast();
@@ -221,7 +212,7 @@ mod tests {
});
let _ = fire.send(());
let all = future::join3(collect(first_rx), collect(second_rx), collect(third_rx));
let all = future::join3(collect(first), collect(second), collect(third));
let (first_results, second_results, third_results) = all.await;
assert_eq!(2, first_results.len());
@@ -235,8 +226,7 @@ mod tests {
fn register_panics_on_invalid_input() {
let registry = Registry::new(vec![EventInfo::default()]);
let (tx, _) = mpsc::channel(1);
registry.register_listener(1, tx);
registry.register_listener(1);
}
#[test]
@@ -245,74 +235,37 @@ mod tests {
registry.record_event(42);
}
#[test]
fn broadcast_cleans_up_disconnected_listeners() {
let 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 = collect(third_rx).await;
assert_eq!(1, results.len());
});
}
#[test]
fn broadcast_returns_if_at_least_one_event_fired() {
let registry = Registry::new(vec![EventInfo::default()]);
let registry = Registry::new(vec![EventInfo::default(), 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);
let first = registry.register_listener(0);
let second = registry.register_listener(1);
registry.record_event(0);
assert_eq!(true, registry.broadcast());
drop(first_rx);
drop(first);
registry.record_event(0);
assert_eq!(false, registry.broadcast());
drop(second_rx);
drop(second);
}
fn rt() -> Runtime {
runtime::Builder::new_current_thread().build().unwrap()
runtime::Builder::new_current_thread()
.enable_time()
.build()
.unwrap()
}
async fn collect(mut rx: crate::sync::mpsc::Receiver<()>) -> Vec<()> {
async fn collect(mut rx: watch::Receiver<()>) -> Vec<()> {
let mut ret = vec![];
while let Some(v) = rx.recv().await {
while let Ok(v) = rx.changed().await {
ret.push(v);
}
tokio/src/signal/reusable_box.rs
+227
View File
@@ -0,0 +1,227 @@
use std::alloc::Layout;
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::pin::Pin;
use std::ptr::{self, NonNull};
use std::task::{Context, Poll};
use std::{fmt, panic};
/// A reusable `Pin<Box<dyn Future<Output = T> + Send>>`.
///
/// This type lets you replace the future stored in the box without
/// reallocating when the size and alignment permits this.
pub(crate) struct ReusableBoxFuture<T> {
boxed: NonNull<dyn Future<Output = T> + Send>,
}
impl<T> ReusableBoxFuture<T> {
/// Create a new `ReusableBoxFuture<T>` containing the provided future.
pub(crate) fn new<F>(future: F) -> Self
where
F: Future<Output = T> + Send + 'static,
{
let boxed: Box<dyn Future<Output = T> + Send> = Box::new(future);
let boxed = Box::into_raw(boxed);
// SAFETY: Box::into_raw does not return null pointers.
let boxed = unsafe { NonNull::new_unchecked(boxed) };
Self { boxed }
}
/// Replace the future currently stored in this box.
///
/// This reallocates if and only if the layout of the provided future is
/// different from the layout of the currently stored future.
pub(crate) fn set<F>(&mut self, future: F)
where
F: Future<Output = T> + Send + 'static,
{
if let Err(future) = self.try_set(future) {
*self = Self::new(future);
}
}
/// Replace the future currently stored in this box.
///
/// This function never reallocates, but returns an error if the provided
/// future has a different size or alignment from the currently stored
/// future.
pub(crate) fn try_set<F>(&mut self, future: F) -> Result<(), F>
where
F: Future<Output = T> + Send + 'static,
{
// SAFETY: The pointer is not dangling.
let self_layout = {
let dyn_future: &(dyn Future<Output = T> + Send) = unsafe { self.boxed.as_ref() };
Layout::for_value(dyn_future)
};
if Layout::new::<F>() == self_layout {
// SAFETY: We just checked that the layout of F is correct.
unsafe {
self.set_same_layout(future);
}
Ok(())
} else {
Err(future)
}
}
/// Set the current future.
///
/// # Safety
///
/// This function requires that the layout of the provided future is the
/// same as `self.layout`.
unsafe fn set_same_layout<F>(&mut self, future: F)
where
F: Future<Output = T> + Send + 'static,
{
// Drop the existing future, catching any panics.
let result = panic::catch_unwind(AssertUnwindSafe(|| {
ptr::drop_in_place(self.boxed.as_ptr());
}));
// Overwrite the future behind the pointer. This is safe because the
// allocation was allocated with the same size and alignment as the type F.
let self_ptr: *mut F = self.boxed.as_ptr() as *mut F;
ptr::write(self_ptr, future);
// Update the vtable of self.boxed. The pointer is not null because we
// just got it from self.boxed, which is not null.
self.boxed = NonNull::new_unchecked(self_ptr);
// If the old future's destructor panicked, resume unwinding.
match result {
Ok(()) => {}
Err(payload) => {
panic::resume_unwind(payload);
}
}
}
/// Get a pinned reference to the underlying future.
pub(crate) fn get_pin(&mut self) -> Pin<&mut (dyn Future<Output = T> + Send)> {
// SAFETY: The user of this box cannot move the box, and we do not move it
// either.
unsafe { Pin::new_unchecked(self.boxed.as_mut()) }
}
/// Poll the future stored inside this box.
pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<T> {
self.get_pin().poll(cx)
}
}
impl<T> Future for ReusableBoxFuture<T> {
type Output = T;
/// Poll the future stored inside this box.
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
Pin::into_inner(self).get_pin().poll(cx)
}
}
// The future stored inside ReusableBoxFuture<T> must be Send.
unsafe impl<T> Send for ReusableBoxFuture<T> {}
// The only method called on self.boxed is poll, which takes &mut self, so this
// struct being Sync does not permit any invalid access to the Future, even if
// the future is not Sync.
unsafe impl<T> Sync for ReusableBoxFuture<T> {}
// Just like a Pin<Box<dyn Future>> is always Unpin, so is this type.
impl<T> Unpin for ReusableBoxFuture<T> {}
impl<T> Drop for ReusableBoxFuture<T> {
fn drop(&mut self) {
unsafe {
drop(Box::from_raw(self.boxed.as_ptr()));
}
}
}
impl<T> fmt::Debug for ReusableBoxFuture<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ReusableBoxFuture").finish()
}
}
#[cfg(test)]
mod test {
use super::ReusableBoxFuture;
use futures::future::FutureExt;
use std::alloc::Layout;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
#[test]
fn test_different_futures() {
let fut = async move { 10 };
// Not zero sized!
assert_eq!(Layout::for_value(&fut).size(), 1);
let mut b = ReusableBoxFuture::new(fut);
assert_eq!(b.get_pin().now_or_never(), Some(10));
b.try_set(async move { 20 })
.unwrap_or_else(|_| panic!("incorrect size"));
assert_eq!(b.get_pin().now_or_never(), Some(20));
b.try_set(async move { 30 })
.unwrap_or_else(|_| panic!("incorrect size"));
assert_eq!(b.get_pin().now_or_never(), Some(30));
}
#[test]
fn test_different_sizes() {
let fut1 = async move { 10 };
let val = [0u32; 1000];
let fut2 = async move { val[0] };
let fut3 = ZeroSizedFuture {};
assert_eq!(Layout::for_value(&fut1).size(), 1);
assert_eq!(Layout::for_value(&fut2).size(), 4004);
assert_eq!(Layout::for_value(&fut3).size(), 0);
let mut b = ReusableBoxFuture::new(fut1);
assert_eq!(b.get_pin().now_or_never(), Some(10));
b.set(fut2);
assert_eq!(b.get_pin().now_or_never(), Some(0));
b.set(fut3);
assert_eq!(b.get_pin().now_or_never(), Some(5));
}
struct ZeroSizedFuture {}
impl Future for ZeroSizedFuture {
type Output = u32;
fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<u32> {
Poll::Ready(5)
}
}
#[test]
fn test_zero_sized() {
let fut = ZeroSizedFuture {};
// Zero sized!
assert_eq!(Layout::for_value(&fut).size(), 0);
let mut b = ReusableBoxFuture::new(fut);
assert_eq!(b.get_pin().now_or_never(), Some(5));
assert_eq!(b.get_pin().now_or_never(), Some(5));
b.try_set(ZeroSizedFuture {})
.unwrap_or_else(|_| panic!("incorrect size"));
assert_eq!(b.get_pin().now_or_never(), Some(5));
assert_eq!(b.get_pin().now_or_never(), Some(5));
}
}
tokio/src/signal/unix.rs
+25 -31
View File
@@ -6,8 +6,8 @@
#![cfg(unix)]
use crate::signal::registry::{globals, EventId, EventInfo, Globals, Init, Storage};
use crate::sync::mpsc::error::TryRecvError;
use crate::sync::mpsc::{channel, Receiver};
use crate::signal::RxFuture;
use crate::sync::watch;
use libc::c_int;
use mio::net::UnixStream;
@@ -222,7 +222,8 @@ fn action(globals: Pin<&'static Globals>, signal: c_int) {
///
/// 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, handle: Handle) -> io::Result<()> {
fn signal_enable(signal: SignalKind, handle: Handle) -> io::Result<()> {
let signal = signal.0;
if signal < 0 || signal_hook_registry::FORBIDDEN.contains(&signal) {
return Err(Error::new(
ErrorKind::Other,
@@ -325,7 +326,7 @@ fn signal_enable(signal: c_int, handle: Handle) -> io::Result<()> {
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct Signal {
rx: Receiver<()>,
inner: RxFuture,
}
/// Creates a new stream which will receive notifications when the current
@@ -351,21 +352,21 @@ pub struct Signal {
/// * If the signal is one of
/// [`signal_hook::FORBIDDEN`](fn@signal_hook_registry::register#panics)
pub fn signal(kind: SignalKind) -> io::Result<Signal> {
signal_with_handle(kind, Handle::current())
let rx = signal_with_handle(kind, Handle::current())?;
Ok(Signal {
inner: RxFuture::new(rx),
})
}
pub(crate) fn signal_with_handle(kind: SignalKind, handle: Handle) -> io::Result<Signal> {
let signal = kind.0;
pub(crate) fn signal_with_handle(
kind: SignalKind,
handle: Handle,
) -> io::Result<watch::Receiver<()>> {
// Turn the signal delivery on once we are ready for it
signal_enable(signal, handle)?;
signal_enable(kind, handle)?;
// 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 { rx })
Ok(globals().register_listener(kind.0 as EventId))
}
impl Signal {
@@ -393,8 +394,7 @@ impl Signal {
/// }
/// ```
pub async fn recv(&mut self) -> Option<()> {
use crate::future::poll_fn;
poll_fn(|cx| self.poll_recv(cx)).await
self.inner.recv().await
}
/// Polls to receive the next signal notification event, outside of an
@@ -432,29 +432,19 @@ impl Signal {
/// }
/// ```
pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
self.rx.poll_recv(cx)
}
/// Try to receive a signal notification without blocking or registering a waker.
pub(crate) fn try_recv(&mut self) -> Result<(), TryRecvError> {
self.rx.try_recv()
self.inner.poll_recv(cx)
}
}
// Work around for abstracting streams internally
pub(crate) trait InternalStream: Unpin {
pub(crate) trait InternalStream {
fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>>;
fn try_recv(&mut self) -> Result<(), TryRecvError>;
}
impl InternalStream for Signal {
fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
self.poll_recv(cx)
}
fn try_recv(&mut self) -> Result<(), TryRecvError> {
self.try_recv()
}
}
pub(crate) fn ctrl_c() -> io::Result<Signal> {
@@ -467,11 +457,15 @@ mod tests {
#[test]
fn signal_enable_error_on_invalid_input() {
signal_enable(-1, Handle::default()).unwrap_err();
signal_enable(SignalKind::from_raw(-1), Handle::default()).unwrap_err();
}
#[test]
fn signal_enable_error_on_forbidden_input() {
signal_enable(signal_hook_registry::FORBIDDEN[0], Handle::default()).unwrap_err();
signal_enable(
SignalKind::from_raw(signal_hook_registry::FORBIDDEN[0]),
Handle::default(),
)
.unwrap_err();
}
}
tokio/src/signal/windows.rs
+10 -14
View File
@@ -8,7 +8,7 @@
#![cfg(windows)]
use crate::signal::registry::{globals, EventId, EventInfo, Init, Storage};
use crate::sync::mpsc::{channel, Receiver};
use crate::signal::RxFuture;
use std::convert::TryFrom;
use std::io;
@@ -76,22 +76,18 @@ impl Init for OsExtraData {
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub(crate) struct Event {
rx: Receiver<()>,
inner: RxFuture,
}
impl Event {
fn new(signum: DWORD) -> io::Result<Self> {
global_init()?;
let (tx, rx) = channel(1);
globals().register_listener(signum as EventId, tx);
let rx = globals().register_listener(signum as EventId);
Ok(Event { rx })
}
pub(crate) async fn recv(&mut self) -> Option<()> {
use crate::future::poll_fn;
poll_fn(|cx| self.rx.poll_recv(cx)).await
Ok(Self {
inner: RxFuture::new(rx),
})
}
}
@@ -195,7 +191,7 @@ impl CtrlC {
/// }
/// ```
pub async fn recv(&mut self) -> Option<()> {
self.inner.recv().await
self.inner.inner.recv().await
}
/// Polls to receive the next signal notification event, outside of an
@@ -227,7 +223,7 @@ impl CtrlC {
/// }
/// ```
pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
self.inner.rx.poll_recv(cx)
self.inner.inner.poll_recv(cx)
}
}
@@ -267,7 +263,7 @@ impl CtrlBreak {
/// }
/// ```
pub async fn recv(&mut self) -> Option<()> {
self.inner.recv().await
self.inner.inner.recv().await
}
/// Polls to receive the next signal notification event, outside of an
@@ -299,7 +295,7 @@ impl CtrlBreak {
/// }
/// ```
pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
self.inner.rx.poll_recv(cx)
self.inner.inner.poll_recv(cx)
}
}
tokio/src/sync/mod.rs
+6 -12
View File
@@ -462,10 +462,8 @@ cfg_sync! {
}
cfg_not_sync! {
#[cfg(any(feature = "fs", feature = "signal", all(unix, feature = "process")))]
pub(crate) mod batch_semaphore;
cfg_fs! {
pub(crate) mod batch_semaphore;
mod mutex;
pub(crate) use mutex::Mutex;
}
@@ -473,20 +471,16 @@ cfg_not_sync! {
#[cfg(any(feature = "rt", feature = "signal", all(unix, feature = "process")))]
pub(crate) mod notify;
#[cfg(any(feature = "rt", all(windows, feature = "process")))]
pub(crate) mod oneshot;
cfg_atomic_waker_impl! {
mod task;
pub(crate) use task::AtomicWaker;
}
#[cfg(any(
feature = "rt",
feature = "process",
feature = "signal"))]
pub(crate) mod oneshot;
cfg_signal_internal! {
pub(crate) mod mpsc;
}
#[cfg(any(feature = "signal", all(unix, feature = "process")))]
pub(crate) mod watch;
}
/// Unit tests
tokio/src/sync/mpsc/bounded.rs
-20
View File
@@ -1,8 +1,5 @@
use crate::sync::batch_semaphore::{self as semaphore, TryAcquireError};
use crate::sync::mpsc::chan;
#[cfg(unix)]
#[cfg(any(feature = "signal", feature = "process"))]
use crate::sync::mpsc::error::TryRecvError;
use crate::sync::mpsc::error::{SendError, TrySendError};
cfg_time! {
@@ -224,23 +221,6 @@ impl<T> Receiver<T> {
crate::future::block_on(self.recv())
}
/// Attempts to return a pending value on this receiver without blocking.
///
/// This method will never block the caller in order to wait for data to
/// become available. Instead, this will always return immediately with
/// a possible option of pending data on the channel.
///
/// This is useful for a flavor of "optimistic check" before deciding to
/// block on a receiver.
///
/// Compared with recv, this function has two failure cases instead of
/// one (one for disconnection, one for an empty buffer).
#[cfg(unix)]
#[cfg(any(feature = "signal", feature = "process"))]
pub(crate) fn try_recv(&mut self) -> Result<T, TryRecvError> {
self.chan.try_recv()
}
/// Closes the receiving half of a channel without dropping it.
///
/// This prevents any further messages from being sent on the channel while
tokio/src/sync/mpsc/chan.rs
-24
View File
@@ -265,30 +265,6 @@ impl<T, S: Semaphore> Rx<T, S> {
}
}
feature! {
#![all(unix, any(feature = "signal", feature = "process"))]
use crate::sync::mpsc::error::TryRecvError;
impl<T, S: Semaphore> Rx<T, S> {
/// Receives the next value without blocking
pub(crate) fn try_recv(&mut self) -> Result<T, TryRecvError> {
use super::block::Read::*;
self.inner.rx_fields.with_mut(|rx_fields_ptr| {
let rx_fields = unsafe { &mut *rx_fields_ptr };
match rx_fields.list.pop(&self.inner.tx) {
Some(Value(value)) => {
self.inner.semaphore.add_permit();
Ok(value)
}
Some(Closed) => Err(TryRecvError::Closed),
None => Err(TryRecvError::Empty),
}
})
}
}
}
impl<T, S: Semaphore> Drop for Rx<T, S> {
fn drop(&mut self) {
use super::block::Read::Value;
tokio/src/sync/mpsc/error.rs
-33
View File
@@ -65,39 +65,6 @@ impl fmt::Display for RecvError {
impl Error for RecvError {}
// ===== TryRecvError =====
feature! {
#![all(unix, any(feature = "signal", feature = "process"))]
/// This enumeration is the list of the possible reasons that try_recv
/// could not return data when called.
#[derive(Debug, PartialEq)]
pub(crate) enum TryRecvError {
/// This channel is currently empty, but the Sender(s) have not yet
/// disconnected, so data may yet become available.
Empty,
/// The channel's sending half has been closed, and there will
/// never be any more data received on it.
Closed,
}
impl fmt::Display for TryRecvError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
fmt,
"{}",
match self {
TryRecvError::Empty => "channel empty",
TryRecvError::Closed => "channel closed",
}
)
}
}
impl Error for TryRecvError {}
}
cfg_time! {
// ===== SendTimeoutError =====
tokio/src/sync/watch.rs
+27 -6
View File
@@ -1,3 +1,5 @@
#![cfg_attr(not(feature = "sync"), allow(dead_code, unreachable_pub))]
//! A single-producer, multi-consumer channel that only retains the *last* sent
//! value.
//!
@@ -51,7 +53,7 @@
//! [`Sender::is_closed`]: crate::sync::watch::Sender::is_closed
//! [`Sender::closed`]: crate::sync::watch::Sender::closed
use crate::sync::Notify;
use crate::sync::notify::Notify;
use crate::loom::sync::atomic::AtomicUsize;
use crate::loom::sync::atomic::Ordering::{Relaxed, SeqCst};
@@ -198,6 +200,14 @@ pub fn channel<T>(init: T) -> (Sender<T>, Receiver<T>) {
}
impl<T> Receiver<T> {
fn from_shared(version: usize, shared: Arc<Shared<T>>) -> Self {
// No synchronization necessary as this is only used as a counter and
// not memory access.
shared.ref_count_rx.fetch_add(1, Relaxed);
Self { version, shared }
}
/// Returns a reference to the most recently sent value
///
/// Outstanding borrows hold a read lock. This means that long lived borrows
@@ -260,6 +270,12 @@ impl<T> Receiver<T> {
// loop around again in case the wake-up was spurious
}
}
cfg_process_driver! {
pub(crate) fn try_has_changed(&mut self) -> Option<Result<(), error::RecvError>> {
maybe_changed(&self.shared, &mut self.version)
}
}
}
fn maybe_changed<T>(
@@ -289,11 +305,7 @@ impl<T> Clone for Receiver<T> {
let version = self.version;
let shared = self.shared.clone();
// No synchronization necessary as this is only used as a counter and
// not memory access.
shared.ref_count_rx.fetch_add(1, Relaxed);
Receiver { shared, version }
Self::from_shared(version, shared)
}
}
@@ -396,6 +408,15 @@ impl<T> Sender<T> {
notified.await;
debug_assert_eq!(0, self.shared.ref_count_rx.load(Relaxed));
}
cfg_signal_internal! {
pub(crate) fn subscribe(&self) -> Receiver<T> {
let shared = self.shared.clone();
let version = shared.version.load(SeqCst);
Receiver::from_shared(version, shared)
}
}
}
impl<T> Drop for Sender<T> {