hotshot_task/

dependency.rs

// Copyright (c) 2021-2024 Espresso Systems (espressosys.com)
// This file is part of the HotShot repository.

// You should have received a copy of the MIT License
// along with the HotShot repository. If not, see <https://mit-license.org/>.

use std::future::Future;

use async_broadcast::{Receiver, RecvError};
use futures::{
    FutureExt,
    future::BoxFuture,
    stream::{FuturesUnordered, StreamExt},
};

/// Type which describes the idea of waiting for a dependency to complete
pub trait Dependency<T> {
    /// Complete will wait until it gets some value `T` then return the value
    fn completed(self) -> impl Future<Output = Option<T>> + Send;
    /// Create an or dependency from this dependency and another
    fn or<D: Dependency<T> + Send + 'static>(self, dep: D) -> OrDependency<T>
    where
        T: Send + Sync + Clone + 'static,
        Self: Sized + Send + 'static,
    {
        let mut or = OrDependency::from_deps(vec![self]);
        or.add_dep(dep);
        or
    }
    /// Create an and dependency from this dependency and another
    fn and<D: Dependency<T> + Send + 'static>(self, dep: D) -> AndDependency<T>
    where
        T: Send + Sync + Clone + 'static,
        Self: Sized + Send + 'static,
    {
        let mut and = AndDependency::from_deps(vec![self]);
        and.add_dep(dep);
        and
    }
}

/// Defines a dependency that completes when all of its deps complete
pub struct AndDependency<T> {
    /// Dependencies being combined
    deps: Vec<BoxFuture<'static, Option<T>>>,
}
impl<T: Clone + Send + Sync> Dependency<Vec<T>> for AndDependency<T> {
    /// Returns a vector of all of the results from it's dependencies.
    /// The results will be in a random order
    async fn completed(self) -> Option<Vec<T>> {
        let futures = FuturesUnordered::from_iter(self.deps);
        futures
            .collect::<Vec<Option<T>>>()
            .await
            .into_iter()
            .collect()
    }
}

impl<T: Clone + Send + Sync + 'static> AndDependency<T> {
    /// Create from a vec of deps
    #[must_use]
    pub fn from_deps(deps: Vec<impl Dependency<T> + Send + 'static>) -> Self {
        let mut pinned = vec![];
        for dep in deps {
            pinned.push(dep.completed().boxed());
        }
        Self { deps: pinned }
    }
    /// Add another dependency
    pub fn add_dep(&mut self, dep: impl Dependency<T> + Send + 'static) {
        self.deps.push(dep.completed().boxed());
    }
    /// Add multiple dependencies
    pub fn add_deps(&mut self, deps: AndDependency<T>) {
        for dep in deps.deps {
            self.deps.push(dep);
        }
    }
}

/// Defines a dependency that completes when one of it's dependencies completes
pub struct OrDependency<T> {
    /// Dependencies being combined
    deps: Vec<BoxFuture<'static, Option<T>>>,
}
impl<T: Clone + Send + Sync> Dependency<T> for OrDependency<T> {
    /// Returns the value of the first completed dependency
    async fn completed(self) -> Option<T> {
        let mut futures = FuturesUnordered::from_iter(self.deps);
        loop {
            let maybe = futures.next().await?;
            if maybe.is_some() {
                return maybe;
            }
        }
    }
}

impl<T: Clone + Send + Sync + 'static> OrDependency<T> {
    /// Creat an `OrDependency` from a vec of dependencies
    #[must_use]
    pub fn from_deps(deps: Vec<impl Dependency<T> + Send + 'static>) -> Self {
        let mut pinned = vec![];
        for dep in deps {
            pinned.push(dep.completed().boxed());
        }
        Self { deps: pinned }
    }
    /// Add another dependency
    pub fn add_dep(&mut self, dep: impl Dependency<T> + Send + 'static) {
        self.deps.push(dep.completed().boxed());
    }
}

/// A dependency that listens on a channel for an event
/// that matches what some value it wants.
pub struct EventDependency<T: Clone + Send + Sync> {
    /// Channel of incoming events
    pub(crate) event_rx: Receiver<T>,

    /// Closure which returns true if the incoming `T` is the
    /// thing that completes this dependency
    pub(crate) match_fn: Box<dyn Fn(&T) -> bool + Send>,

    /// The potentially externally completed dependency. If the dependency was seeded from an event
    /// message, we can mark it as already done in lieu of other events still pending.
    completed_dependency: Option<T>,

    cancel_receiver: Receiver<()>,

    dependency_name: String,
}

