unki/lib/u_lib/src/executor.rs

// list of jobs: job (cmd, args) OR rust fn OR python func + cron-like timing tokoi::time::interval:
// job runner (thread)
// every job runs in other thread/process

use crate::{models::*, UResult, UError, OneOrMany};

use std::collections::HashMap;

use futures::{lock::{Mutex, MutexGuard}, future::BoxFuture};
use lazy_static::lazy_static;
use tokio::{
    spawn,
    task::JoinHandle,
    sync::mpsc::{channel, Receiver, Sender}
};
use uuid::Uuid;

pub type FutRes = UResult<JobResult>;
pub type DynFut = BoxFuture<'static, FutRes>;

lazy_static! {
    static ref FUT_RESULTS: Mutex<HashMap<Uuid, JoinInfo>> = Mutex::new(HashMap::new());
    static ref FUT_CHANNEL: (Mutex<Sender<Uuid>>, Mutex<Receiver<Uuid>>) = {
        spawn(init_receiver());
        let (tx, rx) = channel(100);
        (Mutex::new(tx), Mutex::new(rx))
    };
}

async fn get_static_channel() -> (Sender<Uuid>, MutexGuard<'static, Receiver<Uuid>>) {
    (
        FUT_CHANNEL.0.lock().await.clone(),
        FUT_CHANNEL.1.lock().await
    )
}

struct JoinInfo {
    handle: JoinHandle<FutRes>,
    completed: bool,
    collectable: bool // indicates if future can be popped from pool via pop_task_if_completed
}

pub struct Waiter {
    tasks: Vec<DynFut>,
    fids: Vec<Uuid>
}

impl Waiter {
    pub fn new<S: OneOrMany<DynFut>>(tasks: S) -> Self {
        Self {
            tasks: tasks.into_vec(),
            fids: vec![]
        }
    }

    pub async fn spawn(mut self) -> Self {
        let collectable = self.tasks.len() != 1;
        for f in self.tasks.drain(..) {
            //eprintln!("before static channel");
            let tx = get_static_channel().await.0;
            //eprintln!("after static channel");
            let fid = Uuid::new_v4();
            self.fids.push(fid);
            let task_wrapper = async move {
                //eprintln!("inside wrapper (started): {}", fid);
                let result = f.await;
                tx.send(fid).await.unwrap();
                result
            };
            let handle = JoinInfo {
                handle: spawn(task_wrapper),
                completed: false,
                collectable
            };
            //eprintln!("before push: {}", fid);
            spawn(async {}).await.ok();
            FUT_RESULTS.lock().await.insert(fid, handle);
        }
        self
    }

    // wait until a bunch of tasks is finished
    // NOT GUARANTEED that all tasks will be returned due to
    // possibility to pop them in other places
    pub async fn wait(self) -> Vec<FutRes> {
        let mut result = vec![];
        for fid in self.fids {
            if let Some(task) = pop_task(fid).await {
                let r = task.handle.await;
                result.push(r.unwrap());
            }
        }
        result
    }

    pub async fn run_until_complete(self) -> Vec<FutRes> {
        self.spawn().await.wait().await
    }

    pub async fn run_one_until_complete(self) -> FutRes {
        self.run_until_complete().await.pop().unwrap()
    }
}

async fn pop_task(fid: Uuid) -> Option<JoinInfo> {
    FUT_RESULTS.lock().await.remove(&fid)
}

async fn init_receiver() {
    while let Some(fid) = get_static_channel().await.1.recv().await {
        eprintln!("task {} is done", fid);
        if let Some(j) = FUT_RESULTS.lock().await.get_mut(&fid) {
            j.completed = true;
        }
    }
}

pub async fn pop_task_if_completed(fid: Uuid) -> Option<FutRes> {
    let &mut JoinInfo {handle: _, collectable, completed} = match FUT_RESULTS
        .lock()
        .await
        .get_mut(&fid) {
        Some(t) => t,
        None => return None
    };
    //eprint!("{}, {}: ", &fid, *collectable);
    if collectable && completed {
        let task = pop_task(fid).await.unwrap();
        let result = task.handle.await.unwrap();
        Some(result)
    } else {
        None
    }
}

pub async fn pop_completed() -> Vec<FutRes> {
    let mut completed: Vec<FutRes> = vec![];
    let fids = FUT_RESULTS.lock()
        .await
        .keys()
        .map(|k| *k)
        .collect::<Vec<Uuid>>();
    for fid in fids {
        if let Some(r) = pop_task_if_completed(fid).await {
            completed.push(r);
        }
    }
    completed
}

#[cfg(test)]
mod tests {
    use super::*;

    // WTF
    // WTF
    // WTF
    #[tokio::test]
    async fn test_spawn() {
        use std::sync::Arc;

        let val = Arc::new(Mutex::new(0));
        let t = {
            let v = val.clone();
            spawn(async move {
                *v.lock().await = 5;
            })
        };
        assert_eq!(0, *val.lock().await);
        spawn(async {}).await.ok();
        assert_eq!(5, *val.lock().await);
        t.await.ok();
        assert_eq!(5, *val.lock().await);
    }
}