diff options
Diffstat (limited to 'tokio/src/runtime/blocking/shutdown.rs')
-rw-r--r-- | tokio/src/runtime/blocking/shutdown.rs | 44 |
1 files changed, 44 insertions, 0 deletions
diff --git a/tokio/src/runtime/blocking/shutdown.rs b/tokio/src/runtime/blocking/shutdown.rs new file mode 100644 index 00000000..d9f5eb0f --- /dev/null +++ b/tokio/src/runtime/blocking/shutdown.rs @@ -0,0 +1,44 @@ +//! A shutdown channel. +//! +//! Each worker holds the `Sender` half. When all the `Sender` halves are +//! dropped, the `Receiver` receives a notification. + +use crate::loom::sync::Arc; +use crate::sync::oneshot; + +#[derive(Debug, Clone)] +pub(super) struct Sender { + tx: Arc<oneshot::Sender<()>>, +} + +#[derive(Debug)] +pub(super) struct Receiver { + rx: oneshot::Receiver<()>, +} + +pub(super) fn channel() -> (Sender, Receiver) { + let (tx, rx) = oneshot::channel(); + let tx = Sender { tx: Arc::new(tx) }; + let rx = Receiver { rx }; + + (tx, rx) +} + +impl Receiver { + /// Block the current thread until all `Sender` handles drop. + pub(crate) fn wait(&mut self) { + use crate::runtime::enter::{enter, try_enter}; + + let mut e = if std::thread::panicking() { + match try_enter() { + Some(enter) => enter, + _ => return, + } + } else { + enter() + }; + + // The oneshot completes with an Err + let _ = e.block_on(&mut self.rx); + } +} |