add leak check to the main package

Signed-off-by: Ryan Leung <[email protected]>
tikv · Apr 11, 2024 · 3ad507d · 3ad507d
1 parent f0eb74b
commit 3ad507d
Show file tree

Hide file tree

Showing 35 changed files with 191 additions and 30 deletions.
diff --git a/client/http/client_test.go b/client/http/client_test.go
@@ -24,9 +24,15 @@ import (
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/pd/client/errs"
 	"github.com/tikv/pd/client/retry"
+	"github.com/tikv/pd/client/testutil"
 	"go.uber.org/atomic"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 func TestPDAllowFollowerHandleHeader(t *testing.T) {
 	re := require.New(t)
 	checked := 0

diff --git a/client/resource_group/controller/controller_test.go b/client/resource_group/controller/controller_test.go
@@ -26,8 +26,14 @@ import (
 
 	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/client/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 func createTestGroupCostController(re *require.Assertions) *groupCostController {
 	group := &rmpb.ResourceGroup{
 		Name:     "test",

diff --git a/pkg/balancer/balancer_test.go b/pkg/balancer/balancer_test.go
@@ -19,8 +19,14 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 func TestBalancerPutAndDelete(t *testing.T) {
 	re := require.New(t)
 	balancers := []Balancer[uint32]{

diff --git a/pkg/core/factory_test.go b/pkg/core/factory_test.go
@@ -19,9 +19,15 @@ import (
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/pingcap/kvproto/pkg/metapb"
+	"github.com/tikv/pd/pkg/utils/testutil"
 	"github.com/tikv/pd/pkg/utils/typeutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 func BenchmarkDeepClone(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		src := &metapb.Region{Id: 1}

diff --git a/pkg/core/rangetree/range_tree_test.go b/pkg/core/rangetree/range_tree_test.go
@@ -19,8 +19,14 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 type simpleBucketItem struct {
 	startKey []byte
 	endKey   []byte

diff --git a/pkg/core/storelimit/limit_test.go b/pkg/core/storelimit/limit_test.go
@@ -24,8 +24,14 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/pd/pkg/core/constant"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 func TestStoreLimit(t *testing.T) {
 	re := require.New(t)
 	rate := int64(15)

diff --git a/pkg/id/id_test.go b/pkg/id/id_test.go
@@ -22,8 +22,14 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/pd/pkg/utils/etcdutil"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 const (
 	rootPath   = "/pd"
 	leaderPath = "/pd/leader"

diff --git a/pkg/keyspace/keyspace_test.go b/pkg/keyspace/keyspace_test.go
@@ -32,9 +32,15 @@ import (
 	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/storage/kv"
 	"github.com/tikv/pd/pkg/utils/etcdutil"
+	"github.com/tikv/pd/pkg/utils/testutil"
 	"github.com/tikv/pd/pkg/utils/typeutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 const (
 	testConfig  = "test config"
 	testConfig1 = "config_entry_1"

diff --git a/pkg/ratelimit/concurrency_limiter_test.go b/pkg/ratelimit/concurrency_limiter_test.go
@@ -23,8 +23,14 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 func TestConcurrencyLimiter(t *testing.T) {
 	re := require.New(t)
 	cl := NewConcurrencyLimiter(10)

diff --git a/pkg/schedule/config/store_config_test.go b/pkg/schedule/config/store_config_test.go
@@ -18,8 +18,14 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 func TestMergeCheck(t *testing.T) {
 	re := require.New(t)
 	testdata := []struct {

diff --git a/pkg/schedule/filter/candidates_test.go b/pkg/schedule/filter/candidates_test.go
@@ -22,8 +22,14 @@ import (
 	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/schedule/config"
 	"github.com/tikv/pd/pkg/schedule/plan"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 // A dummy comparer for testing.
 func idComparer(a, b *core.StoreInfo) int {
 	if a.GetID() > b.GetID() {

diff --git a/pkg/schedule/labeler/labeler_test.go b/pkg/schedule/labeler/labeler_test.go
@@ -32,8 +32,13 @@ import (
 	"github.com/tikv/pd/pkg/storage/kv"
 	"github.com/tikv/pd/pkg/utils/etcdutil"
 	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 func TestAdjustRule(t *testing.T) {
 	re := require.New(t)
 	rule := LabelRule{
@@ -95,7 +100,9 @@ func TestAdjustRule2(t *testing.T) {
 func TestGetSetRule(t *testing.T) {
 	re := require.New(t)
 	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
-	labeler, err := NewRegionLabeler(context.Background(), store, time.Millisecond*10)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	labeler, err := NewRegionLabeler(ctx, store, time.Millisecond*10)
 	re.NoError(err)
 	rules := []*LabelRule{
 		{ID: "rule1", Labels: []RegionLabel{{Key: "k1", Value: "v1"}}, RuleType: "key-range", Data: MakeKeyRanges("1234", "5678")},
@@ -148,7 +155,9 @@ func TestTxnWithEtcd(t *testing.T) {
 	_, client, clean := etcdutil.NewTestEtcdCluster(t, 1)
 	defer clean()
 	store := storage.NewStorageWithEtcdBackend(client, "")
-	labeler, err := NewRegionLabeler(context.Background(), store, time.Millisecond*10)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	labeler, err := NewRegionLabeler(ctx, store, time.Millisecond*10)
 	re.NoError(err)
 	// test patch rules in batch
 	rulesNum := 200
@@ -213,7 +222,9 @@ func TestTxnWithEtcd(t *testing.T) {
 func TestIndex(t *testing.T) {
 	re := require.New(t)
 	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
-	labeler, err := NewRegionLabeler(context.Background(), store, time.Millisecond*10)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	labeler, err := NewRegionLabeler(ctx, store, time.Millisecond*10)
 	re.NoError(err)
 	rules := []*LabelRule{
 		{ID: "rule0", Labels: []RegionLabel{{Key: "k1", Value: "v0"}}, RuleType: "key-range", Data: MakeKeyRanges("", "")},
@@ -255,7 +266,9 @@ func TestIndex(t *testing.T) {
 func TestSaveLoadRule(t *testing.T) {
 	re := require.New(t)
 	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
-	labeler, err := NewRegionLabeler(context.Background(), store, time.Millisecond*10)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	labeler, err := NewRegionLabeler(ctx, store, time.Millisecond*10)
 	re.NoError(err)
 	rules := []*LabelRule{
 		{ID: "rule1", Labels: []RegionLabel{{Key: "k1", Value: "v1"}}, RuleType: "key-range", Data: MakeKeyRanges("1234", "5678")},
@@ -266,8 +279,9 @@ func TestSaveLoadRule(t *testing.T) {
 		err := labeler.SetLabelRule(r)
 		re.NoError(err)
 	}
-
-	labeler, err = NewRegionLabeler(context.Background(), store, time.Millisecond*100)
+	ctx1, cancel1 := context.WithCancel(context.Background())
+	defer cancel1()
+	labeler, err = NewRegionLabeler(ctx1, store, time.Millisecond*100)
 	re.NoError(err)
 	for _, r := range rules {
 		r2 := labeler.GetLabelRule(r.ID)
@@ -300,7 +314,9 @@ func expectSameRules(re *require.Assertions, r1, r2 *LabelRule) {
 func TestKeyRange(t *testing.T) {
 	re := require.New(t)
 	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
-	labeler, err := NewRegionLabeler(context.Background(), store, time.Millisecond*10)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	labeler, err := NewRegionLabeler(ctx, store, time.Millisecond*10)
 	re.NoError(err)
 	rules := []*LabelRule{
 		{ID: "rule1", Labels: []RegionLabel{{Key: "k1", Value: "v1"}}, RuleType: "key-range", Data: MakeKeyRanges("1234", "5678")},
@@ -341,7 +357,9 @@ func TestKeyRange(t *testing.T) {
 func TestLabelerRuleTTL(t *testing.T) {
 	re := require.New(t)
 	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
-	labeler, err := NewRegionLabeler(context.Background(), store, time.Millisecond*10)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	labeler, err := NewRegionLabeler(ctx, store, time.Millisecond*10)
 	re.NoError(err)
 	rules := []*LabelRule{
 		{
@@ -416,7 +434,9 @@ func TestGC(t *testing.T) {
 	re := require.New(t)
 	// set gcInterval to 1 hour.
 	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
-	labeler, err := NewRegionLabeler(context.Background(), store, time.Hour)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	labeler, err := NewRegionLabeler(ctx, store, time.Hour)
 	re.NoError(err)
 	ttls := []string{"1ms", "1ms", "1ms", "5ms", "5ms", "10ms", "1h", "24h"}
 	start, _ := hex.DecodeString("1234")

diff --git a/pkg/schedule/operator/builder_test.go b/pkg/schedule/operator/builder_test.go
@@ -25,8 +25,14 @@ import (
 	"github.com/tikv/pd/pkg/mock/mockcluster"
 	"github.com/tikv/pd/pkg/mock/mockconfig"
 	"github.com/tikv/pd/pkg/schedule/config"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 type operatorBuilderTestSuite struct {
 	suite.Suite
 

diff --git a/pkg/schedule/operator/step_test.go b/pkg/schedule/operator/step_test.go
@@ -29,6 +29,7 @@ import (
 type operatorStepTestSuite struct {
 	suite.Suite
 
+	cancel  context.CancelFunc
 	cluster *mockcluster.Cluster
 }
 
@@ -44,7 +45,9 @@ type testCase struct {
 }
 
 func (suite *operatorStepTestSuite) SetupTest() {
-	suite.cluster = mockcluster.NewCluster(context.Background(), mockconfig.NewTestOptions())
+	ctx, cancel := context.WithCancel(context.Background())
+	suite.cancel = cancel
+	suite.cluster = mockcluster.NewCluster(ctx, mockconfig.NewTestOptions())
 	for i := 1; i <= 10; i++ {
 		suite.cluster.PutStoreWithLabels(uint64(i))
 	}
@@ -53,6 +56,12 @@ func (suite *operatorStepTestSuite) SetupTest() {
 	suite.cluster.SetStoreDown(10)
 }
 
+func (suite *operatorStepTestSuite) TearDownTest() {
+	if suite.cancel != nil {
+		suite.cancel()
+	}
+}
+
 func (suite *operatorStepTestSuite) TestTransferLeader() {
 	re := suite.Require()
 	step := TransferLeader{FromStore: 1, ToStore: 2}

diff --git a/pkg/schedule/placement/config_test.go b/pkg/schedule/placement/config_test.go
@@ -18,8 +18,14 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 func TestTrim(t *testing.T) {
 	re := require.New(t)
 	rc := newRuleConfig()

diff --git a/pkg/schedule/plan/balance_plan_test.go b/pkg/schedule/plan/balance_plan_test.go
@@ -21,8 +21,14 @@ import (
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/stretchr/testify/suite"
 	"github.com/tikv/pd/pkg/core"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 type balanceSchedulerPlanAnalyzeTestSuite struct {
 	suite.Suite
 

diff --git a/pkg/schedule/rangelist/range_list_test.go b/pkg/schedule/rangelist/range_list_test.go
@@ -18,8 +18,14 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 func TestRangeList(t *testing.T) {
 	re := require.New(t)
 	rl := NewBuilder().Build()

diff --git a/pkg/schedule/scatter/region_scatterer_test.go b/pkg/schedule/scatter/region_scatterer_test.go
@@ -33,9 +33,15 @@ import (
 	"github.com/tikv/pd/pkg/schedule/hbstream"
 	"github.com/tikv/pd/pkg/schedule/operator"
 	"github.com/tikv/pd/pkg/schedule/placement"
+	"github.com/tikv/pd/pkg/utils/testutil"
 	"github.com/tikv/pd/pkg/versioninfo"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 type sequencer struct {
 	minID uint64
 	maxID uint64

diff --git a/pkg/schedule/schedulers/balance_leader_test.go b/pkg/schedule/schedulers/balance_leader_test.go
@@ -21,8 +21,14 @@ import (
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/mock/mockcluster"
+	"github.com/tikv/pd/pkg/utils/testutil"
+	"go.uber.org/goleak"
 )
 
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m, testutil.LeakOptions...)
+}
+
 func TestBalanceLeaderSchedulerConfigClone(t *testing.T) {
 	re := require.New(t)
 	keyRanges1, _ := getKeyRanges([]string{"a", "b", "c", "d"})