firefox-main/third_party/rust/uniffi_core/src/ffi/rustfuture/future.rs

Source code

Revision control

Copy as Markdown

Other Tools

/* 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`.
//!
use super::{
FutureLowerReturn, RustFutureContinuationCallback, 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(super) struct RustFuture<FfiType> {
// 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: Mutex<Scheduler>,
}
impl<FfiType> RustFuture<FfiType> {
pub(super) 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: Mutex::new(Scheduler::new()),
}
}
pub(super) fn poll(self: Arc<Self>, callback: RustFutureContinuationCallback, data: u64) {
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(data, RustFuturePoll::Ready)
} else {
self.scheduler.lock().unwrap().store(callback, data);
}
}
pub(super) fn is_cancelled(&self) -> bool {
self.scheduler.lock().unwrap().is_cancelled()
}
pub(super) fn cancel(&self) {
self.scheduler.lock().unwrap().cancel();
}
pub(super) fn complete(&self, call_status: &mut RustCallStatus) -> FfiType
where
FfiType: FfiDefault,
{
self.future.lock().unwrap().complete(call_status)
}
pub(super) fn free(self: Arc<Self>) {
// Call cancel() to send any leftover data to the continuation callback
self.scheduler.lock().unwrap().cancel();
// Ensure we drop our inner future, releasing all held references
self.future.lock().unwrap().free();
}
}
// `RawWaker` implementation for RustFuture
impl<FfiType> RustFuture<FfiType>
where
Scheduler: Send + Sync,
{
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.lock().unwrap().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.lock().unwrap().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) }
}
}