From ceb996791cc81d390f2cbc0362a6f3d7fffb99f1 Mon Sep 17 00:00:00 2001 From: Frank Emrich Date: Wed, 14 Feb 2024 17:33:21 +0000 Subject: [PATCH] Backtrace support, part 1: Use StackChain to represent chains of stacks (#98) This is the first PR in a series to add support for generating backtraces in the presence of stack switching. Currently, the `VMContext` contains a (possibly `null`) pointer to the currently running continuation. Every `ContinuationObject` in turn contains a (possibly `null`) pointer to its parent. This effectively yields a linked list of continuations. This PR changes this situation in two ways. See the comment on `wasmtime_continuations::StackChain` for more details. 1. Instead of a (possibly `null`) pointer to a `ContinuationObject`), both `VMContext` and `ContinuationObject` now contain an object of type `StackChain` to represent their "parent". This is still effectively a linked list of stacks. 2. However, by using the new `StackChain` type, it is now possible to have an explicit representation of the main stack at the end of the list of active continuations. In other words, a `StackChain` linked list now ends with a `MainStack` variant. In addition, we now associate a `StackLimits` object with each element in the `StackChain` (including the main stack). This will be used in subsequent PRs to store a subset of the values in `VMRuntimeLimits` for each stack. --- cranelift/wasm/src/environ/spec.rs | 3 +- crates/continuations/src/lib.rs | 101 ++++++- crates/cranelift/src/wasmfx/optimized.rs | 335 +++++++++++++++++------ crates/environ/src/vmoffsets.rs | 41 ++- crates/runtime/src/continuation.rs | 79 ++++-- crates/runtime/src/instance.rs | 34 ++- tests/all/pooling_allocator.rs | 14 +- 7 files changed, 461 insertions(+), 146 deletions(-) diff --git a/cranelift/wasm/src/environ/spec.rs b/cranelift/wasm/src/environ/spec.rs index fea51ca81747..6faf7b522cd3 100644 --- a/cranelift/wasm/src/environ/spec.rs +++ b/cranelift/wasm/src/environ/spec.rs @@ -699,7 +699,8 @@ pub trait FuncEnvironment: TargetEnvironment { /// TODO fn tag_returns(&self, tag_index: u32) -> &[wasmtime_types::WasmValType]; - /// TODO + /// Returns a pointer to the currently running continuation object. + /// Traps if not currently running inside a continuation. fn typed_continuations_load_continuation_object( &mut self, builder: &mut FunctionBuilder, diff --git a/crates/continuations/src/lib.rs b/crates/continuations/src/lib.rs index 48d5ebf98965..7f28e8e7d073 100644 --- a/crates/continuations/src/lib.rs +++ b/crates/continuations/src/lib.rs @@ -30,6 +30,87 @@ pub mod types { pub type ContinuationFiber = Fiber<'static, (), u32, ()>; +/// This type is used to save (and subsequently restore) a subset of the data in +/// `VMRuntimeLimits`. See documentation of `StackChain` for the exact uses. +#[repr(C)] +#[derive(Debug, Clone)] +pub struct StackLimits { + pub stack_limit: usize, + pub last_wasm_exit_fp: usize, + pub last_wasm_exit_pc: usize, + pub last_wasm_entry_sp: usize, +} + +impl Default for StackLimits { + fn default() -> Self { + Self { + stack_limit: 0, + last_wasm_exit_fp: 0, + last_wasm_exit_pc: 0, + last_wasm_entry_sp: 0, + } + } +} + +impl StackLimits { + pub fn with_stack_limit(stack_limit: usize) -> Self { + Self { + stack_limit, + ..Default::default() + } + } +} + +const STACK_CHAIN_ABSENT_DISCRIMINANT: usize = 0; +const STACK_CHAIN_MAIN_STACK_DISCRIMINANT: usize = 1; +const STACK_CHAIN_CONTINUATION_DISCRIMINANT: usize = 2; + +/// This type represents a linked lists of stacks, additionally associating a +/// `StackLimits` object with each element of the list. Here, a "stack" is +/// either a continuation or the main stack. Note that the linked list character +/// arises from the fact that `StackChain::Continuation` variants have a pointer +/// to have `ContinuationObject`, which in turn has a `parent_chain` value of +/// type `StackChain`. +/// +/// There are generally two uses of such chains: +/// +/// 1. The `typed_continuations_chain` field in the VMContext contains such a +/// chain of stacks, where the head of the list denotes the stack that is +/// currently executing (either a continuation or the main stack), as well as +/// the parent stacks, in case of a continuation currently running. Note that in +/// this case, the linked list must contains 0 or more `Continuation` elements, +/// followed by a final `MainStack` element. In particular, this list always +/// ends with `MainStack` and never contains an `Absent` variant. +/// +/// 2. When a continuation is suspended, its chain of parents eventually ends +/// with an `Absent` variant in its `parent_chain` field. Note that a suspended +/// continuation never appears in the stack chain in the VMContext! +/// +/// +/// As mentioned before, each stack in a `StackChain` has a corresponding +/// `StackLimits` object. For continuations, this is stored in the `limits` +/// fields of the corresponding `ContinuationObject`. For the main stack, the +/// `MainStack` variant contains a pointer to the +/// `typed_continuations_main_stack_limits` field of the VMContext. +// FIXME(frank-emrich) Note that the data within the StackLimits objects is +// currently not used or updated in any way. +#[derive(Debug, Clone, PartialEq)] +#[repr(usize, C)] +pub enum StackChain { + /// If stored in the VMContext, used to indicate that the MainStack entry + /// has not been set, yet. If stored in a ContinuationObject's parent_chain + /// field, means that there is currently no parent. + Absent = STACK_CHAIN_ABSENT_DISCRIMINANT, + MainStack(*mut StackLimits) = STACK_CHAIN_MAIN_STACK_DISCRIMINANT, + Continuation(*mut ContinuationObject) = STACK_CHAIN_CONTINUATION_DISCRIMINANT, +} + +impl StackChain { + pub const ABSENT_DISCRIMINANT: usize = STACK_CHAIN_ABSENT_DISCRIMINANT; + pub const MAIN_STACK_DISCRIMINANT: usize = STACK_CHAIN_MAIN_STACK_DISCRIMINANT; + pub const CONTINUATION_DISCRIMINANT: usize = STACK_CHAIN_CONTINUATION_DISCRIMINANT; +} + pub struct Payloads { /// Number of currently occupied slots. pub length: types::payloads::Length, @@ -92,7 +173,9 @@ impl From for i32 { /// TODO #[repr(C)] pub struct ContinuationObject { - pub parent: *mut ContinuationObject, + pub limits: StackLimits, + + pub parent_chain: StackChain, pub fiber: *mut ContinuationFiber, @@ -137,13 +220,25 @@ pub mod offsets { use crate::ContinuationObject; use memoffset::offset_of; + /// Offset of `limits` field + pub const LIMITS: i32 = offset_of!(ContinuationObject, limits) as i32; /// Offset of `args` field pub const ARGS: i32 = offset_of!(ContinuationObject, args) as i32; - /// Offset of `parent` field - pub const PARENT: i32 = offset_of!(ContinuationObject, parent) as i32; + /// Offset of `parent_chain` field + pub const PARENT_CHAIN: i32 = offset_of!(ContinuationObject, parent_chain) as i32; /// Offset of `state` field pub const STATE: i32 = offset_of!(ContinuationObject, state) as i32; /// Offset of `tag_return_values` field pub const TAG_RETURN_VALUES: i32 = offset_of!(ContinuationObject, tag_return_values) as i32; } + + pub mod stack_limits { + use crate::StackLimits; + use memoffset::offset_of; + + pub const STACK_LIMIT: i32 = offset_of!(StackLimits, stack_limit) as i32; + pub const LAST_WASM_EXIT_FP: i32 = offset_of!(StackLimits, last_wasm_exit_fp) as i32; + pub const LAST_WASM_EXIT_PC: i32 = offset_of!(StackLimits, last_wasm_exit_pc) as i32; + pub const LAST_WASM_ENTRY_SP: i32 = offset_of!(StackLimits, last_wasm_entry_sp) as i32; + } } diff --git a/crates/cranelift/src/wasmfx/optimized.rs b/crates/cranelift/src/wasmfx/optimized.rs index 8893f492255a..5cd961333406 100644 --- a/crates/cranelift/src/wasmfx/optimized.rs +++ b/crates/cranelift/src/wasmfx/optimized.rs @@ -15,6 +15,7 @@ pub(crate) use shared::typed_continuations_cont_ref_get_cont_obj; #[allow(unused_imports)] pub(crate) use shared::typed_continuations_new_cont_ref; +#[macro_use] pub(crate) mod typed_continuation_helpers { use cranelift_codegen::ir; use cranelift_codegen::ir::condcodes::IntCC; @@ -317,6 +318,19 @@ pub(crate) mod typed_continuation_helpers { pointer_type: ir::Type, } + /// Size of `wasmtime_continuations::StackChain` in machine words. + /// Used to verify that we have not changed its representation. + const STACK_CHAIN_POINTER_COUNT: usize = + std::mem::size_of::() / std::mem::size_of::(); + + /// Compile-time representation of wasmtime_continuations::StackChain, + /// consisting of two `ir::Value`s. + pub struct StackChain { + discriminant: ir::Value, + payload: ir::Value, + pointer_type: ir::Type, + } + impl ContinuationObject { pub fn new(address: ir::Value, pointer_type: ir::Type) -> ContinuationObject { ContinuationObject { @@ -402,6 +416,19 @@ pub(crate) mod typed_continuation_helpers { } return self.args().get_data(builder); } + + /// Stores the parent of this continuation, which may either be another + /// continuation or the main stack. It is therefore represented as a + /// `StackChain` element. + pub fn set_parent_stack_chain<'a>( + &mut self, + env: &mut crate::func_environ::FuncEnvironment<'a>, + builder: &mut FunctionBuilder, + new_stack_chain: &StackChain, + ) { + let offset = wasmtime_continuations::offsets::continuation_object::PARENT_CHAIN; + new_stack_chain.store(env, builder, self.address, offset) + } } impl Payloads { @@ -722,31 +749,183 @@ pub(crate) mod typed_continuation_helpers { } } - pub fn get_active_continuation<'a>( + /// Returns the stack chain saved in this `VMContext`. Note that the + /// head of the list is the actively running stack (main stack or + /// continuation). + pub fn load_stack_chain<'a>( &self, env: &mut crate::func_environ::FuncEnvironment<'a>, builder: &mut FunctionBuilder, - ) -> ir::Value { + ) -> StackChain { let base_addr = self.address; - let memflags = ir::MemFlags::trusted(); - let offset = i32::try_from(env.offsets.vmctx_typed_continuations_store()).unwrap(); - builder - .ins() - .load(self.pointer_type, memflags, base_addr, offset) + let offset = + i32::try_from(env.offsets.vmctx_typed_continuations_stack_chain()).unwrap(); + StackChain::load(env, builder, base_addr, offset, self.pointer_type) + } + + /// Stores the given stack chain saved in this `VMContext`, overwriting + /// the exsiting one. + pub fn store_stack_chain<'a>( + &self, + env: &mut crate::func_environ::FuncEnvironment<'a>, + builder: &mut FunctionBuilder, + stack_chain: &StackChain, + ) { + let base_addr = self.address; + + let offset = + i32::try_from(env.offsets.vmctx_typed_continuations_stack_chain()).unwrap(); + stack_chain.store(env, builder, base_addr, offset) } + /// Similar to `store_stack_chain`, but instead of storing an arbitrary + /// `StackChain`, stores StackChain::Continuation(contobj)`. pub fn set_active_continuation<'a>( &self, env: &mut crate::func_environ::FuncEnvironment<'a>, builder: &mut FunctionBuilder, contobj: ir::Value, ) { - let base_addr = self.address; + let chain = StackChain::from_continuation(builder, contobj, self.pointer_type); + self.store_stack_chain(env, builder, &chain) + } + } + + impl StackChain { + /// Creates a `Self` corressponding to `StackChain::Continuation(contobj)`. + pub fn from_continuation( + builder: &mut FunctionBuilder, + contobj: ir::Value, + pointer_type: ir::Type, + ) -> StackChain { + debug_assert_eq!(STACK_CHAIN_POINTER_COUNT, 2); + let discriminant = wasmtime_continuations::StackChain::CONTINUATION_DISCRIMINANT; + let discriminant = builder.ins().iconst(pointer_type, discriminant as i64); + StackChain { + discriminant, + payload: contobj, + pointer_type, + } + } + + /// Creates a `Self` corressponding to `StackChain::Absent`. + pub fn absent(builder: &mut FunctionBuilder, pointer_type: ir::Type) -> StackChain { + debug_assert_eq!(STACK_CHAIN_POINTER_COUNT, 2); + let discriminant = wasmtime_continuations::StackChain::ABSENT_DISCRIMINANT; + let discriminant = builder.ins().iconst(pointer_type, discriminant as i64); + let zero_filler = builder.ins().iconst(pointer_type, 0i64); + StackChain { + discriminant, + payload: zero_filler, + pointer_type, + } + } + + /// For debugging purposes. Emits an assertion that `self` does not correspond to + /// `StackChain::Absent`. + pub fn assert_not_absent<'a>( + &self, + env: &mut crate::func_environ::FuncEnvironment<'a>, + builder: &mut FunctionBuilder, + ) { + let discriminant = wasmtime_continuations::StackChain::ABSENT_DISCRIMINANT; + let discriminant = builder.ins().iconst(self.pointer_type, discriminant as i64); + emit_debug_assert_ne!(env, builder, self.discriminant, discriminant); + } + + /// Return the two raw `ir::Value`s that represent this StackChain. + pub fn to_raw_parts(&self) -> [ir::Value; STACK_CHAIN_POINTER_COUNT] { + [self.discriminant, self.payload] + } + /// Construct a `Self` from two raw `ir::Value`s. + pub fn from_raw_parts( + raw_data: [ir::Value; STACK_CHAIN_POINTER_COUNT], + pointer_type: ir::Type, + ) -> StackChain { + StackChain { + discriminant: raw_data[0], + payload: raw_data[1], + pointer_type, + } + } + + /// Load a `StackChain` object from the given address. + pub fn load<'a>( + _env: &mut crate::func_environ::FuncEnvironment<'a>, + builder: &mut FunctionBuilder, + pointer: ir::Value, + initial_offset: i32, + pointer_type: ir::Type, + ) -> StackChain { + let memflags = ir::MemFlags::trusted(); + let mut offset = initial_offset; + let mut data = vec![]; + for _ in 0..STACK_CHAIN_POINTER_COUNT { + data.push(builder.ins().load(pointer_type, memflags, pointer, offset)); + offset += pointer_type.bytes() as i32; + } + let data = <[ir::Value; STACK_CHAIN_POINTER_COUNT]>::try_from(data).unwrap(); + Self::from_raw_parts(data, pointer_type) + } + + /// Store this `StackChain` object at the given address. + pub fn store<'a>( + &self, + _env: &mut crate::func_environ::FuncEnvironment<'a>, + builder: &mut FunctionBuilder, + target_pointer: ir::Value, + initial_offset: i32, + ) { let memflags = ir::MemFlags::trusted(); - let offset = i32::try_from(env.offsets.vmctx_typed_continuations_store()).unwrap(); - builder.ins().store(memflags, contobj, base_addr, offset); + let mut offset = initial_offset; + let data = self.to_raw_parts(); + + for value in data { + debug_assert_eq!( + builder.func.dfg.value_type(value), + Type::int_with_byte_size(self.pointer_type.bytes() as u16).unwrap() + ); + builder.ins().store(memflags, value, target_pointer, offset); + offset += self.pointer_type.bytes() as i32; + } + } + + /// If `self` corresponds to a `StackChain::Continuation`, return the + /// pointer to the continuation object stored in the variant. + /// If `self` corresponds to `StackChain::MainStack`, trap with the + /// given `trap_code`. + /// Calling this if `self` corresponds to `StackChain::Absent` indicates + /// an internal bug. + pub fn unwrap_continuation_or_trap<'a>( + &self, + env: &mut crate::func_environ::FuncEnvironment<'a>, + builder: &mut FunctionBuilder, + trap_code: ir::TrapCode, + ) -> ir::Value { + if cfg!(debug_assertions) { + let absent_discriminant = wasmtime_continuations::StackChain::ABSENT_DISCRIMINANT; + let is_initialized = builder.ins().icmp_imm( + IntCC::NotEqual, + self.discriminant, + absent_discriminant as i64, + ); + emit_debug_assert!(env, builder, is_initialized); + } + + let continuation_discriminant = + wasmtime_continuations::StackChain::CONTINUATION_DISCRIMINANT; + let is_continuation = builder.ins().icmp_imm( + IntCC::Equal, + self.discriminant, + continuation_discriminant as i64, + ); + builder.ins().trapz(is_continuation, trap_code); + + // The representation of StackChain::Continuation stores + // the pointer right after the discriminant. + self.payload } } } @@ -781,31 +960,6 @@ fn typed_continuations_load_return_values<'a>( return values; } -fn typed_continuations_load_parent<'a>( - env: &mut crate::func_environ::FuncEnvironment<'a>, - builder: &mut FunctionBuilder, - contobj: ir::Value, -) -> ir::Value { - let offset = wasmtime_continuations::offsets::continuation_object::PARENT; - let memflags = ir::MemFlags::trusted(); - - builder - .ins() - .load(env.pointer_type(), memflags, contobj, offset) -} - -fn typed_continuations_store_parent<'a>( - _env: &mut crate::func_environ::FuncEnvironment<'a>, - builder: &mut FunctionBuilder, - contobj: ir::Value, - new_parent: ir::Value, -) { - let offset = wasmtime_continuations::offsets::continuation_object::PARENT; - let memflags = ir::MemFlags::trusted(); - - builder.ins().store(memflags, new_parent, contobj, offset); -} - fn typed_continuations_forward_tag_return_values<'a>( env: &mut crate::func_environ::FuncEnvironment<'a>, builder: &mut FunctionBuilder, @@ -974,13 +1128,15 @@ pub(crate) fn typed_continuations_load_continuation_object<'a>( env: &mut crate::func_environ::FuncEnvironment<'a>, builder: &mut FunctionBuilder, ) -> ir::Value { - let pointer_type = env.pointer_type(); let vmctx = env.vmctx(builder.cursor().func); - let base_addr = builder.ins().global_value(pointer_type, vmctx); - - let vmctx = tc::VMContext::new(base_addr, pointer_type); - - vmctx.get_active_continuation(env, builder) + let vmctx = builder.ins().global_value(env.pointer_type(), vmctx); + let vmctx = tc::VMContext::new(vmctx, env.pointer_type()); + let active_stack_chain = vmctx.load_stack_chain(env, builder); + active_stack_chain.unwrap_continuation_or_trap( + env, + builder, + ir::TrapCode::User(crate::DEBUG_ASSERT_TRAP_CODE), + ) } pub(crate) fn translate_cont_new<'a>( @@ -1014,29 +1170,34 @@ pub(crate) fn translate_resume<'a>( let switch_block = builder.create_block(); let forwarding_block = builder.create_block(); + let vmctx = env.vmctx(builder.func); + let vmctx = builder.ins().global_value(env.pointer_type(), vmctx); + // Preamble: Part of previously active block - { - let original_contobj = + + let (resume_contobj, parent_stack_chain) = { + let resume_contobj = shared::typed_continuations_cont_ref_get_cont_obj(env, builder, contref); if resume_args.len() > 0 { // We store the arguments in the continuation object to be resumed. let count = builder.ins().iconst(I64, resume_args.len() as i64); - typed_continuations_store_resume_args( - env, - builder, - resume_args, - count, - original_contobj, - ); + typed_continuations_store_resume_args(env, builder, resume_args, count, resume_contobj); } - // Make the currently running continuation the parent of the one we are about to resume. - let original_running_contobj = typed_continuations_load_continuation_object(env, builder); - typed_continuations_store_parent(env, builder, original_contobj, original_running_contobj); + // Make the currently running continuation (if any) the parent of the one we are about to resume. + let original_stack_chain = + tc::VMContext::new(vmctx, env.pointer_type()).load_stack_chain(env, builder); + original_stack_chain.assert_not_absent(env, builder); + tc::ContinuationObject::new(resume_contobj, env.pointer_type()).set_parent_stack_chain( + env, + builder, + &original_stack_chain, + ); - builder.ins().jump(resume_block, &[original_contobj]); - } + builder.ins().jump(resume_block, &[]); + (resume_contobj, original_stack_chain) + }; // Resume block: actually resume the fiber corresponding to the // continuation object given as a parameter to the block. This @@ -1044,15 +1205,8 @@ pub(crate) fn translate_resume<'a>( // to resume objects other than `original_contobj`. // We make the continuation object that was actually resumed available via // `resumed_contobj`, so that subsequent blocks can refer to it. - let (tag, resumed_contobj) = { + let tag = { builder.switch_to_block(resume_block); - builder.append_block_param(resume_block, env.pointer_type()); - - // The continuation object to actually call resume on. - let resume_contobj = builder.block_params(resume_block)[0]; - - let vmctx = env.vmctx(builder.func); - let vmctx = builder.ins().global_value(env.pointer_type(), vmctx); // We mark `resume_contobj` as the currently running one let vmctx = tc::VMContext::new(vmctx, env.pointer_type()); @@ -1065,9 +1219,8 @@ pub(crate) fn translate_resume<'a>( let (_vmctx, result) = shared::generate_builtin_call!(env, builder, tc_resume, [resume_contobj]); - // Now the parent contobj is active - let parent_contobj = typed_continuations_load_parent(env, builder, resume_contobj); - vmctx.set_active_continuation(env, builder, parent_contobj); + // Now the parent contobj (or main stack) is active again + vmctx.store_stack_chain(env, builder, &parent_stack_chain); // The result encodes whether the return happens via ordinary // means or via a suspend. If the high bit is set, then it is @@ -1097,7 +1250,7 @@ pub(crate) fn translate_resume<'a>( .brif(is_zero, return_block, &[], suspend_block, &[]); // We do not seal this block, yet, because the effect forwarding block has a back edge to it - (tag, resume_contobj) + tag }; // Suspend block. @@ -1147,15 +1300,17 @@ pub(crate) fn translate_resume<'a>( .collect(); let mut args = typed_continuations_load_payloads(env, builder, ¶m_types); - // We have an actual handling block for this tag, rather than just forwarding. - // Detatch the continuation object by setting its parent to NULL. + // We have an actual handling block for this tag, rather than just + // forwarding. Detatch the continuation object by setting its parent + // link to `StackChain::Absent`. let pointer_type = env.pointer_type(); - let null = builder.ins().iconst(pointer_type, 0); - typed_continuations_store_parent(env, builder, resumed_contobj, null); + let chain = tc::StackChain::absent(builder, pointer_type); + tc::ContinuationObject::new(resume_contobj, pointer_type) + .set_parent_stack_chain(env, builder, &chain); // Create and push the continuation reference. We only create // them here because we don't need them when forwarding. - let contref = env.typed_continuations_new_cont_ref(builder, resumed_contobj); + let contref = env.typed_continuations_new_cont_ref(builder, resume_contobj); args.push(contref); @@ -1166,7 +1321,7 @@ pub(crate) fn translate_resume<'a>( } // Note that at this point we haven't actually emitted any - // code for the switching logic itenv, but only filled + // code for the switching logic itself, but only filled // the Switch structure and created the blocks it jumps // to. @@ -1176,11 +1331,11 @@ pub(crate) fn translate_resume<'a>( { builder.switch_to_block(forwarding_block); - let parent_contobj = typed_continuations_load_parent(env, builder, resumed_contobj); - - builder - .ins() - .trapz(parent_contobj, ir::TrapCode::UnhandledTag); + let parent_contobj = parent_stack_chain.unwrap_continuation_or_trap( + env, + builder, + ir::TrapCode::UnhandledTag, + ); // We suspend, thus deferring handling to the parent. // We do nothing about tag *parameters*, these remain unchanged within the @@ -1192,14 +1347,15 @@ pub(crate) fn translate_resume<'a>( // continuation objects, and we need to move them down the chain // back to the continuation object where we originally // suspended. - typed_continuations_forward_tag_return_values( - env, - builder, - parent_contobj, - resumed_contobj, - ); - - builder.ins().jump(resume_block, &[resumed_contobj]); + typed_continuations_forward_tag_return_values(env, builder, parent_contobj, resume_contobj); + + // We create a back edge to the resume block. + // Note that both `resume_cotobj` and `parent_stack_chain` remain unchanged: + // In the current design, where forwarding is implemented by suspending + // up the chain of parent continuations and subsequently resume-ing back + // down the chain, both the continuation being resumed and its parent + // stay the same. + builder.ins().jump(resume_block, &[]); builder.seal_block(resume_block); } @@ -1222,18 +1378,17 @@ pub(crate) fn translate_resume<'a>( builder.switch_to_block(return_block); builder.seal_block(return_block); - let co = tc::ContinuationObject::new(resumed_contobj, env.pointer_type()); + let co = tc::ContinuationObject::new(resume_contobj, env.pointer_type()); co.set_state(builder, wasmtime_continuations::State::Returned); // Load and push the results. let returns = env.continuation_returns(type_index).to_vec(); - let values = - typed_continuations_load_return_values(env, builder, &returns, resumed_contobj); + let values = typed_continuations_load_return_values(env, builder, &returns, resume_contobj); // The continuation has returned and all `ContinuationReferences` // to it should have been be invalidated. We may safely deallocate // it. - shared::typed_continuations_drop_cont_obj(env, builder, resumed_contobj); + shared::typed_continuations_drop_cont_obj(env, builder, resume_contobj); return values; } diff --git a/crates/environ/src/vmoffsets.rs b/crates/environ/src/vmoffsets.rs index dfa8566bc03f..3a358ef656af 100644 --- a/crates/environ/src/vmoffsets.rs +++ b/crates/environ/src/vmoffsets.rs @@ -92,9 +92,13 @@ pub struct VMOffsets

{ defined_func_refs: u32, size: u32, - // NOTE(dhil): The following field is used as "global" to store - // the arguments of continuations and payloads of suspensions. - typed_continuations_store: u32, + // The following field stores a value of type + // `wasmtime_continuations::StackLimits`. + typed_continuations_main_stack_limits: u32, + // The following field stores a value of type + // `wasmtime_continuations::StackChain`. The head of the chain is the + // currently executing stack (main stack or a continuation). + typed_continuations_stack_chain: u32, typed_continuations_payloads: u32, } @@ -358,7 +362,8 @@ impl VMOffsets

{ calculate_sizes! { typed_continuations_payloads: "typed continuations payloads object", - typed_continuations_store: "typed continuations store", + typed_continuations_stack_chain: "typed continuations stack chain", + typed_continuations_main_stack_limits: "typed continuations main stack limits", defined_func_refs: "module functions", defined_globals: "defined globals", owned_memories: "owned memories", @@ -411,7 +416,8 @@ impl From> for VMOffsets

{ defined_globals: 0, defined_func_refs: 0, size: 0, - typed_continuations_store: 0, + typed_continuations_main_stack_limits: 0, + typed_continuations_stack_chain: 0, typed_continuations_payloads: 0, }; @@ -475,8 +481,14 @@ impl From> for VMOffsets

{ ret.num_escaped_funcs, ret.ptr.size_of_vm_func_ref(), ), - size(typed_continuations_store) - = ret.ptr.size(), + + align(std::mem::align_of::() as u32), + size(typed_continuations_main_stack_limits) + = std::mem::size_of::() as u32, + + align(std::mem::align_of::() as u32), + size(typed_continuations_stack_chain) + = std::mem::size_of::() as u32, align(std::mem::align_of::() as u32), size(typed_continuations_payloads) = @@ -734,13 +746,16 @@ impl VMOffsets

{ self.builtin_functions } - /// The offset of the typed continuations store, where we save the currently - /// running continuation (as a pointer to a - /// wasmtime_comtinuations::ContinuationObject, or null if running on main - /// stack). + /// TODO + #[inline] + pub fn vmctx_typed_continuations_main_stack_limits(&self) -> u32 { + self.typed_continuations_main_stack_limits + } + + /// TODO #[inline] - pub fn vmctx_typed_continuations_store(&self) -> u32 { - self.typed_continuations_store + pub fn vmctx_typed_continuations_stack_chain(&self) -> u32 { + self.typed_continuations_stack_chain } /// The offset of the typed continuations payloads object, stored as a as a diff --git a/crates/runtime/src/continuation.rs b/crates/runtime/src/continuation.rs index 8736d409f1a1..f89fb15be6a2 100644 --- a/crates/runtime/src/continuation.rs +++ b/crates/runtime/src/continuation.rs @@ -4,10 +4,10 @@ use crate::vmcontext::{VMFuncRef, VMOpaqueContext, ValRaw}; use crate::{Instance, TrapReason}; use std::cmp; use std::mem; -use std::ptr; use wasmtime_continuations::{debug_println, ENABLE_DEBUG_PRINTING}; pub use wasmtime_continuations::{ - ContinuationFiber, ContinuationObject, ContinuationReference, Payloads, State, + ContinuationFiber, ContinuationObject, ContinuationReference, Payloads, StackChain, + StackLimits, State, }; use wasmtime_fibre::{Fiber, FiberStack, Suspend}; @@ -142,9 +142,11 @@ pub fn cont_new( Box::new(fiber) }; + let tsp = fiber.stack().top().unwrap(); let contobj = Box::new(ContinuationObject { + limits: StackLimits::with_stack_limit(unsafe { tsp.sub(DEFAULT_FIBER_SIZE) } as usize), fiber: Box::into_raw(fiber), - parent: ptr::null_mut(), + parent_chain: StackChain::Absent, args: payload, tag_return_values: Payloads::new(0), state: State::Allocated, @@ -180,12 +182,29 @@ pub fn resume( }; if ENABLE_DEBUG_PRINTING { - let _running_contobj = instance.typed_continuations_store(); - debug_println!( - "Resuming contobj @ {:p}, previously running contobj is {:p}", - contobj, - _running_contobj - ); + let chain = instance.typed_continuations_stack_chain(); + // SAFETY: We maintain as an invariant that the stack chain field in the + // VMContext is non-null and contains a chain of zero or more + // StackChain::Continuation values followed by StackChain::Main. + match unsafe { &*chain } { + StackChain::Continuation(running_contobj) => { + debug_assert_eq!(contobj, *running_contobj); + debug_println!( + "Resuming contobj @ {:p}, previously running contobj is {:p}", + contobj, + running_contobj + ) + } + _ => { + // Before calling this function as a libcall, we must have set + // the parent of the to-be-resumed continuation to the + // previously running one. Hence, we must see a + // `StackChain::Continuation` variant. + return Err(TrapReason::user_without_backtrace(anyhow::anyhow!( + "Invalid StackChain value in VMContext" + ))); + } + } } unsafe { @@ -203,11 +222,14 @@ pub fn resume( // calling trampoline to execute it. if cfg!(debug_assertions) { - let _parent = unsafe { (*contobj).parent }; + // SAFETY: We maintain as an invariant that the stack chain field in the + // VMContext is non-null and contains a chain of zero or more + // StackChain::Continuation values followed by StackChain::Main. + let _parent_chain = unsafe { &(*contobj).parent_chain }; debug_println!( - "Continuation @ {:p} returned normally, setting running continuation in VMContext to {:p}", + "Continuation @ {:p} returned normally, setting running continuation in VMContext to {:?}", contobj, - _parent + _parent_chain ); } Ok(0) // zero value = return normally. @@ -228,18 +250,29 @@ pub fn resume( /// TODO #[inline(always)] pub fn suspend(instance: &mut Instance, tag_index: u32) -> Result<(), TrapReason> { - let running = instance.typed_continuations_store(); - let running = unsafe { - running.as_ref().ok_or_else(|| { - TrapReason::user_without_backtrace(anyhow::anyhow!( - "Calling suspend outside of a continuation" - )) - })? - }; + let chain_ptr = instance.typed_continuations_stack_chain(); + // TODO(dhil): This should be handled in generated code. + // SAFETY: We maintain as an invariant that the stack chain field in the + // VMContext is non-null and contains a chain of zero or more + // StackChain::Continuation values followed by StackChain::Main. + let chain = unsafe { &*chain_ptr }; + let running = match chain { + StackChain::Absent => Err(TrapReason::user_without_backtrace(anyhow::anyhow!( + "Internal error: StackChain not initialised" + ))), + StackChain::MainStack { .. } => Err(TrapReason::user_without_backtrace(anyhow::anyhow!( + "Calling suspend outside of a continuation" + ))), + StackChain::Continuation(running) => { + // SAFETY: See above. + Ok(unsafe { &**running }) + } + }?; let fiber = unsafe { - running.fiber.as_ref().ok_or_else(|| { + // SAFETY: See above. + (*running).fiber.as_ref().ok_or_else(|| { TrapReason::user_without_backtrace(anyhow::anyhow!( "Attempt to dereference null fiber!" )) @@ -250,9 +283,9 @@ pub fn suspend(instance: &mut Instance, tag_index: u32) -> Result<(), TrapReason TrapReason::user_without_backtrace(anyhow::anyhow!("Failed to retrieve stack top pointer!")) })?; debug_println!( - "Suspending while running {:p}, parent is {:p}", + "Suspending while running {:p}, parent is {:?}", running, - running.parent + running.parent_chain ); let suspend = wasmtime_fibre::unix::Suspend::from_top_ptr(stack_ptr); diff --git a/crates/runtime/src/instance.rs b/crates/runtime/src/instance.rs index b2b6cd75fe6b..402f4b143597 100644 --- a/crates/runtime/src/instance.rs +++ b/crates/runtime/src/instance.rs @@ -1133,8 +1133,12 @@ impl Instance { *self.vmctx_plus_offset_mut(offsets.vmctx_builtin_functions()) = &VMBuiltinFunctionsArray::INIT; - *self.vmctx_plus_offset_mut(offsets.vmctx_typed_continuations_store()) = - std::ptr::null_mut::(); + let main_stack_limits_ptr = + self.vmctx_plus_offset_mut(offsets.vmctx_typed_continuations_main_stack_limits()); + *main_stack_limits_ptr = wasmtime_continuations::StackLimits::default(); + + *self.vmctx_plus_offset_mut(offsets.vmctx_typed_continuations_stack_chain()) = + wasmtime_continuations::StackChain::MainStack(main_stack_limits_ptr); // Initialize the Payloads object to be empty let vmctx_payloads: *mut wasmtime_continuations::Payloads = @@ -1279,20 +1283,32 @@ impl Instance { fault } - pub(crate) fn typed_continuations_store( + #[allow(dead_code)] + pub(crate) fn typed_continuations_main_stack_limits( &mut self, - ) -> *mut crate::continuation::ContinuationObject { - unsafe { *self.vmctx_plus_offset_mut(self.offsets().vmctx_typed_continuations_store()) } + ) -> *mut wasmtime_continuations::StackLimits { + unsafe { + self.vmctx_plus_offset_mut(self.offsets().vmctx_typed_continuations_main_stack_limits()) + } + } + + pub(crate) fn typed_continuations_stack_chain( + &mut self, + ) -> *mut wasmtime_continuations::StackChain { + unsafe { + self.vmctx_plus_offset_mut(self.offsets().vmctx_typed_continuations_stack_chain()) + } } #[allow(dead_code)] - pub(crate) fn set_typed_continuations_store( + pub(crate) fn set_typed_continuations_stack_chain( &mut self, - contobj: *mut crate::continuation::ContinuationObject, + chain: *mut wasmtime_continuations::StackChain, ) { unsafe { - let ptr = self.vmctx_plus_offset_mut(self.offsets().vmctx_typed_continuations_store()); - *ptr = contobj; + let ptr = + self.vmctx_plus_offset_mut(self.offsets().vmctx_typed_continuations_stack_chain()); + *ptr = chain; } } diff --git a/tests/all/pooling_allocator.rs b/tests/all/pooling_allocator.rs index 8b6f82274220..df8cdb7c791f 100644 --- a/tests/all/pooling_allocator.rs +++ b/tests/all/pooling_allocator.rs @@ -661,12 +661,12 @@ configured maximum of 16 bytes; breakdown of allocation requirement: " } else { "\ -instance allocation for this module requires 272 bytes which exceeds the \ +instance allocation for this module requires 320 bytes which exceeds the \ configured maximum of 16 bytes; breakdown of allocation requirement: - * 58.82% - 160 bytes - instance state management - * 8.82% - 24 bytes - typed continuations payloads object - * 5.88% - 16 bytes - jit store state + * 50.00% - 160 bytes - instance state management + * 10.00% - 32 bytes - typed continuations payloads object + * 10.00% - 32 bytes - typed continuations main stack limits " }; match Module::new(&engine, "(module)") { @@ -690,11 +690,11 @@ configured maximum of 16 bytes; breakdown of allocation requirement: " } else { "\ -instance allocation for this module requires 1872 bytes which exceeds the \ +instance allocation for this module requires 1920 bytes which exceeds the \ configured maximum of 16 bytes; breakdown of allocation requirement: - * 8.55% - 160 bytes - instance state management - * 85.47% - 1600 bytes - defined globals + * 8.33% - 160 bytes - instance state management + * 83.33% - 1600 bytes - defined globals " }; match Module::new(&engine, &lots_of_globals) {