如何对Rust标准库HashMap进行带/不带2MiB大页的基准测试并控制大页使用？

要在Rust中让指定的HashMap使用Linux大页内存（而非全局透明大页），核心思路是利用Nightly版的分配器API实现自定义大页分配器，为目标HashMap单独指定该分配器。以下是具体实现步骤和注意事项：

前提准备 #

1. 启用Rust Nightly版本：依赖allocator_api特性，该特性仅在Nightly分支中提供。
2. 系统预留大页：Linux需要提前配置预留的大页数量，否则大页分配会失败。以2MB大页为例：

   # 预留1024个2MB大页（共2GB内存）
   sudo sh -c 'echo 1024 > /sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages'
   # 禁用透明大页，避免干扰测试
   sudo sh -c 'echo never > /sys/kernel/mm/transparent_hugepage/enabled'
   

实现自定义大页分配器 #

通过实现GlobalAlloc trait，创建一个优先使用大页分配内存的分配器，分配失败时回退到系统默认分配器：

Cargo.toml配置 #

[package]
name = "hashmap-hugepage-bench"
version = "0.1.0"
edition = "2021"

[dependencies]
libc = "0.2"
# 可选：用于更专业的基准测试
# criterion = "0.5"

[features]
nightly = ["allocator_api"]

核心代码实现 #

#![cfg_attr(feature = "nightly", feature(allocator_api))]
extern crate libc;

use libc::{c_void, mmap, munmap, MAP_ANONYMOUS, MAP_HUGETLB, MAP_PRIVATE, PROT_READ, PROT_WRITE};
use std::collections::HashMap;
use std::alloc::{GlobalAlloc, Layout, System};

// 自定义大页分配器，优先使用2MB大页
struct HugePageAllocator;

const HUGE_PAGE_SIZE: usize = 2 * 1024 * 1024; // 2MB标准大页大小

unsafe impl GlobalAlloc for HugePageAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // 将分配大小向上对齐到大页的整数倍
        let alloc_size = layout.size().next_multiple_of(HUGE_PAGE_SIZE)
            .unwrap_or_else(|_| panic!("Allocation size exceeds huge page limits"));
        // 大页内存的对齐要求为自身大小，取布局对齐和大页大小的最大值
        let align = layout.align().max(HUGE_PAGE_SIZE);

        // 调用mmap申请大页内存
        let ptr = mmap(
            std::ptr::null_mut(),
            alloc_size,
            PROT_READ | PROT_WRITE,
            MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB,
            -1,
            0,
        );

        // 大页分配失败时，回退到系统默认分配器
        if ptr == libc::MAP_FAILED {
            System.alloc(layout)
        } else {
            ptr as *mut u8
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let alloc_size = layout.size().next_multiple_of(HUGE_PAGE_SIZE)
            .unwrap_or_else(|_| panic!("Deallocation size exceeds huge page limits"));
        
        // 释放大页内存，失败则回退到系统分配器
        if munmap(ptr as *mut c_void, alloc_size) != 0 {
            System.dealloc(ptr, layout);
        }
    }
}

fn main() {
    #[cfg(feature = "nightly")]
    {
        // 为单个HashMap指定大页分配器
        let mut hugepage_hashmap: HashMap<i32, String, _, _> = HashMap::with_hasher_and_allocator(
            std::collections::hash_map::RandomState::new(),
            HugePageAllocator,
        );

        // 填充测试数据
        for i in 0..1_000_000 {
            hugepage_hashmap.insert(i, format!("value_{}", i));
        }

        // 测试查找性能
        let start = std::time::Instant::now();
        for i in 0..1_000_000 {
            let _ = hugepage_hashmap.get(&i);
        }
        println!("大页HashMap查找耗时: {:?}", start.elapsed());
    }

    // 默认分配器的HashMap作为性能对照
    let mut std_hashmap = HashMap::new();
    for i in 0..1_000_000 {
        std_hashmap.insert(i, format!("value_{}", i));
    }

    let start = std::time::Instant::now();
    for i in 0..1_000_000 {
        let _ = std_hashmap.get(&i);
    }
    println!("标准HashMap查找耗时: {:?}", start.elapsed());
}

编译与运行 #

# 使用Nightly编译并启用allocator_api特性
cargo +nightly run --features nightly

基准测试优化建议 #

如果需要更精准的性能数据，推荐使用criterion库编写基准测试：

• 自动处理预热、多次运行和统计分析
• 支持对比不同实现的性能差异

示例基准测试代码（需添加criterion依赖）：

#[cfg(feature = "nightly")]
use criterion::{criterion_group, criterion_main, Criterion};

#[cfg(feature = "nightly")]
fn bench_hashmap_lookup(c: &mut Criterion) {
    let mut hugepage_hashmap: HashMap<i32, String, _, _> = HashMap::with_hasher_and_allocator(
        std::collections::hash_map::RandomState::new(),
        HugePageAllocator,
    );
    for i in 0..1_000_000 {
        hugepage_hashmap.insert(i, format!("value_{}", i));
    }

    let mut std_hashmap = HashMap::new();
    for i in 0..1_000_000 {
        std_hashmap.insert(i, format!("value_{}", i));
    }

    c.bench_function("HugePage HashMap Lookup", |b| b.iter(|| {
        for i in 0..1_000 {
            let _ = hugepage_hashmap.get(&i);
        }
    }));

    c.bench_function("Standard HashMap Lookup", |b| b.iter(|| {
        for i in 0..1_000 {
            let _ = std_hashmap.get(&i);
        }
    }));
}

#[cfg(feature = "nightly")]
criterion_group!(benches, bench_hashmap_lookup);
#[cfg(feature = "nightly")]
criterion_main!(benches);

关键注意事项 #

1. 大页对齐要求：所有分配的内存大小必须是大页大小的整数倍，否则mmap会失败，代码中通过next_multiple_of处理对齐。
2. 回退逻辑：添加系统分配器回退，避免因大页不足导致程序崩溃。
3. 测试环境一致性：测试时关闭其他占用内存的程序，确保大页内存充足，避免系统自动换页影响性能。

内容的提问来源于stack exchange，提问作者ajp