impl<T: Clone + Send + Sync + 'static> EventDependency<T> {
    /// Create a new `EventDependency`
    #[must_use]
    pub fn new(
        receiver: Receiver<T>,
        cancel_receiver: Receiver<()>,
        dependency_name: String,
        match_fn: Box<dyn Fn(&T) -> bool + Send>,
    ) -> Self {
        Self {
            event_rx: receiver,
            match_fn: Box::new(match_fn),
            completed_dependency: None,
            cancel_receiver,
            dependency_name,
        }
    }

    /// Mark a dependency as completed.
    pub fn mark_as_completed(&mut self, dependency: T) {
        self.completed_dependency = Some(dependency);
    }
}

impl<T: Clone + Send + Sync + 'static> Dependency<T> for EventDependency<T> {
    async fn completed(mut self) -> Option<T> {
        if let Some(dependency) = self.completed_dependency {
            return Some(dependency);
        }
        loop {
            if let Some(dependency) = self.completed_dependency {
                return Some(dependency);
            }

            tokio::select! {
                recv_event = self.event_rx.recv() => {
                    match recv_event {
                        Ok(event) => {
                            if (self.match_fn)(&event) {
                                return Some(event);
                            }
                        },
                        Err(RecvError::Overflowed(n)) => {
                            tracing::error!("Dependency Task overloaded, skipping {} events", n);
                        },
                        Err(RecvError::Closed) => {
                            return None;
                        },
                    }
                }
                _ = self.cancel_receiver.recv() => {
                   tracing::warn!("{} dependency cancelled", self.dependency_name);
                   return None;
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use async_broadcast::{Receiver, broadcast};

    use super::{AndDependency, Dependency, EventDependency, OrDependency};

    fn eq_dep(
        rx: Receiver<usize>,
        cancel_rx: Receiver<()>,
        dep_name: String,
        val: usize,
    ) -> EventDependency<usize> {
        EventDependency {
            event_rx: rx,
            match_fn: Box::new(move |v| *v == val),
            completed_dependency: None,
            dependency_name: dep_name,
            cancel_receiver: cancel_rx,
        }
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn it_works() {
        let (tx, rx) = broadcast(10);
        let (_cancel_tx, cancel_rx) = broadcast(1);

        let mut deps = vec![];
        for i in 0..5 {
            tx.broadcast(i).await.unwrap();
            deps.push(eq_dep(
                rx.clone(),
                cancel_rx.clone(),
                format!("it_works {i}"),
                5,
            ));
        }

        let and = AndDependency::from_deps(deps);
        tx.broadcast(5).await.unwrap();
        let result = and.completed().await;
        assert_eq!(result, Some(vec![5; 5]));
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn or_dep() {
        let (tx, rx) = broadcast(10);
        let (_cancel_tx, cancel_rx) = broadcast(1);

        tx.broadcast(5).await.unwrap();
        let mut deps = vec![];
        for i in 0..5 {
            deps.push(eq_dep(
                rx.clone(),
                cancel_rx.clone(),
                format!("or_dep {i}"),
                5,
            ));
        }
        let or = OrDependency::from_deps(deps);
        let result = or.completed().await;
        assert_eq!(result, Some(5));
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn and_or_dep() {
        let (tx, rx) = broadcast(10);
        let (_cancel_tx, cancel_rx) = broadcast(1);

        tx.broadcast(1).await.unwrap();
        tx.broadcast(2).await.unwrap();
        tx.broadcast(3).await.unwrap();
        tx.broadcast(5).await.unwrap();
        tx.broadcast(6).await.unwrap();

        let or1 = OrDependency::from_deps(
            [
                eq_dep(
                    rx.clone(),
                    cancel_rx.clone(),
                    format!("and_or_dep or1 {}", 4),
                    4,
                ),
                eq_dep(
                    rx.clone(),
                    cancel_rx.clone(),
                    format!("and_or_dep or1 {}", 6),
                    6,
                ),
            ]
            .into(),
        );
        let or2 = OrDependency::from_deps(
            [
                eq_dep(
                    rx.clone(),
                    cancel_rx.clone(),
                    format!("and_or_dep or2 {}", 4),
                    4,
                ),
                eq_dep(
                    rx.clone(),
                    cancel_rx.clone(),
                    format!("and_or_dep or2 {}", 5),
                    5,
                ),
            ]
            .into(),
        );
        let and = AndDependency::from_deps([or1, or2].into());
        let result = and.completed().await;
        assert_eq!(result, Some(vec![6, 5]));
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn or_and_dep() {
        let (tx, rx) = broadcast(10);
        let (_cancel_tx, cancel_rx) = broadcast(1);

        tx.broadcast(1).await.unwrap();
        tx.broadcast(2).await.unwrap();
        tx.broadcast(3).await.unwrap();
        tx.broadcast(4).await.unwrap();
        tx.broadcast(5).await.unwrap();

        let and1 = eq_dep(
            rx.clone(),
            cancel_rx.clone(),
            format!("or_and_dep and1 {}", 4),
            4,
        )
        .and(eq_dep(
            rx.clone(),
            cancel_rx.clone(),
            format!("or_and_dep and1 {}", 6),
            6,
        ));
        let and2 = eq_dep(
            rx.clone(),
            cancel_rx.clone(),
            format!("or_and_dep and2 {}", 4),
            4,
        )
        .and(eq_dep(
            rx.clone(),
            cancel_rx.clone(),
            format!("or_and_dep and2 {}", 5),
            5,
        ));
        let or = and1.or(and2);
        let result = or.completed().await;
        assert_eq!(result, Some(vec![4, 5]));
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn many_and_dep() {
        let (tx, rx) = broadcast(10);
        let (_cancel_tx, cancel_rx) = broadcast(1);

        tx.broadcast(1).await.unwrap();
        tx.broadcast(2).await.unwrap();
        tx.broadcast(3).await.unwrap();
        tx.broadcast(4).await.unwrap();
        tx.broadcast(5).await.unwrap();
        tx.broadcast(6).await.unwrap();

        let mut and1 = eq_dep(
            rx.clone(),
            cancel_rx.clone(),
            format!("many_and_dep and1 {}", 4),
            4,
        )
        .and(eq_dep(
            rx.clone(),
            cancel_rx.clone(),
            format!("many_and_dep and1 {}", 6),
            6,
        ));
        let and2 = eq_dep(
            rx.clone(),
            cancel_rx.clone(),
            format!("many_and_dep and2 {}", 4),
            4,
        )
        .and(eq_dep(
            rx.clone(),
            cancel_rx.clone(),
            format!("many_and_dep and2 {}", 5),
            5,
        ));
        and1.add_deps(and2);
        let result = and1.completed().await;
        assert_eq!(result, Some(vec![4, 6, 4, 5]));
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn cancel_event_dep() {
        let (tx, rx) = broadcast(10);
        let (cancel_tx, cancel_rx) = broadcast(1);

        for i in 0..=5 {
            tx.broadcast(i).await.unwrap();
        }
        cancel_tx.broadcast(()).await.unwrap();
        let dep = eq_dep(
            rx.clone(),
            cancel_rx.clone(),
            format!("cancel_event_dep {}", 6),
            6,
        );
        let result = dep.completed().await;
        assert_eq!(result, None);
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn drop_cancel_dep() {
        let (tx, rx) = broadcast(10);
        let (cancel_tx, cancel_rx) = broadcast(1);

        for i in 0..=5 {
            tx.broadcast(i).await.unwrap();
        }
        drop(cancel_tx);
        let dep = eq_dep(
            rx.clone(),
            cancel_rx.clone(),
            format!("drop_cancel_dep {}", 6),
            6,
        );
        let result = dep.completed().await;
        assert_eq!(result, None);
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn cancel_and_dep() {
        let (tx, rx) = broadcast(10);
        let (cancel_tx, cancel_rx) = broadcast(1);

        let mut deps = vec![];
        for i in 0..=5 {
            tx.broadcast(i).await.unwrap();
            deps.push(eq_dep(
                rx.clone(),
                cancel_rx.clone(),
                format!("cancel_and_dep {i}"),
                i,
            ))
        }
        deps.push(eq_dep(
            rx.clone(),
            cancel_rx.clone(),
            format!("cancel_and_dep {}", 6),
            6,
        ));
        cancel_tx.broadcast(()).await.unwrap();
        let result = AndDependency::from_deps(deps).completed().await;
        assert_eq!(result, None);
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn cancel_or_dep() {
        let (_, rx) = broadcast(10);
        let (cancel_tx, cancel_rx) = broadcast(1);

        let mut deps = vec![];
        for i in 0..=5 {
            deps.push(eq_dep(
                rx.clone(),
                cancel_rx.clone(),
                format!("cancel_event_dep {i}"),
                i,
            ))
        }
        cancel_tx.broadcast(()).await.unwrap();
        let result = OrDependency::from_deps(deps).completed().await;
        assert_eq!(result, None);
    }
}