From 7535c06fc4d3a26374f2e5b65a29878f06d35331 Mon Sep 17 00:00:00 2001
From: Madhav Bissa <mbissa@gmail.com>
Date: Wed, 25 Mar 2026 16:54:26 +0000
Subject: [PATCH 1/6] balancer/pickfirst: fix flaky metrics tests by awaiting
 async metrics

Fixed a race condition in pickfirst metrics tests where the assertions would
immediately poll the stats.TestMetricsRecorder. Depending on the goroutine
scheduling, asynchronous metrics like 'grpc.subchannel.open_connections'
would occasionally not be fully recorded before the test checked them,
causing flaky failures. This introduces an awaitMetric helper to wait for the
metrics to reach their expected values.
---
 balancer/pickfirst/metrics_test.go | 74 ++++++++++++++----------------
 1 file changed, 34 insertions(+), 40 deletions(-)

diff --git a/balancer/pickfirst/metrics_test.go b/balancer/pickfirst/metrics_test.go
index f5d39d6b3a66..c6ed1a99b208 100644
--- a/balancer/pickfirst/metrics_test.go
+++ b/balancer/pickfirst/metrics_test.go
@@ -97,41 +97,26 @@ func (s) TestPickFirstMetrics(t *testing.T) {
 		t.Fatalf("EmptyCall() failed: %v", err)
 	}
 
-	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_succeeded"); got != 1 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_succeeded", got, 1)
-	}
-	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_failed"); got != 0 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_failed", got, 0)
-	}
-	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 0 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 0)
-	}
-
-	// Checking for subchannel metrics as well
-	if got, _ := tmr.Metric("grpc.subchannel.connection_attempts_succeeded"); got != 1 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.subchannel.connection_attempts_succeeded", got, 1)
-	}
-	if got, _ := tmr.Metric("grpc.subchannel.connection_attempts_failed"); got != 0 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.subchannel.connection_attempts_failed", got, 0)
-	}
-	if got, _ := tmr.Metric("grpc.subchannel.disconnections"); got != 0 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.subchannel.disconnections", got, 0)
-	}
-	if got, _ := tmr.Metric("grpc.subchannel.open_connections"); got != 1 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.subchannel.open_connections", got, 1)
+	for _, metric := range []struct {
+		name string
+		want float64
+	}{
+		{"grpc.lb.pick_first.connection_attempts_succeeded", 1},
+		{"grpc.lb.pick_first.connection_attempts_failed", 0},
+		{"grpc.lb.pick_first.disconnections", 0},
+		{"grpc.subchannel.connection_attempts_succeeded", 1},
+		{"grpc.subchannel.connection_attempts_failed", 0},
+		{"grpc.subchannel.disconnections", 0},
+		{"grpc.subchannel.open_connections", 1},
+	} {
+		awaitMetric(ctx, t, tmr, metric.name, metric.want)
 	}
 
 	ss.Stop()
 	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
-	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 1 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 1)
-	}
-	if got, _ := tmr.Metric("grpc.subchannel.disconnections"); got != 1 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.subchannel.disconnections", got, 1)
-	}
-	if got, _ := tmr.Metric("grpc.subchannel.open_connections"); got != -1 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.subchannel.open_connections", got, -1)
-	}
+	awaitMetric(ctx, t, tmr, "grpc.lb.pick_first.disconnections", 1)
+	awaitMetric(ctx, t, tmr, "grpc.subchannel.disconnections", 1)
+	awaitMetric(ctx, t, tmr, "grpc.subchannel.open_connections", -1)
 }
 
 // TestPickFirstMetricsFailure tests the connection attempts failed metric. It
@@ -161,15 +146,9 @@ func (s) TestPickFirstMetricsFailure(t *testing.T) {
 		t.Fatalf("EmptyCall() passed when expected to fail")
 	}
 
-	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_succeeded"); got != 0 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_succeeded", got, 0)
-	}
-	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_failed"); got != 1 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_failed", got, 1)
-	}
-	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 0 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 0)
-	}
+	awaitMetric(ctx, t, tmr, "grpc.lb.pick_first.connection_attempts_succeeded", 0)
+	awaitMetric(ctx, t, tmr, "grpc.lb.pick_first.connection_attempts_failed", 1)
+	awaitMetric(ctx, t, tmr, "grpc.lb.pick_first.disconnections", 0)
 }
 
 // TestPickFirstMetricsE2E tests the pick first metrics end to end. It
@@ -453,3 +432,18 @@ func runDisconnectLabelTest(t *testing.T, wantLabel string, triggerFunc func(*st
 
 	t.Fatalf("Error waiting for metrics grpc.subchannel.disconnections: %v", ctx.Err())
 }
+
+func awaitMetric(ctx context.Context, t *testing.T, tmr *stats.TestMetricsRecorder, name string, want float64) {
+	t.Helper()
+	for {
+		got, _ := tmr.Metric(name)
+		if got == want {
+			return
+		}
+		select {
+		case <-ctx.Done():
+			t.Fatalf("Timeout waiting for expected data for metric %v, got: %v, want: %v", name, got, want)
+		case <-time.After(10 * time.Millisecond):
+		}
+	}
+}

From 0c936332075a922575badb5dab332f54d504070d Mon Sep 17 00:00:00 2001
From: Madhav Bissa <mbissa@gmail.com>
Date: Thu, 26 Mar 2026 08:48:48 +0000
Subject: [PATCH 2/6] balancer/pickfirst: revert pick_first metrics to
 synchronous checks

---
 balancer/pickfirst/metrics_test.go | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/balancer/pickfirst/metrics_test.go b/balancer/pickfirst/metrics_test.go
index c6ed1a99b208..007b8bfcca5b 100644
--- a/balancer/pickfirst/metrics_test.go
+++ b/balancer/pickfirst/metrics_test.go
@@ -97,13 +97,20 @@ func (s) TestPickFirstMetrics(t *testing.T) {
 		t.Fatalf("EmptyCall() failed: %v", err)
 	}
 
+	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_succeeded"); got != 1 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_succeeded", got, 1)
+	}
+	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_failed"); got != 0 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_failed", got, 0)
+	}
+	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 0 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 0)
+	}
+
 	for _, metric := range []struct {
 		name string
 		want float64
 	}{
-		{"grpc.lb.pick_first.connection_attempts_succeeded", 1},
-		{"grpc.lb.pick_first.connection_attempts_failed", 0},
-		{"grpc.lb.pick_first.disconnections", 0},
 		{"grpc.subchannel.connection_attempts_succeeded", 1},
 		{"grpc.subchannel.connection_attempts_failed", 0},
 		{"grpc.subchannel.disconnections", 0},
@@ -114,7 +121,9 @@ func (s) TestPickFirstMetrics(t *testing.T) {
 
 	ss.Stop()
 	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
-	awaitMetric(ctx, t, tmr, "grpc.lb.pick_first.disconnections", 1)
+	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 1 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 1)
+	}
 	awaitMetric(ctx, t, tmr, "grpc.subchannel.disconnections", 1)
 	awaitMetric(ctx, t, tmr, "grpc.subchannel.open_connections", -1)
 }
