sync.Pool で楽して高速化 - 隙あらば寝る

golang には sync.Pool というライブラリがある。

同じ処理を何度も実行
都度メモリ割り当てが発生

するような場合にサクッとパフォーマンスを向上させられるので紹介する。

メモリ割り当ての処理と、メモリ解放のために動く gc の実行時間を削ることでパフォーマンスを上げることができる。

話を単純にするために以下の仕様で関数を作るとする。

入力: string
出力: 入力を n 回繰り返し、[]byte として返す
- n はグローバル変数、適当に調整する

バージョン1

func func1(in string) (out []byte) {
    buf := &bytes.Buffer{}
    for i := 0; i < n; i++ {
        buf.WriteString(in)
    }
    out = buf.Bytes()
    return
}

bytes.Buffer を作って string を回数分 write する。

最後に Bytes() 関数で方を変換して終わり。

至って普通な実装。

バージョン2

var globalBuf = &bytes.Buffer{}
var globalMutex = sync.Mutex{}
func func2(in string) (out []byte) {
    globalMutex.Lock()
    defer globalMutex.Unlock()
    globalBuf.Reset()
    for i := 0; i < n; i++ {
        globalBuf.WriteString(in)
    }
    out = globalBuf.Bytes()
    return
}

バージョン1では毎回バッファのアロケーションが走るので無駄であると考えて、バッファを共有する実装。

並列実行するとバグってしまうので mutex でロックをとっている。

ここでバージョン1と2を比較する。

バージョン 1 は都度メモリ確保するので遅い
バージョン 2 はロックにより並列処理できないためスケールしない

どちらもイマイチ。

ここで次に考えるのは、バージョン 2 をベースにしたバッファプールの拡張。

バッファとして最初に割り当てておいて、割り当てた数だけ並列実行可能にするという実装が思いつく。

しかし、これは実装が面倒、バッファ管理の struct を作るかな。。。という時に sync.Pool が使える。

バージョン3

var globalPool = sync.Pool{
    New: func() interface{} {
        return &bytes.Buffer{}
    },
}
func func3(in string) (out []byte) {
    buf := globalPool.Get().(*bytes.Buffer)
    defer func() {
        buf.Reset()
        globalPool.Put(buf)
    }()
    for i := 0; i < n; i++ {
        buf.WriteString(in)
    }

    out = buf.Bytes()
    return
}

sync.Pool の New に、バッファを新規作成するための関数を渡す。

sync.Pool は Get したときにプールに空きがなければ New を呼んで作ってくれる。

使い終わったら Put しておくと次の Get で使いまわしてくれる。

内部的には解放したりするが、利用者目線で気にすること無く使える。

以上をベンチにかけると以下の様な結果が得られる。(それぞれ並列で実行)

BenchmarkFunc1-4         1000000              1175 ns/op             132 B/op          1 allocs/op
BenchmarkFunc2-4         2000000               686 ns/op               0 B/op          0 allocs/op
BenchmarkFunc3-4         5000000               407 ns/op               0 B/op          0 allocs/op

ns/op が処理にかかった時間。小さいほど早い。
B/op が割り当てたメモリサイズ。小さい程省メモリ。
allocs/op が割当を行った回数。小さい程アロケーションの回数が少ない。

まとめ

sync.Pool を使う事で簡単に高速化が行える
並列実行できない(シングルスレッドな) v2 のほうが v1 よりも早いので、メモリ操作によるオーバーヘッドは相当に大きい