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.

TODO: Tests

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/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,75 @@
+//! 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 and (2) to pin it down to <= 32 bits, since the format of the
+/// custom-section we're emitting depends on that, (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 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.
 ///