From 067cb1fcbf74396966d96d4727b43ca880e9581a Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Thu, 7 Sep 2017 09:57:31 -0700
Subject: [PATCH] Add RegisterIgnoreGoroutine to leakcheck package (#1507)

---
 test/leakcheck/leakcheck.go      | 74 +++++++++++++++++++++-----------
 test/leakcheck/leakcheck_test.go | 28 ++++++++++++
 2 files changed, 76 insertions(+), 26 deletions(-)

diff --git a/test/leakcheck/leakcheck.go b/test/leakcheck/leakcheck.go
index 84143546..76f9fc54 100644
--- a/test/leakcheck/leakcheck.go
+++ b/test/leakcheck/leakcheck.go
@@ -28,39 +28,61 @@ import (
 	"time"
 )
 
+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",
+}
+
+// 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") {
-		sl := strings.SplitN(g, "\n", 2)
-		if len(sl) != 2 {
-			continue
+		if !ignore(g) {
+			gs = append(gs, g)
 		}
-		stack := strings.TrimSpace(sl[1])
-		if strings.HasPrefix(stack, "testing.RunTests") {
-			continue
-		}
-
-		if stack == "" ||
-			strings.Contains(stack, "testing.Main(") ||
-			strings.Contains(stack, "testing.tRunner(") ||
-			strings.Contains(stack, "testing.(*M).") ||
-			strings.Contains(stack, "runtime.goexit") ||
-			strings.Contains(stack, "created by runtime.gc") ||
-			strings.Contains(stack, "created by runtime/trace.Start") ||
-			strings.Contains(stack, "created by google3/base/go/log.init") ||
-			strings.Contains(stack, "interestingGoroutines") ||
-			strings.Contains(stack, "runtime.MHeap_Scavenger") ||
-			strings.Contains(stack, "signal.signal_recv") ||
-			strings.Contains(stack, "sigterm.handler") ||
-			strings.Contains(stack, "runtime_mcall") ||
-			strings.Contains(stack, "(*loggingT).flushDaemon") ||
-			strings.Contains(stack, "goroutine in C code") {
-			continue
-		}
-		gs = append(gs, g)
 	}
 	sort.Strings(gs)
 	return
diff --git a/test/leakcheck/leakcheck_test.go b/test/leakcheck/leakcheck_test.go
index 9c8cc158..50927e9d 100644
--- a/test/leakcheck/leakcheck_test.go
+++ b/test/leakcheck/leakcheck_test.go
@@ -19,15 +19,19 @@
 package leakcheck
 
 import (
+	"fmt"
+	"strings"
 	"testing"
 	"time"
 )
 
 type testErrorfer struct {
 	errorCount int
+	errors     []string
 }
 
 func (e *testErrorfer) Errorf(format string, args ...interface{}) {
+	e.errors = append(e.errors, fmt.Sprintf(format, args...))
 	e.errorCount++
 }
 
@@ -43,6 +47,30 @@ func TestCheck(t *testing.T) {
 	check(e, time.Second)
 	if e.errorCount != leakCount {
 		t.Errorf("check found %v leaks, want %v leaks", e.errorCount, leakCount)
+		t.Logf("leaked goroutines:\n%v", strings.Join(e.errors, "\n"))
+	}
+	check(t, 3*time.Second)
+}
+
+func ignoredTestingLeak(d time.Duration) {
+	time.Sleep(d)
+}
+
+func TestCheckRegisterIgnore(t *testing.T) {
+	RegisterIgnoreGoroutine("ignoredTestingLeak")
+	const leakCount = 3
+	for i := 0; i < leakCount; i++ {
+		go func() { time.Sleep(2 * time.Second) }()
+	}
+	go func() { ignoredTestingLeak(3 * time.Second) }()
+	if ig := interestingGoroutines(); len(ig) == 0 {
+		t.Error("blah")
+	}
+	e := &testErrorfer{}
+	check(e, time.Second)
+	if e.errorCount != leakCount {
+		t.Errorf("check found %v leaks, want %v leaks", e.errorCount, leakCount)
+		t.Logf("leaked goroutines:\n%v", strings.Join(e.errors, "\n"))
 	}
 	check(t, 3*time.Second)
 }