Simplify the new async system

- Updated the continuation handling code to use a Mutex rather than
  hand-coded locking code.  The performance will be slightly worse, but
  the code simplicity is worth it, and I think I can improve performance
  when implementing mozilla#1730.
- Register the continuation callback at startup rather than passing
  around each call.  Use the callback cell code that I added for this
  purpose, but forgot to use.

uniffi_bindgen/src/bindings/kotlin/gen_kotlin/mod.rs

@@ -196,6 +196,7 @@ pub struct KotlinWrapper<'a> {
     ci: &'a ComponentInterface,
     type_helper_code: String,
     type_imports: BTreeSet<ImportRequirement>,
+    has_async_fns: bool,
 }
 
 impl<'a> KotlinWrapper<'a> {
@@ -208,6 +209,7 @@ impl<'a> KotlinWrapper<'a> {
             ci,
             type_helper_code,
             type_imports,
+            has_async_fns: ci.has_async_fns(),
         }
     }
 
@@ -216,6 +218,10 @@ impl<'a> KotlinWrapper<'a> {
             .iter_types()
             .map(|t| KotlinCodeOracle.find(t))
             .filter_map(|ct| ct.initialization_fn())
+            .chain(
+                self.has_async_fns
+                    .then(|| "uniffiRustFutureContinuationCallback.register".into()),
+            )
             .collect()
     }
 
@@ -301,7 +307,9 @@ impl KotlinCodeOracle {
             FfiType::ForeignExecutorHandle => "USize".to_string(),
             FfiType::ForeignExecutorCallback => "UniFfiForeignExecutorCallback".to_string(),
             FfiType::RustFutureHandle => "Pointer".to_string(),
-            FfiType::RustFutureContinuation => "UniFffiRustFutureContinutationType".to_string(),
+            FfiType::RustFutureContinuationCallback => {
+                "UniFffiRustFutureContinuationCallbackType".to_string()
+            }
             FfiType::RustFutureContinuationData => "USize".to_string(),
         }
     }

uniffi_bindgen/src/bindings/kotlin/templates/Async.kt

@@ -6,10 +6,14 @@ internal const val UNIFFI_RUST_FUTURE_POLL_MAYBE_READY = 1.toShort()
 internal val uniffiContinuationHandleMap = UniFfiHandleMap<CancellableContinuation<Short>>()
 
 // FFI type for Rust future continuations
