Store not just the start but also the end of the load instructions in the custom section.

Load instructions are usually 3 bytes in x64 but extend to 4 if loading from R12 or RSP. We pack the distinction into an array of bits to save perhaps several K of rather hot cache.

Author: erikrose

cranelift/codegen/meta/src/shared/instructions.rs

@@ -393,8 +393,9 @@ fn define_control_flow(
         segfault, which hands off control to a signal handler for further
         action. The handler has access to ``context`` (typically the
         ``VMContext``'s ``vm_store_context``) and can use the second
-        reserved register to store a temp value, as needed on platforms
-        where signal handlers cannot push stack frames.
+        reserved register to store a temp value--like the original return
+        value--as needed on platforms where signal handlers cannot push stack
+        frames.
 
         On x64, RDI holds ``context``, and R10 is used as scratch space.
         "#,

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

@@ -648,17 +648,15 @@ pub(crate) fn emit(
         Inst::DeadLoadWithContext { dst, load_ptr, .. } => {
             let start = sink.cur_offset();
 
-            // Since we're clobbering r10 anyway to store the original return
-            // address, we also use it as a destination for the dead load rather
-            // than sucking up another reg:
             let load_ptr_addr = SyntheticAmode::real(Amode::imm_reg(0, **load_ptr));
+            // Since we're clobbering dst anyway to store the original return
+            // address, also use it as a destination for the dead load rather
+            // than sucking up another reg:
             asm::inst::movq_rm::new(*dst, load_ptr_addr).emit(sink, info, state);
 
-            let end = sink.cur_offset();
-
             // Put the address of this instruction aside so we can later
             // distinguish whether a segfault is its fault.
-            sink.add_epoch_check();
+            sink.add_epoch_check(start, sink.cur_offset());
         }
 
         Inst::JmpKnown { dst } => uncond_jmp(sink, *dst),

cranelift/codegen/src/machinst/buffer.rs

@@ -253,7 +253,7 @@ pub struct MachBuffer<I: VCodeInst> {
     /// 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]>,
+    epoch_checks: SmallVec<[Range<CodeOffset>; 16]>,
     /// Any exception-handler records referred to at call sites.
     exception_handlers: SmallVec<[MachExceptionHandler; 16]>,
     /// Any source location mappings referring to this code.
@@ -384,7 +384,7 @@ pub struct MachBufferFinalized<T: CompilePhase> {
     /// 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]>,
+    pub epoch_checks: SmallVec<[Range<CodeOffset>; 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.
@@ -1704,11 +1704,13 @@ 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
+    /// offset. A signal handler may use 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());
+    /// `start` is the offset of the emitted instruction, and `end` is the
+    /// offset of the immediately following instruction.
+    pub fn add_epoch_check(&mut self, start: CodeOffset, end: CodeOffset) {
+        self.epoch_checks.push(start..end);
     }
 
     /// Add an unwind record at the current offset.
@@ -2213,12 +2215,6 @@ 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/environ/src/compile/epoch_checks.rs

@@ -4,13 +4,13 @@
 
 use crate::obj::ELF_WASMTIME_EPOCH_CHECKS;
 use crate::prelude::*;
+use object::SectionKind;
 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.
+/// Offset of an epoch check (in bytes) within the text section. Specifically,
+/// this points to the load instruction that may trigger the epoch-ending
+/// segfault.
 ///
 /// 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
@@ -21,8 +21,12 @@ type EpochCheckOffset = u32;
 /// 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>>,
+    /// Offset of the start of the load instruction which effects the epoch
+    /// check
+    starts: Vec<u32>,
+    /// Packed bits, parallel to elements of `starts`, which tell the length of
+    /// the load instruction: 0 if 3 bytes, 1 if 4
+    ends: Vec<u8>,
     /// 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
@@ -34,20 +38,29 @@ 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]) {
+    /// offsets must be ordered (within each function) as well.
+    pub fn push(&mut self, func: Range<u64>, checks: &[Range<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;
+        for check in checks {
+            let text_section_relative = func_start + check.start;
             assert!(text_section_relative > self.last_offset);
-            self.return_offsets
-                .push(U32Bytes::new(LittleEndian, text_section_relative));
+            let bit_number = self.starts.len();
+            self.starts.push(text_section_relative);
+            let load_is_long = match check.len() {
+                3 => false,
+                4 => true,
+                _ => panic!(
+                    "Unexpected length of epoch-checking load instruction: {}",
+                    check.len()
+                ),
+            };
+            set_bit(&mut self.ends, bit_number, load_is_long);
             self.last_offset = text_section_relative;
         }
         self.last_offset = func_end;
@@ -61,7 +74,75 @@ impl EpochCheckSection {
             SectionKind::ReadOnlyData,
         );
 
-        // Append offsets.
-        obj.append_section_data(section, object::bytes_of_slice(&self.return_offsets), 4);
+        let num_checks: u32 = self.starts.len().try_into()
+            .expect("there should be few enough epoch checks to be indexed by a u32, as the Wasm itself is only that long");
+        // Number of epoch checks:
+        obj.append_section_data(
+            section,
+            object::bytes_of(&num_checks),
+            4, // For speed, avoid splitting across a cache line.
+        );
+        // Align to 4 so we can use from_raw_parts() when reading:
+        obj.append_section_data(section, object::bytes_of_slice(&self.starts), 4);
+        // We'll be querying this a byte at a time so don't care about alignment:
+        obj.append_section_data(section, object::bytes_of_slice(&self.ends), 1);
+    }
+}
+
+/// Returns the offset at which to resume after the given epoch check.
+///
+/// If there is no epoch check at that offset, return None.
+pub fn return_offset_for_epoch_check(
+    section: &[u8],
+    check: EpochCheckOffset,
+) -> Option<EpochCheckOffset> {
+    let (num_checks, rest) = object::from_bytes::<u32>(section)
+        .expect(".wasmtime.epochchecks section should be long enough to contain count");
+    let (starts, ends) = object::slice_from_bytes::<u32>(rest, *num_checks as usize)
+        .expect(".wasmtime.epochchecks section should be long enough to contain starts");
+
+    starts
+        .binary_search(&check)
+        .ok()
+        .map(|epoch_index| check + (if get_bit(ends, epoch_index) { 4 } else { 3 }))
+}
+
+/// Sets bit `dest_bit_number` (0-based) in `dest_slice` to `value`.
+fn set_bit(dest_slice: &mut Vec<u8>, dest_bit_number: usize, value: bool) {
+    let byte = dest_bit_number / 8;
+    if byte >= dest_slice.len() {
+        dest_slice.resize(byte + 1, 0);
     }
+    let bit = dest_bit_number % 8;
+    if value {
+        dest_slice[byte] |= 1 << bit;
+    } else {
+        dest_slice[byte] &= !(1 << bit);
+    }
+}
+
+/// Returns bit `bit_number` (0-based) of `bytes`.
+fn get_bit(bytes: &[u8], bit_number: usize) -> bool {
+    let byte = bit_number / 8;
+    let bit = bit_number % 8;
+    bytes[byte] & (1 << bit) != 0
+}
+
+#[test]
+fn test_get_and_set_bit() {
+    let mut bits = vec![];
+    set_bit(&mut bits, 0, true); // first
+    set_bit(&mut bits, 2, true); // middle
+    set_bit(&mut bits, 18, true); // something beyond first byte
+    assert_eq!(bits, vec![5, 0, 4]);
+    assert!(get_bit(&bits, 0));
+    assert!(get_bit(&bits, 18));
+    assert!(!get_bit(&bits, 3));
+}
+
+#[test]
+fn test_set_bit_clear() {
+    let mut bits = vec![0xFF];
+    set_bit(&mut bits, 1, false);
+    assert_eq!(bits, vec![0xFD]);
 }

crates/environ/src/obj.rs

@@ -98,17 +98,27 @@ 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 is 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.
+/// A custom section through which we can locate instructions which
+/// check for the end of epochs when using `--epoch-interruption-via-mmu`.
+///
+/// The section allows for finding both the beginnings and ends of such
+/// instructions so the signal handler can identify segfaults which signify the
+/// end of an epoch (vs. ordinary crashes) and also find the address at which to
+/// resume afterward.
+///
+/// The format is architecture-dependent. Because the only forseen need to read
+/// it is on the platform where the binary will be executed and performance is
+/// sensitive, endianness follows whatever is native for the target. For x64,
+/// the only currently supported platform, the section comprises these data:
+/// * A u32 stating how many epoch checks are in each of the following 2 items
+/// * A sorted array of u32s which represent offsets from the beginning of the
+///   text section to the load instruction triggering the epoch end
+/// * A parallel array of bits representing the length of those load
+///   instructions: 0 meaning 3 bytes, 1 meaning 4. These are the only lengths
+///   possible on x64. Resumption should happen directly after the load
+///   instruction. There is no danger of this location pointing beyond the end
+///   of the function; even if a function is empty, there's at least a return
+///   (or, in case of an infinite loop, a jmp) after its prologue.
 ///
 /// The 32-bit encodings herein mean >=4gb text sections are not supported.
 pub const ELF_WASMTIME_EPOCH_CHECKS: &str = ".wasmtime.epochchecks";

crates/environ/src/tunables.rs

@@ -90,7 +90,7 @@ define_tunables! {
         /// Whether or not we use epoch-based interruption.
         pub epoch_interruption: bool,
 
-        /// Whether or not to use MMU tricks to speed epoch deadline checks.
+        /// Whether or not to use MMU tricks to speed epoch-end checks.
         /// TODO: Consider whether this should be orthogonal to
         /// epoch_interruption. If not, combine them into an enum or something.
         pub epoch_interruption_via_mmu: bool,

crates/wasmtime/src/config.rs

@@ -755,7 +755,8 @@ impl Config {
         self
     }
 
-    /// Stuff
+    /// Enables a faster-than-deadlines epoch detection mechanism based on
+    /// memory-page permissions
     pub fn epoch_interruption_via_mmu(&mut self, enable: bool) -> &mut Self {
         self.tunables.epoch_interruption_via_mmu = Some(enable);
         self

crates/wasmtime/src/runtime/code_memory.rs

@@ -8,7 +8,7 @@ use core::ops::Range;
 use object::SectionFlags;
 use object::endian::Endianness;
 use object::read::{Object, ObjectSection, elf::ElfFile64};
-use wasmtime_environ::{Trap, lookup_trap_code, obj};
+use wasmtime_environ::{Trap, lookup_trap_code, obj, return_offset_for_epoch_check};
 use wasmtime_unwinder::ExceptionTable;
 
 /// Management of executable memory within a `MmapVec`
@@ -33,6 +33,7 @@ pub struct CodeMemory {
     text: Range<usize>,
     unwind: Range<usize>,
     trap_data: Range<usize>,
+    epoch_check_data: Range<usize>,
     wasm_data: Range<usize>,
     address_map_data: Range<usize>,
     stack_map_data: Range<usize>,
@@ -127,6 +128,7 @@ impl CodeMemory {
         #[cfg(feature = "debug-builtins")]
         let mut has_native_debug_info = false;
         let mut trap_data = 0..0;
+        let mut epoch_check_data = 0..0;
         let mut exception_data = 0..0;
         let mut frame_tables_data = 0..0;
         let mut wasm_data = 0..0;
@@ -178,6 +180,7 @@ impl CodeMemory {
                 obj::ELF_WASMTIME_ADDRMAP => address_map_data = range,
                 obj::ELF_WASMTIME_STACK_MAP => stack_map_data = range,
                 obj::ELF_WASMTIME_TRAPS => trap_data = range,
+                obj::ELF_WASMTIME_EPOCH_CHECKS => epoch_check_data = range,
                 obj::ELF_WASMTIME_EXCEPTIONS => exception_data = range,
                 obj::ELF_WASMTIME_FRAMES => frame_tables_data = range,
                 obj::ELF_NAME_DATA => func_name_data = range,
@@ -222,6 +225,7 @@ impl CodeMemory {
             text,
             unwind,
             trap_data,
+            epoch_check_data,
             address_map_data,
             stack_map_data,
             exception_data,

tests/all/epoch_mmu.rs

@@ -1,6 +1,6 @@
 #![cfg(not(miri))]
 
-use object::{Object, ObjectSection};
+use object::{LittleEndian, Object, ObjectSection, U32Bytes};
 use wasmtime::{Config, Engine};
 use wasmtime_environ::obj::ELF_WASMTIME_EPOCH_CHECKS;
 
@@ -35,17 +35,28 @@ fn epoch_check_offsets() {
             "{ELF_WASMTIME_EPOCH_CHECKS} section should be present"
         ));
     let data = section.data().unwrap();
-    let offsets: Vec<u32> = data
-        .chunks_exact(4)
-        .map(|c| u32::from_le_bytes(c.try_into().unwrap()))
-        .collect();
+
+    let (count_raw, rest) = object::from_bytes::<U32Bytes<LittleEndian>>(data).expect(
+        ".wasmtime.epochchecks section should be long enough to contain a count of epoch checks",
+    );
+    let count = count_raw.get(LittleEndian) as usize;
+    let (starts_raw, rest) = object::slice_from_bytes::<U32Bytes<LittleEndian>>(rest, count)
+        .expect(".wasmtime.epochchecks section should be long enough to contain a location for each epoch check");
+    let starts: Vec<u32> = starts_raw.iter().map(|b| b.get(LittleEndian)).collect();
+    let (length_bits, _rest) = object::slice_from_bytes::<u8>(rest, count.div_ceil(8))
+        .expect(".wasmtime.epochchecks section should be long enough to contain a length bit for each epoch check");
 
     // 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.
+    // passing, these values remain valid.
+    assert_eq!(
+        starts,
+        vec![12, 15],
+        "There should be 2 epoch checks (function prologue & loop backedge). The offset of the prologue's dead load should be 12, and that of the loop's backedge should be 15."
+    );
     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."
+        length_bits,
+        vec![0],
+        "Neither check's load instruction uses R12 of RSP as its source, so all length bits should be 0."
     );
 }