remora/crates/remora-stream/src/cancel.rs

51 lines
1.2 KiB
Rust

use tokio::sync::watch;
/// Cooperative cancellation token for the copy loop and bounded I/O.
///
/// Cancel wakes waiters so in-flight operations bounded by
/// [`crate::copy_stream`] can return without waiting for their full
/// no-progress timeout.
///
/// A receiver is retained so `watch::Sender::send` always has a live
/// subscriber and the cancelled flag is stored.
#[derive(Clone, Debug)]
pub struct CancelToken {
tx: watch::Sender<bool>,
rx: watch::Receiver<bool>,
}
impl Default for CancelToken {
fn default() -> Self {
Self::new()
}
}
impl CancelToken {
#[must_use]
pub fn new() -> Self {
let (tx, rx) = watch::channel(false);
Self { tx, rx }
}
pub fn cancel(&self) {
let _ = self.tx.send(true);
}
#[must_use]
pub fn is_cancelled(&self) -> bool {
*self.rx.borrow()
}
/// Completes when [`Self::cancel`] has been called.
pub async fn cancelled(&self) {
let mut rx = self.rx.clone();
if *rx.borrow_and_update() {
return;
}
while rx.changed().await.is_ok() {
if *rx.borrow_and_update() {
return;
}
}
}
}