JNI: forward CompletableFuture cancellation to Swift tasks

Samples/SwiftJavaExtractJNISampleApp/Sources/MySwiftLibrary/CancellationProbe.swift

@@ -0,0 +1,40 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2025 Apple Inc. and the Swift.org project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift.org project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+public actor CancellationProbe {
+  private var cancelled = false
+
+  public init() {}
+
+  public func waitForCancellation() async throws -> Int64 {
+    try await suspend()
+    return 1
+  }
+
+  public func waitForOptionalCancellation() async throws -> Int64? {
+    try await suspend()
+    return 1
+  }
+
+  public func wasCancelled() -> Bool { cancelled }
+
+  private func suspend() async throws {
+    do {
+      try await Task.sleep(for: .seconds(60))
+    } catch {
+      cancelled = true
+      throw error
+    }
+  }
+}

Samples/SwiftJavaExtractJNISampleApp/src/test/java/com/example/swift/AsyncCancellationTest.java

@@ -0,0 +1,59 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2025 Apple Inc. and the Swift.org project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift.org project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+package com.example.swift;
+
+import org.junit.jupiter.api.Test;
+import org.swift.swiftkit.core.SwiftArena;
+import java.util.concurrent.*;
+import static org.junit.jupiter.api.Assertions.*;
+
+public class AsyncCancellationTest {
+    @Test
+    void cancellingAConvertedScalarFutureCancelsTheSwiftTask() throws Exception {
+        assertSwiftCancellation(false);
+    }
+
+    @Test
+    void cancellingAnOutParameterResultCancelsTheSwiftTask() throws Exception {
+        assertSwiftCancellation(true);
+    }
+
+    private void assertSwiftCancellation(boolean optional) throws Exception {
+        try (var arena = SwiftArena.ofConfined()) {
+            var probe = CancellationProbe.init(arena);
+            CompletableFuture<?> future = optional
+                ? probe.waitForOptionalCancellation() : probe.waitForCancellation();
+            assertTrue(future.cancel(true));
+            assertThrows(CancellationException.class, future::join);
+            long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+            boolean observed = false;
+            while (System.nanoTime() < deadline && !observed) {
+                observed = probe.wasCancelled().get(1, TimeUnit.SECONDS);
+                if (!observed) Thread.sleep(10);
+            }
+            assertTrue(observed, "Swift must observe cancellation of the returned future");
+            assertTrue(future.cancel(false));
+        }
+    }
+
+    @Test
+    void completedTasksCanBeBoundAndReleasedImmediately() throws Exception {
+        for (int i = 0; i < 100; i++) {
+            var future = MySwiftLibrary.asyncSum(i, 1);
+            assertEquals(i + 1L, future.get(5, TimeUnit.SECONDS));
+            assertFalse(future.cancel(false));
+        }
+    }
+}

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+JavaTranslation.swift

