uniffi_core 0.32.1

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! [`RustFuture`] represents a [`Future`] that can be sent to the foreign code over FFI.
//!
//! This type is not instantiated directly, but via the procedural macros, such as `#[uniffi::export]`.
//!
//! # The big picture
//!
//! We implement async foreign functions using a simplified version of the Future API:
//!
//! 0. At startup, register a [RustFutureContinuationCallback] by calling
//!    rust_future_continuation_callback_set.
//! 1. Call the scaffolding function to get a [Handle]
//! 2. In a loop:
//!   - Call [rust_future_poll]
//!   - Suspend the function until the [rust_future_poll] continuation function is called
//!   - If the continuation was function was called with [RustFuturePoll::Ready], then break
//!     otherwise continue.
//! 3. If the async function is cancelled, then call [rust_future_cancel].  This causes the
//!    continuation function to be called with [RustFuturePoll::Ready] and the [RustFuture] to
//!    enter a cancelled state.
//! 4. Call [rust_future_complete] to get the result of the future.
//! 5. Call [rust_future_free] to free the future, ideally in a finally block.  This:
//!    - Releases any resources held by the future
//!    - Calls any continuation callbacks that have not been called yet
//!
//! Note: Technically, the foreign code calls the scaffolding versions of the `rust_future_*`
//! functions.  These are generated by the scaffolding macro, specially prefixed, and extern "C",
//! and manually monomorphized in the case of [rust_future_complete].  See
//! `uniffi_macros/src/setup_scaffolding.rs` for details.
//!
//! ## How does `Future` work exactly?
//!
//! A [`Future`] in Rust does nothing. When calling an async function, it just
//! returns a `Future` but nothing has happened yet. To start the computation,
//! the future must be polled. It returns [`Poll::Ready(r)`][`Poll::Ready`] if
//! the result is ready, [`Poll::Pending`] otherwise. `Poll::Pending` basically
//! means:
//!
//! > Please, try to poll me later, maybe the result will be ready!
//!
//! This model is very different than what other languages do, but it can actually
//! be translated quite easily, fortunately for us!
//!
//! But… wait a minute… who is responsible to poll the `Future` if a `Future` does
//! nothing? Well, it's _the executor_. The executor is responsible _to drive_ the
//! `Future`: that's where they are polled.
//!
//! But… wait another minute… how does the executor know when to poll a [`Future`]?
//! Does it poll them randomly in an endless loop? Well, no, actually it depends
//! on the executor! A well-designed `Future` and executor work as follows.
//! Normally, when [`Future::poll`] is called, a [`Context`] argument is
//! passed to it. It contains a [`Waker`]. The [`Waker`] is built on top of a
//! [`RawWaker`] which implements whatever is necessary. Usually, a waker will
//! signal the executor to poll a particular `Future`. A `Future` will clone
//! or pass-by-ref the waker to somewhere, as a callback, a completion, a
//! function, or anything, to the system that is responsible to notify when a
//! task is completed. So, to recap, the waker is _not_ responsible for waking the
//! `Future`, it _is_ responsible for _signaling_ the executor that a particular
//! `Future` should be polled again. That's why the documentation of
//! [`Poll::Pending`] specifies:
//!
//! > When a function returns `Pending`, the function must also ensure that the
//! > current task is scheduled to be awoken when progress can be made.
//!
//! “awakening” is done by using the `Waker`.
//!
//! [`Future`]: https://doc.rust-lang.org/std/future/trait.Future.html
//! [`Future::poll`]: https://doc.rust-lang.org/std/future/trait.Future.html#tymethod.poll
//! [`Pol::Ready`]: https://doc.rust-lang.org/std/task/enum.Poll.html#variant.Ready
//! [`Poll::Pending`]: https://doc.rust-lang.org/std/task/enum.Poll.html#variant.Pending
//! [`Context`]: https://doc.rust-lang.org/std/task/struct.Context.html
//! [`Waker`]: https://doc.rust-lang.org/std/task/struct.Waker.html
//! [`RawWaker`]: https://doc.rust-lang.org/std/task/struct.RawWaker.html

