Check for cancellation of Rust task in test

Author: akonradi-signal

java/shared/java/org/signal/libsignal/internal/NativeTesting.java

@@ -95,7 +95,10 @@ private NativeTesting() {}
   public static native long TESTING_FakeChatServer_Create();
   public static native CompletableFuture<Long> TESTING_FakeChatServer_GetNextRemote(long asyncRuntime, long server);
   public static native CompletableFuture<Long> TESTING_FakeRegistrationSession_CreateSession(long asyncRuntime, Object createSession, long chat);
+  public static native long TESTING_FutureCancellationCounter_Create(int initialValue);
+  public static native CompletableFuture TESTING_FutureCancellationCounter_WaitForCount(long asyncRuntime, long count, int target);
   public static native CompletableFuture<Integer> TESTING_FutureFailure(long asyncRuntime, int input);
+  public static native CompletableFuture TESTING_FutureIncrementOnCancel(long asyncRuntime, long guard);
   public static native CompletableFuture<Long> TESTING_FutureProducesOtherPointerType(long asyncRuntime, String input);
   public static native CompletableFuture<Long> TESTING_FutureProducesPointerType(long asyncRuntime, int input);
   public static native CompletableFuture<Integer> TESTING_FutureSuccess(long asyncRuntime, int input);
@@ -104,7 +107,6 @@ private NativeTesting() {}
   public static native String TESTING_JoinStringArray(Object[] array, String joinWith);
   public static native void TESTING_NonSuspendingBackgroundThreadRuntime_Destroy(long handle);
   public static native long TESTING_NonSuspendingBackgroundThreadRuntime_New();
-  public static native CompletableFuture TESTING_OnlyCompletesByCancellation(long asyncRuntime);
   public static native String TESTING_OtherTestingHandleType_getValue(long handle);
   public static native void TESTING_PanicInBodyAsync(Object input);
   public static native CompletableFuture TESTING_PanicInBodyIo(long asyncRuntime, Object input);
@@ -139,6 +141,8 @@ private NativeTesting() {}
   public static native int TESTING_TestingHandleType_getValue(long handle);
   public static native CompletableFuture<Integer> TESTING_TokioAsyncFuture(long asyncRuntime, int input);
 
+  public static native void TestingFutureCancellationCounter_Destroy(long handle);
+
   public static native void TestingHandleType_Destroy(long handle);
 
   public static native int test_only_fn_returns_123();

java/shared/test/java/org/signal/libsignal/internal/FutureTest.java

@@ -68,15 +68,26 @@ public void testFutureFromRustCancel() {
     assertTrue(testFuture.isDone());
   }
 
-  @Test
+  @Test(timeout = 5000)
   @SuppressWarnings("unchecked")
-  public void testFutureOnlyCompletesByCancellation() {
+  public void testFutureOnlyCompletesByCancellation() throws Exception {
     TokioAsyncContext context = new TokioAsyncContext();
+    var counter =
+        new NativeHandleGuard.SimpleOwner(
+            NativeTesting.TESTING_FutureCancellationCounter_Create(0)) {
+          @Override
+          protected void release(long nativeHandle) {
+            NativeTesting.TestingFutureCancellationCounter_Destroy(nativeHandle);
+          }
+        };
     org.signal.libsignal.internal.CompletableFuture<Integer> testFuture =
         context
             .guardedMap(
                 (nativeContextHandle) ->
-                    NativeTesting.TESTING_OnlyCompletesByCancellation(nativeContextHandle))
+                    counter.guardedMap(
+                        counterHandle ->
+                            NativeTesting.TESTING_FutureIncrementOnCancel(
+                                nativeContextHandle, counterHandle)))
             .makeCancelable(context);
     assertTrue(testFuture.cancel(true));
     ExecutionException e = assertThrows(ExecutionException.class, () -> testFuture.get());
@@ -85,6 +96,16 @@ public void testFutureOnlyCompletesByCancellation() {
         e.getCause() instanceof java.util.concurrent.CancellationException);
     assertTrue(testFuture.isCancelled());
     assertTrue(testFuture.isDone());
+
+    // Hangs if the count never gets incremented.
+    context
+        .guardedMap(
+            (nativeContextHandle) ->
+                counter.guardedMap(
+                    counterHandle ->
+                        NativeTesting.TESTING_FutureCancellationCounter_WaitForCount(
+                            nativeContextHandle, counterHandle, 1)))
+        .get();
   }
 
   @Test

node/Native.d.ts

