*: fix panic when to enable lite-init-stats and concurrently-init-stats

pkg/statistics/handle/BUILD.bazel

@@ -35,6 +35,7 @@ go_library(
         "//pkg/types",
         "//pkg/util/chunk",
         "//pkg/util/logutil",
+        "//pkg/util/sqlexec",
         "@com_github_pingcap_errors//:errors",
         "@org_uber_go_zap//:zap",
     ],

pkg/statistics/handle/bootstrap.go

@@ -34,6 +34,7 @@ import (
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util/chunk"
 	"github.com/pingcap/tidb/pkg/util/logutil"
+	"github.com/pingcap/tidb/pkg/util/sqlexec"
 	"go.uber.org/zap"
 )
 
@@ -78,12 +79,6 @@ func (h *Handle) initStatsMeta(is infoschema.InfoSchema) (statstypes.StatsCache,
 	if err != nil {
 		return nil, err
 	}
-	if config.GetGlobalConfig().Performance.ConcurrentlyInitStats {
-		ls := initstats.NewWorker(rc.Next, h.initStatsMeta4Chunk)
-		ls.LoadStats(is, tables, rc)
-		ls.Wait()
-		return tables, nil
-	}
 	req := rc.NewChunk(nil)
 	iter := chunk.NewIterator4Chunk(req)
 	for {
@@ -242,12 +237,6 @@ func (h *Handle) initStatsHistogramsLite(is infoschema.InfoSchema, cache statsty
 		return errors.Trace(err)
 	}
 	defer terror.Call(rc.Close)
-	if config.GetGlobalConfig().Performance.ConcurrentlyInitStats {
-		ls := initstats.NewWorker(rc.Next, h.initStatsHistograms4ChunkLite)
-		ls.LoadStats(is, cache, rc)
-		ls.Wait()
-		return nil
-	}
 	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats)
 	req := rc.NewChunk(nil)
 	iter := chunk.NewIterator4Chunk(req)
@@ -271,12 +260,29 @@ func (h *Handle) initStatsHistograms(is infoschema.InfoSchema, cache statstypes.
 		return errors.Trace(err)
 	}
 	defer terror.Call(rc.Close)
-	if config.GetGlobalConfig().Performance.ConcurrentlyInitStats {
-		ls := initstats.NewWorker(rc.Next, h.initStatsHistograms4Chunk)
-		ls.LoadStats(is, cache, rc)
-		ls.Wait()
-		return nil
+	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats)
+	req := rc.NewChunk(nil)
+	iter := chunk.NewIterator4Chunk(req)
+	for {
+		err := rc.Next(ctx, req)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		if req.NumRows() == 0 {
+			break
+		}
+		h.initStatsHistograms4Chunk(is, cache, iter)
 	}
+	return nil
+}
+
+func (h *Handle) initStatsHistogramsByPaging(is infoschema.InfoSchema, cache statstypes.StatsCache, task initstats.Task) error {
+	sql := "select HIGH_PRIORITY table_id, is_index, hist_id, distinct_count, version, null_count, cm_sketch, tot_col_size, stats_ver, correlation, flag, last_analyze_pos from mysql.stats_histograms where table_id >= %? and table_id < %?"
+	rc, err := util.Exec(h.initStatsCtx, sql, task.StartTid, task.EndTid)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	defer terror.Call(rc.Close)
 	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats)
 	req := rc.NewChunk(nil)
 	iter := chunk.NewIterator4Chunk(req)
@@ -293,6 +299,46 @@ func (h *Handle) initStatsHistograms(is infoschema.InfoSchema, cache statstypes.
 	return nil
 }
 
