Async cancel signals pipeline frees

CHANGELOG.md

@@ -1,6 +1,6 @@
 ## 0.5.2
 
-- fix a narrow race in async cancellation. If `cancel()` on a `CompletableFuture` returned by an async function landed in the few instructions between the pipeline's `isCancelled()` check and its own `rust_future_free`, or between a wake and the re-poll it triggers, the Rust future was freed twice or polled after being freed. Seen once in CI as a glibc `tcache_thread_shutdown()` abort. Every use of the handle now runs under the future's monitor and stops once it has been freed. No measurable change to async call overhead.
+- fix a narrow race in async cancellation. If `cancel()` on a `CompletableFuture` returned by an async function landed in the few instructions between the pipeline's `isCancelled()` check and its own `rust_future_free`, or between a wake and the re-poll it triggers, the Rust future was freed twice or polled after being freed. Seen once in CI as a glibc `tcache_thread_shutdown()` abort. `cancel()` now signals Rust with `rust_future_cancel` and the poll/complete pipeline is the only thing that frees. The Rust future is dropped once the in-flight poll's continuation fires, on whichever thread the executor runs that continuation on; with a pool that is an executor thread rather than the cancelling one. Cancelling with an executor that accepts a task and never runs it now leaks the Rust future, where before `cancel()` freed it inline; see the executor notes in the README.
 - fix generated Java failing to compile when a Rust field or parameter name matched a name the generator used in the same scope, such as an enum variant field named `value` ([#63](https://github.com/IronCoreLabs/uniffi-bindgen-java/issues/63)). Enum and error variant field names, and function, method, constructor and callback-interface parameter names can no longer collide with anything the Java generator emits. Names uniffi-rs itself reserves on the Rust side, such as a field named `buf` or a callback-interface parameter named `uniffi_handle`, are rejected by the Rust derive before bindings are generated and remain unavailable.
 
 ## 0.5.1

Cargo.toml

@@ -56,6 +56,7 @@ uniffi-example-geometry = { git = "https://github.com/mozilla/uniffi-rs.git", br
 uniffi-example-rondpoint = { git = "https://github.com/mozilla/uniffi-rs.git", branch = "release-v0.32.x" }
 uniffi-example-sprites = { git = "https://github.com/mozilla/uniffi-rs.git", branch = "release-v0.32.x" }
 uniffi-example-todolist = { git = "https://github.com/mozilla/uniffi-rs.git", branch = "release-v0.32.x" }
+uniffi-fixture-async-lifecycle = { path = "fixtures/async-lifecycle" }
 uniffi-fixture-benchmarks = { git = "https://github.com/mozilla/uniffi-rs.git", branch = "release-v0.32.x" }
 uniffi-fixture-coverall = { git = "https://github.com/mozilla/uniffi-rs.git", branch = "release-v0.32.x" }
 uniffi-fixture-enum-types = { git = "https://github.com/mozilla/uniffi-rs.git", branch = "release-v0.32.x" }

README.md

@@ -236,6 +236,7 @@ scope).
 
 - a Rust `&[u8]` argument is borrowed by Rust for the duration of the call rather than copied, and takes a **direct** `java.nio.ByteBuffer` (as in Kotlin) instead of `byte[]`. Build one with `ByteBuffer.allocateDirect(arr.length).put(arr).flip()`; a heap buffer throws `IllegalArgumentException`. Reuse the buffer across calls where you can, and don't let another thread write to it while a call is in flight. `Vec<u8>` is unaffected and still maps to `byte[]`.
 - failures in CompletableFutures will cause them to `completeExceptionally`. The error that caused the failure can be checked with `e.getCause()`. When implementing an async Rust trait in Java, you'll need to `completeExceptionally` instead of throwing. See `TestFixtureFutures.java` for an example trait implementation with errors.
+- every async function and method has an overload taking a `java.util.concurrent.Executor`; the other overload uses `ForkJoinPool.commonPool()`. The first poll runs on the calling thread, before the method returns; the executor runs every re-poll and the completion of the returned `CompletableFuture`. Any executor that hands tasks to its own threads works: the common pool, `Executors.newFixedThreadPool`, `newCachedThreadPool`, `newSingleThreadExecutor`, or a virtual-thread executor. Rust invokes the continuation from inside `Waker::wake()`, so an executor that runs tasks on the submitting thread (such as `Runnable::run`) polls the future from within its own waker and deadlocks any future that holds a lock while waking. Cancelling the returned future signals Rust with `rust_future_cancel` and the poll/complete pipeline frees the Rust future once the in-flight poll's continuation fires, on whichever thread the executor runs that continuation on. Two consequences: an executor that accepts a task and never runs it (`ThreadPoolExecutor.DiscardPolicy`, or a queue drained by `shutdownNow()`) leaks the Rust future and everything it owns, and with an inline executor the free happens on the cancelling thread.
 - all primitives are signed in Java by default. Rust correctly interprets the a signed primitive value from Java as unsigned when told to. Callers of Uniffi functions need to be aware when making comparisons (`compareUnsigned`) or printing when a value is actually unsigned to code around footguns on this side.
 - this is an internal note for development but because Enum variants are not cases/hanging off their parent in Java, their named standalone, they can conflict with any/all `java.lang` types. We could do extensive checking and forced renaming around this, but instead we use fully qualified names for all `java.lang` types in all templates. Ensure that when you're making changes you're not dropping those qualified names or adding generated code without them.
 

fixtures/async-lifecycle/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "uniffi-fixture-async-lifecycle"
+version = "0.1.0"
+edition = "2024"
+
+[lib]
+crate-type = ["cdylib", "lib"]
+name = "uniffi_fixture_async_lifecycle"
+
+[dependencies]
+async-trait = "0.1"
+uniffi = { git = "https://github.com/mozilla/uniffi-rs.git", branch = "release-v0.32.x" }

fixtures/async-lifecycle/src/lib.rs

@@ -0,0 +1,141 @@
+/* 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/. */
+
+//! Async functions whose futures count their own construction and drop.
+
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::{
+        Arc, Mutex,
+        atomic::{AtomicU64, Ordering},
+    },
+    task::{Context, Poll, Waker},
+    thread,
+    time::Duration,
+};
+
+uniffi::setup_scaffolding!("async_lifecycle");
+
+static CREATED: AtomicU64 = AtomicU64::new(0);
+static DROPPED: AtomicU64 = AtomicU64::new(0);
+
+struct DropTracker;
+
+impl DropTracker {
+    fn new() -> Self {
+        CREATED.fetch_add(1, Ordering::SeqCst);
+        Self
+    }
+}
+
+impl Drop for DropTracker {
+    fn drop(&mut self) {
+        DROPPED.fetch_add(1, Ordering::SeqCst);
+    }
+}
+
+#[uniffi::export]
+pub fn created_count() -> u64 {
+    CREATED.load(Ordering::SeqCst)
+}
+
+#[uniffi::export]
+pub fn dropped_count() -> u64 {
+    DROPPED.load(Ordering::SeqCst)
+}
+
+#[uniffi::export]
+pub fn reset_counts() {
+    CREATED.store(0, Ordering::SeqCst);
+    DROPPED.store(0, Ordering::SeqCst);
+}
+
+struct TimerState {
+    completed: bool,
+    waker: Option<Waker>,
+}
+
+/// Completes after `duration` on a spawned thread, having woken the future `spurious_wakes` times
+/// first. A wake with no waker stored yet is dropped.
+struct Timer {
+    state: Arc<Mutex<TimerState>>,
+}
+
+impl Timer {
+    fn new(duration: Duration, spurious_wakes: u16) -> Self {
+        let state = Arc::new(Mutex::new(TimerState {
+            completed: false,
+            waker: None,
+        }));
+        let thread_state = Arc::clone(&state);
+        // Wakers are invoked with the lock released: an inline foreign executor re-polls from
+        // inside `wake()`, on this thread.
+        thread::spawn(move || {
+            for _ in 0..spurious_wakes {
+                thread::sleep(Duration::from_millis(1));
+                let waker = thread_state.lock().unwrap().waker.take();
+                if let Some(waker) = waker {
+                    waker.wake();
+                }
+            }
+            thread::sleep(duration);
+            let waker = {
+                let mut state = thread_state.lock().unwrap();
+                state.completed = true;
+                state.waker.take()
+            };
+            if let Some(waker) = waker {
+                waker.wake();
+            }
+        });
+        Self { state }
+    }
+}
+
+impl Future for Timer {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+        let mut state = self.state.lock().unwrap();
+        if state.completed {
+            Poll::Ready(())
+        } else {
+            state.waker = Some(cx.waker().clone());
+            Poll::Pending
+        }
+    }
+}
+
+#[uniffi::export]
+pub async fn tracked_ready() {
+    let _tracker = DropTracker::new();
+}
+
+#[uniffi::export]
+pub async fn tracked_sleep(ms: u16) {
+    let _tracker = DropTracker::new();
+    Timer::new(Duration::from_millis(ms.into()), 0).await;
+}
+
+#[uniffi::export]
+pub async fn tracked_wake_storm(ms: u16, spurious_wakes: u16) {
+    let _tracker = DropTracker::new();
+    Timer::new(Duration::from_millis(ms.into()), spurious_wakes).await;
+}
+
+#[uniffi::export(with_foreign)]
+#[async_trait::async_trait]
+pub trait Notifier: Send + Sync {
+    async fn on_ready(&self);
+}
+
+/// The sleep guarantees the foreign caller holds the future before `on_ready` is upcalled from
+/// a re-poll.
+#[uniffi::export]
+pub async fn sleep_then_notify(ms: u16, notifier: Arc<dyn Notifier>) {
+    let _tracker = DropTracker::new();
+    Timer::new(Duration::from_millis(ms.into()), 0).await;
+    notifier.on_ready().await;
+}

