Go 1.24 の testing.B.Loop() が便利

Benchmark関数を書くとき、重い初期化処理が必要になるケースがある。 例えばデータベースから大量のデータを読み込むベンチマークを書きたい場合、まずはその大きいデータを持ったテーブルを作る必要がある。

Goのベンチマーク関数は、普通は実行時間がターゲット(デフォルトで1秒)に近づくように b.N を変えながら複数回呼び出される。

func BenchmarkBigData(b *testing.B) {  // 複数回呼び出される
    // 重いセットアップ

    b.ResetTimer()
    for range b.N {
        // 測りたいコード
    }
    b.StopTimer()

    // クリーンアップ
}

このようなコードでは複数回「重いセットアップ」「クリーンアップ」がよばれてしまうので、1秒のベンチマークのためのセットアップ&クリーンアップ時間が数十秒かかる可能性が出てくる。

これを避けるために今まで使っていたのは b.Run() を使う方法だ。 b.Run() を呼び出したら、呼び出した側のベンチマーク関数は1度しか実行されなくなり、b.Run() に渡したベンチマーク関数が複数回よばれる。

func BenchmarkBigData(b *testing.B) {  // 1回だけ呼び出される
    // 重いセットアップ

    b.Run("go", func(b *testing.B) {  // 複数回呼び出される
        for range b.N {
            // 測りたいコード
        }
    })

    // クリーンアップ
}

ただ、ベンチマークを実行するときに "BenchmarkBigData/go" みたいにサブテスト名が追加されてしまうのがダサくなってしまう。 Go 1.24 で追加された b.Loop() を使えばサブテストが不要になる。 ただし b.Run() から for b.Loop() に切り替える時は、 b.ResetTimer() や b.StopTimer() を忘れないこと。

func BenchmarkBigData(b *testing.B) {  // 1回だけ呼び出される
    // 重いセットアップ

    b.ResetTimer()
    for b.Loop() {
        // 測りたいコード
    }
    b.StopTimer()

    // クリーンアップ
}
このブログに乗せているコードは引用を除き CC0 1.0 で提供します。