Skip to content

Configure async task context on realloc calls #13949

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
alexcrichton merged 3 commits into from
Jul 24, 2026
Merged

Configure async task context on realloc calls #13949

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
847d857
Configure async task context on `realloc` calls
alexcrichton Jul 23, 2026
33d6d39
Fix conditional builds
alexcrichton Jul 23, 2026
4efdd45
Add a note about restoration
alexcrichton Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions crates/cranelift/src/compiler/component.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -713,7 +713,7 @@ impl<'a> TrampolineCompiler<'a> {
|_, _| {},
);
}
Trampoline::ThreadIndex => {
Trampoline::ThreadIndex { .. } => {
self.translate_libcall(
host::thread_index,
TrapSentinel::NegativeOne,
Expand Down Expand Up @@ -1478,7 +1478,6 @@ impl<'a> TrampolineCompiler<'a> {
let instance = match trampoline {
// These intrinsics explicitly do not check the may-leave flag.
Trampoline::ResourceRep { .. }
| Trampoline::ThreadIndex
| Trampoline::BackpressureInc { .. }
| Trampoline::BackpressureDec { .. } => return,

Expand Down Expand Up @@ -1508,6 +1507,7 @@ impl<'a> TrampolineCompiler<'a> {
| Trampoline::WaitableSetPoll { instance, .. }
| Trampoline::WaitableSetDrop { instance }
| Trampoline::WaitableJoin { instance }
| Trampoline::ThreadIndex { instance }
| Trampoline::ThreadNewIndirect { instance, .. }
| Trampoline::ThreadResumeLater { instance, .. }
| Trampoline::ThreadSuspend { instance, .. }
Expand Down
8 changes: 6 additions & 2 deletions crates/environ/src/component/dfg.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -472,7 +472,9 @@ pub enum Trampoline {
Trap,
EnterSyncCall,
ExitSyncCall,
ThreadIndex,
ThreadIndex {
instance: RuntimeComponentInstanceIndex,
},
ThreadNewIndirect {
instance: RuntimeComponentInstanceIndex,
start_func_ty_idx: ComponentTypeIndex,
Expand Down Expand Up @@ -1156,7 +1158,9 @@ impl LinearizeDfg<'_> {
Trampoline::Trap => info::Trampoline::Trap,
Trampoline::EnterSyncCall => info::Trampoline::EnterSyncCall,
Trampoline::ExitSyncCall => info::Trampoline::ExitSyncCall,
Trampoline::ThreadIndex => info::Trampoline::ThreadIndex,
Trampoline::ThreadIndex { instance } => info::Trampoline::ThreadIndex {
instance: *instance,
},
Trampoline::ThreadNewIndirect {
instance,
start_func_ty_idx,
Expand Down
7 changes: 5 additions & 2 deletions crates/environ/src/component/info.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1109,7 +1109,10 @@ pub enum Trampoline {
ExitSyncCall,

/// Intrinsic used to implement the `thread.index` component model builtin.
ThreadIndex,
ThreadIndex {
/// The specific component instance which is calling the intrinsic.
instance: RuntimeComponentInstanceIndex,
},

/// Intrinsic used to implement the `thread.new-indirect` component model builtin.
ThreadNewIndirect {
Expand Down Expand Up @@ -1247,7 +1250,7 @@ impl Trampoline {
Trap => format!("trap"),
EnterSyncCall => format!("enter-sync-call"),
ExitSyncCall => format!("exit-sync-call"),
ThreadIndex => format!("thread-index"),
ThreadIndex { .. } => format!("thread-index"),
ThreadNewIndirect { .. } => format!("thread-new-indirect"),
ThreadResumeLater { .. } => format!("thread-resume-later"),
ThreadSuspend { .. } => format!("thread-suspend"),
Expand Down
10 changes: 6 additions & 4 deletions crates/environ/src/component/translate/inline.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1092,10 +1092,12 @@ impl<'a> Inliner<'a> {
.push((*func, dfg::CoreDef::UnsafeIntrinsic(*func, intrinsic)));
}
ThreadIndex { func } => {
let index = self
.result
.trampolines
.push((*func, dfg::Trampoline::ThreadIndex));
let index = self.result.trampolines.push((
*func,
dfg::Trampoline::ThreadIndex {
instance: frame.instance,
},
));
frame.funcs.push((*func, dfg::CoreDef::Trampoline(index)));
}
ThreadNewIndirect {
Expand Down
19 changes: 19 additions & 0 deletions crates/wasmtime/src/runtime/component/func/options.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,17 @@ impl<'a, T: 'static> LowerContext<'a, T> {
) -> Result<usize> {
assert!(self.allow_realloc);

// All calls to `realloc` options in the canonical ABI zero out the
// `context.{get,set}` slots for the duration of the call. This sort of
// fakes a "fresh thread" for each call, but this is the only observable
// state so nothing else needs adjusting. Note though that the original
// values are preserved still to get restored after this call.
#[cfg(feature = "component-model-async")]
let orig_context = core::mem::replace(
self.store.0.vm_store_context_mut().component_context_mut(),
Default::default(),
);

let (component, store) = self.instance.component_and_store_mut(self.store.0);
let instance = self.instance.id().get(store);
let options = &component.env_component().options[self.options];
Expand Down Expand Up @@ -185,6 +196,14 @@ impl<'a, T: 'static> LowerContext<'a, T> {
bail!("realloc return: beyond end of memory")
}

// Note that this restoration isn't part of a `Drop` guard which works
// because once a component traps it's locked-down and inaccessible, so
// it's ok if this isn't restored.
#[cfg(feature = "component-model-async")]
{
*self.store.0.vm_store_context_mut().component_context_mut() = orig_context;
}

Ok(result)
}

Expand Down
104 changes: 0 additions & 104 deletions tests/all/component_model/func.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3150,110 +3150,6 @@ async fn thread_index_via_call(style: ApiStyle) -> Result<()> {
Ok(())
}

#[tokio::test]
async fn thread_index_via_post_return_sync() -> Result<()> {
thread_index_via_post_return(ApiStyle::Sync).await
}

#[tokio::test]
async fn thread_index_via_post_return_async() -> Result<()> {
thread_index_via_post_return(ApiStyle::Async).await
}

#[tokio::test]
async fn thread_index_via_post_return_concurrent() -> Result<()> {
thread_index_via_post_return(ApiStyle::Concurrent).await
}

async fn thread_index_via_post_return(style: ApiStyle) -> Result<()> {
let component = r#"
(component
(core module $m
(import "" "thread.index" (func $thread-index (result i32)))
(global $index (mut i32) (i32.const 0))
(func (export "run")
(global.set $index (call $thread-index))
(if (i32.eqz (global.get $index)) (then unreachable))
)
(func (export "run-post-return")
(local $index i32)
(local.set $index (call $thread-index))
(if (i32.eqz (local.get $index)) (then unreachable))
(if (i32.ne (local.get $index) (global.get $index)) (then unreachable))
)
)
(core func $thread-index (canon thread.index))
(core instance $m (instantiate $m (with "" (instance
(export "thread.index" (func $thread-index))
))))
(func (export "run") (canon lift (core func $m "run") (post-return (func $m "run-post-return"))))
)
"#;
let engine = Engine::new(&style.config())?;
let component = Component::new(&engine, component)?;
let mut store = Store::new(&engine, ());
let linker = Linker::new(&engine);
let instance = style.instantiate(&mut store, &linker, &component).await?;
let run = instance.get_typed_func::<(), ()>(&mut store, "run")?;
style.call(&mut store, run, ()).await?;
Ok(())
}

#[tokio::test]
async fn thread_index_via_cabi_realloc_sync() -> Result<()> {
thread_index_via_cabi_realloc(ApiStyle::Sync).await
}

#[tokio::test]
async fn thread_index_via_cabi_realloc_async() -> Result<()> {
thread_index_via_cabi_realloc(ApiStyle::Async).await
}

#[tokio::test]
async fn thread_index_via_cabi_realloc_concurrent() -> Result<()> {
thread_index_via_cabi_realloc(ApiStyle::Concurrent).await
}

async fn thread_index_via_cabi_realloc(style: ApiStyle) -> Result<()> {
let component = r#"
(component
(core module $m
(import "" "thread.index" (func $thread-index (result i32)))
(global $index (mut i32) (i32.const 0))
(memory (export "memory") 1)
(func (export "realloc") (param i32 i32 i32 i32) (result i32)
(global.set $index (call $thread-index))
(if (i32.eqz (global.get $index)) (then unreachable))
(i32.const 100)
)
(func (export "run") (param i32 i32)
(local $index i32)
(local.set $index (call $thread-index))
(if (i32.eqz (local.get $index)) (then unreachable))
(if (i32.ne (local.get $index) (global.get $index)) (then unreachable))
)
)
(core func $thread-index (canon thread.index))
(core instance $m (instantiate $m (with "" (instance
(export "thread.index" (func $thread-index))
))))
(func (export "run") (param "s" string) (canon lift
(core func $m "run")
(memory $m "memory")
(realloc (func $m "realloc"))
))
)
"#;
let engine = Engine::new(&style.config())?;
let component = Component::new(&engine, component)?;
let mut store = Store::new(&engine, ());
let linker = Linker::new(&engine);
let instance = style.instantiate(&mut store, &linker, &component).await?;
let run = instance.get_typed_func::<(String,), ()>(&mut store, "run")?;
style.call(&mut store, run, ("hola".to_string(),)).await?;
Ok(())
}

#[tokio::test]
async fn thread_index_via_resource_drop_sync() -> Result<()> {
thread_index_via_resource_drop(ApiStyle::Sync).await
Expand Down
33 changes: 31 additions & 2 deletions tests/misc_testsuite/component-model-threading/threading-builtins.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@

(assert_return (invoke "run") (u32.const 42))

;; Test that `thread.index` is exempt from may-leave checks
;; Test that `thread.index` is not exempt from may-leave checks
(component
(core func $thread.index (canon thread.index))

Expand All @@ -118,4 +118,33 @@
(canon lift (core func $dm "run") (post-return (func $dm "post-return"))))
)

(assert_return (invoke "run"))
(assert_trap (invoke "run") "cannot leave component instance")

(component
(core module $m
(import "" "thread.index" (func $thread-index (result i32)))
(global $index (mut i32) (i32.const 0))
(memory (export "memory") 1)
(func (export "realloc") (param i32 i32 i32 i32) (result i32)
(global.set $index (call $thread-index))
(if (i32.eqz (global.get $index)) (then unreachable))
(i32.const 100)
)
(func (export "run") (param i32 i32)
(local $index i32)
(local.set $index (call $thread-index))
(if (i32.eqz (local.get $index)) (then unreachable))
(if (i32.ne (local.get $index) (global.get $index)) (then unreachable))
)
)
(core func $thread-index (canon thread.index))
(core instance $m (instantiate $m (with "" (instance
(export "thread.index" (func $thread-index))
))))
(func (export "run") (param "s" string) (canon lift
(core func $m "run")
(memory $m "memory")
(realloc (func $m "realloc"))
))
)
(assert_trap (invoke "run" (str.const "x")) "cannot leave component instance")
22 changes: 11 additions & 11 deletions tests/misc_testsuite/component-model/async/task-builtins.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,8 +180,8 @@
call $backpressure.inc
call $backpressure.dec

;; context.get should be what was set in `realloc`
(if (i32.ne (call $context.get) (i32.const 100)) (then (unreachable)))
;; context set in realloc should be separate from this thread's context
(if (i32.ne (call $context.get) (i32.const 0)) (then (unreachable)))
)
)
(core instance $m (instantiate $M (with "" (instance
Expand Down Expand Up @@ -416,7 +416,7 @@
(if (i32.ne (local.get 2) (i32.const 1)) (then (unreachable)))
(if (i32.ne (local.get 3) (i32.const 2)) (then (unreachable)))

(if (i32.ne (call $context.get) (i32.const 400)) (then (unreachable)))
(if (i32.ne (call $context.get) (i32.const 0)) (then (unreachable)))
(call $context.set (i32.const 500))

call $backpressure.inc
Expand Down Expand Up @@ -448,11 +448,11 @@
(func (export "run")
;; set this tasks's context before calling $run, in calling $run the
;; runtime will then call `realloc` above for the string return value
;; which should see our 400 value. That will then set 500 which we
;; should then see after the return.
;; which should NOT see our 400 value. That will then set 500 which we
;; should NOT then see after the return.
(call $context.set (i32.const 400))
(call $run (i32.const 20))
(if (i32.ne (call $context.get) (i32.const 500)) (then (unreachable)))
(if (i32.ne (call $context.get) (i32.const 400)) (then (unreachable)))
)
)
(core instance $m (instantiate $M (with "" (instance
Expand Down Expand Up @@ -616,7 +616,7 @@
(if (i32.ne (local.get 3) (i32.const 2)) (then (unreachable)))

call $context.get
i32.const 400
i32.const 0
i32.ne
if unreachable end

Expand Down Expand Up @@ -663,15 +663,15 @@
(import "" "task.return" (func $task.return (param i32)))
(import "" "memory" (memory 1))

;; Set context[0] to 400, then read the future, which should call realloc and set
;; context[0] to 500, then check that we see that value.
;; Set context[0] to 400, then read the future, which should call realloc
;; and set context[0] to 500, then check that we DON'T see that value.
(func (export "run-future") (param $fr i32) (result i32)
(local $ret i32)

(call $context.set (i32.const 400))
(local.set $ret (call $future.read (local.get $fr) (i32.const 40)))
(if (i32.ne (i32.const 0 (; COMPLETED ;)) (local.get $ret)) (then (unreachable)))
(if (i32.ne (call $context.get) (i32.const 500)) (then (unreachable)))
(if (i32.ne (call $context.get) (i32.const 400)) (then (unreachable)))

(call $task.return (i32.const 42))
(i32.const 0 (; EXIT ;))
Expand All @@ -684,7 +684,7 @@
(call $context.set (i32.const 400))
(local.set $ret (call $stream.read (local.get $sr) (i32.const 40) (i32.const 1)))
(if (i32.ne (i32.const 0x10 (; COMPLETED | 1<<4 ;)) (local.get $ret)) (then (unreachable)))
(if (i32.ne (call $context.get) (i32.const 500)) (then (unreachable)))
(if (i32.ne (call $context.get) (i32.const 400)) (then (unreachable)))

(call $task.return (i32.const 42))
(i32.const 0 (; EXIT ;))
Expand Down
Loading