@@ -704,7 +704,7 @@ extension JNISwift2JavaGenerator {
 
       switch mode {
       case .completableFuture:
-        nativeFutureType = .completableFuture(nativeFunctionSignature.result.javaType)
+        nativeFutureType = .class(package: "org.swift.swiftkit.core", name: "SwiftCompletableFuture", typeParameters: [nativeFunctionSignature.result.javaType.boxedType])
         translatedFutureType = .completableFuture(translatedFunctionSignature.result.javaType)
         completeMethodID = "_JNIMethodIDCache.CompletableFuture.complete"
         completeExceptionallyMethodID = "_JNIMethodIDCache.CompletableFuture.completeExceptionally"
@@ -730,7 +730,7 @@ extension JNISwift2JavaGenerator {
         outParameters: result.outParameters + [futureOutParameter],
         conversion: .method(
           .constant("future$"),
-          function: "thenApply",
+          function: mode == .completableFuture ? "mapSwiftResult" : "thenApply",
           arguments: [
             .lambda(
               args: ["futureResult$"],

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+NativeTranslation.swift

@@ -2064,6 +2064,10 @@ extension JNISwift2JavaGenerator {
           }
         }
 
+        if completeMethodID == "_JNIMethodIDCache.CompletableFuture.complete" {
+          printer.print("SwiftJavaRuntimeSupport._JNIAsyncTask.bind(task!, to: result_future, in: environment)")
+        }
+
         return ""
 
       case .closure(let args, let body):

Sources/SwiftJavaDocumentation/Documentation.docc/FeaturesJextract.md

@@ -243,6 +243,19 @@ There are two modes of extracting them, configurable using the `asyncFuncMode` s
    }
 }
 
+#### Cancelling JNI async calls
+
+In `completableFuture` mode, cancelling the future returned by a JNI async call
+also cancels its Swift task, including when the generated bridge converts the
+result. Swift cancellation is cooperative: the operation must check cancellation
+or use a cancellation-aware suspension such as `Task.sleep`. Cancellation does
+not interrupt an operating-system thread or wait for the Swift task to finish.
+The `mayInterruptIfRunning` argument does not change this behavior.
+
+This forwarding applies to the future returned by the generated API. Additional
+Java future stages created by callers follow the usual `CompletableFuture`
+cancellation rules. The `legacyFuture` mode keeps its existing behavior.
+
 #### Asynchronous functions and Kotlin async/await
 
 Because Kotlin offers extension `suspending fun` methods on Future types, it is possible to `.await()`

Sources/SwiftJavaRuntimeSupport/_JNIAsyncTask.swift

@@ -0,0 +1,54 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2025 Apple Inc. and the Swift.org project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift.org project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftJava
+
+/// Owns the Swift task until its Java future completes or is cancelled.
+public final class _JNIAsyncTask: Sendable {
+  let task: Task<Void, Never>
+
+  init(_ task: Task<Void, Never>) { self.task = task }
+
+  private static let setTask = _JNIMethodIDCache.Method(name: "setSwiftTask", signature: "(J)V")
+  private static let methods = _JNIMethodIDCache(
+    className: "org/swift/swiftkit/core/SwiftCompletableFuture",
+    methods: [setTask]
+  )
+
+  public static func bind(_ task: Task<Void, Never>, to future: jobject?, in environment: JNIEnvironment) {
+    let handle = Unmanaged.passRetained(_JNIAsyncTask(task)).toOpaque()
+    _ = environment.interface.CallVoidMethodA(
+      environment,
+      future,
+      methods.methods[setTask]!,
+      [jvalue(j: Int64(Int(bitPattern: handle)))]
+    )
+  }
+}
+
+#if compiler(>=6.3)
+@used
+#endif
+@_cdecl("Java_org_swift_swiftkit_core_SwiftCompletableFuture_cancelSwiftTask")
+public func cancelSwiftAsyncTask(environment: UnsafeMutablePointer<JNIEnv?>!, type: jclass, handle: jlong) {
+  Unmanaged<_JNIAsyncTask>.fromOpaque(UnsafeRawPointer(bitPattern: Int(handle))!).takeUnretainedValue().task.cancel()
+}
+
+#if compiler(>=6.3)
+@used
+#endif
+@_cdecl("Java_org_swift_swiftkit_core_SwiftCompletableFuture_releaseSwiftTask")
+public func releaseSwiftAsyncTask(environment: UnsafeMutablePointer<JNIEnv?>!, type: jclass, handle: jlong) {
+  Unmanaged<_JNIAsyncTask>.fromOpaque(UnsafeRawPointer(bitPattern: Int(handle))!).release()
+}

SwiftKitCore/src/main/java/org/swift/swiftkit/core/SwiftCompletableFuture.java

@@ -0,0 +1,89 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2025 Apple Inc. and the Swift.org project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift.org project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+package org.swift.swiftkit.core;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
+
+/** A JNI task future whose cancellation reaches the Swift task, including after result conversion. */
+public final class SwiftCompletableFuture<T> extends CompletableFuture<T> {
+    private final Object taskLock = new Object();
+    private long taskHandle;
+    private final Runnable cancelSource;
+
+    public SwiftCompletableFuture() { this(null); }
+
+    private SwiftCompletableFuture(Runnable cancelSource) { this.cancelSource = cancelSource; }
+
+    public void setSwiftTask(long handle) {
+        synchronized (taskLock) {
+            if (!isDone()) { taskHandle = handle; return; }
+        }
+        if (isCancelled()) cancelSwiftTask(handle);
+        releaseSwiftTask(handle);
+    }
+
+    @Override
+    public boolean cancel(boolean mayInterruptIfRunning) {
+        boolean cancelled = super.cancel(mayInterruptIfRunning);
+        if (cancelled) {
+            finishTask(true);
+            if (cancelSource != null) cancelSource.run();
+        }
+        return cancelled;
+    }
+
+    @Override
+    public boolean complete(T value) {
+        boolean completed = super.complete(value);
+        if (completed) finishTask(false);
+        return completed;
+    }
+
+    @Override
+    public boolean completeExceptionally(Throwable error) {
+        boolean completed = super.completeExceptionally(error);
+        if (completed) finishTask(false);
+        return completed;
+    }
+
+    /** Unlike thenApply, cancelling this result also cancels its source task. */
+    public <U> CompletableFuture<U> mapSwiftResult(Function<? super T, ? extends U> transform) {
+        SwiftCompletableFuture<U> result = new SwiftCompletableFuture<>(() -> cancel(false));
+        whenComplete((value, error) -> {
+            if (error != null) result.completeExceptionally(error);
+            else {
+                try { result.complete(transform.apply(value)); }
+                catch (Throwable failure) { result.completeExceptionally(failure); }
+            }
+        });
+        return result;
+    }
+
+    private void finishTask(boolean cancel) {
+        long handle;
+        synchronized (taskLock) {
+            handle = taskHandle;
+            taskHandle = 0;
+        }
+        if (handle != 0) {
+            if (cancel) cancelSwiftTask(handle);
+            releaseSwiftTask(handle);
+        }
+    }
+
+    private static native void cancelSwiftTask(long handle);
+    private static native void releaseSwiftTask(long handle);
+}

SwiftKitCore/src/test/java/org/swift/swiftkit/core/SwiftCompletableFutureTest.java

@@ -0,0 +1,68 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2025 Apple Inc. and the Swift.org project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift.org project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+package org.swift.swiftkit.core;
+
+import org.junit.jupiter.api.Test;
+import java.util.concurrent.*;
+import static org.junit.jupiter.api.Assertions.*;
+
+public class SwiftCompletableFutureTest {
+    @Test
+    void mappingPreservesResultsAndAlreadyCompletedSources() {
+        var source = new SwiftCompletableFuture<String>();
+        var result = source.mapSwiftResult(String::length);
+        source.complete("hello");
+        assertEquals(5, result.join());
+        assertEquals(5, source.mapSwiftResult(String::length).join());
+        assertFalse(result.cancel(false));
+    }
+
+    @Test
+    void mappingPreservesFailures() {
+        var source = new SwiftCompletableFuture<String>();
+        var result = source.mapSwiftResult(String::length);
+        var error = new IllegalStateException("source failed");
+        source.completeExceptionally(error);
+        assertSame(error, assertThrows(CompletionException.class, result::join).getCause());
+    }
+
+    @Test
+    void mappingPreservesConversionFailures() {
+        var source = new SwiftCompletableFuture<String>();
+        var error = new IllegalArgumentException("conversion failed");
+        var result = source.mapSwiftResult(value -> { throw error; });
+        source.complete("hello");
+        assertSame(error, assertThrows(CompletionException.class, result::join).getCause());
+    }
+
+    @Test
+    void cancellingMappedResultCancelsTheSource() {
+        var source = new SwiftCompletableFuture<String>();
+        var result = source.mapSwiftResult(String::length);
+        assertTrue(result.cancel(false));
+        assertTrue(result.isCancelled());
+        assertTrue(source.isCancelled());
+        assertTrue(result.cancel(true));
+    }
+
+    @Test
+    void cancellingTheSourceCancelsTheMappedResult() {
+        var source = new SwiftCompletableFuture<String>();
+        var result = source.mapSwiftResult(String::length);
+        source.cancel(false);
+        assertTrue(result.isCancelled());
+        assertThrows(CancellationException.class, result::join);
+    }
+}