Skip to content

Refactor defining core wasm async functions #12365

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 2 commits into from
Jan 16, 2026
Merged

Refactor defining core wasm async functions #12365

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c2770ba
Refactor defining core wasm async functions
alexcrichton Jan 16, 2026
bb53cfe
Adjust some features
alexcrichton Jan 16, 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
34 changes: 14 additions & 20 deletions crates/wasmtime/src/runtime/component/linker.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -771,9 +771,9 @@ impl<T: 'static> LinkerInstance<'_, T> {
ty: ResourceType,
dtor: impl Fn(StoreContextMut<'_, T>, u32) -> Result<()> + Send + Sync + 'static,
) -> Result<()> {
let dtor = Arc::new(crate::func::HostFunc::wrap_inner(
let dtor = Arc::new(crate::func::HostFunc::wrap(
&self.engine,
move |mut cx: crate::Caller<'_, T>, (param,): (u32,)| dtor(cx.as_context_mut(), param),
move |mut cx: crate::Caller<'_, T>, param: u32| dtor(cx.as_context_mut(), param),
));
self.insert(name, Definition::Resource(ty, dtor))?;
Ok(())
Expand All @@ -792,12 +792,9 @@ impl<T: 'static> LinkerInstance<'_, T> {
self.engine.config().async_support,
"cannot use `resource_async` without enabling async support in the config"
);
let dtor = Arc::new(crate::func::HostFunc::wrap_inner(
let dtor = Arc::new(crate::func::HostFunc::wrap_async(
&self.engine,
move |mut cx: crate::Caller<'_, T>, (param,): (u32,)| {
cx.as_context_mut()
.block_on(|store| dtor(store, param).into())?
},
move |cx: crate::Caller<'_, T>, (param,): (u32,)| dtor(cx.into(), param),
));
self.insert(name, Definition::Resource(ty, dtor))?;
Ok(())
Expand All @@ -823,23 +820,20 @@ impl<T: 'static> LinkerInstance<'_, T> {
// up the implementation to avoid using e.g. `Accessor::new` and
// `tls::set` directly.
let dtor = Arc::new(dtor);
let dtor = Arc::new(crate::func::HostFunc::wrap_inner(
let dtor = Arc::new(crate::func::HostFunc::wrap_async(
&self.engine,
move |mut cx: crate::Caller<'_, T>, (param,): (u32,)| {
let dtor = dtor.clone();
cx.as_context_mut().block_on(move |mut store| {
Box::pin(async move {
let accessor =
&Accessor::new(crate::store::StoreToken::new(store.as_context_mut()));
let mut future = std::pin::pin!(dtor(accessor, param));
std::future::poll_fn(|cx| {
crate::component::concurrent::tls::set(store.0, || {
future.as_mut().poll(cx)
})
})
.await
Box::new(async move {
let mut store = cx.as_context_mut();
let accessor =
&Accessor::new(crate::store::StoreToken::new(store.as_context_mut()));
let mut future = std::pin::pin!(dtor(accessor, param));
std::future::poll_fn(|cx| {
crate::component::concurrent::tls::set(store.0, || future.as_mut().poll(cx))
})
})?
.await
})
},
));
self.insert(name, Definition::Resource(ty, dtor))?;
Expand Down
160 changes: 87 additions & 73 deletions crates/wasmtime/src/runtime/func.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -371,17 +371,16 @@ impl Func {
/// Panics if the given function type is not associated with this store's
/// engine.
pub fn new<T: 'static>(
store: impl AsContextMut<Data = T>,
mut store: impl AsContextMut<Data = T>,
ty: FuncType,
func: impl Fn(Caller<'_, T>, &[Val], &mut [Val]) -> Result<()> + Send + Sync + 'static,
) -> Self {
assert!(ty.comes_from_same_engine(store.as_context().engine()));
let ty_clone = ty.clone();
unsafe {
Func::new_unchecked(store, ty, move |caller, values| {
Func::invoke_host_func_for_wasm(caller, &ty_clone, values, &func)
})
}
let store = store.as_context_mut().0;
let host = HostFunc::new(store.engine(), ty, func);

// SAFETY: the `T` used by `func` matches the `T` of the store we're
// inserting into via this function's type signature.
unsafe { host.into_func(store) }
}

/// Creates a new [`Func`] with the given arguments, although has fewer
Expand Down Expand Up @@ -417,7 +416,6 @@ impl Func {
ty: FuncType,
func: impl Fn(Caller<'_, T>, &mut [ValRaw]) -> Result<()> + Send + Sync + 'static,
) -> Self {
assert!(ty.comes_from_same_engine(store.as_context().engine()));
Copy link
Contributor

@dicej dicej Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this (and the other removed comes_from_same_engine assertions) no longer needed?

Copy link
Member Author

@alexcrichton alexcrichton Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh still needed, but amongst the layers here this assert was happening at every layer so I deleted all the duplicate ones to delegate to the lowest-level layers to perform the assertion

let store = store.as_context_mut().0;

// SAFETY: the contract required by `new_unchecked` is the same as the
Expand Down Expand Up @@ -502,7 +500,7 @@ impl Func {
/// # }
/// ```
#[cfg(feature = "async")]
pub fn new_async<T, F>(store: impl AsContextMut<Data = T>, ty: FuncType, func: F) -> Func
pub fn new_async<T, F>(mut store: impl AsContextMut<Data = T>, ty: FuncType, func: F) -> Func
where
F: for<'a> Fn(
Caller<'a, T>,
Expand All @@ -514,24 +512,16 @@ impl Func {
+ 'static,
T: 'static,
{
let store = store.as_context_mut().0;
assert!(
store.as_context().async_support(),
store.async_support(),
"cannot use `new_async` without enabling async support in the config"
);
assert!(ty.comes_from_same_engine(store.as_context().engine()));
return Func::new(
store,
ty,
move |Caller { store, caller }, params, results| {
store.with_blocking(|store, cx| {
cx.block_on(core::pin::Pin::from(func(
Caller { store, caller },
params,
results,
)))
})?
},
);
let host = HostFunc::new_async(store.engine(), ty, func);

// SAFETY: the `T` used by `func` matches the `T` of the store we're
// inserting into via this function's type signature.
unsafe { host.into_func(store) }
}

/// Creates a new `Func` from a store and a funcref within that store.
Expand Down Expand Up @@ -812,22 +802,6 @@ impl Func {
unsafe { host.into_func(store) }
}

#[cfg(feature = "async")]
fn wrap_inner<F, T, Params, Results>(mut store: impl AsContextMut<Data = T>, func: F) -> Func
where
F: Fn(Caller<'_, T>, Params) -> Results + Send + Sync + 'static,
Params: WasmTyList,
Results: WasmRet,
T: 'static,
{
let store = store.as_context_mut().0;
let host = HostFunc::wrap_inner(store.engine(), func);

// SAFETY: The `T` the closure takes is the same as the `T` of the store
// we're inserting into via the type signature above.
unsafe { host.into_func(store) }
}

/// Same as [`Func::wrap`], except the closure asynchronously produces the
/// result and the arguments are passed within a tuple. For more information
/// see the [`Func`] documentation.
Expand All @@ -836,7 +810,7 @@ impl Func {
///
/// This function will panic if called with a non-asynchronous store.
#[cfg(feature = "async")]
pub fn wrap_async<T, F, P, R>(store: impl AsContextMut<Data = T>, func: F) -> Func
pub fn wrap_async<T, F, P, R>(mut store: impl AsContextMut<Data = T>, func: F) -> Func
where
F: for<'a> Fn(Caller<'a, T>, P) -> Box<dyn Future<Output = R> + Send + 'a>
+ Send
Expand All @@ -846,16 +820,16 @@ impl Func {
R: WasmRet,
T: 'static,
{
let store = store.as_context_mut().0;
assert!(
store.as_context().async_support(),
store.async_support(),
concat!("cannot use `wrap_async` without enabling async support on the config")
);
Func::wrap_inner(store, move |Caller { store, caller }, args| {
match store.block_on(|store| func(Caller { store, caller }, args).into()) {
Ok(ret) => ret.into_fallible(),
Err(e) => R::fallible_from_error(e),
}
})
let host = HostFunc::wrap_async(store.engine(), func);

// SAFETY: The `T` the closure takes is the same as the `T` of the store
// we're inserting into via the type signature above.
unsafe { host.into_func(store) }
}

/// Returns the underlying wasm type that this `Func` has.
Expand Down Expand Up @@ -1949,7 +1923,7 @@ macro_rules! impl_into_func {
for_each_function_signature!(impl_into_func);

/// Trait implemented for various tuples made up of types which implement
/// [`WasmTy`] that can be passed to [`Func::wrap_inner`] and
/// [`WasmTy`] that can be passed to [`Func::wrap_async`] and
/// [`HostContext::from_closure`].
pub unsafe trait WasmTyList {
/// Get the value type that each Type in the list represents.
Expand Down Expand Up @@ -2028,16 +2002,6 @@ pub struct Caller<'a, T: 'static> {
}

impl<T> Caller<'_, T> {
#[cfg(feature = "async")]
pub(crate) fn new(store: StoreContextMut<'_, T>, caller: Instance) -> Caller<'_, T> {
Caller { store, caller }
}

#[cfg(feature = "async")]
pub(crate) fn caller(&self) -> Instance {
self.caller
}

/// Executes `f` with an appropriate `Caller`.
///
/// This is the entrypoint for host functions in core wasm. The `store` and
Expand Down Expand Up @@ -2257,6 +2221,12 @@ impl<T: 'static> AsContextMut for Caller<'_, T> {
}
}

impl<'a, T: 'static> From<Caller<'a, T>> for StoreContextMut<'a, T> {
fn from(caller: Caller<'a, T>) -> Self {
caller.store
}
}

// State stored inside a `VMArrayCallHostFuncContext`.
struct HostFuncState<F> {
// The actual host function.
Expand Down Expand Up @@ -2457,6 +2427,40 @@ impl HostFunc {
}
}

/// Analog of [`Func::new_async`]
///
/// # Panics
///
/// Panics if the given function type is not associated with the given
/// engine.
#[cfg(feature = "async")]
pub fn new_async<T, F>(engine: &Engine, ty: FuncType, func: F) -> Self
where
F: for<'a> Fn(
Caller<'a, T>,
&'a [Val],
&'a mut [Val],
) -> Box<dyn Future<Output = Result<()>> + Send + 'a>
+ Send
+ Sync
+ 'static,
T: 'static,
{
HostFunc::new(
engine,
ty,
move |Caller { store, caller }, params, results| {
store.with_blocking(|store, cx| {
cx.block_on(core::pin::Pin::from(func(
Caller { store, caller },
params,
results,
)))
})?
},
)
}

/// Analog of [`Func::new_unchecked`]
///
/// # Panics
Expand Down Expand Up @@ -2498,19 +2502,6 @@ impl HostFunc {
HostFunc::_new(engine, ctx.into())
}

/// Analog of [`Func::wrap_inner`]
#[cfg(any(feature = "component-model", feature = "async"))]
pub fn wrap_inner<F, T, Params, Results>(engine: &Engine, func: F) -> Self
where
F: Fn(Caller<'_, T>, Params) -> Results + Send + Sync + 'static,
Params: WasmTyList,
Results: WasmRet,
T: 'static,
{
let ctx = HostContext::from_closure(engine, func);
HostFunc::_new(engine, ctx)
}

/// Analog of [`Func::wrap`]
pub fn wrap<T, Params, Results>(
engine: &Engine,
Expand All @@ -2523,6 +2514,29 @@ impl HostFunc {
HostFunc::_new(engine, ctx)
}

/// Analog of [`Func::wrap_async`]
#[cfg(feature = "async")]
pub fn wrap_async<T, F, P, R>(engine: &Engine, func: F) -> Self
where
F: for<'a> Fn(Caller<'a, T>, P) -> Box<dyn Future<Output = R> + Send + 'a>
+ Send
+ Sync
+ 'static,
P: WasmTyList,
R: WasmRet,
T: 'static,
{
HostFunc::_new(
engine,
HostContext::from_closure(engine, move |Caller { store, caller }, args| {
match store.block_on(|store| func(Caller { store, caller }, args).into()) {
Ok(ret) => ret.into_fallible(),
Err(e) => R::fallible_from_error(e),
}
}),
)
}

/// Requires that this function's signature is already registered within
/// `Engine`. This happens automatically during the above two constructors.
fn _new(engine: &Engine, ctx: HostContext) -> Self {
Expand Down
Loading