1// Copyright 2018 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15// Package leakcheck contains functions to check leaked goroutines.
16//
17// Call "defer leakcheck.Check(t)" at the beginning of tests.
18//
19// This is very closely based on grpc-go's leakcheck.
20package leakcheck
21
22import (
23	"runtime"
24	"sort"
25	"strings"
26	"time"
27)
28
29var goroutinesToIgnore = []string{
30	"net/http/transport.go",
31	"net/http/h2_bundle.go",
32	"src/go.opencensus.io/stats/view/worker.go",
33	"testing.Main(",
34	"testing.tRunner(",
35	"testing.(*M).",
36	"runtime.goexit",
37	"created by runtime.gc",
38	"created by runtime/trace.Start",
39	"interestingGoroutines",
40	"runtime.MHeap_Scavenger",
41	"signal.signal_recv",
42	"sigterm.handler",
43	"runtime_mcall",
44	"(*loggingT).flushDaemon",
45	"goroutine in C code",
46}
47
48// RegisterIgnoreGoroutine appends s into the ignore goroutine list. The
49// goroutines whose stack trace contains s will not be identified as leaked
50// goroutines. Not thread-safe, only call this function in init().
51func RegisterIgnoreGoroutine(s string) {
52	goroutinesToIgnore = append(goroutinesToIgnore, s)
53}
54
55func ignore(g string) bool {
56	sl := strings.SplitN(g, "\n", 2)
57	if len(sl) != 2 {
58		return true
59	}
60	stack := strings.TrimSpace(sl[1])
61	if strings.HasPrefix(stack, "testing.RunTests") {
62		return true
63	}
64
65	if stack == "" {
66		return true
67	}
68
69	for _, s := range goroutinesToIgnore {
70		if strings.Contains(stack, s) {
71			return true
72		}
73	}
74
75	return false
76}
77
78// interestingGoroutines returns all goroutines we care about for the purpose of
79// leak checking. It excludes testing or runtime ones.
80func interestingGoroutines() (gs []string) {
81	buf := make([]byte, 2<<20)
82	buf = buf[:runtime.Stack(buf, true)]
83	for _, g := range strings.Split(string(buf), "\n\n") {
84		if !ignore(g) {
85			gs = append(gs, g)
86		}
87	}
88	sort.Strings(gs)
89	return
90}
91
92// Errorfer is the interface that wraps the Errorf method. It's a subset of
93// testing.TB to make it easy to use Check.
94type Errorfer interface {
95	Errorf(format string, args ...interface{})
96}
97
98func check(efer Errorfer, timeout time.Duration) {
99	// Loop, waiting for goroutines to shut down.
100	// Wait up to timeout, but finish as quickly as possible.
101	deadline := time.Now().Add(timeout)
102	var leaked []string
103	for time.Now().Before(deadline) {
104		if leaked = interestingGoroutines(); len(leaked) == 0 {
105			return
106		}
107		time.Sleep(50 * time.Millisecond)
108	}
109	for _, g := range leaked {
110		efer.Errorf("Leaked goroutine: %v", g)
111	}
112}
113
114// Check looks at the currently-running goroutines and checks if there are any
115// interestring (created by gRPC) goroutines leaked. It waits up to 10 seconds
116// in the error cases.
117func Check(efer Errorfer) {
118	check(efer, 10*time.Second)
119}
120