Refactor stack management on wasip3

cmake/bindings.cmake

@@ -152,7 +152,6 @@ function(wit_bindgen_edit p)
     COMMAND sed ${SED_INPLACE_ARGS} "'s_#include .wasi${p}\.h._#include \"wasi/wasi${p}.h\"_'" ${bottom_half}/sources/wasi${p}.c
     COMMAND sed ${SED_INPLACE_ARGS} "s/extern void exit_exit/_Noreturn extern void exit_exit/" ${bottom_half}/headers/public/wasi/__generated_wasi${p}.h
     COMMAND sed ${SED_INPLACE_ARGS} "s/extern void __wasm_import_exit_exit/_Noreturn extern void __wasm_import_exit_exit/" ${bottom_half}/sources/wasi${p}.c
-    COMMAND sed ${SED_INPLACE_ARGS} "s/__attribute__.*\"cabi_realloc\".*/#include \\\"cabi_realloc_augment.h\\\"/" ${bottom_half}/sources/wasi${p}.c
     DEPENDS bindings-${p}
   )
   add_dependencies(bindings bindings-${p}-edit)

libc-bottom-half/CMakeLists.txt

@@ -184,10 +184,7 @@ if (WASI STREQUAL "p3")
 endif()
 
 if (HAVE_WASM_LIBCALL_THREAD_CONTEXT)
-  list(APPEND bottom_half_sources
-    sources/wasip3_tls.c
-    sources/__wasm_init_task.S
-  )
+  list(APPEND bottom_half_sources sources/__wasm_task_hook.S)
 endif()
 
 # Don't export symbols of generated code in shared libraries, so specifically

libc-bottom-half/crt/wasip3_symbol_references.h

@@ -7,12 +7,9 @@
 #include <wasi/version.h>
 
 // Force some symbols to be linked in for wasip3
-extern void __wasm_init_task(void);
-extern void __wasm_init_async_task(void);
+extern void __wasm_task_hook(uint32_t);
 extern void cabi_realloc(void);
-__attribute__((used)) static void *__wasm_init_task_ref = __wasm_init_task;
-__attribute__((used)) static void *__wasm_init_async_task_ref =
-    __wasm_init_async_task;
+__attribute__((used)) static void *__wasm_task_hook_ref = __wasm_task_hook;
 __attribute__((used)) static void *cabi_realloc_ref = cabi_realloc;
 
 // Force `__wasm_{g,s}et_{stack_pointer,tls_base}` to exist as defined symbols.

libc-bottom-half/sources/__cabi_realloc_wrapper.S

