Fix leaking trait futures nondeterministically

Cargo.toml

@@ -31,7 +31,6 @@ path = "src/main.rs"
 
 [dependencies]
 anyhow = "1"
-rinja = { version = "0.3", default-features = false, features = ["config"] }
 camino = "1.1.6"
 cargo_metadata = "0.15"
 clap = { version = "4", default-features = false, features = [
@@ -44,6 +43,7 @@ heck = "0.5"
 once_cell = "1.19.0"
 paste = "1"
 regex = "1.10.4"
+rinja = { version = "0.3", default-features = false, features = ["config"] }
 serde = "1"
 textwrap = "0.16.1"
 toml = "0.5" # can't be on 8, `Value` is part of public interface

src/gen_java/callback_interface.rs

@@ -2,7 +2,7 @@
  * 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/. */
 
-use super::{potentially_add_external_package, CodeType, Config};
+use super::{CodeType, Config, potentially_add_external_package};
 use crate::ComponentInterface;
 
 #[derive(Debug)]

src/gen_java/compounds.rs

@@ -4,8 +4,8 @@
 
 use super::{AsCodeType, CodeType, Config};
 use uniffi_bindgen::{
-    backend::{Literal, Type},
     ComponentInterface,
+    backend::{Literal, Type},
 };
 
 #[derive(Debug)]

src/gen_java/custom.rs

@@ -2,7 +2,7 @@
  * 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/. */
 
-use super::{potentially_add_external_package, CodeType, Config};
+use super::{CodeType, Config, potentially_add_external_package};
 use crate::ComponentInterface;
 
 #[derive(Debug)]

src/gen_java/enum_.rs

@@ -2,7 +2,7 @@
  * 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/. */
 
-use super::{potentially_add_external_package, CodeType, Config};
+use super::{CodeType, Config, potentially_add_external_package};
 use crate::ComponentInterface;
 use uniffi_bindgen::backend::Literal;
 

src/gen_java/object.rs

@@ -2,8 +2,8 @@
  * 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/. */
 
-use super::{potentially_add_external_package, CodeType, Config};
-use uniffi_bindgen::{interface::ObjectImpl, ComponentInterface};
+use super::{CodeType, Config, potentially_add_external_package};
+use uniffi_bindgen::{ComponentInterface, interface::ObjectImpl};
 
 #[derive(Debug)]
 pub struct ObjectCodeType {

src/gen_java/record.rs

@@ -2,7 +2,7 @@
  * 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/. */
 
-use super::{potentially_add_external_package, CodeType, Config};
+use super::{CodeType, Config, potentially_add_external_package};
 use uniffi_bindgen::ComponentInterface;
 
 #[derive(Debug)]

src/gen_java/variant.rs

@@ -2,7 +2,7 @@
  * 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/. */
 
-use super::{potentially_add_external_package, AsCodeType, CodeType, Config, JavaCodeOracle};
+use super::{AsCodeType, CodeType, Config, JavaCodeOracle, potentially_add_external_package};
 use uniffi_bindgen::interface::{ComponentInterface, Variant};
 
 #[derive(Debug)]

src/templates/Async.java

@@ -13,7 +13,7 @@ public final class UniffiAsyncHelpers {
     static final byte UNIFFI_RUST_FUTURE_POLL_READY = (byte) 0;
     static final byte UNIFFI_RUST_FUTURE_POLL_MAYBE_READY = (byte) 1;
     static final UniffiHandleMap<CompletableFuture<Byte>> uniffiContinuationHandleMap = new UniffiHandleMap<>();
-    static final UniffiHandleMap<CompletableFuture<Void>> uniffiForeignFutureHandleMap = new UniffiHandleMap<>();
+    static final UniffiHandleMap<CancelableForeignFuture> uniffiForeignFutureHandleMap = new UniffiHandleMap<>();
 
     // FFI type for Rust future continuations
     enum UniffiRustFutureContinuationCallbackImpl implements UniffiRustFutureContinuationCallback {
@@ -49,6 +49,26 @@ public boolean cancel(boolean ignored) {
         }
     } 
 
+    // helper so both the Java completable future and the job that handles it finishing and reports to Rust can be
+    // retrieved (and potentially cancelled) by handle. This allows our FreeImpl to be a parameterless singleton,
+    // preventing #19, which was caused by our FreeImpls being GCd before Rust called back into them.
+    static class CancelableForeignFuture {
+        private CompletableFuture<?> childFuture;
+        private CompletableFuture<Void> childFutureHandler;
+
+        CancelableForeignFuture(CompletableFuture<?> childFuture, CompletableFuture<Void> childFutureHandler) {
+            this.childFuture = childFuture;
+            this.childFutureHandler = childFutureHandler;
+        }
+
+        public void cancel() {
+            var successfullyCancelled = this.childFutureHandler.cancel(true);
+            if(successfullyCancelled) {
+                childFuture.cancel(true);
+            }
+        }
+    }
+
     static <T, F, E extends Exception> CompletableFuture<T> uniffiRustCallAsync(
         long rustFuture,
         PollingFunction pollFunc,
@@ -131,11 +151,7 @@ private static byte poll(long rustFuture, PollingFunction pollFunc) throws Inter
         CompletableFuture<Byte> pollFuture = new CompletableFuture<>();
         var handle = uniffiContinuationHandleMap.insert(pollFuture);
         pollFunc.apply(rustFuture, UniffiRustFutureContinuationCallbackImpl.INSTANCE, handle);
-
-        // busy-wait until the poll completes
-        // TODO(java): may be more efficient to use a CountdownLatch here instead of a CF we end up busy-waiting
-        //     because of Java bugs.
-        do {} while (!pollFuture.isDone());
+        do {} while (!pollFuture.isDone()); // removing this makes futures not cancel (sometimes)
         return pollFuture.get();
     }
 
@@ -148,7 +164,7 @@ static <T> UniffiForeignFuture uniffiTraitInterfaceCallAsync(
         // Uniffi does its best to support structured concurrency across the FFI.
         // If the Rust future is dropped, `UniffiForeignFutureFreeImpl` is called, which will cancel the Java completable future if it's still running.
         var foreignFutureCf = makeCall.get();
-        CompletableFuture<Void> job = CompletableFuture.supplyAsync(() -> {
+        CompletableFuture<Void> ffHandler = CompletableFuture.supplyAsync(() -> {
             try {
                 foreignFutureCf.thenAcceptAsync(handleSuccess).get();
             } catch(Throwable e) {
@@ -166,9 +182,8 @@ static <T> UniffiForeignFuture uniffiTraitInterfaceCallAsync(
 
             return null;
         });
-
-        long handle = uniffiForeignFutureHandleMap.insert(job);
-        return new UniffiForeignFuture(handle, new UniffiForeignFutureFreeImpl<T>(foreignFutureCf));
+        long handle = uniffiForeignFutureHandleMap.insert(new CancelableForeignFuture(foreignFutureCf, ffHandler));
+        return new UniffiForeignFuture(handle, UniffiForeignFutureFreeImpl.INSTANCE);
     }
 
     @SuppressWarnings("unchecked")
@@ -180,7 +195,7 @@ static <T, E extends Throwable> UniffiForeignFuture uniffiTraitInterfaceCallAsyn
         Class<E> errorClass
     ){
         var foreignFutureCf = makeCall.get();
-        CompletableFuture<Void> job = CompletableFuture.supplyAsync(() -> {
+        CompletableFuture<Void> ffHandler = CompletableFuture.supplyAsync(() -> {
             try {
                 foreignFutureCf.thenAcceptAsync(handleSuccess).get();
             } catch (Throwable e) {
@@ -208,24 +223,17 @@ static <T, E extends Throwable> UniffiForeignFuture uniffiTraitInterfaceCallAsyn
             return null;
         });
 
-        long handle = uniffiForeignFutureHandleMap.insert(job);
-        return new UniffiForeignFuture(handle, new UniffiForeignFutureFreeImpl<T>(foreignFutureCf));
+        long handle = uniffiForeignFutureHandleMap.insert(new CancelableForeignFuture(foreignFutureCf, ffHandler));
+        return new UniffiForeignFuture(handle, UniffiForeignFutureFreeImpl.INSTANCE);
     }
 
-    static class UniffiForeignFutureFreeImpl<T> implements UniffiForeignFutureFree {
-        private CompletableFuture<T> childFuture;
-
-        UniffiForeignFutureFreeImpl(CompletableFuture<T> childFuture) {
-            this.childFuture = childFuture;
-        }
+    enum UniffiForeignFutureFreeImpl implements UniffiForeignFutureFree {
+        INSTANCE;
 
         @Override
         public void callback(long handle) {
-            var job = uniffiForeignFutureHandleMap.remove(handle);
-            var successfullyCancelled = job.cancel(true);
-            if(successfullyCancelled) {
-                childFuture.cancel(true);
-            }
+            var futureWithHandler = uniffiForeignFutureHandleMap.remove(handle);
+            futureWithHandler.cancel();
         }
     }
 

tests/scripts/TestFixtureFutures/TestFixtureFutures.java

@@ -38,7 +38,7 @@ public static long measureTimeMillis(FutureRunnable r) {
 
   public static void assertReturnsImmediately(long actualTime, String testName) {
     // TODO(java): 4ms limit in Kotlin
-    assert actualTime <= 20 : MessageFormat.format("unexpected {0} time: {1}ms", testName, actualTime);
+    assert actualTime <= 40 : MessageFormat.format("unexpected {0} time: {1}ms", testName, actualTime);
   }
 
   public static void assertApproximateTime(long actualTime, int expectedTime, String testName) {
@@ -230,6 +230,8 @@ public CompletableFuture<Void> tryDelay(String delayMs) {
         }
 
         var traitObj = new JavaAsyncParser();
+        var startingHandleCount = UniffiAsyncHelpers.uniffiForeignFutureHandleCount();
+        assert startingHandleCount == 0 : MessageFormat.format("{0} starting handle count != 0", startingHandleCount);
         assert Futures.asStringUsingTrait(traitObj, 1, 42).get().equals("42");
         assert Futures.tryFromStringUsingTrait(traitObj, 1, "42").get().equals(42);
         try {
@@ -271,7 +273,6 @@ public CompletableFuture<Void> tryDelay(String delayMs) {
         assert traitObj.completedDelays == completedDelaysBefore : MessageFormat.format("{0} current delays != {1} delays before", traitObj.completedDelays, completedDelaysBefore);
 
         // Test that all handles were cleaned up
-        System.gc();
         var endingHandleCount = UniffiAsyncHelpers.uniffiForeignFutureHandleCount();
         assert endingHandleCount == 0 : MessageFormat.format("{0} current handle count != 0", endingHandleCount);
       }

tests/tests.rs

@@ -2,7 +2,7 @@
 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/. */
 
-use anyhow::{bail, Context, Result};
+use anyhow::{Context, Result, bail};
 use camino::{Utf8Path, Utf8PathBuf};
 use cargo_metadata::{MetadataCommand, Package, Target};
 use std::env::consts::ARCH;