@@ -630,14 +630,16 @@ export function TESTING_FakeChatSentRequest_TakeHttpRequest(request: Wrapper<Fak
 export function TESTING_FakeChatServer_Create(): FakeChatServer;
 export function TESTING_FakeChatServer_GetNextRemote(asyncRuntime: Wrapper<TokioAsyncContext>, server: Wrapper<FakeChatServer>): CancellablePromise<FakeChatRemoteEnd>;
 export function TESTING_FakeRegistrationSession_CreateSession(asyncRuntime: Wrapper<TokioAsyncContext>, createSession: RegistrationCreateSessionRequest, chat: Wrapper<FakeChatServer>): CancellablePromise<RegistrationService>;
+export function TESTING_FutureCancellationCounter_Create(initialValue: number): TestingFutureCancellationCounter;
+export function TESTING_FutureCancellationCounter_WaitForCount(asyncRuntime: Wrapper<TokioAsyncContext>, count: Wrapper<TestingFutureCancellationCounter>, target: number): CancellablePromise<void>;
 export function TESTING_FutureFailure(asyncRuntime: Wrapper<NonSuspendingBackgroundThreadRuntime>, _input: number): CancellablePromise<number>;
+export function TESTING_FutureIncrementOnCancel(asyncRuntime: Wrapper<TokioAsyncContext>, _guard: TestingFutureCancellationGuard): CancellablePromise<void>;
 export function TESTING_FutureProducesOtherPointerType(asyncRuntime: Wrapper<NonSuspendingBackgroundThreadRuntime>, input: string): CancellablePromise<OtherTestingHandleType>;
 export function TESTING_FutureProducesPointerType(asyncRuntime: Wrapper<NonSuspendingBackgroundThreadRuntime>, input: number): CancellablePromise<TestingHandleType>;
 export function TESTING_FutureSuccess(asyncRuntime: Wrapper<NonSuspendingBackgroundThreadRuntime>, input: number): CancellablePromise<number>;
 export function TESTING_InputStreamReadIntoZeroLengthSlice(capsAlphabetInput: InputStream): Promise<Buffer>;
 export function TESTING_JoinStringArray(array: string[], joinWith: string): string;
 export function TESTING_NonSuspendingBackgroundThreadRuntime_New(): NonSuspendingBackgroundThreadRuntime;
-export function TESTING_OnlyCompletesByCancellation(asyncRuntime: Wrapper<TokioAsyncContext>): CancellablePromise<void>;
 export function TESTING_OtherTestingHandleType_getValue(handle: Wrapper<OtherTestingHandleType>): string;
 export function TESTING_PanicInBodyAsync(_input: null): Promise<void>;
 export function TESTING_PanicInBodyIo(asyncRuntime: Wrapper<NonSuspendingBackgroundThreadRuntime>, _input: null): CancellablePromise<void>;
@@ -771,6 +773,7 @@ interface SessionRecord { readonly __type: unique symbol; }
 interface SgxClientState { readonly __type: unique symbol; }
 interface SignalMessage { readonly __type: unique symbol; }
 interface SignedPreKeyRecord { readonly __type: unique symbol; }
+interface TestingFutureCancellationCounter { readonly __type: unique symbol; }
 interface TestingHandleType { readonly __type: unique symbol; }
 interface TokioAsyncContext { readonly __type: unique symbol; }
 interface UnauthenticatedChatConnection { readonly __type: unique symbol; }

node/ts/test/FutureTest.ts

@@ -18,6 +18,22 @@ function makeAsyncRuntime(): Native.Wrapper<Native.NonSuspendingBackgroundThread
   };
 }
 
