Merge pull request #523 from morrisonlevi/oversized-allocations

feat: recognize and use over sized allocations

Author: Amanieu

Cargo.toml

@@ -47,6 +47,9 @@ serde_test = "1.0"
 doc-comment = "0.3.1"
 bumpalo = { version = "3.13.0", features = ["allocator-api2"] }
 
+[target.'cfg(unix)'.dev-dependencies]
+libc = "0.2.155"
+
 [features]
 default = ["default-hasher", "inline-more", "allocator-api2", "equivalent", "raw-entry"]
 

benches/with_capacity.rs

@@ -0,0 +1,38 @@
+#![feature(test)]
+
+extern crate test;
+
+use hashbrown::HashMap;
+use test::{black_box, Bencher};
+
+type Map<K, V> = HashMap<K, V>;
+
+macro_rules! bench_with_capacity {
+    ($name:ident, $cap:expr) => {
+        #[bench]
+        fn $name(b: &mut Bencher) {
+            b.iter(|| {
+                // Construct a new empty map with a given capacity and return it to avoid
+                // being optimized away. Dropping it measures allocation + minimal setup.
+                let m: Map<usize, usize> = Map::with_capacity($cap);
+                black_box(m)
+            });
+        }
+    };
+}
+
+bench_with_capacity!(with_capacity_000000, 0);
+bench_with_capacity!(with_capacity_000001, 1);
+bench_with_capacity!(with_capacity_000003, 3);
+bench_with_capacity!(with_capacity_000007, 7);
+bench_with_capacity!(with_capacity_000008, 8);
+bench_with_capacity!(with_capacity_000016, 16);
+bench_with_capacity!(with_capacity_000032, 32);
+bench_with_capacity!(with_capacity_000064, 64);
+bench_with_capacity!(with_capacity_000128, 128);
+bench_with_capacity!(with_capacity_000256, 256);
+bench_with_capacity!(with_capacity_000512, 512);
+bench_with_capacity!(with_capacity_001024, 1024);
+bench_with_capacity!(with_capacity_004096, 4096);
+bench_with_capacity!(with_capacity_016384, 16384);
+bench_with_capacity!(with_capacity_065536, 65536);

src/map.rs

@@ -6631,3 +6631,136 @@ mod test_map {
         );
     }
 }
+
+#[cfg(all(test, unix, any(feature = "nightly", feature = "allocator-api2")))]
+mod test_map_with_mmap_allocations {
+    use super::HashMap;
+    use crate::raw::prev_pow2;
+    use core::alloc::Layout;
+    use core::ptr::{null_mut, NonNull};
+
+    #[cfg(feature = "nightly")]
+    use core::alloc::{AllocError, Allocator};
+
+    #[cfg(all(feature = "allocator-api2", not(feature = "nightly")))]
+    use allocator_api2::alloc::{AllocError, Allocator};
+
+    /// This is not a production quality allocator, just good enough for
+    /// some basic tests.
+    #[derive(Clone, Copy, Debug)]
+    struct MmapAllocator {
+        /// Guarantee this is a power of 2.
+        page_size: usize,
+    }
+
+    impl MmapAllocator {
+        fn new() -> Result<Self, AllocError> {
+            let result = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
+            if result < 1 {
+                return Err(AllocError);
+            }
+
+            let page_size = result as usize;
+            if !page_size.is_power_of_two() {
+                Err(AllocError)
+            } else {
+                Ok(Self { page_size })
+            }
+        }
+
+        fn fit_to_page_size(&self, n: usize) -> Result<usize, AllocError> {
+            // If n=0, give a single page (wasteful, I know).
+            let n = if n == 0 { self.page_size } else { n };
+
+            match n & (self.page_size - 1) {
+                0 => Ok(n),
+                rem => n.checked_add(self.page_size - rem).ok_or(AllocError),
+            }
+        }
+    }
+
+    unsafe impl Allocator for MmapAllocator {
+        fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+            if layout.align() > self.page_size {
+                return Err(AllocError);
+            }
+
+            let null = null_mut();
+            let len = self.fit_to_page_size(layout.size())? as libc::size_t;
+            let prot = libc::PROT_READ | libc::PROT_WRITE;
+            let flags = libc::MAP_PRIVATE | libc::MAP_ANON;
+            let addr = unsafe { libc::mmap(null, len, prot, flags, -1, 0) };
+
+            // mmap returns MAP_FAILED on failure, not Null.
+            if addr == libc::MAP_FAILED {
+                return Err(AllocError);
+            }
+
+            match NonNull::new(addr.cast()) {
+                Some(data) => {
+                    // SAFETY: this is NonNull::slice_from_raw_parts.
+                    Ok(unsafe {
+                        NonNull::new_unchecked(core::ptr::slice_from_raw_parts_mut(
+                            data.as_ptr(),
+                            len,
+                        ))
+                    })
+                }
+
+                // This branch shouldn't be taken in practice, but since we
+                // cannot return null as a valid pointer in our type system,
+                // we attempt to handle it.
+                None => {
+                    _ = unsafe { libc::munmap(addr, len) };
+                    Err(AllocError)
+                }
+            }
+        }
+
+        unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+            // If they allocated it with this layout, it must round correctly.
+            let size = self.fit_to_page_size(layout.size()).unwrap();
+            let _result = libc::munmap(ptr.as_ptr().cast(), size);
+            debug_assert_eq!(0, _result)
+        }
+    }
+
+    #[test]
+    fn test_tiny_allocation_gets_rounded_to_page_size() {
+        let alloc = MmapAllocator::new().unwrap();
+        let mut map: HashMap<usize, (), _, _> = HashMap::with_capacity_in(1, alloc);
+
+        // Size of an element plus its control byte.
+        let rough_bucket_size = core::mem::size_of::<(usize, ())>() + 1;
+
+        // Accounting for some misc. padding that's likely in the allocation
+        // due to rounding to group width, etc.
+        let overhead = 3 * core::mem::size_of::<usize>();
+        let num_buckets = (alloc.page_size - overhead) / rough_bucket_size;
+        // Buckets are always powers of 2.
+        let min_elems = prev_pow2(num_buckets);
+        // Real load-factor is 7/8, but this is a lower estimation, so 1/2.
+        let min_capacity = min_elems >> 1;
+        let capacity = map.capacity();
+        assert!(
+            capacity >= min_capacity,
+            "failed: {capacity} >= {min_capacity}"
+        );
+
+        // Fill it up.
+        for i in 0..capacity {
+            map.insert(i, ());
+        }
+        // Capacity should not have changed and it should be full.
+        assert_eq!(capacity, map.len());
+        assert_eq!(capacity, map.capacity());
+
+        // Alright, make it grow.
+        map.insert(capacity, ());
+        assert!(
+            capacity < map.capacity(),
+            "failed: {capacity} < {}",
+            map.capacity()
+        );
+    }
+}

