JNI: cancelling an async function's future cancels the Swift task

Samples/SwiftJavaExtractJNIKotlinSampleApp/Sources/KotlinSampleLibrary/Cancellation.swift

@@ -0,0 +1,36 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2026 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 Synchronization
+
+/// Long-running async work: cancelling the Kotlin coroutine awaiting it cancels the Swift task.
+public final class Download: Sendable {
+  private let cancelled = Mutex(false)
+
+  public init() {}
+
+  public func run() async -> Bool {
+    do {
+      try await Task.sleep(for: .seconds(30))
+      return false
+    } catch {
+      cancelled.withLock { $0 = true }
+      return true
+    }
+  }
+
+  public func wasCancelled() -> Bool {
+    cancelled.withLock { $0 }
+  }
+}

Samples/SwiftJavaExtractJNIKotlinSampleApp/src/test/kotlin/com/example/kotlin/CancellationTest.kt

@@ -0,0 +1,36 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2026 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.kotlin
+
+import kotlinx.coroutines.future.await
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withTimeoutOrNull
+import kotlin.test.Test
+import kotlin.test.assertNull
+import kotlin.test.assertTrue
+
+class CancellationTest {
+    @Test
+    fun cancellingTheCoroutineCancelsTheSwiftTask() = runBlocking {
+        val download = Download.init()
+        // await() cancels the future when its coroutine is cancelled, here by the timeout.
+        val result = withTimeoutOrNull(100) { download.run().await() }
+        assertNull(result)
+
+        val deadline = System.currentTimeMillis() + 5_000
+        while (!download.wasCancelled() && System.currentTimeMillis() < deadline) Thread.sleep(20)
+        assertTrue(download.wasCancelled())
+    }
+}

Samples/SwiftJavaExtractJNISampleApp/Sources/MySwiftLibrary/AsyncCancellation.swift

@@ -0,0 +1,37 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2026 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 Synchronization
+
+/// Cancelling the Java future returned for an `async` function cancels the Swift task.
+public final class CancellableWork: Sendable {
+  private let cancelled = Mutex(false)
+
+  public init() {}
+
+  /// Sleeps until cancelled; returns whether it was.
+  public func waitUntilCancelled() async -> Bool {
+    do {
+      try await Task.sleep(for: .seconds(30))
+      return false
+    } catch {
+      cancelled.withLock { $0 = true }
+      return true
+    }
+  }
+
+  public func wasCancelled() -> Bool {
+    cancelled.withLock { $0 }
+  }
+}

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

@@ -87,4 +87,20 @@ void asyncString() throws Exception {
         Future<String> future = MySwiftLibrary.asyncString("hey");
         assertEquals("hey", future.get());
     }
-}
+
+    @Test
+    void cancellingTheFutureCancelsTheSwiftTask() throws Exception {
+        try (var arena = SwiftArena.ofConfined()) {
+            CancellableWork work = CancellableWork.init(arena);
+            var future = work.waitUntilCancelled();
+            Thread.sleep(100);
+            assertTrue(future.cancel(true));
+
+            long deadline = System.currentTimeMillis() + 5_000;
+            while (!work.wasCancelled() && System.currentTimeMillis() < deadline) {
+                Thread.sleep(20);
+            }
+            assertTrue(work.wasCancelled());
+        }
+    }
+}

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+JavaTranslation.swift

@@ -741,21 +741,29 @@ extension JNISwift2JavaGenerator {
       )
 
       let result = translatedFunctionSignature.result
+      let convertedFuture = JavaNativeConversionStep.method(
+        .constant("future$"),
+        function: "thenApply",
+        arguments: [
+          .lambda(
+            args: ["futureResult$"],
+            body: .replacingPlaceholder(result.conversion, placeholder: "futureResult$")
+          )
+        ]
+      )
       translatedFunctionSignature.result = TranslatedResult(
         javaType: translatedFutureType,
         nativeJavaType: .void,
         annotations: result.annotations,
         outParameters: result.outParameters + [futureOutParameter],
-        conversion: .method(
-          .constant("future$"),
-          function: "thenApply",
-          arguments: [
-            .lambda(
-              args: ["futureResult$"],
-              body: .replacingPlaceholder(result.conversion, placeholder: "futureResult$")
-            )
-          ]
-        )
+        // Cancelling the returned (derived) future must reach `future$`, which the Swift task watches.
+        conversion: mode == .completableFuture
+          ? .method(
+            .constant("org.swift.swiftkit.core.CompletableFutureCancellation"),
+            function: "forwardCancellation",
+            arguments: [convertedFuture, .constant("future$")]
+          )
+          : convertedFuture
       )
 
       // Update native function
