snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

internal/leakcheck/leakcheck.go

@@ -0,0 +1,489 @@
+/*
+ *
+ * Copyright 2017 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+// Package leakcheck contains functions to check leaked goroutines and buffers.
+//
+// Call the following at the beginning of test:
+//
+//	defer leakcheck.NewLeakChecker(t).Check()
+package leakcheck
+
+import (
+	"context"
+	"fmt"
+	"runtime"
+	"runtime/debug"
+	"slices"
+	"sort"
+	"strconv"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/mem"
+)
+
+// failTestsOnLeakedBuffers is a special flag that will cause tests to fail if
+// leaked buffers are detected, instead of simply logging them as an
+// informational failure. This can be enabled with the "checkbuffers" compile
+// flag, e.g.:
+//
+//	go test -tags=checkbuffers
+var failTestsOnLeakedBuffers = false
+
+func init() {
+	defaultPool := mem.DefaultBufferPool()
+	globalPool.Store(&defaultPool)
+	internal.SetDefaultBufferPool.(func(mem.BufferPool))(&globalPool)
+}
+
+var globalPool swappableBufferPool
+var globalTimerTracker *timerFactory
+
+type swappableBufferPool struct {
+	atomic.Pointer[mem.BufferPool]
+}
+
+func (b *swappableBufferPool) Get(length int) *[]byte {
+	return (*b.Load()).Get(length)
+}
+
+func (b *swappableBufferPool) Put(buf *[]byte) {
+	(*b.Load()).Put(buf)
+}
+
+// SetTrackingBufferPool replaces the default buffer pool in the mem package to
+// one that tracks where buffers are allocated. CheckTrackingBufferPool should
+// then be invoked at the end of the test to validate that all buffers pulled
+// from the pool were returned.
+func SetTrackingBufferPool(logger Logger) {
+	newPool := mem.BufferPool(&trackingBufferPool{
+		pool:             *globalPool.Load(),
+		logger:           logger,
+		allocatedBuffers: make(map[*[]byte][]uintptr),
+	})
+	globalPool.Store(&newPool)
+}
+
+// CheckTrackingBufferPool undoes the effects of SetTrackingBufferPool, and fails
+// unit tests if not all buffers were returned. It is invalid to invoke this
+// function without previously having invoked SetTrackingBufferPool.
+func CheckTrackingBufferPool() {
+	p := (*globalPool.Load()).(*trackingBufferPool)
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	globalPool.Store(&p.pool)
+
+	type uniqueTrace struct {
+		stack []uintptr
+		count int
+	}
+
+	var totalLeakedBuffers int
+	var uniqueTraces []uniqueTrace
+	for _, stack := range p.allocatedBuffers {
+		idx, ok := slices.BinarySearchFunc(uniqueTraces, stack, func(trace uniqueTrace, stack []uintptr) int {
+			return slices.Compare(trace.stack, stack)
+		})
+		if !ok {
+			uniqueTraces = slices.Insert(uniqueTraces, idx, uniqueTrace{stack: stack})
+		}
+		uniqueTraces[idx].count++
+		totalLeakedBuffers++
+	}
+
+	for _, ut := range uniqueTraces {
+		frames := runtime.CallersFrames(ut.stack)
+		var trace strings.Builder
+		for {
+			f, ok := frames.Next()
+			if !ok {
+				break
+			}
+			trace.WriteString(f.Function)
+			trace.WriteString("\n\t")
+			trace.WriteString(f.File)
+			trace.WriteString(":")
+			trace.WriteString(strconv.Itoa(f.Line))
+			trace.WriteString("\n")
+		}
+		format := "%d allocated buffers never freed:\n%s"
+		args := []any{ut.count, trace.String()}
+		if failTestsOnLeakedBuffers {
+			p.logger.Errorf(format, args...)
+		} else {
+			p.logger.Logf("WARNING "+format, args...)
+		}
+	}
+
+	if totalLeakedBuffers > 0 {
+		p.logger.Logf("%g%% of buffers never freed", float64(totalLeakedBuffers)/float64(p.bufferCount))
+	}
+}
+
+type trackingBufferPool struct {
+	pool   mem.BufferPool
+	logger Logger
+
+	lock             sync.Mutex
+	bufferCount      int
+	allocatedBuffers map[*[]byte][]uintptr
+}
+
+func (p *trackingBufferPool) Get(length int) *[]byte {
+	p.lock.Lock()
+	defer p.lock.Unlock()
+	p.bufferCount++
+	buf := p.pool.Get(length)
+	p.allocatedBuffers[buf] = currentStack(2)
+	return buf
+}
+
+func (p *trackingBufferPool) Put(buf *[]byte) {
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	if _, ok := p.allocatedBuffers[buf]; !ok {
+		p.logger.Errorf("Unknown buffer freed:\n%s", string(debug.Stack()))
+	} else {
+		delete(p.allocatedBuffers, buf)
+	}
+	p.pool.Put(buf)
+}
+
+var goroutinesToIgnore = []string{
+	"testing.Main(",
+	"testing.tRunner(",
+	"testing.(*M).",
+	"runtime.goexit",
+	"created by runtime.gc",
+	"created by runtime/trace.Start",
+	"interestingGoroutines",
+	"runtime.MHeap_Scavenger",
+	"signal.signal_recv",
+	"sigterm.handler",
+	"runtime_mcall",
+	"(*loggingT).flushDaemon",
+	"goroutine in C code",
+	// Ignore the http read/write goroutines. gce metadata.OnGCE() was leaking
+	// these, root cause unknown.
+	//
+	// https://github.com/grpc/grpc-go/issues/5171
+	// https://github.com/grpc/grpc-go/issues/5173
+	"created by net/http.(*Transport).dialConn",
+}
+
+// RegisterIgnoreGoroutine appends s into the ignore goroutine list. The
+// goroutines whose stack trace contains s will not be identified as leaked
+// goroutines. Not thread-safe, only call this function in init().
+func RegisterIgnoreGoroutine(s string) {
+	goroutinesToIgnore = append(goroutinesToIgnore, s)
+}
+
+func ignore(g string) bool {
+	sl := strings.SplitN(g, "\n", 2)
+	if len(sl) != 2 {
+		return true
+	}
+	stack := strings.TrimSpace(sl[1])
+	if strings.HasPrefix(stack, "testing.RunTests") {
+		return true
+	}
+
+	if stack == "" {
+		return true
+	}
+
+	for _, s := range goroutinesToIgnore {
+		if strings.Contains(stack, s) {
+			return true
+		}
+	}
+
+	return false
+}
+
+// interestingGoroutines returns all goroutines we care about for the purpose of
+// leak checking. It excludes testing or runtime ones.
+func interestingGoroutines() (gs []string) {
+	buf := make([]byte, 2<<20)
+	buf = buf[:runtime.Stack(buf, true)]
+	for _, g := range strings.Split(string(buf), "\n\n") {
+		if !ignore(g) {
+			gs = append(gs, g)
+		}
+	}
+	sort.Strings(gs)
+	return
+}
+
+// Logger is the interface that wraps the Logf and Errorf method. It's a subset
+// of testing.TB to make it easy to use this package.
+type Logger interface {
+	Logf(format string, args ...any)
+	Errorf(format string, args ...any)
+}
+
+// CheckGoroutines looks at the currently-running goroutines and checks if there
+// are any interesting (created by gRPC) goroutines leaked. It waits up to 10
+// seconds in the error cases.
+func CheckGoroutines(ctx context.Context, logger Logger) {
+	// Loop, waiting for goroutines to shut down.
+	// Wait up to timeout, but finish as quickly as possible.
+	var leaked []string
+	for ctx.Err() == nil {
+		if leaked = interestingGoroutines(); len(leaked) == 0 {
+			return
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+	for _, g := range leaked {
+		logger.Errorf("Leaked goroutine: %v", g)
+	}
+}
+
+// LeakChecker captures a Logger and is returned by NewLeakChecker as a
+// convenient method to set up leak check tests in a unit test.
+type LeakChecker struct {
+	logger Logger
+}
+
+// NewLeakChecker offers a convenient way to set up the leak checks for a
+// specific unit test. It can be used as follows, at the beginning of tests:
+//
+//	defer leakcheck.NewLeakChecker(t).Check()
+//
+// It initially invokes SetTrackingBufferPool to set up buffer tracking, then the
+// deferred LeakChecker.Check call will invoke CheckTrackingBufferPool and
+// CheckGoroutines with a default timeout of 10 seconds.
+func NewLeakChecker(logger Logger) *LeakChecker {
+	SetTrackingBufferPool(logger)
+	return &LeakChecker{logger: logger}
+}
+
+type timerFactory struct {
+	mu              sync.Mutex
+	allocatedTimers map[internal.Timer][]uintptr
+}
+
+func (tf *timerFactory) timeAfterFunc(d time.Duration, f func()) internal.Timer {
+	tf.mu.Lock()
+	defer tf.mu.Unlock()
+	ch := make(chan internal.Timer, 1)
+	timer := time.AfterFunc(d, func() {
+		f()
+		tf.remove(<-ch)
+	})
+	ch <- timer
+	tf.allocatedTimers[timer] = currentStack(2)
+	return &trackingTimer{
+		Timer:  timer,
+		parent: tf,
+	}
+}
+
+func (tf *timerFactory) remove(timer internal.Timer) {
+	tf.mu.Lock()
+	defer tf.mu.Unlock()
+	delete(tf.allocatedTimers, timer)
+}
+
+func (tf *timerFactory) pendingTimers() []string {
+	tf.mu.Lock()
+	defer tf.mu.Unlock()
+	leaked := []string{}
+	for _, stack := range tf.allocatedTimers {
+		leaked = append(leaked, fmt.Sprintf("Allocated timer never cancelled:\n%s", traceToString(stack)))
+	}
+	return leaked
+}
+
+type trackingTimer struct {
+	internal.Timer
+	parent *timerFactory
+}
+
+func (t *trackingTimer) Stop() bool {
+	t.parent.remove(t.Timer)
+	return t.Timer.Stop()
+}
+
+// TrackTimers replaces internal.TimerAfterFunc with one that tracks timer
+// creations, stoppages and expirations. CheckTimers should then be invoked at
+// the end of the test to validate that all timers created have either executed
+// or are cancelled.
+func TrackTimers() {
+	globalTimerTracker = &timerFactory{
+		allocatedTimers: make(map[internal.Timer][]uintptr),
+	}
+	internal.TimeAfterFunc = globalTimerTracker.timeAfterFunc
+}
+
+// CheckTimers undoes the effects of TrackTimers, and fails unit tests if not
+// all timers were cancelled or executed. It is invalid to invoke this function
+// without previously having invoked TrackTimers.
+func CheckTimers(ctx context.Context, logger Logger) {
+	tt := globalTimerTracker
+
+	// Loop, waiting for timers to be cancelled.
+	// Wait up to timeout, but finish as quickly as possible.
+	var leaked []string
+	for ctx.Err() == nil {
+		if leaked = tt.pendingTimers(); len(leaked) == 0 {
+			return
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+	for _, g := range leaked {
+		logger.Errorf("Leaked timers: %v", g)
+	}
+
+	// Reset the internal function.
+	internal.TimeAfterFunc = func(d time.Duration, f func()) internal.Timer {
+		return time.AfterFunc(d, f)
+	}
+}
+
+func currentStack(skip int) []uintptr {
+	var stackBuf [16]uintptr
+	var stack []uintptr
+	skip++
+	for {
+		n := runtime.Callers(skip, stackBuf[:])
+		stack = append(stack, stackBuf[:n]...)
+		if n < len(stackBuf) {
+			break
+		}
+		skip += len(stackBuf)
+	}
+	return stack
+}
+
+func traceToString(stack []uintptr) string {
+	frames := runtime.CallersFrames(stack)
+	var trace strings.Builder
+	for {
+		f, ok := frames.Next()
+		if !ok {
+			break
+		}
+		trace.WriteString(f.Function)
+		trace.WriteString("\n\t")
+		trace.WriteString(f.File)
+		trace.WriteString(":")
+		trace.WriteString(strconv.Itoa(f.Line))
+		trace.WriteString("\n")
+	}
+	return trace.String()
+}
+
+// Async Reporter Leak Checking
+
+var asyncReporterTracker *reporterTracker
+
+type reporterTracker struct {
+	mu          sync.Mutex
+	allocations map[*int][]uintptr
+}
+
+func newReporterTracker() *reporterTracker {
+	return &reporterTracker{
+		allocations: make(map[*int][]uintptr),
+	}
+}
+
+// register records the stack trace.
+func (rt *reporterTracker) register() *int {
+	rt.mu.Lock()
+	defer rt.mu.Unlock()
+
+	id := new(int)
+	// Skip 4 frames: register -> internal.Delegate -> stats.RegisterAsyncReporter -> Caller
+	rt.allocations[id] = currentStack(4)
+	return id
+}
+
+// unregister removes the ID.
+func (rt *reporterTracker) unregister(id *int) {
+	rt.mu.Lock()
+	defer rt.mu.Unlock()
+	delete(rt.allocations, id)
+}
+
+// leakedStackTraces returns formatted stack traces for all currently registered
+// reporters.
+func (rt *reporterTracker) leakedStackTraces() []string {
+	rt.mu.Lock()
+	defer rt.mu.Unlock()
+
+	var traces []string
+	for _, pcs := range rt.allocations {
+		msg := "\n--- Leaked Async Reporter Registration ---\n" + traceToString(pcs)
+		traces = append(traces, msg)
+	}
+	return traces
+}
+
+// TrackAsyncReporters installs the tracking delegate.
+func TrackAsyncReporters() {
+	asyncReporterTracker = newReporterTracker()
+
+	// Swap the delegate: Replace the default pass-through with tracking logic.
+	internal.AsyncReporterCleanupDelegate = func(originalCleanup func()) func() {
+		// 1. Capture Stack Trace (happens during Registration)
+		token := asyncReporterTracker.register()
+
+		// 2. Return Wrapped Cleanup
+		return func() {
+			// Defer unregister to ensure we stop tracking even if the original cleanup panics.
+			defer asyncReporterTracker.unregister(token)
+
+			if originalCleanup != nil {
+				originalCleanup()
+			}
+		}
+	}
+}
+
+// CheckAsyncReporters verifies that no leaks exist and restores the default delegate.
+func CheckAsyncReporters(logger Logger) {
+	// Restore the delegate: Reset to the default pass-through behavior.
+	internal.AsyncReporterCleanupDelegate = func(cleanup func()) func() {
+		return cleanup
+	}
+
+	if asyncReporterTracker == nil {
+		return
+	}
+
+	leaks := asyncReporterTracker.leakedStackTraces()
+	if len(leaks) > 0 {
+		// Join all stack traces into one message
+		allTraces := ""
+		for _, trace := range leaks {
+			allTraces += trace
+		}
+		logger.Errorf("Found %d leaked async reporters:%s", len(leaks), allTraces)
+	}
+
+	// Clean up global state
+	asyncReporterTracker = nil
+}

internal/leakcheck/leakcheck_enabled.go

@@ -0,0 +1,25 @@
+//go:build checkbuffers
+
+/*
+ *
+ * Copyright 2017 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package leakcheck
+
+func init() {
+	failTestsOnLeakedBuffers = true
+}

internal/leakcheck/leakcheck_test.go

@@ -0,0 +1,173 @@
+/*
+ *
+ * Copyright 2017 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package leakcheck
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/internal"
+)
+
+type testLogger struct {
+	errorCount int
+	errors     []string
+}
+
+func (e *testLogger) Logf(string, ...any) {
+}
+
+func (e *testLogger) Errorf(format string, args ...any) {
+	e.errors = append(e.errors, fmt.Sprintf(format, args...))
+	e.errorCount++
+}
+
+func TestCheck(t *testing.T) {
+	const leakCount = 3
+	ch := make(chan struct{})
+	for i := 0; i < leakCount; i++ {
+		go func() { <-ch }()
+	}
+	if leaked := interestingGoroutines(); len(leaked) != leakCount {
+		t.Errorf("interestingGoroutines() = %d, want length %d", len(leaked), leakCount)
+	}
+	e := &testLogger{}
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if CheckGoroutines(ctx, e); e.errorCount < leakCount {
+		t.Errorf("CheckGoroutines() = %d, want count %d", e.errorCount, leakCount)
+		t.Logf("leaked goroutines:\n%v", strings.Join(e.errors, "\n"))
+	}
+	close(ch)
+	ctx, cancel = context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	CheckGoroutines(ctx, t)
+}
+
+func ignoredTestingLeak(d time.Duration) {
+	time.Sleep(d)
+}
+
+func TestCheckRegisterIgnore(t *testing.T) {
+	RegisterIgnoreGoroutine("ignoredTestingLeak")
+	go ignoredTestingLeak(3 * time.Second)
+	const leakCount = 3
+	ch := make(chan struct{})
+	for i := 0; i < leakCount; i++ {
+		go func() { <-ch }()
+	}
+	if leaked := interestingGoroutines(); len(leaked) != leakCount {
+		t.Errorf("interestingGoroutines() = %d, want length %d", len(leaked), leakCount)
+	}
+	e := &testLogger{}
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if CheckGoroutines(ctx, e); e.errorCount < leakCount {
+		t.Errorf("CheckGoroutines() = %d, want count %d", e.errorCount, leakCount)
+		t.Logf("leaked goroutines:\n%v", strings.Join(e.errors, "\n"))
+	}
+	close(ch)
+	ctx, cancel = context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	CheckGoroutines(ctx, t)
+}
+
+// TestTrackTimers verifies that only leaked timers are reported and expired,
+// stopped timers are ignored.
+func TestTrackTimers(t *testing.T) {
+	TrackTimers()
+	const leakCount = 3
+	for i := 0; i < leakCount; i++ {
+		internal.TimeAfterFunc(2*time.Second, func() {
+			t.Logf("Timer %d fired.", i)
+		})
+	}
+	wg := sync.WaitGroup{}
+	// Let a couple of timers expire.
+	for i := 0; i < 2; i++ {
+		wg.Add(1)
+		internal.TimeAfterFunc(time.Millisecond, func() {
+			wg.Done()
+		})
+	}
+	wg.Wait()
+
+	// Stop a couple of timers.
+	for i := 0; i < leakCount; i++ {
+		t := internal.TimeAfterFunc(time.Hour, func() {
+			t.Error("Timer fired before test ended.")
+		})
+		t.Stop()
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	e := &testLogger{}
+	CheckTimers(ctx, e)
+	if e.errorCount != leakCount {
+		t.Errorf("CheckTimers found %v leaks, want %v leaks", e.errorCount, leakCount)
+		t.Logf("leaked timers:\n%v", strings.Join(e.errors, "\n"))
+	}
+	ctx, cancel = context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	CheckTimers(ctx, t)
+}
+
+func TestLeakChecker_DetectsLeak(t *testing.T) {
+	TrackAsyncReporters()
+	// Safety defer: ensure we restore the default delegate even if the test crashes
+	// before CheckAsyncReporters is called.
+	defer func() {
+		internal.AsyncReporterCleanupDelegate = func(f func()) func() { return f }
+	}()
+
+	// We utilize the internal delegate directly to simulate stats.RegisterAsyncReporter behavior.
+	noOpCleanup := func() {}
+	wrappedCleanup := internal.AsyncReporterCleanupDelegate(noOpCleanup)
+
+	// Create a leak: We discard 'wrappedCleanup' without calling it.
+	_ = wrappedCleanup
+	tl := &testLogger{}
+	CheckAsyncReporters(tl)
+	if tl.errorCount == 0 {
+		t.Error("Expected leak checker to report a leak, but it succeeded silently.")
+	}
+	if asyncReporterTracker != nil {
+		t.Error("Expected CheckAsyncReporters to cleanup global tracker, but it was not nil.")
+	}
+}
+
+func TestLeakChecker_PassesOnCleanup(t *testing.T) {
+	TrackAsyncReporters()
+	defer func() {
+		internal.AsyncReporterCleanupDelegate = func(f func()) func() { return f }
+	}()
+	noOpCleanup := func() {}
+	wrappedCleanup := internal.AsyncReporterCleanupDelegate(noOpCleanup)
+	wrappedCleanup()
+	tl := &testLogger{}
+	CheckAsyncReporters(tl)
+	if tl.errorCount > 0 {
+		t.Errorf("Expected no leaks, but got errors: %v", tl.errors)
+	}
+}