@@ -146,9 +155,15 @@ func (s) TestPickFirstMetricsFailure(t *testing.T) {
 		t.Fatalf("EmptyCall() passed when expected to fail")
 	}
 
-	awaitMetric(ctx, t, tmr, "grpc.lb.pick_first.connection_attempts_succeeded", 0)
-	awaitMetric(ctx, t, tmr, "grpc.lb.pick_first.connection_attempts_failed", 1)
-	awaitMetric(ctx, t, tmr, "grpc.lb.pick_first.disconnections", 0)
+	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_succeeded"); got != 0 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_succeeded", got, 0)
+	}
+	if got, _ := tmr.Metric("grpc.lb.pick_first.connection_attempts_failed"); got != 1 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.connection_attempts_failed", got, 1)
+	}
+	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 0 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 0)
+	}
 }
 
 // TestPickFirstMetricsE2E tests the pick first metrics end to end. It

From 2848dac1a972c909d4a4eaa3db4aebfc2319490a Mon Sep 17 00:00:00 2001
From: Madhav Bissa <mbissa@gmail.com>
Date: Thu, 26 Mar 2026 12:55:07 +0000
Subject: [PATCH 3/6] review comments

---
 balancer/pickfirst/metrics_test.go | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/balancer/pickfirst/metrics_test.go b/balancer/pickfirst/metrics_test.go
index 007b8bfcca5b..d1f942d3dd98 100644
--- a/balancer/pickfirst/metrics_test.go
+++ b/balancer/pickfirst/metrics_test.go
@@ -107,13 +107,18 @@ func (s) TestPickFirstMetrics(t *testing.T) {
 		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 0)
 	}
 
+	if got, _ := tmr.Metric("grpc.subchannel.connection_attempts_failed"); got != 0 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.subchannel.connection_attempts_failed", got, 0)
+	}
+	if got, _ := tmr.Metric("grpc.subchannel.disconnections"); got != 0 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.subchannel.disconnections", got, 0)
+	}
+
 	for _, metric := range []struct {
 		name string
 		want float64
 	}{
 		{"grpc.subchannel.connection_attempts_succeeded", 1},
-		{"grpc.subchannel.connection_attempts_failed", 0},
-		{"grpc.subchannel.disconnections", 0},
 		{"grpc.subchannel.open_connections", 1},
 	} {
 		awaitMetric(ctx, t, tmr, metric.name, metric.want)
@@ -124,7 +129,9 @@ func (s) TestPickFirstMetrics(t *testing.T) {
 	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 1 {
 		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 1)
 	}
-	awaitMetric(ctx, t, tmr, "grpc.subchannel.disconnections", 1)
+	if got, _ := tmr.Metric("grpc.subchannel.disconnections"); got != 1 {
+		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.subchannel.disconnections", got, 1)
+	}
 	awaitMetric(ctx, t, tmr, "grpc.subchannel.open_connections", -1)
 }
 

From c969a89c6dcfdef3f48967a7b6643886a219e544 Mon Sep 17 00:00:00 2001
From: Madhav Bissa <mbissa@gmail.com>
Date: Wed, 29 Apr 2026 10:04:02 +0000
Subject: [PATCH 4/6] review comments

---
 balancer/pickfirst/metrics_test.go            |  50 +++---
 .../testutils/stats/test_metrics_recorder.go  | 167 +++++++++++++-----
 2 files changed, 144 insertions(+), 73 deletions(-)

diff --git a/balancer/pickfirst/metrics_test.go b/balancer/pickfirst/metrics_test.go
index d1f942d3dd98..3bf3a4e90b58 100644
--- a/balancer/pickfirst/metrics_test.go
+++ b/balancer/pickfirst/metrics_test.go
@@ -41,6 +41,7 @@ import (
 	"google.golang.org/grpc/resolver"
 	"google.golang.org/grpc/resolver/manual"
 	"google.golang.org/grpc/serviceconfig"
+	estats "google.golang.org/grpc/experimental/stats"
 	"google.golang.org/grpc/stats/opentelemetry"
 
 	"go.opentelemetry.io/otel/attribute"
@@ -107,32 +108,23 @@ func (s) TestPickFirstMetrics(t *testing.T) {
 		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 0)
 	}
 
+	// Checking for subchannel metrics as well
+	waitForInt64Count(ctx, t, tmr, "grpc.subchannel.connection_attempts_succeeded", 1)
 	if got, _ := tmr.Metric("grpc.subchannel.connection_attempts_failed"); got != 0 {
 		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.subchannel.connection_attempts_failed", got, 0)
 	}
 	if got, _ := tmr.Metric("grpc.subchannel.disconnections"); got != 0 {
 		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.subchannel.disconnections", got, 0)
 	}
-
-	for _, metric := range []struct {
-		name string
-		want float64
-	}{
-		{"grpc.subchannel.connection_attempts_succeeded", 1},
-		{"grpc.subchannel.open_connections", 1},
-	} {
-		awaitMetric(ctx, t, tmr, metric.name, metric.want)
-	}
+	waitForInt64UpDownCount(ctx, t, tmr, "grpc.subchannel.open_connections", 1)
 
 	ss.Stop()
 	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
 	if got, _ := tmr.Metric("grpc.lb.pick_first.disconnections"); got != 1 {
 		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.lb.pick_first.disconnections", got, 1)
 	}
-	if got, _ := tmr.Metric("grpc.subchannel.disconnections"); got != 1 {
-		t.Errorf("Unexpected data for metric %v, got: %v, want: %v", "grpc.subchannel.disconnections", got, 1)
-	}
-	awaitMetric(ctx, t, tmr, "grpc.subchannel.open_connections", -1)
+	waitForInt64Count(ctx, t, tmr, "grpc.subchannel.disconnections", 1)
+	waitForInt64UpDownCount(ctx, t, tmr, "grpc.subchannel.open_connections", -1)
 }
 
 // TestPickFirstMetricsFailure tests the connection attempts failed metric. It