+func (h *Handle) initStatsHistogramsConcurrency(is infoschema.InfoSchema, cache statstypes.StatsCache) error {
+	sql := "select HIGH_PRIORITY min(table_id), max(table_id)  from mysql.stats_histograms"
+	rs, err := util.Exec(h.initStatsCtx, sql)
+	if err != nil {
+		return err
+	}
+	var rows []chunk.Row
+	rows, err = sqlexec.DrainRecordSet(context.Background(), rs, 64)
+	if err != nil {
+		return err
+	}
+	err = rs.Close()
+	if err != nil {
+		return err
+	}
+	var maxTid, minTid int64
+	if len(rows) > 0 {
+		minTid = rows[0].GetInt64(0)
+		maxTid = rows[0].GetInt64(1)
+	}
+	ls := initstats.NewWorker(func(task initstats.Task) error {
+		return h.initStatsHistogramsByPaging(is, cache, task)
+	})
+	ls.LoadStats()
+	tid := minTid
+	for tid <= maxTid {
+		ls.SendTask(initstats.Task{
+			StartTid: tid,
+			EndTid:   tid + 1000,
+		})
+		tid += 1000
+	}
+	ls.SendTask(initstats.Task{
+		StartTid: tid,
+		EndTid:   tid + 1000,
+	})
+	ls.Wait()
+	return nil
+}
+
 func (*Handle) initStatsTopN4Chunk(cache statstypes.StatsCache, iter *chunk.Iterator4Chunk) {
 	affectedIndexes := make(map[*statistics.Index]struct{})
 	for row := iter.Begin(); row != iter.End(); row = iter.Next() {
@@ -326,14 +372,6 @@ func (h *Handle) initStatsTopN(cache statstypes.StatsCache) error {
 		return errors.Trace(err)
 	}
 	defer terror.Call(rc.Close)
-	if config.GetGlobalConfig().Performance.ConcurrentlyInitStats {
-		ls := initstats.NewWorker(rc.Next, func(_ infoschema.InfoSchema, cache statstypes.StatsCache, iter *chunk.Iterator4Chunk) {
-			h.initStatsTopN4Chunk(cache, iter)
-		})
-		ls.LoadStats(nil, cache, rc)
-		ls.Wait()
-		return nil
-	}
 	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats)
 	req := rc.NewChunk(nil)
 	iter := chunk.NewIterator4Chunk(req)
@@ -458,26 +496,18 @@ func (h *Handle) initStatsBuckets(cache statstypes.StatsCache) error {
 		return errors.Trace(err)
 	}
 	defer terror.Call(rc.Close)
-	if config.GetGlobalConfig().Performance.ConcurrentlyInitStats {
-		ls := initstats.NewWorker(rc.Next, func(_ infoschema.InfoSchema, cache statstypes.StatsCache, iter *chunk.Iterator4Chunk) {
-			h.initStatsBuckets4Chunk(cache, iter)
-		})
-		ls.LoadStats(nil, cache, rc)
-		ls.Wait()
-	} else {
-		ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats)
-		req := rc.NewChunk(nil)
-		iter := chunk.NewIterator4Chunk(req)
-		for {
-			err := rc.Next(ctx, req)
-			if err != nil {
-				return errors.Trace(err)
-			}
-			if req.NumRows() == 0 {
-				break
-			}
-			h.initStatsBuckets4Chunk(cache, iter)
+	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats)
+	req := rc.NewChunk(nil)
+	iter := chunk.NewIterator4Chunk(req)
+	for {
+		err := rc.Next(ctx, req)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		if req.NumRows() == 0 {
+			break
 		}
+		h.initStatsBuckets4Chunk(cache, iter)
 	}
 	tables := cache.Values()
 	for _, table := range tables {
@@ -544,9 +574,16 @@ func (h *Handle) InitStats(is infoschema.InfoSchema) (err error) {
 	if err != nil {
 		return errors.Trace(err)
 	}
-	err = h.initStatsHistograms(is, cache)
-	if err != nil {
-		return errors.Trace(err)
+	if config.GetGlobalConfig().Performance.ConcurrentlyInitStats {
+		err = h.initStatsHistogramsConcurrency(is, cache)
+		if err != nil {
+			return errors.Trace(err)
+		}
+	} else {
+		err = h.initStatsHistograms(is, cache)
+		if err != nil {
+			return errors.Trace(err)
+		}
 	}
 	err = h.initStatsTopN(cache)
 	if err != nil {

pkg/statistics/handle/initstats/BUILD.bazel

@@ -6,13 +6,8 @@ go_library(
     importpath = "github.com/pingcap/tidb/pkg/statistics/handle/initstats",
     visibility = ["//visibility:public"],
     deps = [
-        "//pkg/infoschema",
-        "//pkg/kv",
-        "//pkg/statistics/handle/logutil",
-        "//pkg/statistics/handle/types",
         "//pkg/util",
-        "//pkg/util/chunk",
-        "//pkg/util/sqlexec",
+        "//pkg/util/logutil",
         "@org_uber_go_zap//:zap",
     ],
 )

pkg/statistics/handle/initstats/load_stats.go

@@ -15,72 +15,60 @@
 package initstats
 
 import (
-	"context"
 	"runtime"
-	"sync"
 
-	"github.com/pingcap/tidb/pkg/infoschema"
-	"github.com/pingcap/tidb/pkg/kv"
-	"github.com/pingcap/tidb/pkg/statistics/handle/logutil"
-	statstypes "github.com/pingcap/tidb/pkg/statistics/handle/types"
 	"github.com/pingcap/tidb/pkg/util"
-	"github.com/pingcap/tidb/pkg/util/chunk"
-	"github.com/pingcap/tidb/pkg/util/sqlexec"
+	"github.com/pingcap/tidb/pkg/util/logutil"
 	"go.uber.org/zap"
 )
 
+// Task represents the range of the table for loading stats.
+type Task struct {
+	StartTid int64
+	EndTid   int64
+}
+
 // Worker is used to load stats concurrently.
 type Worker struct {
-	taskFunc func(ctx context.Context, req *chunk.Chunk) error
-	dealFunc func(is infoschema.InfoSchema, cache statstypes.StatsCache, iter *chunk.Iterator4Chunk)
-	mu       sync.Mutex
-	wg       util.WaitGroupWrapper
+	dealFunc func(task Task) error
+	taskChan chan Task
+
+	wg util.WaitGroupWrapper
 }
 
 // NewWorker creates a new Worker.
-func NewWorker(
-	taskFunc func(ctx context.Context, req *chunk.Chunk) error,
-	dealFunc func(is infoschema.InfoSchema, cache statstypes.StatsCache, iter *chunk.Iterator4Chunk)) *Worker {
+func NewWorker(dealFunc func(task Task) error) *Worker {
 	return &Worker{
-		taskFunc: taskFunc,
 		dealFunc: dealFunc,
+		taskChan: make(chan Task, 1),
 	}
 }
 
 // LoadStats loads stats concurrently when to init stats
-func (ls *Worker) LoadStats(is infoschema.InfoSchema, cache statstypes.StatsCache, rc sqlexec.RecordSet) {
+func (ls *Worker) LoadStats() {
 	concurrency := runtime.GOMAXPROCS(0)
 	for n := 0; n < concurrency; n++ {
 		ls.wg.Run(func() {
-			req := rc.NewChunk(nil)
-			ls.loadStats(is, cache, req)
+			ls.loadStats()
 		})
 	}
 }
 
-func (ls *Worker) loadStats(is infoschema.InfoSchema, cache statstypes.StatsCache, req *chunk.Chunk) {
-	iter := chunk.NewIterator4Chunk(req)
-	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats)
-	for {
-		err := ls.getTask(ctx, req)
-		if err != nil {
-			logutil.StatsLogger().Error("load stats failed", zap.Error(err))
-			return
-		}
-		if req.NumRows() == 0 {
-			return
+func (ls *Worker) loadStats() {
+	for task := range ls.taskChan {
+		if err := ls.dealFunc(task); err != nil {
+			logutil.BgLogger().Error("load stats failed", zap.Error(err))
 		}
-		ls.dealFunc(is, cache, iter)
 	}
 }
 
-func (ls *Worker) getTask(ctx context.Context, req *chunk.Chunk) error {
-	ls.mu.Lock()
-	defer ls.mu.Unlock()
-	return ls.taskFunc(ctx, req)
+// SendTask sends a task to the load stats worker.
+func (ls *Worker) SendTask(task Task) {
+	ls.taskChan <- task
 }
 
 // Wait closes the load stats worker.
 func (ls *Worker) Wait() {
+	close(ls.taskChan)
 	ls.wg.Wait()
 }