+class CancelCounter {
+  readonly _nativeHandle: Native.TestingFutureCancellationCounter;
+  constructor(initialValue: number = 0) {
+    this._nativeHandle =
+      Native.TESTING_FutureCancellationCounter_Create(initialValue);
+  }
+
+  public async waitForCount(context: TokioAsyncContext, target: number) {
+    await Native.TESTING_FutureCancellationCounter_WaitForCount(
+      context,
+      this,
+      target
+    );
+  }
+}
+
 describe('Async runtime not on the Node executor', () => {
   it('handles success', async () => {
     const runtime = makeAsyncRuntime();
@@ -42,28 +58,36 @@ describe('TokioAsyncContext', () => {
   it('supports cancellation of running future', async () => {
     const runtime = new TokioAsyncContext(Native.TokioAsyncContext_new());
     const abortController = new AbortController();
+    const counter = new CancelCounter();
     const pending = runtime.makeCancellable(
       abortController.signal,
-      Native.TESTING_OnlyCompletesByCancellation(runtime)
+      Native.TESTING_FutureIncrementOnCancel(runtime, counter)
     );
     const timeout = setTimeout(200, 'timed out');
     assert.equal('timed out', await Promise.race([pending, timeout]));
     abortController.abort();
-    return expect(pending)
-      .to.eventually.be.rejectedWith(LibSignalErrorBase)
-      .and.have.property('code', ErrorCode.Cancelled);
+    return Promise.all([
+      expect(pending)
+        .to.eventually.be.rejectedWith(LibSignalErrorBase)
+        .and.have.property('code', ErrorCode.Cancelled),
+      expect(counter.waitForCount(runtime, 1)).to.be.fulfilled,
+    ]);
   });
 
   it('supports pre-cancellation of not-yet-running future', async () => {
     const runtime = new TokioAsyncContext(Native.TokioAsyncContext_new());
     const abortController = new AbortController();
+    const counter = new CancelCounter();
     abortController.abort();
     const pending = runtime.makeCancellable(
       abortController.signal,
-      Native.TESTING_OnlyCompletesByCancellation(runtime)
+      Native.TESTING_FutureIncrementOnCancel(runtime, counter)
     );
-    return expect(pending)
-      .to.eventually.be.rejectedWith(LibSignalErrorBase)
-      .and.have.property('code', ErrorCode.Cancelled);
+    return Promise.all([
+      expect(pending)
+        .to.eventually.be.rejectedWith(LibSignalErrorBase)
+        .and.have.property('code', ErrorCode.Cancelled),
+      expect(counter.waitForCount(runtime, 1)).to.be.fulfilled,
+    ]);
   });
 });

rust/bridge/shared/macros/src/jni.rs

@@ -158,10 +158,16 @@ fn bridge_io_body(
                 // Wrap the actual work to catch any panics.
                 let __future = jni::catch_unwind(std::panic::AssertUnwindSafe(async {
                     #(#input_loading)*
-                    let __result = #orig_name(#(#input_names),*).await;
-                    // If the original function can't fail, wrap the result in Ok for uniformity.
-                    // See TransformHelper::ok_if_needed.
-                    Ok(TransformHelper(__result).ok_if_needed()?.0)
+                        ::tokio::select! {
+                            __result = #orig_name(#(#input_names),*) => {
+                                // If the original function can't fail, wrap the result in Ok for uniformity.
+                                // See TransformHelper::ok_if_needed.
+                                Ok(TransformHelper(__result).ok_if_needed()?.0)
+                            }
+                            _ = __cancel => {
+                                Err(jni::FutureCancelled.into())
+                            }
+                        }
                 }));
                 // Pass the stored inputs to the reporter to drop them while attached to the JVM.
 

rust/bridge/shared/testing/src/convert.rs

@@ -74,6 +74,30 @@ async fn TESTING_FutureFailure(_input: u8) -> Result<i32, SignalProtocolError> {
     Err(SignalProtocolError::InvalidArgument("failure".to_string()))
 }
 
+bridge_handle_fns!(TestingFutureCancellationCounter, clone = false);
+
+#[bridge_fn]
+fn TESTING_FutureCancellationCounter_Create(initial_value: u8) -> TestingFutureCancellationCounter {
+    TestingFutureCancellationCounter(tokio::sync::Semaphore::new(initial_value.into()).into())
+}
+
+#[bridge_io(TokioAsyncContext)]
+async fn TESTING_FutureCancellationCounter_WaitForCount(
+    count: &TestingFutureCancellationCounter,
+    target: u8,
+) {
+    let _permits = count
+        .0
+        .acquire_many(target.into())
+        .await
+        .expect("not closed");
+}
+
+#[bridge_io(TokioAsyncContext)]
+async fn TESTING_FutureIncrementOnCancel(_guard: TestingFutureCancellationGuard) {
+    std::future::pending().await
+}
+
 #[bridge_io(TokioAsyncContext)]
 async fn TESTING_TokioAsyncFuture(input: u8) -> i32 {
     i32::from(input) * 3

rust/bridge/shared/testing/src/net.rs

@@ -49,11 +49,6 @@ async fn TESTING_CdsiLookupResponseConvert() -> LookupResponse {
     }
 }
 
-#[bridge_io(TokioAsyncContext)]
-async fn TESTING_OnlyCompletesByCancellation() {
-    std::future::pending::<()>().await
-}
-
 macro_rules! make_error_testing_enum {
     (enum $name:ident for $orig:ident {
         $($orig_case:ident => $case:ident,)*

rust/bridge/shared/testing/src/types.rs

@@ -3,6 +3,9 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 //
 
+use std::panic::{RefUnwindSafe, UnwindSafe};
+use std::sync::Arc;
+
 #[allow(unused_imports)]
 use libsignal_protocol::SignalProtocolError;
 
@@ -347,3 +350,74 @@ impl<'a> node::ResultTypeInfo<'a> for PanicOnReturn {
         panic!("deliberate panic");
     }
 }
+
+/// Counter for future cancellations
+pub struct TestingFutureCancellationCounter(pub(crate) Arc<tokio::sync::Semaphore>);
+
+impl UnwindSafe for TestingFutureCancellationCounter {}
+impl RefUnwindSafe for TestingFutureCancellationCounter {}
+
+/// RAII guard that increments a counter on `Drop`.
+///
+/// This is bridged as a reference to a [`TestingFutureCancellationCounter`].
+pub struct TestingFutureCancellationGuard {
+    increment_on_drop: Arc<tokio::sync::Semaphore>,
+}
+
+impl Drop for TestingFutureCancellationGuard {
+    fn drop(&mut self) {
+        self.increment_on_drop.add_permits(1);
+    }
+}
+
+#[cfg(feature = "ffi")]
+impl ffi::SimpleArgTypeInfo for TestingFutureCancellationGuard {
+    type ArgType = <&'static TestingFutureCancellationCounter as ffi::SimpleArgTypeInfo>::ArgType;
+
+    fn convert_from(foreign: Self::ArgType) -> ffi::SignalFfiResult<Self> {
+        <&TestingFutureCancellationCounter as ffi::SimpleArgTypeInfo>::convert_from(foreign).map(
+            |TestingFutureCancellationCounter(counter)| TestingFutureCancellationGuard {
+                increment_on_drop: Arc::clone(counter),
+            },
+        )
+    }
+}
+
+#[cfg(feature = "jni")]
+impl<'a> jni::SimpleArgTypeInfo<'a> for TestingFutureCancellationGuard {
+    type ArgType = <&'a TestingFutureCancellationCounter as jni::SimpleArgTypeInfo<'a>>::ArgType;
+
+    fn convert_from(
+        env: &mut jni::JNIEnv<'a>,
+        foreign: &Self::ArgType,
+    ) -> Result<Self, jni::BridgeLayerError> {
+        <&TestingFutureCancellationCounter as jni::SimpleArgTypeInfo>::convert_from(env, foreign)
+            .map(
+                |TestingFutureCancellationCounter(counter)| TestingFutureCancellationGuard {
+                    increment_on_drop: Arc::clone(counter),
+                },
+            )
+    }
+}
+
+#[cfg(feature = "node")]
+impl<'storage> node::AsyncArgTypeInfo<'storage> for TestingFutureCancellationGuard {
+    type ArgType = node::JsObject;
+    type StoredType = Option<node::DefaultFinalize<Self>>;
+    fn save_async_arg(
+        cx: &mut neon::prelude::FunctionContext,
+        foreign: neon::prelude::Handle<Self::ArgType>,
+    ) -> neon::prelude::NeonResult<Self::StoredType> {
+        <&TestingFutureCancellationCounter as node::AsyncArgTypeInfo>::save_async_arg(cx, foreign)
+            .map(move |handle| {
+                Some(node::DefaultFinalize(TestingFutureCancellationGuard {
+                    increment_on_drop: Arc::clone(&handle.0),
+                }))
+            })
+    }
+    fn load_async_arg(stored: &'storage mut Self::StoredType) -> Self {
+        stored.take().unwrap().0
+    }
+}
+
+bridge_as_handle!(TestingFutureCancellationCounter);

rust/bridge/shared/types/src/ffi/convert.rs

@@ -1082,6 +1082,8 @@ macro_rules! ffi_arg_type {
     (Ignored<$typ:ty>) => (*const std::ffi::c_void);
     (AsType<$typ:ident, $bridged:ident>) => (ffi_arg_type!($bridged));
 
+    (TestingFutureCancellationGuard) => (ffi_arg_type!(&TestingFutureCancellationCounter));
+
     // In order to provide a fixed-sized array of the correct length,
     // a serialized type FooBar must have a constant FOO_BAR_LEN that's in scope (and exposed to C).
     (Serialized<$typ:ident>) => (*const [std::ffi::c_uchar; ::paste::paste!([<$typ:snake:upper _LEN>])]);

rust/bridge/shared/types/src/jni/convert.rs

@@ -1933,6 +1933,7 @@ macro_rules! jni_arg_type {
     (CreateSession) => {
         $crate::jni::JObject<'local>
     };
+    (TestingFutureCancellationGuard) => { ::jni::sys::jlong };
 
     (Ignored<$typ:ty>) => (::jni::objects::JObject<'local>);
 }