parent
583b5fd0e6
commit
6a2570a0c4
@ -0,0 +1,292 @@
|
||||
//! Async progression
|
||||
use super::*;
|
||||
use termprogress::{
|
||||
prelude::*,
|
||||
};
|
||||
use futures::{
|
||||
prelude::*,
|
||||
};
|
||||
use tokio::{
|
||||
sync::{
|
||||
mpsc,
|
||||
RwLock,
|
||||
watch,
|
||||
oneshot,
|
||||
|
||||
RwLockReadGuard,
|
||||
RwLockWriteGuard,
|
||||
},
|
||||
task::{self, JoinHandle},
|
||||
};
|
||||
use std::{
|
||||
sync::{
|
||||
Weak,
|
||||
},
|
||||
fmt,
|
||||
error,
|
||||
};
|
||||
|
||||
/// Command to send to worker task.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum CommandKind
|
||||
{
|
||||
Line(String),
|
||||
LineErr(String),
|
||||
|
||||
Bump(isize),
|
||||
BumpHigh(isize),
|
||||
|
||||
Shutdown,
|
||||
|
||||
Many(Vec<CommandKind>),
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum CommandIter
|
||||
{
|
||||
One(std::iter::Once<CommandKind>),
|
||||
Many(std::vec::IntoIter<CommandKind>),
|
||||
}
|
||||
impl ExactSizeIterator for CommandIter{}
|
||||
|
||||
impl Iterator for CommandIter
|
||||
{
|
||||
type Item = CommandKind;
|
||||
fn next(&mut self) -> Option<Self::Item>
|
||||
{
|
||||
match self {
|
||||
Self::One(one) => one.next(),
|
||||
Self::Many(many) => many.next(),
|
||||
}
|
||||
}
|
||||
|
||||
fn size_hint(&self) -> (usize, Option<usize>)
|
||||
{
|
||||
let sz = match self {
|
||||
Self::One(_) => 1,
|
||||
Self::Many(m) => m.len(),
|
||||
};
|
||||
(sz, Some(sz))
|
||||
}
|
||||
}
|
||||
impl std::iter::FusedIterator for CommandIter{}
|
||||
|
||||
|
||||
impl CommandKind
|
||||
{
|
||||
/// Enumerate all possible commands if this is `Many`.
|
||||
///
|
||||
/// The outputs may still contain `Many`.
|
||||
fn enumerate(self) -> CommandIter
|
||||
{
|
||||
match self {
|
||||
Self::Many(many) => CommandIter::Many(many.into_iter()),
|
||||
other => CommandIter::One(std::iter::once(other)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BarRef<B>(Arc<RwLock<B>>);
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Command(CommandKind, oneshot::Sender<()>);
|
||||
|
||||
/// The bar's state
|
||||
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
|
||||
pub struct State
|
||||
{
|
||||
max: usize,
|
||||
cur: usize,
|
||||
//TODO: Tasks
|
||||
}
|
||||
|
||||
impl State
|
||||
{
|
||||
/// The current progress
|
||||
pub fn prog(&self) -> f64
|
||||
{
|
||||
(self.cur as f64) / (self.max as f64)
|
||||
}
|
||||
}
|
||||
|
||||
/// A handle to a running async progress bar
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Handle<B = Bar>
|
||||
where B: ProgressBar,
|
||||
{
|
||||
// Channel to send commands to the worker
|
||||
chan: mpsc::Sender<Command>,
|
||||
// A weak reference to the worker's bar itself
|
||||
bar: Weak<RwLock<B>>,
|
||||
// A strong reference to the bar's state
|
||||
state: Arc<RwLock<State>>,
|
||||
// Has the worker shut down?
|
||||
dead: watch::Receiver<bool>,
|
||||
}
|
||||
|
||||
impl<B: ProgressBar> Handle<B>
|
||||
{
|
||||
/// Is the worker alive?
|
||||
pub fn is_alive(&self) -> bool
|
||||
{
|
||||
self.bar.strong_count()>0 && !*self.dead.borrow()
|
||||
}
|
||||
|
||||
/// Yields until the worker shutds down gracefully
|
||||
pub async fn closed(&mut self) -> Result<(),WorkerCommError>
|
||||
{
|
||||
loop {
|
||||
match self.dead.recv().await {
|
||||
Some(true) => return Ok(()),
|
||||
None => return Err(WorkerCommError),
|
||||
_ => continue,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a command to the worker.
|
||||
///
|
||||
/// Returns a future that completes to `Ok` when the worker successfully processes the command, and `Err` if the worker exits before processing it
|
||||
pub async fn send_command(&mut self, command: CommandKind) -> Result<impl Future<Output=Result<(), WorkerCommError>>, WorkerCommError>
|
||||
{
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.chan.send(Command(command, tx)).await.map_err(|_| WorkerCommError)?;
|
||||
|
||||
Ok(rx.map(|res| res.map_err(|_| WorkerCommError)))
|
||||
}
|
||||
|
||||
/// Send a command to the worker and then wait for it to be processed
|
||||
pub async fn send_command_and_wait(&mut self, command: CommandKind) -> Result<(), WorkerCommError>
|
||||
{
|
||||
self.send_command(command).await?.await
|
||||
}
|
||||
|
||||
/// Send a command to the worker but do not wait for it to be processed
|
||||
pub async fn send_command_and_detach(&mut self, command: CommandKind) -> Result<(), WorkerCommError>
|
||||
{
|
||||
let _ = self.send_command(command).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get a reference to the state
|
||||
pub async fn state(&self) -> RwLockReadGuard<'_, State>
|
||||
{
|
||||
self.state.read().await
|
||||
}
|
||||
|
||||
/// Get an upgraded reference to the bar itself, if it still exists.
|
||||
///
|
||||
/// This kinda messes things up... Hiding for now.
|
||||
async fn bar_ref(&self) -> Result<BarRef<B>, WorkerCommError>
|
||||
{
|
||||
Ok(BarRef(self.bar.upgrade().ok_or(WorkerCommError)?))
|
||||
}
|
||||
}
|
||||
|
||||
/// Error communicating with worker
|
||||
#[derive(Debug)]
|
||||
pub struct WorkerCommError;
|
||||
|
||||
impl fmt::Display for WorkerCommError
|
||||
{
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
|
||||
{
|
||||
write!(f, "failed to communicate with worker.")
|
||||
}
|
||||
}
|
||||
|
||||
/// Host a progress bar detached
|
||||
pub fn host<B: ProgressBar + Send + Sync + 'static>(bar: B) -> (Handle<B>, JoinHandle<B>)
|
||||
{
|
||||
let state = Arc::new(RwLock::new(Default::default()));
|
||||
let (mut rx, death, bar, handle) = {
|
||||
let (tx, rx) = mpsc::channel(24);
|
||||
let (death, dead) = watch::channel(false);
|
||||
let bar = Arc::new(RwLock::new(bar));
|
||||
let handle = Handle {
|
||||
chan: tx,
|
||||
dead,
|
||||
bar: Arc::downgrade(&bar),
|
||||
state: Arc::clone(&state),
|
||||
};
|
||||
(rx,death,bar,handle)
|
||||
};
|
||||
(handle, tokio::spawn(async move {
|
||||
({
|
||||
macro_rules! update_bar {
|
||||
(to $state:ident $($tt:tt)*) => {
|
||||
{
|
||||
let mut bar = bar.write().await;
|
||||
bar.set_progress($state.prog());
|
||||
update_bar!($($tt)*);
|
||||
}
|
||||
};
|
||||
(write error $line:ident $($tt:tt)*) => {
|
||||
{
|
||||
let bar = bar.read().await;
|
||||
let string = &$line[..];
|
||||
bar.eprintln(string);
|
||||
update_bar!($($tt)*);
|
||||
}
|
||||
};
|
||||
(write $(std)? $line:ident $($tt:tt)*) => {
|
||||
{
|
||||
let bar = bar.read().await;
|
||||
let string = &$line[..];
|
||||
bar.println(string);
|
||||
update_bar!($($tt)*);
|
||||
}
|
||||
};
|
||||
() => {};
|
||||
}
|
||||
while let Some(Command(command, response)) = rx.recv().await {
|
||||
let _response = util::defer(move || response.send(()).ignore());
|
||||
match command {
|
||||
CommandKind::Shutdown => break,
|
||||
CommandKind::BumpHigh(sz) if sz >= 0 => {
|
||||
let mut state = state.write().await;
|
||||
state.max = state.max.saturating_add(sz as usize);
|
||||
|
||||
update_bar!(to state);
|
||||
},
|
||||
CommandKind::BumpHigh(sz) => {
|
||||
debug_assert!(sz <0);
|
||||
let mut state = state.write().await;
|
||||
state.max = state.max.saturating_sub(sz.abs() as usize);
|
||||
|
||||
update_bar!(to state);
|
||||
},
|
||||
CommandKind::Bump(sz) if sz >= 0 => {
|
||||
let mut state = state.write().await;
|
||||
state.cur = state.cur.saturating_add(sz as usize);
|
||||
|
||||
update_bar!(to state);
|
||||
},
|
||||
CommandKind::Bump(sz) => {
|
||||
debug_assert!(sz <0);
|
||||
let mut state = state.write().await;
|
||||
state.cur = state.cur.saturating_sub(sz.abs() as usize);
|
||||
|
||||
update_bar!(to state);
|
||||
},
|
||||
CommandKind::Line(line) => update_bar!(write line),
|
||||
CommandKind::LineErr(line) => update_bar!(write error line),
|
||||
CommandKind::Many(_) => unimplemented!(),
|
||||
}
|
||||
}
|
||||
|
||||
// Consume the bar and return
|
||||
{
|
||||
let mut bar = bar;
|
||||
loop {
|
||||
bar = match Arc::try_unwrap(bar) {
|
||||
Ok(bar) => break bar,
|
||||
Err(bar) => bar,
|
||||
};
|
||||
task::yield_now().await;
|
||||
}.into_inner()
|
||||
}
|
||||
}, death.broadcast(true)).0
|
||||
}))
|
||||
}
|
Loading…
Reference in new issue