Stop GC while having a sandwitched JS frame

This makes nested eval GC safe

Author: kateinoigakukun

ext/witapi/bindgen/rb-abi-guest.c

@@ -205,3 +205,13 @@ __attribute__((export_name("rb-vm-bugreport: func() -> ()")))
 void __wasm_export_rb_abi_guest_rb_vm_bugreport(void) {
   rb_abi_guest_rb_vm_bugreport();
 }
+__attribute__((export_name("rb-gc-enable: func() -> bool")))
+int32_t __wasm_export_rb_abi_guest_rb_gc_enable(void) {
+  bool ret = rb_abi_guest_rb_gc_enable();
+  return ret;
+}
+__attribute__((export_name("rb-gc-disable: func() -> bool")))
+int32_t __wasm_export_rb_abi_guest_rb_gc_disable(void) {
+  bool ret = rb_abi_guest_rb_gc_disable();
+  return ret;
+}

ext/witapi/bindgen/rb-abi-guest.h

@@ -68,6 +68,8 @@ extern "C"
   void rb_abi_guest_rb_clear_errinfo(void);
   void rb_abi_guest_rstring_ptr(rb_abi_guest_rb_abi_value_t value, rb_abi_guest_string_t *ret0);
   void rb_abi_guest_rb_vm_bugreport(void);
+  bool rb_abi_guest_rb_gc_enable(void);
+  bool rb_abi_guest_rb_gc_disable(void);
   #ifdef __cplusplus
 }
 #endif

ext/witapi/bindgen/rb-abi-guest.wit

@@ -18,3 +18,6 @@ rb-clear-errinfo: func()
 rstring-ptr: func(value: rb-abi-value) -> string
 
 rb-vm-bugreport: func()
+
+rb-gc-enable: func() -> bool
+rb-gc-disable: func() -> bool

ext/witapi/witapi-core.c

@@ -303,4 +303,11 @@ void rb_vm_bugreport(const void *);
 
 void rb_abi_guest_rb_vm_bugreport(void) { rb_vm_bugreport(NULL); }
 
+bool rb_abi_guest_rb_gc_enable(void) {
+  return rb_gc_enable() == Qtrue;
+}
+bool rb_abi_guest_rb_gc_disable(void) {
+  return rb_gc_disable() == Qtrue;
+}
+
 void Init_witapi(void) {}

packages/npm-packages/ruby-wasm-wasi/src/bindgen/rb-abi-guest.d.ts

