Rust 大文件处理对比：标准库、Tokio 与内存映射的性能分析

先来看看下面这个代码：

let vb = unsafe {
        VarBuilder::from_mmaped_safetensors(&["./model.safetensors"], DType::F32, &Device::Cpu)?
    };

上面是 candle 加载大模型权重文件的一个用法。这段代码通过内存映射（mmap）和safetensors格式实现了大模型文件的零拷贝加载。这引申出一个问题：

如何优雅地读一个超大文件?

假设一个服务有读文件（比如上传）的需求，一次性把文件读进内存是一个不可取的操作。常规文件上传会将文件从磁盘读到用户空间的缓冲区，这直接造成了内存的消耗。

（正常的大文件上传，应该采用分块上传。）

下面对比一下不同方法处理大文件所占用的内存使用。

统计内存使用量：

struct TrackingAllocator;

static ALLOCATED: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for TrackingAllocator {
    unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
        let ret = unsafe { System.alloc(layout) };
        if !ret.is_null() {
            ALLOCATED.fetch_add(layout.size(), Ordering::SeqCst);
        }
        ret
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
        unsafe { System.dealloc(ptr, layout) };
        ALLOCATED.fetch_sub(layout.size(), Ordering::SeqCst);
    }
}

#[global_allocator]
static GLOBAL: TrackingAllocator = TrackingAllocator;

fn print_memory_usage() {
    println!("Allocated memory: {} bytes", ALLOCATED.load(Ordering::SeqCst));
}

生成 1GB 的文件：

vdd if=/dev/urandom of=testfile.bin bs=1M count=1024
# -rw-rw-r-- 1 staf staf 1.0G  3月 29 07:38 testfile.bin

使用标准库：

/// 该函数用于读取指定文件的内容，并将其写入到另一个文件中，同时打印内存使用情况。
fn normal(file_path: &str)->anyhow::Result<()> {
    let mut file = File::open(file_path)?;
    let mut buffer = Vec::new();
    file.read_to_end(&mut buffer)?; // 读取文件内容到缓冲区

    let mut wf = File::create("testfile222.bin").expect("failed to create the file");
    wf.write_all(&buffer)?;
    print_memory_usage();

    Ok(())
}

调用得到内存使用量如下：

Allocated memory: 1073741824 bytes

使用 异步库 tokio

async fn tokio_io(file_path: &str) -> anyhow::Result<()> {
    // 异步打开文件
    let mut file = tokio::fs::File::open(file_path).await?;
    
    let mut wf = tokio::fs::File::create("testfile222.bin").await?;
    let _ = tokio::io::copy(&mut file, &mut wf).await?;
    print_memory_usage();

    Ok(())
}

调用得到内存使用量如下：

Allocated memory: 172952 bytes

tokio 内部应该有做分块循环缓冲区优化，所以内存使用量要比标准库小得多。

mmap

Memory-mapped（内存映射）是一种通过操作系统将文件或设备直接映射到进程虚拟内存空间的技术 ，使得文件内容可像内存一样直接访问，无需显式调用read/write系统调用。

复制文件大体流程：

[源文件页缓存] → [进程虚拟内存] → [目标文件页缓存]
       ↑--------- 用户态拷贝 ---------↑

由于它不经过用户空间缓冲，所以几乎是没什么内存占用的。

使用 mmap :

cargo add memmap2

/// 文件件映射到内存，并将其内容写入到新文件中。
fn io_mmap(file_path: &str)->anyhow::Result<()> {

    let file = File::open(file_path).expect("failed to open the file");
    let mmap = unsafe { Mmap::map(&file).expect("failed to map the file") };

    let mut wf = File::create("testfile222.bin").expect("failed to create the file");
        let _ = wf.write_all(&mmap[..]);
        print_memory_usage();

    Ok(())
}

调用得到内存使用量如下：

Allocated memory: 1024 bytes # 内存映射

对比（1GB 文件）