Rework to minimize mmap/munmap

These are fine for single-threaded but were skewing results
when doing parallelism.

Author: posborne

src/main.rs

@@ -2,6 +2,7 @@ use clap::Parser;
 use nix::sys::mman::{mmap_anonymous, MapFlags, ProtFlags};
 use rayon::prelude::*;
 use serde::Serialize;
+use std::marker::PhantomData;
 use std::mem::MaybeUninit;
 use std::slice;
 use std::time::{Duration, Instant};
@@ -47,7 +48,6 @@ enum Strategy {
 struct BenchArgs {
     total_size: usize,
     dirty_fraction: f64,
-    quiet: bool,
     threads: usize,
     processes: usize,
 }
@@ -62,6 +62,72 @@ struct BenchResult {
     pub processes: usize,
 }
 
+impl BenchResult {
+    fn new(args: &BenchArgs, strategy: Strategy, duration: Duration) -> Self {
+        let BenchArgs {
+            total_size,
+            dirty_fraction,
+            threads,
+            processes,
+            ..
+        } = *args;
+        BenchResult {
+            strategy,
+            total_size,
+            dirty_fraction,
+            duration,
+            threads,
+            processes,
+        }
+    }
+}
+
+struct MemoryRegion<'a> {
+    ptr: *mut u8,
+    size: usize,
+    dirty_pct: f64,
+    phantom: PhantomData<&'a [u8]>,
+}
+
+impl<'a> MemoryRegion<'a> {
+    pub fn new(size: usize, dirty_pct: f64, force_resident: bool) -> anyhow::Result<Self> {
+        let prot = ProtFlags::PROT_READ | ProtFlags::PROT_WRITE;
+        let flags = MapFlags::MAP_PRIVATE | MapFlags::MAP_ANONYMOUS;
+        let map = unsafe { mmap_anonymous(None, size.try_into()?, prot, flags) }?;
+        let map = map.as_ptr() as *mut u8;
+
+        if force_resident {
+            let keep_res_slice = unsafe { slice::from_raw_parts_mut(map, size) };
+            keep_res_slice.fill(0);
+        }
+
+        Ok(MemoryRegion {
+            ptr: map,
+            size,
+            dirty_pct,
+            phantom: PhantomData,
+        })
+    }
+
+    pub fn as_mut_slice(&mut self) -> &'a mut [u8] {
+        unsafe { slice::from_raw_parts_mut(self.ptr, self.size) }
+    }
+
+    pub fn make_dirty(&mut self) {
+        let dirty_bytes = (self.size as f64 * self.dirty_pct).round() as usize;
+        if dirty_bytes > 0 {
+            let dirty_slice = unsafe { slice::from_raw_parts_mut(self.ptr, dirty_bytes) };
+            dirty_slice.fill(0xAA);
+        }
+    }
+}
+
+impl<'a> Drop for MemoryRegion<'a> {
+    fn drop(&mut self) {
+        unsafe { libc::munmap(self.ptr as *mut libc::c_void, self.size) };
+    }
+}
+
 macro_rules! qprintln {
     ($condition:expr, $($arg:tt)*) => {
         if !$condition {
@@ -99,7 +165,6 @@ fn main() -> anyhow::Result<()> {
     let bench_args = BenchArgs {
         total_size,
         dirty_fraction,
-        quiet,
         threads: args.threads,
         processes: args.processes,
     };
@@ -124,174 +189,97 @@ fn main() -> anyhow::Result<()> {
     );
     qprintln!(quiet, "------------------------------\n");
 
-    rayon::ThreadPoolBuilder::new()
-        .num_threads(args.threads)
-        .build_global()?;
+    // we want to reduce the number of new regions we create
+    // while still creating enough work to be meaningful
+    let do_memset = || -> anyhow::Result<Vec<BenchResult>> {
+        let mut region = MemoryRegion::new(total_size, args.dirty_fraction, true)?;
+        (0..args.iterations)
+            .map(|_i| run_benchmark_memset(&bench_args, &mut region))
+            .collect::<anyhow::Result<Vec<BenchResult>>>()
+    };
 
-    let results = if args.threads > 1 {
+    let do_madvise = || {
+        let mut region = MemoryRegion::new(total_size, args.dirty_fraction, false)?;
         (0..args.iterations)
-            .into_par_iter()
-            .map(|_i| {
-                [
-                    run_benchmark_memset(&bench_args),
-                    run_benchmark_madvise(&bench_args),
-                    run_benchmark_pagemap_scan(&bench_args),
-                ]
-            })
-            .flatten()
-            .collect::<anyhow::Result<Vec<BenchResult>>>()?
-    } else {
+            .map(|_i| run_benchmark_madvise(&bench_args, &mut region))
+            .collect::<anyhow::Result<Vec<BenchResult>>>()
+    };
+
+    let do_pagemap_scan = || {
+        let mut region = MemoryRegion::new(total_size, args.dirty_fraction, false)?;
         (0..args.iterations)
-            .into_iter()
-            .map(|_i| {
-                [
-                    run_benchmark_memset(&bench_args),
-                    run_benchmark_madvise(&bench_args),
-                    run_benchmark_pagemap_scan(&bench_args),
-                ]
-            })
-            .flatten()
-            .collect::<anyhow::Result<Vec<BenchResult>>>()?
+            .map(|_i| run_benchmark_pagemap_scan(&bench_args, &mut region))
+            .collect::<anyhow::Result<Vec<BenchResult>>>()
     };
 
+    let results: Vec<BenchResult> = (0..args.threads)
+        .into_par_iter()
+        .map(|_| [do_memset(), do_madvise(), do_pagemap_scan()])
+        .flatten()
+        .flatten()
+        .flatten()
+        .collect();
+
     if args.json {
         println!("{}", serde_json::to_string(&results)?);
     }
 
     Ok(())
 }
 
-/// Allocates and dirties memory for a test scenario.
-fn setup_memory(total_size: usize, dirty_fraction: f64, warmup: bool) -> anyhow::Result<*mut u8> {
-    let prot = ProtFlags::PROT_READ | ProtFlags::PROT_WRITE;
-    let flags = MapFlags::MAP_PRIVATE | MapFlags::MAP_ANONYMOUS;
-    let map = unsafe { mmap_anonymous(None, total_size.try_into()?, prot, flags) }?;
-    let map = map.as_ptr() as *mut u8;
-
-    if warmup {
-        let keep_res_slice = unsafe { slice::from_raw_parts_mut(map, total_size) };
-        keep_res_slice.fill(0);
-    }
-
-    // Dirty a fraction of the memory
-    let dirty_bytes = (total_size as f64 * dirty_fraction).round() as usize;
-    if dirty_bytes > 0 {
-        let dirty_slice = unsafe { slice::from_raw_parts_mut(map, dirty_bytes) };
-        dirty_slice.fill(1); // Write something to make pages dirty
-    }
-
-    Ok(map)
-}
-
-fn run_benchmark_memset(args: &BenchArgs) -> anyhow::Result<BenchResult> {
-    let BenchArgs {
-        total_size,
-        dirty_fraction,
-        quiet,
-        threads,
-        processes,
-    } = *args;
-    qprintln!(quiet, "Scenario 1: Naive memset on all pages");
-    let map = setup_memory(total_size, dirty_fraction, true)?;
-    let slice = unsafe { slice::from_raw_parts_mut(map, total_size) };
-
+fn run_benchmark_memset(
+    args: &BenchArgs,
+    region: &mut MemoryRegion,
+) -> anyhow::Result<BenchResult> {
+    region.make_dirty();
     let start = Instant::now();
-    slice.fill(0);
+    region.as_mut_slice().fill(0);
     let duration = start.elapsed();
 
-    qprintln!(quiet, "  Zeroed all {} bytes.", total_size);
-    qprintln!(quiet, "  Time taken: {:?}\n", duration);
-
-    unsafe { libc::munmap(map as *mut libc::c_void, total_size) };
-
-    Ok(BenchResult {
-        strategy: Strategy::MemZero,
-        total_size,
-        dirty_fraction,
-        duration,
-        threads,
-        processes,
-    })
+    Ok(BenchResult::new(args, Strategy::MemZero, duration))
 }
 
-fn run_benchmark_madvise(args: &BenchArgs) -> anyhow::Result<BenchResult> {
-    let BenchArgs {
-        total_size,
-        dirty_fraction,
-        quiet,
-        threads,
-        processes,
-    } = *args;
-    qprintln!(quiet, "Scenario 2: use madvise on all pages");
-    let map = setup_memory(total_size, dirty_fraction, false)?;
-
+fn run_benchmark_madvise(
+    args: &BenchArgs,
+    region: &mut MemoryRegion,
+) -> anyhow::Result<BenchResult> {
+    region.make_dirty();
     let start = Instant::now();
-    let ret = unsafe { libc::madvise(map as *mut libc::c_void, total_size, libc::MADV_DONTNEED) };
+    let ret = unsafe {
+        libc::madvise(
+            region.ptr as *mut libc::c_void,
+            args.total_size,
+            libc::MADV_DONTNEED,
+        )
+    };
     let duration = start.elapsed();
 
     if ret != 0 {
         return Err(std::io::Error::last_os_error().into());
     }
 
-    qprintln!(quiet, "  Called madvise on {total_size} bytes.");
-    qprintln!(quiet, "  Time taken: {:?}\n", duration);
-
-    unsafe { libc::munmap(map as *mut libc::c_void, total_size) };
-
-    Ok(BenchResult {
-        strategy: Strategy::Madvise,
-        total_size,
-        dirty_fraction,
-        duration,
-        threads,
-        processes,
-    })
+    Ok(BenchResult::new(args, Strategy::Madvise, duration))
 }
 
-fn run_benchmark_pagemap_scan(args: &BenchArgs) -> anyhow::Result<BenchResult> {
-    let BenchArgs {
-        total_size,
-        dirty_fraction,
-        quiet,
-        threads,
-        processes,
-    } = *args;
-    qprintln!(quiet, "Scenario 3: Only memset dirty pages");
-    assert_eq!(total_size % rustix::param::page_size(), 0);
-    let map = setup_memory(total_size, dirty_fraction, false)?;
-    let pages = total_size / rustix::param::page_size();
+fn run_benchmark_pagemap_scan(
+    args: &BenchArgs,
+    region: &mut MemoryRegion,
+) -> anyhow::Result<BenchResult> {
+    let pages = args.total_size / rustix::param::page_size();
 
+    region.make_dirty();
     let start = Instant::now();
 
     let mut regions: Box<[MaybeUninit<pagemap::PageRegion>]> = Box::new_uninit_slice(pages);
-    let dirty_pages = pagemap::dirty_pages_in_region(map, total_size, regions.as_mut())?;
-
-    let mut total_zeroed = 0;
+    let dirty_pages =
+        pagemap::dirty_pages_in_region(region.ptr, args.total_size, regions.as_mut())?;
     for dirty_region in dirty_pages.regions {
         let start_ptr = dirty_region.start as *mut u8;
         let len = usize::try_from(dirty_region.end - dirty_region.start)?;
         let region_slice = unsafe { slice::from_raw_parts_mut(start_ptr, len) };
         region_slice.fill(0);
-        total_zeroed += len;
     }
     let duration = start.elapsed();
 
-    qprintln!(
-        quiet,
-        "  Found {} dirty page ranges.",
-        dirty_pages.regions.len()
-    );
-    qprintln!(quiet, "  Zeroed {} bytes.", total_zeroed);
-    qprintln!(quiet, "  Time taken: {:?}\n", duration);
-
-    unsafe { libc::munmap(map as *mut libc::c_void, total_size) };
-
-    Ok(BenchResult {
-        strategy: Strategy::PagemapScan,
-        total_size,
-        dirty_fraction,
-        duration,
-        threads,
-        processes,
-    })
+    Ok(BenchResult::new(args, Strategy::PagemapScan, duration))
 }