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

io: refactor out usage of Weak in the io handle (#4656)

Browse Source
This commit is contained in:
盏一
2022-05-19 10:24:27 +02:00
committed by GitHub
parent 3fce06c2cc
commit 931a7773de
3 changed files with 79 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/io/driver/mod.rs
+74 -63
View File
@@ -18,13 +18,13 @@ mod metrics;
use crate::park::{Park, Unpark};
use crate::util::slab::{self, Slab};
use crate::{loom::sync::Mutex, util::bit};
use crate::{loom::sync::RwLock, util::bit};
use metrics::IoDriverMetrics;
use std::fmt;
use std::io;
use std::sync::{Arc, Weak};
use std::sync::Arc;
use std::time::Duration;
/// I/O driver, backed by Mio.
@@ -37,10 +37,8 @@ pub(crate) struct Driver {
events: Option<mio::Events>,
/// Primary slab handle containing the state for each resource registered
/// with this driver. During Drop this is moved into the Inner structure, so
/// this is an Option to allow it to be vacated (until Drop this is always
/// Some).
resources: Option<Slab<ScheduledIo>>,
/// with this driver.
resources: Slab<ScheduledIo>,
/// The system event queue.
poll: mio::Poll,
@@ -52,7 +50,7 @@ pub(crate) struct Driver {
/// A reference to an I/O driver.
#[derive(Clone)]
pub(crate) struct Handle {
inner: Weak<Inner>,
pub(super) inner: Arc<Inner>,
}
#[derive(Debug)]
@@ -61,20 +59,17 @@ pub(crate) struct ReadyEvent {
pub(crate) ready: Ready,
}
pub(super) struct Inner {
/// Primary slab handle containing the state for each resource registered
/// with this driver.
///
/// The ownership of this slab is moved into this structure during
/// `Driver::drop`, so that `Inner::drop` can notify all outstanding handles
/// without risking new ones being registered in the meantime.
resources: Mutex<Option<Slab<ScheduledIo>>>,
struct IoDispatcher {
allocator: slab::Allocator<ScheduledIo>,
is_shutdown: bool,
}
pub(super) struct Inner {
/// Registers I/O resources.
registry: mio::Registry,
/// Allocates `ScheduledIo` handles when creating new resources.
pub(super) io_dispatch: slab::Allocator<ScheduledIo>,
io_dispatch: RwLock<IoDispatcher>,
/// Used to wake up the reactor from a call to `turn`.
waker: mio::Waker,
@@ -130,11 +125,10 @@ impl Driver {
tick: 0,
events: Some(mio::Events::with_capacity(1024)),
poll,
resources: Some(slab),
resources: slab,
inner: Arc::new(Inner {
resources: Mutex::new(None),
registry,
io_dispatch: allocator,
io_dispatch: RwLock::new(IoDispatcher::new(allocator)),
waker,
metrics: IoDriverMetrics::default(),
}),
@@ -149,7 +143,7 @@ impl Driver {
/// to bind them to this event loop.
pub(crate) fn handle(&self) -> Handle {
Handle {
inner: Arc::downgrade(&self.inner),
inner: Arc::clone(&self.inner),
}
}
@@ -160,7 +154,7 @@ impl Driver {
self.tick = self.tick.wrapping_add(1);
if self.tick == COMPACT_INTERVAL {
self.resources.as_mut().unwrap().compact()
self.resources.compact()
}
let mut events = self.events.take().expect("i/o driver event store missing");
@@ -194,7 +188,7 @@ impl Driver {
fn dispatch(&mut self, token: mio::Token, ready: Ready) {
let addr = slab::Address::from_usize(ADDRESS.unpack(token.0));
let resources = self.resources.as_mut().unwrap();
let resources = &mut self.resources;
let io = match resources.get(addr) {
Some(io) => io,
@@ -214,22 +208,7 @@ impl Driver {
impl Drop for Driver {
fn drop(&mut self) {
(*self.inner.resources.lock()) = self.resources.take();
}
}
impl Drop for Inner {
fn drop(&mut self) {
let resources = self.resources.lock().take();
if let Some(mut slab) = resources {
slab.for_each(|io| {
// If a task is waiting on the I/O resource, notify it. The task
// will then attempt to use the I/O resource and fail due to the
// driver being shutdown.
io.shutdown();
});
}
self.shutdown();
}
}
@@ -251,7 +230,16 @@ impl Park for Driver {
Ok(())
}
fn shutdown(&mut self) {}
fn shutdown(&mut self) {
if self.inner.shutdown() {
self.resources.for_each(|io| {
// If a task is waiting on the I/O resource, notify it. The task
// will then attempt to use the I/O resource and fail due to the
// driver being shutdown. And shutdown will clear all wakers.
io.shutdown();
});
}
}
}
impl fmt::Debug for Driver {
@@ -292,18 +280,11 @@ cfg_not_rt! {
cfg_net! {
cfg_metrics! {
impl Handle {
// TODO: Remove this when handle contains `Arc<Inner>` so that we can return
// &IoDriverMetrics instead of using a closure.
//
// Related issue: https://github.com/tokio-rs/tokio/issues/4509
pub(crate) fn with_io_driver_metrics<F, R>(&self, f: F) -> Option<R>
where
F: Fn(&IoDriverMetrics) -> R,
{
self.inner().map(|inner| f(&inner.metrics))
impl Handle {
pub(crate) fn metrics(&self) -> &IoDriverMetrics {
&self.inner.metrics
}
}
}
}
}
@@ -318,13 +299,7 @@ impl Handle {
/// 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.waker.wake().expect("failed to wake I/O driver");
}
}
pub(super) fn inner(&self) -> Option<Arc<Inner>> {
self.inner.upgrade()
self.inner.waker.wake().expect("failed to wake I/O driver");
}
}
@@ -340,6 +315,17 @@ impl fmt::Debug for Handle {
}
}
// ===== impl IoDispatcher =====
impl IoDispatcher {
fn new(allocator: slab::Allocator<ScheduledIo>) -> Self {
Self {
allocator,
is_shutdown: false,
}
}
}
// ===== impl Inner =====
impl Inner {
@@ -351,12 +337,7 @@ impl Inner {
source: &mut impl mio::event::Source,
interest: Interest,
) -> io::Result<slab::Ref<ScheduledIo>> {
let (address, shared) = self.io_dispatch.allocate().ok_or_else(|| {
io::Error::new(
io::ErrorKind::Other,
"reactor at max registered I/O resources",
)
})?;
let (address, shared) = self.allocate()?;
let token = GENERATION.pack(shared.generation(), ADDRESS.pack(address.as_usize(), 0));
@@ -376,6 +357,36 @@ impl Inner {
Ok(())
}
/// shutdown the dispatcher.
fn shutdown(&self) -> bool {
let mut io = self.io_dispatch.write().unwrap();
if io.is_shutdown {
return false;
}
io.is_shutdown = true;
true
}
fn is_shutdown(&self) -> bool {
return self.io_dispatch.read().unwrap().is_shutdown;
}
fn allocate(&self) -> io::Result<(slab::Address, slab::Ref<ScheduledIo>)> {
let io = self.io_dispatch.read().unwrap();
if io.is_shutdown {
return Err(io::Error::new(
io::ErrorKind::Other,
"failed to find event loop",
));
}
io.allocator.allocate().ok_or_else(|| {
io::Error::new(
io::ErrorKind::Other,
"reactor at max registered I/O resources",
)
})
}
}
impl Direction {
tokio/src/io/driver/registration.rs
+4 -15
View File
@@ -72,14 +72,7 @@ impl Registration {
interest: Interest,
handle: Handle,
) -> io::Result<Registration> {
let shared = if let Some(inner) = handle.inner() {
inner.add_source(io, interest)?
} else {
return Err(io::Error::new(
io::ErrorKind::Other,
"failed to find event loop",
));
};
let shared = handle.inner.add_source(io, interest)?;
Ok(Registration { handle, shared })
}
@@ -101,11 +94,7 @@ impl Registration {
///
/// `Err` is returned if an error is encountered.
pub(crate) fn deregister(&mut self, io: &mut impl Source) -> io::Result<()> {
let inner = match self.handle.inner() {
Some(inner) => inner,
None => return Err(io::Error::new(io::ErrorKind::Other, "reactor gone")),
};
inner.deregister_source(io)
self.handle.inner.deregister_source(io)
}
pub(crate) fn clear_readiness(&self, event: ReadyEvent) {
@@ -157,7 +146,7 @@ impl Registration {
let coop = ready!(crate::coop::poll_proceed(cx));
let ev = ready!(self.shared.poll_readiness(cx, direction));
if self.handle.inner().is_none() {
if self.handle.inner.is_shutdown() {
return Poll::Ready(Err(gone()));
}
@@ -235,7 +224,7 @@ cfg_io_readiness! {
pin!(fut);
crate::future::poll_fn(|cx| {
if self.handle.inner().is_none() {
if self.handle.inner.is_shutdown() {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::Other,
crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR
tokio/src/runtime/metrics/runtime.rs
+1 -1
View File
@@ -529,7 +529,7 @@ cfg_net! {
.as_inner()
.io_handle
.as_ref()
.and_then(|h| h.with_io_driver_metrics(f))
.map(|h| f(h.metrics()))
.unwrap_or(0)
}
}