future_handles 0.2.0

use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Weak};
use std::task::{Context, Poll};

use spinning_top::Spinlock;

use crate::state::State;
use crate::{HandleError, HandleResult};

/// Creates a new [`CompletableFuture`] and the associated [`CompleteHandle`].
///
/// ```rust
/// # use future_handles::sync;
/// # async fn deadlocks_before_result() -> u32 {
/// let (future, handle) = sync::create();
///
/// func_with_callback(|res| {
///     handle.complete(res);
/// });
///
/// match future.await {
///     // The callback was invoked and the result set via the handle.
///     Ok(res) => res,
///     // The callback was never invoked, but the handle has been dropped.
///     Err(err) => panic!("Handle was dropped without setting a value")
/// }
/// # }
/// # fn func_with_callback<F>(func: F)
/// #    where F: FnOnce(u32) {
/// #    func(1);
/// # }
/// ```
///
/// # Danger!
///
/// Be careful to not await the future **before** setting the complete value or dropping the handle
/// while in the same async block, or you will cause a **deadlock**!
///
/// For a safer API, see [`scoped`].
///
/// ## Deadlock Examples
///
/// Setting the result:
/// ```rust
/// # use future_handles::sync;
/// async fn deadlocks_before_result() {
///     let (future, handle) = sync::create();
///
///     // Start awaiting here...
///     future.await.unwrap();
///
///     // The result is set here, but we'll never be able to reach it!
///     handle.complete(1);
/// }
/// ```
///
/// Dropping the [`CompleteHandle`]. Be careful as this is more subtle, and **MAY** cause a deadlock
/// depending on your compiler's implementation, as Rust is under no obligation to drop the
/// [`CompleteHandle`] before you await the future, as it's lifetime extends until the end of the
/// block:
/// ```rust
/// # use future_handles::sync;
/// async fn may_deadlock_before_drop() {
///     let (future, handle) = sync::create::<()>();
///     // The handle could be dropped immediately here and never deadlock.
///
///     // Start awaiting here...
///     future.await.unwrap();
///
///     // Or it could be dropped here, but we'll never reach it, deadlocking!
/// }
/// ```
///
/// Introducing an inverse dependency between the [`CompletableFuture`] and it's [`CompleteHandle`]
/// will always deadlock:
/// ```rust
/// # use future_handles::sync;
/// async fn may_deadlock_before_drop() {
///     let (future, handle) = sync::create::<bool>();
///
///     // Making the completion depend on the result of the computation itself will cause a deadlock.
///     if future.await.unwrap() {
///         handle.complete(false);
///     }
/// }
/// ```
///
/// [`CompletableFuture`]: CompletableFuture
/// [`CompleteHandle`]: CompleteHandle
/// [`scoped`]: super::scoped::scoped
pub fn create<T: Send>() -> (CompletableFuture<T>, CompleteHandle<T>) {
    let payload = Arc::new(Spinlock::new(State::new()));
    let weak = Arc::downgrade(&payload);
    (
        CompletableFuture { inner: payload },
        CompleteHandle { inner: weak },
    )
}

/// A thread-safe handle to complete the associated [`CompletableFuture`].
/// It can be safely dropped without setting a completion value.
///
/// ```rust
/// # use future_handles::sync::CompleteHandle;
/// fn func(complete_handle: CompleteHandle<u32>) {
///     if let Some(res) = func_that_may_fail() {
///         // Set the result.
///         complete_handle.complete(res);
///     }
///
///     // Or just drop the handle.
/// }
/// # fn func_that_may_fail() -> Option<u32> {
/// #    Some(1)
/// # }
/// ```
///
/// If cloned, the handles race to complete the future.
///
/// ```rust
/// # use futures;
/// # use future_handles::sync::CompleteHandle;
/// async fn func(complete_handle: CompleteHandle<u32>) {
///     let clone = complete_handle.clone();
///
///     let a = async { complete_handle.complete(1); };
///     let b = async { clone.complete(2); };
///
///     futures::join!(a, b); // The handles race to set the result.
/// }
/// ```
///
/// [`CompletableFuture`]: CompletableFuture
pub struct CompleteHandle<T> {
    inner: Weak<Spinlock<State<HandleResult<T>>>>,
}

impl<T> CompleteHandle<T> {
    fn try_set(&self, value: HandleResult<T>) -> bool {
        if let Some(arc) = self.inner.upgrade() {
            let (res, maybe_waker) = {
                let mut guard = arc.lock();
                guard.try_set_result(value)
            };

            if let Some(waker) = maybe_waker {
                waker.wake();
            }
            return res;
        }
        false
    }

    /// Complete the future consuming the handle.
    /// Returns `true` if this handle has set the complete value.
    pub fn complete(self, value: T) -> bool {
        self.try_set(Ok(value))
    }
}

impl<T> Drop for CompleteHandle<T> {
    fn drop(&mut self) {
        self.try_set(Err(HandleError::DroppedBeforeComplete));
    }
}

impl<T> Clone for CompleteHandle<T> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }

    fn clone_from(&mut self, source: &Self) {
        self.inner = source.inner.clone();
    }
}

/// A thread-safe future that can only be completed by the associated [`CompleteHandle`].
///
/// Since the [`CompleteHandle`] can be dropped without setting a completion value,
/// `CompletableFuture` always wraps the return value in a [`HandleResult`].
/// ```rust
/// # use futures::Future;
/// # use future_handles::HandleResult;
/// # async fn func(completable_future: impl Future<Output = HandleResult<()>>) { ///
/// match completable_future.await {
///     Ok(res) => res,
///     Err(err) => panic!("Handle was dropped without setting a value")
/// }
/// # }
/// ```
///
/// [`CompleteHandle`]: CompleteHandle
/// [`HandleResult`]: crate::HandleResult
pub struct CompletableFuture<T> {
    inner: Arc<Spinlock<State<HandleResult<T>>>>,
}

impl<T: Send> Future for CompletableFuture<T> {
    /// The associated [`CompleteHandle`] may be dropped before setting a result, as such, it must
    /// be handled by returning a [`HandleResult`].
    type Output = HandleResult<T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut guard = self.inner.lock();
        match guard.try_consume() {
            None => {
                guard.update_waker(cx.waker().clone());
                Poll::Pending
            }
            Some(res) => Poll::Ready(res),
        }
    }
}

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

    use futures::executor::ThreadPool;
    use futures::task::SpawnExt;

    use super::*;

    fn get_test_objects() -> (
        ThreadPool,
        impl Future<Output = HandleResult<()>>,
        CompleteHandle<()>,
    ) {
        let (fut, comp) = create::<()>();
        (ThreadPool::new().unwrap(), fut, comp)
    }

    #[test]
    fn return_value() {
        let (pool, fut, comp) = get_test_objects();

        pool.spawn(async move {
            fut.await.unwrap();
        })
        .unwrap();

        pool.spawn(async {
            comp.complete(());
        })
        .unwrap();
    }

    #[test]
    fn drop() {
        let (pool, fut, comp) = get_test_objects();

        pool.spawn(async move {
            match fut.await {
                Err(err) => {
                    match err {
                        HandleError::DroppedBeforeComplete => {} // Ok
                    }
                }
                Ok(_) => panic!("Completer dropped but future returned Ok"),
            }
        })
        .unwrap();

        pool.spawn(async {
            mem::drop(comp);
        })
        .unwrap();
    }

    #[test]
    fn complete_before_await() {
        let (pool, fut, comp) = get_test_objects();

        comp.complete(());
        pool.spawn(async move {
            fut.await.unwrap();
        })
        .unwrap();
    }
}