proto: un-flake TestHasExtensionNoAlloc

Before this change, this test had a flakiness rate of 3 out of 1000 runs
(on Google’s internal testing infrastructure).

After this change, the test is not flaky, even when running with 50_000 runs.

testing.AllocsPerRun is meant to be robust against individual allocations
from background goroutines, but not in the way we called the function.

Change-Id: If99573734fb2de5d7128b5e61c2019867ccbb9b2
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/716520
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Christian Höppner <hoeppi@google.com>

Author: stapelberg

proto/extension_test.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"reflect"
 	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -120,38 +121,51 @@ func TestHasExtensionNoAlloc(t *testing.T) {
 		Corecursive: &testpb.TestAllExtensions{},
 	})
 
+	// Number of runs that testing.AllocsPerRun will do.
+	const runs = 100
+
+	for _, tc := range []struct {
+		name string
+		m    proto.Message
+	}{
+		{name: "Nil", m: nil},
+		{name: "Eager", m: mEager},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			avg := testing.AllocsPerRun(runs, func() {
+				proto.HasExtension(tc.m, testpb.E_OptionalNestedMessage)
+			})
+			if avg != 0 {
+				t.Errorf("proto.HasExtension should not allocate, but allocated %.2fx per run", avg)
+			}
+		})
+	}
+
+	// Lazy initialization only happens once, so we need to allocate enough
+	// copies for all testing.AllocsPerRun callbacks.
+	const copies = runs + 1 // + 1 because testing.AllocsPerRun does a warmup call
 	b, err := proto.Marshal(mEager)
 	if err != nil {
 		t.Fatal(err)
 	}
-	mLazy := &testpb.TestAllExtensions{}
-	if err := proto.Unmarshal(b, mLazy); err != nil {
-		t.Fatal(err)
+	mLazy := make([]*testpb.TestAllExtensions, copies)
+	for i := range copies {
+		mLazy[i] = new(testpb.TestAllExtensions)
+		if err := proto.Unmarshal(b, mLazy[i]); err != nil {
+			t.Fatal(err)
+		}
 	}
-
 	for _, tc := range []struct {
 		name string
-		m    proto.Message
+		m    []*testpb.TestAllExtensions
 	}{
-		{name: "Nil", m: nil},
-		{name: "Eager", m: mEager},
 		{name: "Lazy", m: mLazy},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			// Testing for allocations can be done with `testing.AllocsPerRun`, but it
-			// has some snags that complicate its use for us:
-			//  - It performs a warmup invocation before starting the measurement. We
-			//    want to skip this because lazy initialization only happens once.
-			//  - Despite returning a float64, the returned value is an integer, so <1
-			//    allocations per operation are returned as 0. Therefore, pass runs =
-			//    1.
-			warmup := true
-			avg := testing.AllocsPerRun(1, func() {
-				if warmup {
-					warmup = false
-					return
-				}
-				proto.HasExtension(tc.m, testpb.E_OptionalNestedMessage)
+			var idx atomic.Int64
+			avg := testing.AllocsPerRun(runs, func() {
+				i := int(idx.Add(1))
+				proto.HasExtension(tc.m[i-1], testpb.E_OptionalNestedMessage)
 			})
 			if avg != 0 {
 				t.Errorf("proto.HasExtension should not allocate, but allocated %.2fx per run", avg)