infra part 2

crates/continuations/src/lib.rs

@@ -92,8 +92,40 @@ const STACK_CHAIN_CONTINUATION_DISCRIMINANT: usize = 2;
 /// 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.
+///
+/// The following invariants hold for these `StackLimits` objects, and the data
+/// in `VMRuntimeLimits`.
+///
+/// Currently executing stack:
+/// For the currently executing stack (i.e., the stack that is at the head of
+/// the VMContext's `typed_continuations_chain` list), the associated
+/// `StackLimits` object contains stale/undefined data. Instead, the live data
+/// describing the limits for the currently executing stack is always maintained
+/// in `VMRuntimeLimits`. Note that as a general rule independently from any
+/// execution of continuations, the `last_wasm_exit*` fields in the
+/// `VMRuntimeLimits` contain undefined values while executing wasm.
+///
+/// Parents of currently executing stack:
+/// For stacks that appear in the tail of the VMContext's
+/// `typed_continuations_chain` list (i.e., stacks that are not currently
+/// executing themselves, but are a parent of the currently executing stack), we
+/// have the following: All the fields in the stack's StackLimits are valid,
+/// describing the stack's stack limit, and pointers where executing for that
+/// stack entered and exited WASM.
+///
+/// Suspended continuations:
+/// For suspended continuations (including their parents), we have the
+/// following. Note that the main stack can never be in this state. The
+/// `stack_limit` and `last_enter_wasm_sp` fields of the corresponding
+/// `StackLimits` object contain valid data, while the `last_exit_wasm_*` fields
+/// contain arbitrary values.
+/// There is only one exception to this: Note that a continuation that has been
+/// created with cont.new, but never been resumed so far, is considered
+/// "suspended". However, its `last_enter_wasm_sp` field contains undefined
+/// data. This is justified, because when resume-ing a continuation for the
+/// first time, a native-to-wasm trampoline is called, which sets up the
+/// `last_wasm_entry_sp` in the `VMRuntimeLimits` with the correct value, thus
+/// restoring the necessary invariant.
 #[derive(Debug, Clone, PartialEq)]
 #[repr(usize, C)]
 pub enum StackChain {

crates/cranelift/src/func_environ.rs

@@ -158,7 +158,8 @@ pub struct FuncEnvironment<'module_environment> {
     /// VMRuntimeLimits` for this function's vmctx argument. This pointer is stored
     /// in the vmctx itself, but never changes for the lifetime of the function,
     /// so if we load it up front we can continue to use it throughout.
-    vmruntime_limits_ptr: cranelift_frontend::Variable,
+    /// NOTE(frank-emrich) pub for use in crate::wasmfx::* modules
+    pub(crate) vmruntime_limits_ptr: cranelift_frontend::Variable,
 
     /// A cached epoch deadline value, when performing epoch-based
     /// interruption. Loaded from `VMRuntimeLimits` and reloaded after
@@ -2646,7 +2647,14 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
     ) -> WasmResult<()> {
         // If the `vmruntime_limits_ptr` variable will get used then we initialize
         // it here.
-        if self.tunables.consume_fuel || self.tunables.epoch_interruption {
+        if true || self.tunables.consume_fuel || self.tunables.epoch_interruption {
+            // TODO(frank-emrich) This is now done unconditionally because we
+            // need the `vmruntime_limits_ptr` variable when translating resume.
+            // This has no runtime overhead: We are adding a load to every
+            // function, but if it is not actually used, cranelift's DCE pass
+            // will remove it. However, it would be nicer to check if the
+            // function actually contains resume instructions, and only run
+            // `declare_vmruntime_limits_ptr` then.
             self.declare_vmruntime_limits_ptr(builder);
         }
         // Additionally we initialize `fuel_var` if it will get used.

crates/cranelift/src/wasmfx/optimized.rs

@@ -24,6 +24,7 @@ pub(crate) mod typed_continuation_helpers {
     use cranelift_frontend::FunctionBuilder;
     use std::mem;
     use wasmtime_environ::BuiltinFunctionIndex;
+    use wasmtime_environ::PtrSize;
 
     // This is a reference to this very module.
     // We need it so that we can refer to the functions inside this module from
@@ -927,6 +928,106 @@ pub(crate) mod typed_continuation_helpers {
             // the pointer right after the discriminant.
             self.payload
         }
+
+        /// 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
+        /// `ContinuationObject`.
+        pub 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;
+
+            self.assert_not_absent(env, builder);
+
+            // `self` corresponds to a StackChain::MainStack or
+            // StackChain::Continuation.
+            // In both cases, the payload is a pointer.
+            let ptr = self.payload;
+
+            // `obj` is now a pointer to the beginning of either
+            // 1. A ContinuationObject object (in the case of a
+            // StackChain::Continuation)
+            // 2. A StackLimits object (in the case of
+            // StackChain::MainStack)
+            //
+            // Since a ContinuationObject starts with an (inlined) StackLimits
+            // 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::continuation_object::LIMITS, 0);
+            ptr
+        }
+
+        /// Sets `last_wasm_entry_sp` and `stack_limit` fields in
+        /// `VMRuntimelimits` using the values from the `StackLimits` object
+        /// associated with this stack chain.
+        pub fn write_limits_to_vmcontext<'a>(
+            &self,
+            env: &mut crate::func_environ::FuncEnvironment<'a>,
+            builder: &mut FunctionBuilder,
+            vmruntime_limits: cranelift_frontend::Variable,
+        ) {
+            use wasmtime_continuations::offsets as o;
+
+            let stack_limits_ptr = self.get_stack_limits_ptr(env, builder);
+            let vmruntime_limits_ptr = builder.use_var(vmruntime_limits);
+
+            let memflags = ir::MemFlags::trusted();
+
+            let mut copy_to_vm_runtime_limits = |our_offset, their_offset| {
+                let our_value =
+                    builder
+                        .ins()
+                        .load(self.pointer_type, memflags, stack_limits_ptr, our_offset);
+                builder
+                    .ins()
+                    .store(memflags, our_value, vmruntime_limits_ptr, their_offset);
+            };
+
+            let pointer_size = self.pointer_type.bytes() as u8;
+            copy_to_vm_runtime_limits(
+                o::stack_limits::STACK_LIMIT,
+                pointer_size.vmruntime_limits_stack_limit(),
+            );
+            copy_to_vm_runtime_limits(
+                o::stack_limits::LAST_WASM_ENTRY_SP,
+                pointer_size.vmruntime_limits_last_wasm_entry_sp(),
+            );
+        }
+
+        /// Overwrites the `last_wasm_entry_sp` field of the `StackLimits`
+        /// object associated with this stack chain by loading the corresponding
+        /// field from the `VMRuntimeLimits`.
+        pub fn load_limits_from_vmcontext<'a>(
+            &self,
+            env: &mut crate::func_environ::FuncEnvironment<'a>,
+            builder: &mut FunctionBuilder,
+            vmruntime_limits: cranelift_frontend::Variable,
+        ) {
+            use wasmtime_continuations::offsets as o;
+
+            let stack_limits_ptr = self.get_stack_limits_ptr(env, builder);
+            let vmruntime_limits_ptr = builder.use_var(vmruntime_limits);
+
+            let memflags = ir::MemFlags::trusted();
+            let pointer_size = self.pointer_type.bytes() as u8;
+
+            let last_wasm_entry_sp = builder.ins().load(
+                self.pointer_type,
+                memflags,
+                vmruntime_limits_ptr,
+                pointer_size.vmruntime_limits_last_wasm_entry_sp(),
+            );
+            builder.ins().store(
+                memflags,
+                last_wasm_entry_sp,
+                stack_limits_ptr,
+                o::stack_limits::LAST_WASM_ENTRY_SP,
+            );
+        }
     }
 }
 
