Add JS task cancellation via an AbortSignal

crates/ubrn_bindgen/src/bindings/react_native/gen_typescript/templates/CallbackInterfaceImpl.ts

@@ -22,13 +22,21 @@ const {{ trait_impl }}: { vtable: {{ vtable|ffi_type_name }}; register: () => vo
             {%- endif %}
         ) => {
             const uniffiMakeCall = {# space #}
-            {%- call ts::async(meth) -%}
-            (): {% call ts::return_type(meth) %} => {
+            {%- if meth.is_async() %}
+            async (signal: AbortSignal)
+            {%- else %}
+            ()
+            {%- endif %}
+            : {% call ts::return_type(meth) %} => {
                 const jsCallback = {{ ffi_converter_name }}.lift(uniffiHandle);
                 return {% call ts::await(meth) %}jsCallback.{{ meth.name()|fn_name }}(
                     {%- for arg in meth.arguments() %}
                     {{ arg|ffi_converter_name(self) }}.lift({{ arg.name()|var_name }}){% if !loop.last %}, {% endif %}
                     {%- endfor %}
+                    {%- if meth.is_async() -%}
+                    {%-   if !meth.arguments().is_empty() %}, {% endif -%}
+                    { signal }
+                    {%- endif %}
                 )
             }
             {%- if !meth.is_async() %}
@@ -60,7 +68,7 @@ const {{ trait_impl }}: { vtable: {{ vtable|ffi_type_name }}; register: () => vo
                 /*lowerString:*/ FfiConverterString.lower
             )
             {%- endmatch %}
-            {%- else %}
+            {%- else %} {# // is_async = true #}
 
             const uniffiHandleSuccess = (returnValue: {% call ts::raw_return_type(meth) %}) => {
                 uniffiFutureCallback(

crates/ubrn_bindgen/src/bindings/react_native/gen_typescript/templates/macros.ts

@@ -215,11 +215,20 @@ import { decl_type_name } from "./EnumTemplate"
 {%- endif -%}
 {%- endmacro %}
 
+{#-
+// This macros is almost identical to `arg_list_decl`,
+// but is for interface methods, which do not allow
+// default values for arguments.
+#}
 {% macro arg_list_protocol(func) %}
     {%- for arg in func.arguments() -%}
         {{ arg.name()|var_name }}: {{ arg|type_name(self) -}}
         {%- if !loop.last %}, {% endif -%}
     {%- endfor %}
+    {%- if func.is_async() %}
+    {%-   if !func.arguments().is_empty() %}, {% endif -%}
+    asyncOpts_?: { signal: AbortSignal }
+    {%- endif %}
 {%- endmacro %}
 
 {%- macro async(func) %}

fixtures/futures/tests/bindings/test_futures.ts

@@ -49,6 +49,18 @@ function delayPromise(delayMs: number): Promise<void> {
   });
 }
 
+function cancellableDelayPromise(
+  delayMs: number,
+  abortSignal: AbortSignal,
+): Promise<void> {
+  return new Promise((resolve, reject) => {
+    const timer = setTimeout(resolve, delayMs);
+    abortSignal.addEventListener("abort", () => {
+      reject(abortSignal.reason);
+    });
+  });
+}
+
 function checkRemainingFutures(t: Asserts) {
   t.assertEqual(
     0,
@@ -282,6 +294,38 @@ function checkRemainingFutures(t: Asserts) {
     t.end();
   });
 
+  class CancellableTsAsyncParser extends TsAsyncParser {
+    /**
+     * Each async callback method has an additional optional argument
+     * `asyncOptions_`. This contains an `AbortSignal`.
+     *
+     * If the Rust task is cancelled, then this abort signal is
+     * told, which can be used to co-operatively cancel the
+     * async callback.
+     *
+     * @param delayMs
+     * @param asyncOptions_
+     */
+    async delay(
+      delayMs: number,
+      asyncOptions_?: { signal: AbortSignal },
+    ): Promise<void> {
+      await this.doCancellableDelay(delayMs, asyncOptions_?.signal);
+    }
+
+    private async doCancellableDelay(
+      ms: number,
+      signal?: AbortSignal,
+    ): Promise<void> {
+      if (signal) {
+        await cancellableDelayPromise(ms, signal);
+      } else {
+        await delayPromise(ms);
+      }
+      this.completedDelays += 1;
+    }
+  }
+
   /**
    * Rust supports task cancellation, but it's not automatic. It is rather like
    * Javascript's.
@@ -460,28 +504,61 @@ function checkRemainingFutures(t: Asserts) {
     },
   );
 
-  await xasyncTest(
+  class Counter {
+    expectedCount = 0;
+    unexpectedCount = 0;
+    ok() {
+      return () => this.expectedCount++;
+    }
+    wat() {
+      return () => this.unexpectedCount++;
+    }
+  }
+
+  await asyncTest(
     "a future that uses a lock and that is cancelled from JS",
     async (t) => {
+      const errors = new Counter();
+      const success = new Counter();
+
+      // Task 1 should hold the resource for 100 seconds.
+      // We make an abort controller and get the signal from it, and pass it to
+      // Rust.
+      // Cancellation is done by dropping the future, so the Rust should be prepared
+      // for that.
+      const abortController = new AbortController();
       const task1 = useSharedResource(
         SharedResourceOptions.create({
-          releaseAfterMs: 5000,
+          releaseAfterMs: 100000,
           timeoutMs: 100,
         }),
-      );
-      // #RUST_TASK_CANCELLATION
-      //
-      // Again this test is not really applicable for JS, as it has no standard way of
-      // cancelling a task.
-      // task1.cancel()
-
-      // Try accessing the shared resource again.  The initial task should release the shared resource
-      // before the timeout expires.
+        { signal: abortController.signal },
+      ).then(success.wat(), errors.ok());
+
+      // Task 2 should try to grab the resource, but timeout after 1 second.
+      // Unless we abort task 1, then task 1 will hold on, but task 2 will timeout and
+      // fail.
       const task2 = useSharedResource(
         SharedResourceOptions.create({ releaseAfterMs: 0, timeoutMs: 1000 }),
+      ).then(success.ok(), errors.wat());
+
+      // We wait for 500 ms, then call the abortController.abort().
+      const delay = delayPromise(500).then(() => abortController.abort());
+
+      await Promise.allSettled([task1, task2, delay]);
+      t.assertEqual(errors.expectedCount, 1, "only task1 should have failed");
+      t.assertEqual(
+        success.expectedCount,
+        1,
+        "only task2 should have succeeded",
       );
 
-      await Promise.allSettled([task1, task2]);
+      t.assertEqual(errors.unexpectedCount, 0, "task2 should not have failed");
+      t.assertEqual(
+        success.unexpectedCount,
+        0,
+        "task1 should not have succeeded",
+      );
       checkRemainingFutures(t);
       t.end();
     },

typescript/src/async-callbacks.ts

@@ -6,13 +6,28 @@
 import { CALL_ERROR, CALL_UNEXPECTED_ERROR } from "./rust-call";
 import { type UniffiHandle, UniffiHandleMap } from "./handle-map";
 
-const UNIFFI_FOREIGN_FUTURE_HANDLE_MAP = new UniffiHandleMap<Promise<any>>();
+// Some additional data we hold for each in-flight promise.
+type PromiseHelper = {
+  // The promise itself, so it doesn't get GCd by mistake.
+  promise: Promise<any>;
+  // The abort controller we will use to cancel, if necessary.
+  abortController: AbortController;
+  // A mutable object which gets set when the promise has succeeded or errored.
+  // If uniffiForeignFutureFree gets called before settled is turned to true,
+  // then we know that it is a call to cancel the task.
+  settledHolder: {
+    settled: boolean;
+  };
+};
+const UNIFFI_FOREIGN_FUTURE_HANDLE_MAP = new UniffiHandleMap<PromiseHelper>();
 
+// Some degenerate functions used for default arguments.
 const notExpectedError = (err: any) => false;
 function emptyLowerError<E>(e: E): ArrayBuffer {
   throw new Error("Unreachable");
 }
 
+// Callbacks passed into Rust.
 export type UniffiForeignFutureFree = (handle: bigint) => void;
 
 export type UniffiForeignFuture = {
@@ -21,7 +36,7 @@ export type UniffiForeignFuture = {
 };
 
 export function uniffiTraitInterfaceCallAsync<T>(
-  makeCall: () => Promise<T>,
+  makeCall: (signal: AbortSignal) => Promise<T>,
   handleSuccess: (value: T) => void,
   handleError: (callStatus: /*i8*/ number, errorBuffer: ArrayBuffer) => void,
   lowerString: (str: string) => ArrayBuffer,
@@ -37,51 +52,54 @@ export function uniffiTraitInterfaceCallAsync<T>(
 }
 
 export function uniffiTraitInterfaceCallAsyncWithError<T, E>(
-  makeCall: () => Promise<T>,
+  makeCall: (signal: AbortSignal) => Promise<T>,
   handleSuccess: (value: T) => void,
   handleError: (callStatus: /*i8*/ number, errorBuffer: ArrayBuffer) => void,
   isErrorType: (error: any) => boolean,
   lowerError: (error: E) => ArrayBuffer,
   lowerString: (str: string) => ArrayBuffer,
 ): UniffiForeignFuture {
-  const promise = makeCall().then(handleSuccess, (error: any) => {
-    let message = error.message ? error.message : error.toString();
-    if (isErrorType(error)) {
-      try {
-        handleError(CALL_ERROR, lowerError(error as E));
-        return;
-      } catch (e: any) {
-        // Fall through to unexpected error handling below.
-        message = `Error handling error "${e}", originally: "${message}"`;
+  const settledHolder: { settled: boolean } = { settled: false };
+  const abortController = new AbortController();
+  const promise = makeCall(abortController.signal)
+    // Before doing anything else, we record that the promise has been settled.
+    // Doing this after the `then` call means we only do this once all of that has finished,
+    // which is way too late.
+    .finally(() => (settledHolder.settled = true))
+    .then(handleSuccess, (error: any) => {
+      let message = error.message ? error.message : error.toString();
+      if (isErrorType(error)) {
+        try {
+          handleError(CALL_ERROR, lowerError(error as E));
+          return;
+        } catch (e: any) {
+          // Fall through to unexpected error handling below.
+          message = `Error handling error "${e}", originally: "${message}"`;
+        }
       }
-    }
-    // This is the catch all:
-    // 1. if there was an unexpected error causing a rejection
-    // 2. if there was an unexpected error in the handleError function.
-    handleError(CALL_UNEXPECTED_ERROR, lowerString(message));
-  });
-  const handle = UNIFFI_FOREIGN_FUTURE_HANDLE_MAP.insert(promise);
+      // This is the catch all:
+      // 1. if there was an unexpected error causing a rejection
+      // 2. if there was an unexpected error in the handleError function.
+      handleError(CALL_UNEXPECTED_ERROR, lowerString(message));
+    });
+
+  const promiseHelper = { abortController, settledHolder, promise };
+  const handle = UNIFFI_FOREIGN_FUTURE_HANDLE_MAP.insert(promiseHelper);
   return /* UniffiForeignFuture */ {
     handle,
     free: uniffiForeignFutureFree,
   };
 }
 
 function uniffiForeignFutureFree(handle: UniffiHandle) {
-  const promise = UNIFFI_FOREIGN_FUTURE_HANDLE_MAP.remove(handle);
+  const helper = UNIFFI_FOREIGN_FUTURE_HANDLE_MAP.remove(handle);
   // #JS_TASK_CANCELLATION
   //
   // This would be where the request from Rust to cancel a JS task would come out.
   // Check if the promise has been settled, and if not, cancel it.
-  //
-  // In the future, we might use an AbortController here, but hermes doesn't implement it yet.
-  //
-  // We would need to:
-  //   - figure out a way of passing the abortController signal to the JS callback.
-  //   - storing a holder with the promise and the abortController in this map.
-  //   - call the abortController in this method.
-  //
-  // abortController.abort();
+  if (!helper.settledHolder.settled) {
+    helper.abortController.abort();
+  }
 }
 
 // For testing