use std::time::Duration;
use tokio::sync::watch;
pub struct Drain(watch::Sender<bool>);
pub struct Draining(watch::Receiver<bool>);
impl Default for Drain {
fn default() -> Self {
Self(watch::Sender::new(false))
}
}
impl Drain {
pub fn subscribe(&self) -> Draining {
Draining(self.0.subscribe())
}
pub async fn drain(&self, timeout: Duration) {
self.0.send_replace(true);
let _ = tokio::time::timeout(timeout, self.0.closed()).await;
}
}
impl Draining {
pub async fn started(&mut self) {
// a dropped drain never starts and this pends forever instead of returning
if self.0.wait_for(|draining| *draining).await.is_err() {
std::future::pending::<()>().await;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Instant;
#[tokio::test]
async fn drain_waits_for_subscribers_to_go() {
let drain = Drain::default();
let mut draining = drain.subscribe();
let subscriber = tokio::spawn(async move {
draining.started().await;
tokio::time::sleep(Duration::from_millis(50)).await;
});
let started = Instant::now();
drain.drain(Duration::from_secs(5)).await;
assert!(started.elapsed() >= Duration::from_millis(50));
assert!(started.elapsed() < Duration::from_secs(5));
subscriber.await.unwrap();
}
#[tokio::test]
async fn drain_gives_up_on_a_subscriber_that_never_goes() {
let drain = Drain::default();
let _stuck = drain.subscribe();
let started = Instant::now();
drain.drain(Duration::from_millis(50)).await;
assert!(started.elapsed() < Duration::from_secs(5));
}
#[tokio::test]
async fn a_subscriber_after_the_drain_started_sees_it_right_away() {
let drain = Drain::default();
drain.drain(Duration::ZERO).await;
let mut late = drain.subscribe();
tokio::time::timeout(Duration::from_secs(1), late.started())
.await
.expect("a late subscriber must see the drain");
}
}