Stow epoch-check code locations in a custom section of the native binary.

This will let the signal handler distinguish epoch interrupts from other segfaults.

* Fix ISLE bug that was causing emit.rs code not to be triggered.
* Bless inconsequential destination register reassignment in existing disas test.

Author: erikrose

cranelift/codegen/src/isa/x64/inst/emit.rs

@@ -649,8 +649,7 @@ pub(crate) fn emit(
             // The ISLE has already emitted the dead load. Put the address of
             // this instruction aside so we can later distinguish whether a
             // segfault is its fault.
-
-            // Search for "let pc_offset = layout.ip_offset as i32;" as a string to pull on.
+            sink.add_epoch_check();
         }
 
         Inst::JmpKnown { dst } => uncond_jmp(sink, *dst),

cranelift/codegen/src/isa/x64/lower.isle

@@ -3592,7 +3592,7 @@
                                              load_ptr
                                              (zero_offset))
                                   (ExtKind.None)))
-            (_ SideEffectNoResult (x64_dead_load_with_context load_ptr context)))
+            (_ Unit (emit_side_effect (x64_dead_load_with_context load_ptr context))))
        (output_none)))
 
 ;;;; Rules for `get_{frame,stack}_pointer` and `get_return_address` ;;;;;;;;;;;;

cranelift/codegen/src/machinst/buffer.rs