@@ -1212,12 +1313,26 @@ pub(crate) fn translate_resume<'a>(
         let vmctx = tc::VMContext::new(vmctx, env.pointer_type());
         vmctx.set_active_continuation(env, builder, resume_contobj);
 
+        // Note that the resume_contobj libcall a few lines further below
+        // manipulates the stack limits as follows:
+        // 1. Copy stack_limit, last_wasm_entry_sp and last_wasm_exit* values from
+        // VMRuntimeLimits into the currently active continuation (i.e., the
+        // one that will become the parent of the to-be-resumed one)
+        //
+        // 2. Copy `stack_limit` and `last_wasm_entry_sp` in the
+        // `StackLimits` of `resume_contobj` into the `VMRuntimeLimits`.
+        let parent_stacks_limit_pointer = parent_stack_chain.get_stack_limits_ptr(env, builder);
+
         // We mark `resume_contobj` to be invoked
         let co = tc::ContinuationObject::new(resume_contobj, env.pointer_type());
         co.set_state(builder, wasmtime_continuations::State::Invoked);
 
-        let (_vmctx, result) =
-            shared::generate_builtin_call!(env, builder, tc_resume, [resume_contobj]);
+        let (_vmctx, result) = shared::generate_builtin_call!(
+            env,
+            builder,
+            tc_resume,
+            [resume_contobj, parent_stacks_limit_pointer]
+        );
 
         // Now the parent contobj (or main stack) is active again
         vmctx.store_stack_chain(env, builder, &parent_stack_chain);
