third_party/rust/uniffi_core/src/ffi/foreignfuture.rs

   1 /* This Source Code Form is subject to the terms of the Mozilla Public
   2  * License, v. 2.0. If a copy of the MPL was not distributed with this
   3  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
   4
   5 //! This module defines a Rust Future that wraps an async foreign function call.
   6 //!
   7 //! The general idea is to create a [oneshot::Channel], hand the sender to the foreign side, and
   8 //! await the receiver side on the Rust side.
   9 //!
  10 //! The foreign side should:
  11 //!   * Input a [ForeignFutureCallback] and a `u64` handle in their scaffolding function.
  12 //!     This is the sender, converted to a raw pointer, and an extern "C" function that sends the result.
  13 //!   * Return a [ForeignFuture], which represents the foreign task object corresponding to the async function.
  14 //!   * Call the [ForeignFutureCallback] when the async function completes with:
  15 //!     * The `u64` handle initially passed in
  16 //!     * The `ForeignFutureResult` for the call
  17 //!   * Wait for the [ForeignFutureHandle::free] function to be called to free the task object.
  18 //!     If this is called before the task completes, then the task will be cancelled.
  19
  20 use crate::{LiftReturn, RustCallStatus, UnexpectedUniFFICallbackError};
  21
  22 /// Handle for a foreign future
  23 pub type ForeignFutureHandle = u64;
  24
  25 /// Handle for a callback data associated with a foreign future.
  26 pub type ForeignFutureCallbackData = *mut ();
  27
  28 /// Callback that's passed to a foreign async functions.
  29 ///
  30 /// See `LiftReturn` trait for how this is implemented.
  31 pub type ForeignFutureCallback<FfiType> =
  32     extern "C" fn(oneshot_handle: u64, ForeignFutureResult<FfiType>);
  33
  34 /// C struct that represents the result of a foreign future
  35 #[repr(C)]
  36 pub struct ForeignFutureResult<T> {
  37     // Note: for void returns, T is `()`, which isn't directly representable with C since it's a ZST.
  38     // Foreign code should treat that case as if there was no `return_value` field.
  39     return_value: T,
  40     call_status: RustCallStatus,
  41 }
  42
  43 /// Perform a call to a foreign async method
  44
  45 /// C struct that represents the foreign future.
  46 ///
  47 /// This is what's returned by the async scaffolding functions.
  48 #[repr(C)]
  49 pub struct ForeignFuture {
  50     pub handle: ForeignFutureHandle,
  51     pub free: extern "C" fn(handle: ForeignFutureHandle),
  52 }
  53
  54 impl Drop for ForeignFuture {
  55     fn drop(&mut self) {
  56         (self.free)(self.handle)
  57     }
  58 }
  59
  60 unsafe impl Send for ForeignFuture {}
  61
  62 pub async fn foreign_async_call<F, T, UT>(call_scaffolding_function: F) -> T
  63 where
  64     F: FnOnce(ForeignFutureCallback<T::ReturnType>, u64) -> ForeignFuture,
  65     T: LiftReturn<UT>,
  66 {
  67     let (sender, receiver) = oneshot::channel::<ForeignFutureResult<T::ReturnType>>();
  68     // Keep the ForeignFuture around, even though we don't ever use it.
  69     // The important thing is that the ForeignFuture will be dropped when this Future is.
  70     let _foreign_future =
  71         call_scaffolding_function(foreign_future_complete::<T, UT>, sender.into_raw() as u64);
  72     match receiver.await {
  73         Ok(result) => T::lift_foreign_return(result.return_value, result.call_status),
  74         Err(e) => {
  75             // This shouldn't happen in practice, but we can do our best to recover
  76             T::handle_callback_unexpected_error(UnexpectedUniFFICallbackError::new(format!(
  77                 "Error awaiting foreign future: {e}"
  78             )))
  79         }
  80     }
  81 }
  82
  83 pub extern "C" fn foreign_future_complete<T: LiftReturn<UT>, UT>(
  84     oneshot_handle: u64,
  85     result: ForeignFutureResult<T::ReturnType>,
  86 ) {
  87     let channel = unsafe { oneshot::Sender::from_raw(oneshot_handle as *mut ()) };
  88     // Ignore errors in send.
  89     //
  90     // Error means the receiver was already dropped which will happen when the future is cancelled.
  91     let _ = channel.send(result);
  92 }
  93
  94 #[cfg(test)]
  95 mod test {
  96     use super::*;
  97     use crate::{Lower, RustBuffer};
  98     use once_cell::sync::OnceCell;
  99     use std::{
 100         future::Future,
 101         pin::Pin,
 102         sync::{
 103             atomic::{AtomicU32, Ordering},
 104             Arc,
 105         },
 106         task::{Context, Poll, Wake},
 107     };
 108
 109     struct MockForeignFuture {
 110         freed: Arc<AtomicU32>,
 111         callback_info: Arc<OnceCell<(ForeignFutureCallback<RustBuffer>, u64)>>,
 112         rust_future: Option<Pin<Box<dyn Future<Output = String>>>>,
 113     }
 114
 115     impl MockForeignFuture {
 116         fn new() -> Self {
 117             let callback_info = Arc::new(OnceCell::new());
 118             let freed = Arc::new(AtomicU32::new(0));
 119
 120             let rust_future: Pin<Box<dyn Future<Output = String>>> = {
 121                 let callback_info = callback_info.clone();
 122                 let freed = freed.clone();
 123                 Box::pin(foreign_async_call::<_, String, crate::UniFfiTag>(
 124                     move |callback, data| {
 125                         callback_info.set((callback, data)).unwrap();
 126                         ForeignFuture {
 127                             handle: Arc::into_raw(freed) as *mut () as u64,
 128                             free: Self::free,
 129                         }
 130                     },
 131                 ))
 132             };
 133             let rust_future = Some(rust_future);
 134             let mut mock_foreign_future = Self {
 135                 freed,
 136                 callback_info,
 137                 rust_future,
 138             };
 139             // Poll the future once, to start it up.   This ensures that `callback_info` is set.
 140             let _ = mock_foreign_future.poll();
 141             mock_foreign_future
 142         }
 143
 144         fn poll(&mut self) -> Poll<String> {
 145             let waker = Arc::new(NoopWaker).into();
 146             let mut context = Context::from_waker(&waker);
 147             self.rust_future
 148                 .as_mut()
 149                 .unwrap()
 150                 .as_mut()
 151                 .poll(&mut context)
 152         }
 153
 154         fn complete_success(&self, value: String) {
 155             let (callback, data) = self.callback_info.get().unwrap();
 156             callback(
 157                 *data,
 158                 ForeignFutureResult {
 159                     return_value: <String as Lower<crate::UniFfiTag>>::lower(value),
 160                     call_status: RustCallStatus::new(),
 161                 },
 162             );
 163         }
 164
 165         fn complete_error(&self, error_message: String) {
 166             let (callback, data) = self.callback_info.get().unwrap();
 167             callback(
 168                 *data,
 169                 ForeignFutureResult {
 170                     return_value: RustBuffer::default(),
 171                     call_status: RustCallStatus::error(error_message),
 172                 },
 173             );
 174         }
 175
 176         fn drop_future(&mut self) {
 177             self.rust_future = None
 178         }
 179
 180         fn free_count(&self) -> u32 {
 181             self.freed.load(Ordering::Relaxed)
 182         }
 183
 184         extern "C" fn free(handle: u64) {
 185             let flag = unsafe { Arc::from_raw(handle as *mut AtomicU32) };
 186             flag.fetch_add(1, Ordering::Relaxed);
 187         }
 188     }
 189
 190     struct NoopWaker;
 191
 192     impl Wake for NoopWaker {
 193         fn wake(self: Arc<Self>) {}
 194     }
 195
 196     #[test]
 197     fn test_foreign_future() {
 198         let mut mock_foreign_future = MockForeignFuture::new();
 199         assert_eq!(mock_foreign_future.poll(), Poll::Pending);
 200         mock_foreign_future.complete_success("It worked!".to_owned());
 201         assert_eq!(
 202             mock_foreign_future.poll(),
 203             Poll::Ready("It worked!".to_owned())
 204         );
 205         // Since the future is complete, it should free the foreign future
 206         assert_eq!(mock_foreign_future.free_count(), 1);
 207     }
 208
 209     #[test]
 210     #[should_panic]
 211     fn test_foreign_future_error() {
 212         let mut mock_foreign_future = MockForeignFuture::new();
 213         assert_eq!(mock_foreign_future.poll(), Poll::Pending);
 214         mock_foreign_future.complete_error("It Failed!".to_owned());
 215         let _ = mock_foreign_future.poll();
 216     }
 217
 218     #[test]
 219     fn test_drop_after_complete() {
 220         let mut mock_foreign_future = MockForeignFuture::new();
 221         mock_foreign_future.complete_success("It worked!".to_owned());
 222         assert_eq!(mock_foreign_future.free_count(), 0);
 223         assert_eq!(
 224             mock_foreign_future.poll(),
 225             Poll::Ready("It worked!".to_owned())
 226         );
 227         // Dropping the future after it's complete should not panic, and not cause a double-free
 228         mock_foreign_future.drop_future();
 229         assert_eq!(mock_foreign_future.free_count(), 1);
 230     }
 231
 232     #[test]
 233     fn test_drop_before_complete() {
 234         let mut mock_foreign_future = MockForeignFuture::new();
 235         mock_foreign_future.complete_success("It worked!".to_owned());
 236         // Dropping the future before it's complete should cancel the future
 237         assert_eq!(mock_foreign_future.free_count(), 0);
 238         mock_foreign_future.drop_future();
 239         assert_eq!(mock_foreign_future.free_count(), 1);
 240     }
 241 }