@@ -1,129 +0,0 @@
-#include <wasi/version.h>
-
-#ifdef __wasm_libcall_thread_context__
-// The size, in bytes, of the stack statically allocated for `cabi_realloc`.
-// This doesn't need to be too big as it's just running `cabi_realloc` as defined
-// in wasi-libc. The current number is determined by:
-//
-// * Start with 16 bytes.
-// * Double until tests don't crash in a debug build.
-// * Double again.
-//
-// The goal is to be pretty modest since this is part of every module.
-#define REALLOC_STACK_SIZE 512
-
-// Declare the `__wasm_{g,s}et_stack_pointer` functions.
-.functype   __wasm_set_stack_pointer (i32) -> ()
-.functype   __wasm_get_stack_pointer () -> (i32)
-.type       __wasm_set_stack_pointer,@function
-.type       __wasm_get_stack_pointer,@function
-.functype   __cabi_realloc_impl (i32, i32, i32, i32) -> (i32)
-.type       __cabi_realloc_impl,@function
-
-// Go ahead and generate an `-fPIC`-compatible addressing mode for
-// `cabi_realloc_stack` so this is compatible with both PIC and not builds.
-.type        __memory_base,@global
-.globaltype  __memory_base, i32, immutable
-
-// For coop-threads builds we need to configure TLS to the initial TLS
-// value for all tasks/etc, so import those symbols here.
-//
-// Note that this saves and restores context slot 1 itself rather than this
-// module's TLS base. When several modules are linked together the slot holds an
-// array of per-module base pointers, so clobbering just libc's entry would both
-// leave the other modules pointing at the wrong TLS and scribble through
-// whatever stale pointer the slot happened to contain.
-#ifdef __wasi_cooperative_threads__
-.functype   wasip3_context_set_1 (i32) -> ()
-.functype   wasip3_context_get_1 () -> (i32)
-.functype   __wasilibc_set_task_tls (i32) -> ()
-.type       wasip3_context_set_1,@function
-.type       wasip3_context_get_1,@function
-.type       __wasilibc_set_task_tls,@function
-.import_module wasip3_context_get_1, "$root"
-.import_name wasip3_context_get_1, "[context-get-1]"
-.import_module wasip3_context_set_1, "$root"
-.import_name wasip3_context_set_1, "[context-set-1]"
-.type       __init_tls_base,@global
-#ifdef __PIC__
-.globaltype __init_tls_base, i32
-#else
-.globaltype __init_tls_base, i32, immutable
-#endif
-#endif
-
-// This wrapper takes over the `cabi_realloc` symbol itself and the
-// `cabi_realloc` core module export. The C implementation it forwards to is
-// renamed to `__cabi_realloc_impl` by `cabi_realloc_augment.h`.
-//
-// Note that it's important that the symbol name matches the export name here.
-// In a `BUILD_SHARED=ON` build `libc.so` is dynamically linked against by other
-// shared libraries and by `crt1-{command,reactor}.o`, and the name a symbol is
-// resolved by in that context is its wasm export name. If `.export_name` were
-// used to rename the export then this function would be unresolvable by name
-// from anything linking against `libc.so`.
-.globl cabi_realloc
-.export_name cabi_realloc, "cabi_realloc"
-.type cabi_realloc,@function
-cabi_realloc:
-  .functype cabi_realloc (i32, i32, i32, i32) -> (i32)
-
-  // Prior to bytecodealliance/wasmtime#13949 wasmtime's handling of
-  // `cabi_realloc` and context slots was a bit buggy. That means that prior
-  // to that PR the slots aren't guaranteed to be 0 and the value present will
-  // be used by some future task. This export is going to clobber these
-  // slots, so while that Wasmtime change is percolating this preserves the
-  // slots around the invocation of `cabi_realloc`. Once #13949 percolates and
-  // ships then this workaround code, and the restore down below, can be
-  // deleted.
-  .local i32, i32
-  call __wasm_get_stack_pointer
-  local.set 4
-#ifdef __wasi_cooperative_threads__
-  call wasip3_context_get_1
-  local.set 5
-#endif
-
-  // First up, configure the stack pointer. Get the base of the stack from the
-  // data symbol defined below, add the stack size, and that's what we're
-  // starting with.
-  i32.const cabi_realloc_stack@MBREL
-  global.get __memory_base
-  i32.add
-  i32.const  REALLOC_STACK_SIZE
-  i32.add
-  call __wasm_set_stack_pointer
-
-  // Next up configure TLS. This is only required in coop-threads builds.
-#ifdef __wasi_cooperative_threads__
-  global.get   __init_tls_base
-  call __wasilibc_set_task_tls
-#endif
-
-  // And finally forward to the actual implementation of `cabi_realloc`
-  // defined in `wasip3.c`
-  local.get 0
-  local.get 1
-  local.get 2
-  local.get 3
-  call __cabi_realloc_impl
-
-  local.get 4
-  call __wasm_set_stack_pointer
-#ifdef __wasi_cooperative_threads__
-  local.get 5
-  call wasip3_context_set_1
-#endif
-
-  end_function
-
-// Define a data symbol which is the stack that `cabi_realloc` runs on.
-// Note that this is aligned to 16-bytes, the expected stack alignment.
-.type cabi_realloc_stack,@object
-.section .bss.cabi_realloc_stack,"",@
-.p2align 4, 0x0
-cabi_realloc_stack:
-  .skip  REALLOC_STACK_SIZE
-  .size cabi_realloc_stack,  REALLOC_STACK_SIZE
-
-#endif // __wasm_libcall_thread_context__