@@ -298,6 +290,7 @@ func metricsDataFromReader(ctx context.Context, reader *metric.ManualReader) map
 	return gotMetrics
 }
 
+<<<<<<< HEAD
 // TestDisconnectLabel tests the disconnect label metric plumbing.
 // Separately, e2e tests are more exhaustive and check for all disconnect reasons.
 func (s) TestDisconnectLabel(t *testing.T) {
@@ -455,17 +448,24 @@ func runDisconnectLabelTest(t *testing.T, wantLabel string, triggerFunc func(*st
 	t.Fatalf("Error waiting for metrics grpc.subchannel.disconnections: %v", ctx.Err())
 }
 
-func awaitMetric(ctx context.Context, t *testing.T, tmr *stats.TestMetricsRecorder, name string, want float64) {
+func waitForInt64Count(ctx context.Context, t *testing.T, tmr *stats.TestMetricsRecorder, name string, want int64) {
 	t.Helper()
-	for {
-		got, _ := tmr.Metric(name)
-		if got == want {
-			return
-		}
-		select {
-		case <-ctx.Done():
-			t.Fatalf("Timeout waiting for expected data for metric %v, got: %v, want: %v", name, got, want)
-		case <-time.After(10 * time.Millisecond):
-		}
+	err := tmr.WaitForInt64Count(ctx, stats.MetricsData{
+		Handle:  &estats.MetricDescriptor{Name: name},
+		IntIncr: want,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func waitForInt64UpDownCount(ctx context.Context, t *testing.T, tmr *stats.TestMetricsRecorder, name string, want int64) {
+	t.Helper()
+	err := tmr.WaitForInt64UpDownCount(ctx, stats.MetricsData{
+		Handle:  &estats.MetricDescriptor{Name: name},
+		IntIncr: want,
+	})
+	if err != nil {
+		t.Fatal(err)
 	}
 }
diff --git a/internal/testutils/stats/test_metrics_recorder.go b/internal/testutils/stats/test_metrics_recorder.go
index 1c9149539f8b..8766ade05670 100644
--- a/internal/testutils/stats/test_metrics_recorder.go
+++ b/internal/testutils/stats/test_metrics_recorder.go
@@ -46,7 +46,7 @@ type TestMetricsRecorder struct {
 	// mu protects data.
 	mu sync.Mutex
 	// data is the most recent update for each metric name.
-	data map[string]float64
+	data map[string]MetricsData
 }
 
 // NewTestMetricsRecorder returns a new TestMetricsRecorder.
@@ -59,7 +59,7 @@ func NewTestMetricsRecorder() *TestMetricsRecorder {
 		intGaugeCh:       testutils.NewChannelWithSize(10),
 		intUpDownCountCh: testutils.NewChannelWithSize(10),
 
-		data: make(map[string]float64),
+		data: make(map[string]MetricsData),
 	}
 }
 
@@ -69,14 +69,24 @@ func (r *TestMetricsRecorder) Metric(name string) (float64, bool) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	data, ok := r.data[name]
-	return data, ok
+	if !ok {
+		return 0, false
+	}
+	switch data.Handle.Type {
+	case estats.MetricTypeIntCount, estats.MetricTypeIntGauge, estats.MetricTypeIntUpDownCount, estats.MetricTypeIntAsyncGauge:
+		return float64(data.IntIncr), true
+	case estats.MetricTypeFloatCount:
+		return data.FloatIncr, true
+	default:
+		return 0, false
+	}
 }
 
 // ClearMetrics clears the metrics data store of the test metrics recorder.
 func (r *TestMetricsRecorder) ClearMetrics() {
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data = make(map[string]float64)
+	r.data = make(map[string]MetricsData)
 }
 
 // MetricsData represents data associated with a metric.
@@ -92,17 +102,75 @@ type MetricsData struct {
 	LabelVals []string
 }
 
+func (r *TestMetricsRecorder) waitForMetric(ctx context.Context, ch *testutils.Channel, name string) (MetricsData, error) {
+	r.mu.Lock()
+	if data, ok := r.data[name]; ok {
+		r.mu.Unlock()
+		return data, nil
+	}
+	r.mu.Unlock()
+
+	for {
+		got, err := ch.Receive(ctx)
+		if err != nil {
+			return MetricsData{}, err
+		}
+		md := got.(MetricsData)
+
+		r.mu.Lock()
+		r.data[md.Handle.Name] = md
+		r.mu.Unlock()
+
+		if md.Handle.Name == name {
+			return md, nil
+		}
+	}
+}
+
 // WaitForInt64Count waits for an int64 count metric to be recorded and verifies
 // that the recorded metrics data matches the expected metricsDataWant. Returns
 // an error if failed to wait or received wrong data.
 func (r *TestMetricsRecorder) WaitForInt64Count(ctx context.Context, metricsDataWant MetricsData) error {
-	got, err := r.intCountCh.Receive(ctx)
+	got, err := r.waitForMetric(ctx, r.intCountCh, metricsDataWant.Handle.Name)
 	if err != nil {
 		return fmt.Errorf("timeout waiting for int64Count")
 	}
-	metricsDataGot := got.(MetricsData)
-	if diff := cmp.Diff(metricsDataGot, metricsDataWant); diff != "" {
-		return fmt.Errorf("int64count metricsData received unexpected value (-got, +want): %v", diff)
+	if got.IntIncr != metricsDataWant.IntIncr {
+		return fmt.Errorf("int64count metricsData received unexpected value: got %v, want %v", got.IntIncr, metricsDataWant.IntIncr)
+	}
+	if metricsDataWant.LabelKeys != nil {
+		if diff := cmp.Diff(got.LabelKeys, metricsDataWant.LabelKeys); diff != "" {
+			return fmt.Errorf("int64count metricsData received unexpected label keys: %v", diff)
+		}
+	}
+	if metricsDataWant.LabelVals != nil {
+		if diff := cmp.Diff(got.LabelVals, metricsDataWant.LabelVals); diff != "" {
+			return fmt.Errorf("int64count metricsData received unexpected label values: %v", diff)
+		}
+	}
+	return nil
+}
+
+// WaitForInt64UpDownCount waits for an int64 up-down count metric to be recorded and verifies
+// that the recorded metrics data matches the expected metricsDataWant. Returns
+// an error if failed to wait or received wrong data.
+func (r *TestMetricsRecorder) WaitForInt64UpDownCount(ctx context.Context, metricsDataWant MetricsData) error {
+	got, err := r.waitForMetric(ctx, r.intUpDownCountCh, metricsDataWant.Handle.Name)
+	if err != nil {
+		return fmt.Errorf("timeout waiting for int64UpDownCount")
+	}
+	if got.IntIncr != metricsDataWant.IntIncr {
+		return fmt.Errorf("int64UpDownCount metricsData received unexpected value: got %v, want %v", got.IntIncr, metricsDataWant.IntIncr)
+	}
+	if metricsDataWant.LabelKeys != nil {
+		if diff := cmp.Diff(got.LabelKeys, metricsDataWant.LabelKeys); diff != "" {
+			return fmt.Errorf("int64UpDownCount metricsData received unexpected label keys: %v", diff)
+		}
+	}
+	if metricsDataWant.LabelVals != nil {
+		if diff := cmp.Diff(got.LabelVals, metricsDataWant.LabelVals); diff != "" {
+			return fmt.Errorf("int64UpDownCount metricsData received unexpected label values: %v", diff)
+		}
 	}
 	return nil
 }
@@ -125,45 +193,46 @@ func (r *TestMetricsRecorder) WaitForInt64CountIncr(ctx context.Context, incrWan
 // RecordInt64Count sends the metrics data to the intCountCh channel and updates
 // the internal data map with the recorded value.
 func (r *TestMetricsRecorder) RecordInt64Count(handle *estats.Int64CountHandle, incr int64, labels ...string) {
-	r.intCountCh.ReceiveOrFail()
-	r.intCountCh.Send(MetricsData{
+	md := MetricsData{
 		Handle:    handle.Descriptor(),
 		IntIncr:   incr,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
-	})
+	}
+	r.intCountCh.ReceiveOrFail()
+	r.intCountCh.Send(md)
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = float64(incr)
+	r.data[handle.Name] = md
 }
 
 // RecordInt64UpDownCount sends the metrics data to the intUpDownCountCh channel and updates
 // the internal data map with the recorded value.
 func (r *TestMetricsRecorder) RecordInt64UpDownCount(handle *estats.Int64UpDownCountHandle, incr int64, labels ...string) {
-	r.intUpDownCountCh.ReceiveOrFail()
-	r.intUpDownCountCh.Send(MetricsData{
+	md := MetricsData{
 		Handle:    handle.Descriptor(),
 		IntIncr:   incr,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
-	})
+	}
+	r.intUpDownCountCh.ReceiveOrFail()
+	r.intUpDownCountCh.Send(md)
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = float64(incr)
+	r.data[handle.Name] = md
 }
 
 // WaitForFloat64Count waits for a float count metric to be recorded and
 // verifies that the recorded metrics data matches the expected metricsDataWant.
 // Returns an error if failed to wait or received wrong data.
 func (r *TestMetricsRecorder) WaitForFloat64Count(ctx context.Context, metricsDataWant MetricsData) error {
-	got, err := r.floatCountCh.Receive(ctx)
+	got, err := r.waitForMetric(ctx, r.floatCountCh, metricsDataWant.Handle.Name)
 	if err != nil {
 		return fmt.Errorf("timeout waiting for float64Count")
 	}
-	metricsDataGot := got.(MetricsData)
-	if diff := cmp.Diff(metricsDataGot, metricsDataWant); diff != "" {
+	if diff := cmp.Diff(got, metricsDataWant); diff != "" {
 		return fmt.Errorf("float64count metricsData received unexpected value (-got, +want): %v", diff)
 	}
 	return nil
@@ -172,29 +241,29 @@ func (r *TestMetricsRecorder) WaitForFloat64Count(ctx context.Context, metricsDa
 // RecordFloat64Count sends the metrics data to the floatCountCh channel and
 // updates the internal data map with the recorded value.
 func (r *TestMetricsRecorder) RecordFloat64Count(handle *estats.Float64CountHandle, incr float64, labels ...string) {
-	r.floatCountCh.ReceiveOrFail()
-	r.floatCountCh.Send(MetricsData{
+	md := MetricsData{
 		Handle:    handle.Descriptor(),
 		FloatIncr: incr,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
-	})
+	}
+	r.floatCountCh.ReceiveOrFail()
+	r.floatCountCh.Send(md)
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = incr
+	r.data[handle.Name] = md
 }
 
 // WaitForInt64Histo waits for an int histo metric to be recorded and verifies
 // that the recorded metrics data matches the expected metricsDataWant. Returns
 // an error if failed to wait or received wrong data.
 func (r *TestMetricsRecorder) WaitForInt64Histo(ctx context.Context, metricsDataWant MetricsData) error {
-	got, err := r.intHistoCh.Receive(ctx)
+	got, err := r.waitForMetric(ctx, r.intHistoCh, metricsDataWant.Handle.Name)
 	if err != nil {
 		return fmt.Errorf("timeout waiting for int64Histo")
 	}
-	metricsDataGot := got.(MetricsData)
-	if diff := cmp.Diff(metricsDataGot, metricsDataWant); diff != "" {
+	if diff := cmp.Diff(got, metricsDataWant); diff != "" {
 		return fmt.Errorf("int64Histo metricsData received unexpected value (-got, +want): %v", diff)
 	}
 	return nil
@@ -203,29 +272,29 @@ func (r *TestMetricsRecorder) WaitForInt64Histo(ctx context.Context, metricsData
 // RecordInt64Histo sends the metrics data to the intHistoCh channel and updates
 // the internal data map with the recorded value.
 func (r *TestMetricsRecorder) RecordInt64Histo(handle *estats.Int64HistoHandle, incr int64, labels ...string) {
-	r.intHistoCh.ReceiveOrFail()
-	r.intHistoCh.Send(MetricsData{
+	md := MetricsData{
 		Handle:    handle.Descriptor(),
 		IntIncr:   incr,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
-	})
+	}
+	r.intHistoCh.ReceiveOrFail()
+	r.intHistoCh.Send(md)
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = float64(incr)
+	r.data[handle.Name] = md
 }
 
 // WaitForFloat64Histo waits for a float histo metric to be recorded and
 // verifies that the recorded metrics data matches the expected metricsDataWant.
 // Returns an error if failed to wait or received wrong data.
 func (r *TestMetricsRecorder) WaitForFloat64Histo(ctx context.Context, metricsDataWant MetricsData) error {
-	got, err := r.floatHistoCh.Receive(ctx)
+	got, err := r.waitForMetric(ctx, r.floatHistoCh, metricsDataWant.Handle.Name)
 	if err != nil {
 		return fmt.Errorf("timeout waiting for float64Histo")
 	}
-	metricsDataGot := got.(MetricsData)
-	if diff := cmp.Diff(metricsDataGot, metricsDataWant); diff != "" {
+	if diff := cmp.Diff(got, metricsDataWant); diff != "" {
 		return fmt.Errorf("float64Histo metricsData received unexpected value (-got, +want): %v", diff)
 	}
 	return nil
@@ -234,28 +303,28 @@ func (r *TestMetricsRecorder) WaitForFloat64Histo(ctx context.Context, metricsDa
 // RecordFloat64Histo sends the metrics data to the floatHistoCh channel and
 // updates the internal data map with the recorded value.
 func (r *TestMetricsRecorder) RecordFloat64Histo(handle *estats.Float64HistoHandle, incr float64, labels ...string) {
-	r.floatHistoCh.ReceiveOrFail()
-	r.floatHistoCh.Send(MetricsData{
+	md := MetricsData{
 		Handle:    handle.Descriptor(),
 		FloatIncr: incr,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
-	})
+	}
+	r.floatHistoCh.ReceiveOrFail()
+	r.floatHistoCh.Send(md)
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = incr
+	r.data[handle.Name] = md
 }
 
 // WaitForInt64Gauge waits for a int gauge metric to be recorded and verifies
 // that the recorded metrics data matches the expected metricsDataWant.
 func (r *TestMetricsRecorder) WaitForInt64Gauge(ctx context.Context, metricsDataWant MetricsData) error {
-	got, err := r.intGaugeCh.Receive(ctx)
+	got, err := r.waitForMetric(ctx, r.intGaugeCh, metricsDataWant.Handle.Name)
 	if err != nil {
 		return fmt.Errorf("timeout waiting for int64Gauge")
 	}
-	metricsDataGot := got.(MetricsData)
-	if diff := cmp.Diff(metricsDataGot, metricsDataWant); diff != "" {
+	if diff := cmp.Diff(got, metricsDataWant); diff != "" {
 		return fmt.Errorf("int64Gauge metricsData received unexpected value (-got, +want): %v", diff)
 	}
 	return nil
@@ -264,17 +333,18 @@ func (r *TestMetricsRecorder) WaitForInt64Gauge(ctx context.Context, metricsData
 // RecordInt64Gauge sends the metrics data to the intGaugeCh channel and updates
 // the internal data map with the recorded value.
 func (r *TestMetricsRecorder) RecordInt64Gauge(handle *estats.Int64GaugeHandle, incr int64, labels ...string) {
-	r.intGaugeCh.ReceiveOrFail()
-	r.intGaugeCh.Send(MetricsData{
+	md := MetricsData{
 		Handle:    handle.Descriptor(),
 		IntIncr:   incr,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
-	})
+	}
+	r.intGaugeCh.ReceiveOrFail()
+	r.intGaugeCh.Send(md)
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = float64(incr)
+	r.data[handle.Name] = md
 }
 
 // To implement a estats.AsyncMetricsRecorder, which allows it to be used in async metrics:
@@ -282,17 +352,18 @@ func (r *TestMetricsRecorder) RecordInt64Gauge(handle *estats.Int64GaugeHandle,
 // RecordInt64AsyncGauge sends the metrics data to the intGaugeCh channel and updates
 // the internal data map with the recorded value.
 func (r *TestMetricsRecorder) RecordInt64AsyncGauge(handle *estats.Int64AsyncGaugeHandle, incr int64, labels ...string) {
-	r.intGaugeCh.ReceiveOrFail()
-	r.intGaugeCh.Send(MetricsData{
+	md := MetricsData{
 		Handle:    handle.Descriptor(),
 		IntIncr:   incr,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
-	})
+	}
+	r.intGaugeCh.ReceiveOrFail()
+	r.intGaugeCh.Send(md)
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = float64(incr)
+	r.data[handle.Name] = md
 }
 
 // To implement a stats.Handler, which allows it to be set as a dial option:

From 7d50135608f58baa0b322aacd4512a0d7a47470f Mon Sep 17 00:00:00 2001
From: Madhav Bissa <mbissa@gmail.com>
Date: Wed, 10 Jun 2026 09:46:13 +0000
Subject: [PATCH 5/6] testutils/stats: store metric slice and support sequence
 verification in TestMetricsRecorder

---
 balancer/pickfirst/metrics_test.go            |   3 +-
 .../testutils/stats/test_metrics_recorder.go  | 144 +++++++++---------
 2 files changed, 70 insertions(+), 77 deletions(-)

diff --git a/balancer/pickfirst/metrics_test.go b/balancer/pickfirst/metrics_test.go
index 3bf3a4e90b58..8a0d3f9b25b0 100644
--- a/balancer/pickfirst/metrics_test.go
+++ b/balancer/pickfirst/metrics_test.go
@@ -32,6 +32,7 @@ import (
 	"google.golang.org/grpc/balancer/pickfirst"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/credentials/insecure"
+	estats "google.golang.org/grpc/experimental/stats"
 	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/internal/stubserver"
 	"google.golang.org/grpc/internal/testutils"
@@ -41,7 +42,6 @@ import (
 	"google.golang.org/grpc/resolver"
 	"google.golang.org/grpc/resolver/manual"
 	"google.golang.org/grpc/serviceconfig"
-	estats "google.golang.org/grpc/experimental/stats"
 	"google.golang.org/grpc/stats/opentelemetry"
 
 	"go.opentelemetry.io/otel/attribute"
@@ -290,7 +290,6 @@ func metricsDataFromReader(ctx context.Context, reader *metric.ManualReader) map
 	return gotMetrics
 }
 
-<<<<<<< HEAD
 // TestDisconnectLabel tests the disconnect label metric plumbing.
 // Separately, e2e tests are more exhaustive and check for all disconnect reasons.
 func (s) TestDisconnectLabel(t *testing.T) {
diff --git a/internal/testutils/stats/test_metrics_recorder.go b/internal/testutils/stats/test_metrics_recorder.go
index 8766ade05670..ec69ab367320 100644
--- a/internal/testutils/stats/test_metrics_recorder.go
+++ b/internal/testutils/stats/test_metrics_recorder.go
@@ -43,10 +43,12 @@ type TestMetricsRecorder struct {
 	intGaugeCh       *testutils.Channel
 	intUpDownCountCh *testutils.Channel
 
-	// mu protects data.
+	// mu protects data and consumed.
 	mu sync.Mutex
-	// data is the most recent update for each metric name.
-	data map[string]MetricsData
+	// data contains all recorded updates for each metric name in order.
+	data map[string][]MetricsData
+	// consumed tracks the number of updates consumed per metric name.
+	consumed map[string]int
 }
 
 // NewTestMetricsRecorder returns a new TestMetricsRecorder.
@@ -59,7 +61,8 @@ func NewTestMetricsRecorder() *TestMetricsRecorder {
 		intGaugeCh:       testutils.NewChannelWithSize(10),
 		intUpDownCountCh: testutils.NewChannelWithSize(10),
 
-		data: make(map[string]MetricsData),
+		data:     make(map[string][]MetricsData),
+		consumed: make(map[string]int),
 	}
 }
 
@@ -68,10 +71,11 @@ func NewTestMetricsRecorder() *TestMetricsRecorder {
 func (r *TestMetricsRecorder) Metric(name string) (float64, bool) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	data, ok := r.data[name]
-	if !ok {
+	slice, ok := r.data[name]
+	if !ok || len(slice) == 0 {
 		return 0, false
 	}
+	data := slice[len(slice)-1]
 	switch data.Handle.Type {
 	case estats.MetricTypeIntCount, estats.MetricTypeIntGauge, estats.MetricTypeIntUpDownCount, estats.MetricTypeIntAsyncGauge:
 		return float64(data.IntIncr), true
@@ -86,7 +90,8 @@ func (r *TestMetricsRecorder) Metric(name string) (float64, bool) {
 func (r *TestMetricsRecorder) ClearMetrics() {
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data = make(map[string]MetricsData)
+	r.data = make(map[string][]MetricsData)
+	r.consumed = make(map[string]int)
 }
 
 // MetricsData represents data associated with a metric.
@@ -103,72 +108,64 @@ type MetricsData struct {
 }
 
 func (r *TestMetricsRecorder) waitForMetric(ctx context.Context, ch *testutils.Channel, name string) (MetricsData, error) {
-	r.mu.Lock()
-	if data, ok := r.data[name]; ok {
-		r.mu.Unlock()
-		return data, nil
-	}
-	r.mu.Unlock()
-
 	for {
-		got, err := ch.Receive(ctx)
-		if err != nil {
-			return MetricsData{}, err
-		}
-		md := got.(MetricsData)
-
 		r.mu.Lock()
-		r.data[md.Handle.Name] = md
+		if slice, ok := r.data[name]; ok && len(slice) > r.consumed[name] {
+			md := slice[r.consumed[name]]
+			r.consumed[name]++
+			r.mu.Unlock()
+			return md, nil
+		}
 		r.mu.Unlock()
 
-		if md.Handle.Name == name {
-			return md, nil
+		_, err := ch.Receive(ctx)
+		if err != nil {
+			return MetricsData{}, err
 		}
 	}
 }
 
-// WaitForInt64Count waits for an int64 count metric to be recorded and verifies
-// that the recorded metrics data matches the expected metricsDataWant. Returns
-// an error if failed to wait or received wrong data.
-func (r *TestMetricsRecorder) WaitForInt64Count(ctx context.Context, metricsDataWant MetricsData) error {
-	got, err := r.waitForMetric(ctx, r.intCountCh, metricsDataWant.Handle.Name)
+// WaitForInt64Count waits for an int64 count metric and verifies that it
+// matches want. Returns an error if failed to wait or received wrong data.
+func (r *TestMetricsRecorder) WaitForInt64Count(ctx context.Context, want MetricsData) error {
+	got, err := r.waitForMetric(ctx, r.intCountCh, want.Handle.Name)
 	if err != nil {
 		return fmt.Errorf("timeout waiting for int64Count")
 	}
-	if got.IntIncr != metricsDataWant.IntIncr {
-		return fmt.Errorf("int64count metricsData received unexpected value: got %v, want %v", got.IntIncr, metricsDataWant.IntIncr)
+	if got.IntIncr != want.IntIncr {
+		return fmt.Errorf("int64count metricsData received unexpected value: got %v, want %v", got.IntIncr, want.IntIncr)
 	}
-	if metricsDataWant.LabelKeys != nil {
-		if diff := cmp.Diff(got.LabelKeys, metricsDataWant.LabelKeys); diff != "" {
+	if want.LabelKeys != nil {
+		if diff := cmp.Diff(got.LabelKeys, want.LabelKeys); diff != "" {
 			return fmt.Errorf("int64count metricsData received unexpected label keys: %v", diff)
 		}
 	}
-	if metricsDataWant.LabelVals != nil {
-		if diff := cmp.Diff(got.LabelVals, metricsDataWant.LabelVals); diff != "" {
+	if want.LabelVals != nil {
+		if diff := cmp.Diff(got.LabelVals, want.LabelVals); diff != "" {
 			return fmt.Errorf("int64count metricsData received unexpected label values: %v", diff)
 		}
 	}
 	return nil
 }
 
-// WaitForInt64UpDownCount waits for an int64 up-down count metric to be recorded and verifies
-// that the recorded metrics data matches the expected metricsDataWant. Returns
-// an error if failed to wait or received wrong data.
-func (r *TestMetricsRecorder) WaitForInt64UpDownCount(ctx context.Context, metricsDataWant MetricsData) error {
-	got, err := r.waitForMetric(ctx, r.intUpDownCountCh, metricsDataWant.Handle.Name)
+// WaitForInt64UpDownCount waits for an int64 up-down count metric and
+// verifies that it matches want. Returns an error if failed to wait or
+// received wrong data.
+func (r *TestMetricsRecorder) WaitForInt64UpDownCount(ctx context.Context, want MetricsData) error {
+	got, err := r.waitForMetric(ctx, r.intUpDownCountCh, want.Handle.Name)
 	if err != nil {
 		return fmt.Errorf("timeout waiting for int64UpDownCount")
 	}
-	if got.IntIncr != metricsDataWant.IntIncr {
-		return fmt.Errorf("int64UpDownCount metricsData received unexpected value: got %v, want %v", got.IntIncr, metricsDataWant.IntIncr)
+	if got.IntIncr != want.IntIncr {
+		return fmt.Errorf("int64UpDownCount metricsData received unexpected value: got %v, want %v", got.IntIncr, want.IntIncr)
 	}
-	if metricsDataWant.LabelKeys != nil {
-		if diff := cmp.Diff(got.LabelKeys, metricsDataWant.LabelKeys); diff != "" {
+	if want.LabelKeys != nil {
+		if diff := cmp.Diff(got.LabelKeys, want.LabelKeys); diff != "" {
 			return fmt.Errorf("int64UpDownCount metricsData received unexpected label keys: %v", diff)
 		}
 	}
-	if metricsDataWant.LabelVals != nil {
-		if diff := cmp.Diff(got.LabelVals, metricsDataWant.LabelVals); diff != "" {
+	if want.LabelVals != nil {
+		if diff := cmp.Diff(got.LabelVals, want.LabelVals); diff != "" {
 			return fmt.Errorf("int64UpDownCount metricsData received unexpected label values: %v", diff)
 		}
 	}
@@ -204,7 +201,7 @@ func (r *TestMetricsRecorder) RecordInt64Count(handle *estats.Int64CountHandle,
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = md
+	r.data[handle.Name] = append(r.data[handle.Name], md)
 }
 
 // RecordInt64UpDownCount sends the metrics data to the intUpDownCountCh channel and updates
@@ -221,18 +218,17 @@ func (r *TestMetricsRecorder) RecordInt64UpDownCount(handle *estats.Int64UpDownC
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = md
+	r.data[handle.Name] = append(r.data[handle.Name], md)
 }
 
-// WaitForFloat64Count waits for a float count metric to be recorded and
-// verifies that the recorded metrics data matches the expected metricsDataWant.
-// Returns an error if failed to wait or received wrong data.
-func (r *TestMetricsRecorder) WaitForFloat64Count(ctx context.Context, metricsDataWant MetricsData) error {
-	got, err := r.waitForMetric(ctx, r.floatCountCh, metricsDataWant.Handle.Name)
+// WaitForFloat64Count waits for a float count metric and verifies that it
+// matches want. Returns an error if failed to wait or received wrong data.
+func (r *TestMetricsRecorder) WaitForFloat64Count(ctx context.Context, want MetricsData) error {
+	got, err := r.waitForMetric(ctx, r.floatCountCh, want.Handle.Name)
 	if err != nil {
 		return fmt.Errorf("timeout waiting for float64Count")
 	}
-	if diff := cmp.Diff(got, metricsDataWant); diff != "" {
+	if diff := cmp.Diff(got, want); diff != "" {
 		return fmt.Errorf("float64count metricsData received unexpected value (-got, +want): %v", diff)
 	}
 	return nil
@@ -252,18 +248,17 @@ func (r *TestMetricsRecorder) RecordFloat64Count(handle *estats.Float64CountHand
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = md
+	r.data[handle.Name] = append(r.data[handle.Name], md)
 }
 
-// WaitForInt64Histo waits for an int histo metric to be recorded and verifies
-// that the recorded metrics data matches the expected metricsDataWant. Returns
-// an error if failed to wait or received wrong data.
-func (r *TestMetricsRecorder) WaitForInt64Histo(ctx context.Context, metricsDataWant MetricsData) error {
-	got, err := r.waitForMetric(ctx, r.intHistoCh, metricsDataWant.Handle.Name)
+// WaitForInt64Histo waits for an int histo metric and verifies that it matches
+// want. Returns an error if failed to wait or received wrong data.
+func (r *TestMetricsRecorder) WaitForInt64Histo(ctx context.Context, want MetricsData) error {
+	got, err := r.waitForMetric(ctx, r.intHistoCh, want.Handle.Name)
 	if err != nil {
 		return fmt.Errorf("timeout waiting for int64Histo")
 	}
-	if diff := cmp.Diff(got, metricsDataWant); diff != "" {
+	if diff := cmp.Diff(got, want); diff != "" {
 		return fmt.Errorf("int64Histo metricsData received unexpected value (-got, +want): %v", diff)
 	}
 	return nil
@@ -283,18 +278,17 @@ func (r *TestMetricsRecorder) RecordInt64Histo(handle *estats.Int64HistoHandle,
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = md
+	r.data[handle.Name] = append(r.data[handle.Name], md)
 }
 
-// WaitForFloat64Histo waits for a float histo metric to be recorded and
-// verifies that the recorded metrics data matches the expected metricsDataWant.
-// Returns an error if failed to wait or received wrong data.
-func (r *TestMetricsRecorder) WaitForFloat64Histo(ctx context.Context, metricsDataWant MetricsData) error {
-	got, err := r.waitForMetric(ctx, r.floatHistoCh, metricsDataWant.Handle.Name)
+// WaitForFloat64Histo waits for a float histo metric and verifies that it
+// matches want. Returns an error if failed to wait or received wrong data.
+func (r *TestMetricsRecorder) WaitForFloat64Histo(ctx context.Context, want MetricsData) error {
+	got, err := r.waitForMetric(ctx, r.floatHistoCh, want.Handle.Name)
 	if err != nil {
 		return fmt.Errorf("timeout waiting for float64Histo")
 	}
-	if diff := cmp.Diff(got, metricsDataWant); diff != "" {
+	if diff := cmp.Diff(got, want); diff != "" {
 		return fmt.Errorf("float64Histo metricsData received unexpected value (-got, +want): %v", diff)
 	}
 	return nil
@@ -314,17 +308,17 @@ func (r *TestMetricsRecorder) RecordFloat64Histo(handle *estats.Float64HistoHand
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = md
+	r.data[handle.Name] = append(r.data[handle.Name], md)
 }
 
-// WaitForInt64Gauge waits for a int gauge metric to be recorded and verifies
-// that the recorded metrics data matches the expected metricsDataWant.
-func (r *TestMetricsRecorder) WaitForInt64Gauge(ctx context.Context, metricsDataWant MetricsData) error {
-	got, err := r.waitForMetric(ctx, r.intGaugeCh, metricsDataWant.Handle.Name)
+// WaitForInt64Gauge waits for an int gauge metric and verifies that it matches
+// want.
+func (r *TestMetricsRecorder) WaitForInt64Gauge(ctx context.Context, want MetricsData) error {
+	got, err := r.waitForMetric(ctx, r.intGaugeCh, want.Handle.Name)
 	if err != nil {
 		return fmt.Errorf("timeout waiting for int64Gauge")
 	}
-	if diff := cmp.Diff(got, metricsDataWant); diff != "" {
+	if diff := cmp.Diff(got, want); diff != "" {
 		return fmt.Errorf("int64Gauge metricsData received unexpected value (-got, +want): %v", diff)
 	}
 	return nil
@@ -344,7 +338,7 @@ func (r *TestMetricsRecorder) RecordInt64Gauge(handle *estats.Int64GaugeHandle,
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = md
+	r.data[handle.Name] = append(r.data[handle.Name], md)
 }
 
 // To implement a estats.AsyncMetricsRecorder, which allows it to be used in async metrics:
@@ -363,7 +357,7 @@ func (r *TestMetricsRecorder) RecordInt64AsyncGauge(handle *estats.Int64AsyncGau
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.data[handle.Name] = md
+	r.data[handle.Name] = append(r.data[handle.Name], md)
 }
 
 // To implement a stats.Handler, which allows it to be set as a dial option:

From 2669e2e831f5ccecc2a4d510dfcde5c49dbcdd5f Mon Sep 17 00:00:00 2001
From: Madhav Bissa <mbissa@gmail.com>
Date: Mon, 15 Jun 2026 07:39:57 +0000
Subject: [PATCH 6/6] fixing race review comments

---
 .../testutils/stats/test_metrics_recorder.go  | 49 +++++++++++--------
 1 file changed, 28 insertions(+), 21 deletions(-)

diff --git a/internal/testutils/stats/test_metrics_recorder.go b/internal/testutils/stats/test_metrics_recorder.go
index ec69ab367320..205d398ae1a1 100644
--- a/internal/testutils/stats/test_metrics_recorder.go
+++ b/internal/testutils/stats/test_metrics_recorder.go
@@ -196,12 +196,13 @@ func (r *TestMetricsRecorder) RecordInt64Count(handle *estats.Int64CountHandle,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
 	}
-	r.intCountCh.ReceiveOrFail()
-	r.intCountCh.Send(md)
 
 	r.mu.Lock()
-	defer r.mu.Unlock()
 	r.data[handle.Name] = append(r.data[handle.Name], md)
+	r.mu.Unlock()
+
+	r.intCountCh.ReceiveOrFail()
+	r.intCountCh.Send(md)
 }
 
 // RecordInt64UpDownCount sends the metrics data to the intUpDownCountCh channel and updates
@@ -213,12 +214,13 @@ func (r *TestMetricsRecorder) RecordInt64UpDownCount(handle *estats.Int64UpDownC
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
 	}
-	r.intUpDownCountCh.ReceiveOrFail()
-	r.intUpDownCountCh.Send(md)
 
 	r.mu.Lock()
-	defer r.mu.Unlock()
 	r.data[handle.Name] = append(r.data[handle.Name], md)
+	r.mu.Unlock()
+
+	r.intUpDownCountCh.ReceiveOrFail()
+	r.intUpDownCountCh.Send(md)
 }
 
 // WaitForFloat64Count waits for a float count metric and verifies that it
@@ -243,12 +245,13 @@ func (r *TestMetricsRecorder) RecordFloat64Count(handle *estats.Float64CountHand
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
 	}
-	r.floatCountCh.ReceiveOrFail()
-	r.floatCountCh.Send(md)
 
 	r.mu.Lock()
-	defer r.mu.Unlock()
 	r.data[handle.Name] = append(r.data[handle.Name], md)
+	r.mu.Unlock()
+
+	r.floatCountCh.ReceiveOrFail()
+	r.floatCountCh.Send(md)
 }
 
 // WaitForInt64Histo waits for an int histo metric and verifies that it matches
@@ -273,12 +276,13 @@ func (r *TestMetricsRecorder) RecordInt64Histo(handle *estats.Int64HistoHandle,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
 	}
-	r.intHistoCh.ReceiveOrFail()
-	r.intHistoCh.Send(md)
 
 	r.mu.Lock()
-	defer r.mu.Unlock()
 	r.data[handle.Name] = append(r.data[handle.Name], md)
+	r.mu.Unlock()
+
+	r.intHistoCh.ReceiveOrFail()
+	r.intHistoCh.Send(md)
 }
 
 // WaitForFloat64Histo waits for a float histo metric and verifies that it
@@ -303,12 +307,13 @@ func (r *TestMetricsRecorder) RecordFloat64Histo(handle *estats.Float64HistoHand
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
 	}
-	r.floatHistoCh.ReceiveOrFail()
-	r.floatHistoCh.Send(md)
 
 	r.mu.Lock()
-	defer r.mu.Unlock()
 	r.data[handle.Name] = append(r.data[handle.Name], md)
+	r.mu.Unlock()
+
+	r.floatHistoCh.ReceiveOrFail()
+	r.floatHistoCh.Send(md)
 }
 
 // WaitForInt64Gauge waits for an int gauge metric and verifies that it matches
@@ -333,12 +338,13 @@ func (r *TestMetricsRecorder) RecordInt64Gauge(handle *estats.Int64GaugeHandle,
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
 	}
-	r.intGaugeCh.ReceiveOrFail()
-	r.intGaugeCh.Send(md)
 
 	r.mu.Lock()
-	defer r.mu.Unlock()
 	r.data[handle.Name] = append(r.data[handle.Name], md)
+	r.mu.Unlock()
+
+	r.intGaugeCh.ReceiveOrFail()
+	r.intGaugeCh.Send(md)
 }
 
 // To implement a estats.AsyncMetricsRecorder, which allows it to be used in async metrics:
@@ -352,12 +358,13 @@ func (r *TestMetricsRecorder) RecordInt64AsyncGauge(handle *estats.Int64AsyncGau
 		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
 		LabelVals: labels,
 	}
-	r.intGaugeCh.ReceiveOrFail()
-	r.intGaugeCh.Send(md)
 
 	r.mu.Lock()
-	defer r.mu.Unlock()
 	r.data[handle.Name] = append(r.data[handle.Name], md)
+	r.mu.Unlock()
+
+	r.intGaugeCh.ReceiveOrFail()
+	r.intGaugeCh.Send(md)
 }
 
 // To implement a stats.Handler, which allows it to be set as a dial option: