allow multiple seeds

Author: nia-e

src/alloc/alloc_bytes.rs

@@ -12,7 +12,7 @@ use crate::helpers::ToU64 as _;
 #[derive(Clone, Copy, Debug)]
 pub enum MiriAllocParams {
     Global,
-    Isolated,
+    Isolated(u64),
 }
 
 /// Allocation bytes that explicitly handle the layout of the data they're storing.
@@ -53,10 +53,10 @@ impl Drop for MiriAllocBytes {
         unsafe {
             match self.params {
                 MiriAllocParams::Global => alloc::dealloc(self.ptr, alloc_layout),
-                MiriAllocParams::Isolated => {
+                MiriAllocParams::Isolated(id) => {
                     #[cfg(target_os = "linux")]
                     {
-                        IsolatedAlloc::dealloc(self.ptr, alloc_layout)
+                        IsolatedAlloc::dealloc(self.ptr, alloc_layout, id)
                     }
                     #[cfg(not(target_os = "linux"))]
                     {
@@ -127,10 +127,10 @@ impl AllocBytes for MiriAllocBytes {
         let alloc_fn = |layout| unsafe {
             match params {
                 MiriAllocParams::Global => alloc::alloc(layout),
-                MiriAllocParams::Isolated => {
+                MiriAllocParams::Isolated(id) => {
                     #[cfg(target_os = "linux")]
                     {
-                        IsolatedAlloc::alloc(layout)
+                        IsolatedAlloc::alloc(layout, id)
                     }
                     #[cfg(not(target_os = "linux"))]
                     {
@@ -156,10 +156,10 @@ impl AllocBytes for MiriAllocBytes {
         let alloc_fn = |layout| unsafe {
             match params {
                 MiriAllocParams::Global => alloc::alloc_zeroed(layout),
-                MiriAllocParams::Isolated => {
+                MiriAllocParams::Isolated(id) => {
                     #[cfg(target_os = "linux")]
                     {
-                        IsolatedAlloc::alloc_zeroed(layout)
+                        IsolatedAlloc::alloc_zeroed(layout, id)
                     }
                     #[cfg(not(target_os = "linux"))]
                     {

src/alloc/isolated_alloc.rs

@@ -1,12 +1,21 @@
 use std::alloc::{self, Layout};
+use std::collections::HashMap;
+use std::hash::{BuildHasherDefault, DefaultHasher};
 use std::sync;
 
 static ALLOCATOR: sync::Mutex<IsolatedAlloc> = sync::Mutex::new(IsolatedAlloc::empty());
 
+pub struct IsolatedAlloc {
+    #[allow(rustc::default_hash_types)]
+    allocators: HashMap<u64, IsolatedAllocInner, BuildHasherDefault<DefaultHasher>>,
+    /// The host (not emulated) page size, or 0 if it has not yet been set.
+    page_size: usize,
+}
+
 /// A dedicated allocator for interpreter memory contents, ensuring they are stored on dedicated
 /// pages (not mixed with Miri's own memory). This is very useful for native-lib mode.
 #[derive(Debug)]
-pub struct IsolatedAlloc {
+pub struct IsolatedAllocInner {
     /// Pointers to page-aligned memory that has been claimed by the allocator.
     /// Every pointer here must point to a page-sized allocation claimed via
     /// the global allocator.
@@ -23,8 +32,6 @@ pub struct IsolatedAlloc {
     /// in 8-byte chunks currently, so the `u8`s are only ever 0 (fully free) or
     /// 255 (fully allocated).
     page_infos: Vec<Box<[u8]>>,
-    /// The host (not emulated) page size.
-    page_size: usize,
 }
 
 // SAFETY: We only point to heap-allocated data
@@ -35,19 +42,80 @@ impl IsolatedAlloc {
     /// allow this function to be `const`; it is updated to its real value on
     /// the first call to `alloc()` or `alloc_zeroed()`.
     const fn empty() -> Self {
-        Self { page_ptrs: Vec::new(), huge_ptrs: Vec::new(), page_infos: Vec::new(), page_size: 0 }
+        #[allow(rustc::default_hash_types)]
+        Self { allocators: HashMap::with_hasher(BuildHasherDefault::new()), page_size: 0 }
+    }
+
+    /// Allocates memory as described in `Layout`, from the pool marked by
+    /// `id`. Note that the same `id` must be used upon calling `dealloc`.
+    ///
+    /// SAFETY: See `alloc::alloc()`.
+    pub unsafe fn alloc(layout: Layout, id: u64) -> *mut u8 {
+        unsafe { Self::alloc_inner(layout, id, false) }
+    }
+
+    /// Same as `alloc()`, but zeroes out data before allocating.
+    ///
+    /// SAFETY: See `alloc::alloc_zeroed()`.
+    pub unsafe fn alloc_zeroed(layout: Layout, id: u64) -> *mut u8 {
+        unsafe { Self::alloc_inner(layout, id, true) }
+    }
+
+    /// Abstracts over `alloc` and `alloc_zeroed`.
+    ///
+    /// SAFETY: See `alloc::alloc()`/`alloc::alloc_zeroed()`.
+    unsafe fn alloc_inner(layout: Layout, id: u64, zeroed: bool) -> *mut u8 {
+        let mut alloc = ALLOCATOR.lock().unwrap();
+        if alloc.page_size == 0 {
+            unsafe {
+                alloc.page_size = libc::sysconf(libc::_SC_PAGESIZE).try_into().unwrap();
+            }
+        }
+        // Store this this AFTER setting the page size
+        let page_size = alloc.page_size;
+
+        match alloc.allocators.get_mut(&id) {
+            Some(alloc_inner) => unsafe { alloc_inner.allocate(layout, page_size, zeroed) },
+            None => {
+                let mut new_inner = IsolatedAllocInner::new();
+                let ret = unsafe { new_inner.allocate(layout, page_size, zeroed) };
+                alloc.allocators.insert(id, new_inner);
+                ret
+            }
+        }
+    }
+
+    /// Deallocates a pointer from the memory pool associated with a given `id`.
+    ///
+    /// SAFETY: See `alloc::dealloc()`, with the extra caveat that `id` must
+    /// correspond to the `id` used upon first allocating the memory.
+    pub unsafe fn dealloc(ptr: *mut u8, layout: Layout, id: u64) {
+        let mut alloc = ALLOCATOR.lock().unwrap();
+        let page_size = alloc.page_size;
+        let alloc_inner = alloc.allocators.get_mut(&id).unwrap();
+        unsafe { alloc_inner.deallocate(ptr, layout, page_size) };
+        // Remove if the machine with that id no longer has any memory
+        if alloc_inner.huge_ptrs.is_empty() && alloc_inner.page_ptrs.is_empty() {
+            alloc.allocators.remove(&id);
+        }
+    }
+}
+
+impl IsolatedAllocInner {
+    /// Creates an empty allocator.
+    const fn new() -> Self {
+        Self { page_ptrs: Vec::new(), huge_ptrs: Vec::new(), page_infos: Vec::new() }
     }
 
     /// Expands the available memory pool by adding one page.
-    fn add_page(&mut self) -> (*mut u8, &mut Box<[u8]>) {
-        assert_ne!(self.page_size, 0);
+    fn add_page(&mut self, page_size: usize) -> (*mut u8, &mut Box<[u8]>) {
+        assert_ne!(page_size, 0);
 
-        let page_layout =
-            unsafe { Layout::from_size_align_unchecked(self.page_size, self.page_size) };
+        let page_layout = unsafe { Layout::from_size_align_unchecked(page_size, page_size) };
         // We don't overwrite the bytes we hand out so make sure they're zeroed by default!
         let page_ptr = unsafe { alloc::alloc(page_layout) };
         // `page_infos` has to be one-eighth of the pagesize per the field docs
-        self.page_infos.push(vec![0u8; self.page_size / 8].into_boxed_slice());
+        self.page_infos.push(vec![0u8; page_size / 8].into_boxed_slice());
         self.page_ptrs.push(page_ptr);
         (page_ptr, self.page_infos.last_mut().unwrap())
     }
@@ -61,38 +129,15 @@ impl IsolatedAlloc {
         (size, align)
     }
 
-    /// Allocates memory as described in `Layout`.
-    ///
-    /// SAFETY: `See alloc::alloc()`
-    #[inline]
-    pub unsafe fn alloc(layout: Layout) -> *mut u8 {
-        let mut alloc = ALLOCATOR.lock().unwrap();
-        unsafe { alloc.alloc_inner(layout, false) }
-    }
-
-    /// Same as `alloc()`, but zeroes out data before allocating.
-    ///
-    /// SAFETY: See `alloc::alloc_zeroed()`
-    pub unsafe fn alloc_zeroed(layout: Layout) -> *mut u8 {
-        let mut alloc = ALLOCATOR.lock().unwrap();
-        unsafe { alloc.alloc_inner(layout, true) }
-    }
-
     /// Abstracts over the logic of `alloc_zeroed` vs `alloc`, as determined by
     /// the `zeroed` argument.
     ///
-    /// SAFETY: See `alloc::alloc()`
-    unsafe fn alloc_inner(&mut self, layout: Layout, zeroed: bool) -> *mut u8 {
-        if self.page_size == 0 {
-            unsafe {
-                self.page_size = libc::sysconf(libc::_SC_PAGESIZE).try_into().unwrap();
-            }
-        }
-
-        if layout.align() > self.page_size || layout.size() > self.page_size {
+    /// SAFETY: See `alloc::alloc()`, with the added restriction that `page_size`
+    /// corresponds to the host pagesize.
+    unsafe fn allocate(&mut self, layout: Layout, page_size: usize, zeroed: bool) -> *mut u8 {
+        if layout.align() > page_size || layout.size() > page_size {
             unsafe { self.alloc_multi_page(layout, zeroed) }
         } else {
-            let page_size = self.page_size;
             for (&mut page, pinfo) in std::iter::zip(&mut self.page_ptrs, &mut self.page_infos) {
                 if let Some(ptr) =
                     unsafe { Self::alloc_from_page(page_size, layout, page, pinfo, zeroed) }
@@ -102,12 +147,12 @@ impl IsolatedAlloc {
             }
 
             // We get here only if there's no space in our existing pages
-            let (page, pinfo) = self.add_page();
+            let (page, pinfo) = self.add_page(page_size);
             unsafe { Self::alloc_from_page(page_size, layout, page, pinfo, zeroed).unwrap() }
         }
     }
 
-    /// Used internally by `alloc_inner` to abstract over some logic.
+    /// Used internally by `allocate` to abstract over some logic.
     ///
     /// SAFETY: `page` must be a page-aligned pointer to an allocated page,
     /// where the allocation is (at least) `page_size` bytes.
@@ -118,7 +163,7 @@ impl IsolatedAlloc {
         pinfo: &mut Box<[u8]>,
         zeroed: bool,
     ) -> Option<*mut u8> {
-        let (size, align) = IsolatedAlloc::normalized_layout(layout);
+        let (size, align) = IsolatedAllocInner::normalized_layout(layout);
 
         for idx in (0..page_size).step_by(align) {
             let idx_pinfo = idx / 8;
@@ -151,31 +196,19 @@ impl IsolatedAlloc {
         ret
     }
 
-    /// Deallocates a pointer from the isolated allocator.
-    ///
-    /// SAFETY: This pointer must have been allocated with `IsolatedAlloc::alloc()`
-    /// (or `alloc_zeroed()`) with the same layout as the one passed.
-    pub unsafe fn dealloc(ptr: *mut u8, layout: Layout) {
-        let mut alloc = ALLOCATOR.lock().unwrap();
-        unsafe {
-            alloc.dealloc_inner(ptr, layout);
-        }
-    }
-
-    /// Same as `dealloc`, but from a specific allocator. Useful for tests.
+    /// Deallocates a pointer from this allocator.
     ///
-    /// SAFETY: See `dealloc` above.
-    unsafe fn dealloc_inner(&mut self, ptr: *mut u8, layout: Layout) {
-        let (size, align) = IsolatedAlloc::normalized_layout(layout);
-
-        if size == 0 || ptr.is_null() {
-            return;
-        }
-
-        let ptr_idx = ptr.addr() % self.page_size;
+    /// SAFETY: This pointer must have been allocated by calling `alloc()` (or
+    /// `alloc_zeroed()`) with the same layout as the one passed on this same
+    /// `IsolatedAllocInner`, and `page_size` must correspond to the host
+    /// pagesize.
+    unsafe fn deallocate(&mut self, ptr: *mut u8, layout: Layout, page_size: usize) {
+        let (size, align) = IsolatedAllocInner::normalized_layout(layout);
+
+        let ptr_idx = ptr.addr() % page_size;
         let page_addr = ptr.addr() - ptr_idx;
 
-        if align > self.page_size || size > self.page_size {
+        if align > page_size || size > page_size {
             unsafe {
                 self.dealloc_multi_page(ptr, layout);
             }
@@ -195,8 +228,7 @@ impl IsolatedAlloc {
         }
 
         let mut free = vec![];
-        let page_layout =
-            unsafe { Layout::from_size_align_unchecked(self.page_size, self.page_size) };
+        let page_layout = unsafe { Layout::from_size_align_unchecked(page_size, page_size) };
         for (idx, pinfo) in self.page_infos.iter().enumerate() {
             if pinfo.iter().all(|p| *p == 0) {
                 free.push(idx);
@@ -240,54 +272,44 @@ mod tests {
 
     #[test]
     fn small_zeroes() {
-        // We can't use the global one because it'll be shared between tests
-        let mut allocator = IsolatedAlloc::empty();
-
         let layout = Layout::from_size_align(256, 32).unwrap();
         // allocate_zeroed
-        let ptr = unsafe { allocator.alloc_inner(layout, true) };
+        let ptr = unsafe { IsolatedAlloc::alloc_zeroed(layout, 0) };
         assert_zeroes(ptr, layout);
         unsafe {
-            allocator.dealloc_inner(ptr, layout);
+            IsolatedAlloc::dealloc(ptr, layout, 0);
         }
     }
 
     #[test]
     fn big_zeroes() {
-        let mut allocator = IsolatedAlloc::empty();
-
         let layout = Layout::from_size_align(16 * 1024, 128).unwrap();
-        let ptr = unsafe { allocator.alloc_inner(layout, true) };
+        let ptr = unsafe { IsolatedAlloc::alloc_zeroed(layout, 1) };
         assert_zeroes(ptr, layout);
         unsafe {
-            allocator.dealloc_inner(ptr, layout);
+            IsolatedAlloc::dealloc(ptr, layout, 1);
         }
     }
 
     #[test]
     fn repeated_allocs() {
-        let mut allocator = IsolatedAlloc::empty();
-
         for sz in (1..=(16 * 1024)).step_by(128) {
             let layout = Layout::from_size_align(sz, 1).unwrap();
-            let ptr = unsafe { allocator.alloc_inner(layout, true) };
+            let ptr = unsafe { IsolatedAlloc::alloc_zeroed(layout, 2) };
             assert_zeroes(ptr, layout);
-            eprintln!("Success at size {sz}");
             unsafe {
                 ptr.write_bytes(255, sz);
-                allocator.dealloc_inner(ptr, layout);
+                IsolatedAlloc::dealloc(ptr, layout, 2);
             }
         }
     }
 
     #[test]
     fn no_overlaps() {
-        let mut allocator = IsolatedAlloc::empty();
-        no_overlaps_inner(&mut allocator);
+        no_overlaps_inner(3);
     }
 
-    // So we can reuse this in `check_leaks`
-    fn no_overlaps_inner(allocator: &mut IsolatedAlloc) {
+    fn no_overlaps_inner(id: u64) {
         // Some random sizes and aligns
         let mut sizes = vec![32; 10];
         sizes.append(&mut vec![15; 4]);
@@ -304,29 +326,28 @@ mod tests {
         let layouts: Vec<_> = std::iter::zip(sizes, aligns)
             .map(|(sz, al)| Layout::from_size_align(sz, al).unwrap())
             .collect();
-        let ptrs: Vec<_> =
-            layouts.iter().map(|layout| unsafe { allocator.alloc_inner(*layout, true) }).collect();
+        let ptrs: Vec<_> = layouts
+            .iter()
+            .map(|layout| unsafe { IsolatedAlloc::alloc_zeroed(*layout, id) })
+            .collect();
 
         for (&ptr, &layout) in std::iter::zip(&ptrs, &layouts) {
             // Make sure we don't allocate overlapping ranges
             unsafe {
                 assert_zeroes(ptr, layout);
                 ptr.write_bytes(255, layout.size());
-                allocator.dealloc_inner(ptr, layout);
+                IsolatedAlloc::dealloc(ptr, layout, id);
             }
         }
     }
 
     #[test]
     fn check_leaks() {
         // Generate some noise first
-        let mut allocator = IsolatedAlloc::empty();
-        no_overlaps_inner(&mut allocator);
+        no_overlaps_inner(4);
+        let alloc = ALLOCATOR.lock().unwrap();
 
-        for pinfo in &allocator.page_infos {
-            for eight_bytes in 0..pinfo.len() {
-                assert_eq!(eight_bytes, 0);
-            }
-        }
+        // Should get auto-deleted if the allocations are empty
+        assert!(!alloc.allocators.contains_key(&4));
     }
 }