use tokio::sync::{broadcast, mpsc};
#[derive(Debug)]
pub struct Shutdown {
shutdown: bool,
notify: broadcast::Receiver<()>,
}
impl Shutdown {
pub fn new(notify: broadcast::Receiver<()>) -> Shutdown {
Shutdown {
shutdown: false,
notify,
}
}
pub fn is_shutdown(&self) -> bool {
self.shutdown
}
pub async fn recv(&mut self) {
if self.shutdown {
return;
}
let _ = self.notify.recv().await;
self.shutdown = true;
}
}
pub struct ShutdownNotifier {
pub notify_shutdown: broadcast::Sender<()>,
pub shutdown_complete_rx: mpsc::Receiver<()>,
pub shutdown_complete_tx: mpsc::Sender<()>,
}
impl ShutdownNotifier {
pub fn new() -> Self {
let (notify_shutdown, _) = broadcast::channel(1);
let (shutdown_complete_tx, shutdown_complete_rx) = mpsc::channel(1);
Self {
notify_shutdown,
shutdown_complete_rx,
shutdown_complete_tx,
}
}
pub fn subscribe_for_shutdown(&self) -> Shutdown {
Shutdown::new(self.notify_shutdown.subscribe())
}
}
impl Default for ShutdownNotifier {
fn default() -> Self {
ShutdownNotifier::new()
}
}