【Go】优化文件下载处理：从多级复制到零拷贝流式处理

news2025/5/14 2:52:24

在开发音频处理服务过程中，我们面临一个常见需求：从网络下载音频文件并保存到本地。这个看似简单的操作，实际上有很多优化空间。本文将分享一个逐步优化的过程，展示如何从一个基础实现逐步改进到高效的流式下载方案。

初始实现：读取全部再写入

最初的实现相对简单，但存在明显的内存和IO效率问题：

body, _ := io.ReadAll(res.Body)
defer res.Body.Close()

saveDir, _ := path.Split(writePath)
_ = os.MkdirAll(saveDir, 0776)

fd, _ := os.OpenFile(writePath, os.O_CREATE|os.O_RDWR, 0776)
_, _ = fd.Write(body)
fd.Close()

这个实现存在几个问题：

将整个响应体一次性读入内存，对于大文件会消耗大量内存
忽略了错误处理
执行了两次IO操作：先读取到内存，再写入文件

第一次优化：预分配缓冲区

第一步优化是基于已知文件大小预分配缓冲区，减少内存重新分配：

var buf []byte
if size > 0 {
    buf = make([]byte, 0, size)
    buf, err = io.ReadAll(io.LimitReader(res.Body, size))
} else {
    buf, err = io.ReadAll(res.Body)
}

if err != nil {
    res.Body.Close()
    return define.ReturnError(ecode.CODE_HTTP_DO_REQUEST_FAIL, err, "读取响应内容失败")
}

defer res.Body.Close()
// 使用os.WriteFile简化写入操作
if err = os.WriteFile(writePath, buf, 0776); err != nil {
    return define.ReturnError(ecode.CODE_FILE_WRITE_FAIL, err, "保存语音失败")
}

改进点：

使用io.LimitReader限制读取大小，防止恶意服务器发送过大内容
预分配合适大小的缓冲区，减少内存重新分配
添加错误处理
使用os.WriteFile简化写入操作

第二次优化：使用bytes.Buffer

进一步优化，使用bytes.Buffer更高效地处理内存缓冲：

var buf *bytes.Buffer
if size > 0 {
    buf = bytes.NewBuffer(make([]byte, 0, size))
    _, err = io.Copy(buf, io.LimitReader(res.Body, size))
} else {
    buf = &bytes.Buffer{}
    _, err = io.Copy(buf, res.Body)
}

if err != nil {
    res.Body.Close()
    return define.ReturnError(ecode.CODE_HTTP_DO_REQUEST_FAIL, err, "读取响应内容失败")
}

// 将缓冲区内容写入文件
if err = os.WriteFile(writePath, buf.Bytes(), 0776); err != nil {
    return define.ReturnError(ecode.CODE_FILE_WRITE_FAIL, err, "保存语音失败")
}

改进点：

使用bytes.Buffer替代原始切片，更适合变长数据
使用io.Copy可能比io.ReadAll更高效
保留了预分配缓冲区的优势

最终优化：零拷贝流式处理

最后的优化彻底消除了中间缓冲区，实现了真正的流式处理：

// 创建目录
saveDir, _ := path.Split(writePath)
if err = os.MkdirAll(saveDir, 0776); err != nil {
    return define.ReturnError(ecode.CODE_FILE_WRITE_FAIL, err, "创建目录失败")
}

// 创建文件
fd, err := os.OpenFile(writePath, os.O_CREATE|os.O_WRONLY, 0776)
if err != nil {
    return define.ReturnError(ecode.CODE_FILE_WRITE_FAIL, err, "创建文件失败")
}
defer fd.Close()

// 直接从HTTP响应流写入文件，根据size决定是否限制读取大小
var reader io.Reader = res.Body
if size > 0 {
    reader = io.LimitReader(res.Body, size)
}

_, err = io.Copy(fd, reader)
if err != nil {
    return define.ReturnError(ecode.CODE_FILE_WRITE_FAIL, err, "写入文件失败")
}