#![doc = include_str!("../README.md")]
use futures::{Stream, ready};
use pin_project_lite::pin_project;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use thiserror::Error;
use tokio::sync::broadcast;
use tokio_stream::wrappers::BroadcastStream;
use tokio_stream::wrappers::errors::BroadcastStreamRecvError;
use tokio_util::sync::{
CancellationToken, WaitForCancellationFuture, WaitForCancellationFutureOwned,
};
#[must_use = "if unused, the progress token will be completed immediately"]
pub struct CompleteGuard<'a, S: Clone + Send + 'static> {
token: &'a ProgressToken<S>,
}
impl<'a, S: Clone + Send + 'static> CompleteGuard<'a, S> {
pub fn forget(self) {
std::mem::forget(self);
}
}
impl<'a, S: Clone + Send + 'static> Drop for CompleteGuard<'a, S> {
fn drop(&mut self) {
self.token.complete();
}
}
#[derive(Debug, Clone, Copy)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum Progress {
Determinate(f64),
Indeterminate,
}
impl Progress {
pub fn as_f64(&self) -> Option<f64> {
match self {
Progress::Determinate(v) => Some(*v),
Progress::Indeterminate => None,
}
}
}
#[derive(Debug, Clone, Copy, Error)]
pub enum ProgressError {
#[error("progress updates lagged")]
Lagged,
#[error("the operation has been cancelled")]
Cancelled,
}
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ProgressUpdate<S> {
pub progress: Progress,
pub statuses: Vec<S>,
pub is_cancelled: bool,
}
impl<S> ProgressUpdate<S> {
pub fn status(&self) -> &S {
self.statuses.last().unwrap()
}
}
struct ProgressNodeInner<S> {
parent: Option<Arc<ProgressNode<S>>>,
children: Vec<(Arc<ProgressNode<S>>, f64)>,
progress: Progress,
status: S,
is_completed: bool,
update_sender: broadcast::Sender<ProgressUpdate<S>>,
}
struct ProgressNode<S> {
inner: Mutex<ProgressNodeInner<S>>,
}
impl<S: Clone + Send> ProgressNode<S> {
fn new(status: S) -> Self {
let (tx, _) = broadcast::channel(16);
Self {
inner: Mutex::new(ProgressNodeInner {
parent: None,
children: Vec::new(),
progress: Progress::Determinate(0.0),
status,
is_completed: false,
update_sender: tx,
}),
}
}
fn child(parent: &Arc<Self>, weight: f64, status: S) -> Arc<Self> {
let mut parent_inner = parent.inner.lock().unwrap();
let (tx, _) = broadcast::channel(16);
let child = Self {
inner: Mutex::new(ProgressNodeInner {
parent: Some(parent.clone()),
children: Vec::new(),
progress: Progress::Determinate(0.0),
status,
is_completed: false,
update_sender: tx,
}),
};
let child = Arc::new(child);
parent_inner.children.push((child.clone(), weight));
child
}
fn calculate_progress(node: &Arc<Self>) -> Progress {
let inner = node.inner.lock().unwrap();
if matches!(inner.progress, Progress::Indeterminate) {
return Progress::Indeterminate;
}
if inner.children.is_empty() {
return inner.progress;
}
let has_indeterminate = inner
.children
.iter()
.filter(|(child, _)| {
let child_inner = child.inner.lock().unwrap();
!child_inner.is_completed
})
.any(|(child, _)| matches!(Self::calculate_progress(child), Progress::Indeterminate));
if has_indeterminate {
return Progress::Indeterminate;
}
let total: f64 = inner
.children
.iter()
.map(|(child, weight)| {
match Self::calculate_progress(child) {
Progress::Determinate(p) => p * weight,
Progress::Indeterminate => 0.0, }
})
.sum();
Progress::Determinate(total)
}
fn get_status_hierarchy(node: &Arc<Self>) -> Vec<S> {
let inner = node.inner.lock().unwrap();
let mut result = vec![inner.status.clone()];
if !inner.children.is_empty() {
let active_child = inner
.children
.iter()
.filter(|(child, _)| {
let child_inner = child.inner.lock().unwrap();
!child_inner.is_completed
})
.next();
if let Some((child, _)) = active_child {
let child_statuses = Self::get_status_hierarchy(child);
result.extend(child_statuses);
}
}
result
}
fn notify_subscribers(node: &Arc<Self>, is_cancelled: bool) {
let update = ProgressUpdate {
progress: Self::calculate_progress(node),
statuses: Self::get_status_hierarchy(node),
is_cancelled,
};
{
let inner = node.inner.lock().unwrap();
let _ = inner.update_sender.send(update);
};
let parent = {
let inner = node.inner.lock().unwrap();
inner.parent.clone()
};
if let Some(parent) = parent {
Self::notify_subscribers(&parent, false);
}
}
}
#[derive(Clone)]
pub struct ProgressToken<S> {
node: Arc<ProgressNode<S>>,
cancel_token: CancellationToken,
}
impl<S: Default + Clone + Send + 'static> Default for ProgressToken<S> {
fn default() -> Self {
Self::new(S::default())
}
}
impl<S: std::fmt::Debug> std::fmt::Debug for ProgressToken<S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ProgressToken")
.field("is_cancelled", &self.cancel_token.is_cancelled())
.finish()
}
}
impl<S: Clone + Send + 'static> ProgressToken<S> {
pub fn new(status: impl Into<S>) -> Self {
let node = Arc::new(ProgressNode::new(status.into()));
Self {
node,
cancel_token: CancellationToken::new(),
}
}
pub fn child(&self, weight: f64, status: impl Into<S>) -> Self {
let node = ProgressNode::child(&self.node, weight, status.into());
Self {
node,
cancel_token: self.cancel_token.child_token(),
}
}
pub fn update_progress(&self, progress: f64) {
if self.is_cancelled() {
return;
}
let is_completed = {
let inner = self.node.inner.lock().unwrap();
inner.is_completed
};
if is_completed {
return;
}
let mut inner = self.node.inner.lock().unwrap();
inner.progress = Progress::Determinate(progress.max(0.0).min(1.0));
drop(inner);
ProgressNode::notify_subscribers(&self.node, false);
}
pub fn update_indeterminate(&self) {
if self.is_cancelled() {
return;
}
let mut inner = self.node.inner.lock().unwrap();
if inner.is_completed {
return;
}
inner.progress = Progress::Indeterminate;
drop(inner);
ProgressNode::notify_subscribers(&self.node, false);
}
pub fn update_status(&self, status: impl Into<S>) {
if self.is_cancelled() {
return;
}
let mut inner = self.node.inner.lock().unwrap();
if inner.is_completed {
return;
}
inner.status = status.into();
drop(inner);
ProgressNode::notify_subscribers(&self.node, false);
}
pub fn update(&self, progress: Progress, status: impl Into<S>) {
if self.is_cancelled() {
return;
}
let mut inner = self.node.inner.lock().unwrap();
if inner.is_completed {
return;
}
inner.status = status.into();
inner.progress = progress;
drop(inner);
ProgressNode::notify_subscribers(&self.node, false);
}
pub fn complete(&self) {
if self.is_cancelled() {
return;
}
let mut inner = self.node.inner.lock().unwrap();
if !inner.is_completed {
inner.is_completed = true;
inner.progress = Progress::Determinate(1.0);
drop(inner);
ProgressNode::notify_subscribers(&self.node, false);
}
}
pub fn check(&self) -> Result<(), ProgressError> {
if self.is_cancelled() {
Err(ProgressError::Cancelled)
} else {
Ok(())
}
}
pub fn is_cancelled(&self) -> bool {
self.cancel_token.is_cancelled()
}
pub fn cancel(&self) {
if !self.cancel_token.is_cancelled() {
self.cancel_token.cancel();
ProgressNode::notify_subscribers(&self.node, true);
}
}
pub fn state(&self) -> Progress {
ProgressNode::calculate_progress(&self.node)
}
pub fn statuses(&self) -> Vec<S> {
ProgressNode::get_status_hierarchy(&self.node)
}
pub fn cancelled(&self) -> WaitForCancellationFuture {
self.cancel_token.cancelled()
}
pub fn cancelled_owned(self) -> WaitForCancellationFutureOwned {
self.cancel_token.cancelled_owned()
}
pub async fn updated(&self) -> Result<ProgressUpdate<S>, ProgressError> {
let mut rx = {
let inner = self.node.inner.lock().unwrap();
inner.update_sender.subscribe()
};
tokio::select! {
_ = self.cancel_token.cancelled() => {
Err(ProgressError::Cancelled)
}
result = rx.recv() => {
match result {
Ok(update) => Ok(update),
Err(broadcast::error::RecvError::Closed) => Err(ProgressError::Cancelled),
Err(broadcast::error::RecvError::Lagged(_)) => Err(ProgressError::Lagged),
}
}
}
}
pub fn subscribe(&self) -> ProgressStream<'_, S> {
let rx = {
let inner = self.node.inner.lock().unwrap();
inner.update_sender.subscribe()
};
ProgressStream {
token: self,
rx: BroadcastStream::new(rx),
}
}
pub fn complete_guard(&self) -> CompleteGuard<'_, S> {
CompleteGuard { token: self }
}
}
pin_project! {
#[must_use = "futures do nothing unless polled"]
pub struct WaitForUpdateFuture<'a, S> {
token: &'a ProgressToken<S>,
#[pin]
future: tokio::sync::futures::Notified<'a>,
}
}
impl<'a, S: Clone + Send + 'static> Future for WaitForUpdateFuture<'a, S> {
type Output = Option<ProgressUpdate<S>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut this = self.project();
if this.token.cancel_token.is_cancelled() {
return Poll::Ready(None);
}
ready!(this.future.as_mut().poll(cx));
Poll::Ready(Some(ProgressUpdate {
progress: this.token.state(),
statuses: this.token.statuses(),
is_cancelled: false,
}))
}
}
pin_project! {
#[must_use = "streams do nothing unless polled"]
pub struct ProgressStream<'a, S> {
token: &'a ProgressToken<S>,
#[pin]
rx: BroadcastStream<ProgressUpdate<S>>,
}
}
impl<'a, S: Clone + Send + 'static> Stream for ProgressStream<'a, S> {
type Item = Result<ProgressUpdate<S>, ProgressError>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.project().rx.poll_next(cx).map(|opt| {
opt.and_then(|res| match res {
Ok(update) => Some(Ok(update)),
Err(BroadcastStreamRecvError::Lagged(_)) => Some(Err(ProgressError::Lagged)),
})
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures::StreamExt;
use std::time::Duration;
use tokio::time::sleep;
async fn create_test_hierarchy() -> (
ProgressToken<String>,
ProgressToken<String>,
ProgressToken<String>,
) {
let root = ProgressToken::new("root".to_string());
let child1 = root.child(0.6, "child1".to_string());
let child2 = root.child(0.4, "child2".to_string());
(root, child1, child2)
}
#[tokio::test]
async fn test_basic_progress_updates() {
let token: ProgressToken<String> = ProgressToken::new("test".to_string());
token.update_progress(0.5);
assert!(
matches!(token.state(), Progress::Determinate(p) if (p - 0.5).abs() < f64::EPSILON)
);
token.update_progress(1.0);
assert!(
matches!(token.state(), Progress::Determinate(p) if (p - 1.0).abs() < f64::EPSILON)
);
token.update_progress(1.5);
assert!(
matches!(token.state(), Progress::Determinate(p) if (p - 1.0).abs() < f64::EPSILON)
);
token.update_progress(-0.5);
assert!(matches!(token.state(), Progress::Determinate(p) if p.abs() < f64::EPSILON));
}
#[tokio::test]
async fn test_hierarchical_progress() {
let (root, child1, child2) = create_test_hierarchy().await;
child1.update_progress(0.5);
child2.update_progress(0.5);
assert!(matches!(root.state(), Progress::Determinate(p) if (p - 0.5).abs() < f64::EPSILON));
child1.update_progress(1.0);
assert!(matches!(root.state(), Progress::Determinate(p) if (p - 0.8).abs() < f64::EPSILON));
}
#[tokio::test]
async fn test_indeterminate_state() {
let (root, child1, child2) = create_test_hierarchy().await;
child1.update_indeterminate();
child2.update_progress(0.5);
assert!(matches!(root.state(), Progress::Indeterminate));
child1.update_progress(0.5);
assert!(matches!(root.state(), Progress::Determinate(_)));
}
#[tokio::test]
async fn test_status_updates() {
let token: ProgressToken<String> = ProgressToken::new("initial status".to_string());
let statuses = token.statuses();
assert_eq!(statuses, vec!["initial status".to_string()]);
token.update_status("updated status".to_string());
let statuses = token.statuses();
assert_eq!(statuses, vec!["updated status".to_string()]);
}
#[tokio::test]
async fn test_status_hierarchy() {
let (root, child1, _) = create_test_hierarchy().await;
let statuses = root.statuses();
assert_eq!(statuses, vec!["root".to_string(), "child1".to_string()]);
child1.update_status("updated child1".to_string());
let statuses = root.statuses();
assert_eq!(
statuses,
vec!["root".to_string(), "updated child1".to_string()]
);
}
#[tokio::test]
async fn test_cancellation() {
let (root, child1, child2) = create_test_hierarchy().await;
root.cancel();
assert!(root.cancel_token.is_cancelled());
assert!(child1.cancel_token.is_cancelled());
assert!(child2.cancel_token.is_cancelled());
child1.update_progress(0.5);
assert!(matches!(child1.state(), Progress::Determinate(p) if p.abs() < f64::EPSILON));
}
#[tokio::test]
async fn test_complete_guard() {
let token: ProgressToken<String> = ProgressToken::new("test".to_string());
{
let _guard = token.complete_guard();
token.update_progress(0.5);
assert!(
matches!(token.state(), Progress::Determinate(p) if (p - 0.5).abs() < f64::EPSILON)
);
}
assert!(
matches!(token.state(), Progress::Determinate(p) if (p - 1.0).abs() < f64::EPSILON)
);
token.update_progress(0.5);
assert!(
matches!(token.state(), Progress::Determinate(p) if (p - 1.0).abs() < f64::EPSILON)
);
let token: ProgressToken<String> = ProgressToken::new("test2".to_string());
{
let guard = token.complete_guard();
token.update_progress(0.5);
guard.forget(); }
assert!(
matches!(token.state(), Progress::Determinate(p) if (p - 0.5).abs() < f64::EPSILON)
);
}
#[tokio::test]
async fn test_subscription() {
let token: ProgressToken<String> = ProgressToken::new("test".to_string());
let mut subscription = token.subscribe();
token.update_progress(0.5);
let update = subscription.next().await.unwrap().unwrap();
assert!(
matches!(update.progress, Progress::Determinate(p) if (p - 0.5).abs() < f64::EPSILON)
);
}
#[tokio::test]
async fn test_multiple_subscribers() {
let token: ProgressToken<String> = ProgressToken::new("test".to_string());
let mut sub1 = token.subscribe();
let mut sub2 = token.subscribe();
token.update_progress(0.5);
let update1 = sub1.next().await.unwrap().unwrap();
let update2 = sub2.next().await.unwrap().unwrap();
assert!(
matches!(update1.progress, Progress::Determinate(p) if (p - 0.5).abs() < f64::EPSILON),
"{update1:?}"
);
assert!(
matches!(update2.progress, Progress::Determinate(p) if (p - 0.5).abs() < f64::EPSILON),
"{update2:?}"
);
token.update_progress(0.75);
let update1 = sub1.next().await.unwrap().unwrap();
let update2 = sub2.next().await.unwrap().unwrap();
assert!(
matches!(update1.progress, Progress::Determinate(p) if (p - 0.75).abs() < f64::EPSILON),
"{update1:?}"
);
assert!(
matches!(update2.progress, Progress::Determinate(p) if (p - 0.75).abs() < f64::EPSILON),
"{update2:?}"
);
}
#[tokio::test]
async fn test_concurrent_updates() {
let token: ProgressToken<String> = ProgressToken::new("test".to_string());
let mut handles = vec![];
for i in 0..10 {
let token = token.clone();
handles.push(tokio::spawn(async move {
sleep(Duration::from_millis(i * 10)).await;
token.update_progress(i as f64 / 10.0);
}));
}
for handle in handles {
handle.await.unwrap();
}
assert!(
matches!(token.state(), Progress::Determinate(p) if (p - 0.9).abs() < f64::EPSILON)
);
}
#[tokio::test]
async fn test_edge_cases() {
let token: ProgressToken<String> = ProgressToken::new("single".to_string());
token.update_progress(0.5);
assert!(
matches!(token.state(), Progress::Determinate(p) if (p - 0.5).abs() < f64::EPSILON)
);
let mut current: ProgressToken<String> = ProgressToken::new("root".to_string());
for i in 0..10 {
current = current.child(1.0, format!("child{}", i));
}
current.update_progress(1.0);
assert!(
matches!(current.state(), Progress::Determinate(p) if (p - 1.0).abs() < f64::EPSILON)
);
}
#[tokio::test]
async fn test_three_level_hierarchy_progress() {
let root: ProgressToken<String> = ProgressToken::new("root".to_string());
let child1 = root.child(0.7, "child1".to_string());
let child2 = root.child(0.3, "child2".to_string());
let grandchild1_1 = child1.child(0.6, "grandchild1_1".to_string());
let grandchild1_2 = child1.child(0.4, "grandchild1_2".to_string());
let grandchild2_1 = child2.child(1.0, "grandchild2_1".to_string());
grandchild1_1.update_progress(0.5); grandchild1_2.update_progress(1.0); grandchild2_1.update_progress(0.6);
assert!(
matches!(child1.state(), Progress::Determinate(p) if (p - 0.7).abs() < f64::EPSILON),
"child1 progress incorrect"
);
assert!(
matches!(child2.state(), Progress::Determinate(p) if (p - 0.6).abs() < f64::EPSILON),
"child2 progress incorrect"
);
assert!(
matches!(root.state(), Progress::Determinate(p) if (p - 0.67).abs() < f64::EPSILON),
"root progress incorrect"
);
}
#[tokio::test]
async fn test_completion_hierarchy() {
let root: ProgressToken<String> = ProgressToken::new("root".to_string());
let child1 = root.child(0.6, "child1".to_string());
let child2 = root.child(0.4, "child2".to_string());
let grandchild1 = child1.child(1.0, "grandchild1".to_string());
grandchild1.update_progress(0.5);
grandchild1.complete();
assert!(
matches!(grandchild1.state(), Progress::Determinate(p) if (p - 1.0).abs() < f64::EPSILON),
"completed grandchild should be at 100%"
);
assert!(
matches!(child1.state(), Progress::Determinate(p) if (p - 1.0).abs() < f64::EPSILON),
"child1 progress should reflect completed grandchild"
);
child2.update_progress(0.5);
assert!(
matches!(root.state(), Progress::Determinate(p) if (p - 0.8).abs() < f64::EPSILON),
"root progress incorrect after child completion"
);
child2.complete();
assert!(
matches!(root.state(), Progress::Determinate(p) if (p - 1.0).abs() < f64::EPSILON),
"root progress should be 100% when all children complete"
);
let mut root_inner = root.node.inner.lock().unwrap();
assert!(
!root_inner.is_completed,
"root should not be auto-completed when children complete"
);
}
#[tokio::test]
async fn test_mixed_completion_states() {
let root: ProgressToken<String> = ProgressToken::new("root".to_string());
let child1 = root.child(0.5, "child1".to_string());
let child2 = root.child(0.5, "child2".to_string());
let grandchild1_1 = child1.child(0.7, "grandchild1_1".to_string());
let grandchild1_2 = child1.child(0.3, "grandchild1_2".to_string());
grandchild1_1.complete();
grandchild1_2.update_progress(0.5);
assert!(
matches!(child1.state(), Progress::Determinate(p) if (p - 0.85).abs() < f64::EPSILON),
"child1 progress incorrect with mixed completion"
);
child2.update_progress(0.4);
assert!(
matches!(root.state(), Progress::Determinate(p) if (p - 0.625).abs() < f64::EPSILON),
"root progress incorrect with mixed completion states"
);
grandchild1_2.complete();
child2.complete();
assert!(
matches!(root.state(), Progress::Determinate(p) if (p - 1.0).abs() < f64::EPSILON),
"root progress should be 100% when all descendants complete"
);
assert!(
matches!(child1.state(), Progress::Determinate(p) if (p - 1.0).abs() < f64::EPSILON),
"child1 progress should be 100% when all grandchildren complete"
);
}
#[tokio::test]
async fn test_status_propagation() {
let root: ProgressToken<String> = ProgressToken::new("root".to_string());
let child1 = root.child(0.6, "child1".to_string());
let child2 = root.child(0.4, "child2".to_string());
let grandchild1 = child1.child(1.0, "grandchild1".to_string());
let statuses = root.statuses();
assert_eq!(
statuses,
vec![
"root".to_string(),
"child1".to_string(),
"grandchild1".to_string()
]
);
grandchild1.update_status("updated grandchild".to_string());
let statuses = root.statuses();
assert_eq!(
statuses,
vec![
"root".to_string(),
"child1".to_string(),
"updated grandchild".to_string()
]
);
child1.update_status("updated child1".to_string());
let statuses = root.statuses();
assert_eq!(
statuses,
vec![
"root".to_string(),
"updated child1".to_string(),
"updated grandchild".to_string()
]
);
root.update_status("updated root".to_string());
let statuses = root.statuses();
assert_eq!(
statuses,
vec![
"updated root".to_string(),
"updated child1".to_string(),
"updated grandchild".to_string()
]
);
}
#[tokio::test]
async fn test_status_propagation_with_multiple_children() {
let root: ProgressToken<String> = ProgressToken::new("root".to_string());
let child1 = root.child(0.5, "child1".to_string());
let child2 = root.child(0.5, "child2".to_string());
let grandchild1_1 = child1.child(0.7, "grandchild1_1".to_string());
let grandchild1_2 = child1.child(0.3, "grandchild1_2".to_string());
let grandchild2_1 = child2.child(1.0, "grandchild2_1".to_string());
let statuses = root.statuses();
assert_eq!(
statuses,
vec![
"root".to_string(),
"child1".to_string(),
"grandchild1_1".to_string()
]
);
grandchild1_2.update_status("updated grandchild1_2".to_string());
let statuses = root.statuses();
assert_eq!(
statuses,
vec![
"root".to_string(),
"child1".to_string(),
"grandchild1_1".to_string()
]
);
grandchild1_1.update_status("updated grandchild1_1".to_string());
let statuses = root.statuses();
assert_eq!(
statuses,
vec![
"root".to_string(),
"child1".to_string(),
"updated grandchild1_1".to_string()
]
);
grandchild2_1.update_status("updated grandchild2_1".to_string());
let statuses = root.statuses();
assert_eq!(
statuses,
vec![
"root".to_string(),
"child1".to_string(),
"updated grandchild1_1".to_string()
]
);
child2.update_status("updated child2".to_string());
let statuses = root.statuses();
assert_eq!(
statuses,
vec![
"root".to_string(),
"child1".to_string(),
"updated grandchild1_1".to_string()
]
);
}
#[tokio::test]
async fn test_status_propagation_with_completion() {
let root: ProgressToken<String> = ProgressToken::new("root".to_string());
let child1 = root.child(0.6, "child1".to_string());
let child2 = root.child(0.4, "child2".to_string());
let grandchild1 = child1.child(1.0, "grandchild1".to_string());
let statuses = root.statuses();
assert_eq!(
statuses,
vec![
"root".to_string(),
"child1".to_string(),
"grandchild1".to_string()
]
);
grandchild1.update_status("completed grandchild".to_string());
grandchild1.complete();
let statuses = root.statuses();
assert_eq!(statuses, vec!["root".to_string(), "child1".to_string()]);
child1.update_status("completed child1".to_string());
child1.complete();
let statuses = root.statuses();
assert_eq!(statuses, vec!["root".to_string(), "child2".to_string()]);
child2.update_status("updated child2".to_string());
let statuses = root.statuses();
assert_eq!(
statuses,
vec!["root".to_string(), "updated child2".to_string()]
);
}
}