github.com/searKing/golang/go@v1.2.117/testing/leakcheck/leakcheck.go (about)

     1  // Copyright 2020 The searKing Author. 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  // Package leakcheck contains functions to check leaked goroutines.
     6  //
     7  // Call "defer leakcheck.Check(t)" at the beginning of tests.
     8  // it's borrowed from https://github.com/grpc/grpc-go/blob/master/internal/leakcheck/leakcheck.go
     9  package leakcheck
    10  
    11  import (
    12  	"runtime"
    13  	"sort"
    14  	"strings"
    15  	"time"
    16  )
    17  
    18  var goroutinesToIgnore = []string{
    19  	"testing.Main(",
    20  	"testing.tRunner(",
    21  	"testing.(*M).",
    22  	"runtime.goexit",
    23  	"created by runtime.gc",
    24  	"created by runtime/trace.Start",
    25  	"interestingGoroutines",
    26  	"runtime.MHeap_Scavenger",
    27  	"signal.signal_recv",
    28  	"sigterm.handler",
    29  	"runtime_mcall",
    30  	"(*loggingT).flushDaemon",
    31  	"goroutine in C code",
    32  }
    33  
    34  // RegisterIgnoreGoroutine appends s into the ignore goroutine list. The
    35  // goroutines whose stack trace contains s will not be identified as leaked
    36  // goroutines. Not thread-safe, only call this function in init().
    37  func RegisterIgnoreGoroutine(s string) {
    38  	goroutinesToIgnore = append(goroutinesToIgnore, s)
    39  }
    40  
    41  func ignore(g string) bool {
    42  	sl := strings.SplitN(g, "\n", 2)
    43  	if len(sl) != 2 {
    44  		return true
    45  	}
    46  	stack := strings.TrimSpace(sl[1])
    47  	if strings.HasPrefix(stack, "testing.RunTests") {
    48  		return true
    49  	}
    50  
    51  	if stack == "" {
    52  		return true
    53  	}
    54  
    55  	for _, s := range goroutinesToIgnore {
    56  		if strings.Contains(stack, s) {
    57  			return true
    58  		}
    59  	}
    60  
    61  	return false
    62  }
    63  
    64  // interestingGoroutines returns all goroutines we care about for the purpose of
    65  // leak checking. It excludes testing or runtime ones.
    66  func interestingGoroutines() (gs []string) {
    67  	buf := make([]byte, 2<<20)
    68  	buf = buf[:runtime.Stack(buf, true)]
    69  	for _, g := range strings.Split(string(buf), "\n\n") {
    70  		if !ignore(g) {
    71  			gs = append(gs, g)
    72  		}
    73  	}
    74  	sort.Strings(gs)
    75  	return
    76  }
    77  
    78  // Errorfer is the interface that wraps the Errorf method. It's a subset of
    79  // testing.TB to make it easy to use Check.
    80  type Errorfer interface {
    81  	Errorf(format string, args ...any)
    82  }
    83  
    84  func check(efer Errorfer, timeout time.Duration) {
    85  	// Loop, waiting for goroutines to shut down.
    86  	// Wait up to timeout, but finish as quickly as possible.
    87  	deadline := time.Now().Add(timeout)
    88  	var leaked []string
    89  	for time.Now().Before(deadline) {
    90  		if leaked = interestingGoroutines(); len(leaked) == 0 {
    91  			return
    92  		}
    93  		time.Sleep(50 * time.Millisecond)
    94  	}
    95  	for _, g := range leaked {
    96  		efer.Errorf("Leaked goroutine: %v", g)
    97  	}
    98  }
    99  
   100  // Check looks at the currently-running goroutines and checks if there are any
   101  // interesting (created by gRPC) goroutines leaked. It waits up to 10 seconds
   102  // in the error cases.
   103  func Check(efer Errorfer) {
   104  	check(efer, 10*time.Second)
   105  }