github.com/XiaoMi/Gaea@v1.2.5/util/testleak/leaktest.go (about)

     1  // Copyright 2013 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Copyright 2016 PingCAP, Inc.
     6  //
     7  // Licensed under the Apache License, Version 2.0 (the "License");
     8  // you may not use this file except in compliance with the License.
     9  // You may obtain a copy of the License at
    10  //
    11  //     http://www.apache.org/licenses/LICENSE-2.0
    12  //
    13  // Unless required by applicable law or agreed to in writing, software
    14  // distributed under the License is distributed on an "AS IS" BASIS,
    15  // See the License for the specific language governing permissions and
    16  // limitations under the License.
    17  // +build leak
    18  
    19  package testleak
    20  
    21  import (
    22  	"runtime"
    23  	"sort"
    24  	"strings"
    25  	"testing"
    26  	"time"
    27  
    28  	"github.com/pingcap/check"
    29  )
    30  
    31  func interestingGoroutines() (gs []string) {
    32  	buf := make([]byte, 2<<20)
    33  	buf = buf[:runtime.Stack(buf, true)]
    34  	for _, g := range strings.Split(string(buf), "\n\n") {
    35  		sl := strings.SplitN(g, "\n", 2)
    36  		if len(sl) != 2 {
    37  			continue
    38  		}
    39  		stack := strings.TrimSpace(sl[1])
    40  		if stack == "" ||
    41  			strings.Contains(stack, "created by github.com/pingcap/tidb.init") ||
    42  			strings.Contains(stack, "testing.RunTests") ||
    43  			strings.Contains(stack, "check.(*resultTracker).start") ||
    44  			strings.Contains(stack, "check.(*suiteRunner).runFunc") ||
    45  			strings.Contains(stack, "check.(*suiteRunner).parallelRun") ||
    46  			strings.Contains(stack, "localstore.(*dbStore).scheduler") ||
    47  			strings.Contains(stack, "tikv.(*noGCHandler).Start") ||
    48  			strings.Contains(stack, "ddl.(*ddl).start") ||
    49  			strings.Contains(stack, "ddl.(*delRange).startEmulator") ||
    50  			strings.Contains(stack, "domain.NewDomain") ||
    51  			strings.Contains(stack, "testing.(*T).Run") ||
    52  			strings.Contains(stack, "domain.(*Domain).LoadPrivilegeLoop") ||
    53  			strings.Contains(stack, "domain.(*Domain).UpdateTableStatsLoop") ||
    54  			strings.Contains(stack, "testing.Main(") ||
    55  			strings.Contains(stack, "runtime.goexit") ||
    56  			strings.Contains(stack, "created by runtime.gc") ||
    57  			strings.Contains(stack, "interestingGoroutines") ||
    58  			strings.Contains(stack, "runtime.MHeap_Scavenger") {
    59  			continue
    60  		}
    61  		gs = append(gs, stack)
    62  	}
    63  	sort.Strings(gs)
    64  	return
    65  }
    66  
    67  var beforeTestGorountines = map[string]bool{}
    68  
    69  // BeforeTest gets the current goroutines.
    70  // It's used for check.Suite.SetUpSuite() function.
    71  // Now it's only used in the tidb_test.go.
    72  func BeforeTest() {
    73  	for _, g := range interestingGoroutines() {
    74  		beforeTestGorountines[g] = true
    75  	}
    76  }
    77  
    78  const defaultCheckCnt = 50
    79  
    80  func checkLeakAfterTest(errorFunc func(cnt int, g string)) func() {
    81  	if len(beforeTestGorountines) == 0 {
    82  		for _, g := range interestingGoroutines() {
    83  			beforeTestGorountines[g] = true
    84  		}
    85  	}
    86  
    87  	cnt := defaultCheckCnt
    88  	return func() {
    89  		defer func() {
    90  			beforeTestGorountines = map[string]bool{}
    91  		}()
    92  
    93  		var leaked []string
    94  		for i := 0; i < cnt; i++ {
    95  			leaked = leaked[:0]
    96  			for _, g := range interestingGoroutines() {
    97  				if !beforeTestGorountines[g] {
    98  					leaked = append(leaked, g)
    99  				}
   100  			}
   101  			// Bad stuff found, but goroutines might just still be
   102  			// shutting down, so give it some time.
   103  			if len(leaked) != 0 {
   104  				time.Sleep(50 * time.Millisecond)
   105  				continue
   106  			}
   107  
   108  			return
   109  		}
   110  		for _, g := range leaked {
   111  			errorFunc(cnt, g)
   112  		}
   113  	}
   114  }
   115  
   116  // AfterTest gets the current goroutines and runs the returned function to
   117  // get the goroutines at that time to contrast whether any goroutines leaked.
   118  // Usage: defer testleak.AfterTest(c)()
   119  // It can call with BeforeTest() at the beginning of check.Suite.TearDownSuite() or
   120  // call alone at the beginning of each test.
   121  func AfterTest(c *check.C) func() {
   122  	errorFunc := func(cnt int, g string) {
   123  		c.Errorf("Test check-count %d appears to have leaked: %v", cnt, g)
   124  	}
   125  	return checkLeakAfterTest(errorFunc)
   126  }
   127  
   128  // AfterTestT is used after all the test cases is finished.
   129  func AfterTestT(t *testing.T) func() {
   130  	errorFunc := func(cnt int, g string) {
   131  		t.Errorf("Test check-count %d appears to have leaked: %v", cnt, g)
   132  	}
   133  	return checkLeakAfterTest(errorFunc)
   134  }