-internal object uniffiRustFutureContinuation: UniFffiRustFutureContinutationType {
+internal object uniffiRustFutureContinuationCallback: UniFffiRustFutureContinuationCallbackType {
     override fun callback(continuationHandle: USize, pollResult: Short) {
         uniffiContinuationHandleMap.remove(continuationHandle)?.resume(pollResult)
     }
+
+    internal fun register(lib: _UniFFILib) {
+        lib.{{ ci.ffi_rust_future_continuation_callback_set().name() }}(this)
+    }
 }
 
 internal suspend fun<T, F, E: Exception> uniffiRustCallAsync(
@@ -23,7 +27,6 @@ internal suspend fun<T, F, E: Exception> uniffiRustCallAsync(
             val pollResult = suspendCancellableCoroutine<Short> { continuation ->
                 _UniFFILib.INSTANCE.{{ ci.ffi_rust_future_poll().name() }}(
                     rustFuture,
-                    uniffiRustFutureContinuation,
                     uniffiContinuationHandleMap.insert(continuation)
                 )
             }
@@ -36,3 +39,4 @@ internal suspend fun<T, F, E: Exception> uniffiRustCallAsync(
         _UniFFILib.INSTANCE.{{ ci.ffi_rust_future_free().name() }}(rustFuture)
     }
 }
+

uniffi_bindgen/src/bindings/kotlin/templates/Helpers.kt

@@ -156,6 +156,6 @@ internal class UniFfiHandleMap<T: Any> {
 }
 
 // FFI type for Rust future continuations
-internal interface UniFffiRustFutureContinutationType : com.sun.jna.Callback {
+internal interface UniFffiRustFutureContinuationCallbackType : com.sun.jna.Callback {
     fun callback(continuationHandle: USize, pollResult: Short);
 }

uniffi_bindgen/src/bindings/python/gen_python/mod.rs

@@ -312,7 +312,7 @@ impl PythonCodeOracle {
             FfiType::ForeignExecutorHandle => "ctypes.c_size_t".to_string(),
             FfiType::ForeignExecutorCallback => "_UNIFFI_FOREIGN_EXECUTOR_CALLBACK_T".to_string(),
             FfiType::RustFutureHandle => "ctypes.c_void_p".to_string(),
-            FfiType::RustFutureContinuation => "_UNIFFI_FUTURE_CONTINUATION_T".to_string(),
+            FfiType::RustFutureContinuationCallback => "_UNIFFI_FUTURE_CONTINUATION_T".to_string(),
             FfiType::RustFutureContinuationData => "ctypes.c_size_t".to_string(),
         }
     }

uniffi_bindgen/src/bindings/python/templates/Async.py

@@ -2,13 +2,13 @@
 _UNIFFI_RUST_FUTURE_POLL_READY = 0
 _UNIFFI_RUST_FUTURE_POLL_MAYBE_READY = 1
 
-# Stores futures for _uniffi_continuation_func
+# Stores futures for _uniffi_continuation_callback
 _UniffiContinuationPointerManager = _UniffiPointerManager()
 
 # Continuation callback for async functions
 # lift the return value or error and resolve the future, causing the async function to resume.
 @_UNIFFI_FUTURE_CONTINUATION_T
-def _uniffi_continuation_func(future_ptr, poll_code):
+def _uniffi_continuation_callback(future_ptr, poll_code):
     (eventloop, future) = _UniffiContinuationPointerManager.release_pointer(future_ptr)
     eventloop.call_soon_threadsafe(_uniffi_set_future_result, future, poll_code)
 
@@ -25,7 +25,6 @@ async def _uniffi_rust_call_async(rust_future, ffi_complete, lift_func, error_ff
             future = eventloop.create_future()
             _UniffiLib.{{ ci.ffi_rust_future_poll().name() }}(
                 rust_future,
-                _uniffi_continuation_func,
                 _UniffiContinuationPointerManager.new_pointer((eventloop, future)),
             )
             poll_code = await future
@@ -37,3 +36,5 @@ async def _uniffi_rust_call_async(rust_future, ffi_complete, lift_func, error_ff
         )
     finally:
         _UniffiLib.{{ ci.ffi_rust_future_free().name() }}(rust_future)
+
+_UniffiLib.{{ ci.ffi_rust_future_continuation_callback_set().name() }}(_uniffi_continuation_callback)

uniffi_bindgen/src/bindings/ruby/gen_ruby/mod.rs

@@ -163,7 +163,7 @@ mod filters {
                 unimplemented!("Foreign executors are not implemented")
             }
             FfiType::RustFutureHandle
-            | FfiType::RustFutureContinuation
+            | FfiType::RustFutureContinuationCallback
             | FfiType::RustFutureContinuationData => {
                 unimplemented!("Async functions are not implemented")
             }

uniffi_bindgen/src/bindings/swift/gen_swift/mod.rs

@@ -342,6 +342,7 @@ pub struct SwiftWrapper<'a> {
     ci: &'a ComponentInterface,
     type_helper_code: String,
     type_imports: BTreeSet<String>,
+    has_async_fns: bool,
 }
 impl<'a> SwiftWrapper<'a> {
     pub fn new(config: Config, ci: &'a ComponentInterface) -> Self {
@@ -353,6 +354,7 @@ impl<'a> SwiftWrapper<'a> {
             ci,
             type_helper_code,
             type_imports,
+            has_async_fns: ci.has_async_fns(),
         }
     }
 
@@ -365,6 +367,10 @@ impl<'a> SwiftWrapper<'a> {
             .iter_types()
             .map(|t| SwiftCodeOracle.find(t))
             .filter_map(|ct| ct.initialization_fn())
+            .chain(
+                self.has_async_fns
+                    .then(|| "uniffiInitContinuationCallback".into()),
+            )
             .collect()
     }
 }
@@ -463,7 +469,7 @@ impl SwiftCodeOracle {
             FfiType::ForeignCallback => "ForeignCallback".into(),
             FfiType::ForeignExecutorHandle => "Int".into(),
             FfiType::ForeignExecutorCallback => "ForeignExecutorCallback".into(),
-            FfiType::RustFutureContinuation => "UniFfiRustFutureContinuation".into(),
+            FfiType::RustFutureContinuationCallback => "UniFfiRustFutureContinuation".into(),
             FfiType::RustFutureHandle | FfiType::RustFutureContinuationData => {
                 "UnsafeMutableRawPointer".into()
             }
@@ -475,7 +481,7 @@ impl SwiftCodeOracle {
             FfiType::ForeignCallback
             | FfiType::ForeignExecutorCallback
             | FfiType::RustFutureHandle
-            | FfiType::RustFutureContinuation
+            | FfiType::RustFutureContinuationCallback
             | FfiType::RustFutureContinuationData => {
                 format!("{} _Nonnull", self.ffi_type_label_raw(ffi_type))
             }
@@ -560,7 +566,9 @@ pub mod filters {
             FfiType::ForeignCallback => "ForeignCallback _Nonnull".into(),
             FfiType::ForeignExecutorCallback => "UniFfiForeignExecutorCallback _Nonnull".into(),
             FfiType::ForeignExecutorHandle => "size_t".into(),
-            FfiType::RustFutureContinuation => "UniFfiRustFutureContinuation _Nonnull".into(),
+            FfiType::RustFutureContinuationCallback => {
+                "UniFfiRustFutureContinuation _Nonnull".into()
+            }
             FfiType::RustFutureHandle | FfiType::RustFutureContinuationData => {
                 "void* _Nonnull".into()
             }

uniffi_bindgen/src/bindings/swift/templates/Async.swift

@@ -20,7 +20,6 @@ internal func uniffiRustCallAsync<F, T>(
             pollResult = await withUnsafeContinuation {
                 {{ ci.ffi_rust_future_poll().name() }}(
                     rustFuture,
-                    uniffiFutureContinuation,
                     ContinuationHolder($0).toOpaque()
                 )
             }
@@ -37,7 +36,7 @@ internal func uniffiRustCallAsync<F, T>(
 
 // Callback handlers for an async calls.  These are invoked by Rust when the future is ready.  They
 // lift the return value or error and resume the suspended function.
-fileprivate func uniffiFutureContinuation(ptr: UnsafeMutableRawPointer, pollResult: Int8) {
+fileprivate func uniffiFutureContinuationCallback(ptr: UnsafeMutableRawPointer, pollResult: Int8) {
     ContinuationHolder.fromOpaque(ptr).resume(pollResult)
 }
 
@@ -62,3 +61,7 @@ class ContinuationHolder {
         return Unmanaged<ContinuationHolder>.fromOpaque(ptr).takeRetainedValue()
     }
 }
+
+fileprivate func uniffiInitContinuationCallback() {
+    {{ ci.ffi_rust_future_continuation_callback_set().name() }}(uniffiFutureContinuationCallback)
+}

uniffi_bindgen/src/interface/ffi.rs

@@ -57,7 +57,7 @@ pub enum FfiType {
     /// Pointer to a Rust future
     RustFutureHandle,
     /// Continuation function for a Rust future
-    RustFutureContinuation,
+    RustFutureContinuationCallback,
     RustFutureContinuationData,
     // TODO: you can imagine a richer structural typesystem here, e.g. `Ref<String>` or something.
     // We don't need that yet and it's possible we never will, so it isn't here for now.

uniffi_bindgen/src/interface/mod.rs

@@ -413,6 +413,24 @@ impl ComponentInterface {
         }
     }
 
+    /// Builtin FFI function to set the Rust Future continuation callback
+    pub fn ffi_rust_future_continuation_callback_set(&self) -> FfiFunction {
+        FfiFunction {
+            name: format!(
+                "ffi_{}_rust_future_continuation_callback_set",
+                self.ffi_namespace()
+            ),
+            arguments: vec![FfiArgument {
+                name: "callback".into(),
+                type_: FfiType::RustFutureContinuationCallback,
+            }],
+            return_type: None,
+            is_async: false,
+            has_rust_call_status_arg: false,
+            is_object_free_function: false,
+        }
+    }
+
     /// Builtin FFI function to poll a Rust future.
     pub fn ffi_rust_future_poll(&self) -> FfiFunction {
         FfiFunction {
@@ -423,11 +441,6 @@ impl ComponentInterface {
                     name: "handle".to_string(),
                     type_: FfiType::RustFutureHandle,
                 },
-                // Continuation to call when the future can make progress
-                FfiArgument {
-                    name: "continuation".into(),
-                    type_: FfiType::RustFutureContinuation,
-                },
                 // Data to pass to the continuation
                 FfiArgument {
                     name: "uniffi_callback".into(),
@@ -597,6 +610,7 @@ impl ComponentInterface {
     /// List all FFI functions definitions for async functionality.
     pub fn iter_futures_ffi_function_definitons(&self) -> impl Iterator<Item = FfiFunction> {
         [
+            self.ffi_rust_future_continuation_callback_set(),
             self.ffi_rust_future_poll(),
             self.ffi_rust_future_cancel(),
             self.ffi_rust_future_free(),

uniffi_bindgen/src/scaffolding/mod.rs

@@ -84,7 +84,7 @@ mod filters {
             FfiType::ForeignBytes => "::uniffi::ForeignBytes".into(),
             FfiType::ForeignCallback => "::uniffi::ForeignCallback".into(),
             FfiType::RustFutureHandle => "::uniffi::RustFutureHandle".into(),
-            FfiType::RustFutureContinuation => "::uniffi::RustFutureContinuation".into(),
+            FfiType::RustFutureContinuationCallback => "::uniffi::RustFutureContinuation".into(),
             FfiType::RustFutureContinuationData => "*const ()".into(),
             FfiType::ForeignExecutorHandle => "::uniffi::ForeignExecutorHandle".into(),
             FfiType::ForeignExecutorCallback => "::uniffi::ForeignExecutorCallback".into(),

uniffi_core/src/ffi/foreigncallbacks.rs

@@ -10,7 +10,7 @@
 
 use std::sync::atomic::{AtomicUsize, Ordering};
 
-use crate::{ForeignExecutorHandle, RustBuffer, RustTaskCallback};
+use crate::{ForeignExecutorHandle, RustBuffer, RustFuturePoll, RustTaskCallback};
 
 /// ForeignCallback is the Rust representation of a foreign language function.
 /// It is the basis for all callbacks interfaces. It is registered exactly once per callback interface,
@@ -56,12 +56,21 @@ pub type ForeignExecutorCallback = extern "C" fn(
     task_data: *const (),
 ) -> i8;
 
+/// Foreign callback that's passed to [rust_future_poll]
+///
+/// The Rust side of things calls this when the foreign side should call [rust_future_poll] again
+/// to continue progress on the future.
+pub type RustFutureContinuationCallback = extern "C" fn(callback_data: *const (), RustFuturePoll);
+
 /// Store a [ForeignCallback] pointer
 pub(crate) struct ForeignCallbackCell(AtomicUsize);
 
 /// Store a [ForeignExecutorCallback] pointer
 pub(crate) struct ForeignExecutorCallbackCell(AtomicUsize);
 
+/// Store a [RustFutureContinuationCallback] pointer
+pub(crate) struct RustFutureContinuationCallbackCell(AtomicUsize);
+
 /// Macro to define foreign callback types as well as the callback cell.
 macro_rules! impl_foreign_callback_cell {
     ($callback_type:ident, $cell_type:ident) => {
@@ -101,3 +110,7 @@ macro_rules! impl_foreign_callback_cell {
 
 impl_foreign_callback_cell!(ForeignCallback, ForeignCallbackCell);
 impl_foreign_callback_cell!(ForeignExecutorCallback, ForeignExecutorCallbackCell);
+impl_foreign_callback_cell!(
+    RustFutureContinuationCallback,
+    RustFutureContinuationCallbackCell
+);