From 24a4a79c087a337a21b1d53d37f651f51930d4b1 Mon Sep 17 00:00:00 2001 From: Frank Emrich Date: Thu, 3 Oct 2024 12:13:54 +0100 Subject: [PATCH] More fine-grained tracking of continuations' states (#228) This PR is part of an effort to implement more efficient effect forwarding. This particular PR makes our tracking of the state of continuations more fine-grained. Currently, the `State` enum only allows us to distinguish three cases: 1. A continuation was allocated, but never invoked (`State::Allocated`). 2. A continuation has been invoked at some point, and not returned, yet (`State::Invoked`). 3. A continuation has returned (`State::Returned`). This PR replaces `State::Invoked` with three new states, indicating more closely what's going on: - `State::Running` indicates that the continuation is the one currently executing code - `State::Parent` indicates that a continuation is suspended because it `resumed` another one. - `State::Suspended` indicates that a continuation was suspended directly (i.e., by calling `suspend`, or in the future `switch`) This PR incurrs a few additional stores per stack switch, for example because on `resume` we now need to update the state of the resumee *and* its new parent. On our usual benchmarks (in particular, `state`), this has no measurable effect. However, the two benchmarks in the `micro` subfolder see a 5-10% regression. Interestingly, this doesn't seem to be caused by the instructions responsible for the additional state updates (I experimented with *not* updating the parent states), but just some result of instructions being scheduled slightly differently after applying this PR. In any case, this regression seems acceptable anyway. This PR also extends the information we track about the main stack to carry a `State`. Otherwise we would need to check whenever we want to update the `State` of the parent if the parent is the main stack or another continuation. --- crates/continuations/src/lib.rs | 59 ++++- crates/cranelift/src/wasmfx/optimized.rs | 235 ++++++++++++------ crates/wasmtime/src/runtime/store.rs | 14 +- .../wasmtime/src/runtime/vm/continuation.rs | 61 +++-- 4 files changed, 245 insertions(+), 124 deletions(-) diff --git a/crates/continuations/src/lib.rs b/crates/continuations/src/lib.rs index fb6e01f81457..4da1e21f2e9d 100644 --- a/crates/continuations/src/lib.rs +++ b/crates/continuations/src/lib.rs @@ -71,6 +71,27 @@ pub struct StackLimits { pub last_wasm_entry_sp: usize, } +/// This type represents "common" information that we need to save both for the +/// main stack and each continuation. +#[repr(C)] +#[derive(Debug, Clone)] +pub struct CommonStackInformation { + pub limits: StackLimits, + /// For the main stack, this field must only have one of the following values: + /// - Running + /// - Parent + pub state: State, +} + +impl CommonStackInformation { + pub fn running_default() -> Self { + Self { + limits: StackLimits::default(), + state: State::Running, + } + } +} + impl StackLimits { pub fn with_stack_limit(stack_limit: usize) -> Self { Self { @@ -171,12 +192,17 @@ pub const STACK_CHAIN_CONTINUATION_DISCRIMINANT: usize = 2; pub enum State { /// The `VMContRef` has been created, but `resume` has never been /// called on it. During this stage, we may add arguments using `cont.bind`. - Allocated, - /// `resume` has been invoked at least once on the `VMContRef`, - /// meaning that the function passed to `cont.new` has started executing. - /// Note that this state does not indicate whether the execution of this - /// function is currently suspended or not. - Invoked, + Fresh, + /// The continuation is running, meaning that it is the one currently + /// executing code. + Running, + /// The continuation is suspended because it executed a resume instruction + /// that has not finished yet. In other words, it became the parent of + /// another continuation (which may itself be `Running`, a `Parent`, or + /// `Suspended`). + Parent, + /// The continuation was suspended by a `suspend` instruction. + Suspended, /// The function originally passed to `cont.new` has returned normally. /// Note that there is no guarantee that a VMContRef will ever /// reach this status, as it may stay suspended until being dropped. @@ -214,22 +240,21 @@ pub mod offsets { /// Offsets of fields in `wasmtime_runtime::continuation::VMContRef`. /// We uses tests there to ensure these values are correct. pub mod vm_cont_ref { - use crate::Payloads; + use crate::{CommonStackInformation, Payloads}; - /// Offset of `limits` field - pub const LIMITS: usize = 0; + /// Offset of `common_stack_information` field + pub const COMMON_STACK_INFORMATION: usize = 0; /// Offset of `parent_chain` field - pub const PARENT_CHAIN: usize = LIMITS + 4 * core::mem::size_of::(); + pub const PARENT_CHAIN: usize = + COMMON_STACK_INFORMATION + core::mem::size_of::(); /// Offset of `stack` field pub const STACK: usize = PARENT_CHAIN + 2 * core::mem::size_of::(); /// Offset of `args` field pub const ARGS: usize = STACK + super::FIBER_STACK_SIZE; /// Offset of `tag_return_values` field pub const TAG_RETURN_VALUES: usize = ARGS + core::mem::size_of::(); - /// Offset of `state` field - pub const STATE: usize = TAG_RETURN_VALUES + core::mem::size_of::(); /// Offset of `revision` field - pub const REVISION: usize = STATE + core::mem::size_of::(); + pub const REVISION: usize = TAG_RETURN_VALUES + core::mem::size_of::(); } pub mod stack_limits { @@ -242,6 +267,14 @@ pub mod offsets { pub const LAST_WASM_ENTRY_SP: usize = offset_of!(StackLimits, last_wasm_entry_sp); } + pub mod common_stack_information { + use crate::CommonStackInformation; + use memoffset::offset_of; + + pub const LIMITS: usize = offset_of!(CommonStackInformation, limits); + pub const STATE: usize = offset_of!(CommonStackInformation, state); + } + /// Size of wasmtime_runtime::continuation::FiberStack. /// We test there that this value is correct. pub const FIBER_STACK_SIZE: usize = 3 * core::mem::size_of::(); diff --git a/crates/cranelift/src/wasmfx/optimized.rs b/crates/cranelift/src/wasmfx/optimized.rs index bb97e4f2ca28..b46436604e3d 100644 --- a/crates/cranelift/src/wasmfx/optimized.rs +++ b/crates/cranelift/src/wasmfx/optimized.rs @@ -322,6 +322,10 @@ pub(crate) mod typed_continuation_helpers { pointer_type: ir::Type, } + pub struct CommonStackInformation { + pub address: ir::Value, + } + /// Compile-time representation of `crate::runtime::vm::fibre::FiberStack`. pub struct FiberStack { /// This is NOT the "top of stack" address of the stack itself. In line @@ -348,58 +352,14 @@ pub(crate) mod typed_continuation_helpers { Payloads::new(self.address, offset as i32, self.pointer_type) } - /// Loads the value of the `state` field of the VMContRef, - /// which is represented using the `State` enum. - fn load_state(&self, builder: &mut FunctionBuilder) -> ir::Value { - let mem_flags = ir::MemFlags::trusted(); - let offset = wasmtime_continuations::offsets::vm_cont_ref::STATE as i32; - - // Let's make sure that we still represent the State enum as i32. - debug_assert!(mem::size_of::() == mem::size_of::()); - - builder.ins().load(I32, mem_flags, self.address, offset) - } - - /// Sets the value of the `state` field of the `VMContRef`, - pub fn set_state( + pub fn common_stack_information<'a>( &self, + _env: &mut crate::func_environ::FuncEnvironment<'a>, builder: &mut FunctionBuilder, - state: wasmtime_continuations::State, - ) { - let mem_flags = ir::MemFlags::trusted(); - let offset = wasmtime_continuations::offsets::vm_cont_ref::STATE as i32; - - // Let's make sure that we still represent the State enum as i32. - debug_assert!(mem::size_of::() == mem::size_of::()); - - let v = builder.ins().iconst(I32, state.discriminant() as i64); - builder.ins().store(mem_flags, v, self.address, offset); - } - - /// Checks whether the `VMContRef` is invoked (i.e., `resume` - /// was called at least once on the continuation). - pub fn is_invoked(&self, builder: &mut FunctionBuilder) -> ir::Value { - // TODO(frank-emrich) In the future, we may get rid of the State field - // in `VMContRef` and try to infer the state by other means. - // For example, we may alllocate the `ContinuationFiber` lazily, doing - // so only at the point when a continuation is actualy invoked, meaning - // that we can use the null-ness of the `fiber` field as an indicator - // for invokedness. - let actual_state = self.load_state(builder); - let invoked: i32 = i32::from(wasmtime_continuations::State::Invoked); - builder - .ins() - .icmp_imm(IntCC::Equal, actual_state, invoked as i64) - } - - /// Checks whether the `VMContRef` has returned (i.e., the - /// function used as continuation has returned normally). - pub fn has_returned(&self, builder: &mut FunctionBuilder) -> ir::Value { - let actual_state = self.load_state(builder); - let returned: i32 = i32::from(wasmtime_continuations::State::Returned); - builder - .ins() - .icmp_imm(IntCC::Equal, actual_state, returned as i64) + ) -> CommonStackInformation { + let offset = wasmtime_continuations::offsets::vm_cont_ref::COMMON_STACK_INFORMATION; + let address = builder.ins().iadd_imm(self.address, offset as i64); + CommonStackInformation { address } } /// Returns pointer to buffer where results are stored after a @@ -410,7 +370,11 @@ pub(crate) mod typed_continuation_helpers { builder: &mut FunctionBuilder, ) -> ir::Value { if cfg!(debug_assertions) { - let has_returned = self.has_returned(builder); + let has_returned = self.common_stack_information(env, builder).has_state( + env, + builder, + wasmtime_continuations::State::Returned, + ); emit_debug_assert!(env, builder, has_returned); } return self.args().get_data(builder); @@ -1022,15 +986,13 @@ pub(crate) mod typed_continuation_helpers { } /// Must only be called if `self` represents a `MainStack` or `Continuation` variant. - /// Returns a pointer to the associated `StackLimits` object (i.e., in - /// the former case, the pointer directly stored in the variant, or in - /// the latter case a pointer to the `StackLimits` data within the - /// `VMContRef`. - pub fn get_stack_limits_ptr<'a>( + /// Returns a pointer to the associated `CommonStackInformation` object either stored in + /// the `MainStackInfo` object (if `MainStack`) or the `VMContRef` (if `Continuation`) + pub fn get_common_stack_information<'a>( &self, env: &mut crate::func_environ::FuncEnvironment<'a>, builder: &mut FunctionBuilder, - ) -> ir::Value { + ) -> CommonStackInformation { use wasmtime_continuations::offsets as o; self.assert_not_absent(env, builder); @@ -1038,24 +1000,115 @@ pub(crate) mod typed_continuation_helpers { // `self` corresponds to a StackChain::MainStack or // StackChain::Continuation. // In both cases, the payload is a pointer. - let ptr = self.payload; + let address = self.payload; // `obj` is now a pointer to the beginning of either // 1. A `VMContRef` struct (in the case of a // StackChain::Continuation) - // 2. A StackLimits struct (in the case of + // 2. A CommonStackInformation struct (in the case of // StackChain::MainStack) // - // Since a `VMContRef` starts with an (inlined) StackLimits + // Since a `VMContRef` starts with an (inlined) CommonStackInformation // object at offset 0, we actually have in both cases that `ptr` is // now the address of the beginning of a StackLimits object. - debug_assert_eq!(o::vm_cont_ref::LIMITS, 0); - ptr + debug_assert_eq!(o::vm_cont_ref::COMMON_STACK_INFORMATION, 0); + CommonStackInformation { address } + } + } + + impl CommonStackInformation { + fn get_state_ptr<'a>( + &self, + _env: &mut crate::func_environ::FuncEnvironment<'a>, + builder: &mut FunctionBuilder, + ) -> ir::Value { + use wasmtime_continuations::offsets as o; + + builder + .ins() + .iadd_imm(self.address, o::common_stack_information::STATE as i64) + } + + fn get_stack_limits_ptr<'a>( + &self, + _env: &mut crate::func_environ::FuncEnvironment<'a>, + builder: &mut FunctionBuilder, + ) -> ir::Value { + use wasmtime_continuations::offsets as o; + + builder + .ins() + .iadd_imm(self.address, o::common_stack_information::LIMITS as i64) + } + + fn load_state<'a>( + &self, + env: &mut crate::func_environ::FuncEnvironment<'a>, + builder: &mut FunctionBuilder, + ) -> ir::Value { + // Let's make sure that we still represent the State enum as i32. + debug_assert!(mem::size_of::() == mem::size_of::()); + + let mem_flags = ir::MemFlags::trusted(); + let state_ptr = self.get_state_ptr(env, builder); + + builder.ins().load(I32, mem_flags, state_ptr, 0) + } + + pub fn set_state<'a>( + &self, + env: &mut crate::func_environ::FuncEnvironment<'a>, + builder: &mut FunctionBuilder, + state: wasmtime_continuations::State, + ) { + // Let's make sure that we still represent the State enum as i32. + debug_assert!(mem::size_of::() == mem::size_of::()); + + let discriminant = builder.ins().iconst(I32, state.discriminant() as i64); + emit_debug_println!( + env, + builder, + "setting state of CommonStackInformation {:p} to {}", + self.address, + discriminant + ); + + let mem_flags = ir::MemFlags::trusted(); + let state_ptr = self.get_state_ptr(env, builder); + + builder.ins().store(mem_flags, discriminant, state_ptr, 0); + } + + pub fn has_state<'a>( + &self, + env: &mut crate::func_environ::FuncEnvironment<'a>, + builder: &mut FunctionBuilder, + state: wasmtime_continuations::State, + ) -> ir::Value { + let actual_state = self.load_state(env, builder); + + builder + .ins() + .icmp_imm(IntCC::Equal, actual_state, state.discriminant() as i64) + } + + /// Checks whether the `State` reflects that the stack has ever been + /// active (instead of just having been allocated, but never resumed). + pub fn was_invoked<'a>( + &self, + env: &mut crate::func_environ::FuncEnvironment<'a>, + builder: &mut FunctionBuilder, + ) -> ir::Value { + let actual_state = self.load_state(env, builder); + let allocated: i32 = i32::from(wasmtime_continuations::State::Fresh); + builder + .ins() + .icmp_imm(IntCC::NotEqual, actual_state, allocated as i64) } /// Sets `last_wasm_entry_sp` and `stack_limit` fields in - /// `VMRuntimelimits` using the values from the `StackLimits` object - /// associated with this stack chain. + /// `VMRuntimelimits` using the values from the `StackLimits` of this + /// object. pub fn write_limits_to_vmcontext<'a>( &self, env: &mut crate::func_environ::FuncEnvironment<'a>, @@ -1070,7 +1123,7 @@ pub(crate) mod typed_continuation_helpers { let mut copy_to_vm_runtime_limits = |our_offset, their_offset| { let our_value = builder.ins().load( - self.pointer_type, + env.pointer_type(), memflags, stack_limits_ptr, our_offset as i32, @@ -1083,7 +1136,7 @@ pub(crate) mod typed_continuation_helpers { ); }; - let pointer_size = self.pointer_type.bytes() as u8; + let pointer_size = env.pointer_type().bytes() as u8; copy_to_vm_runtime_limits( o::stack_limits::STACK_LIMIT, pointer_size.vmruntime_limits_stack_limit(), @@ -1095,7 +1148,7 @@ pub(crate) mod typed_continuation_helpers { } /// Overwrites the `last_wasm_entry_sp` field of the `StackLimits` - /// object associated with this stack chain by loading the corresponding + /// object in the `StackLimits` of this object by loading the corresponding /// field from the `VMRuntimeLimits`. /// If `load_stack_limit` is true, we do the same for the `stack_limit` /// field. @@ -1115,11 +1168,11 @@ pub(crate) mod typed_continuation_helpers { let stack_limits_ptr = self.get_stack_limits_ptr(env, builder); let memflags = ir::MemFlags::trusted(); - let pointer_size = self.pointer_type.bytes() as u8; + let pointer_size = env.pointer_type().bytes() as u8; let mut copy = |runtime_limits_offset, stack_limits_offset| { let from_vm_runtime_limits = builder.ins().load( - self.pointer_type, + env.pointer_type(), memflags, vmruntime_limits_ptr, runtime_limits_offset, @@ -1314,10 +1367,11 @@ pub(crate) fn typed_continuations_store_resume_args<'a>( builder.append_block_param(store_data_block, env.pointer_type()); let co = tc::VMContRef::new(contref, env.pointer_type()); - let is_invoked = co.is_invoked(builder); + let csi = co.common_stack_information(env, builder); + let was_invoked = csi.was_invoked(env, builder); builder .ins() - .brif(is_invoked, use_payloads_block, &[], use_args_block, &[]); + .brif(was_invoked, use_payloads_block, &[], use_args_block, &[]); { builder.switch_to_block(use_args_block); @@ -1543,7 +1597,8 @@ pub(crate) fn translate_resume<'a>( // We keep this check mostly for the test that runs a continuation // twice with `unsafe_disable_continuation_linearity_check` enabled. let zero = builder.ins().iconst(I8, 0); - let has_returned = vmcontref.has_returned(builder); + let csi = vmcontref.common_stack_information(env, builder); + let has_returned = csi.has_state(env, builder, wasmtime_continuations::State::Returned); emit_debug_assert_eq!(env, builder, has_returned, zero); } @@ -1589,9 +1644,12 @@ pub(crate) fn translate_resume<'a>( // description of the invariants that we maintain for the various stack // limits. - // We mark `resume_contref` to be invoked + // `resume_contref` is now active, and its parent is suspended. let co = tc::VMContRef::new(resume_contref, env.pointer_type()); - co.set_state(builder, wasmtime_continuations::State::Invoked); + let resume_csi = co.common_stack_information(env, builder); + let parent_csi = parent_stack_chain.get_common_stack_information(env, builder); + resume_csi.set_state(env, builder, wasmtime_continuations::State::Running); + parent_csi.set_state(env, builder, wasmtime_continuations::State::Parent); // We update the `StackLimits` of the parent of the continuation to be resumed // as well as the `VMRuntimeLimits`. @@ -1612,7 +1670,7 @@ pub(crate) fn translate_resume<'a>( let vm_runtime_limits_ptr = vmctx.load_vm_runtime_limits_ptr(env, builder); let last_wasm_exit_fp = builder.ins().get_frame_pointer(env.pointer_type()); let last_wasm_exit_pc = builder.ins().get_instruction_pointer(env.pointer_type()); - parent_stack_chain.load_limits_from_vmcontext( + parent_csi.load_limits_from_vmcontext( env, builder, vm_runtime_limits_ptr, @@ -1620,9 +1678,7 @@ pub(crate) fn translate_resume<'a>( Some(last_wasm_exit_fp), Some(last_wasm_exit_pc), ); - let resume_stackchain = - tc::StackChain::from_continuation(builder, resume_contref, env.pointer_type()); - resume_stackchain.write_limits_to_vmcontext(env, builder, vm_runtime_limits_ptr); + resume_csi.write_limits_to_vmcontext(env, builder, vm_runtime_limits_ptr); let fiber_stack = co.get_fiber_stack(env, builder); let control_context_ptr = fiber_stack.load_control_context(env, builder); @@ -1651,6 +1707,8 @@ pub(crate) fn translate_resume<'a>( // Now the parent contref (or main stack) is active again vmctx.store_stack_chain(env, builder, &parent_stack_chain); + parent_csi.set_state(env, builder, wasmtime_continuations::State::Running); + resume_csi.set_state(env, builder, wasmtime_continuations::State::Suspended); // Extract the result and signal bit. let result = ControlEffect::new(result); @@ -1683,7 +1741,9 @@ pub(crate) fn translate_resume<'a>( // We store parts of the VMRuntimeLimits into the continuation that just suspended. let suspended_chain = tc::StackChain::from_continuation(builder, resume_contref, env.pointer_type()); - suspended_chain.load_limits_from_vmcontext( + let suspended_csi = suspended_chain.get_common_stack_information(env, builder); + let parent_csi = parent_stack_chain.get_common_stack_information(env, builder); + suspended_csi.load_limits_from_vmcontext( env, builder, vm_runtime_limits_ptr, @@ -1694,7 +1754,7 @@ pub(crate) fn translate_resume<'a>( // Afterwards (!), restore parts of the VMRuntimeLimits from the // parent of the suspended continuation (which is now active). - parent_stack_chain.write_limits_to_vmcontext(env, builder, vm_runtime_limits_ptr); + parent_csi.write_limits_to_vmcontext(env, builder, vm_runtime_limits_ptr); // Extract the tag let tag = ControlEffect::value(resume_result, env, builder); @@ -1830,10 +1890,12 @@ pub(crate) fn translate_resume<'a>( // Restore parts of the VMRuntimeLimits from the parent of the // returned continuation (which is now active). - parent_stack_chain.write_limits_to_vmcontext(env, builder, vm_runtime_limits_ptr); + let parent_csi = parent_stack_chain.get_common_stack_information(env, builder); + parent_csi.write_limits_to_vmcontext(env, builder, vm_runtime_limits_ptr); let co = tc::VMContRef::new(resume_contref, env.pointer_type()); - co.set_state(builder, wasmtime_continuations::State::Returned); + let resume_csi = co.common_stack_information(env, builder); + resume_csi.set_state(env, builder, wasmtime_continuations::State::Returned); // Load and push the results. let returns = env.continuation_returns(type_index).to_vec(); @@ -1872,6 +1934,15 @@ pub(crate) fn translate_suspend<'a>( let suspend_payload = ControlEffect::make_suspend(env, builder, tag_addr).0 .0; + if cfg!(debug_assertions) { + let is_running = cref.common_stack_information(env, builder).has_state( + env, + builder, + wasmtime_continuations::State::Running, + ); + emit_debug_assert!(env, builder, is_running); + } + builder .ins() .stack_switch(control_context_ptr, control_context_ptr, suspend_payload); diff --git a/crates/wasmtime/src/runtime/store.rs b/crates/wasmtime/src/runtime/store.rs index b1b306cdadcc..b92e7388d7fa 100644 --- a/crates/wasmtime/src/runtime/store.rs +++ b/crates/wasmtime/src/runtime/store.rs @@ -80,7 +80,7 @@ use crate::instance::InstanceData; use crate::linker::Definition; use crate::module::RegisteredModuleId; use crate::prelude::*; -use crate::runtime::vm::continuation::stack_chain::{StackChain, StackChainCell, StackLimits}; +use crate::runtime::vm::continuation::stack_chain::{StackChain, StackChainCell}; use crate::runtime::vm::mpk::{self, ProtectionKey, ProtectionMask}; use crate::runtime::vm::{ Backtrace, ExportGlobal, GcHeapAllocationIndex, GcRootsList, GcStore, @@ -104,7 +104,7 @@ use core::ops::{Deref, DerefMut}; use core::pin::Pin; use core::ptr; use core::task::{Context, Poll}; -use wasmtime_continuations::WasmFXConfig; +use wasmtime_continuations::{CommonStackInformation, WasmFXConfig}; mod context; pub use self::context::*; @@ -325,7 +325,7 @@ pub struct StoreOpaque { // Finally, observe that the stack chain adds more internal self references: // The stack chain always contains a `MainStack` element at the ends which // has a pointer to the `main_stack_limits` field of the same `StoreOpaque`. - main_stack_limits: StackLimits, + main_stack_information: CommonStackInformation, stack_chain: StackChainCell, instances: Vec, @@ -533,7 +533,7 @@ impl Store { _marker: marker::PhantomPinned, engine: engine.clone(), runtime_limits: Default::default(), - main_stack_limits: Default::default(), + main_stack_information: CommonStackInformation::running_default(), stack_chain: StackChainCell::absent(), instances: Vec::new(), #[cfg(feature = "component-model")] @@ -625,7 +625,7 @@ impl Store { // `Store` indicates that `inner` is supposed to be at a stable // location at this point, without explicitly being `Pin`-ed. let stack_chain = inner.stack_chain.0.get(); - *stack_chain = StackChain::MainStack(inner.main_stack_limits()); + *stack_chain = StackChain::MainStack(inner.main_stack_information()); } Self { @@ -1905,10 +1905,10 @@ impl StoreOpaque { } #[inline] - pub fn main_stack_limits(&self) -> *mut StackLimits { + pub fn main_stack_information(&self) -> *mut CommonStackInformation { // NOTE(frank-emrich) This looks dogdy, but follows the same pattern as // `vmruntime_limits()` above. - &self.main_stack_limits as *const StackLimits as *mut StackLimits + &self.main_stack_information as *const CommonStackInformation as *mut CommonStackInformation } #[inline] diff --git a/crates/wasmtime/src/runtime/vm/continuation.rs b/crates/wasmtime/src/runtime/vm/continuation.rs index dfa799b0793b..585a4af140ab 100644 --- a/crates/wasmtime/src/runtime/vm/continuation.rs +++ b/crates/wasmtime/src/runtime/vm/continuation.rs @@ -80,7 +80,7 @@ pub mod optimized { use core::cmp; use core::mem; #[allow(unused)] - use wasmtime_continuations::{debug_println, ENABLE_DEBUG_PRINTING}; + use wasmtime_continuations::{debug_println, CommonStackInformation, ENABLE_DEBUG_PRINTING}; pub use wasmtime_continuations::{Payloads, StackLimits, State}; /// Fibers used for continuations @@ -89,8 +89,8 @@ pub mod optimized { /// TODO #[repr(C)] pub struct VMContRef { - /// The limits of this continuation's stack. - pub limits: StackLimits, + /// The `CommonStackInformation` of this continuation's stack. + pub common_stack_information: CommonStackInformation, /// The parent of this continuation, which may be another continuation, the /// main stack, or absent (in case of a suspended continuation). @@ -110,9 +110,6 @@ pub mod optimized { /// In particular, this may only be Some when `state` is `Invoked`. pub tag_return_values: Payloads, - /// Indicates the state of this continuation. - pub state: State, - /// Revision counter. pub revision: u64, } @@ -130,20 +127,20 @@ pub mod optimized { /// so. Used to create `VMContRef`s when initializing pooling allocator. pub fn empty() -> Self { let limits = StackLimits::with_stack_limit(Default::default()); + let state = State::Fresh; + let common_stack_information = CommonStackInformation { limits, state }; let parent_chain = StackChain::Absent; let stack = FiberStack::unallocated(); let args = Payloads::new(0); let tag_return_values = Payloads::new(0); - let state = State::Allocated; let revision = 0; Self { - limits, + common_stack_information, parent_chain, stack, args, tag_return_values, - state, revision, } } @@ -190,9 +187,9 @@ pub mod optimized { )) })? }; - assert!(parent.state == State::Invoked); - assert!(child.state == State::Invoked); - assert!(child.tag_return_values.length == 0); + debug_assert!(parent.common_stack_information.state == State::Running); + debug_assert!(child.common_stack_information.state == State::Suspended); + debug_assert!(child.tag_return_values.length == 0); mem::swap(&mut child.tag_return_values, &mut parent.tag_return_values); Ok(()) @@ -216,7 +213,7 @@ pub mod optimized { { let contref = unsafe { contref.as_mut().unwrap() }; // A continuation must have run to completion before dropping it. - assert!(contref.state == State::Returned); + debug_assert!(contref.common_stack_information.state == State::Returned); // Note that we *could* deallocate the `Payloads` (i.e., `args` and // `tag_return_values`) here, but choose not to: @@ -265,9 +262,10 @@ pub mod optimized { { let contref = unsafe { contref.as_mut().unwrap() }; - contref.limits = limits; + let csi = &mut contref.common_stack_information; + csi.limits = limits; + csi.state = State::Fresh; contref.parent_chain = StackChain::Absent; - contref.state = State::Allocated; contref.args.ensure_capacity(capacity); // In order to give the pool a uniform interface for the optimized @@ -300,7 +298,10 @@ pub mod optimized { use core::mem::offset_of; use wasmtime_continuations::offsets::*; - assert_eq!(offset_of!(VMContRef, limits), vm_cont_ref::LIMITS); + assert_eq!( + offset_of!(VMContRef, common_stack_information), + vm_cont_ref::COMMON_STACK_INFORMATION + ); assert_eq!( offset_of!(VMContRef, parent_chain), vm_cont_ref::PARENT_CHAIN @@ -311,12 +312,14 @@ pub mod optimized { offset_of!(VMContRef, tag_return_values), vm_cont_ref::TAG_RETURN_VALUES ); - assert_eq!(offset_of!(VMContRef, state), vm_cont_ref::STATE); + + assert_eq!(offset_of!(VMContRef, revision), vm_cont_ref::REVISION); assert_eq!(core::mem::size_of::(), FIBER_STACK_SIZE); assert_eq!(core::mem::size_of::(), STACK_CHAIN_SIZE); - assert_eq!(offset_of!(VMContRef, revision), vm_cont_ref::REVISION); + // `CommonStackInformation` and `StackLimits` offsets don't need tests because + // they are defined diretly with `offset_of!` } } @@ -687,6 +690,7 @@ pub mod baseline { pub mod stack_chain { use super::imp::VMContRef; use core::cell::UnsafeCell; + use wasmtime_continuations::CommonStackInformation; pub use wasmtime_continuations::StackLimits; /// This type represents a linked lists of stacks, additionally associating a @@ -758,7 +762,8 @@ pub mod stack_chain { /// field, means that there is currently no parent. Absent = wasmtime_continuations::STACK_CHAIN_ABSENT_DISCRIMINANT, /// Represents the main stack. - MainStack(*mut StackLimits) = wasmtime_continuations::STACK_CHAIN_MAIN_STACK_DISCRIMINANT, + MainStack(*mut CommonStackInformation) = + wasmtime_continuations::STACK_CHAIN_MAIN_STACK_DISCRIMINANT, /// Represents a continuation's stack. Continuation(*mut VMContRef) = wasmtime_continuations::STACK_CHAIN_CONTINUATION_DISCRIMINANT, @@ -788,27 +793,39 @@ pub mod stack_chain { /// Each stack is represented by a tuple `(co_opt, sl)`, where sl is a pointer /// to the stack's `StackLimits` object and `co_opt` is a pointer to the /// corresponding `VMContRef`, or None for the main stack. + #[cfg_attr(feature = "wasmfx_baseline", allow(dead_code))] pub struct ContinuationChainIterator(StackChain); impl Iterator for ContinuationChainIterator { type Item = (Option<*mut VMContRef>, *mut StackLimits); + #[cfg(not(feature = "wasmfx_baseline"))] fn next(&mut self) -> Option { match self.0 { StackChain::Absent => None, - StackChain::MainStack(ms) => { - let next = (None, ms); + StackChain::MainStack(csi) => { + let stack_limits = unsafe { &mut (*csi).limits } as *mut StackLimits; + + let next = (None, stack_limits); self.0 = StackChain::Absent; Some(next) } StackChain::Continuation(ptr) => { let continuation = unsafe { ptr.as_mut().unwrap() }; - let next = (Some(ptr), (&mut continuation.limits) as *mut StackLimits); + let next = ( + Some(ptr), + (&mut continuation.common_stack_information.limits) as *mut StackLimits, + ); self.0 = continuation.parent_chain.clone(); Some(next) } } } + + #[cfg(feature = "wasmfx_baseline")] + fn next(&mut self) -> Option { + unimplemented!() + } } #[repr(transparent)]