hans 0.1.0

use std::{future::Future, marker::PhantomData, pin::Pin, time::Duration};

use ::tokio::time::sleep;
use anyhow::Result;
use chrono::{DateTime, Utc};

/// A connector to interract with a provider.
#[cfg_attr(test, mockall::automock)]
pub trait Executor: Send + Sync {}

/// An orchestrator.
pub struct Orchestrator<EXEC: Executor> {
    exec: EXEC,
    delay: Duration,
}

impl<EXEC: Executor> Orchestrator<EXEC> {
    /// Creates an orchestrator.
    pub fn new(exec: EXEC) -> Self {
        Self {
            exec,
            delay: Duration::from_secs(5),
        }
    }

    /// Orchestrates a task chain until its completion.
    ///
    /// You can configure a delay between two loops execution with
    /// [`Self::with_delay`]. By default, the delay is set to 5 seconds.
    ///
    /// Returns the status of the task chain.
    /// It fails if the executor failed.
    pub async fn orchestrate<TASK: Task<EXEC>>(
        &self,
        chain: &TaskChain<EXEC, TASK>,
    ) -> Result<TaskChainStatus> {
        let mut status = self.orchestrate_once(chain).await?;
        while !status.is_finished() {
            sleep(self.delay).await;
            status = self.orchestrate_once(chain).await?;
        }
        Ok(status)
    }

    /// Orchestrates a task chain only once.
    ///
    /// Returns the status of the task chain.
    /// It fails if the executor failed.
    pub fn orchestrate_once<'a, TASK: Task<EXEC>>(
        &'a self,
        chain: &'a TaskChain<EXEC, TASK>,
    ) -> Pin<Box<dyn Future<Output = Result<TaskChainStatus>> + 'a>> {
        Box::pin(async {
            let mut finished = 0;
            let mut succeeded = 0;
            for task in &chain.tasks {
                let status = task.status(&self.exec).await?;
                self.handle_status(&status, task, &mut finished, &mut succeeded)
                    .await?;
            }
            if finished == chain.tasks.len() {
                if succeeded == finished {
                    if let Some(next) = &chain.next {
                        self.orchestrate_once(next.as_ref()).await
                    } else {
                        Ok(TaskChainStatus::Succeeded)
                    }
                } else {
                    Ok(TaskChainStatus::Failed)
                }
            } else {
                Ok(TaskChainStatus::Running)
            }
        })
    }

    /// Sets the delay used between two loops of orchestration.
    pub fn with_delay(mut self, delay: Duration) -> Self {
        self.delay = delay;
        self
    }

    fn handle_status<'a, TASK: Task<EXEC>>(
        &'a self,
        status: &'a TaskStatus,
        task: &'a TASK,
        finished: &'a mut usize,
        succeeded: &'a mut usize,
    ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
        Box::pin(async move {
            match status {
                TaskStatus::Failed { .. } => {
                    if !task.is_deleted(&self.exec).await? {
                        task.delete(&self.exec).await?;
                    }
                    *finished += 1;
                }
                TaskStatus::Succeeded { .. } => {
                    if !task.is_deleted(&self.exec).await? {
                        task.delete(&self.exec).await?;
                    }
                    *finished += 1;
                    *succeeded += 1;
                }
                TaskStatus::Pending => {
                    let status = task.start(&self.exec).await?;
                    self.handle_status(&status, task, finished, succeeded)
                        .await?;
                }
                _ => {}
            }
            Ok(())
        })
    }
}

/// A task.
#[cfg_attr(test, mockall::automock)]
pub trait Task<EXEC: Executor> {
    /// Deletes this task.
    ///
    /// It fails if the executor failed.
    fn delete(&self, exec: &EXEC) -> impl Future<Output = Result<()>>;

    /// Returns `true` if this task was deleted, `false` otherwise.
    ///
    /// It fails if the executor failed.
    fn is_deleted(&self, exec: &EXEC) -> impl Future<Output = Result<bool>>;

    /// Starts this task.
    ///
    /// Returns the task status.
    /// It fails if the executor failed.
    fn start(&self, exec: &EXEC) -> impl Future<Output = Result<TaskStatus>>;

