use crate::{
StepID,
error::AnnattoError,
workflow::{StatusMessage, StatusSender},
};
use log::{info, warn};
use std::sync::{Arc, Mutex};
struct ProgressState {
tx: Option<StatusSender>,
accumulated_finished_work: usize,
}
pub struct ProgressReporter {
state: Arc<Mutex<ProgressState>>,
total_work: Option<usize>,
step_id: StepID,
}
impl ProgressReporter {
pub fn new(
tx: Option<StatusSender>,
step_id: StepID,
total_work: usize,
) -> Result<ProgressReporter, AnnattoError> {
let state = ProgressState {
tx,
accumulated_finished_work: 0,
};
let reporter = ProgressReporter {
state: Arc::new(Mutex::new(state)),
step_id,
total_work: Some(total_work),
};
reporter.worked(0)?;
Ok(reporter)
}
pub fn new_unknown_total_work(
tx: Option<StatusSender>,
step_id: StepID,
) -> Result<ProgressReporter, AnnattoError> {
let state = ProgressState {
tx,
accumulated_finished_work: 0,
};
let reporter = ProgressReporter {
state: Arc::new(Mutex::new(state)),
step_id,
total_work: None,
};
reporter.worked(0)?;
Ok(reporter)
}
pub fn info<S: ToString>(&self, msg: S) -> Result<(), AnnattoError> {
let msg = msg.to_string();
let state = self.state.lock()?;
if let Some(ref tx) = state.tx {
tx.send(StatusMessage::Info(msg))?;
} else {
info!("{msg}");
}
Ok(())
}
pub fn warn<S: ToString>(&self, msg: S) -> Result<(), AnnattoError> {
let msg = msg.to_string();
let state = self.state.lock()?;
if let Some(ref tx) = state.tx {
tx.send(StatusMessage::Warning(msg))?;
} else {
warn!("{msg}");
}
Ok(())
}
pub fn worked(&self, finished_work: usize) -> Result<(), AnnattoError> {
let mut state = self.state.lock()?;
state.accumulated_finished_work += finished_work;
if let Some(ref tx) = state.tx {
tx.send(StatusMessage::Progress {
id: self.step_id.clone(),
total_work: self.total_work,
finished_work: state.accumulated_finished_work,
})?;
}
Ok(())
}
}