@@ -1258,6 +1373,15 @@ pub(crate) fn translate_resume<'a>(
         builder.switch_to_block(suspend_block);
         builder.seal_block(suspend_block);
 
+        // We store parts of the VMRuntimeLimits into the continuation that just suspended.
+        let suspended_chain =
+            tc::StackChain::from_continuation(builder, resume_contobj, env.pointer_type());
+        suspended_chain.load_limits_from_vmcontext(env, builder, env.vmruntime_limits_ptr);
+
+        // 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, env.vmruntime_limits_ptr);
+
         // We need to terminate this block before being allowed to switch to another one
         builder.ins().jump(switch_block, &[]);
     };
@@ -1378,6 +1502,10 @@ pub(crate) fn translate_resume<'a>(
         builder.switch_to_block(return_block);
         builder.seal_block(return_block);
 
+        // 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, env.vmruntime_limits_ptr);
+
         let co = tc::ContinuationObject::new(resume_contobj, env.pointer_type());
         co.set_state(builder, wasmtime_continuations::State::Returned);
 

crates/environ/src/builtin.rs

@@ -55,7 +55,7 @@ macro_rules! foreach_builtin_function {
             /// Creates a new continuation from a funcref.
             tc_cont_new(vmctx: vmctx, r: pointer, param_count: i64, result_count: i64) -> pointer;
             /// Resumes a continuation.
-            tc_resume(vmctx: vmctx, contobj: pointer) -> i32;
+            tc_resume(vmctx: vmctx, contobj: pointer, parent_stack_limits: pointer) -> i32;
             /// Suspends a continuation.
             tc_suspend(vmctx: vmctx, tag: i32);
             /// Returns the continuation object corresponding to the given continuation reference.

crates/runtime/src/continuation.rs

@@ -164,6 +164,7 @@ pub fn cont_new(
 pub fn resume(
     instance: &mut Instance,
     contobj: *mut ContinuationObject,
+    parent_stack_limits: *mut StackLimits,
 ) -> Result<u32, TrapReason> {
     let cont = unsafe {
         contobj.as_ref().ok_or_else(|| {
@@ -207,6 +208,20 @@ pub fn resume(
         }
     }
 
+    unsafe {
+        let runtime_limits = &**instance.runtime_limits();
+
+        (*parent_stack_limits).stack_limit = *runtime_limits.stack_limit.get();
+        (*parent_stack_limits).last_wasm_entry_sp = *runtime_limits.last_wasm_entry_sp.get();
+        // These last two values were only just updated in the `runtime_limits`
+        // because we entered the current libcall.
+        (*parent_stack_limits).last_wasm_exit_fp = *runtime_limits.last_wasm_exit_fp.get();
+        (*parent_stack_limits).last_wasm_exit_pc = *runtime_limits.last_wasm_exit_pc.get();
+
+        *runtime_limits.stack_limit.get() = (*contobj).limits.stack_limit;
+        *runtime_limits.last_wasm_entry_sp.get() = (*contobj).limits.last_wasm_entry_sp;
+    }
+
     unsafe {
         (*(*(*instance.store()).vmruntime_limits())
             .stack_limit

crates/runtime/src/libcalls.rs

@@ -788,10 +788,15 @@ fn tc_cont_new(
     Ok(ans.cast::<u8>())
 }
 
-fn tc_resume(instance: &mut Instance, contobj: *mut u8) -> Result<u32, TrapReason> {
+fn tc_resume(
+    instance: &mut Instance,
+    contobj: *mut u8,
+    parent_stack_limits: *mut u8,
+) -> Result<u32, TrapReason> {
     crate::continuation::resume(
         instance,
         contobj.cast::<crate::continuation::ContinuationObject>(),
+        parent_stack_limits.cast::<crate::continuation::StackLimits>(),
     )
 }