Justin Azoff

Hi! Random things are here :-)

Ensuring zero allocations in Go tests

A common goal in Go is wanting to ensure that calling a function results in a small number of allocations. Often you want to go even further and ensure that calling a function results in zero allocations.

There are guides on how to write benchmarks that can tell you how many allocations a function is making. However, once you have such a benchmark, it’s up to you to notice if the allocation amount ever changes. There are more complicated tools that can help compare and track benchmark results over time, but those are overkill for ensuring a function does not allocate. The Go testing package includes functionality for this, but accomplishing this is not as obvious as it could be.

Let us consider the following two functions and associated benchmarks:

func TotalLengthGood(ss []string) int {
	n := 0
	for _, s := range ss {
		n += len(s)
	}
	return n
}

func TotalLengthBad(ss []string) int {
	return len(strings.Join(ss, ""))
}

var Result int

func BenchmarkTotalLengthGood(b *testing.B) {
	var r int
	ss := []string{"hello", "world"}
	b.ReportAllocs()
	for n := 0; n < b.N; n++ {
		r = TotalLengthGood(ss)
	}
	Result = r
}

func BenchmarkTotalLengthBad(b *testing.B) {
	var r int
	ss := []string{"hello", "world"}
	b.ReportAllocs()
	for n := 0; n < b.N; n++ {
		r = TotalLengthBad(ss)
	}
	Result = r
}

Running these benchmarks gives these results

BenchmarkTotalLengthGood-4      216503835     5.518 ns/op      0 B/op       0 allocs/op
BenchmarkTotalLengthBad-4       14567442      76.68 ns/op      16 B/op      1 allocs/op

Not only is TotalLengthBad slower, each call to it makes one allocation. TotalLengthGood is faster and has zero allocations, but how do we ensure that this doesn’t break at some point? The Go testing package provides a function called testing.Benchmark that can be used to run an existing benchmark function and return a BenchmarkResult struct. This struct has a few fields and methods, including AllocsPerOp that returns the same value that go test -bench outputs.

We can add new tests that run the existing benchmarks, then assert that AllocsPerOp is zero:

func TestTotalLengthGoodZeroAllocs(t *testing.T) {
	res := testing.Benchmark(BenchmarkTotalLengthGood)
	allocs := res.AllocsPerOp()
	if allocs != 0 {
		t.Fatalf("Expected 0 AllocsPerOp, got %d", allocs)
	}
}

func TestTotalLengthBadZeroAllocs(t *testing.T) {
	res := testing.Benchmark(BenchmarkTotalLengthBad)
	allocs := res.AllocsPerOp()
	if allocs != 0 {
		t.Fatalf("Expected 0 AllocsPerOp, got %d", allocs)
	}
}

Running these new tests results in

=== RUN   TestTotalLengthGoodZeroAllocs
--- PASS: TestTotalLengthGoodZeroAllocs (2.02s)
=== RUN   TestTotalLengthBadZeroAllocs
    main_test.go:57: Expected 0 AllocsPerOp, got 1
--- FAIL: TestTotalLengthBadZeroAllocs (1.50s)

Extracting a reusable helper

The TestTotalLengthGoodZeroAllocs and TestTotalLengthBadZeroAllocs functions above are essentially the same. The common task of running an existing benchmark function and ensuring that the number of allocations is under a threshold can be pulled out into a helper function:

func testBenchmarkAllocs(t *testing.T, f func(b *testing.B), threshold int64) {
	res := testing.Benchmark(f)
	allocs := res.AllocsPerOp()
	if allocs > threshold {
		t.Fatalf("Expected AllocsPerOp <= %d, got %d", threshold, allocs)
	}
}

func TestTotalLengthGoodZeroAllocs(t *testing.T) {
	testBenchmarkAllocs(t, BenchmarkTotalLengthGood, 0)
}

func TestTotalLengthBadZeroAllocs(t *testing.T) {
	testBenchmarkAllocs(t, BenchmarkTotalLengthBad, 0)
}

See also: testing.AllocsPerRun

The testing.AllocsPerRun function can also be used to verify zero allocations. The benefit of using testing.Benchmark is that an existing Benchmark function can be used. The potential downside of using a Benchmark function is the default value for the -benchtime flag is 1s and this will cause each call to testing.Benchmark to take at least 1s. This can be avoided by running tests with go test -benchtime .01s or go test -benchtime 100x.