    /// Returns the status of this task.
    ///
    /// It fails if the executor failed.
    fn status(&self, exec: &EXEC) -> impl Future<Output = Result<TaskStatus>>;
}

/// A task chain.
///
/// It's composed by node.
/// Each node has tasks that can be executed in parallel.
/// If all tasks of the node are succeeded, the next node is started.
///
/// ![task-chain](https://raw.githubusercontent.com/leroyguillaume/hans/main/assets/task-chain.svg)
pub struct TaskChain<EXEC: Executor, TASK: Task<EXEC>> {
    next: Option<Box<TaskChain<EXEC, TASK>>>,
    tasks: Vec<TASK>,
    _exec: PhantomData<EXEC>,
}

impl<EXEC: Executor, TASK: Task<EXEC>> TaskChain<EXEC, TASK> {
    /// Creates a task chain builder.
    ///
    /// It's an alias for [TaskChainBuilder::new].
    ///
    /// - `tasks`: the tasks of the first node of the chain.
    pub fn builder<TASKS: IntoIterator<Item = TASK>>(tasks: TASKS) -> TaskChainBuilder<EXEC, TASK> {
        TaskChainBuilder::new(tasks)
    }
}

/// A task chain builder.
pub struct TaskChainBuilder<EXEC: Executor, TASK: Task<EXEC>> {
    parent: Option<Box<TaskChainBuilder<EXEC, TASK>>>,
    tasks: Vec<TASK>,
    _exec: PhantomData<EXEC>,
}

impl<EXEC: Executor, TASK: Task<EXEC>> TaskChainBuilder<EXEC, TASK> {
    /// Creates a task chain builder.
    ///
    /// - `tasks`: the tasks of the first node of the chain.
    pub fn new<TASKS: IntoIterator<Item = TASK>>(tasks: TASKS) -> Self {
        Self {
            parent: None,
            tasks: tasks.into_iter().collect(),
            _exec: PhantomData,
        }
    }

    /// Buils the task chain.
    pub fn build(self) -> TaskChain<EXEC, TASK> {
        self.build_recursively(None)
    }

    /// Creates a new node that depends of the last one.
    ///
    /// - `tasks`: the tasks of the new node of the chain.
    pub fn then<TASKS: IntoIterator<Item = TASK>>(self, tasks: TASKS) -> Self {
        Self {
            parent: Some(Box::new(self)),
            tasks: tasks.into_iter().collect(),
            _exec: PhantomData,
        }
    }

    fn build_recursively(self, next: Option<TaskChain<EXEC, TASK>>) -> TaskChain<EXEC, TASK> {
        let chain = TaskChain {
            next: next.map(Box::new),
            tasks: self.tasks,
            _exec: PhantomData,
        };
        if let Some(parent) = self.parent {
            parent.build_recursively(Some(chain))
        } else {
            chain
        }
    }
}

/// A status of a task chain.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TaskChainStatus {
    /// The task chain was failed.
    Failed,
    /// The task chain is running.
    Running,
    /// The task chain was succeeded.
    Succeeded,
}

impl TaskChainStatus {
    /// Returns `true` if this task status is `Succeeded` or `Failed`, `false` otherwise.
    pub fn is_finished(&self) -> bool {
        matches!(self, TaskChainStatus::Failed | TaskChainStatus::Succeeded)
    }
}

/// A status for a task.
///
/// ![task-status](https://raw.githubusercontent.com/leroyguillaume/hans/main/assets/task-status.svg)
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TaskStatus {
    /// The task was failed.
    Failed {
        finished_at: DateTime<Utc>,
        started_at: DateTime<Utc>,
    },
    /// The task is waiting for its dependencies.
    Pending,
    /// The task is waiting for be started.
    Provisioning,
    /// The task is running.
    Running { started_at: DateTime<Utc> },
    /// The task was succeeded.
    Succeeded {
        finished_at: DateTime<Utc>,
        started_at: DateTime<Utc>,
    },
}

pub mod tokio;

#[cfg(test)]
mod test {
    use chrono::{Duration, Utc};

    use super::{MockExecutor, MockTask, Orchestrator, TaskChain, TaskChainStatus, TaskStatus};

    mod orchestrator {
        use super::*;

        mod orchestrate_once {
            use super::*;

