domain: fix memory leak for stats

domain/domain.go

@@ -24,6 +24,7 @@ import (
 	"github.com/coreos/etcd/clientv3"
 	"github.com/grpc-ecosystem/go-grpc-prometheus"
 	"github.com/ngaut/pools"
+	"github.com/ngaut/sync2"
 	"github.com/pingcap/tidb/ast"
 	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/ddl"
@@ -54,6 +55,7 @@ type Domain struct {
 	privHandle      *privileges.Handle
 	statsHandle     unsafe.Pointer
 	statsLease      time.Duration
+	statsUpdating   sync2.AtomicInt32
 	ddl             ddl.DDL
 	info            *InfoSyncer
 	m               sync.Mutex
@@ -664,6 +666,20 @@ func (do *Domain) CreateStatsHandle(ctx sessionctx.Context) {
 	atomic.StorePointer(&do.statsHandle, unsafe.Pointer(statistics.NewHandle(ctx, do.statsLease)))
 }
 
+// StatsUpdating checks if the stats worker is updating.
+func (do *Domain) StatsUpdating() bool {
+	return do.statsUpdating.Get() > 0
+}
+
+// SetStatsUpdating sets the value of stats updating.
+func (do *Domain) SetStatsUpdating(val bool) {
+	if val {
+		do.statsUpdating.Set(1)
+	} else {
+		do.statsUpdating.Set(0)
+	}
+}
+
 // RunAutoAnalyze indicates if this TiDB server starts auto analyze worker and can run auto analyze job.
 var RunAutoAnalyze = true
 
@@ -680,6 +696,7 @@ func (do *Domain) UpdateTableStatsLoop(ctx sessionctx.Context) error {
 	}
 	owner := do.newStatsOwner()
 	do.wg.Add(1)
+	do.SetStatsUpdating(true)
 	go do.updateStatsWorker(ctx, owner)
 	if RunAutoAnalyze {
 		do.wg.Add(1)
@@ -729,6 +746,7 @@ func (do *Domain) updateStatsWorker(ctx sessionctx.Context, owner owner.Manager)
 		log.Info("[stats] init stats info takes ", time.Now().Sub(t))
 	}
 	defer func() {
+		do.SetStatsUpdating(false)
 		recoverInDomain("updateStatsWorker", false)
 		do.wg.Done()
 	}()

executor/executor_test.go

@@ -113,6 +113,7 @@ func (s *testSuite) SetUpSuite(c *C) {
 	}
 	d, err := session.BootstrapSession(s.store)
 	c.Assert(err, IsNil)
+	d.SetStatsUpdating(true)
 	s.domain = d
 }
 

infoschema/tables_test.go

@@ -93,6 +93,7 @@ func (s *testSuite) TestDataForTableStatsField(c *C) {
 	session.SetStatsLease(0)
 	do, err := session.BootstrapSession(store)
 	c.Assert(err, IsNil)
+	do.SetStatsUpdating(true)
 	defer do.Close()
 	oldExpiryTime := infoschema.TableStatsCacheExpiry
 	infoschema.TableStatsCacheExpiry = 0

planner/core/cbo_test.go

@@ -693,6 +693,7 @@ func newStoreWithBootstrap() (kv.Storage, *domain.Domain, error) {
 	session.SetSchemaLease(0)
 	session.SetStatsLease(0)
 	dom, err := session.BootstrapSession(store)
+	dom.SetStatsUpdating(true)
 	return store, dom, errors.Trace(err)
 }
 

server/statistics_handler_test.go

@@ -49,6 +49,7 @@ func (ds *testDumpStatsSuite) startServer(c *C) {
 	session.SetStatsLease(0)
 	ds.domain, err = session.BootstrapSession(ds.store)
 	c.Assert(err, IsNil)
+	ds.domain.SetStatsUpdating(true)
 	tidbdrv := NewTiDBDriver(ds.store)
 
 	cfg := config.NewConfig()

session/session.go

@@ -1065,8 +1065,9 @@ func CreateSession(store kv.Storage) (Session, error) {
 	}
 	privilege.BindPrivilegeManager(s, pm)
 
-	// Add statsUpdateHandle.
-	if do.StatsHandle() != nil {
+	// Add stats collector, and it will be freed by background stats worker
+	// which periodically updates stats using the collected data.
+	if do.StatsHandle() != nil && do.StatsUpdating() {
 		s.statsCollector = do.StatsHandle().NewSessionStatsCollector()
 	}
 

statistics/handle_test.go

@@ -445,5 +445,6 @@ func newStoreWithBootstrap(statsLease time.Duration) (kv.Storage, *domain.Domain
 	session.SetStatsLease(statsLease)
 	domain.RunAutoAnalyze = false
 	do, err := session.BootstrapSession(store)
+	do.SetStatsUpdating(true)
 	return store, do, errors.Trace(err)
 }