use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, OnceLock};
use std::time::Duration;
static GLOBAL_PROGRESS_TRACKER: OnceLock<Arc<ProgressTracker>> = OnceLock::new();
pub fn set_tracker_once(tracker: Arc<ProgressTracker>) {
let _ = GLOBAL_PROGRESS_TRACKER.set(tracker);
}
pub(crate) fn current_global_progress_tracker() -> Arc<ProgressTracker> {
GLOBAL_PROGRESS_TRACKER.get().cloned().unwrap_or_default()
}
pub struct WorkInfo {
pub completed: AtomicU64,
pub total: u64,
pub units: String,
}
impl WorkInfo {
pub fn is_in_progress(&self) -> bool {
self.completed.load(Ordering::Relaxed) < self.total
}
}
pub enum ProgressUpdate {
Step(String),
BeginStep(String),
EndStep(String),
Work(Arc<WorkInfo>),
}
pub struct ProgressListener {
receiver: Receiver<ProgressUpdate>,
}
#[derive(Default, Debug)]
pub struct ProgressTracker {
subscribers: Vec<Sender<ProgressUpdate>>,
}
pub(crate) struct WorkGuard {
info: Arc<WorkInfo>,
}
impl WorkGuard {
pub fn inc(&self, amount: u64) {
self.info.completed.fetch_add(amount, Ordering::Relaxed);
}
}
impl Drop for WorkGuard {
fn drop(&mut self) {
self.info
.completed
.store(self.info.total, Ordering::Relaxed);
}
}
impl ProgressTracker {
pub fn new() -> ProgressTracker {
ProgressTracker {
subscribers: vec![],
}
}
pub(crate) fn begin_step(&self, message: String) {
self.subscribers.iter().for_each(|sender| {
let _ = sender.send(ProgressUpdate::BeginStep(message.clone()));
});
}
pub(crate) fn end_step(&self, message: String) {
self.subscribers.iter().for_each(|sender| {
let _ = sender.send(ProgressUpdate::EndStep(message.clone()));
});
}
pub(crate) fn doing_work(&self, total: u64, units: String) -> WorkGuard {
let info = Arc::new(WorkInfo {
completed: AtomicU64::new(0),
total,
units,
});
self.subscribers.iter().for_each(|sender| {
let _ = sender.send(ProgressUpdate::Work(Arc::clone(&info)));
});
WorkGuard {
info: Arc::clone(&info),
}
}
pub fn subscribe(&mut self) -> ProgressListener {
let (sender, receiver) = channel();
self.subscribers.push(sender);
ProgressListener { receiver }
}
}
impl Iterator for ProgressListener {
type Item = ProgressUpdate;
fn next(&mut self) -> Option<Self::Item> {
self.receiver.iter().next()
}
}
impl ProgressListener {
pub fn try_next_for(&self, timeout: Duration) -> Option<ProgressUpdate> {
self.receiver.recv_timeout(timeout).ok()
}
}