address review pt 2

Author: nia-e

src/alloc/isolated_alloc.rs

@@ -36,17 +36,22 @@ impl IsolatedAlloc {
             page_ptrs: Vec::new(),
             huge_ptrs: Vec::new(),
             page_infos: Vec::new(),
+            // SAFETY: `sysconf(_SC_PAGESIZE)` is always safe to call at runtime
+            // See https://www.man7.org/linux/man-pages/man3/sysconf.3.html
             page_size: unsafe { libc::sysconf(libc::_SC_PAGESIZE).try_into().unwrap() },
         }
     }
 
     /// Expands the available memory pool by adding one page.
-    fn add_page(&mut self, page_size: usize) -> (*mut u8, &mut DenseBitSet<usize>) {
-        let page_layout = unsafe { Layout::from_size_align_unchecked(page_size, page_size) };
+    fn add_page(&mut self) -> (*mut u8, &mut DenseBitSet<usize>) {
+        // SAFETY: the system page size must always be a power of 2 and nonzero,
+        // and cannot overflow an isize on the host.
+        let page_layout = unsafe { Layout::from_size_align_unchecked(self.page_size, self.page_size) };
+        // SAFETY: The system page size, which is the layout size, cannot be 0
         let page_ptr = unsafe { alloc::alloc(page_layout) };
         // `page_infos` has to have one bit for each `COMPRESSION_FACTOR`-sized chunk of bytes in the page.
-        assert!(page_size % COMPRESSION_FACTOR == 0);
-        self.page_infos.push(DenseBitSet::new_empty(page_size / COMPRESSION_FACTOR));
+        assert!(self.page_size % COMPRESSION_FACTOR == 0);
+        self.page_infos.push(DenseBitSet::new_empty(self.page_size / COMPRESSION_FACTOR));
         self.page_ptrs.push(page_ptr);
         (page_ptr, self.page_infos.last_mut().unwrap())
     }
@@ -85,13 +90,15 @@ impl IsolatedAlloc {
     ///
     /// SAFETY: See `alloc::alloc()`
     pub unsafe fn alloc(&mut self, layout: Layout) -> *mut u8 {
+        // SAFETY: Upheld by caller
         unsafe { self.allocate(layout, false) }
     }
 
     /// Same as `alloc`, but zeroes out the memory.
     ///
     /// SAFETY: See `alloc::alloc_zeroed()`
     pub unsafe fn alloc_zeroed(&mut self, layout: Layout) -> *mut u8 {
+        // SAFETY: Upheld by caller
         unsafe { self.allocate(layout, true) }
     }
 
@@ -103,9 +110,13 @@ impl IsolatedAlloc {
     unsafe fn allocate(&mut self, layout: Layout, zeroed: bool) -> *mut u8 {
         let (size, align) = IsolatedAlloc::normalized_layout(layout);
         if IsolatedAlloc::is_huge_alloc(size, align, self.page_size) {
+            // SAFETY: Validity of `layout` upheld by caller; we checked that
+            // the size and alignment are appropriate for being a huge alloc
             unsafe { self.alloc_huge(layout, zeroed) }
         } else {
             for (&mut page, pinfo) in std::iter::zip(&mut self.page_ptrs, &mut self.page_infos) {
+                // SAFETY: The value in `self.page_size` is used to allocate
+                // `page`, with page alignment
                 if let Some(ptr) =
                     unsafe { Self::alloc_from_page(self.page_size, layout, page, pinfo, zeroed) }
                 {
@@ -117,7 +128,9 @@ impl IsolatedAlloc {
             let page_size = self.page_size;
             // Add another page and allocate from it; this cannot fail since the
             // new page is empty and we already asserted it fits into a page
-            let (page, pinfo) = self.add_page(page_size);
+            let (page, pinfo) = self.add_page();
+
+            // SAFETY: See comment on `alloc_from_page` above
             unsafe { Self::alloc_from_page(page_size, layout, page, pinfo, zeroed).unwrap() }
         }
     }
@@ -149,6 +162,11 @@ impl IsolatedAlloc {
             let range_avail = !(idx_pinfo..idx_pinfo + size_pinfo).any(|idx| pinfo.contains(idx));
             if range_avail {
                 pinfo.insert_range(idx_pinfo..idx_pinfo + size_pinfo);
+                // SAFETY: We checked the available bytes after `idx` in the call
+                // to `domain_size` above and asserted there are at least `idx +
+                // layout.size()` bytes available and unallocated after it.
+                // `page` must point to the start of the page, so adding `idx`
+                // is safe per the above.
                 unsafe {
                     let ptr = page.add(idx);
                     if zeroed {
@@ -168,6 +186,7 @@ impl IsolatedAlloc {
     /// SAFETY: Same as `alloc()`.
     unsafe fn alloc_huge(&mut self, layout: Layout, zeroed: bool) -> *mut u8 {
         let layout = IsolatedAlloc::huge_normalized_layout(layout, self.page_size);
+        // SAFETY: Upheld by caller
         let ret =
             unsafe { if zeroed { alloc::alloc_zeroed(layout) } else { alloc::alloc(layout) } };
         self.huge_ptrs.push((ret, layout.size()));
@@ -183,6 +202,8 @@ impl IsolatedAlloc {
         let (size, align) = IsolatedAlloc::normalized_layout(layout);
 
         if IsolatedAlloc::is_huge_alloc(size, align, self.page_size) {
+            // SAFETY: Partly upheld by caller, and we checked that the size
+            // and align mean this must have been allocated via `alloc_huge`
             unsafe {
                 self.dealloc_huge(ptr, layout);
             }
@@ -195,8 +216,9 @@ impl IsolatedAlloc {
             // Find the page this allocation belongs to.
             // This could be made faster if the list was sorted -- the allocator isn't fully optimized at the moment.
             let pinfo = std::iter::zip(&mut self.page_ptrs, &mut self.page_infos)
-                .find(|(page, _)| page.addr() == page_addr);
-            let Some((_, pinfo)) = pinfo else {
+                .enumerate()
+                .find(|(_, (page, _))| page.addr() == page_addr);
+            let Some((idx_of_pinfo, (_, pinfo))) = pinfo else {
                 panic!(
                     "Freeing in an unallocated page: {ptr:?}\nHolding pages {:?}",
                     self.page_ptrs
@@ -209,15 +231,18 @@ impl IsolatedAlloc {
                 pinfo.remove(idx);
             }
 
-            // We allocated all the pages with this layout
-            let page_layout =
-                unsafe { Layout::from_size_align_unchecked(self.page_size, self.page_size) };
             // Only 1 page might have been freed after a dealloc, so if it exists,
             // find it and free it (and adjust the vectors)
-            if let Some(free_idx) = self.page_infos.iter().position(|pinfo| pinfo.is_empty()) {
-                self.page_infos.remove(free_idx);
+            if pinfo.is_empty() {
+                // SAFETY: `self.page_size` is always a power of 2 and does not overflow an isize
+                let page_layout =
+                    unsafe { Layout::from_size_align_unchecked(self.page_size, self.page_size) };
+                self.page_infos.remove(idx_of_pinfo);
+                // SAFETY: We checked that there are no outstanding allocations
+                // from us pointing to this page, and we know it was allocated
+                // with this layout
                 unsafe {
-                    alloc::dealloc(self.page_ptrs.remove(free_idx), page_layout);
+                    alloc::dealloc(self.page_ptrs.remove(idx_of_pinfo), page_layout);
                 }
             }
         }
@@ -235,6 +260,7 @@ impl IsolatedAlloc {
             .expect("Freeing unallocated pages");
         // And kick it from the list
         self.huge_ptrs.remove(idx);
+        // SAFETY: Caller ensures validity of the layout
         unsafe {
             alloc::dealloc(ptr, layout);
         }
@@ -247,7 +273,10 @@ mod tests {
 
     /// Helper function to assert that all bytes from `ptr` to `ptr.add(layout.size())`
     /// are zeroes.
-    fn assert_zeroes(ptr: *mut u8, layout: Layout) {
+    /// 
+    /// SAFETY: `ptr` must have been allocated with `layout`.
+    unsafe fn assert_zeroes(ptr: *mut u8, layout: Layout) {
+        // SAFETY: Caller ensures this is valid
         unsafe {
             for ofs in 0..layout.size() {
                 assert_eq!(0, ptr.add(ofs).read());
@@ -261,9 +290,11 @@ mod tests {
         let mut alloc = IsolatedAlloc::new();
         // 256 should be less than the pagesize on *any* system
         let layout = Layout::from_size_align(256, 32).unwrap();
+        // SAFETY: layout size is the constant above, not 0
         let ptr = unsafe { alloc.alloc_zeroed(layout) };
-        assert_zeroes(ptr, layout);
+        // SAFETY: `ptr` was just allocated with `layout`
         unsafe {
+            assert_zeroes(ptr, layout);
             alloc.dealloc(ptr, layout);
         }
     }
@@ -274,9 +305,11 @@ mod tests {
         let mut alloc = IsolatedAlloc::new();
         // 16k is about as big as pages get e.g. on macos aarch64
         let layout = Layout::from_size_align(16 * 1024, 128).unwrap();
+        // SAFETY: layout size is the constant above, not 0
         let ptr = unsafe { alloc.alloc_zeroed(layout) };
-        assert_zeroes(ptr, layout);
+        // SAFETY: `ptr` was just allocated with `layout`
         unsafe {
+            assert_zeroes(ptr, layout);
             alloc.dealloc(ptr, layout);
         }
     }
@@ -289,24 +322,21 @@ mod tests {
         // Try both sub-pagesize allocs and those larger than / equal to a page
         for sz in (1..=(16 * 1024)).step_by(128) {
             let layout = Layout::from_size_align(sz, 1).unwrap();
+            // SAFETY: all sizes in the range above are nonzero as we start from 1
             let ptr = unsafe { alloc.alloc_zeroed(layout) };
-            assert_zeroes(ptr, layout);
+            // SAFETY: `ptr` was just allocated with `layout`, which was used
+            // to bound the access size
             unsafe {
+                assert_zeroes(ptr, layout);
                 ptr.write_bytes(255, sz);
                 alloc.dealloc(ptr, layout);
             }
         }
     }
 
-    /// Checks that allocations of different sizes do not overlap.
-    #[test]
-    fn no_overlaps() {
-        let mut alloc = IsolatedAlloc::new();
-        no_overlaps_inner(&mut alloc);
-    }
-
-    /// Allows us to reuse this bit for `no_overlaps` and `check_leaks`.
-    fn no_overlaps_inner(alloc: &mut IsolatedAlloc) {
+    /// Checks that allocations of different sizes do not overlap. Not a freestanding
+    /// test because we use it as part of `check_leaks()` also.
+    fn no_overlaps(alloc: &mut IsolatedAlloc) {
         // Some random sizes and aligns
         let mut sizes = vec![32; 10];
         sizes.append(&mut vec![15; 4]);
@@ -327,13 +357,18 @@ mod tests {
         let layouts: Vec<_> = std::iter::zip(sizes, aligns)
             .map(|(sz, al)| Layout::from_size_align(sz, al).unwrap())
             .collect();
+        // SAFETY: all sizes specified in `sizes` are nonzero
         let ptrs: Vec<_> =
             layouts.iter().map(|layout| unsafe { alloc.alloc_zeroed(*layout) }).collect();
 
         for (&ptr, &layout) in std::iter::zip(&ptrs, &layouts) {
             // We requested zeroed allocations, so check that that's true
             // Then write to the end of the current size, so if the allocs
-            // overlap (or the zeroing is wrong) then `assert_zeroes` will panic
+            // overlap (or the zeroing is wrong) then `assert_zeroes` will panic.
+            // Also check that the alignment we asked for was respected
+            assert_eq!(ptr.addr().rem_euclid(layout.align()), 0);
+            // SAFETY: each `ptr` was allocated with its corresponding `layout`,
+            // which is used to bound the access size
             unsafe {
                 assert_zeroes(ptr, layout);
                 ptr.write_bytes(255, layout.size());
@@ -346,9 +381,9 @@ mod tests {
     #[test]
     fn check_leaks() {
         let mut alloc = IsolatedAlloc::new();
-
         // Generate some noise first so leaks can manifest
-        no_overlaps_inner(&mut alloc);
+        no_overlaps(&mut alloc);
+
         // And then verify that no memory was leaked
         assert!(alloc.page_ptrs.is_empty() && alloc.huge_ptrs.is_empty());
     }