@@ -80,6 +80,8 @@ export class RbAbiGuest {
   rbClearErrinfo(): void;
   rstringPtr(value: RbAbiValue): string;
   rbVmBugreport(): void;
+  rbGcEnable(): boolean;
+  rbGcDisable(): boolean;
 }
 
 export class RbIseq {

packages/npm-packages/ruby-wasm-wasi/src/bindgen/rb-abi-guest.js

@@ -1,4 +1,4 @@
-import { data_view, to_uint32, UTF8_DECODER, utf8_encode, UTF8_ENCODED_LEN, Slab } from './intrinsics.js';
+import { data_view, to_uint32, UTF8_DECODER, utf8_encode, UTF8_ENCODED_LEN, Slab, throw_invalid_bool } from './intrinsics.js';
 export class RbAbiGuest {
   constructor() {
     this._resource0_slab = new Slab();
@@ -161,6 +161,16 @@ export class RbAbiGuest {
   rbVmBugreport() {
     this._exports['rb-vm-bugreport: func() -> ()']();
   }
+  rbGcEnable() {
+    const ret = this._exports['rb-gc-enable: func() -> bool']();
+    const bool0 = ret;
+    return bool0 == 0 ? false : (bool0 == 1 ? true : throw_invalid_bool());
+  }
+  rbGcDisable() {
+    const ret = this._exports['rb-gc-disable: func() -> bool']();
+    const bool0 = ret;
+    return bool0 == 0 ? false : (bool0 == 1 ? true : throw_invalid_bool());
+  }
 }
 
 export class RbIseq {

packages/npm-packages/ruby-wasm-wasi/src/index.ts

@@ -1,5 +1,6 @@
 import * as RbAbi from "./bindgen/rb-abi-guest";
 import {
+  RbJsAbiHost,
   addRbJsAbiHostToImports,
   JsAbiResult,
   JsAbiValue,
@@ -29,9 +30,41 @@ export class RubyVM {
   private instance: WebAssembly.Instance | null = null;
   private transport: JsValueTransport;
   private exceptionFormatter: RbExceptionFormatter;
+  private interfaceState: RbAbiInterfaceState = {
+    hasJSFrameAfterRbFrame: false,
+  }
 
   constructor() {
-    this.guest = new RbAbi.RbAbiGuest();
+    // Wrap exported functions from Ruby VM to stop GC during the
+    // function call if the call stack has sandwitched JS frames
+    // like JS -> Ruby -> JS -> Ruby.
+    const proxyExports = (exports: RbAbi.RbAbiGuest) => {
+      const excludedMethods: (keyof RbAbi.RbAbiGuest)[] = ["addToImports", "instantiate", "rbGcEnable", "rbGcDisable"];
+      const excluded = ["constructor"].concat(excludedMethods);
+      // wrap all methods in RbAbi.RbAbiGuest class
+      for (const key of Object.getOwnPropertyNames(RbAbi.RbAbiGuest.prototype)) {
+        if (excluded.includes(key)) {
+          continue;
+        }
+        const value = exports[key];
+        if (typeof value === "function") {
+          exports[key] = (...args: any[]) => {
+            const shouldStopGC = this.interfaceState.hasJSFrameAfterRbFrame;
+            let isAlreadyDisabled = false;
+            if (shouldStopGC) {
+              isAlreadyDisabled = exports.rbGcDisable();
+            }
+            const result = Reflect.apply(value, exports, args)
+            if (shouldStopGC && !isAlreadyDisabled) {
+              exports.rbGcEnable();
+            }
+            return result;
+          }
+        }
+      }
+      return exports;
+    }
+    this.guest = proxyExports(new RbAbi.RbAbiGuest());
     this.transport = new JsValueTransport();
     this.exceptionFormatter = new RbExceptionFormatter();
   }
@@ -80,9 +113,28 @@ export class RubyVM {
         }
       };
     }
+    // NOTE: The imported functions must disable Ruby GC if they call
+    // Ruby API that may trigger GC. Otherwise, the GC may collect
+    // objects that are still referenced by Wasm locals because Asyncify
+    // cannot scan the Wasm stack above the JS frame.
+    const proxyImports = (imports: RbJsAbiHost) => {
+      for (const [key, value] of Object.entries(imports)) {
+        if (typeof value === "function") {
+          imports[key] = (...args: any[]) => {
+            const oldValue = this.interfaceState.hasJSFrameAfterRbFrame;
+            this.interfaceState.hasJSFrameAfterRbFrame = true;
+            const result = Reflect.apply(value, imports, args);
+            this.interfaceState.hasJSFrameAfterRbFrame = oldValue;
+            return result;
+          }
+        }
+      }
+      return imports;
+    };
+
     addRbJsAbiHostToImports(
       imports,
-      {
+      proxyImports({
         evalJs: wrapTry((code) => {
           return Function(code)();
         }),
@@ -201,7 +253,7 @@ export class RubyVM {
         reflectSetPrototypeOf: function (target, prototype): boolean {
           throw new Error("Function not implemented.");
         },
-      },
+      }),
       (name) => {
         return this.instance.exports[name];
       }
@@ -294,6 +346,15 @@ export class RubyVM {
   }
 }
 
+type RbAbiInterfaceState = {
+  /**
+   * Track if the last JS frame that was created by a Ruby frame
+   * to determine if we need to disable Ruby GC during nested
+   * JS->Ruby->JS->Ruby calls
+  **/
+  hasJSFrameAfterRbFrame: boolean;
+}
+
 /**
  * Export a JS value held by the Ruby VM to the JS environment.
  * This is implemented in a dirty way since wit cannot reference resources

packages/npm-packages/ruby-wasm-wasi/test/gc.test.ts

@@ -82,4 +82,23 @@ describe("GC integration", () => {
     expect(o1.call("hash").toString()).toBe(o2.call("hash").toString());
     expect(o2.call("hash").toString()).toBe(o3.call("hash").toString());
   });
+
+  test("stop GC while having a sandwitched JS frame", async () => {
+    const vm = await initRubyVM();
+    const o = vm.eval(`
+    require "js"
+
+    JS.eval(<<~JS)
+      return {
+        takeVM(vm) {
+          return vm.eval("GC.disable").toJS();
+        }
+      }
+    JS
+    `);
+    const wasDisabled = o.call("takeVM", vm.wrap(vm));
+    expect(wasDisabled.toJS()).toBe(true);
+    const isNotEnabledBack = vm.eval("GC.enable");
+    expect(isNotEnabledBack.toJS()).toBe(false);
+  });
 });