diff --git a/crates/environ/src/component/types.rs b/crates/environ/src/component/types.rs index 56c2874fa8d6..9aefd48d8d75 100644 --- a/crates/environ/src/component/types.rs +++ b/crates/environ/src/component/types.rs @@ -585,6 +585,10 @@ pub struct TypeComponentInstance { pub struct TypeFunc { /// Whether or not this is an async function. pub async_: bool, + /// Whether any parameter (transitively) contains a `borrow` handle, + /// letting the runtime skip borrow-scope tracking when lending is + /// impossible. (Borrows are only valid in parameters, not results.) + pub contains_borrow: bool, /// Names of parameters. pub param_names: Vec, /// Parameters to the function represented as a tuple. diff --git a/crates/environ/src/component/types_builder.rs b/crates/environ/src/component/types_builder.rs index a9822ea1afdc..a1b4da20885d 100644 --- a/crates/environ/src/component/types_builder.rs +++ b/crates/environ/src/component/types_builder.rs @@ -230,16 +230,18 @@ impl ComponentTypesBuilder { .params .iter() .map(|(_name, ty)| self.valtype(types, ty)) - .collect::>()?; + .collect::>>()?; let results = ty .result .iter() .map(|ty| self.valtype(types, ty)) .collect::>()?; - let params = self.new_tuple_type(params); + let contains_borrow = params.iter().any(|ty| self.ty_contains_borrow_resource(ty)); + let params = self.new_tuple_type(params.into()); let results = self.new_tuple_type(results); let ty = TypeFunc { async_: ty.async_, + contains_borrow, param_names, params, results, diff --git a/crates/wasmtime/src/runtime/component/concurrent.rs b/crates/wasmtime/src/runtime/component/concurrent.rs index a5a34e481557..a92e63868d47 100644 --- a/crates/wasmtime/src/runtime/component/concurrent.rs +++ b/crates/wasmtime/src/runtime/component/concurrent.rs @@ -823,18 +823,52 @@ pub(crate) enum WaitResult { Completed, } +/// Poll `future` once with a no-op waker in the store's tls context, +/// before any task machinery exists. For pending futures, +/// `Instance::continue_polling` should be used to add them to +/// `ConcurrentState::futures` and have them polled by the event loop. +#[inline] +pub(crate) fn poll_once(store: &mut dyn VMStore, future: &mut F) -> Poll +where + F: Future + Unpin, +{ + tls::set(store, || { + Pin::new(future).poll(&mut Context::from_waker(&Waker::noop())) + }) +} + /// Poll the specified future until it completes on behalf of a guest->host call /// using a sync-lowered import. /// -/// This is similar to `Instance::first_poll` except it's for sync-lowered -/// imports, meaning we don't need to handle cancellation and we can block the -/// caller until the task completes, at which point the caller can handle -/// lowering the result to the guest's stack and linear memory. +/// This is similar to `poll_once` + `Instance::continue_polling` except it's +/// for sync-lowered imports, meaning we don't need to handle cancellation and +/// we can block the caller until the task completes, at which point the caller +/// can handle lowering the result to the guest's stack and linear memory. pub(crate) fn poll_and_block( store: &mut dyn VMStore, future: impl Future> + Send + 'static, ) -> Result { - let task = store.current_host_thread()?; + // Fast path: poll the future once before materializing any `HostTask`. + // Most host implementations complete immediately, in which case the + // result is returned without any task-table traffic or result boxing. + let mut inner = Box::pin(future); + let poll = tls::set(store, || { + inner + .as_mut() + .poll(&mut Context::from_waker(&Waker::noop())) + }); + if let Poll::Ready(result) = poll { + return result; + } + + // Slow path: the future is genuinely pending. Materialize the host task + // now if the entrypoint deferred it (in which case this function also + // owns its cleanup below); otherwise use the already-current task. + let (task, materialized) = match store.current_thread()? { + CurrentThread::Host(task) => (task, false), + _ => (store.materialize_host_task()?, true), + }; + let future = inner; // Wrap the future in a closure which will take care of stashing the result // in `GuestTask::result` and resuming this fiber when the host task @@ -898,13 +932,21 @@ pub(crate) fn poll_and_block( // Retrieve and return the result. let host_state = &mut store.concurrent_state_mut()?.get_mut(task)?.state; - match mem::replace(host_state, HostTaskState::CalleeDone { cancelled: false }) { + let result = match mem::replace(host_state, HostTaskState::CalleeDone { cancelled: false }) { HostTaskState::CalleeFinished(result) => Ok(match result.downcast() { Ok(result) => *result, Err(_) => bail_bug!("host task finished with wrong type of result"), }), _ => bail_bug!("unexpected host task state after completion"), + }; + + // A task materialized here (rather than by the entrypoint) is also + // cleaned up here, mirroring `host_task_delete` for the eager case. + if materialized { + store.host_task_reenter_caller()?; + store.concurrent_state_mut()?.delete(task)?; } + result } /// Execute the specified guest call. @@ -1601,6 +1643,7 @@ impl StoreContextMut<'_, T> { impl StoreOpaque { /// Returns the currently-running thread, promoting any deferred lazy thread /// into a fully-materialized `CurrentThread`. + #[inline] pub(crate) fn current_thread(&mut self) -> Result { // Without concurrency support there is nothing to force. if !self.concurrency_support() { @@ -1619,6 +1662,13 @@ impl StoreOpaque { .unforced_current_thread); } + self.force_deferred_current_thread() + } + + /// Slow path of [`Self::current_thread`]: promote the deferred lazy + /// thread into a fully-materialized `CurrentThread`. + #[cold] + fn force_deferred_current_thread(&mut self) -> Result { // The component instance whose adapters pushed the deferred frames; all // frames in a guest-to-guest, sync-to-sync call chain of fused adapters // live within a single `wasmtime::component::Instance` (because @@ -1690,13 +1740,6 @@ impl StoreOpaque { } } - fn current_host_thread(&mut self) -> Result> { - match self.current_thread()?.host() { - Some(id) => Ok(id), - None => bail_bug!("current thread is not a host thread"), - } - } - /// Returns whether there's a pending cancellation on the current guest thread, /// consuming the event if so. fn take_pending_cancellation(&mut self) -> Result { @@ -1812,17 +1855,38 @@ impl StoreOpaque { /// relatively expensive table manipulations. This would ideally be /// optimized to avoid the full allocation of a `HostTask` in at least some /// situations. - pub(crate) fn host_task_create(&mut self) -> Result>> { + #[inline] + pub(crate) fn host_task_create( + &mut self, + track_scope: bool, + defer: bool, + ) -> Result>> { if !self.concurrency_support() { - self.enter_call_not_concurrent()?; + // Resource-borrow scope tracking is only needed when the + // function's parameters can actually contain `borrow` handles; + // this is precomputed per function type. + if track_scope { + self.enter_call_not_concurrent()?; + } + return Ok(None); + } + // Deferred (borrow-free) calls run as the caller's thread; a task is + // only materialized if the host future actually suspends. + if defer { return Ok(None); } + Ok(Some(self.materialize_host_task()?)) + } + + /// Push a new `HostTask` for a host call made by the current guest + /// thread, and make it the current thread. + pub(crate) fn materialize_host_task(&mut self) -> Result> { let caller = self.current_guest_thread()?; let state = self.concurrent_state_mut()?; let task = state.push(HostTask::new(caller, HostTaskState::CalleeStarted))?; log::trace!("new host task {task:?}"); self.set_thread(task)?; - Ok(Some(task)) + Ok(task) } /// Invoked before lowering the results of a host task to the guest. @@ -1830,11 +1894,16 @@ impl StoreOpaque { /// This is used to update the current thread annotations within the store /// to ensure that it reflects the guest task, not the host task, since /// lowering may execute guest code. + #[inline] pub fn host_task_reenter_caller(&mut self) -> Result<()> { if !self.concurrency_support() { return Ok(()); } - let task = self.current_host_thread()?; + // For deferred (task-less) host calls the current thread already is + // the caller; nothing to restore. + let CurrentThread::Host(task) = self.current_thread()? else { + return Ok(()); + }; let caller = self.concurrent_state_mut()?.get_mut(task)?.caller; self.set_thread(caller)?; Ok(()) @@ -1846,14 +1915,21 @@ impl StoreOpaque { /// Note that this isn't invoked when the host is invoked asynchronously and /// the host isn't complete yet. In that situation the host task persists /// and will be cleaned up separately in `subtask_drop` - pub(crate) fn host_task_delete(&mut self, task: Option>) -> Result<()> { + #[inline] + pub(crate) fn host_task_delete( + &mut self, + task: Option>, + track_scope: bool, + ) -> Result<()> { match task { Some(task) => { log::trace!("delete host task {task:?}"); self.concurrent_state_mut()?.delete(task)?; } None => { - self.exit_call_not_concurrent(); + if track_scope { + self.exit_call_not_concurrent(); + } } } Ok(()) @@ -2341,6 +2417,7 @@ impl StoreOpaque { /// Used by `ResourceTables` to record the scope of a borrow to get undone /// in the future. + #[inline] pub(crate) fn current_scope_id(&mut self) -> Result> { if !self.concurrency_support() { return self.current_scope_id_not_concurrent(); @@ -3182,65 +3259,31 @@ impl Instance { Ok(status.pack(waitable)) } - /// Poll the specified future once on behalf of a guest->host call using an - /// async-lowered import. - /// - /// If it returns `Ready`, return `Ok(None)`. Otherwise, if it returns - /// `Pending`, add it to the set of futures to be polled as part of this - /// instance's event loop until it completes, and then return - /// `Ok(Some(handle))` where `handle` is the waitable handle to return. - /// - /// Whether the future returns `Ready` immediately or later, the `lower` - /// function will be used to lower the result, if any, into the guest caller's - /// stack and linear memory. The `lower` function is invoked with `None` if - /// the future is cancelled. - pub(crate) fn first_poll( + /// Continue an async-lowered host call whose future was already given + /// its first poll (see `poll_once`) and is pending: wrap it in an + /// abortable `JoinHandle` (for cancellation), materialize the `HostTask` + /// if the entrypoint deferred it, and register the future with the event + /// loop, returning the waitable handle for the guest. + pub(crate) fn continue_polling( self, mut store: StoreContextMut<'_, T>, future: impl Future> + Send + 'static, lower: impl FnOnce(StoreContextMut, Option) -> Result<()> + Send + 'static, ) -> Result> { let token = StoreToken::new(store.as_context_mut()); - let task = store.0.current_host_thread()?; - let state = store.0.concurrent_state_mut()?; - // Create an abortable future which hooks calls to poll and manages call - // context state for the future. let (join_handle, future) = JoinHandle::run(future); + let future = Box::pin(future); + let task = match store.0.current_thread()? { + CurrentThread::Host(task) => task, + _ => store.0.materialize_host_task()?, + }; { - let state = &mut state.get_mut(task)?.state; + let state = &mut store.0.concurrent_state_mut()?.get_mut(task)?.state; assert!(matches!(state, HostTaskState::CalleeStarted)); *state = HostTaskState::CalleeRunning(join_handle); } - let mut future = Box::pin(future); - - // Finally, poll the future. We can use a dummy `Waker` here because - // we'll add the future to `ConcurrentState::futures` and poll it - // automatically from the event loop if it doesn't complete immediately - // here. - let poll = tls::set(store.0, || { - future - .as_mut() - .poll(&mut Context::from_waker(&Waker::noop())) - }); - - match poll { - // It finished immediately; lower the result and delete the task. - Poll::Ready(result) => { - let result = result.transpose()?; - lower(store.as_context_mut(), result)?; - return Ok(None); - } - - // Future isn't ready yet, so fall through. - Poll::Pending => {} - } - - // It hasn't finished yet; add the future to - // `ConcurrentState::futures` so it will be polled by the event - // loop and allocate a waitable handle to return to the guest. - // Wrap the future in a closure responsible for lowering the result into // the guest's stack and memory, as well as notifying any waiters that // the task returned. @@ -5230,13 +5273,6 @@ impl CurrentThread { } } - fn host(&self) -> Option> { - match self { - Self::Host(id) => Some(*id), - _ => None, - } - } - fn is_none(&self) -> bool { matches!(self, Self::None) } diff --git a/crates/wasmtime/src/runtime/component/concurrent_disabled.rs b/crates/wasmtime/src/runtime/component/concurrent_disabled.rs index a476bfb869ad..2b4473f10910 100644 --- a/crates/wasmtime/src/runtime/component/concurrent_disabled.rs +++ b/crates/wasmtime/src/runtime/component/concurrent_disabled.rs @@ -162,16 +162,22 @@ impl StoreOpaque { Ok(self.exit_call_not_concurrent()) } - pub(crate) fn host_task_create(&mut self) -> Result<()> { - self.enter_call_not_concurrent() + pub(crate) fn host_task_create(&mut self, track_scope: bool, _defer: bool) -> Result<()> { + if track_scope { + self.enter_call_not_concurrent()?; + } + Ok(()) } pub(crate) fn host_task_reenter_caller(&mut self) -> Result<()> { Ok(()) } - pub(crate) fn host_task_delete(&mut self, (): ()) -> Result<()> { - Ok(self.exit_call_not_concurrent()) + pub(crate) fn host_task_delete(&mut self, (): (), track_scope: bool) -> Result<()> { + if track_scope { + self.exit_call_not_concurrent(); + } + Ok(()) } pub(crate) fn check_blocking(&mut self) -> crate::Result<()> { diff --git a/crates/wasmtime/src/runtime/component/func/host.rs b/crates/wasmtime/src/runtime/component/func/host.rs index b1d8a931f8bd..620069c1e113 100644 --- a/crates/wasmtime/src/runtime/component/func/host.rs +++ b/crates/wasmtime/src/runtime/component/func/host.rs @@ -22,6 +22,8 @@ use core::mem::{self, MaybeUninit}; #[cfg(feature = "async")] use core::pin::Pin; use core::ptr::NonNull; +#[cfg(feature = "component-model-async")] +use core::task::Poll; use wasmtime_environ::component::{ CanonicalAbiInfo, InterfaceType, MAX_FLAT_PARAMS, MAX_FLAT_RESULTS, OptionsIndex, TypeFuncIndex, }; @@ -359,6 +361,7 @@ where /// "Rust" entrypoint after panic-handling infrastructure is set up and raw /// arguments are translated to Rust types. + #[inline] fn entrypoint( &self, mut store: StoreContextMut<'_, T>, @@ -369,6 +372,9 @@ where ) -> Result<()> { let vminstance = instance.id().get(store.0); let async_ = vminstance.component().env_component().options[options].async_; + // Skip resource-borrow scope tracking when the signature makes + // lending impossible (precomputed at compile time). + let track_scope = vminstance.component().types()[ty].contains_borrow; // If this is a synchronous-lower of a host-async function, then the // guest is blocking. Test, in the context of the guest task, if that's @@ -378,12 +384,23 @@ where } // Enter the host by pushing a `HostTask` into the concurrent state. - let host_task = store.0.host_task_create()?; + // Borrow-free calls defer task creation entirely (see + // `host_task_create`); borrow-ful calls keep the eager task as their + // resource scope identity. + let defer_task = !track_scope; + let host_task = store.0.host_task_create(track_scope, defer_task)?; let host_task_complete = if async_ { #[cfg(feature = "component-model-async")] { - self.call_async_lower(store.as_context_mut(), instance, ty, options, storage)? + self.call_async_lower( + store.as_context_mut(), + instance, + ty, + options, + storage, + track_scope, + )? } #[cfg(not(feature = "component-model-async"))] unreachable!( @@ -391,7 +408,14 @@ where when `component-model-async` feature disabled" ); } else { - self.call_sync_lower(store.as_context_mut(), instance, ty, options, storage)?; + self.call_sync_lower( + store.as_context_mut(), + instance, + ty, + options, + storage, + track_scope, + )?; true }; @@ -401,7 +425,7 @@ where // function transitively would have updated the current guest thread to // the caller of this host function. if host_task_complete { - store.0.host_task_delete(host_task)?; + store.0.host_task_delete(host_task, track_scope)?; } Ok(()) @@ -414,6 +438,7 @@ where /// the `async` option when lowered. Note that the host function itself /// can still be async, in which case this will block here waiting for it /// to finish. + #[inline] fn call_sync_lower( &self, mut store: StoreContextMut<'_, T>, @@ -421,8 +446,13 @@ where ty: TypeFuncIndex, options: OptionsIndex, storage: &mut [MaybeUninit], + track_scope: bool, ) -> Result<()> { - let mut lift = LiftContext::new(store.0.store_opaque_mut(), options, instance)?; + let mut lift = if track_scope { + LiftContext::new(store.0.store_opaque_mut(), options, instance)? + } else { + LiftContext::new_without_scope(store.0.store_opaque_mut(), options, instance)? + }; let (params, rest) = self.load_params(&mut lift, ty, MAX_FLAT_PARAMS, storage)?; let ret = match self.run(store.as_context_mut(), params) { @@ -463,6 +493,7 @@ where ty: TypeFuncIndex, options: OptionsIndex, storage: &mut [MaybeUninit], + track_scope: bool, ) -> Result { use wasmtime_environ::component::MAX_FLAT_ASYNC_PARAMS; @@ -473,7 +504,11 @@ where // Lift the parameters, either from flat storage or from linear // memory. - let mut lift = LiftContext::new(store.0.store_opaque_mut(), options, instance)?; + let mut lift = if track_scope { + LiftContext::new(store.0.store_opaque_mut(), options, instance)? + } else { + LiftContext::new_without_scope(store.0.store_opaque_mut(), options, instance)? + }; let (params, rest) = self.load_params(&mut lift, ty, MAX_FLAT_ASYNC_PARAMS, storage)?; // Load/validate the return pointer, if present. @@ -505,8 +540,24 @@ where None } #[cfg(feature = "component-model-async")] - HostResult::Future(future) => { - instance.first_poll(store, future, move |store, ret| { + HostResult::Future(mut future) => { + // Ready results are lowered exactly like the `Done` arm + // above, with no task, no `JoinHandle`, and no re-boxing + // (the future is already heap-pinned, hence `Unpin`). + match concurrent::poll_once(store.0, &mut future) { + Poll::Ready(result) => { + Self::lower_result_and_exit_call( + &mut LowerContext::new(store, options, instance), + ty, + Some(result?), + Destination::Memory(retptr), + )?; + storage[0].write(ValRaw::u32(Status::Returned.pack(None))); + return Ok(true); + } + Poll::Pending => {} + } + instance.continue_polling(store, future, move |store, ret| { Self::lower_result_and_exit_call( &mut LowerContext::new(store, options, instance), ty, @@ -530,6 +581,7 @@ where /// /// This will internally decide the ABI source of the parameters and use /// `storage` appropriately. + #[inline] fn load_params<'a>( &self, lift: &mut LiftContext<'_>, @@ -565,6 +617,7 @@ where } /// Stores the result `ret` into `dst` which is calculated per the ABI. + #[inline] fn lower_result_and_exit_call( lower: &mut LowerContext<'_, T>, ty: TypeFuncIndex, @@ -572,8 +625,11 @@ where dst: Destination<'_>, ) -> Result<()> { // Before lowering below semantically ensure that the caller has dropped - // all of its borrows and such. - lower.validate_scope_exit()?; + // all of its borrows and such. Skipped when the signature cannot + // contain borrows (in which case no scope was entered either). + if lower.types[ty].contains_borrow { + lower.validate_scope_exit()?; + } // At this point we're transitioning back to the caller task which means // that the current task needs to be updated. This will restore the diff --git a/crates/wasmtime/src/runtime/component/func/options.rs b/crates/wasmtime/src/runtime/component/func/options.rs index 6538ab0837cf..0bea3eb14a26 100644 --- a/crates/wasmtime/src/runtime/component/func/options.rs +++ b/crates/wasmtime/src/runtime/component/func/options.rs @@ -338,10 +338,36 @@ impl<'a> LiftContext<'a> { store: &'a mut StoreOpaque, options: OptionsIndex, instance_handle: Instance, + ) -> Result> { + Self::new_impl(store, options, instance_handle, true) + } + + /// Same as [`Self::new`] but without resolving the current resource + /// scope, for calls whose signature statically cannot contain `borrow` + /// handles (the only consumers of the scope id). + #[inline] + pub(crate) fn new_without_scope( + store: &'a mut StoreOpaque, + options: OptionsIndex, + instance_handle: Instance, + ) -> Result> { + Self::new_impl(store, options, instance_handle, false) + } + + #[inline] + fn new_impl( + store: &'a mut StoreOpaque, + options: OptionsIndex, + instance_handle: Instance, + want_scope: bool, ) -> Result> { let store_id = store.id(); let hostcall_fuel = store.hostcall_fuel(); - let current_scope_id = store.current_scope_id()?; + let current_scope_id = if want_scope { + store.current_scope_id()? + } else { + None + }; // From `&mut StoreOpaque` provided the goal here is to project out // three different disjoint fields owned by the store: memory, // `CallContexts`, and `HandleTable`. There's no native API for that diff --git a/crates/wasmtime/src/runtime/component/store.rs b/crates/wasmtime/src/runtime/component/store.rs index 2dccbb41d712..6a670aaa259b 100644 --- a/crates/wasmtime/src/runtime/component/store.rs +++ b/crates/wasmtime/src/runtime/component/store.rs @@ -218,6 +218,7 @@ impl StoreComponentInstanceId { /// # Panics /// /// Panics if `self` does not belong to `store`. + #[inline] pub(crate) fn get<'a>(&self, store: &'a StoreOpaque) -> &'a ComponentInstance { self.assert_belongs_to(store.id()); store.component_instance(self.instance) @@ -228,6 +229,7 @@ impl StoreComponentInstanceId { /// # Panics /// /// Panics if `self` does not belong to `store`. + #[inline] pub(crate) fn get_mut<'a>(&self, store: &'a mut StoreOpaque) -> Pin<&'a mut ComponentInstance> { self.from_data_get_mut(store.store_data_mut()) } @@ -356,6 +358,7 @@ impl StoreOpaque { support } + #[inline] pub(crate) fn lift_context_parts( &mut self, instance: Instance, @@ -420,6 +423,7 @@ impl StoreOpaque { )) } + #[inline] pub(crate) fn enter_call_not_concurrent(&mut self) -> Result<()> { let state = match &mut self.component_data_mut().task_state { ComponentTaskState::NotConcurrent(state) => state, @@ -430,6 +434,7 @@ impl StoreOpaque { Ok(()) } + #[inline] pub(crate) fn exit_call_not_concurrent(&mut self) { let state = match &mut self.component_data_mut().task_state { ComponentTaskState::NotConcurrent(state) => state, @@ -459,6 +464,7 @@ impl StoreOpaque { } } + #[inline] pub(crate) fn current_scope_id_not_concurrent(&mut self) -> Result> { match &mut self.component_data_mut().task_state { ComponentTaskState::NotConcurrent(state) => match state.scopes.len().checked_sub(1) {