[Web] Retain packed functions across Asyncify calls

Asyncify can suspend a prepared call after the JavaScript wrapper has been disposed through another alias. Take an extra reference on the PackedFunc when the call first unwinds and release it with the frame, so the rewound call still has a live callee. Synchronous calls do not take the reference.

Release the retained handle with the frame, including when cleanup fails.

Author: akaashrp

web/src/asyncify.ts

@@ -146,7 +146,10 @@ export class AsyncifyHandler {
    * @param func The input function
    * @returns The wrapped async function
    */
-  wrapExport(func: (...args: Array<any>) => any): (...args: Array<any>) => Promise<any> {
+  wrapExport(
+    func: (...args: Array<any>) => any,
+    onSuspend?: () => void
+  ): (...args: Array<any>) => Promise<any> {
     return async (...args: Array<any>) => {
       assert(this.getState() == AsyncifyStateKind.None);
 
@@ -159,6 +162,11 @@ export class AsyncifyHandler {
       // the if (isPromise(value)) condition in wrapImport
       while (this.getState() == AsyncifyStateKind.Unwinding) {
         this.stopUnwind();
+        // The stack is unwound and no other JavaScript has run yet, so the
+        // caller can still take references before the await below.
+        if (onSuspend !== undefined) {
+          onSuspend();
+        }
         // try to resolve the promise that the internal requested
         // we then store it into the temp value in storedValueBeforeRewind
         // which then get passed onto the function(see wrapImport)

web/src/ctypes.ts

@@ -176,6 +176,11 @@ export type FTVMFFIWasmFunctionCreate = (
  */
 export type FTVMFFIWasmFunctionDeleter = (self: Pointer) => void;
 
+/**
+ * int TVMFFIObjectIncRef(TVMFFIObjectHandle obj);
+ */
+export type FTVMFFIObjectIncRef = (obj: Pointer) => number;
+
 /**
  * int TVMFFIObjectDecRef(TVMFFIObjectHandle obj);
  */

web/src/runtime.ts

@@ -509,6 +509,8 @@ class PackedFuncCell extends TVMObject {
 
 interface PackedCallFrame {
   cell: PackedFuncCell;
+  /** Extra reference taken once the call suspends in Asyncify. */
+  retainedCell?: PackedFuncCell;
   stack: CachedCallStack;
   argsOffset: PtrOffset;
   numArgs: number;
@@ -1975,17 +1977,21 @@ export class Instance implements Disposable {
         // return storage.  In particular, the Wasm stack may retain pointers
         // into transient byte arguments while its execution is suspended.
         const run = this.asyncifyHandler.wrapExport(
-          () => this.invokePackedCall(frame as PackedCallFrame)
+          () => this.invokePackedCall(frame as PackedCallFrame),
+          () => this.retainPackedCall(frame as PackedCallFrame)
         );
         return await run();
       } finally {
-        if (frame !== undefined) {
-          this.releasePackedCall(frame);
-        }
-        callInProgress = false;
-        this.asyncifyCallInProgress = false;
-        if (disposeRequested) {
-          func.dispose();
+        try {
+          if (frame !== undefined) {
+            this.releasePackedCall(frame);
+          }
+        } finally {
+          callInProgress = false;
+          this.asyncifyCallInProgress = false;
+          if (disposeRequested) {
+            func.dispose();
+          }
         }
       }
     }) as AsyncPackedFunc;
@@ -2498,13 +2504,31 @@ export class Instance implements Disposable {
     }
   }
 
+  /**
+   * Take an extra reference on the function once a call has suspended.
+   *
+   * While the Wasm stack is unwound, other JavaScript can dispose the
+   * function through another wrapper. The frame keeps the callee alive
+   * until the rewound call finishes. Synchronous calls do not take it.
+   */
+  private retainPackedCall(frame: PackedCallFrame): void {
+    if (frame.released || frame.retainedCell !== undefined) {
+      return;
+    }
+    const handle = frame.cell.getHandle();
+    this.lib.checkCall(
+      (this.exports.TVMFFIObjectIncRef as ctypes.FTVMFFIObjectIncRef)(handle)
+    );
+    frame.retainedCell = new PackedFuncCell(handle, this.lib, this.ctx);
+  }
+
   private invokePackedCall(frame: PackedCallFrame): any {
     if (frame.released) {
       throw new Error("Cannot invoke a released packed-call frame");
     }
     this.lib.checkCall(
       (this.exports.TVMFFIFunctionCall as ctypes.FTVMFFIFunctionCall)(
-        frame.cell.getHandle(),
+        (frame.retainedCell ?? frame.cell).getHandle(),
         frame.stack.ptrFromOffset(frame.argsOffset),
         frame.numArgs,
         frame.stack.ptrFromOffset(frame.retOffset)
@@ -2519,7 +2543,11 @@ export class Instance implements Disposable {
   private releasePackedCall(frame: PackedCallFrame): void {
     if (!frame.released) {
       frame.released = true;
-      this.lib.recycleCallStack(frame.stack);
+      try {
+        this.lib.recycleCallStack(frame.stack);
+      } finally {
+        frame.retainedCell?.dispose();
+      }
     }
   }
 

web/tests/node/test_ffi_bytes.js

@@ -123,16 +123,23 @@ test("PackedFunc captures Wasm-backed byte sources automatically", () => {
   }
   try {
     const input = new Uint8Array(tvm.memory.memory.buffer, ptr, 32);
-    input.set(makeBytes(input.length));
-    expect(tvm.captureWasmByteSources([input])).toEqual([{
-      buffer: input.buffer,
-      byteOffset: input.byteOffset,
-      byteLength: input.byteLength,
-    }]);
-    tvm.withNewScope(() => {
-      const echo = tvm.getGlobalFunc("testing.echo");
-      expectExactBytes(echo(input), makeBytes(input.length));
-    });
+    const expected = makeBytes(input.length);
+    input.set(expected);
+    const capture = tvm.captureWasmByteSources;
+    tvm.captureWasmByteSources = function(args) {
+      const sources = capture.call(this, args);
+      this.memory.memory.grow(1);
+      return sources;
+    };
+    try {
+      tvm.withNewScope(() => {
+        const echo = tvm.getGlobalFunc("testing.echo");
+        expectExactBytes(echo(input), expected);
+      });
+      expect(input.byteLength).toBe(0);
+    } finally {
+      tvm.captureWasmByteSources = capture;
+    }
   } finally {
     tvm.exports.TVMWasmFreeSpace(ptr);
   }
@@ -214,6 +221,60 @@ test("Asyncify rejects overlapping calls from different wrappers", async () => {
   }
 });
 
+test("Asyncify call retains the function across alias disposal", async () => {
+  tvm.beginScope();
+  let release;
+  try {
+    const input = makeBytes(64 * 1024 + 17);
+    const gate = new Promise((resolve) => {
+      release = resolve;
+    });
+    tvm.registerAsyncifyFunc("tvmjs.testing.wait_for_alias_disposal", async () => {
+      await gate;
+    }, true);
+
+    const callback = tvm.getGlobalFunc("tvmjs.testing.wait_for_alias_disposal");
+    const original = tvm.getGlobalFunc(
+      "tvmjs.testing.check_byte_array_across_callbacks"
+    );
+    const active = tvm.wrapAsyncifyPackedFunc(original);
+    const alias = tvm.wrapAsyncifyPackedFunc(original);
+
+    const pending = active(input, callback);
+    alias.dispose();
+    expect(original._tvmPackedCell.getHandle(false)).toBe(0);
+    release();
+    expect(await pending).toBe(0);
+  } finally {
+    if (release !== undefined) {
+      release();
+    }
+    tvm.endScope();
+  }
+});
+
+test("Asyncify clears call state when frame cleanup throws", async () => {
+  tvm.beginScope();
+  const releasePackedCall = tvm.releasePackedCall;
+  try {
+    const original = tvm.getGlobalFunc("testing.echo");
+    const wrapped = tvm.wrapAsyncifyPackedFunc(original);
+    tvm.releasePackedCall = function(frame) {
+      releasePackedCall.call(this, frame);
+      throw new Error("expected frame cleanup failure");
+    };
+
+    const pending = wrapped(1);
+    wrapped.dispose();
+    await expect(pending).rejects.toThrow("expected frame cleanup failure");
+    expect(original._tvmPackedCell.getHandle(false)).toBe(0);
+    expect(tvm.asyncifyCallInProgress).toBe(false);
+  } finally {
+    tvm.releasePackedCall = releasePackedCall;
+    tvm.endScope();
+  }
+});
+
 test("Asyncify does not consume an owned byte result before completion", async () => {
   if (!tvm.asyncifyEnabled()) {
     throw new Error("The test runtime must be built with Asyncify enabled");