use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
#[derive(Clone, Debug, Default)]
pub struct CancelToken {
flag: Arc<AtomicBool>,
}
impl CancelToken {
pub fn new() -> Self {
Self::default()
}
pub fn cancel(&self) {
self.flag.store(true, Ordering::SeqCst);
}
pub fn is_cancelled(&self) -> bool {
self.flag.load(Ordering::SeqCst)
}
pub fn check(&self) -> Result<(), Cancelled> {
if self.is_cancelled() {
Err(Cancelled)
} else {
Ok(())
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Cancelled;
impl std::fmt::Display for Cancelled {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "cancelled")
}
}
impl std::error::Error for Cancelled {}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Progress {
pub phase: String,
pub message: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub detail: Option<String>,
pub current: u64,
pub total: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub operation_id: Option<String>,
}
pub trait ProgressSink {
fn progress(&mut self, progress: &Progress);
}
#[derive(Clone, Copy, Debug, Default)]
pub struct NullSink;
impl ProgressSink for NullSink {
fn progress(&mut self, _progress: &Progress) {}
}
#[derive(Clone, Debug, Default)]
pub struct RecordingSink {
pub events: Vec<Progress>,
}
impl ProgressSink for RecordingSink {
fn progress(&mut self, progress: &Progress) {
self.events.push(progress.clone());
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WorkBudget {
chunk: u64,
used: u64,
next_boundary: u64,
}
impl Default for WorkBudget {
fn default() -> Self {
Self::new(Self::DEFAULT_CHUNK)
}
}
impl WorkBudget {
pub const DEFAULT_CHUNK: u64 = 1024;
pub fn new(chunk: u64) -> Self {
let chunk = chunk.max(1);
Self { chunk, used: 0, next_boundary: chunk }
}
pub fn chunk(&self) -> u64 {
self.chunk
}
pub fn used(&self) -> u64 {
self.used
}
pub fn step(&mut self, units: u64) -> bool {
self.used = self.used.saturating_add(units);
if self.used >= self.next_boundary {
let crossed = (self.used - self.next_boundary) / self.chunk + 1;
self.next_boundary = self.next_boundary.saturating_add(crossed * self.chunk);
true
} else {
false
}
}
}
pub struct Runner<'a> {
token: CancelToken,
sink: &'a mut dyn ProgressSink,
budget: WorkBudget,
operation_id: Option<String>,
}
impl<'a> Runner<'a> {
pub fn new(token: CancelToken, sink: &'a mut dyn ProgressSink, budget: WorkBudget) -> Self {
Self { token, sink, budget, operation_id: None }
}
pub fn detached(sink: &'a mut dyn ProgressSink) -> Self {
Self::new(CancelToken::new(), sink, WorkBudget::default())
}
pub fn with_operation(mut self, operation_id: impl Into<String>) -> Self {
self.operation_id = Some(operation_id.into());
self
}
pub fn token(&self) -> &CancelToken {
&self.token
}
pub fn budget(&self) -> &WorkBudget {
&self.budget
}
pub fn stage(&mut self, phase: &str, message: &str, current: u64, total: u64) -> Result<(), Cancelled> {
self.publish(phase, message, None, current, total);
self.token.check()
}
pub fn stage_detail(
&mut self,
phase: &str,
message: &str,
detail: &str,
current: u64,
total: u64,
) -> Result<(), Cancelled> {
self.publish(phase, message, Some(detail.to_string()), current, total);
self.token.check()
}
pub fn work(&mut self, units: u64, phase: &str, current: u64, total: u64) -> Result<(), Cancelled> {
if self.budget.step(units) {
self.publish(phase, phase, None, current, total);
self.token.check()?;
}
Ok(())
}
fn publish(&mut self, phase: &str, message: &str, detail: Option<String>, current: u64, total: u64) {
let progress = Progress {
phase: phase.to_string(),
message: message.to_string(),
detail,
current,
total,
operation_id: self.operation_id.clone(),
};
self.sink.progress(&progress);
}
}