Rust如何兼容读取未压缩文件与gzip压缩文件 避免重复读取

你现有实现的问题根源是：检测gzip头时，GzDecoder外层的BufReader已经预读了文件的部分字节，即使你对原File执行seek回起始位置，也无法恢复已经被BufReader消耗的内容，所以才被迫重复打开文件。

下面给出两种常见的优化实现： #

方案1：全量缓存文件内容（适合绝大多数场景） #

逻辑最简单，只打开一次文件，将内容先读到字节缓冲区再判断格式，对于GB级以下的文件性能无感知：

use std::fs::File;
use std::io::{self, Read};
use flate2::read::GzDecoder;

fn read_file(path: &str) -> io::Result<()> {
    // 单次读取整个文件到内存缓冲区
    let mut file_content = Vec::new();
    File::open(path)?.read_to_end(&mut file_content)?;

    let mut text = String::new();
    // 基于内存缓冲区判断是否为gzip格式
    let mut gz_decoder = GzDecoder::new(&file_content[..]);
    match gz_decoder.header() {
        Some(_) => gz_decoder.read_to_string(&mut text)?,
        None => text = String::from_utf8(file_content).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
    }

    for line in text.lines() {
        println!("{:?}", line);
    }
    Ok(())
}

fn main() {
    read_file("file.txt.gz").expect("文件读取失败");
}

方案2：流式读取（适合超大文件场景） #

如果要处理几十GB级的大文件，不想全量加载到内存，可以通过判断gzip固定魔法数（头两个字节为0x1f 0x8b），再用Chain拼接已经读取的头字节和剩余文件流，全程流式读取内存占用极低：

use std::fs::File;
use std::io::{self, Read, Cursor};
use std::io::Chain;
use flate2::read::GzDecoder;

// 统一封装两种Reader，避免分支重复写读取逻辑
enum GenericReader<R: Read> {
    Gzipped(GzDecoder<Chain<Cursor<[u8; 2]>, R>>),
    Plain(Chain<Cursor<[u8; 2]>, R>)
}

impl<R: Read> Read for GenericReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self {
            GenericReader::Gzipped(r) => r.read(buf),
            GenericReader::Plain(r) => r.read(buf)
        }
    }
}

fn read_file(path: &str) -> io::Result<()> {
    let mut file = File::open(path)?;
    // 先读取前两个字节判断gzip魔法数
    let mut magic_header = [0u8; 2];
    let read_len = file.read(&mut magic_header)?;

    // 构造统一的Reader实例
    let mut reader = if read_len == 2 && magic_header == [0x1f, 0x8b] {
        GenericReader::Gzipped(GzDecoder::new(Cursor::new(magic_header).chain(file)))
    } else {
        GenericReader::Plain(Cursor::new(magic_header).chain(file))
    };

    // 统一读取内容到字符串
    let mut text = String::new();
    reader.read_to_string(&mut text)?;

    for line in text.lines() {
        println!("{:?}", line);
    }
    Ok(())
}

fn main() {
    read_file("file.txt.gz").expect("文件读取失败");
}

两种方案都只需要打开一次文件，没有重复IO开销，你可以根据自己的使用场景选择。

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