cs-utils 0.21.1

use std::pin::Pin;

use anyhow::{Result, anyhow};
use futures::{future, Future};
use tokio::sync::oneshot::{self, error::TryRecvError};

use super::wait;

/// Wrap a future into one that can be stopped at any given
/// moment by invoking a returned `stop` function.
/// 
/// Wraps the original result of a future into an `Option` to
/// indicate if the future completed (`Some(original result)`) or
/// was stopped explicitelly (`None`).
///
/// ### Examples
/// 
/// ```
/// use std::pin::Pin;
///
/// use futures::{future, Future};
/// use cs_utils::futures::{wait, with_stop};
/// 
/// #[tokio::main]
/// async fn main() {
///     // as future that never completes
///     let forever_future = async move {
///         loop {
///             wait(5).await;
///         }
///     };
///     
///     // wrap the original future to get the `stop` function
///     let (stop, forever_future) = with_stop(forever_future);
///     
///     // create futures for testing purposes
///     let futures: Vec<Pin<Box<dyn Future<Output = ()>>>> = vec![
///         Box::pin(async move {
///             // run the forever future
///             let result = forever_future.await;
/// 
///             assert!(
///                 result.is_none(),
///                 "Must complete with `None` result since it was stopped prematurely.",
///             );
///         }),
///         Box::pin(async move {
///             wait(25).await;
///
///             // stop the forever future
///             stop()
///                 .expect("Cannot stop the future.");
///             
///             // wait for some more time so the assertion in
///             // the boxed future above has a chance to run
///             wait(25).await;
///         }),
///     ];
///     
///     // race the futures to completion
///     future::select_all(futures).await;
/// }
/// ```
pub fn with_stop<
    T: Send + 'static,
    TFuture: Future<Output = T> + Send + 'static,
>(
    original_future: TFuture,
) -> (Box<dyn FnOnce() -> Result<()> + Send>, Pin<Box<dyn Future<Output = Option<T>> + Send + 'static>>) {
    let (test_end_sender, mut test_end_receiver) = oneshot::channel::<()>();
    let futures: Vec<Pin<Box<dyn Future<Output = Option<T>> + Send + 'static>>> = vec![
        Box::pin(async move {
            while let Err(TryRecvError::Empty) = test_end_receiver.try_recv() {
                // common thread time slice is `~100ms`, so
                // `5ms` delay should be granular enough here 
                wait(5).await;
            }

            return None;
        }),
        
        Box::pin(async move {
            return Some(
                original_future.await,
            );
        }),
    ];

    let stop: Box<dyn FnOnce() -> Result<()> + Send> = Box::new(
        move || {
            match test_end_sender.send(()) {
                Ok(_) => Ok(()),
                Err(_) => Err(anyhow!("Cannot stop the future, might be already stopped.")),
            }
        },
    );

    let result_future: Pin<Box<dyn Future<Output = Option<T>> + Send + 'static>> = Box::pin(async move {
        let (result, _, _) = future::select_all(futures).await;

        result
    });

    return (stop, result_future);
}

#[cfg(test)]
mod tests {
    use std::pin::Pin;
    use futures::{future, Future};

    use cs_utils::{
        random_number,
        futures::{wait, with_stop},
    };
    
    #[tokio::test]
    async fn can_stop_a_future() {
        // as future that never completes
        let forever_future = async move {
            loop {
                wait(5).await;
            }
        };
        
        // wrap the original future to get the `stop` function
        let (stop, forever_future) = with_stop(forever_future);
        
        // create futures for testing purposes
        let futures: Vec<Pin<Box<dyn Future<Output = ()>>>> = vec![
            Box::pin(async move {
                // run the forever future
                let result = forever_future.await;
    
                assert!(
                    result.is_none(),
                    "Stopped future must complete with `None` result.",
                );
            }),
            Box::pin(async move {
                wait(25).await;

                // stop the forever future
                stop()
                    .expect("Cannot stop the future.");
                
                // wait for some more time so the assertion in
                // the boxed future above has a chance to run
                wait(25).await;
            }),
        ];
        
        // race the futures to completion
        future::select_all(futures).await;
    }

    #[tokio::test]
    async fn completed_future_returns_some_result() {
        let completing_future_result = random_number(0..=i128::MAX);
    
        // as future that never completes
        let completing_future = async move {
            wait(5).await;

            return completing_future_result;
        };
        
        // wrap the original future to get the `stop` function
        let (_stop, completing_future) = with_stop(completing_future);

        let result = completing_future.await
            .expect("Must complete with `Some(original result)`.");

        assert_eq!(
            result,
            completing_future_result,
            "Original and received future results must match.",
        );
    }
}