rustocat/src/shutdown.rs

63 lines
1.3 KiB
Rust

use tokio::sync::broadcast;
pub(crate) struct Shutdown {
shutdown: bool,
notify: broadcast::Receiver<()>,
sender: broadcast::Sender<()>,
}
impl Shutdown {
pub(crate) fn new() -> Shutdown {
let (sender, notify) = broadcast::channel(1);
Shutdown {
shutdown: false,
notify,
sender,
}
}
pub(crate) fn shutdown(&mut self) {
if self.shutdown {
return;
}
let _ = self.sender.send(());
self.shutdown = true;
}
pub(crate) fn receiver(&self) -> ShutdownReceiver {
ShutdownReceiver::new(self.notify.resubscribe())
}
}
#[derive(Debug)]
pub(crate) struct ShutdownReceiver {
shutdown: bool,
notify: broadcast::Receiver<()>,
}
impl ShutdownReceiver {
pub(crate) fn new(notify: broadcast::Receiver<()>) -> ShutdownReceiver {
ShutdownReceiver {
shutdown: false,
notify,
}
}
pub(crate) async fn recv(&mut self) {
if self.shutdown {
return;
}
let _ = self.notify.recv().await;
self.shutdown = true;
}
}
impl Clone for ShutdownReceiver {
fn clone(&self) -> Self {
ShutdownReceiver {
shutdown: self.shutdown,
notify: self.notify.resubscribe(),
}
}
}