firefox-main/third_party/rust/uniffi_core/src/ffi/foreignfuture.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/. */
//! This module defines a Rust Future that wraps an async foreign function call.
//!
//! The general idea is to create a oneshot channel, hand the sender to the foreign side, and
//! await the receiver side on the Rust side.
//!
//! The foreign side should:
//! * Input a [ForeignFutureCallback] and a `u64` handle in their scaffolding function.
//! This is the sender, converted to a raw pointer, and an extern "C" function that sends the result.
//! * Call the [ForeignFutureCallback] when the async function completes with:
//! * The `u64` handle initially passed in
//! * The `ForeignFutureResult` struct for the call
//! * Optionally, set the `ForeignFutureDroppedCallback` value if you want to be notified when the
//! Future is dropped in Rust. For languages that support it, this can be hooked up to
//! cancelling the async task for the method.
use crate::{oneshot, LiftReturn, RustCallStatus};
/// Callback that's passed to a foreign async functions.
///
/// See `LiftReturn` trait for how this is implemented.
pub type ForeignFutureCallback<FfiType> =
extern "C" fn(oneshot_handle: u64, ForeignFutureResult<FfiType>);
/// C struct that represents the result of a foreign future
#[repr(C)]
pub struct ForeignFutureResult<T> {
// Note: for void returns, T is `()`, which isn't directly representable with C since it's a ZST.
// Foreign code should treat that case as if there was no `return_value` field.
return_value: T,
call_status: RustCallStatus,
}
/// C callback function that's called when the Rust side of a foreign future is dropped
pub type ForeignFutureDroppedCallback = extern "C" fn(data: u64);
/// C struct that represents a foreign future dropped callback.
#[repr(C)]
pub struct ForeignFutureDroppedCallbackStruct {
pub callback_data: u64,
pub callback: ForeignFutureDroppedCallback,
}
impl Default for ForeignFutureDroppedCallbackStruct {
/// The default value implements a no-op callback.
/// This will be used for languages that don't set their own callbacks.
fn default() -> Self {
extern "C" fn callback(_data: u64) {}
Self {
callback_data: 0,
callback,
}
}
}
impl Drop for ForeignFutureDroppedCallbackStruct {
fn drop(&mut self) {
(self.callback)(self.callback_data)
}
}
unsafe impl Send for ForeignFutureDroppedCallbackStruct {}
pub async fn foreign_async_call<F, T, UT>(call_scaffolding_function: F) -> T
where
F: FnOnce(ForeignFutureCallback<T::ReturnType>, u64, &mut ForeignFutureDroppedCallbackStruct),
T: LiftReturn<UT>,
{
// Create a oneshot channel that will receive the result of the callback method
let (sender, receiver) = oneshot::channel::<ForeignFutureResult<T::ReturnType>>();
// Create complete callback/data from the oneshot channel
let complete_callback = foreign_future_complete::<T, UT>;
let complete_callback_data = sender.into_raw() as u64;
// Create a `ForeignFutureDroppedCallback`.
// Since this is owned by the `Future` that we're generating, when the future is dropped it
// will be dropped to.
// future is.
let mut foreign_future_dropped_callback = ForeignFutureDroppedCallbackStruct::default();
// Call the async method
call_scaffolding_function(
complete_callback,
complete_callback_data,
&mut foreign_future_dropped_callback,
);
// Await the result and use it to return a value
let result = receiver.await;
T::lift_foreign_return(result.return_value, result.call_status)
}
pub extern "C" fn foreign_future_complete<T: LiftReturn<UT>, UT>(
oneshot_handle: u64,
result: ForeignFutureResult<T::ReturnType>,
) {
let channel = unsafe { oneshot::Sender::from_raw(oneshot_handle as *mut ()) };
channel.send(result);
}
#[cfg(test)]
mod test {
use super::*;
use crate::{Lower, RustBuffer};
use once_cell::sync::OnceCell;
use std::{
future::Future,
pin::Pin,
sync::{
atomic::{AtomicU32, Ordering},
Arc,
},
task::{Context, Poll, Wake},
};
struct MockForeignFuture {
freed: Arc<AtomicU32>,
callback_info: Arc<OnceCell<(ForeignFutureCallback<RustBuffer>, u64)>>,
rust_future: Option<Pin<Box<dyn Future<Output = String>>>>,
}
impl MockForeignFuture {
fn new() -> Self {
let callback_info = Arc::new(OnceCell::new());
let future_dropped_call_count = Arc::new(AtomicU32::new(0));
let rust_future: Pin<Box<dyn Future<Output = String>>> = {
let callback_info = callback_info.clone();
let future_dropped_call_count = future_dropped_call_count.clone();
Box::pin(foreign_async_call::<_, String, crate::UniFfiTag>(
move |callback, data, out_dropped_callback| {
callback_info.set((callback, data)).unwrap();
*out_dropped_callback = ForeignFutureDroppedCallbackStruct {
callback_data: Arc::into_raw(future_dropped_call_count) as *mut ()
as u64,
callback: Self::future_drapped_callback,
};
},
))
};
let rust_future = Some(rust_future);
let mut mock_foreign_future = Self {
freed: future_dropped_call_count,
callback_info,
rust_future,
};
// Poll the future once, to start it up. This ensures that `callback_info` is set.
let _ = mock_foreign_future.poll();
mock_foreign_future
}
fn poll(&mut self) -> Poll<String> {
let waker = Arc::new(NoopWaker).into();
let mut context = Context::from_waker(&waker);
self.rust_future
.as_mut()
.unwrap()
.as_mut()
.poll(&mut context)
}
fn complete_success(&self, value: String) {
let (callback, data) = self.callback_info.get().unwrap();
callback(
*data,
ForeignFutureResult {
return_value: <String as Lower<crate::UniFfiTag>>::lower(value),
call_status: RustCallStatus::default(),
},
);
}
fn complete_error(&self, error_message: String) {
let (callback, data) = self.callback_info.get().unwrap();
callback(
*data,
ForeignFutureResult {
return_value: RustBuffer::default(),
call_status: RustCallStatus::error(error_message),
},
);
}
fn drop_future(&mut self) {
self.rust_future = None
}
fn free_count(&self) -> u32 {
self.freed.load(Ordering::Relaxed)
}
extern "C" fn future_drapped_callback(handle: u64) {
let flag = unsafe { Arc::from_raw(handle as *mut AtomicU32) };
flag.fetch_add(1, Ordering::Relaxed);
}
}
struct NoopWaker;
impl Wake for NoopWaker {
fn wake(self: Arc<Self>) {}
}
#[test]
fn test_foreign_future() {
let mut mock_foreign_future = MockForeignFuture::new();
assert_eq!(mock_foreign_future.poll(), Poll::Pending);
mock_foreign_future.complete_success("It worked!".to_owned());
assert_eq!(
mock_foreign_future.poll(),
Poll::Ready("It worked!".to_owned())
);
// Since the future is complete, it should free the foreign future
assert_eq!(mock_foreign_future.free_count(), 1);
}
#[test]
#[should_panic]
fn test_foreign_future_error() {
let mut mock_foreign_future = MockForeignFuture::new();
assert_eq!(mock_foreign_future.poll(), Poll::Pending);
mock_foreign_future.complete_error("It Failed!".to_owned());
let _ = mock_foreign_future.poll();
}
#[test]
fn test_drop_after_complete() {
let mut mock_foreign_future = MockForeignFuture::new();
mock_foreign_future.complete_success("It worked!".to_owned());
assert_eq!(mock_foreign_future.free_count(), 0);
assert_eq!(
mock_foreign_future.poll(),
Poll::Ready("It worked!".to_owned())
);
// Dropping the future after it's complete should not panic, and not cause a double-free
mock_foreign_future.drop_future();
assert_eq!(mock_foreign_future.free_count(), 1);
}
#[test]
fn test_drop_before_complete() {
let mut mock_foreign_future = MockForeignFuture::new();
mock_foreign_future.complete_success("It worked!".to_owned());
// Dropping the future before it's complete should cancel the future
assert_eq!(mock_foreign_future.free_count(), 0);
mock_foreign_future.drop_future();
assert_eq!(mock_foreign_future.free_count(), 1);
}
}