cs-utils 0.21.1

use std::pin::Pin;

use futures::{future, Future};

use super::wait;

/// Runs a future to completion or to timeout, whatever happens first.
/// 
/// Wraps the original result of a future into an `Option` to indicate if
/// the future completed (`Some(original result)`) or timed out (`None`).
///
/// ### Examples
/// 
/// ```
/// use std::time::Instant;
///
/// use cs_utils::{
///     random_number,
///     futures::{wait, with_timeout},
/// };
///
/// #[tokio::main]
/// async fn main() {
///     let timeout: u64 = 25;
///
///     // as future that never completes
///     let forever_future = async move {
///         loop {
///             wait(5).await;
///         }
///     };
///
///     // wrap the original future to get the `stop` function
///     let with_timeout_future = with_timeout(forever_future, timeout);
///
///     // record starting time
///     let start_time = Instant::now();
/// 
///     // wait for the future to complete
///     let result = with_timeout_future.await;
///         
///     // calculate elapsed time
///     let time_delta_ms = (Instant::now() - start_time).as_millis();
///
///     assert!(
///         result.is_none(),
///         "Timed out future must complete with `None` result.",
///     );
///
///     // assert that the completion time of the future is close to the `timeout`
///     assert!(
///         time_delta_ms >= (timeout - 2) as u128,
///         "Must have waited for at least duration of the timeout.",
///     );
///     assert!(
///         time_delta_ms <= (timeout + 2) as u128,
///         "Must have waited for at most duration of the timeout.",
///     );
/// }
/// ```
pub fn with_timeout<
    T: Send + 'static,
    TFuture: Future<Output = T> + Send + 'static,
>(
    original_future: TFuture,
    timeout_ms: u64,
) -> Pin<Box<dyn Future<Output = Option<T>> + Send + 'static>> {
    let futures: Vec<Pin<Box<dyn Future<Output = Option<T>> + Send + 'static>>> = vec![
        Box::pin(async move {
            wait(timeout_ms).await;

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

    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 result_future;
}

#[cfg(test)]
mod tests {
    use std::time::Instant;

    use cs_utils::{
        random_number,
        futures::{wait, with_timeout},
    };
    
    #[tokio::test]
    async fn can_stop_a_future_after_a_timeout() {
        let timeout: u64 = 25;

        // as future that never completes
        let forever_future = async move {
            loop {
                wait(5).await;
            }
        };
        
        // wrap the original future to get the `stop` function
        let with_timeout_future = with_timeout(forever_future, timeout);

        // record starting time
        let start_time = Instant::now();
        
        // wait for the future to complete
        let result = with_timeout_future.await;
        
        // calculate elapsed time
        let time_delta_ms = (Instant::now() - start_time).as_millis();

        assert!(
            result.is_none(),
            "Timed out future must complete with `None` result.",
        );

        // assert that the completion time of the future is close to the `timeout`
        assert!(
            time_delta_ms >= (timeout - 2) as u128,
            "Must have waited for at least duration of the timeout (\"{delta}\" vs \"{timeout}\").",
            delta = time_delta_ms,
            timeout = timeout,
        );
        assert!(
            time_delta_ms <= (timeout + 2) as u128,
            "Must have waited for at most duration of the timeout (\"{delta}\" vs \"{timeout}\").",
            delta = time_delta_ms,
            timeout = timeout,
        );
    }

    #[tokio::test]
    async fn completed_future_returns_some_result() {
        let timeout: u64 = 25;
        let completion_delay: u64 = 5;
        let completing_future_result = random_number(0..=i128::MAX);

        // as future that never completes
        let completing_future = async move {
            wait(completion_delay).await;

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

        // record starting time
        let start_time = Instant::now();
        
        // wait for the future to complete
        let result = with_timeout_future.await
            .expect("Completed future must complete with `Some(original result)`.");

        // calculate elapsed time
        let time_delta_ms = (Instant::now() - start_time).as_millis();
    
        assert_eq!(
            result,
            completing_future_result,
            "Original and received future results must match.",
        );

        // assert that the completion time of the future is close to the `completion_delay`
        assert!(
            time_delta_ms >= (completion_delay - 2) as u128,
            "Must have waited for at least duration of the completion delay (\"{delta}\" vs \"{delay}\").",
            delta = time_delta_ms,
            delay = completion_delay,
        );
        assert!(
            time_delta_ms <= (completion_delay + 2) as u128,
            "Must have waited for at most duration of the completion delay (\"{delta}\" vs \"{delay}\").",
            delta = time_delta_ms,
            delay = completion_delay,
        );
    }
}