@@ -765,6 +773,7 @@ extension JNISwift2JavaGenerator {
         isThrowing: originalFunctionSignature.isThrowing || originalFunctionSignature.isImplicitlyThrowing,
         completeMethodID: completeMethodID,
         completeExceptionallyMethodID: completeExceptionallyMethodID,
+        cancelsTaskWithFuture: mode == .completableFuture,
       )
       nativeFunctionSignature.result.javaType = .void
       nativeFunctionSignature.result.outParameters.append(.init(name: "result_future", type: nativeFutureType))

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+NativeTranslation.swift

@@ -1465,7 +1465,8 @@ extension JNISwift2JavaGenerator {
       nativeFunctionSignature: NativeFunctionSignature,
       isThrowing: Bool,
       completeMethodID: String,
-      completeExceptionallyMethodID: String
+      completeExceptionallyMethodID: String,
+      cancelsTaskWithFuture: Bool
     )
 
     /// `{ (args) -> return body }`
@@ -1968,7 +1969,8 @@ extension JNISwift2JavaGenerator {
         let nativeFunctionSignature,
         let isThrowing,
         let completeMethodID,
-        let completeExceptionallyMethodID
+        let completeExceptionallyMethodID,
+        let cancelsTaskWithFuture
       ):
         var globalRefs: [String] = ["globalFuture"]
 
@@ -2101,6 +2103,16 @@ extension JNISwift2JavaGenerator {
           }
         }
 
+        if cancelsTaskWithFuture {
+          // Cancelling the Java future cancels the task. Registered with the call's own (local)
+          // reference, which stays valid even if an immediate task already finished.
+          printer.printIfBlock("let task") { printer in
+            printer.print(
+              "SwiftJavaRuntimeSupport._cancelTaskWhenFutureIsCancelled(result_future, task: task, environment: environment)"
+            )
+          }
+        }
+
         return ""
 
       case .closure(let args, let body):

Sources/SwiftJavaRuntimeSupport/CompletableFutureAwait.swift

@@ -102,3 +102,58 @@ public func Java_org_swift_swiftkit_core_CompletableFutureContinuation__00024res
     )
   )
 }
+
+// ==== -----------------------------------------------------------------------
+// MARK: Cancelling the Swift task behind an exported `async` function
+
+@JavaClass("org.swift.swiftkit.core.CompletableFutureCancellation")
+class JavaCompletableFutureCancellation: JavaObject {}
+
+extension JavaClass<JavaCompletableFutureCancellation> {
+  @JavaStaticMethod
+  func cancelTaskWhenCancelled(_ future: JavaCompletableFuture?, _ task: Int64)
+}
+
+private final class SwiftTaskHandle: @unchecked Sendable {
+  let task: Task<Void, Never>
+
+  init(_ task: Task<Void, Never>) {
+    self.task = task
+  }
+}
+
+/// Cancels `task` when the Java `future` it completes is cancelled.
+///
+/// Used by generated code for `async` functions exported to Java. Pass the future the native call
+/// received, which stays valid for the duration of the call even if the task has already finished.
+public func _cancelTaskWhenFutureIsCancelled(
+  _ future: jobject?,
+  task: Task<Void, Never>,
+  environment: JNIEnvironment
+) {
+  guard let future else { return }
+  let handle = Unmanaged.passRetained(SwiftTaskHandle(task)).toOpaque()
+  let cancellationClass = try! JavaClass<JavaCompletableFutureCancellation>(environment: environment)
+  cancellationClass.cancelTaskWhenCancelled(
+    JavaCompletableFuture(javaThis: future, environment: environment),
+    Int64(Int(bitPattern: handle))
+  )
+}
+
+@_cdecl("Java_org_swift_swiftkit_core_CompletableFutureCancellation__00024cancelTask")
+public func Java_org_swift_swiftkit_core_CompletableFutureCancellation__00024cancelTask(
+  environment: UnsafeMutablePointer<JNIEnv?>!,
+  thisClass: jclass,
+  task: jlong
+) {
+  Unmanaged<SwiftTaskHandle>.fromOpaque(UnsafeRawPointer(bitPattern: Int(task))!).takeUnretainedValue().task.cancel()
+}
+
+@_cdecl("Java_org_swift_swiftkit_core_CompletableFutureCancellation__00024releaseTask")
+public func Java_org_swift_swiftkit_core_CompletableFutureCancellation__00024releaseTask(
+  environment: UnsafeMutablePointer<JNIEnv?>!,
+  thisClass: jclass,
+  task: jlong
+) {
+  Unmanaged<SwiftTaskHandle>.fromOpaque(UnsafeRawPointer(bitPattern: Int(task))!).release()
+}

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