@@ -252,6 +252,8 @@ pub struct MachBuffer<I: VCodeInst> {
     call_sites: SmallVec<[MachCallSite; 16]>,
     /// Any patchable call site locations.
     patchable_call_sites: SmallVec<[MachPatchableCallSite; 16]>,
+    /// Any locations which do an MMU-based check for the end of an epoch.
+    epoch_checks: SmallVec<[EpochCheckOffset; 16]>,
     /// Any exception-handler records referred to at call sites.
     exception_handlers: SmallVec<[MachExceptionHandler; 16]>,
     /// Any source location mappings referring to this code.
@@ -343,6 +345,7 @@ impl MachBufferFinalized<Stencil> {
             traps: self.traps,
             call_sites: self.call_sites,
             patchable_call_sites: self.patchable_call_sites,
+            epoch_checks: self.epoch_checks,
             exception_handlers: self.exception_handlers,
             srclocs: self
                 .srclocs
@@ -380,6 +383,8 @@ pub struct MachBufferFinalized<T: CompilePhase> {
     pub(crate) call_sites: SmallVec<[MachCallSite; 16]>,
     /// Any patchable call site locations refering to this code.
     pub(crate) patchable_call_sites: SmallVec<[MachPatchableCallSite; 16]>,
+    /// Any locations which do an MMU-based check for the end of an epoch.
+    pub epoch_checks: SmallVec<[EpochCheckOffset; 16]>,
     /// Any exception-handler records referred to at call sites.
     pub(crate) exception_handlers: SmallVec<[FinalizedMachExceptionHandler; 16]>,
     /// Any source location mappings referring to this code.
@@ -480,6 +485,7 @@ impl<I: VCodeInst> MachBuffer<I> {
             traps: SmallVec::new(),
             call_sites: SmallVec::new(),
             patchable_call_sites: SmallVec::new(),
+            epoch_checks: SmallVec::new(),
             exception_handlers: SmallVec::new(),
             srclocs: SmallVec::new(),
             debug_tags: vec![],
@@ -1581,6 +1587,7 @@ impl<I: VCodeInst> MachBuffer<I> {
             traps: self.traps,
             call_sites: self.call_sites,
             patchable_call_sites: self.patchable_call_sites,
+            epoch_checks: self.epoch_checks,
             exception_handlers: finalized_exception_handlers,
             srclocs,
             debug_tags: self.debug_tags,
@@ -1696,6 +1703,14 @@ impl<I: VCodeInst> MachBuffer<I> {
         });
     }
 
+    /// Record that an MMU-based epoch interruption check occurs at the current
+    /// offset. The signal handler uses these annotations to distinguish that a
+    /// segfault is actually an epoch interruption in disguise. The
+    /// DeadLoadWithContext instruction is assumed to have already been emitted.
+    pub fn add_epoch_check(&mut self) {
+        self.epoch_checks.push(self.cur_offset());
+    }
+
     /// Add an unwind record at the current offset.
     pub fn add_unwind(&mut self, unwind: UnwindInst) {
         self.unwind_info.push((self.cur_offset(), unwind));
@@ -2198,6 +2213,12 @@ pub struct MachPatchableCallSite {
     pub len: u32,
 }
 
+/// The location of an epoch-end check, when using MMU-based epoch interruption.
+///
+/// Specifically, this points to the instruction after the one that does the
+/// epoch-end check: the one at which to resume execution.
+pub type EpochCheckOffset = CodeOffset;
+
 /// A source-location mapping resulting from a compilation.
 #[derive(PartialEq, Debug, Clone)]
 #[cfg_attr(

crates/cranelift/src/compiler.rs

@@ -35,10 +35,11 @@ use wasmparser::{FuncValidatorAllocations, FunctionBody};
 use wasmtime_environ::obj::{ELF_WASMTIME_EXCEPTIONS, ELF_WASMTIME_FRAMES};
 use wasmtime_environ::{
     Abi, AddressMapSection, BuiltinFunctionIndex, CacheStore, CompileError, CompiledFunctionBody,
-    DefinedFuncIndex, FlagValue, FrameInstPos, FrameStackShape, FrameStateSlotBuilder,
-    FrameTableBuilder, FuncKey, FunctionBodyData, FunctionLoc, HostCall, InliningCompiler,
-    ModuleTranslation, ModuleTypesBuilder, PtrSize, StackMapSection, StaticModuleIndex,
-    TrapEncodingBuilder, TrapSentinel, TripleExt, Tunables, WasmFuncType, WasmValType,
+    DefinedFuncIndex, EpochCheckSection, FlagValue, FrameInstPos, FrameStackShape,
+    FrameStateSlotBuilder, FrameTableBuilder, FuncKey, FunctionBodyData, FunctionLoc, HostCall,
+    InliningCompiler, ModuleTranslation, ModuleTypesBuilder, PtrSize, StackMapSection,
+    StaticModuleIndex, TrapEncodingBuilder, TrapSentinel, TripleExt, Tunables, WasmFuncType,
+    WasmValType,
 };
 use wasmtime_unwinder::ExceptionTableBuilder;
 
@@ -468,6 +469,7 @@ impl wasmtime_environ::Compiler for Compiler {
         let mut stack_maps = StackMapSection::default();
         let mut exception_tables = ExceptionTableBuilder::default();
         let mut frame_tables = FrameTableBuilder::default();
+        let mut epoch_checks = EpochCheckSection::default();
 
         let funcs = funcs
             .iter()
@@ -536,6 +538,9 @@ impl wasmtime_environ::Compiler for Compiler {
                 )?;
                 nop_units.get_or_insert_with(|| func.buffer.nop_units.clone());
             }
+            if self.tunables.epoch_interruption_via_mmu {
+                epoch_checks.push(range.clone(), &func.buffer.epoch_checks);
+            }
             builder.append_padding(self.linkopts.padding_between_functions);
 
             let info = FunctionLoc {
@@ -582,6 +587,9 @@ impl wasmtime_environ::Compiler for Compiler {
         }
         stack_maps.append_to(obj);
         traps.append_to(obj);
+        if self.tunables.epoch_interruption_via_mmu {
+            epoch_checks.append_to(obj);
+        }
 
         let exception_section = obj.add_section(
             obj.segment_name(StandardSegment::Data).to_vec(),

crates/environ/src/compile/epoch_checks.rs

@@ -0,0 +1,76 @@
+//! Emission of compiled-artifact metadata describing where epoch-end checks
+//! occur in the code when using MMU-based epoch interruption. This lets the
+//! signal handler distinguish epoch interruptions from general segfaults.
+
+use crate::obj::ELF_WASMTIME_EPOCH_CHECKS;
+use crate::prelude::*;
+use object::write::{Object, StandardSegment};
+use object::{LittleEndian, SectionKind, U32Bytes};
+use std::ops::Range;
+
+/// Offset of an epoch check within its function, in bytes. Specifically, this
+/// points to The instruction after the one that does the epoch-end check: the
+/// one at which to resume execution.
+///
+/// This is parallel to cranelift's CodeOffset and exists to (1) avoid making it
+/// a dependency, (2) pin it down to <= 32 bits, since the format of the
+/// custom-section we're emitting depends on that, and (3) hold documentation.
+type EpochCheckOffset = u32;
+
+/// A builder and emitter of the custom section which houses MMU-based
+/// epoch-end-check locations in a native binary.
+#[derive(Default)]
+pub struct EpochCheckSection {
+    /// Offset of the instruction to resume at after the epoch end and task switch
+    return_offsets: Vec<U32Bytes<LittleEndian>>,
+    /// The largest (and most recent, because we accept them only in order)
+    /// offset received so far for enforcing ordering. This is relative to the
+    /// start of the code (text) section so we can make sure functions are
+    /// ordered too.
+    last_offset: u32,
+}
+
+impl EpochCheckSection {
+    /// Adds an epoch-check location to the section.
+    ///
+    /// Calls to this must be ordered by the location of `func`, and
+    /// `check_offsets` must be ordered (within each function) as well.
+    pub fn push(&mut self, func: Range<u64>, check_offsets: &[EpochCheckOffset]) {
+        // Check that functions have been pushed in order so our section is
+        // sorted for free.
+        let func_start = u32::try_from(func.start).unwrap();
+        let func_end = u32::try_from(func.end).unwrap();
+        assert!(func_start >= self.last_offset);
+
+        // Remember each offset, ensuring they are in order.
+        for offset in check_offsets {
+            let text_section_relative = func_start + offset;
+            assert!(text_section_relative > self.last_offset);
+            self.return_offsets
+                .push(U32Bytes::new(LittleEndian, text_section_relative));
+            self.last_offset = text_section_relative;
+        }
+        self.last_offset = func_end;
+    }
+
+    /// Encodes this section into an object.
+    pub fn append_to(self, obj: &mut Object) {
+        let section = obj.add_section(
+            obj.segment_name(StandardSegment::Data).to_vec(),
+            ELF_WASMTIME_EPOCH_CHECKS.as_bytes().to_vec(),
+            SectionKind::ReadOnlyData,
+        );
+
+        // Append length.
+        obj.append_section_data(
+            section,
+            &u32::try_from(self.return_offsets.len())
+                .unwrap()
+                .to_le_bytes(),
+            1,
+        );
+
+        // Append offsets.
+        obj.append_section_data(section, object::bytes_of_slice(&self.return_offsets), 1);
+    }
+}

crates/environ/src/compile/mod.rs

@@ -16,6 +16,7 @@ use std::path;
 use std::sync::Arc;
 
 mod address_map;
+mod epoch_checks;
 mod frame_table;
 mod module_artifacts;
 mod module_environ;
@@ -24,6 +25,7 @@ mod stack_maps;
 mod trap_encoding;
 
 pub use self::address_map::*;
+pub use self::epoch_checks::*;
 pub use self::frame_table::*;
 pub use self::module_artifacts::*;
 pub use self::module_environ::*;

crates/environ/src/obj.rs

@@ -98,6 +98,25 @@ pub const ELF_WASMTIME_STACK_MAP: &str = ".wasmtime.stackmap";
 /// to the 32-bit encodings for offsets this doesn't support images >=4gb.
 pub const ELF_WASMTIME_TRAPS: &str = ".wasmtime.traps";
 
+/// A custom section which contains the offsets of instructions which check for
+/// the end of epochs when using `--epoch-interruption-via-mmu`.
+///
+/// The contents are examined at runtime by the signal handler to determine
+/// whether a segfault is due to an epoch ending (vs. a legitimate crash).
+///
+/// This section has a custom binary encoding:
+///
+/// * The section starts with a 32-bit little-endian integer which tell the
+///   number of items in the following array.
+/// * Next comes a sorted array of 32-bit unsigned little-endian integers which
+///   represent offsets from the beginning of the text section to the
+///   instruction following the load which triggers epoch-ending segfaults.
+///   TODO: We may point elsewhere if it's more useful to the signal handler. Be
+///   careful if this could end up pointing off the end of the text section.
+///
+/// The 32-bit encodings herein mean that >=4gb text sections are not supported.
+pub const ELF_WASMTIME_EPOCH_CHECKS: &str = ".wasmtime.epochchecks";
+
 /// A custom binary-encoded section of the wasmtime compilation
 /// artifacts which encodes exception tables.
 ///

tests/all/epoch_mmu.rs

@@ -0,0 +1,58 @@
+#![cfg(not(miri))]
+
+use object::{Object, ObjectSection};
+use wasmtime::{Config, Engine};
+use wasmtime_environ::obj::ELF_WASMTIME_EPOCH_CHECKS;
+
+/// Asserts that each epoch-check offset encoded into the binary points to the
+/// byte after its corresponding dead load.
+#[test]
+fn epoch_check_offsets() {
+    let mut config = Config::new();
+    config.target("x86_64").unwrap();
+    config.epoch_interruption_via_mmu(true);
+    let engine = Engine::new(&config).unwrap();
+
+    // A function with an infinite loop contains two epoch checks: one in the
+    // function prologue and another at the loop backedge.
+    let elf_bytes = engine
+        .precompile_module(
+            // If you change this wat, change it in
+            // epoch-interruption-mmu-compile-loop.wat, too.
+            r#"(module
+             (memory 0)
+             (func (loop (br 0)))
+           )"#
+            .as_bytes(),
+        )
+        .unwrap();
+
+    let elf = object::read::elf::ElfFile64::<object::Endianness>::parse(&*elf_bytes)
+        .expect("ELF should be parseable");
+    let section = elf
+        .section_by_name(ELF_WASMTIME_EPOCH_CHECKS)
+        .expect(&format!(
+            "{ELF_WASMTIME_EPOCH_CHECKS} section should be present"
+        ));
+    let data = section.data().unwrap();
+    assert!(data.len() >= 4, "section should at least contain a count");
+    let count = u32::from_le_bytes(data[..4].try_into().unwrap()) as usize;
+    assert_eq!(
+        data.len(),
+        4 + count * 4,
+        "section should be the right size to hold the offsets it claims to contain"
+    );
+    let offsets: Vec<u32> = data[4..]
+        .chunks_exact(4)
+        .map(|c| u32::from_le_bytes(c.try_into().unwrap()))
+        .collect();
+
+    // The emitted machine code is nailed down by the
+    // epoch-interruption-mmu-compile-loop.wat disas test. As long as that keeps
+    // passing, these offsets remain valid.
+    assert_eq!(
+        offsets,
+        vec![15, 18],
+        "There should be 2 epoch checks (function prologue & loop backedge). The offset after the prologue's dead load should be 15, and the one after the loop's backedge should be 18."
+    );
+}

tests/all/main.rs

@@ -15,6 +15,7 @@ mod debug;
 mod defaults;
 mod engine;
 mod epoch_interruption;
+mod epoch_mmu;
 mod eqref;
 mod exceptions;
 mod exnrefs;

tests/disas/epoch-interruption-mmu-compile-loop.wat

@@ -0,0 +1,19 @@
+;;! target = "x86_64"
+;;! test = "compile"
+;;! flags = ["-Wepoch-interruption-via-mmu=y"]
+
+;; Nail down codegen for the snippet in epoch_check_offsets() test. If this
+;; starts failing, that may need the offsets in its assert reexamined.
+
+(module
+  (memory 0)
+  (func (loop (br 0)))
+)
+;; wasm[0]::function[0]:
+;;       pushq   %rbp
+;;       movq    %rsp, %rbp
+;;       movq    8(%rdi), %r8
+;;       movq    0x10(%r8), %r8
+;;       movq    (%r8), %r9
+;;       movq    (%r8), %r10
+;;       jmp     0xf