Skip to content

Generalize __wasm_init_task to a "task hook" - #2603

Merged
alexcrichton merged 1 commit into
bytecodealliance:mainfrom
alexcrichton:new-task-intrinsics
Aug 14, 2026
Merged

alexcrichton merged 1 commit into
bytecodealliance:mainfrom
alexcrichton:new-task-intrinsics

Conversation

@alexcrichton

Copy link
Copy Markdown
Member

This commit refactors wit-component's handling of __wasm_init_task and __wasm_init_async_task into a "task hook" and adds it to more locations. The purpose of this refactor is to enable correct stack management in wasi-libc in the wasip3 ABI where stacks may need to be dynamically allocated in some situations. For this use case the previous init hooks were not sufficient because they didn't run in all contexts (such as post-return) nor did they hook completion of a task to deallocate a stack.

The new __wasm_task_hook function is still invoked in the same place as the previous initialization hooks, but it's additionally invoked at the end of a task and in more places. This intrinsic takes a single i32 parameter which is where the hook is being invoked. The locations are:

  • 0 - at the start of a synchronously lifted function.
  • 1 - at the end of a synchronously lifted function (before post-return if that's specified).
  • 2 - at the start of an async-lifted function, either stackful or not.
  • 3 - at the start of an async-lifted's callback option.
  • 4 - when an async-lifted entrypoint returns, or a callback option returns, and the task is blocking or yielding.
  • 5 - when an async-lifted function completes.
  • 6 - when _initialize is originally invoked.
  • 7 - when _initialize completes.
  • 8 - when a resource destructor is invoked.
  • 9 - when a resource destructor completes.
  • 10 - when a post-return hook is invoked.
  • 11 - when a post-return hook completes.

The goal of this change is to enable wasi-libc to correctly manage stacks in a WASIp3 world. This isn't necessary in a WASIp2 world with a global stack pointer because everything "falls out" of the stack pointer being in the main module, threads not existing, and the component model's reentrancy/blocking rules enable safe usage of the global. For WASIp3, however, a global stack pointer cannot be used due to threads, and so context.{get,set} slots are used instead, and these are fresh for all tasks and need initialization. The above hooks enable wasi-libc to primarily use the "main stack" where possible, but there are scenarios where that isn't possible such as for cabi_realloc and for sync-reentrant tasks.

The majority of this implementation is in the fixup module and this has additionally been tested against a local copy of wasi-libc. I've written previously-failing tests in wasi-libc which expose where stack management today is insufficient, and with this change, and an updated wasi-libc with this new intrinsic, all the tests now pass.

One final change I plan to do after this is to additionally add hooks for the cabi_realloc function. Right now that's required to do its own stack management, but it would be more consistent if it were folded into the same infrastructure for hooks. That'll be a bit of a broader change, though, so it's deferred to a future commit.

This commit refactors `wit-component`'s handling of `__wasm_init_task`
and `__wasm_init_async_task` into a "task hook" and adds it to more
locations. The purpose of this refactor is to enable correct stack
management in wasi-libc in the wasip3 ABI where stacks may need to be
dynamically allocated in some situations. For this use case the previous
init hooks were not sufficient because they didn't run in all contexts
(such as post-return) nor did they hook completion of a task to
deallocate a stack.

The new `__wasm_task_hook` function is still invoked in the same place
as the previous initialization hooks, but it's additionally invoked at
the end of a task and in more places. This intrinsic takes a single
`i32` parameter which is where the hook is being invoked. The locations
are:

* 0 - at the start of a synchronously lifted function.
* 1 - at the end of a synchronously lifted function (before post-return
  if that's specified).
* 2 - at the start of an async-lifted function, either stackful or not.
* 3 - at the start of an async-lifted's `callback` option.
* 4 - when an async-lifted entrypoint returns, or a `callback` option
  returns, and the task is blocking or yielding.
* 5 - when an async-lifted function completes.
* 6 - when `_initialize` is originally invoked.
* 7 - when `_initialize` completes.
* 8 - when a resource destructor is invoked.
* 9 - when a resource destructor completes.
* 10 - when a post-return hook is invoked.
* 11 - when a post-return hook completes.

The goal of this change is to enable wasi-libc to correctly manage
stacks in a WASIp3 world. This isn't necessary in a WASIp2 world with a
`global` stack pointer because everything "falls out" of the stack
pointer being in the main module, threads not existing, and the
component model's reentrancy/blocking rules enable safe usage of the
global. For WASIp3, however, a `global` stack pointer cannot be used due
to threads, and so `context.{get,set}` slots are used instead, and these
are fresh for all tasks and need initialization. The above hooks enable
wasi-libc to primarily use the "main stack" where possible, but there
are scenarios where that isn't possible such as for `cabi_realloc` and
for sync-reentrant tasks.

The majority of this implementation is in the `fixup` module and this
has additionally been tested against a local copy of wasi-libc. I've
written previously-failing tests in wasi-libc which expose where stack
management today is insufficient, and with this change, and an updated
wasi-libc with this new intrinsic, all the tests now pass.

One final change I plan to do after this is to additionally add hooks
for the `cabi_realloc` function. Right now that's required to do its own
stack management, but it would be more consistent if it were folded into
the same infrastructure for hooks. That'll be a bit of a broader change,
though, so it's deferred to a future commit.
@alexcrichton
alexcrichton requested a review from dicej August 14, 2026 15:55
@alexcrichton
alexcrichton requested a review from a team as a code owner August 14, 2026 15:55
@alexcrichton
alexcrichton added this pull request to the merge queue Aug 14, 2026
Merged via the queue into bytecodealliance:main with commit 5201ffe Aug 14, 2026
37 checks passed
@alexcrichton
alexcrichton deleted the new-task-intrinsics branch August 14, 2026 17:35
alexcrichton added a commit to alexcrichton/wasm-tools that referenced this pull request Aug 17, 2026
* Use `__wasm_task_hook` for `realloc` options

This commit updates the `__wasm_task_hook` intrinsic, added in bytecodealliance#2603, to
additionally get called for the `realloc` function, typically exported
as `cabi_realloc` today. The goal of this commit is to uniformly use
this hook for wasip3 task initialization/configuration as opposed to the
current state of affairs in wasi-libc where `cabi_realloc` is special.

The implementation here is more involved than bytecodealliance#2603 `realloc` options
are required to create the fixup module and thus come from the original
instance raw. When using task hooks, however, the desired state is that
the lowerings created use the hooked version of `realloc`. As with all
problems in computer science, this is solved with another layer of
indirection.

Specifically the shim module may not contain shims for `realloc` options
if used for lowerings. These are then hooked up through the fixup module
so it's instantiated with both lowered functions and the "real"
implementation so everything can get wired up at the same time.

While implementing this I went ahead an implemented a minor optimization
where for modules that import linear memory they no longer need an
shim-per-import-using-realloc and instead just have a single shim for
realloc. This means that when importing memory generated components
generally have even fewer shims than before.

This commit then attempts to add a variety of tests for various shapes
of this new hook and various other options. This was developed in tandem
with changes to wasi-libc and additionally verified against those.

* Update crates/wit-component/src/encoding/world.rs

Co-authored-by: Joel Dice <joel.dice@akamai.com>

---------

Co-authored-by: Joel Dice <joel.dice@akamai.com>
alexcrichton added a commit to WebAssembly/wasi-libc that referenced this pull request Aug 19, 2026
This commit is a refactor of how wasip3 works specifically with respect
to management of the stack and intrinsics used. The end result of this
work is that a number of bugs, present in wasi-libc today, are all
fixed:

* Sync-reentrant calls, allowed in the component model, previously used
the same stack and stomped over each other.
* Resource destructors, executed with fresh task state, did not
correctly setup a stack to execute on.

The fix for the first issue here is particularly gnarly because it means
that stacks are required to be dynamically allocated in some situations.
This in turn means that those same stacks need to be deallocated, and
prior to this commit there was no great place to doing that. The fix for
the second issue is a bit broader than wasi-libc as well as it couldn't
be solved exclusively here.

To resolve these problems this commit updates the `wasm-tools` version
required, notably including recent PRs to adjust how linking works to
transform `__wasm_init_task` and `__wasm_init_async_task` into a
more-generic `__wasm_task_hook` function. This hook function is invoked
in more locations, such as resource destructors, and additionally
subsumes previous custom logic with `cabi_realloc`. Effectively
`__wasm_task_hook` is now the sole location concerned with configuring a
stack and TLS for the task-to-run. This handles all situations such as
sync/async tasks, `_initialize`, `cabi_realloc`, resource destructors,
and post-return. This was implemented in
bytecodealliance/wasm-tools#2603 and bytecodealliance/wasm-tools#2605.

To get existing tests all passing after updating `wasm-tools` many of
the changes here were needed no matter what. To fully exercise the new
functionality, however, I've also started adding a composition-based
test suite which creates two components, composes them together, and
runs that as a test. This creates the ability to test situations such as
resource destructors and reentrant behavior purely within this
repository. My hope is that as we discover bugs/issues elsewhere it'll
be possible to add a small reduced test case to this repo as a
composition. For now the composition tests are pretty simple and only
cover the cases that failed before this PR and now fail after this PR.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants