Cherry pick #7571 and #7579 to v1.66.x release branch

experimental/stats/metricregistry.go

@@ -20,7 +20,6 @@ package stats
 
 import (
 	"maps"
-	"testing"
 
 	"google.golang.org/grpc/grpclog"
 	"google.golang.org/grpc/internal"
@@ -250,9 +249,9 @@ func RegisterInt64Gauge(descriptor MetricDescriptor) *Int64GaugeHandle {
 }
 
 // snapshotMetricsRegistryForTesting snapshots the global data of the metrics
-// registry. Registers a cleanup function on the provided testing.T that sets
-// the metrics registry to its original state. Only called in testing functions.
-func snapshotMetricsRegistryForTesting(t *testing.T) {
+// registry. Returns a cleanup function that sets the metrics registry to its
+// original state.
+func snapshotMetricsRegistryForTesting() func() {
 	oldDefaultMetrics := DefaultMetrics
 	oldRegisteredMetrics := registeredMetrics
 	oldMetricsRegistry := metricsRegistry
@@ -262,9 +261,9 @@ func snapshotMetricsRegistryForTesting(t *testing.T) {
 	maps.Copy(registeredMetrics, registeredMetrics)
 	maps.Copy(metricsRegistry, metricsRegistry)
 
-	t.Cleanup(func() {
+	return func() {
 		DefaultMetrics = oldDefaultMetrics
 		registeredMetrics = oldRegisteredMetrics
 		metricsRegistry = oldMetricsRegistry
-	})
+	}
 }

experimental/stats/metricregistry_test.go

@@ -38,7 +38,9 @@ func Test(t *testing.T) {
 // TestPanic tests that registering two metrics with the same name across any
 // type of metric triggers a panic.
 func (s) TestPanic(t *testing.T) {
-	snapshotMetricsRegistryForTesting(t)
+	cleanup := snapshotMetricsRegistryForTesting()
+	defer cleanup()
+
 	want := "metric simple counter already registered"
 	defer func() {
 		if r := recover(); !strings.Contains(fmt.Sprint(r), want) {
@@ -64,7 +66,9 @@ func (s) TestPanic(t *testing.T) {
 // this tests the interactions between the metrics recorder and the metrics
 // registry.
 func (s) TestMetricRegistry(t *testing.T) {
-	snapshotMetricsRegistryForTesting(t)
+	cleanup := snapshotMetricsRegistryForTesting()
+	defer cleanup()
+
 	intCountHandle1 := RegisterInt64Count(MetricDescriptor{
 		Name:           "simple counter",
 		Description:    "sum of all emissions from tests",
@@ -141,7 +145,9 @@ func (s) TestMetricRegistry(t *testing.T) {
 // metric registry. A component (simulated by test) should be able to record on
 // the different registered int count metrics.
 func TestNumerousIntCounts(t *testing.T) {
-	snapshotMetricsRegistryForTesting(t)
+	cleanup := snapshotMetricsRegistryForTesting()
+	defer cleanup()
+
 	intCountHandle1 := RegisterInt64Count(MetricDescriptor{
 		Name:           "int counter",
 		Description:    "sum of all emissions from tests",

internal/internal.go

@@ -217,10 +217,9 @@ var (
 	SetConnectedAddress any // func(scs *SubConnState, addr resolver.Address)
 
 	// SnapshotMetricRegistryForTesting snapshots the global data of the metric
-	// registry. Registers a cleanup function on the provided testing.T that
-	// sets the metric registry to its original state. Only called in testing
-	// functions.
-	SnapshotMetricRegistryForTesting any // func(t *testing.T)
+	// registry. Returns a cleanup function that sets the metric registry to its
+	// original state. Only called in testing functions.
+	SnapshotMetricRegistryForTesting func() func()
 
 	// SetDefaultBufferPoolForTesting updates the default buffer pool, for
 	// testing purposes.

internal/stats/metrics_recorder_list_test.go

@@ -132,7 +132,9 @@ type recordingLoadBalancer struct {
 // test then asserts that the recorded metrics show up on both configured stats
 // handlers.
 func (s) TestMetricsRecorderList(t *testing.T) {
-	internal.SnapshotMetricRegistryForTesting.(func(t *testing.T))(t)
+	cleanup := internal.SnapshotMetricRegistryForTesting()
+	defer cleanup()
+
 	mr := manual.NewBuilderWithScheme("test-metrics-recorder-list")
 	defer mr.Close()
 
@@ -212,7 +214,9 @@ func (s) TestMetricsRecorderList(t *testing.T) {
 // TestMetricRecorderListPanic tests that the metrics recorder list panics if
 // received the wrong number of labels for a particular metric.
 func (s) TestMetricRecorderListPanic(t *testing.T) {
-	internal.SnapshotMetricRegistryForTesting.(func(t *testing.T))(t)
+	cleanup := internal.SnapshotMetricRegistryForTesting()
+	defer cleanup()
+
 	intCountHandle := estats.RegisterInt64Count(estats.MetricDescriptor{
 		Name:           "simple counter",
 		Description:    "sum of all emissions from tests",

mem/buffer_slice.go

@@ -92,9 +92,11 @@ func (s BufferSlice) Materialize() []byte {
 }
 
 // MaterializeToBuffer functions like Materialize except that it writes the data
-// to a single Buffer pulled from the given BufferPool. As a special case, if the
-// input BufferSlice only actually has one Buffer, this function has nothing to
-// do and simply returns said Buffer.
+// to a single Buffer pulled from the given BufferPool.
+//
+// As a special case, if the input BufferSlice only actually has one Buffer, this
+// function simply increases the refcount before returning said Buffer. Freeing this
+// buffer won't release it until the BufferSlice is itself released.
 func (s BufferSlice) MaterializeToBuffer(pool BufferPool) Buffer {
 	if len(s) == 1 {
 		s[0].Ref()

server.go

@@ -1359,6 +1359,7 @@ func (s *Server) processUnaryRPC(ctx context.Context, t transport.ServerTranspor
 		}
 		return err
 	}
+	defer d.Free()
 	if channelz.IsOn() {
 		t.IncrMsgRecv()
 	}

stats/opentelemetry/metricsregistry_test.go

@@ -61,7 +61,9 @@ func newServerStatsHandler(options MetricsOptions) metricsRecorderForTest {
 // the expected metrics emissions, which includes default metrics and optional
 // label assertions.
 func (s) TestMetricsRegistryMetrics(t *testing.T) {
-	internal.SnapshotMetricRegistryForTesting.(func(t *testing.T))(t)
+	cleanup := internal.SnapshotMetricRegistryForTesting()
+	defer cleanup()
+
 	intCountHandle1 := estats.RegisterInt64Count(estats.MetricDescriptor{
 		Name:           "int-counter-1",
 		Description:    "Sum of calls from test",