*: make test leak great again (pingcap#4543)

Makefile

@@ -120,10 +120,7 @@ race: parserlib
 
 leak: parserlib
 	@export log_level=debug; \
-	for dir in $(PACKAGES); do \
-		echo $$dir; \
-		$(GOTEST) -tags leak $$dir | awk 'END{if($$1=="FAIL") {exit 1}}' || exit 1; \
-	done;
+	$(GOTEST) -tags leak $(PACKAGES)
 
 tikv_integration_test: parserlib
 	$(GOTEST) ./store/tikv/. -with-tikv=true

bootstrap_test.go

@@ -41,7 +41,9 @@ func (s *testBootstrapSuite) SetUpSuite(c *C) {
 
 func (s *testBootstrapSuite) TestBootstrap(c *C) {
 	defer testleak.AfterTest(c)()
-	store := newStoreWithBootstrap(c, s.dbName)
+	store, dom := newStoreWithBootstrap(c, s.dbName)
+	defer dom.Close()
+	defer store.Close()
 	se := newSession(c, store, s.dbName)
 	mustExecSQL(c, se, "USE mysql;")
 	r := mustExecSQL(c, se, `select * from user;`)
@@ -168,7 +170,7 @@ func (s *testBootstrapSuite) testBootstrapWithError(c *C) {
 // TestUpgrade tests upgrading
 func (s *testBootstrapSuite) TestUpgrade(c *C) {
 	defer testleak.AfterTest(c)()
-	store := newStoreWithBootstrap(c, s.dbName)
+	store, _ := newStoreWithBootstrap(c, s.dbName)
 	defer store.Close()
 	se := newSession(c, store, s.dbName)
 	mustExecSQL(c, se, "USE mysql;")
@@ -211,7 +213,9 @@ func (s *testBootstrapSuite) TestUpgrade(c *C) {
 	c.Assert(ver, Equals, int64(0))
 
 	// Create a new session then upgrade() will run automatically.
-	BootstrapSession(store)
+	dom1, err := BootstrapSession(store)
+	c.Assert(err, IsNil)
+	defer dom1.Close()
 	se2 := newSession(c, store, s.dbName)
 	r = mustExecSQL(c, se2, `SELECT VARIABLE_VALUE from mysql.TiDB where VARIABLE_NAME="tidb_server_version";`)
 	row, err = r.Next()

ddl/ddl_db_test.go

@@ -101,7 +101,6 @@ func (s *testDBSuite) testErrorCode(c *C, sql string, errCode int) {
 }
 
 func (s *testDBSuite) TestMySQLErrorCode(c *C) {
-	defer testleak.AfterTest(c)
 	s.tk = testkit.NewTestKit(c, s.store)
 	s.tk.MustExec("use " + s.schemaName)
 
@@ -598,7 +597,6 @@ func (s *testDBSuite) TestIssue2293(c *C) {
 }
 
 func (s *testDBSuite) TestCreateIndexType(c *C) {
-	defer testleak.AfterTest(c)()
 	s.tk = testkit.NewTestKit(c, s.store)
 	s.tk.MustExec("use " + s.schemaName)
 	sql := `CREATE TABLE test_index (
@@ -1019,11 +1017,13 @@ func match(c *C, row []interface{}, expected ...interface{}) {
 }
 
 func (s *testDBSuite) TestUpdateMultipleTable(c *C) {
-	defer testleak.AfterTest(c)
+	defer testleak.AfterTest(c)()
 	store, err := tidb.NewStore("memory://update_multiple_table")
 	c.Assert(err, IsNil)
-	_, err = tidb.BootstrapSession(store)
+	defer store.Close()
+	dom, err := tidb.BootstrapSession(store)
 	c.Assert(err, IsNil)
+	defer dom.Close()
 	tk := testkit.NewTestKit(c, store)
 	tk.MustExec("use test")
 	tk.MustExec("create table t1 (c1 int, c2 int)")
@@ -1081,7 +1081,6 @@ func (s *testDBSuite) TestUpdateMultipleTable(c *C) {
 }
 
 func (s *testDBSuite) TestCreateTableTooLarge(c *C) {
-	defer testleak.AfterTest(c)
 	s.tk = testkit.NewTestKit(c, s.store)
 	s.tk.MustExec("use test")
 
@@ -1104,12 +1103,14 @@ func (s *testDBSuite) TestCreateTableTooLarge(c *C) {
 }
 
 func (s *testDBSuite) TestCreateTableWithLike(c *C) {
-	defer testleak.AfterTest(c)
+	defer testleak.AfterTest(c)()
 	store, err := tidb.NewStore("memory://create_table_like")
 	c.Assert(err, IsNil)
+	defer store.Close()
 	s.tk = testkit.NewTestKit(c, store)
-	_, err = tidb.BootstrapSession(store)
+	dom, err := tidb.BootstrapSession(store)
 	c.Assert(err, IsNil)
+	defer dom.Close()
 
 	// for the same database
 	s.tk.MustExec("use test")
@@ -1153,7 +1154,6 @@ func (s *testDBSuite) TestCreateTableWithLike(c *C) {
 }
 
 func (s *testDBSuite) TestCreateTable(c *C) {
-	defer testleak.AfterTest(c)
 	store, err := tidb.NewStore("memory://create_table")
 	c.Assert(err, IsNil)
 	s.tk = testkit.NewTestKit(c, store)
@@ -1176,11 +1176,13 @@ func (s *testDBSuite) TestCreateTable(c *C) {
 }
 
 func (s *testDBSuite) TestTruncateTable(c *C) {
-	defer testleak.AfterTest(c)
+	defer testleak.AfterTest(c)()
 	store, err := tidb.NewStore("memory://truncate_table")
 	c.Assert(err, IsNil)
-	_, err = tidb.BootstrapSession(store)
+	defer store.Close()
+	dom, err := tidb.BootstrapSession(store)
 	c.Assert(err, IsNil)
+	defer dom.Close()
 	tk := testkit.NewTestKit(c, store)
 	tk.MustExec("use test")
 	tk.MustExec("create table t (c1 int, c2 int)")
@@ -1236,11 +1238,13 @@ func (s *testDBSuite) TestAlterTableRenameTable(c *C) {
 }
 
 func (s *testDBSuite) testRenameTable(c *C, storeStr, sql string) {
-	defer testleak.AfterTest(c)
+	defer testleak.AfterTest(c)()
 	store, err := tidb.NewStore("memory://" + storeStr)
 	c.Assert(err, IsNil)
-	_, err = tidb.BootstrapSession(store)
+	defer store.Close()
+	dom, err := tidb.BootstrapSession(store)
 	c.Assert(err, IsNil)
+	defer dom.Close()
 	s.tk = testkit.NewTestKit(c, store)
 	s.tk.MustExec("use test")
 
@@ -1287,11 +1291,13 @@ func (s *testDBSuite) testRenameTable(c *C, storeStr, sql string) {
 }
 
 func (s *testDBSuite) TestRenameMultiTables(c *C) {
-	defer testleak.AfterTest(c)
+	defer testleak.AfterTest(c)()
 	store, err := tidb.NewStore("memory://rename_multi_tables")
 	c.Assert(err, IsNil)
-	_, err = tidb.BootstrapSession(store)
+	defer store.Close()
+	dom, err := tidb.BootstrapSession(store)
 	c.Assert(err, IsNil)
+	defer dom.Close()
 	s.tk = testkit.NewTestKit(c, store)
 	s.tk.MustExec("use test")
 	s.tk.MustExec("create table t1(id int)")
@@ -1305,7 +1311,7 @@ func (s *testDBSuite) TestRenameMultiTables(c *C) {
 }
 
 func (s *testDBSuite) TestAddNotNullColumn(c *C) {
-	defer testleak.AfterTest(c)
+	defer testleak.AfterTest(c)()
 	s.tk = testkit.NewTestKit(c, s.store)
 	s.tk.MustExec("use test_db")
 	// for different databases
@@ -1431,9 +1437,6 @@ func (s *testDBSuite) TestChangeColumnPosition(c *C) {
 }
 
 func (s *testDBSuite) TestGeneratedColumnDDL(c *C) {
-	defer func() {
-		testleak.AfterTest(c)()
-	}()
 	s.tk = testkit.NewTestKit(c, s.store)
 	s.tk.MustExec("use test")
 

domain/domain.go

@@ -51,6 +51,7 @@ type Domain struct {
 	sysSessionPool  *pools.ResourcePool
 	exit            chan struct{}
 	etcdClient      *clientv3.Client
+	wg              sync.WaitGroup
 
 	MockReloadFailed MockFailure // It mocks reload failed.
 }
@@ -348,6 +349,7 @@ func (do *Domain) Close() {
 		do.etcdClient.Close()
 	}
 	do.sysSessionPool.Close()
+	do.wg.Wait()
 }
 
 type ddlCallback struct {
@@ -539,6 +541,7 @@ func (do *Domain) UpdateTableStatsLoop(ctx context.Context) error {
 	if lease <= 0 {
 		return nil
 	}
+	do.wg.Add(1)
 	go do.updateStatsWorker(ctx, lease)
 	return nil
 }
@@ -558,6 +561,7 @@ func (do *Domain) updateStatsWorker(ctx context.Context, lease time.Duration) {
 				log.Error("[stats] update stats info fail: ", errors.ErrorStack(err))
 			}
 		case <-do.exit:
+			do.wg.Done()
 			return
 			// This channel is sent only by ddl owner or the drop stats executor.
 		case t := <-statsHandle.DDLEventCh():

executor/executor_test.go

@@ -100,15 +100,19 @@ func (s *testSuite) TearDownSuite(c *C) {
 	atomic.StoreInt32(&expression.TurnOnNewExprEval, 0)
 }
 
+func (s *testSuite) SetUpTest(c *C) {
+	testleak.BeforeTest()
+}
+
 func (s *testSuite) TearDownTest(c *C) {
-	testleak.AfterTest(c)()
 	tk := testkit.NewTestKit(c, s.store)
 	tk.MustExec("use test")
 	r := tk.MustQuery("show tables")
 	for _, tb := range r.Rows() {
 		tableName := tb[0]
 		tk.MustExec(fmt.Sprintf("drop table %v", tableName))
 	}
+	testleak.AfterTest(c)()
 }
 
 func (s *testSuite) TestAdmin(c *C) {

expression/integration_test.go

@@ -22,6 +22,7 @@ import (
 	. "github.com/pingcap/check"
 	"github.com/pingcap/tidb"
 	"github.com/pingcap/tidb/context"
+	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/mysql"
 	"github.com/pingcap/tidb/terror"
@@ -37,6 +38,7 @@ var _ = Suite(&testIntegrationSuite{})
 
 type testIntegrationSuite struct {
 	store kv.Storage
+	dom   *domain.Domain
 	ctx   context.Context
 }
 
@@ -51,10 +53,17 @@ func (s *testIntegrationSuite) cleanEnv(c *C) {
 }
 
 func (s *testIntegrationSuite) SetUpSuite(c *C) {
-	s.store, _ = newStoreWithBootstrap()
+	var err error
+	s.store, s.dom, err = newStoreWithBootstrap()
+	c.Assert(err, IsNil)
 	s.ctx = mock.NewContext()
 }
 
+func (s *testIntegrationSuite) TearDownSuite(c *C) {
+	s.dom.Close()
+	s.store.Close()
+}
+
 func (s *testIntegrationSuite) TestFuncREPEAT(c *C) {
 	tk := testkit.NewTestKit(c, s.store)
 	defer func() {

expression/typeinferer_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/pingcap/tidb"
 	"github.com/pingcap/tidb/ast"
 	"github.com/pingcap/tidb/context"
+	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/expression"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/model"
@@ -39,9 +40,10 @@ type testTypeInferrerSuite struct {
 
 func (ts *testTypeInferrerSuite) TestInferType(c *C) {
 	c.Skip("we re-implement this test in plan/typeinfer_test.go")
-	store, err := newStoreWithBootstrap()
+	store, dom, err := newStoreWithBootstrap()
 	c.Assert(err, IsNil)
 	defer store.Close()
+	defer dom.Close()
 	testKit := testkit.NewTestKit(c, store)
 	testKit.MustExec("use test")
 	sql := `create table t (
@@ -370,9 +372,10 @@ func (ts *testTypeInferrerSuite) TestInferType(c *C) {
 
 func (s *testTypeInferrerSuite) TestColumnInfoModified(c *C) {
 	defer testleak.AfterTest(c)()
-	store, err := newStoreWithBootstrap()
+	store, dom, err := newStoreWithBootstrap()
 	c.Assert(err, IsNil)
 	defer store.Close()
+	defer dom.Close()
 	testKit := testkit.NewTestKit(c, store)
 	testKit.MustExec("use test")
 	testKit.MustExec("drop table if exists tab0")
@@ -387,9 +390,10 @@ func (s *testTypeInferrerSuite) TestColumnInfoModified(c *C) {
 
 func (s *testTypeInferrerSuite) TestIsHybridType(c *C) {
 	defer testleak.AfterTest(c)()
-	store, err := newStoreWithBootstrap()
+	store, dom, err := newStoreWithBootstrap()
 	c.Assert(err, IsNil)
 	defer store.Close()
+	defer dom.Close()
 	testKit := testkit.NewTestKit(c, store)
 	testKit.MustExec("use test")
 	testKit.MustExec("drop table if exists t")
@@ -436,12 +440,12 @@ func (s *testTypeInferrerSuite) TestIsHybridType(c *C) {
 	}
 }
 
-func newStoreWithBootstrap() (kv.Storage, error) {
+func newStoreWithBootstrap() (kv.Storage, *domain.Domain, error) {
 	store, err := tikv.NewMockTikvStore()
 	if err != nil {
-		return nil, errors.Trace(err)
+		return nil, nil, errors.Trace(err)
 	}
 	tidb.SetSchemaLease(0)
-	_, err = tidb.BootstrapSession(store)
-	return store, errors.Trace(err)
+	dom, err := tidb.BootstrapSession(store)
+	return store, dom, errors.Trace(err)
 }

new_session_test.go

@@ -44,9 +44,11 @@ type testSessionSuite struct {
 	cluster   *mocktikv.Cluster
 	mvccStore *mocktikv.MvccStore
 	store     kv.Storage
+	dom       *domain.Domain
 }
 
 func (s *testSessionSuite) SetUpSuite(c *C) {
+	testleak.BeforeTest()
 	s.cluster = mocktikv.NewCluster()
 	mocktikv.BootstrapWithSingleStore(s.cluster)
 	s.mvccStore = mocktikv.NewMvccStore()
@@ -58,12 +60,17 @@ func (s *testSessionSuite) SetUpSuite(c *C) {
 	s.store = store
 	tidb.SetSchemaLease(0)
 	tidb.SetStatsLease(0)
-	_, err = tidb.BootstrapSession(s.store)
+	s.dom, err = tidb.BootstrapSession(s.store)
 	c.Assert(err, IsNil)
 }
 
-func (s *testSessionSuite) TearDownTest(c *C) {
+func (s *testSessionSuite) TearDownSuite(c *C) {
+	s.dom.Close()
+	s.store.Close()
 	testleak.AfterTest(c)()
+}
+
+func (s *testSessionSuite) TearDownTest(c *C) {
 	tk := testkit.NewTestKitWithInit(c, s.store)
 	r := tk.MustQuery("show tables")
 	for _, tb := range r.Rows() {
@@ -563,10 +570,11 @@ type testSchemaSuite struct {
 	mvccStore *mocktikv.MvccStore
 	store     kv.Storage
 	lease     time.Duration
+	dom       *domain.Domain
+	checkLeak func()
 }
 
 func (s *testSchemaSuite) TearDownTest(c *C) {
-	testleak.AfterTest(c)()
 	tk := testkit.NewTestKitWithInit(c, s.store)
 	r := tk.MustQuery("show tables")
 	for _, tb := range r.Rows() {
@@ -576,6 +584,7 @@ func (s *testSchemaSuite) TearDownTest(c *C) {
 }
 
 func (s *testSchemaSuite) SetUpSuite(c *C) {
+	testleak.BeforeTest()
 	s.cluster = mocktikv.NewCluster()
 	mocktikv.BootstrapWithSingleStore(s.cluster)
 	s.mvccStore = mocktikv.NewMvccStore()
@@ -588,8 +597,15 @@ func (s *testSchemaSuite) SetUpSuite(c *C) {
 	s.lease = 20 * time.Millisecond
 	tidb.SetSchemaLease(s.lease)
 	tidb.SetStatsLease(0)
-	_, err = tidb.BootstrapSession(s.store)
+	dom, err := tidb.BootstrapSession(s.store)
 	c.Assert(err, IsNil)
+	s.dom = dom
+}
+
+func (s *testSchemaSuite) TearDownSuite(c *C) {
+	s.dom.Close()
+	s.store.Close()
+	testleak.AfterTest(c)()
 }
 
 func (s *testSchemaSuite) TestSchemaCheckerSQL(c *C) {