自定义默克尔树优化
1. 简介
在文章 技术-文件哈希 中，我用 go 语言实现了基础的默克尔树，并提到了几个优化点。在这篇文章中，我将会实现具体的优化代码：
2. 使用通道代替数组
在这里我表示歉意，之前提到的这个优化点不是很现实，我在实现的时候发现如果采用通道来代替数组，就会导致 合并哈希时乱序，为了避免乱序会用到更多的内存（目前我还没有想出一种占用更少内存的方法，如果大家有什么思路，欢迎讨论），不如就采取这种简单的实现。
3. 使用 sync.Pool 复用缓冲区
3.1 原先的做法
每个 goroutine 都会创建一个大小为 1Mi 的缓冲区，若文件分片数较多（比如 1024 个分片），会瞬间分配 1Gi 内存，导致内存峰值过高，且 频繁的内存分配/回收会增加 GC 压力。
3.2 优化思路
sync.Pool 是 Go 内置的对象池，可复用临时对象，减少内存分配次数和 GC 开销，尤其适合 “创建成本高、复用性强” 的临时对象。
3.3 代码实现
var bufferPool = sync.Pool{
New: func() interface{} {
return make([]byte, chunkSize)
},
}
func MerkleHash(filePath string) (string, error) {
// ...
for i := 0; i < numChunks; i++ {
// ...
buffer := bufferPool.Get().([]byte)
defer bufferPool.Put(buffer)
// 仅读取实际需要的字节数，避免读取多余数据
_, err := file.ReadAt(buffer[:size], offset)
// ...
}
// ...
}
4. 合并哈希时避免临时切片分配
4.1 原先的做法
mergeHashes 中使用 append(a, b...) 会创建临时切片（长度为 32 字节，因为 MD5 哈希是 16 字节，两个拼接后是 32 字节），频繁合并时会产生大量临时对象，增加 GC 压力。
4.2 优化思路
可以预分配固定大小的合并缓冲区（32 字节，16+16），直接将两个哈希写入缓冲区，避免 append 产生的临时切片。由于合并哈希时是单线程操作，这样做也不会产生并发问题。
4.3 代码实现
// 预分配合并用的固定大小缓冲区（32 字节 = 16 字节 MD5 + 16 字节 MD5）
var mergeBuf = make([]byte, 32)
func mergeHashes(a, b []byte) []byte {
// 直接将 a 和 b 写入预分配的缓冲区，无临时切片分配
copy(mergeBuf[:16], a)
copy(mergeBuf[16:], b)
return hashChunk(mergeBuf)
}
5. 控制并发度
5.1 原先的做法
原代码中 goroutine 的数量和和分片的数量一致，若文件极大，会创建大量 goroutine，从而造成以下问题：
- 内存峰值过高：即使复用缓冲区，大量 goroutine 的栈内存 + 哈希结果内存也会很大。
- 操作系统线程调度压力增大：goroutine 过多会导致上下文切换频繁。
- 磁盘 I/O 竞争激烈：过多 goroutine 同时调用
ReadAt会导致磁盘寻道频繁，反而降低读取效率。
5.2 优化思路
使用带缓冲的通道（信号量）控制最大并发数：并发数 = CPU 核心数 × 2。这里与 CPU 核心数相乘的因子最好根据磁盘 I/O 性能调整，比如机械硬盘适合低并发，SSD 可适当提高。
5.3 代码实现
import "runtime"
// 可配置的最大并发数
var maxConcurrency = runtime.NumCPU() * 2
func MerkleHash(filePath string) (string, error) {
// ...
// 信号量通道：控制最大并发数
sem := make(chan struct{}, maxConcurrency)
for i := 0; i < numChunks; i++ {
go func(idx int) {
defer wg.Done()
// 占用信号量（若已达最大并发，会阻塞在此）
sem <- struct{}{}
defer func() { <-sem }() // 释放信号量
// ...
}
}
// ...
}
6. 优化后的整体代码
import (
"crypto/md5"
"encoding/hex"
"os"
"runtime"
"sync"
)
// chunkSize 是分片大小的最大值，我这里设置成 1Mi 了
const chunkSize = 1 * 1024 * 1024
var bufferPool = sync.Pool{
New: func() interface{} {
return make([]byte, chunkSize)
},
}
func hashChunk(data []byte) []byte {
h := md5.New()
h.Write(data)
return h.Sum(nil)
}
// 预分配合并用的固定大小缓冲区（32 字节 = 16 字节 MD5 + 16 字节 MD5）
var mergeBuf = make([]byte, 32)
func mergeHashes(a, b []byte) []byte {
// 直接将 a 和 b 写入预分配的缓冲区，无临时切片分配
copy(mergeBuf[:16], a)
copy(mergeBuf[16:], b)
return hashChunk(mergeBuf)
}
// 可配置的最大并发数
var maxConcurrency = runtime.NumCPU() * 2
func MerkleHash(filePath string) (string, error) {
file, err := os.Open(filePath)
if err != nil {
return "", err
}
defer file.Close()
fileInfo, _ := file.Stat()
fileSize := fileInfo.Size()
numChunks := (int(fileSize) + chunkSize - 1) / chunkSize
// 多线程计算分片哈希
leafHashes := make([][]byte, numChunks)
var wg sync.WaitGroup
wg.Add(numChunks)
// 信号量通道：控制最大并发数
sem := make(chan struct{}, maxConcurrency)
for i := 0; i < numChunks; i++ {
go func(idx int) {
defer wg.Done()
// 占用信号量（若已达最大并发，会阻塞在此）
sem <- struct{}{}
defer func() { <-sem }() // 释放信号量
offset := int64(idx * chunkSize)
size := int64(chunkSize)
if offset+size > fileSize {
size = fileSize - offset
}
buffer := bufferPool.Get().([]byte)
defer bufferPool.Put(buffer)
// 仅读取实际需要的字节数，避免读取多余数据
_, err := file.ReadAt(buffer[:size], offset)
if err != nil {
panic(err)
}
leafHashes[idx] = hashChunk(buffer)
}(i)
}
wg.Wait()
// 单线程两两合并哈希
currLevel := leafHashes
for len(currLevel) > 1 {
nextLevel := make([][]byte, (len(currLevel)+1)/2)
for i := 0; i < len(currLevel); i += 2 {
if i+1 < len(currLevel) {
// 如果剩余大于一个节点
nextLevel[i/2] = mergeHashes(currLevel[i], currLevel[i+1])
} else {
// 如果只剩余一个节点
nextLevel[(i+1)/2] = currLevel[i]
}
}
currLevel = nextLevel
}
rootHash := currLevel[0]
return hex.EncodeToString(rootHash), nil
}
7. 基准测试
对于这次的优化，我们可以运行基准测试来验证优化是否产生效果（基准测试仍然用上一个文章的代码），测试结果如下：
goos: darwin
goarch: arm64
pkg: xxxxxxxxxxxxxxxx/xxxx/xxxx
BenchmarkMerkleHash/size=1Mi-16 801 1328113 ns/op 674 B/op 13 allocs/op
BenchmarkMerkleHash/size=2Mi-16 890 1352327 ns/op 2076 B/op 19 allocs/op
BenchmarkMerkleHash/size=5Mi-16 793 1507144 ns/op 2878 B/op 36 allocs/op
BenchmarkMerkleHash/size=10Mi-16 762 1587954 ns/op 2508 B/op 62 allocs/op
BenchmarkMerkleHash/size=100Mi-16 103 11413764 ns/op 20195 B/op 515 allocs/op
BenchmarkMerkleHash/size=1Gi-16 10 105901417 ns/op 198765 B/op 5141 allocs/op
PASS
ok xxxxxxxxxxxxxxxx/xxxx/xxxx 10.089s
与之前的测试结果相比，发现单次操作的耗时下降了，而且内存的分配次数和分配大小也有所下降，尤其是内存分配大小，这就是使用 sync.Pool 和控制并发度的好处。
8. 总结
本文实现了自定义默克尔树的三个优化点，有一个优化点我没有想到应该如何实现，不过实现其他三个已经使我们的自定义默克尔树更快、使用的内存更少了。我希望本文能够作为一个使用 sync.Pool 和控制并发度的例子，在一定程度上启发看到本文的各位。
