Make pagemap integration resilient across forks

Author: alexcrichton

crates/wasmtime/Cargo.toml

@@ -113,6 +113,10 @@ wasmtime-versioned-export-macros = { workspace = true, optional = true }
 name = "host_segfault"
 harness = false
 
+[[test]]
+name = "engine_across_forks"
+harness = false
+
 # =============================================================================
 #
 # Features for the Wasmtime crate.

crates/wasmtime/src/runtime/vm/sys/unix/pagemap.rs

@@ -7,18 +7,100 @@ use crate::prelude::*;
 
 use self::ioctl::{Categories, PageMapScanBuilder};
 use crate::runtime::vm::{HostAlignedByteCount, host_page_size};
+use rustix::fd::AsRawFd;
 use rustix::ioctl::ioctl;
 use std::fs::File;
 use std::mem::MaybeUninit;
 use std::ptr;
+use std::sync::LazyLock;
+
+/// A static file-per-process which represents this process's page map file.
+///
+/// Note that this is required to be updated on a fork because otherwise this'll
+/// refer to the parent process's page map instead of the child process's page
+/// map. Thus when first initializing this file the `pthread_atfork` function is
+/// used to hook the child process to update this.
+///
+/// Also note that updating this is not done via mutation but rather it's done
+/// with `dup2` to replace the file descriptor that `File` points to in-place.
+/// The local copy of of `File` is then closed in the atfork handler.
+static PROCESS_PAGEMAP: LazyLock<Option<File>> = LazyLock::new(|| {
+    let pagemap = File::open("/proc/self/pagemap").ok()?;
+
+    // SAFETY: all libc functions are unsafe by default, and we're basically
+    // going to do our damndest to make sure this invocation of `pthread_atfork`
+    // is safe, namely the handler registered here is intentionally quite
+    // minimal and only accesses the `PROCESS_PAGEMAP`.
+    let rc = unsafe { libc::pthread_atfork(None, None, Some(after_fork_in_child)) };
+    if rc != 0 {
+        return None;
+    }
+
+    return Some(pagemap);
+
+    /// Hook executed as part of `pthread_atfork` in the child process after a
+    /// fork.
+    ///
+    /// # Safety
+    ///
+    /// This function is not safe to call in general and additionally has its
+    /// own stringent safety requirements. This is after a fork but before exec
+    /// so all the safety requirements of `Command::pre_exec` in the standard
+    /// library apply here. Effectively the standard library primitives are
+    /// avoided here as they aren't necessarily safe to execute in this context.
+    unsafe extern "C" fn after_fork_in_child() {
+        let Some(parent_pagemap) = PROCESS_PAGEMAP.as_ref() else {
+            // This should not be reachable, but to avoid panic infrastructure
+            // here this is just skipped instead.
+            return;
+        };
+
+        // SAFETY: see function documentation.
+        //
+        // Here `/proc/self/pagemap` is opened in the child. If that fails for
+        // whatever reason then the pagemap is replaced with `/dev/null` which
+        // means that all future ioctls for `PAGEMAP_SCAN` will fail. If that
+        // fails then that's left to abort the process for now. If that's
+        // problematic we may want to consider opening a local pipe and then
+        // installing that here? Unsure.
+        //
+        // Once a fd is opened the `dup2` syscall is used to replace the
+        // previous file descriptor stored in `parent_pagemap`. That'll update
+        // the pagemap in-place in this child for all future use in case this is
+        // further used in the child.
+        //
+        // And finally once that's all done the `child_pagemap` is itself
+        // closed since we have no more need for it.
+        unsafe {
+            let flags = libc::O_CLOEXEC | libc::O_RDONLY;
+            let mut child_pagemap = libc::open(c"/proc/self/pagemap".as_ptr(), flags);
+            if child_pagemap == -1 {
+                child_pagemap = libc::open(c"/dev/null".as_ptr(), flags);
+            }
+            if child_pagemap == -1 {
+                libc::abort();
+            }
+
+            let rc = libc::dup2(child_pagemap, parent_pagemap.as_raw_fd());
+            if rc == -1 {
+                libc::abort();
+            }
+            let rc = libc::close(child_pagemap);
+            if rc == -1 {
+                libc::abort();
+            }
+        }
+    }
+});
 
 #[derive(Debug)]