@@ -0,0 +1,55 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2026 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;
+
+/**
+ * Propagates cancellation of the {@link CompletableFuture} returned for an {@code async} Swift
+ * function to the Swift {@code Task} running it.
+ */
+public final class CompletableFutureCancellation {
+    private CompletableFutureCancellation() {}
+
+    /**
+     * Cancels {@code source} when {@code result}, a future derived from it, is cancelled.
+     * Cancelling a derived future (e.g. from {@code thenApply}) does not otherwise reach its source.
+     */
+    public static <T> CompletableFuture<T> forwardCancellation(CompletableFuture<T> result, CompletableFuture<?> source) {
+        result.whenComplete((value, error) -> {
+            if (result.isCancelled()) {
+                source.cancel(true);
+            }
+        });
+        return result;
+    }
+
+    /**
+     * Cancels the Swift task at {@code task} if {@code future} is cancelled, and releases the task
+     * handle once {@code future} completes in any way.
+     */
+    public static void cancelTaskWhenCancelled(CompletableFuture<?> future, long task) {
+        future.whenComplete((value, error) -> {
+            if (future.isCancelled()) {
+                $cancelTask(task);
+            }
+            $releaseTask(task);
+        });
+    }
+
+    private static native void $cancelTask(long task);
+
+    private static native void $releaseTask(long task);
+}

Tests/JExtractSwiftTests/JNI/JNIActorTests.swift

@@ -45,10 +45,10 @@ struct JNIActor1MethodsTests {
         public java.util.concurrent.CompletableFuture<java.lang.Void> hello() {
           java.util.concurrent.CompletableFuture<java.lang.Void> future$ = new java.util.concurrent.CompletableFuture<java.lang.Void>();
           K.$hello(this.$memoryAddress(), future$);
-          return future$.thenApply((futureResult$) -> {
+          return org.swift.swiftkit.core.CompletableFutureCancellation.forwardCancellation(future$.thenApply((futureResult$) -> {
             return futureResult$;
           }
-          );
+          ), future$);
         }
         """,
         """
@@ -93,10 +93,10 @@ struct JNIActor1MethodsTests {
         public java.util.concurrent.CompletableFuture<java.lang.Void> hi() {
           java.util.concurrent.CompletableFuture<java.lang.Void> future$ = new java.util.concurrent.CompletableFuture<java.lang.Void>();
           K.$hi(this.$memoryAddress(), future$);
-          return future$.thenApply((futureResult$) -> {
+          return org.swift.swiftkit.core.CompletableFutureCancellation.forwardCancellation(future$.thenApply((futureResult$) -> {
             return futureResult$;
           }
-          );
+          ), future$);
         }
         """,
         """
@@ -237,10 +237,10 @@ struct JNIActor1MethodsTests {
         public java.util.concurrent.CompletableFuture<java.lang.Void> concurrentlyAsync() {
           java.util.concurrent.CompletableFuture<java.lang.Void> future$ = new java.util.concurrent.CompletableFuture<java.lang.Void>();
           K.$concurrentlyAsync(this.$memoryAddress(), future$);
-          return future$.thenApply((futureResult$) -> {
+          return org.swift.swiftkit.core.CompletableFutureCancellation.forwardCancellation(future$.thenApply((futureResult$) -> {
             return futureResult$;
           }
-          );
+          ), future$);
         }
         """,
         """