libc-bottom-half/sources/__wasm_task_hook.S

@@ -0,0 +1,96 @@
+#include <wasi/version.h>
+
+// The size, in bytes, of the stack statically allocated for `__wasilibc_task_hook`.
+// This doesn't need to be too big as it's just running `malloc` as defined
+// in wasi-libc. The current number is determined by:
+//
+// * Start with 16 bytes.
+// * Double until tests don't crash in a debug build.
+// * Double again.
+//
+// The goal is to be pretty modest since this is part of every module.
+#define HOOK_STACK_SIZE 512
+
+.text
+
+.globl       __wasm_set_stack_pointer
+.functype    __wasm_set_stack_pointer (i32) -> ()
+.globl       __wasm_get_stack_pointer
+.functype    __wasm_get_stack_pointer () -> (i32)
+.functype    __wasilibc_task_hook (i32, i32, i32, i32, i32) -> (i32)
+
+.globaltype  __init_stack_pointer, i32, immutable
+.globaltype  __memory_base, i32, immutable
+
+#ifdef __wasi_cooperative_threads__
+#ifdef __PIC__
+.globaltype __init_tls_base, i32
+#else
+.globaltype __init_tls_base, i32, immutable
+#endif
+#endif
+
+// Hook point for component model task lifecycle events which is invoked by
+// `wit-component`'s componentization process.
+//
+// Calls to this function are injected by `wit-component` when a component is
+// created from a core module. This is invoked with single `i32` parameter
+// indicating what's being hooked. Hooks include situations such as a
+// synchronous task starting, an asynchronous task blocking, cabi_realloc
+// returning, etc. The goal of this hook is to manage the `context.{get,set}`
+// slots for this task, notably the stack pointer and TLS base pointer.
+//
+// When this function is invoked the context slots may both be 0, indicating
+// that there is not a stack, hence the reason for this being an assembly
+// function. This function quickly sets up a small temporary stack, allocated
+// below as `hook_stack`, and then delegates to the `__wasilibc_task_hook`
+// function which finishes the task started here. The return value of
+// `__wasilibc_task_hook` becomes the stack pointer after this function returns.
+.globl  __wasm_task_hook
+.export_name  __wasm_task_hook, __wasm_task_hook
+.type  __wasm_task_hook,@function
+__wasm_task_hook:
+  .functype  __wasm_task_hook (i32) -> ()
+  .local i32
+
+  // prep arguments to `__wasilibc_task_hook`, notably:
+  //
+  // * The hook being configured (this function's argument)
+  // * The TLS base for the main thread (to configure TLS)
+  // * The initial stack pointer for the main thread (if that'll get used)
+  // * The current stack pointer (if we're deallocating a stack)
+  // * The hook stack pointer (passed further down below)
+  local.get 0
+#ifdef __wasi_cooperative_threads__
+  global.get __init_tls_base
+#else
+  i32.const 0
+#endif
+  global.get __init_stack_pointer
+  call __wasm_get_stack_pointer
+
+  // setup a temporary stack
+  i32.const hook_stack@MBREL
+  global.get __memory_base
+  i32.add
+  i32.const HOOK_STACK_SIZE
+  i32.add
+  local.tee 1
+  call __wasm_set_stack_pointer
+
+  // figure out what our stack should be after this hook invocation and
+  // configure it.
+  local.get 1
+  call __wasilibc_task_hook
+  call __wasm_set_stack_pointer
+
+  end_function
+
+// Define a data symbol which is the stack that `__wasilibc_task_hook` runs on.
+// Note that this is aligned to 16-bytes, the expected stack alignment.
+.type hook_stack,@object
+.section .bss.hook_stack,"",@
+.p2align 4, 0x0
+hook_stack:
+  .skip  HOOK_STACK_SIZE
+  .size hook_stack,  HOOK_STACK_SIZE