use super::{
    FutureLowerReturn, RustFutureCallback, RustFutureContinuationBoundCallback, RustFuturePoll,
    Scheduler, UniffiCompatibleFuture,
};
use crate::{try_rust_call, FfiDefault, LiftArgsError, RustCallResult, RustCallStatus};
use std::{
    future, panic,
    pin::{pin, Pin},
    sync::{Arc, Mutex},
    task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
};

/// Wraps the actual future we're polling
///
/// * Lowers the future result into an FFI type
/// * Splits `Future::poll` into 2 parts.  WrappedFuture::poll(), which returns true if the future
///   is ready and `WrappedFuture::complete` which extracts the result.  This matches how our FFI
///   functions work.
/// * Polls the future inside of `rust_call_with_out_status` so that panics can be caught.
struct WrappedFuture<FfiType> {
    // Create an option for the future and result.
    // This could be a single enum, but makes future pinning harder.
    // For example if you call `std::mem::take()` while a future was stored in the enum, it would break the pinning guarantee.
    //
    // The output is `Result<FfiType, RustCallStatus>`.
    // Both the `Ok` and `Err` side are easy to pass back across the FFI.
    future: Option<Pin<Box<dyn UniffiCompatibleFuture<RustCallResult<FfiType>>>>>,
    result: Option<Result<FfiType, RustCallStatus>>,
}

impl<FfiType> WrappedFuture<FfiType> {
    fn new<F, T, UT>(future: F) -> Self
    where
        F: UniffiCompatibleFuture<Result<T, LiftArgsError>> + 'static,
        T: FutureLowerReturn<UT, ReturnType = FfiType>,
    {
        let wrapped_future = async {
            let mut future = pin!(future);
            future::poll_fn(move |cx| {
                let call_result = try_rust_call(
                    // This closure uses a `&mut F` value, which means it's not UnwindSafe by
                    // default.  If the future panics, it may be in an invalid state.
                    //
                    // However, we can safely use `AssertUnwindSafe` since a panic will lead the
                    // `Err` case below and we will never poll the future again.
                    panic::AssertUnwindSafe(|| match future.as_mut().poll(cx) {
                        Poll::Pending => Ok(Poll::Pending),
                        Poll::Ready(Ok(v)) => T::lower_return(v).map(Poll::Ready),
                        Poll::Ready(Err(e)) => T::handle_failed_lift(e).map(Poll::Ready),
                    }),
                );
                match call_result {
                    Ok(Poll::Pending) => Poll::Pending,
                    Ok(Poll::Ready(v)) => Poll::Ready(Ok(v)),
                    Err(call_status) => Poll::Ready(Err(call_status)),
                }
            })
            .await
        };
        Self {
            future: Some(Box::pin(wrapped_future)),
            result: None,
        }
    }

    // Poll the future and check if it's ready or not
    fn poll(&mut self, context: &mut Context<'_>) -> bool {
        if self.result.is_some() {
            true
        } else if let Some(future) = &mut self.future {
            match future.as_mut().poll(context) {
                Poll::Pending => false,
                Poll::Ready(result) => {
                    self.future = None;
                    self.result = Some(result);
                    true
                }
            }
        } else {
            trace!("poll with neither future nor result set");
            true
        }
    }

    fn complete(&mut self, out_status: &mut RustCallStatus) -> FfiType
    where
        FfiType: FfiDefault,
    {
        let mut return_value = FfiType::ffi_default();
        match self.result.take() {
            Some(Ok(v)) => return_value = v,
            Some(Err(call_status)) => *out_status = call_status,
            None => *out_status = RustCallStatus::cancelled(),
        }
        self.free();
        return_value
    }

    fn free(&mut self) {
        self.future = None;
        self.result = None;
    }
}

/// Future that the foreign code is awaiting
pub struct RustFuture<FfiType, Callback = RustFutureContinuationBoundCallback> {
    // This Mutex should never block if our code is working correctly, since there should not be
    // multiple threads calling [Self::poll] and/or [Self::complete] at the same time.
    future: Mutex<WrappedFuture<FfiType>>,
    scheduler: Scheduler<Callback>,
}

