use std::{
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
time::{Duration, Instant},
};
#[derive(Clone, Debug, Default)]
pub struct CancellationToken(Arc<AtomicBool>);
impl CancellationToken {
pub fn cancel(&self) {
self.0.store(true, Ordering::Relaxed);
}
pub fn is_cancelled(&self) -> bool {
self.0.load(Ordering::Relaxed)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum ScanAborted {
#[error("scan cancelled")]
Cancelled,
#[error("scan deadline exceeded")]
TimedOut,
}
#[derive(Clone, Debug, Default)]
pub struct ScanControl {
deadline: Option<Instant>,
cancellation: Option<CancellationToken>,
}
impl ScanControl {
pub fn with_deadline(mut self, deadline: Instant) -> Self {
self.deadline = Some(deadline);
self
}
pub fn with_timeout(self, timeout: Duration) -> anyhow::Result<Self> {
let deadline = Instant::now()
.checked_add(timeout)
.ok_or_else(|| anyhow::anyhow!("scan timeout is too large"))?;
Ok(self.with_deadline(deadline))
}
pub fn with_cancellation(mut self, token: CancellationToken) -> Self {
self.cancellation = Some(token);
self
}
pub(crate) fn is_limited(&self) -> bool {
self.deadline.is_some() || self.cancellation.is_some()
}
#[inline]
pub fn check(&self) -> Result<(), ScanAborted> {
if self.cancellation.as_ref().is_some_and(CancellationToken::is_cancelled) {
return Err(ScanAborted::Cancelled);
}
if self.deadline.is_some_and(|deadline| Instant::now() >= deadline) {
return Err(ScanAborted::TimedOut);
}
Ok(())
}
}