src/gen_java/mod.rs

@@ -1909,6 +1909,16 @@ mod filters {
         Ok(format!("(_allocator, future, continuation) -> {call}"))
     }
 
+    #[askama::filter_fn]
+    pub fn async_cancel(
+        callable: impl Callable,
+        _v: &dyn askama::Values,
+        ci: &ComponentInterface,
+    ) -> Result<String, askama::Error> {
+        let ffi_func = callable.ffi_rust_future_cancel(ci);
+        Ok(format!("(future) -> UniffiLib.{ffi_func}(future)"))
+    }
+
     #[askama::filter_fn]
     pub fn async_free(
         callable: impl Callable,
@@ -1936,11 +1946,51 @@ mod filters {
         _v: &dyn askama::Values,
         spaces: &i32,
     ) -> Result<String, askama::Error> {
-        let middle = textwrap::indent(&textwrap::dedent(docstring.as_ref()), " * ");
-        let wrapped = format!("/**\n{middle}\n */");
+        Ok(javadoc(&textwrap::dedent(docstring.as_ref()), *spaces))
+    }
 
-        let spaces = usize::try_from(*spaces).unwrap_or_default();
-        Ok(textwrap::indent(&wrapped, &" ".repeat(spaces)))
+    /// Javadoc for one overload of an async callable.
+    #[askama::filter_fn]
+    pub fn async_docstring(
+        callable: impl Callable,
+        _v: &dyn askama::Values,
+        spaces: &i32,
+        with_executor: bool,
+    ) -> Result<String, askama::Error> {
+        let mut body = String::new();
+        if let Some(docstring) = callable.docstring() {
+            body.push_str(&textwrap::dedent(docstring));
+            body.push_str("\n\n");
+        }
+        body.push_str(if with_executor {
+            ASYNC_EXECUTOR_PARAM_DOC
+        } else {
+            ASYNC_COMMON_POOL_DOC
+        });
+        Ok(javadoc(&body, *spaces))
+    }
+
+    const ASYNC_COMMON_POOL_DOC: &str = "\
+Re-polls and completes on {@link java.util.concurrent.ForkJoinPool#commonPool()}. See the overload
+taking an {@link java.util.concurrent.Executor} for the contract an executor must meet.";
+
+    const ASYNC_EXECUTOR_PARAM_DOC: &str = "\
+@param _uniffiExecutor runs every re-poll and the completion of the returned future. The first
+    poll runs on the calling thread, before this method returns. Any executor that hands tasks to
+    its own threads works: {@link java.util.concurrent.ForkJoinPool#commonPool()}, a fixed, cached
+    or single-thread pool, or a virtual-thread executor. Rust invokes the continuation from inside
+    {@code Waker::wake()}, so an executor that runs tasks on the submitting thread, such as
+    {@code Runnable::run}, polls the future from within its own waker and deadlocks any future
+    that holds a lock while waking. Cancelling the returned future signals Rust; the pipeline then
+    frees the Rust future, so an executor that accepts a task and never runs it, such as one using
+    {@link java.util.concurrent.ThreadPoolExecutor.DiscardPolicy} or one whose queue was drained
+    by {@code shutdownNow()}, leaks the Rust future and everything it owns.";
+
+    fn javadoc(body: &str, spaces: i32) -> String {
+        let middle = textwrap::indent(body, " * ");
+        let wrapped = format!("/**\n{middle}\n */");
+        let spaces = usize::try_from(spaces).unwrap_or_default();
+        textwrap::indent(&wrapped, &" ".repeat(spaces))
     }
 
     /// Returns the type name suitable for use in field declarations, method parameters, and return types.