src/raw/alloc.rs

@@ -15,9 +15,9 @@ mod inner {
     use core::ptr::NonNull;
 
     #[allow(clippy::map_err_ignore)]
-    pub(crate) fn do_alloc<A: Allocator>(alloc: &A, layout: Layout) -> Result<NonNull<u8>, ()> {
+    pub(crate) fn do_alloc<A: Allocator>(alloc: &A, layout: Layout) -> Result<NonNull<[u8]>, ()> {
         match alloc.allocate(layout) {
-            Ok(ptr) => Ok(ptr.as_non_null_ptr()),
+            Ok(ptr) => Ok(ptr),
             Err(_) => Err(()),
         }
     }
@@ -38,9 +38,9 @@ mod inner {
     use core::ptr::NonNull;
 
     #[allow(clippy::map_err_ignore)]
-    pub(crate) fn do_alloc<A: Allocator>(alloc: &A, layout: Layout) -> Result<NonNull<u8>, ()> {
+    pub(crate) fn do_alloc<A: Allocator>(alloc: &A, layout: Layout) -> Result<NonNull<[u8]>, ()> {
         match alloc.allocate(layout) {
-            Ok(ptr) => Ok(ptr.cast()),
+            Ok(ptr) => Ok(ptr),
             Err(_) => Err(()),
         }
     }
@@ -61,7 +61,7 @@ mod inner {
 
     #[allow(clippy::missing_safety_doc)] // not exposed outside of this crate
     pub unsafe trait Allocator {
-        fn allocate(&self, layout: Layout) -> Result<NonNull<u8>, ()>;
+        fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, ()>;
         unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout);
     }
 
@@ -70,8 +70,19 @@ mod inner {
 
     unsafe impl Allocator for Global {
         #[inline]
-        fn allocate(&self, layout: Layout) -> Result<NonNull<u8>, ()> {
-            unsafe { NonNull::new(alloc(layout)).ok_or(()) }
+        fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, ()> {
+            match unsafe { NonNull::new(alloc(layout)) } {
+                Some(data) => {
+                    // SAFETY: this is NonNull::slice_from_raw_parts.
+                    Ok(unsafe {
+                        NonNull::new_unchecked(core::ptr::slice_from_raw_parts_mut(
+                            data.as_ptr(),
+                            layout.size(),
+                        ))
+                    })
+                }
+                None => Err(()),
+            }
         }
         #[inline]
         unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
@@ -86,7 +97,7 @@ mod inner {
         }
     }
 
-    pub(crate) fn do_alloc<A: Allocator>(alloc: &A, layout: Layout) -> Result<NonNull<u8>, ()> {
+    pub(crate) fn do_alloc<A: Allocator>(alloc: &A, layout: Layout) -> Result<NonNull<[u8]>, ()> {
         alloc.allocate(layout)
     }
 }