Merge pull request #302 from ruby/katei/mitigate-stackoverflow-part2

Mitigate stack-overflow issue take 2
ruby · Nov 13, 2023 · 99f5000 · 99f5000
2 parents c5052f2 + d85b9ea
commit 99f5000
Show file tree

Hide file tree

Showing 6 changed files with 193 additions and 24 deletions.
diff --git a/ext/witapi/witapi-core.c b/ext/witapi/witapi-core.c
@@ -60,8 +60,20 @@ rb_eval_string_value_protect(VALUE str, int *pstate) {
 
 __attribute__((import_module("asyncify"), import_name("start_unwind"))) void
 asyncify_start_unwind(void *buf);
+#define asyncify_start_unwind(buf)                                             \
+  do {                                                                         \
+    extern void *rb_asyncify_unwind_buf;                                       \
+    rb_asyncify_unwind_buf = (buf);                                            \
+    asyncify_start_unwind((buf));                                              \
+  } while (0)
 __attribute__((import_module("asyncify"), import_name("stop_unwind"))) void
 asyncify_stop_unwind(void);
+#define asyncify_stop_unwind()                                                 \
+  do {                                                                         \
+    extern void *rb_asyncify_unwind_buf;                                       \
+    rb_asyncify_unwind_buf = NULL;                                             \
+    asyncify_stop_unwind();                                                    \
+  } while (0)
 __attribute__((import_module("asyncify"), import_name("start_rewind"))) void
 asyncify_start_rewind(void *buf);
 __attribute__((import_module("asyncify"), import_name("stop_rewind"))) void
@@ -114,6 +126,8 @@ rb_wasm_throw_prohibit_rewind_exception(const char *c_msg, size_t msg_len);
       void *asyncify_buf = NULL;                                               \
       extern void *rb_asyncify_unwind_buf;                                     \
       void *asyncify_unwound_buf = rb_asyncify_unwind_buf;                     \
+      if (asyncify_unwound_buf == NULL)                                        \
+        break;                                                                 \
       asyncify_stop_unwind();                                                  \
                                                                                \
       if ((asyncify_buf = rb_wasm_handle_jmp_unwind()) != NULL) {              \

diff --git a/packages/npm-packages/ruby-3_2-wasm-wasi/src/browser.script.js b/packages/npm-packages/ruby-3_2-wasm-wasi/src/browser.script.js
@@ -1,4 +1,12 @@
 import { main } from "@ruby/wasm-wasi/dist/browser.script"
 import * as pkg from "../package.json"
 
-main(pkg)
+main(pkg, {
+  env: {
+    // WORKAROUND(katei): setjmp consumes a LOT of stack in Ruby 3.2,
+    // so we extend default Fiber stack size as well as main stack
+    // size allocated by wasm-ld's --stack-size. 8MB is enough for
+    // most cases. See https://github.com/ruby/ruby.wasm/issues/273
+    "RUBY_FIBER_MACHINE_STACK_SIZE": "8388608"
+  }
+})
diff --git a/packages/npm-packages/ruby-wasm-wasi/src/browser.script.ts b/packages/npm-packages/ruby-wasm-wasi/src/browser.script.ts
@@ -1,11 +1,14 @@
 import { DefaultRubyVM } from "./browser";
 
-export const main = async (pkg: { name: string; version: string }) => {
+export const main = async (
+  pkg: { name: string; version: string },
+  options?: Parameters<typeof DefaultRubyVM>[1],
+) => {
   const response = fetch(
     `https://cdn.jsdelivr.net/npm/${pkg.name}@${pkg.version}/dist/ruby+stdlib.wasm`,
   );
   const module = await compileWebAssemblyModule(response);
-  const { vm } = await DefaultRubyVM(module);
+  const { vm } = await DefaultRubyVM(module, options);
 
   vm.printVersion();
 

diff --git a/packages/npm-packages/ruby-wasm-wasi/src/browser.ts b/packages/npm-packages/ruby-wasm-wasi/src/browser.ts
@@ -62,28 +62,23 @@ const consolePrinter = () => {
 
 export const DefaultRubyVM = async (
   rubyModule: WebAssembly.Module,
-  options: { consolePrint: boolean } = { consolePrint: true },
+  options: {
+    consolePrint?: boolean;
+    env?: Record<string, string> | undefined;
+  } = {},
 ): Promise<{
   vm: RubyVM;
   wasi: WASI;
   instance: WebAssembly.Instance;
 }> => {
   await init();
 
-  const wasi = new WASI({
-    env: {
-      // FIXME(katei): setjmp consumes a LOT of stack now, so we extend
-      // default Fiber stack size as well as main stack size allocated
-      // by wasm-ld's --stack-size. The ideal solution is to reduce
-      // stack consumption in setjmp.
-      RUBY_FIBER_MACHINE_STACK_SIZE: "16777216",
-    },
-  });
+  const wasi = new WASI({ env: options.env });
   const vm = new RubyVM();
 
   const imports = wasi.getImports(rubyModule) as WebAssembly.Imports;
   vm.addToImports(imports);
-  const printer = options.consolePrint ? consolePrinter() : undefined;
+  const printer = (options.consolePrint ?? true) ? consolePrinter() : undefined;
   printer?.addToImports(imports);
 
   const instance = await WebAssembly.instantiate(rubyModule, imports);

diff --git a/packages/npm-packages/ruby-wasm-wasi/src/node.ts b/packages/npm-packages/ruby-wasm-wasi/src/node.ts
@@ -1,16 +1,11 @@
 import { WASI } from "wasi";
 import { RubyVM } from "./index";
 
-export const DefaultRubyVM = async (rubyModule: WebAssembly.Module) => {
-  const wasi = new WASI({
-    env: {
-      // FIXME(katei): setjmp consumes a LOT of stack now, so we extend
-      // default Fiber stack size as well as main stack size allocated
-      // by wasm-ld's --stack-size. The ideal solution is to reduce
-      // stack consumption in setjmp.
-      RUBY_FIBER_MACHINE_STACK_SIZE: "16777216",
-    },
-  });
+export const DefaultRubyVM = async (
+  rubyModule: WebAssembly.Module,
+  options: { env?: Record<string, string> | undefined } = {},
+) => {
+  const wasi = new WASI({ env: options.env });
   const vm = new RubyVM();
   const imports = {
     wasi_snapshot_preview1: wasi.wasiImport,

diff --git a/patches/0001-wasm-allocate-Asyncify-setjmp-buffer-in-heap.patch b/patches/0001-wasm-allocate-Asyncify-setjmp-buffer-in-heap.patch
@@ -0,0 +1,154 @@
+From 253a1835b07a291a84348ae812e61c8fde222962 Mon Sep 17 00:00:00 2001
+From: Yuta Saito <[email protected]>
+Date: Sun, 12 Nov 2023 07:18:01 +0900
+Subject: [PATCH] [wasm] allocate Asyncify setjmp buffer in heap
+
+`rb_jmpbuf_t` type is considerably large due to inline-allocated
+Asyncify buffer, and it leads to stack overflow even with small number
+of C-method call frames. This commit allocates the Asyncify buffer used
+by `rb_wasm_setjmp` in heap to mitigate the issue.
+
+This patch introduces a new type `rb_vm_tag_jmpbuf_t` to abstract the
+representation of a jump buffer, and init/deinit hook points to manage
+lifetime of the buffer. These changes are effectively NFC for non-wasm
+platforms.
+---
+ eval_intern.h |  6 ++++--
+ vm.c          |  2 +-
+ vm_core.h     | 59 +++++++++++++++++++++++++++++++++++++++++++++++++--
+ 3 files changed, 62 insertions(+), 5 deletions(-)
+
+diff --git a/eval_intern.h b/eval_intern.h
+index 778b63e0ea..d008b17ca1 100644
+--- a/eval_intern.h
++++ b/eval_intern.h
+@@ -110,9 +110,11 @@ extern int select_large_fdset(int, fd_set *, fd_set *, fd_set *, struct timeval
+   _tag.tag = Qundef; \
+   _tag.prev = _ec->tag; \
+   _tag.lock_rec = rb_ec_vm_lock_rec(_ec); \
++  rb_vm_tag_jmpbuf_init(&_tag.buf); \
+
+ #define EC_POP_TAG() \
+   _ec->tag = _tag.prev; \
++  rb_vm_tag_jmpbuf_deinit(&_tag.buf); \
+ } while (0)
+
+ #define EC_TMPPOP_TAG() \
+@@ -161,7 +163,7 @@ rb_ec_tag_jump(const rb_execution_context_t *ec, enum ruby_tag_type st)
+ {
+     RUBY_ASSERT(st != TAG_NONE);
+     ec->tag->state = st;
+-    ruby_longjmp(ec->tag->buf, 1);
++    ruby_longjmp(RB_VM_TAG_JMPBUF_GET(ec->tag->buf), 1);
+ }
+
+ /*
+@@ -169,7 +171,7 @@ rb_ec_tag_jump(const rb_execution_context_t *ec, enum ruby_tag_type st)
+   [ISO/IEC 9899:1999] 7.13.1.1
+ */
+ #define EC_EXEC_TAG() \
+-    (UNLIKELY(ruby_setjmp(_tag.buf)) ? rb_ec_tag_state(VAR_FROM_MEMORY(_ec)) : (EC_REPUSH_TAG(), 0))
++    (UNLIKELY(ruby_setjmp(RB_VM_TAG_JMPBUF_GET(_tag.buf))) ? rb_ec_tag_state(VAR_FROM_MEMORY(_ec)) : (EC_REPUSH_TAG(), 0))
+
+ #define EC_JUMP_TAG(ec, st) rb_ec_tag_jump(ec, st)
+
+diff --git a/vm.c b/vm.c
+index 7f43484905..789e0956be 100644
+--- a/vm.c
++++ b/vm.c
+@@ -2462,7 +2462,7 @@ vm_exec(rb_execution_context_t *ec)
+
+     rb_wasm_try_catch_init(&try_catch, vm_exec_bottom_main, vm_exec_bottom_rescue, &ctx);
+
+-    rb_wasm_try_catch_loop_run(&try_catch, &_tag.buf);
++    rb_wasm_try_catch_loop_run(&try_catch, &RB_VM_TAG_JMPBUF_GET(_tag.buf));
+
+     result = ctx.result;
+ #else
+diff --git a/vm_core.h b/vm_core.h
+index acad6280be..45290c21f7 100644
+--- a/vm_core.h
++++ b/vm_core.h
+@@ -884,6 +884,61 @@ typedef RUBY_JMP_BUF rb_jmpbuf_t;
+ typedef void *rb_jmpbuf_t[5];
+ #endif
+
++/*
++  `rb_vm_tag_jmpbuf_t` type represents a buffer used to
++  long jump to a C frame associated with `rb_vm_tag`.
++
++  Use-site of `rb_vm_tag_jmpbuf_t` is responsible for calling the
++  following functions:
++  - `rb_vm_tag_jmpbuf_init` once `rb_vm_tag_jmpbuf_t` is allocated.
++  - `rb_vm_tag_jmpbuf_deinit` once `rb_vm_tag_jmpbuf_t` is no longer necessary.
++
++  `RB_VM_TAG_JMPBUF_GET` transforms a `rb_vm_tag_jmpbuf_t` into a
++  `rb_jmpbuf_t` to be passed to `rb_setjmp/rb_longjmp`.
++*/
++#if defined(__wasm__) && !defined(__EMSCRIPTEN__)
++/*
++  WebAssembly target with Asyncify-based SJLJ needs
++  to capture the execution context by unwind/rewind-ing
++  call frames into a jump buffer. The buffer space tends
++  to be considerably large unlike other architectures'
++  register-based buffers.
++  Therefore, we allocates the buffer on the heap on such
++  environments.
++*/
++typedef rb_jmpbuf_t *rb_vm_tag_jmpbuf_t;
++
++#define RB_VM_TAG_JMPBUF_GET(buf) (*buf)
++
++inline void
++rb_vm_tag_jmpbuf_init(rb_vm_tag_jmpbuf_t *jmpbuf)
++{
++  *jmpbuf = malloc(sizeof(rb_jmpbuf_t));
++}
++
++inline void
++rb_vm_tag_jmpbuf_deinit(const rb_vm_tag_jmpbuf_t *jmpbuf)
++{
++  free(*jmpbuf);
++}
++#else
++typedef rb_jmpbuf_t rb_vm_tag_jmpbuf_t;
++
++#define RB_VM_TAG_JMPBUF_GET(buf) (buf)
++
++inline void
++rb_vm_tag_jmpbuf_init(rb_vm_tag_jmpbuf_t *jmpbuf)
++{
++  // no-op
++}
++
++inline void
++rb_vm_tag_jmpbuf_deinit(const rb_vm_tag_jmpbuf_t *jmpbuf)
++{
++  // no-op
++}
++#endif
++
+ /*
+   the members which are written in EC_PUSH_TAG() should be placed at
+   the beginning and the end, so that entire region is accessible.
+@@ -891,7 +946,7 @@ typedef void *rb_jmpbuf_t[5];
+ struct rb_vm_tag {
+     VALUE tag;
+     VALUE retval;
+-    rb_jmpbuf_t buf;
++    rb_vm_tag_jmpbuf_t buf;
+     struct rb_vm_tag *prev;
+     enum ruby_tag_type state;
+     unsigned int lock_rec;
+@@ -899,7 +954,7 @@ struct rb_vm_tag {
+
+ STATIC_ASSERT(rb_vm_tag_buf_offset, offsetof(struct rb_vm_tag, buf) > 0);
+ STATIC_ASSERT(rb_vm_tag_buf_end,
+-              offsetof(struct rb_vm_tag, buf) + sizeof(rb_jmpbuf_t) <
++              offsetof(struct rb_vm_tag, buf) + sizeof(rb_vm_tag_jmpbuf_t) <
+               sizeof(struct rb_vm_tag));
+
+ struct rb_unblock_callback {
+-- 
+2.39.3 (Apple Git-145)
+