impl<FfiType, Callback> RustFuture<FfiType, Callback>
where
    Callback: RustFutureCallback,
{
    pub fn new<F, T, UT>(future: F, _tag: UT) -> Self
    where
        F: UniffiCompatibleFuture<Result<T, LiftArgsError>> + 'static,
        T: FutureLowerReturn<UT, ReturnType = FfiType>,
    {
        Self {
            future: Mutex::new(WrappedFuture::new(future)),
            scheduler: Scheduler::new(),
        }
    }

    pub fn poll(self: Arc<Self>, callback: Callback) {
        let cancelled = self.is_cancelled();
        let ready = cancelled || {
            let mut locked = self.future.lock().unwrap();
            let waker = Arc::clone(&self).into_waker();
            locked.poll(&mut Context::from_waker(&waker))
        };
        if ready {
            trace!("RustFuture::poll is ready (cancelled: {cancelled})");
            callback.invoke(RustFuturePoll::Ready)
        } else {
            self.scheduler.store(callback);
        }
    }

    pub fn is_cancelled(&self) -> bool {
        self.scheduler.is_cancelled()
    }

    pub fn cancel(&self) {
        self.scheduler.cancel();
    }

    pub fn complete(&self, call_status: &mut RustCallStatus) -> FfiType
    where
        FfiType: FfiDefault,
    {
        self.future.lock().unwrap().complete(call_status)
    }

    pub fn free(&self) {
        // Call cancel() to send any leftover data to the continuation callback
        self.scheduler.cancel();
        // Ensure we drop our inner future, releasing all held references
        self.future.lock().unwrap().free();
    }
}

// `RawWaker` implementation for RustFuture
impl<FfiType, Callback> RustFuture<FfiType, Callback>
where
    Scheduler: Send + Sync,
    Callback: RustFutureCallback,
{
    unsafe fn waker_clone(ptr: *const ()) -> RawWaker {
        trace!("RustFuture::waker_clone called ({ptr:?})");
        Arc::<Self>::increment_strong_count(ptr.cast::<Self>());
        RawWaker::new(
            ptr,
            &RawWakerVTable::new(
                Self::waker_clone,
                Self::waker_wake,
                Self::waker_wake_by_ref,
                Self::waker_drop,
            ),
        )
    }

    unsafe fn waker_wake(ptr: *const ()) {
        trace!("RustFuture::waker_wake called ({ptr:?})");
        Self::recreate_arc(ptr).scheduler.wake();
    }

    unsafe fn waker_wake_by_ref(ptr: *const ()) {
        trace!("RustFuture::waker_wake_by_ref called ({ptr:?})");
        // For wake_by_ref, we can use the pointer directly, without consuming it to re-create the
        // arc.
        let ptr = ptr.cast::<Self>();
        (*ptr).scheduler.wake();
    }

    unsafe fn waker_drop(ptr: *const ()) {
        trace!("RustFuture::waker_drop called ({ptr:?})");
        drop(Self::recreate_arc(ptr));
    }

    /// Recreate an Arc<Self> from a raw pointer
    ///
    /// # Safety
    /// * ptr must have been created from `Arc::into_raw()`
    /// * ptr must only be used once to re-create the arc
    unsafe fn recreate_arc(ptr: *const ()) -> Arc<Self> {
        let ptr = ptr.cast::<Self>();
        Arc::<Self>::from_raw(ptr)
    }

    fn into_waker(self: Arc<Self>) -> Waker {
        trace!("RustFuture::creating waker ({:?})", Arc::as_ptr(&self));
        let raw_waker = RawWaker::new(
            Arc::into_raw(self).cast::<()>(),
            &RawWakerVTable::new(
                Self::waker_clone,
                Self::waker_wake,
                Self::waker_wake_by_ref,
                Self::waker_drop,
            ),
        );

        // Safety:
        //
        // * The vtable functions are thread-safe, since we only access `Self::scheduler` and that
        //   has a `Send` + `Sync` bound.
        // * We manage the raw arc pointers correctly
        unsafe { Waker::from_raw(raw_waker) }
    }
}