|
|
@@ -0,0 +1,266 @@
|
|
|
+use anyhow::{Result, anyhow };
|
|
|
+use futures::{
|
|
|
+ stream::{AbortHandle, Abortable},
|
|
|
+ Future,
|
|
|
+};
|
|
|
+use std::{
|
|
|
+ collections::HashMap,
|
|
|
+ sync::{Arc, Mutex},
|
|
|
+};
|
|
|
+use tokio::{sync::mpsc::{self, Sender}, task::JoinHandle};
|
|
|
+
|
|
|
+use crate::utils::Run;
|
|
|
+
|
|
|
+pub struct Runner<R: Run> {
|
|
|
+ task_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
|
|
|
+ abort_handle: Option<AbortHandle>,
|
|
|
+ progress: Arc<Mutex<Box<dyn Inc + Send>>>,
|
|
|
+ result_sender: mpsc::Sender<anyhow::Result<R>>,
|
|
|
+ // result: Arc<Mutex<Box<Option<anyhow::Result<R>>>>>,
|
|
|
+}
|
|
|
+
|
|
|
+pub struct BasicProgress {
|
|
|
+ pub length: usize,
|
|
|
+ pub step: usize,
|
|
|
+}
|
|
|
+
|
|
|
+pub trait Inc {
|
|
|
+ fn inc(&mut self, delta: usize);
|
|
|
+}
|
|
|
+
|
|
|
+impl Inc for BasicProgress {
|
|
|
+ fn inc(&mut self, delta: usize) {
|
|
|
+ if self.step < self.length {
|
|
|
+ self.step += delta;
|
|
|
+ }
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+pub struct IndProgress {
|
|
|
+ inner: indicatif::ProgressBar,
|
|
|
+}
|
|
|
+
|
|
|
+impl Inc for IndProgress {
|
|
|
+ fn inc(&mut self, delta: usize) {
|
|
|
+ self.inner.inc(delta as u64)
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+impl IndProgress {
|
|
|
+ pub fn new(len: u64) -> Self {
|
|
|
+ Self {
|
|
|
+ inner: indicatif::ProgressBar::new(len),
|
|
|
+ }
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+impl<R: 'static + Run> Runner<R> {
|
|
|
+ pub fn new<P: 'static + Inc + Send>(progress: P) -> Self {
|
|
|
+ let (result_sender, _) = mpsc::channel(1);
|
|
|
+ Runner {
|
|
|
+ task_handle: Arc::new(Mutex::new(None)),
|
|
|
+ abort_handle: None,
|
|
|
+ progress: Arc::new(Mutex::new(Box::new(progress))),
|
|
|
+ result_sender,
|
|
|
+ // result: Arc::new(Mutex::new(Box::new(None))),
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ pub fn spawn<F, Fut>(&mut self, task: F)
|
|
|
+ where
|
|
|
+ F: FnOnce(Arc<Mutex<Box<dyn Inc + Send>>>, Sender<Result<R>>) ,
|
|
|
+ Fut: Future<Output = ()> + Send + 'static,
|
|
|
+ {
|
|
|
+ let (abort_handle, abort_registration) = AbortHandle::new_pair();
|
|
|
+ let progress = self.progress.clone();
|
|
|
+ let result_sender = self.result_sender.clone();
|
|
|
+
|
|
|
+ let task = Abortable::new(
|
|
|
+ tokio::task::spawn_blocking( move || {
|
|
|
+ let result_sender = self.result_sender.clone();
|
|
|
+ task(progress, result_sender);
|
|
|
+ }),
|
|
|
+ abort_registration,
|
|
|
+ );
|
|
|
+
|
|
|
+ let task_handle = self.task_handle.clone();
|
|
|
+ let handle = tokio::spawn(async move {
|
|
|
+ match task.await {
|
|
|
+ Ok(_) => println!("Task completed successfully"),
|
|
|
+ Err(_) => println!("Task aborted"),
|
|
|
+ }
|
|
|
+ *task_handle.lock().unwrap() = None;
|
|
|
+ });
|
|
|
+
|
|
|
+ *self.task_handle.lock().unwrap() = Some(handle);
|
|
|
+ self.abort_handle = Some(abort_handle);
|
|
|
+ }
|
|
|
+
|
|
|
+ pub async fn has_finished(&self) -> bool {
|
|
|
+ let mut guard = self.task_handle.lock().unwrap();
|
|
|
+ if let Some(handle) = guard.take() {
|
|
|
+ if handle.await.is_ok() {
|
|
|
+ return true;
|
|
|
+ }
|
|
|
+ }
|
|
|
+ false
|
|
|
+ }
|
|
|
+
|
|
|
+ pub fn abort(&mut self) {
|
|
|
+ if let Some(abort_handle) = self.abort_handle.take() {
|
|
|
+ abort_handle.abort();
|
|
|
+ }
|
|
|
+ *self.task_handle.lock().unwrap() = None;
|
|
|
+ }
|
|
|
+
|
|
|
+ pub fn get_progress(&self) -> Arc<Mutex<Box<dyn Inc + Send>>> {
|
|
|
+ self.progress.clone()
|
|
|
+ }
|
|
|
+
|
|
|
+ fn get_result(&self) -> anyhow::Result<R> {
|
|
|
+ if let Some(r) = Arc::into_inner(self.result.clone()) {
|
|
|
+ if let Ok(r) = r.into_inner() {
|
|
|
+ if let Some(r) = r {
|
|
|
+ return r;
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ Err(anyhow!("Error while getting the results."))
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+pub struct Runners<R: Run> {
|
|
|
+ runners: HashMap<String, Runner<R>>,
|
|
|
+}
|
|
|
+
|
|
|
+impl<R: 'static + Run + Send> Runners<R> {
|
|
|
+ pub fn new() -> Self {
|
|
|
+ Runners {
|
|
|
+ runners: HashMap::new(),
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ pub fn add_runner(&mut self, id: &str, runner: Runner<R>) {
|
|
|
+ self.runners.insert(id.to_string(), runner);
|
|
|
+ }
|
|
|
+
|
|
|
+ pub fn spawn_task<F, Fut>(&mut self, id: &str, task: F)
|
|
|
+ where
|
|
|
+ F: FnOnce(Arc<Mutex<Box<dyn Inc + Send>>>) -> Fut + Send + 'static,
|
|
|
+ Fut: Future<Output = anyhow::Result<R>> + Send + 'static,
|
|
|
+ {
|
|
|
+ if let Some(actor) = self.runners.get_mut(id) {
|
|
|
+ actor.spawn(task);
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ pub async fn has_finished(&self, id: &str) -> bool {
|
|
|
+ if let Some(runner) = self.runners.get(id) {
|
|
|
+ return runner.has_finished().await;
|
|
|
+ }
|
|
|
+ false
|
|
|
+ }
|
|
|
+
|
|
|
+ pub fn abort_task(&mut self, id: &str) {
|
|
|
+ if let Some(runner) = self.runners.get_mut(id) {
|
|
|
+ runner.abort();
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ pub fn get_progress(&self, id: &str) -> Option<Arc<Mutex<Box<dyn Inc + Send>>>> {
|
|
|
+ self.runners.get(id).map(|runner| runner.get_progress())
|
|
|
+ }
|
|
|
+
|
|
|
+ pub async fn remove_finished_tasks(&mut self) {
|
|
|
+ let mut finished_ids = Vec::new();
|
|
|
+
|
|
|
+ for (id, runner) in &self.runners {
|
|
|
+ if runner.has_finished().await {
|
|
|
+ finished_ids.push(id.clone());
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ for id in finished_ids {
|
|
|
+ self.runners.remove(&id);
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ pub fn get_result(&self, id: &str) -> anyhow::Result<R> {
|
|
|
+ if let Some(r) = self.runners.get(id).map(|actor| actor.get_result()) {
|
|
|
+ return r;
|
|
|
+ } else {
|
|
|
+ return Err(anyhow!("Error while getting results."));
|
|
|
+ }
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+// #[tokio::main]
|
|
|
+// async fn main() {
|
|
|
+// let mut runners = Runners::new();
|
|
|
+//
|
|
|
+// let progress1 = Box::new(Progress { length: 10, step: 0 });
|
|
|
+// let actor1 = Runner::new(progress1);
|
|
|
+// runners.add_actor("actor1".to_string(), actor1);
|
|
|
+//
|
|
|
+// let progress2 = Box::new(Counter { count: 0, max: 10 });
|
|
|
+// let actor2 = Runner::new(progress2);
|
|
|
+// runners.add_actor("actor2".to_string(), actor2);
|
|
|
+//
|
|
|
+// runners.spawn_task("actor1", |progress| {
|
|
|
+// println!("Task for actor1 started");
|
|
|
+// for _ in 0..10 {
|
|
|
+// thread::sleep(Duration::from_secs(1));
|
|
|
+// let mut progress = progress.lock().unwrap();
|
|
|
+// progress.inc();
|
|
|
+// if let Some(p) = progress.as_any().downcast_ref::<Progress>() {
|
|
|
+// println!("actor1 Progress: {}/{}", p.step, p.length);
|
|
|
+// }
|
|
|
+// }
|
|
|
+// println!("Task for actor1 finished");
|
|
|
+// });
|
|
|
+//
|
|
|
+// runners.spawn_task("actor2", |progress| {
|
|
|
+// println!("Task for actor2 started");
|
|
|
+// for _ in 0..10 {
|
|
|
+// thread::sleep(Duration::from_secs(1));
|
|
|
+// let mut progress = progress.lock().unwrap();
|
|
|
+// progress.inc();
|
|
|
+// if let Some(p) = progress.as_any().downcast_ref::<Counter>() {
|
|
|
+// println!("actor2 Counter: {}/{}", p.count, p.max);
|
|
|
+// }
|
|
|
+// }
|
|
|
+// println!("Task for actor2 finished");
|
|
|
+// });
|
|
|
+//
|
|
|
+// tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
|
|
|
+//
|
|
|
+// if let Some(progress) = runners.get_progress("actor1") {
|
|
|
+// let progress = progress.lock().unwrap();
|
|
|
+// if let Some(p) = progress.as_any().downcast_ref::<Progress>() {
|
|
|
+// println!("Current progress for actor1: {}/{}", p.step, p.length);
|
|
|
+// }
|
|
|
+// }
|
|
|
+//
|
|
|
+// if let Some(progress) = runners.get_progress("actor2") {
|
|
|
+// let progress = progress.lock().unwrap();
|
|
|
+// if let Some(p) = progress.as_any().downcast_ref::<Counter>()) {
|
|
|
+// println!("Current progress for actor2: {}/{}", p.count, p.max);
|
|
|
+// }
|
|
|
+// }
|
|
|
+//
|
|
|
+// if runners.has_finished("actor1").await {
|
|
|
+// println!("Task for actor1 has already finished");
|
|
|
+// } else {
|
|
|
+// println!("Task for actor1 is still running");
|
|
|
+// runners.abort_task("actor1");
|
|
|
+// println!("Task for actor1 aborted");
|
|
|
+// }
|
|
|
+//
|
|
|
+// if runners.has_finished("actor2").await {
|
|
|
+// println!("Task for actor2 has already finished");
|
|
|
+// } else {
|
|
|
+// println!("Task for actor2 is still running");
|
|
|
+// runners.abort_task("actor2");
|
|
|
+// println!("Task for actor2 aborted");
|
|
|
+// }
|
|
|
+// }
|