Goでハッシュ値計算時にメモリ消費量を抑える方法

初めに

先日、Lambdaでファイルのチェックサムをしていたら、ファイルサイズが大きすぎてout of memoryを起こしたという事象に遭遇しました。
その原因と対策について、書いていきます。

Goでのハッシュ値計算

Goでsha256などを使ってファイルをチェックサムをする時に、ファイルをバイト列にしてcryptoパッケージを使ってハッシュを計算します。
しかし、ファイルの中身をすべてメモリに展開すると、ファイルのサイズ分のメモリを消費してしまうという問題があります。

b, err := ioutil.ReadFile("video.mp4")
if err != nil {
	// error handling
}

hash := sha256.New()
hash.Write(b)
v := hash.Sum(nil)

そこでio.Copyを使ってメモリ消費を抑えてハッシュ値を計算できます。

r, err := os.Open("video.mp4")
if err != nil {
	// error handling
}

hash := sha256.New()
if _, err := io.Copy(hash, r); err != nil {
	// error handling
}

v := hash.Sum(nil)

ベンチマーク

crypto/sha256はio.Writerを実装しているのため、io.Copyを使って効率よくデータを受け取り計算できます。
実際どれくらいメモリ消費量を抑えられるか、次のコードでベンチマークを取ってみます。

package main

import (
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"testing"
)

var good = "a6f1de4f1aba03ff704abfe8264d6b3d5dc4f6a256792f759a80df08b1dfcc42"

func BenchmarkHashBuffer(b *testing.B) {
	for i := 0; i < b.N; i++ {
		input, err := ioutil.ReadFile("xxx.log")
		if err != nil {
			b.Fatal(err)
		}
		hash := sha256.New()
		hash.Write(input)
		got := fmt.Sprintf("%x", hash.Sum(nil))
		if good != got {
			b.Fatal("unexpected sum")
		}
	}

}

func BenchmarkHashStream(b *testing.B) {
	for i := 0; i < b.N; i++ {
		r, err := os.Open("xxx.log")
		if err != nil {
			b.Fatal(err)
		}
		defer r.Close()
		hash := sha256.New()
		if _, err := io.Copy(hash, r); err != nil {
			b.Fatal(err)
		}
		got := fmt.Sprintf("%x", hash.Sum(nil))
		if good != got {
			b.Fatal("unexpected sum")
		}
	}
}

次がベンチマークの結果です。1.2GBのファイルのハッシュ値を算出するのにio.Copyを使った方がメモリ使用量を結構抑えらています。

MacbookPro13% go test -bench . -benchmem
goos: darwin
goarch: arm64
pkg: github.com/skanehira/test/hash
BenchmarkHashBuffer-8          2         800078625 ns/op        1274898200 B/op       16 allocs/op
BenchmarkHashStream-8                  2         661733333 ns/op           33912 B/op         13 allocs/op
PASS
ok      github.com/skanehira/test/hash  4.739s

まとめ

sha256以外にもmd5も同様にio.Copyが使えることは確認しています。他のハッシュ値計算はできるかは確認していないんですが、おそらく可能だと思います。
基本的にio.Copyを使ったほうが省メモリですので、可能な場合は使った方がよいと思います。