            #[tokio::test]
            async fn start_a_b() {
                let mut a = MockTask::new();
                a.expect_start()
                    .times(1)
                    .returning(move |_| Box::pin(async move { Ok(TaskStatus::Provisioning) }));
                a.expect_status()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(TaskStatus::Pending) }));
                let mut b = MockTask::new();
                b.expect_start()
                    .times(1)
                    .returning(move |_| Box::pin(async move { Ok(TaskStatus::Provisioning) }));
                b.expect_status()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(TaskStatus::Pending) }));
                let c = MockTask::new();
                let chain = TaskChain::builder([a, b]).then([c]).build();
                let orch = Orchestrator::new(MockExecutor::new());
                let status = orch
                    .orchestrate_once(&chain)
                    .await
                    .expect("failed to orchestrate");
                assert_eq!(status, TaskChainStatus::Running);
            }

            #[tokio::test]
            async fn a_b_succeeded_instantly() {
                let mut a = MockTask::new();
                a.expect_delete()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(()) }));
                a.expect_is_deleted()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(false) }));
                a.expect_start().times(1).returning(move |_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Succeeded {
                            finished_at: Utc::now(),
                            started_at: Utc::now() - Duration::seconds(5),
                        })
                    })
                });
                a.expect_status()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(TaskStatus::Pending) }));
                let mut b = MockTask::new();
                b.expect_is_deleted()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(true) }));
                b.expect_start().times(1).returning(move |_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Succeeded {
                            finished_at: Utc::now(),
                            started_at: Utc::now() - Duration::seconds(5),
                        })
                    })
                });
                b.expect_status()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(TaskStatus::Pending) }));
                let mut c = MockTask::new();
                c.expect_start()
                    .times(1)
                    .returning(move |_| Box::pin(async move { Ok(TaskStatus::Provisioning) }));
                c.expect_status()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(TaskStatus::Pending) }));
                let chain = TaskChain::builder([a, b]).then([c]).build();
                let orch = Orchestrator::new(MockExecutor::new());
                let status = orch
                    .orchestrate_once(&chain)
                    .await
                    .expect("failed to orchestrate");
                assert_eq!(status, TaskChainStatus::Running);
            }

            #[tokio::test]
            async fn a_failed_instantly_b_succeeded() {
                let mut a = MockTask::new();
                a.expect_delete()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(()) }));
                a.expect_is_deleted()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(false) }));
                a.expect_start().times(1).returning(move |_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Failed {
                            finished_at: Utc::now(),
                            started_at: Utc::now() - Duration::seconds(5),
                        })
                    })
                });
                a.expect_status()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(TaskStatus::Pending) }));
                let mut b = MockTask::new();
                b.expect_is_deleted()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(false) }));
                b.expect_status().times(1).returning(|_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Succeeded {
                            finished_at: Utc::now(),
                            started_at: Utc::now() - Duration::seconds(5),
                        })
                    })
                });
                b.expect_delete()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(()) }));
                let c = MockTask::new();
                let chain = TaskChain::builder([a, b]).then([c]).build();
                let orch = Orchestrator::new(MockExecutor::new());
                let status = orch
                    .orchestrate_once(&chain)
                    .await
                    .expect("failed to orchestrate");
                assert_eq!(status, TaskChainStatus::Failed);
            }

            #[tokio::test]
            async fn a_still_running_start_b() {
                let mut a = MockTask::new();
                a.expect_status().times(1).returning(move |_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Running {
                            started_at: Utc::now(),
                        })
                    })
                });
                let mut b = MockTask::new();
                b.expect_start()
                    .times(1)
                    .returning(move |_| Box::pin(async move { Ok(TaskStatus::Provisioning) }));
                b.expect_status()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(TaskStatus::Pending) }));
                let c = MockTask::new();
                let chain = TaskChain::builder([a, b]).then([c]).build();
                let orch = Orchestrator::new(MockExecutor::new());
                let status = orch
                    .orchestrate_once(&chain)
                    .await
                    .expect("failed to orchestrate");
                assert_eq!(status, TaskChainStatus::Running);
            }

            #[tokio::test]
            async fn a_still_running_b_succeeded() {
                let mut a = MockTask::new();
                a.expect_status().times(1).returning(move |_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Running {
                            started_at: Utc::now(),
                        })
                    })
                });
                let mut b = MockTask::new();
                b.expect_is_deleted()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(true) }));
                b.expect_status().times(1).returning(move |_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Succeeded {
                            finished_at: Utc::now(),
                            started_at: Utc::now() - Duration::seconds(5),
                        })
                    })
                });
                let c = MockTask::new();
                let chain = TaskChain::builder([a, b]).then([c]).build();
                let orch = Orchestrator::new(MockExecutor::new());
                let status = orch
                    .orchestrate_once(&chain)
                    .await
                    .expect("failed to orchestrate");
                assert_eq!(status, TaskChainStatus::Running);
            }

            #[tokio::test]
            async fn start_c() {
                let mut a = MockTask::new();
                a.expect_is_deleted()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(true) }));
                a.expect_status().times(1).returning(move |_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Succeeded {
                            finished_at: Utc::now(),
                            started_at: Utc::now() - Duration::seconds(5),
                        })
                    })
                });
                let mut b = MockTask::new();
                b.expect_is_deleted()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(true) }));
                b.expect_status().times(1).returning(move |_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Succeeded {
                            finished_at: Utc::now(),
                            started_at: Utc::now() - Duration::seconds(5),
                        })
                    })
                });
                let mut c = MockTask::new();
                c.expect_start()
                    .times(1)
                    .returning(move |_| Box::pin(async move { Ok(TaskStatus::Provisioning) }));
                c.expect_status()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(TaskStatus::Pending) }));
                let chain = TaskChain::builder([a, b]).then([c]).build();
                let orch = Orchestrator::new(MockExecutor::new());
                let status = orch
                    .orchestrate_once(&chain)
                    .await
                    .expect("failed to orchestrate");
                assert_eq!(status, TaskChainStatus::Running);
            }

            #[tokio::test]
            async fn a_failed() {
                let mut a = MockTask::new();
                a.expect_delete()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(()) }));
                a.expect_is_deleted()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(false) }));
                a.expect_status().times(1).returning(move |_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Failed {
                            finished_at: Utc::now(),
                            started_at: Utc::now() - Duration::seconds(5),
                        })
                    })
                });
                let mut b = MockTask::new();
                b.expect_delete()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(()) }));
                b.expect_is_deleted()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(false) }));
                b.expect_status().times(1).returning(move |_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Succeeded {
                            finished_at: Utc::now(),
                            started_at: Utc::now() - Duration::seconds(5),
                        })
                    })
                });
                let c = MockTask::new();
                let chain = TaskChain::builder([a, b]).then([c]).build();
                let orch = Orchestrator::new(MockExecutor::new());
                let status = orch
                    .orchestrate_once(&chain)
                    .await
                    .expect("failed to orchestrate");
                assert_eq!(status, TaskChainStatus::Failed);
            }

            #[tokio::test]
            async fn c_succeeded() {
                let mut a = MockTask::new();
                a.expect_is_deleted()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(true) }));
                a.expect_status().times(1).returning(move |_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Succeeded {
                            finished_at: Utc::now(),
                            started_at: Utc::now() - Duration::seconds(5),
                        })
                    })
                });
                let mut b = MockTask::new();
                b.expect_is_deleted()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(true) }));
                b.expect_status().times(1).returning(move |_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Succeeded {
                            finished_at: Utc::now(),
                            started_at: Utc::now() - Duration::seconds(5),
                        })
                    })
                });
                let mut c = MockTask::new();
                c.expect_delete()
                    .times(1)
                    .returning(|_| Box::pin(async move { Ok(()) }));
                c.expect_is_deleted()
                    .times(1)
                    .returning(|_| Box::pin(async { Ok(false) }));
                c.expect_status().times(1).returning(|_| {
                    Box::pin(async move {
                        Ok(TaskStatus::Succeeded {
                            finished_at: Utc::now(),
                            started_at: Utc::now() - Duration::seconds(5),
                        })
                    })
                });
                let chain = TaskChain::builder([a, b]).then([c]).build();
                let orch = Orchestrator::new(MockExecutor::new());
                let status = orch
                    .orchestrate_once(&chain)
                    .await
                    .expect("failed to orchestrate");
                assert_eq!(status, TaskChainStatus::Succeeded);
            }
        }
    }
}