diff --git a/benches/call.rs b/benches/call.rs index be24a286c430..f0293fa8d10a 100644 --- a/benches/call.rs +++ b/benches/call.rs @@ -904,12 +904,12 @@ mod component { #[cfg(feature = "component-model-async")] mod component_async { use super::*; - use wasmtime::component::{Component, Linker}; + use wasmtime::component::{Component, Linker, Resource, ResourceType}; pub fn measure_execution_time(c: &mut Criterion) { let mut group = c.benchmark_group("component-async"); - host_to_guest(&mut group); - guest_to_host(&mut group); + host_to_wasm(&mut group); + wasm_to_host(&mut group); } fn engine() -> Engine { @@ -918,7 +918,7 @@ mod component_async { Engine::new(&config).unwrap() } - fn host_to_guest(group: &mut BenchmarkGroup<'_, WallTime>) { + fn host_to_wasm(group: &mut BenchmarkGroup<'_, WallTime>) { let engine = engine(); let component = Component::new( &engine, @@ -957,7 +957,7 @@ mod component_async { .get_typed_func::<(), ()>(&mut store, "nop") .unwrap(); - group.bench_function("host-to-guest", |b| { + group.bench_function("host-to-wasm", |b| { b.iter(|| { run_await(store.run_concurrent(async |accessor| { nop.call_concurrent(accessor, ()).await.unwrap() @@ -967,16 +967,46 @@ mod component_async { }); } - fn guest_to_host(group: &mut BenchmarkGroup<'_, WallTime>) { - let engine = engine(); + fn wasm_to_host(group: &mut BenchmarkGroup<'_, WallTime>) { + let mut config = Config::new(); + config.wasm_component_model_async(true); + config.wasm_component_model_async_stackful(true); + let engine = Engine::new(&config).unwrap(); let component = Component::new( &engine, r#" (component + (import "r" (type $r (sub resource))) (import "nop" (func $nop async)) + (import "new" (func $new (result (own $r)))) + (import "borrow" (func $borrow async (param "r" (borrow $r)))) + (import "pending-once" (func $pending-once async)) + (core module $memory (memory (export "memory") 1)) + (core instance $memory (instantiate $memory)) (core func $nop (canon lower (func $nop) async)) + (core func $new (canon lower (func $new))) + (core func $borrow (canon lower (func $borrow) async)) + (core func $pending-once (canon lower (func $pending-once) async)) + (core func $drop (canon resource.drop $r)) + (core func $waitable-set-new (canon waitable-set.new)) + (core func $waitable-join (canon waitable.join)) + (core func $waitable-set-wait + (canon waitable-set.wait (memory (core memory $memory "memory")))) + (core func $waitable-set-drop (canon waitable-set.drop)) + (core func $subtask-drop (canon subtask.drop)) (core module $m + (import "" "memory" (memory 1)) (import "" "nop" (func $nop (result i32))) + (import "" "new" (func $new (result i32))) + (import "" "borrow" (func $borrow (param i32) (result i32))) + (import "" "pending-once" (func $pending-once (result i32))) + (import "" "drop" (func $drop (param i32))) + (import "" "waitable-set.new" (func $waitable-set-new (result i32))) + (import "" "waitable.join" (func $waitable-join (param i32 i32))) + (import "" "waitable-set.wait" + (func $waitable-set-wait (param i32 i32) (result i32))) + (import "" "waitable-set.drop" (func $waitable-set-drop (param i32))) + (import "" "subtask.drop" (func $subtask-drop (param i32))) (import "" "task.return" (func $task-return)) (func (export "run") (param $iters i64) (result i32) loop $l @@ -990,6 +1020,49 @@ mod component_async { call $task-return i32.const 0 ) + (func (export "run-borrow") (param $iters i64) (result i32) + (local $r i32) + (local.set $r (call $new)) + loop $l + (drop (call $borrow (local.get $r))) + (local.tee $iters (i64.add (local.get $iters) (i64.const -1))) + i64.const 0 + i64.ne + br_if $l + end + + (call $drop (local.get $r)) + call $task-return + i32.const 0 + ) + (func (export "run-pending-once") (param $iters i64) + (local $ret i32) + (local $task i32) + (local $set i32) + (local.set $set (call $waitable-set-new)) + loop $l + (local.set $ret (call $pending-once)) + ;; Verify that the first poll returned Pending. + (if (i32.ne + (i32.and (local.get $ret) (i32.const 0xf)) + (i32.const 1)) + (then unreachable)) + (local.set $task + (i32.shr_u (local.get $ret) (i32.const 4))) + (call $waitable-join + (local.get $task) (local.get $set)) + (drop (call $waitable-set-wait + (local.get $set) (i32.const 0))) + (call $subtask-drop (local.get $task)) + (local.tee $iters + (i64.sub (local.get $iters) (i64.const 1))) + i64.const 0 + i64.ne + br_if $l + end + (call $waitable-set-drop (local.get $set)) + call $task-return + ) (func (export "callback") (param i32 i32 i32) (result i32) unreachable ) @@ -997,7 +1070,17 @@ mod component_async { (core func $task-return (canon task.return)) (core instance $i (instantiate $m (with "" (instance + (export "memory" (memory $memory "memory")) (export "nop" (func $nop)) + (export "new" (func $new)) + (export "borrow" (func $borrow)) + (export "pending-once" (func $pending-once)) + (export "drop" (func $drop)) + (export "waitable-set.new" (func $waitable-set-new)) + (export "waitable.join" (func $waitable-join)) + (export "waitable-set.wait" (func $waitable-set-wait)) + (export "waitable-set.drop" (func $waitable-set-drop)) + (export "subtask.drop" (func $subtask-drop)) (export "task.return" (func $task-return)) )) )) @@ -1007,22 +1090,59 @@ mod component_async { (callback (core func $i "callback")) ) ) + (func (export "run-borrow") async (param "iterations" u64) + (canon lift (core func $i "run-borrow") + async + (callback (core func $i "callback")) + ) + ) + (func (export "run-pending-once") async (param "iterations" u64) + (canon lift (core func $i "run-pending-once") async) + ) ) "#, ) .unwrap(); let mut store = Store::new(&engine, ()); let mut linker = Linker::new(&engine); + linker + .root() + .resource("r", ResourceType::host::(), |_, _| Ok(())) + .unwrap(); linker .root() .func_wrap_concurrent("nop", |_, ()| Box::pin(async { Ok(()) })) .unwrap(); + linker + .root() + .func_wrap("new", |_, ()| Ok((Resource::::new_own(0),))) + .unwrap(); + linker + .root() + .func_wrap_concurrent("borrow", |_, (_r,): (Resource,)| { + Box::pin(async { Ok(()) }) + }) + .unwrap(); + linker + .root() + .func_wrap_concurrent("pending-once", |_, ()| { + let mut pending = true; + Box::pin(std::future::poll_fn(move |cx| { + if std::mem::take(&mut pending) { + cx.waker().wake_by_ref(); + Poll::Pending + } else { + Poll::Ready(Ok(())) + } + })) + }) + .unwrap(); let instance = run_await(linker.instantiate_async(&mut store, &component)).unwrap(); let run = instance .get_typed_func::<(u64,), ()>(&mut store, "run") .unwrap(); - group.bench_function("guest-to-host", |b| { + group.bench_function("wasm-to-host", |b| { b.iter_custom(|iterations| { let start = Instant::now(); run_await(store.run_concurrent(async |accessor| { @@ -1032,6 +1152,41 @@ mod component_async { start.elapsed() }); }); + + let run_borrow = instance + .get_typed_func::<(u64,), ()>(&mut store, "run-borrow") + .unwrap(); + group.bench_function("wasm-to-host-borrow", |b| { + b.iter_custom(|iterations| { + let start = Instant::now(); + run_await(store.run_concurrent(async |accessor| { + run_borrow + .call_concurrent(accessor, (iterations,)) + .await + .unwrap() + })) + .unwrap(); + start.elapsed() + }); + }); + + let run_pending_once = instance + .get_typed_func::<(u64,), ()>(&mut store, "run-pending-once") + .unwrap(); + + group.bench_function("wasm-to-host-pending-once", |b| { + b.iter_custom(|iterations| { + let start = Instant::now(); + run_await(store.run_concurrent(async |accessor| { + run_pending_once + .call_concurrent(accessor, (iterations,)) + .await + .unwrap() + })) + .unwrap(); + start.elapsed() + }); + }); } } diff --git a/crates/wasmtime/src/runtime/component/concurrent.rs b/crates/wasmtime/src/runtime/component/concurrent.rs index b87c1403a9b4..34784e5d006f 100644 --- a/crates/wasmtime/src/runtime/component/concurrent.rs +++ b/crates/wasmtime/src/runtime/component/concurrent.rs @@ -63,7 +63,7 @@ use crate::prelude::*; use crate::store::{Store, StoreId, StoreInner, StoreOpaque, StoreToken}; #[cfg(feature = "gc")] use crate::vm::GcRootsList; -use crate::vm::component::{CallContext, ComponentInstance, InstanceState}; +use crate::vm::component::{CallContext, ComponentInstance, CurrentScope, InstanceState, Scope}; use crate::vm::{AlwaysMut, SendSyncPtr, VMFuncRef, VMLazyThread, VMMemoryDefinition, VMStore}; use crate::{ AsContext, AsContextMut, FuncType, Result, StoreContext, StoreContextMut, ValRaw, ValType, bail, @@ -857,33 +857,13 @@ pub(crate) fn poll_and_block( host_task: EnteredHostTask, future: impl Future> + Send + 'static, ) -> Result { - let task = store.current_host_thread()?; - - // 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 - // completes. - let mut future = Box::pin(async move { - let result = run_with_host_task_set(task, future).await??; - tls::get(move |store| { - let state = store.concurrent_state_mut()?; - let host_state = &mut state.get_mut(task)?.state; - assert!(matches!(host_state, HostTaskState::CalleeStarted)); - *host_state = HostTaskState::CalleeFinished(Box::new(result)); - - Waitable::Host(task).set_event( - state, - Some(Event::Subtask { - status: Status::Returned, - }), - )?; - - Ok(()) - }) - }) as HostTaskFuture; - - // 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. + // Poll the future once before creating a host task. The host task will be + // created lazily if it's needed during the poll and otherwise will be + // created if the future suspends. 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 mut future = Box::pin(future); let poll = tls::set(store, || { future .as_mut() @@ -891,19 +871,44 @@ pub(crate) fn poll_and_block( }); let caller = match host_task { - Some(pair) => pair.1, + Some(caller) => caller, None => bail_bug!("host task wasn't created but should have been"), }; - match poll { - // It completed immediately; check the result and delete the task. - Poll::Ready(result) => result?, + let task = match poll { + // It completed immediately, so no persistent host task is needed. + Poll::Ready(result) => return result, - // It did not complete immediately; add it to + // It did not complete immediately; create the host task and add it to // `ConcurrentState::futures` so it will be polled via the event loop; // then use `GuestThread::sync_call_set` to wait for the task to // complete, suspending the current fiber until it does so. Poll::Pending => { + let Some(task) = store.materialize_host_task_id()? else { + bail_bug!("current thread is not a host thread") + }; + + // Wrap the future in a closure which will stash its result in the + // host task and resume this fiber when it completes. + let future = Box::pin(async move { + let result = run_with_host_task_set(task, future).await??; + tls::get(move |store| { + let state = store.concurrent_state_mut()?; + let host_state = &mut state.get_mut(task)?.state; + assert!(matches!(host_state, HostTaskState::CalleeStarted)); + *host_state = HostTaskState::CalleeFinished(Box::new(result)); + + Waitable::Host(task).set_event( + state, + Some(Event::Subtask { + status: Status::Returned, + }), + )?; + + Ok(()) + }) + }) as HostTaskFuture; + let caller_instance = store.concurrent_state_mut()?.get_mut(caller.task)?.instance; store.switch_or_trap_if_may_not_suspend(caller_instance)?; @@ -922,8 +927,9 @@ pub(crate) fn poll_and_block( // this function returns and the task is deleted that there are no // more lingering references to this host task. Waitable::Host(task).join(store.concurrent_state_mut()?, None)?; + task } - } + }; // Retrieve and return the result. let host_state = &mut store.concurrent_state_mut()?.get_mut(task)?.state; @@ -1057,6 +1063,7 @@ impl StoreContextMut<'_, T> { assert!(state.high_priority.is_empty()); assert!(state.low_priority.is_empty()); assert!(state.unforced_current_thread.is_none()); + assert!(state.deferred_host_call_context.is_none()); assert!(state.futures_mut().unwrap().is_empty()); assert!(state.global_error_context_ref_counts.is_empty()); } @@ -1696,13 +1703,15 @@ impl StoreContextMut<'_, T> { /// Return value of [`StoreOpaque::host_task_create`]. /// /// This is an `Option` to handle the dynamic `store.concurrency_support()` -/// property, and when set this returns the host task that was created in -/// addition to the previously running guest thread. -pub type EnteredHostTask = Option<(TableId, QualifiedThreadId)>; +/// property. When present this records the guest thread to restore when the +/// host call exits. The corresponding [`HostTask`] will need to be lazily +/// created if needed via [`StoreOpaque::materialize_host_task_id`]. +pub type EnteredHostTask = Option; impl StoreOpaque { - /// Returns the currently-running thread, promoting any deferred lazy thread - /// into a fully-materialized `CurrentThread`. + /// Returns the currently-running thread, promoting any deferred lazy guest + /// thread into a fully-materialized `CurrentThread`. Deferred [`HostTask`]s + /// are not materialized. #[inline] pub(crate) fn current_thread(&mut self) -> Result { // Without concurrency support there is nothing to force. @@ -1800,13 +1809,25 @@ 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"), + // A result of `None` may indicate that this is either the top-level event + // loop, a deferred hast task, or concurrency support is disabled. In all + // cases we don't have an ID for the task. + pub(crate) fn current_materialized_host_task(&mut self) -> Result>> { + match self.current_thread()? { + CurrentThread::Host(id) => Ok(Some(id)), + CurrentThread::DeferredHost(_) | CurrentThread::None => Ok(None), + _ => bail_bug!("current thread is not a host thread"), } } + /// Returns the current host task ID, materializing a deferred host task if + /// one is active. `None` represents a call from the top-level host. + fn materialize_host_task_id(&mut self) -> Result>> { + Ok(self + .concurrent_state_mut()? + .materialize_current_host_task_id()?) + } + fn enter_sync_call(&mut self, callee: RuntimeInstance) -> Result<()> { log::trace!("enter sync-typed call {callee:?}"); let state = self.instance_state(callee).concurrent_state(); @@ -1858,6 +1879,15 @@ impl StoreOpaque { } let thread = self.current_thread()?; + let caller = if let Some(thread) = thread.guest() { + Caller::Guest { thread: *thread } + } else { + Caller::Host { + tx: None, + host_future_present: false, + caller: self.materialize_host_task_id()?, + } + }; let state = self.concurrent_state_mut()?; let guest_thread = GuestTask::new( state, @@ -1868,15 +1898,7 @@ impl StoreOpaque { memory: None, string_encoding: StringEncoding::Utf8, }, - if let Some(thread) = thread.guest() { - Caller::Guest { thread: *thread } - } else { - Caller::Host { - tx: None, - host_future_present: false, - caller: thread, - } - }, + caller, None, callee, callee_async_typed, @@ -1918,7 +1940,9 @@ impl StoreOpaque { let caller = match &task.caller { &Caller::Guest { thread } => thread.into(), - &Caller::Host { caller, .. } => caller, + &Caller::Host { caller, .. } => caller + .map(CurrentThread::Host) + .unwrap_or(CurrentThread::None), }; task.lift_result = None; task.exited = true; @@ -1948,21 +1972,18 @@ impl StoreOpaque { /// Similar to `enter_guest_sync_call` except for when the guest makes a /// transition to the host. /// - /// FIXME: this is called for all guest->host transitions and performs some - /// relatively expensive table manipulations. This would ideally be - /// optimized to avoid the full allocation of a `HostTask` in at least some - /// situations. + /// This initially records a deferred host call. A full [`HostTask`] should + /// be allocated later if needed via + /// [`StoreOpaque::materialize_host_task_id`]. pub(crate) fn host_task_create(&mut self) -> Result { if !self.concurrency_support() { self.enter_call_not_concurrent()?; return Ok(None); } let caller = self.current_guest_thread()?; - let state = self.concurrent_state_mut()?; - let task = state.push(HostTask::new(caller.task, HostTaskState::CalleeStarted))?; - log::trace!("new host task {task:?}"); - self.set_thread(task)?; - Ok(Some((task, caller))) + log::trace!("new deferred host task with caller {caller:?}"); + self.set_thread(CurrentThread::DeferredHost(caller))?; + Ok(Some(caller)) } /// Dual of `host_task_create` and signifies that the host has finished and @@ -1971,12 +1992,20 @@ 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: EnteredHostTask) -> Result<()> { - match task { - Some((task, caller)) => { + pub(crate) fn host_task_delete( + &mut self, + original_task: EnteredHostTask, + materialized_task: Option>, + ) -> Result<()> { + match original_task { + Some(caller) => { self.set_thread(caller)?; - log::trace!("delete host task {task:?}"); - self.concurrent_state_mut()?.delete(task)?; + log::trace!( + "delete host task with caller {original_task:?} and materialized as {materialized_task:?}" + ); + if let Some(task) = materialized_task { + self.concurrent_state_mut()?.delete(task)?; + } } None => { self.exit_call_not_concurrent(); @@ -2000,8 +2029,24 @@ impl StoreOpaque { fn set_thread(&mut self, thread: impl Into) -> Result { let thread = thread.into(); let state = self.concurrent_state_mut()?; + state.debug_assert_deferred_host_invariant(); let old_thread = mem::replace(&mut state.unforced_current_thread, thread); + // Ensure `deferred_host_call_context` invariant is maintained when + // switching threads and that we aren't dropping a non-empty + // `CallContext`. + if let CurrentThread::DeferredHost(_) = old_thread { + let context = state + .deferred_host_call_context + .take() + .expect("deferred host call context should be present"); + debug_assert!(context.is_empty()); + }; + if let CurrentThread::DeferredHost(_) = thread { + state.deferred_host_call_context = Some(CallContext::default()); + }; + state.debug_assert_deferred_host_invariant(); + // First thing to do after swapping threads is updating the context // slots for this thread within the store. This restores the behavior of // `context.{get,set}`. This involves taking the old state out of the @@ -2429,18 +2474,20 @@ impl StoreOpaque { /// Used by `ResourceTables` to record the scope of a borrow to get undone /// in the future. - pub(crate) fn current_scope_id(&mut self) -> Result> { + pub(crate) fn current_scope(&mut self) -> Result> { if !self.concurrency_support() { - return self.current_scope_id_not_concurrent(); + return Ok(self + .current_scope_id_not_concurrent()? + .map(|id| CurrentScope::Id(Scope::Id(id)))); } - let (bits, is_host) = match self.current_thread()? { - CurrentThread::Guest(id) => (id.task.rep(), false), - CurrentThread::GuestTask(id) => (id.rep(), false), - CurrentThread::Host(id) => (id.rep(), true), + + Ok(match self.current_thread()? { + CurrentThread::Guest(id) => Some(CurrentScope::Id(Scope::Id(id.task.rep()))), + CurrentThread::GuestTask(id) => Some(CurrentScope::Id(Scope::Id(id.rep()))), + CurrentThread::Host(id) => Some(CurrentScope::Id(Scope::HostId(id.rep()))), + CurrentThread::DeferredHost(_) => Some(CurrentScope::DeferredHost), CurrentThread::None => return Ok(None), - }; - assert_eq!((bits << 1) >> 1, bits); - Ok(Some((bits << 1) | u32::from(is_host))) + }) } fn queue_task( @@ -3369,28 +3416,25 @@ impl Instance { /// /// 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. + /// stack and linear memory. The `lower` function is invoked with the + /// `Option` param `None` if the future is cancelled. The + /// `Option>` is passed as `Some` if the host task was + /// materialized during execution and allows `lower` to delete the task if + /// needed. pub(crate) fn first_poll( self, mut store: StoreContextMut<'_, T>, host_task: EnteredHostTask, future: impl Future> + Send + 'static, - lower: impl FnOnce(StoreContextMut, Option, bool) -> Result<()> + Send + 'static, + lower: impl FnOnce(StoreContextMut, Option, bool, 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 state = &mut state.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 @@ -3407,7 +3451,10 @@ impl Instance { // It finished immediately; lower the result and delete the task. Poll::Ready(result) => { let result = result.transpose()?; - lower(store.as_context_mut(), result, true)?; + // Check if the host task was materialized so that it can be + // deleted in `lower`. + let task = store.0.current_materialized_host_task()?; + lower(store.as_context_mut(), result, true, task)?; return Ok(Status::Returned.pack(None)); } @@ -3415,6 +3462,18 @@ impl Instance { Poll::Pending => {} } + // The future will outlive this call frame, so materialize the deferred + // host task and attach its cancellation handle before publishing it to + // the event loop. + let Some(task) = store.0.materialize_host_task_id()? else { + bail_bug!("current thread is not a host thread") + }; + { + let state = &mut store.0.concurrent_state_mut()?.get_mut(task)?.state; + assert!(matches!(state, HostTaskState::CalleeStarted)); + *state = HostTaskState::CalleeRunning(join_handle); + } + // 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. @@ -3440,7 +3499,7 @@ impl Instance { Status::ReturnCancelled }; - lower(store.as_context_mut(), result, false)?; + lower(store.as_context_mut(), result, false, Some(task))?; let state = store.0.concurrent_state_mut()?; match &mut state.get_mut(task)?.state { // The task is already flagged as finished because it was @@ -3473,7 +3532,7 @@ impl Instance { // Make this task visible to the guest and then record what it // was made visible as. let caller = match host_task { - Some(pair) => pair.1, + Some(caller) => caller, None => bail_bug!("host task wasn't created but should have been"), }; let state = store.0.concurrent_state_mut()?; @@ -4883,10 +4942,10 @@ enum Caller { /// If true, there's a host future that must be dropped before the task /// can be deleted. host_future_present: bool, - /// Represents the caller of the host function which called back into a - /// guest. Note that this thread could belong to an entirely unrelated + /// The host task which called into the guest, or `None` for a call from + /// the top-level host. The task may belong to an entirely unrelated /// top-level component instance than the one the host called into. - caller: CurrentThread, + caller: Option>, }, /// Another guest thread called the guest task Guest { @@ -5521,11 +5580,15 @@ pub(crate) enum CurrentThread { Guest(QualifiedThreadId), /// The currently running thread is a host task. Host(TableId), + /// The currently running thread is a host call whose task has not yet been + /// materialized. The contained ID identifies its guest caller. + DeferredHost(QualifiedThreadId), /// A bit of a kludge to get `StoreOpaque::parent` working with backtraces /// and this serves as the parent node of a `Host` task. This ideally would /// get removed in favor of separate backtrace storage. GuestTask(TableId), - /// There is no currently running thread. + /// There is no currently running thread because we are in the main event + /// loop or concurrency is disabled. None, } @@ -5545,13 +5608,6 @@ impl CurrentThread { } } - fn host(&self) -> Option> { - match self { - Self::Host(id) => Some(*id), - _ => None, - } - } - fn is_none(&self) -> bool { matches!(self, Self::None) } @@ -5584,6 +5640,13 @@ pub struct ConcurrentState { /// be preferred over directly accessing this field. unforced_current_thread: CurrentThread, + /// Borrow state for the deferred host call, if any. + /// + /// This is `Some` if and only if [`Self::unforced_current_thread`] is + /// [`CurrentThread::DeferredHost`]. Materializing the host task moves this + /// context into that task. + deferred_host_call_context: Option, + /// The set of pending host and background tasks, if any. /// /// See `ComponentInstance::poll_until` for where we temporarily take this @@ -5662,6 +5725,7 @@ impl Default for ConcurrentState { fn default() -> Self { Self { unforced_current_thread: CurrentThread::None, + deferred_host_call_context: None, table: AlwaysMut::new(ResourceTable::new()), futures: AlwaysMut::new(Some(FuturesUnordered::new())), switch_item: None, @@ -5785,6 +5849,7 @@ impl ConcurrentState { // These fields do not contain GC references. worker_item: _, unforced_current_thread: _, + deferred_host_call_context: _, suspend_reason: _, global_error_context_ref_counts: _, interesting_tasks: _, @@ -5985,20 +6050,23 @@ impl ConcurrentState { /// Used by `ResourceTables` to acquire the current `CallContext` for the /// specified task. - /// - /// The `task` is bit-packed as returned by `current_call_context_scope_id` - /// below. - pub fn call_context(&mut self, task: u32) -> Result<&mut CallContext> { - let (task, is_host) = (task >> 1, task & 1 == 1); - if is_host { - let task: TableId = TableId::new(task); - Ok(&mut self.get_mut(task)?.call_context) - } else { - let task: TableId = TableId::new(task); - Ok(&mut self.get_mut(task)?.call_context) + pub fn call_context(&mut self, task: Scope) -> Result<&mut CallContext> { + match task { + Scope::HostId(task) => { + let task: TableId = TableId::new(task); + Ok(&mut self.get_mut(task)?.call_context) + } + Scope::Id(task) => { + let task: TableId = TableId::new(task); + Ok(&mut self.get_mut(task)?.call_context) + } } } + pub(crate) fn deferred_host_call_context(&mut self) -> Option<&mut CallContext> { + self.deferred_host_call_context.as_mut() + } + fn futures_mut(&mut self) -> Result<&mut FuturesUnordered> { match self.futures.get_mut().as_mut() { Some(f) => Ok(f), @@ -6018,14 +6086,67 @@ impl ConcurrentState { CurrentThread::Host(id) => { return Some(CurrentThread::GuestTask(self.get_mut(id).ok()?.caller)); } + CurrentThread::DeferredHost(caller) => return Some(caller.into()), CurrentThread::None => return None, }; let task = self.get_mut(task).ok()?; Some(match task.caller { - Caller::Host { caller, .. } => caller, + Caller::Host { caller, .. } => caller.map_or(CurrentThread::None, CurrentThread::Host), Caller::Guest { thread } => thread.into(), }) } + + fn debug_assert_deferred_host_invariant(&self) { + debug_assert_eq!( + self.deferred_host_call_context.is_some(), + matches!(self.unforced_current_thread, CurrentThread::DeferredHost(_)), + "a deferred host thread and call context must exist together", + ); + } + + fn materialize_host_task(&mut self) -> Result { + self.debug_assert_deferred_host_invariant(); + let caller = match self.unforced_current_thread { + CurrentThread::DeferredHost(caller) => caller, + thread => return Ok(thread), + }; + + // Push first so allocation failure leaves the deferred state intact. + let task = self.push(HostTask::new(caller.task, HostTaskState::CalleeStarted))?; + let call_context = self + .deferred_host_call_context + .take() + .expect("deferred host call context should be present"); + self.get_mut(task) + .expect("newly inserted host task should be present") + .call_context = call_context; + self.unforced_current_thread = CurrentThread::Host(task); + self.debug_assert_deferred_host_invariant(); + log::trace!("new host task materialized {task:?}"); + Ok(CurrentThread::Host(task)) + } + + fn materialize_current_host_task_id(&mut self) -> Result>> { + match self.materialize_host_task()? { + CurrentThread::Host(id) => Ok(Some(id)), + CurrentThread::None => Ok(None), + CurrentThread::Guest(_) | CurrentThread::GuestTask(_) => { + bail_bug!("tried to materialize a host task id from a guest thread") + } + CurrentThread::DeferredHost(_) => { + bail_bug!( + "current thread is a deferred host thread which should have been materialized" + ) + } + } + } + + pub(crate) fn materialize_current_scope(&mut self) -> Result { + match self.materialize_host_task()? { + CurrentThread::Host(id) => Ok(Scope::HostId(id.rep())), + _ => bail_bug!("current scope is not a deferred host scope"), + } + } } /// Provide a type hint to compiler about the shape of a parameter lower @@ -6173,6 +6294,10 @@ pub(crate) fn prepare_call( + Sync + 'static, ) -> Result> { + if !store.0.may_enter() { + bail!(Trap::CannotEnterComponent); + } + let (options, _flags, ty, raw_options) = handle.abi_info(store.0); let instance = handle.instance().id().get(store.0); @@ -6189,7 +6314,7 @@ pub(crate) fn prepare_call( .map(SendSyncPtr::new); let string_encoding = options.string_encoding; let token = StoreToken::new(store.as_context_mut()); - let caller = store.0.current_thread()?; + let caller = store.0.materialize_host_task_id()?; let state = store.0.concurrent_state_mut()?; let (tx, rx) = oneshot::channel(); @@ -6228,10 +6353,6 @@ pub(crate) fn prepare_call( async_lifted, )?; - if !store.0.may_enter() { - bail!(Trap::CannotEnterComponent); - } - Ok(PreparedCall { handle, thread, diff --git a/crates/wasmtime/src/runtime/component/concurrent_disabled.rs b/crates/wasmtime/src/runtime/component/concurrent_disabled.rs index 2ad8e26a4c74..2e7f88562781 100644 --- a/crates/wasmtime/src/runtime/component/concurrent_disabled.rs +++ b/crates/wasmtime/src/runtime/component/concurrent_disabled.rs @@ -2,6 +2,7 @@ use crate::component::func::{LiftContext, LowerContext}; use crate::component::matching::InstanceType; use crate::component::{ComponentType, Lift, Lower, RuntimeInstance, Val}; use crate::store::StoreOpaque; +use crate::vm::component::{CurrentScope, Scope}; use crate::{Result, bail, error::format_err}; use core::convert::Infallible; use core::mem::MaybeUninit; @@ -165,11 +166,17 @@ impl StoreOpaque { self.enter_call_not_concurrent() } - pub(crate) fn host_task_delete(&mut self, (): ()) -> Result<()> { + pub(crate) fn host_task_delete(&mut self, (): (), (): ()) -> Result<()> { Ok(self.exit_call_not_concurrent()) } - pub(crate) fn current_scope_id(&mut self) -> Result> { - self.current_scope_id_not_concurrent() + pub(crate) fn current_materialized_host_task(&mut self) -> Result<()> { + Ok(()) + } + + pub(crate) fn current_scope(&mut self) -> Result> { + Ok(self + .current_scope_id_not_concurrent()? + .map(|id| CurrentScope::Id(Scope::Id(id)))) } } diff --git a/crates/wasmtime/src/runtime/component/func/host.rs b/crates/wasmtime/src/runtime/component/func/host.rs index f5064e576a71..fd889f9d8a9e 100644 --- a/crates/wasmtime/src/runtime/component/func/host.rs +++ b/crates/wasmtime/src/runtime/component/func/host.rs @@ -419,7 +419,12 @@ where )?) }; lower.validate_scope_exit()?; - lower.store.0.host_task_delete(entered_host_task)?; + // Check if running the future created an actual host task in the store. + let materialized_host_task = lower.store.0.current_materialized_host_task()?; + lower + .store + .0 + .host_task_delete(entered_host_task, materialized_host_task)?; Self::lower_raw(&mut lower, ty, ret, dst) } @@ -473,7 +478,12 @@ where let result = result?; let mut lower = LowerContext::new(store, options, instance); lower.validate_scope_exit()?; - lower.store.0.host_task_delete(entered_host_task)?; + // Check if running the future created an actual host task in the store. + let materialized_host_task = lower.store.0.current_materialized_host_task()?; + lower + .store + .0 + .host_task_delete(entered_host_task, materialized_host_task)?; Self::lower_raw(&mut lower, ty, result, Destination::Memory(retptr))?; Status::Returned.pack(None) } @@ -481,11 +491,14 @@ where store.as_context_mut(), entered_host_task, future, - move |store, ret, immediate| { + move |store, ret, immediate, materialized_host_task| { let mut lower = LowerContext::new(store, options, instance); lower.validate_scope_exit()?; if immediate { - lower.store.0.host_task_delete(entered_host_task)?; + lower + .store + .0 + .host_task_delete(entered_host_task, materialized_host_task)?; } // FIXME(WebAssembly/component-model#678) the currently // running thread for this exit lower is wrong. This happens diff --git a/crates/wasmtime/src/runtime/component/func/options.rs b/crates/wasmtime/src/runtime/component/func/options.rs index ce4e96566ba3..9279b538ca7e 100644 --- a/crates/wasmtime/src/runtime/component/func/options.rs +++ b/crates/wasmtime/src/runtime/component/func/options.rs @@ -7,7 +7,7 @@ use crate::component::store::ComponentTaskState; use crate::component::{Instance, ResourceType, RuntimeInstance}; use crate::prelude::*; use crate::runtime::vm::VMFuncRef; -use crate::runtime::vm::component::{ComponentInstance, HandleTable, ResourceTables}; +use crate::runtime::vm::component::{ComponentInstance, CurrentScope, HandleTable, ResourceTables}; use crate::store::{StoreId, StoreOpaque}; use alloc::sync::Arc; use core::fmt; @@ -325,7 +325,7 @@ impl<'a, T: 'static> LowerContext<'a, T> { #[doc(hidden)] pub struct LiftContext<'a> { store_id: StoreId, - current_scope_id: Option, + current_scope: Option, /// Like lowering, lifting always has options configured. options: OptionsIndex, @@ -360,7 +360,7 @@ impl<'a> LiftContext<'a> { ) -> Result> { let store_id = store.id(); let hostcall_fuel = store.hostcall_fuel(); - let current_scope_id = store.current_scope_id()?; + let current_scope = store.current_scope()?; // 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 @@ -375,7 +375,7 @@ impl<'a> LiftContext<'a> { Ok(LiftContext { store_id, - current_scope_id, + current_scope, memory, options, types: component.types(), @@ -493,7 +493,7 @@ impl<'a> LiftContext<'a> { host_table: self.host_table, task_state: self.task_state, guest: Some(self.instance.as_mut().instance_states()), - current_scope_id: self.current_scope_id, + current_scope: self.current_scope, }, self.host_resource_data, ) diff --git a/crates/wasmtime/src/runtime/component/store.rs b/crates/wasmtime/src/runtime/component/store.rs index 2343ac5eb198..8006c4581fca 100644 --- a/crates/wasmtime/src/runtime/component/store.rs +++ b/crates/wasmtime/src/runtime/component/store.rs @@ -2,10 +2,10 @@ use crate::prelude::*; use crate::runtime::component::{HostResourceData, Instance}; use crate::runtime::vm; use crate::runtime::vm::component::{ - CallContext, ComponentInstance, HandleTable, OwnedComponentInstance, + CallContext, ComponentInstance, HandleTable, OwnedComponentInstance, Scope, }; use crate::store::{StoreData, StoreId, StoreOpaque}; -use crate::{AsContext, AsContextMut, Engine, Store, StoreContextMut}; +use crate::{AsContext, AsContextMut, Engine, Store, StoreContextMut, bail_bug}; use core::pin::Pin; use wasmtime_environ::component::RuntimeComponentInstanceIndex; use wasmtime_environ::prelude::TryPrimaryMap; @@ -402,7 +402,7 @@ impl StoreOpaque { vm::component::ResourceTables<'_>, &mut crate::component::HostResourceData, )> { - let current_scope_id = self.current_scope_id()?; + let current_scope = self.current_scope()?; let store_id = self.id(); let data = self.component_data_mut(); @@ -421,7 +421,7 @@ impl StoreOpaque { host_table: &mut data.component_host_table, task_state: &mut data.task_state, guest, - current_scope_id, + current_scope, }, &mut data.host_resource_data, )) @@ -553,14 +553,35 @@ pub struct ComponentTasksNotConcurrent { } impl ComponentTaskState { - pub fn call_context(&mut self, id: u32) -> Result<&mut CallContext> { + pub fn call_context(&mut self, id: Scope) -> Result<&mut CallContext> { match self { - ComponentTaskState::NotConcurrent(state) => Ok(&mut state.scopes[id as usize]), + ComponentTaskState::NotConcurrent(state) => match id { + Scope::Id(id) => Ok(&mut state.scopes[id as usize]), + Scope::HostId(_) => bail_bug!("non-concurrent scope cannot be a host ID"), + }, #[cfg(feature = "component-model-async")] ComponentTaskState::Concurrent(state) => state.call_context(id), } } + pub(crate) fn materialize_current_scope(&mut self) -> Result { + match self { + ComponentTaskState::NotConcurrent(_) => { + bail_bug!("a non-concurrent scope cannot be deferred") + } + #[cfg(feature = "component-model-async")] + ComponentTaskState::Concurrent(state) => state.materialize_current_scope(), + } + } + + pub(crate) fn deferred_host_call_context(&mut self) -> Option<&mut CallContext> { + match self { + ComponentTaskState::NotConcurrent(_) => None, + #[cfg(feature = "component-model-async")] + ComponentTaskState::Concurrent(state) => state.deferred_host_call_context(), + } + } + #[cfg(feature = "component-model-async")] pub fn concurrent_state_mut(&mut self) -> &mut ConcurrentState { match self { diff --git a/crates/wasmtime/src/runtime/vm/component.rs b/crates/wasmtime/src/runtime/vm/component.rs index e1599ae34f72..94b81b9e2f5f 100644 --- a/crates/wasmtime/src/runtime/vm/component.rs +++ b/crates/wasmtime/src/runtime/vm/component.rs @@ -45,7 +45,9 @@ mod resources; pub use self::handle_table::{HandleTable, RemovedResource}; #[cfg(feature = "component-model-async")] pub use self::handle_table::{ThreadHandleTable, TransmitLocalState, Waitable}; -pub use self::resources::{CallContext, ResourceTables, TypedResource, TypedResourceIndex}; +pub use self::resources::{ + CallContext, CurrentScope, ResourceTables, Scope, TypedResource, TypedResourceIndex, +}; /// Represents the state of a (sub-)component instance. #[derive(Default)] diff --git a/crates/wasmtime/src/runtime/vm/component/handle_table.rs b/crates/wasmtime/src/runtime/vm/component/handle_table.rs index d3140af9ed32..4f48010acda2 100644 --- a/crates/wasmtime/src/runtime/vm/component/handle_table.rs +++ b/crates/wasmtime/src/runtime/vm/component/handle_table.rs @@ -1,4 +1,4 @@ -use super::{TypedResource, TypedResourceIndex}; +use super::{Scope, TypedResource, TypedResourceIndex}; use crate::prelude::TryVec; use crate::{Result, bail}; use core::mem; @@ -36,7 +36,7 @@ pub enum RemovedResource { /// An `own` resource was removed with the specified `rep` Own { rep: u32 }, /// A `borrow` resource was removed originally created within `scope`. - Borrow { scope: u32 }, + Borrow { scope: Scope }, } /// Different kinds of waitables returned by [`HandleTable::waitable_rep`]. @@ -68,7 +68,7 @@ enum Slot { /// count of the `scope`. ResourceBorrow { resource: TypedResource, - scope: u32, + scope: Scope, }, /// Represents a host task handle. @@ -198,7 +198,7 @@ impl HandleTable { /// Inserts a new `borrow` resource into this table whose type/rep are /// specified by `resource`. The `scope` specified is used by /// `CallContexts` to manage lending information. - pub fn resource_borrow_insert(&mut self, resource: TypedResource, scope: u32) -> Result { + pub fn resource_borrow_insert(&mut self, resource: TypedResource, scope: Scope) -> Result { self.insert(Slot::ResourceBorrow { resource, scope }) } diff --git a/crates/wasmtime/src/runtime/vm/component/resources.rs b/crates/wasmtime/src/runtime/vm/component/resources.rs index 202b7c33520d..777a81ad9bb8 100644 --- a/crates/wasmtime/src/runtime/vm/component/resources.rs +++ b/crates/wasmtime/src/runtime/vm/component/resources.rs @@ -73,9 +73,27 @@ pub struct ResourceTables<'a> { /// as borrow counts. pub task_state: &'a mut ComponentTaskState, - /// Identifier for the current "scope" which is used for various functions - /// on `task_state` above to mutate borrows/etc of the current scope. - pub current_scope_id: Option, + /// The current scope, used to mutate borrows and lenders for the call. + pub current_scope: Option, +} + +/// The resource-borrow scope associated with the current component call. +#[derive(Debug, Copy, Clone, Eq, PartialEq)] +pub enum CurrentScope { + /// A scope which already has an ID in the component task state. + Id(Scope), + /// A host scope whose task has not yet been materialized. + #[cfg_attr(not(feature = "component-model-async"), allow(dead_code))] + DeferredHost, +} + +/// Identifier for a component call's resource-borrow scope. +#[derive(Debug, Copy, Clone, Eq, PartialEq)] +pub enum Scope { + /// A non-concurrent scope or a concurrent guest-task scope. + Id(u32), + /// A concurrent host-task scope. + HostId(u32), } /// Typed representation of a "rep" for a resource. @@ -178,6 +196,13 @@ pub struct CallContext { borrow_count: u32, } +impl CallContext { + #[cfg_attr(not(feature = "component-model-async"), allow(dead_code))] + pub(crate) fn is_empty(&self) -> bool { + self.lenders.is_empty() && self.borrow_count == 0 + } +} + impl ResourceTables<'_> { fn table_for_resource(&mut self, resource: &TypedResource) -> &mut HandleTable { match resource { @@ -266,9 +291,29 @@ impl ResourceTables<'_> { } } - fn current_scope_id(&self) -> Result { - match self.current_scope_id { - Some(id) => Ok(id), + fn materialize_current_scope(&mut self) -> Result { + let id = match self.current_scope { + Some(CurrentScope::Id(id)) => return Ok(id), + Some(CurrentScope::DeferredHost) => self.task_state.materialize_current_scope()?, + None => bail_bug!("no current scope"), + }; + self.current_scope = Some(CurrentScope::Id(id)); + Ok(id) + } + + /// Returns the current call's resource-borrow scope. + /// + /// Unlike [`Self::materialize_current_scope`], this does not materialize a + /// deferred host task. + fn current_scope(&mut self) -> Result<&mut CallContext> { + match self.current_scope { + Some(CurrentScope::Id(id)) => self.task_state.call_context(id), + Some(CurrentScope::DeferredHost) => { + match self.task_state.deferred_host_call_context() { + Some(cx) => Ok(cx), + None => bail_bug!("deferred host scope has no call context"), + } + } None => bail_bug!("no current scope"), } } @@ -285,8 +330,7 @@ impl ResourceTables<'_> { pub fn resource_lift_borrow(&mut self, index: TypedResourceIndex) -> Result { let (rep, is_own) = self.table_for_index(&index).resource_lend(index)?; if is_own { - let current = self.current_scope_id()?; - self.task_state.call_context(current)?.lenders.push(index); + self.current_scope()?.lenders.push(index); } Ok(rep) } @@ -304,7 +348,7 @@ impl ResourceTables<'_> { /// `VMComponentContext` which handles the special case of avoiding borrow /// tracking entirely. pub fn resource_lower_borrow(&mut self, resource: TypedResource) -> Result { - let scope = self.current_scope_id()?; + let scope = self.materialize_current_scope()?; let cx = self.task_state.call_context(scope)?; cx.borrow_count = cx.borrow_count.checked_add(1).unwrap(); self.table_for_resource(&resource) @@ -318,7 +362,7 @@ impl ResourceTables<'_> { /// resources that were originally passed in. #[inline] pub fn validate_scope_exit(&mut self) -> Result<()> { - let cx = self.task_state.call_context(self.current_scope_id()?)?; + let cx = self.current_scope()?; if cx.borrow_count > 0 { bail!("borrow handles still remain at the end of the call") }