-pub struct PageMap(File);
+pub struct PageMap(&'static File);
 
 impl PageMap {
     #[cfg(feature = "pooling-allocator")]
     pub fn new() -> Option<PageMap> {
-        let file = File::open("/proc/self/pagemap").ok()?;
+        let file = PROCESS_PAGEMAP.as_ref()?;
+
         // Check if the `pagemap_scan` ioctl is supported.
         let mut regions = vec![MaybeUninit::uninit(); 1];
         let pm_scan = PageMapScanBuilder::new(ptr::slice_from_raw_parts(ptr::null_mut(), 0))

crates/wasmtime/tests/engine_across_forks.rs

@@ -0,0 +1,132 @@
+use wasmtime::Result;
+
+fn main() -> Result<()> {
+    #[cfg(unix)]
+    if true {
+        use libtest_mimic::{Arguments, Trial};
+
+        let mut trials = Vec::new();
+        for (name, test) in linux::TESTS {
+            trials.push(Trial::test(*name, || {
+                test().map_err(|e| format!("{e:?}").into())
+            }));
+        }
+
+        let mut args = Arguments::from_args();
+        // I'll be honest, I'm scared of threads + fork, so I'm just
+        // preemptively disabling threads here.
+        args.test_threads = Some(1);
+        libtest_mimic::run(&args, trials).exit()
+    }
+
+    Ok(())
+}
+
+mod linux {
+    use rustix::fd::AsRawFd;
+    use rustix::process::{Pid, WaitOptions, waitpid};
+    use std::io::{self, BufRead, BufReader};
+    use wasmtime::*;
+
+    pub const TESTS: &[(&str, fn() -> Result<()>)] = &[
+        ("smoke", smoke),
+        ("pooling_allocator_reset", pooling_allocator_reset),
+    ];
+
+    fn smoke() -> Result<()> {
+        let engine = Engine::default();
+        let module = Module::new(&engine, r#"(module (func (export "")))"#)?;
+        run_in_child(|| {
+            let mut store = Store::new(&engine, ());
+            let instance = Instance::new(&mut store, &module, &[])?;
+            let export = instance.get_typed_func::<(), ()>(&mut store, "")?;
+            export.call(&mut store, ())?;
+            Ok(())
+        })?;
+        Ok(())
+    }
+
+    fn pooling_allocator_reset() -> Result<()> {
+        let mut pooling = PoolingAllocationConfig::new();
+        pooling.linear_memory_keep_resident(4096);
+        let mut config = Config::new();
+        config.allocation_strategy(pooling);
+        let engine = Engine::new(&config)?;
+        let module = Module::new(
+            &engine,
+            r#"
+                (module
+                    (memory (export "") 1 1)
+                    (data (i32.const 0) "\0a")
+                )
+            "#,
+        )?;
+
+        let assert_pristine = || {
+            let mut store = Store::new(&engine, ());
+            let instance = Instance::new(&mut store, &module, &[])?;
+            let memory = instance.get_memory(&mut store, "").unwrap();
+            let data = memory.data(&store);
+            assert_eq!(data[0], 0x0a);
+            anyhow::Ok((store, memory))
+        };
+        run_in_child(|| {
+            // Allocate a memory, and then mutate it.
+            let (mut store, memory) = assert_pristine()?;
+            let data = memory.data_mut(&mut store);
+            data[0] = 0;
+            drop(store);
+
+            // Allocating the memory again should reuse the same pooling
+            // allocator slot but it should be reset correctly.
+            assert_pristine()?;
+            assert_pristine()?;
+            Ok(())
+        })?;
+        Ok(())
+    }
+
+    fn run_in_child(closure: impl FnOnce() -> Result<()>) -> Result<()> {
+        let (read, write) = io::pipe()?;
+        let child = match unsafe { libc::fork() } {
+            -1 => return Err(io::Error::last_os_error().into()),
+
+            0 => {
+                // If a panic happens, don't let it go above this stack frame.
+                let _bomb = Bomb;
+
+                drop(read);
+                unsafe {
+                    assert!(libc::dup2(write.as_raw_fd(), 1) == 1);
+                    assert!(libc::dup2(write.as_raw_fd(), 2) == 2);
+                }
+                drop(write);
+
+                closure().unwrap();
+                std::process::exit(0);
+            }
+
+            pid => pid,
+        };
+
+        drop(write);
+
+        for line in BufReader::new(read).lines() {
+            println!("CHILD: {}", line?);
+        }
+
+        let (_pid, status) =
+            waitpid(Some(Pid::from_raw(child).unwrap()), WaitOptions::empty())?.unwrap();
+        assert_eq!(status.as_raw(), 0);
+
+        Ok(())
+    }
+
+    struct Bomb;
+
+    impl Drop for Bomb {
+        fn drop(&mut self) {
+            std::process::exit(1);
+        }
+    }
+}