br: concurrently repairing indexes (pingcap#59159)

close pingcap#59158

Author: Leavrth

br/pkg/checkpoint/log_restore.go

@@ -304,6 +304,9 @@ type CheckpointIngestIndexRepairSQL struct {
 	IndexName  string    `json:"index-name"`
 	AddSQL     string    `json:"add-sql"`
 	AddArgs    []any     `json:"add-args"`
+
+	OldIndexIDFound bool `json:"-"`
+	IndexRepaired   bool `json:"-"`
 }
 
 type CheckpointIngestIndexRepairSQLs struct {

br/pkg/glue/progressing.go

@@ -7,12 +7,15 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"sync/atomic"
 	"time"
 
 	"github.com/fatih/color"
+	"github.com/pingcap/log"
 	"github.com/pingcap/tidb/br/pkg/utils"
 	"github.com/vbauerster/mpb/v7"
 	"github.com/vbauerster/mpb/v7/decor"
+	"go.uber.org/zap"
 	"golang.org/x/term"
 )
 
@@ -188,3 +191,80 @@ func buildOneTaskBar(pb *mpb.Progress, title string, total int) *mpb.Bar {
 			color.RedString("ABORTED"))),
 	)
 }
+
+type ProgressBar interface {
+	Increment()
+	Done()
+}
+
+type MultiProgress interface {
+	AddTextBar(string, int64) ProgressBar
+	Wait()
+}
+
+func (ops ConsoleOperations) StartMultiProgress() MultiProgress {
+	if !ops.OutputIsTTY() {
+		return &NopMultiProgress{}
+	}
+	pb := mpb.New(mpb.WithOutput(ops.Out()), mpb.WithRefreshRate(400*time.Millisecond))
+	return &TerminalMultiProgress{
+		progress: pb,
+	}
+}
+
+type NopMultiProgress struct{}
+
+type LogBar struct {
+	name  string
+	total int64
+}
+
+func (nmp *NopMultiProgress) AddTextBar(name string, total int64) ProgressBar {
+	log.Info("progress start", zap.String("name", name))
+	return &LogBar{
+		name:  name,
+		total: total,
+	}
+}
+
+func (nmp *NopMultiProgress) Wait() {}
+
+func (lb *LogBar) Increment() {
+	if atomic.AddInt64(&lb.total, -1) <= 0 {
+		log.Info("progress done", zap.String("name", lb.name))
+	}
+}
+
+func (lb *LogBar) Done() {}
+
+type TerminalBar struct {
+	bar *mpb.Bar
+}
+
+func (tb *TerminalBar) Increment() {
+	tb.bar.Increment()
+}
+
+func (tb *TerminalBar) Done() {
+	tb.bar.Abort(false)
+	tb.bar.Wait()
+}
+
+type TerminalMultiProgress struct {
+	progress *mpb.Progress
+}
+
+func (tmp *TerminalMultiProgress) AddTextBar(name string, total int64) ProgressBar {
+	bar := tmp.progress.New(total,
+		mpb.NopStyle(),
+		mpb.PrependDecorators(decor.Name(name)),
+		mpb.AppendDecorators(decor.OnAbort(decor.OnComplete(decor.Spinner(spinnerText), spinnerDoneText),
+			color.RedString("ABORTED"),
+		)),
+	)
+	return &TerminalBar{bar: bar}
+}
+
+func (tmp *TerminalMultiProgress) Wait() {
+	tmp.progress.Wait()
+}

br/pkg/restore/log_client/BUILD.bazel

@@ -95,11 +95,13 @@ go_test(
     flaky = True,
     shard_count = 50,
     deps = [
+        "//br/pkg/checkpoint",
         "//br/pkg/errors",
         "//br/pkg/glue",
         "//br/pkg/gluetidb",
         "//br/pkg/mock",
         "//br/pkg/restore",
+        "//br/pkg/restore/ingestrec",
         "//br/pkg/restore/internal/import_client",
         "//br/pkg/restore/internal/rawkv",
         "//br/pkg/restore/split",
@@ -111,6 +113,8 @@ go_test(
         "//br/pkg/utiltest",
         "//pkg/domain",
         "//pkg/kv",
+        "//pkg/meta/model",
+        "//pkg/parser/ast",
         "//pkg/planner/core/resolve",
         "//pkg/session",
         "//pkg/sessionctx",

br/pkg/restore/log_client/client.go

@@ -84,6 +84,10 @@ const maxSplitKeysOnce = 10240
 // rawKVBatchCount specifies the count of entries that the rawkv client puts into TiKV.
 const rawKVBatchCount = 64
 
+// session count for repairing ingest indexes. Currently only one TiDB node executes adding index jobs
+// at the same time and the add-index job concurrency is about min(10, `TiDB CPUs / 4`).
+const defaultRepairIndexSessionCount uint = 10
+
 // LogRestoreManager is a comprehensive wrapper that encapsulates all logic related to log restoration,
 // including concurrency management, checkpoint handling, and file importing for efficient log processing.
 type LogRestoreManager struct {
@@ -456,16 +460,48 @@ func (rc *LogClient) CleanUpKVFiles(
 	return rc.logRestoreManager.fileImporter.ClearFiles(ctx, rc.pdClient, "v1")
 }
 
-// Init create db connection and domain for storage.
-func (rc *LogClient) Init(ctx context.Context, g glue.Glue, store kv.Storage) error {
-	var err error
-	rc.unsafeSession, err = g.CreateSession(store)
+func createSession(ctx context.Context, g glue.Glue, store kv.Storage) (glue.Session, error) {
+	unsafeSession, err := g.CreateSession(store)
 	if err != nil {
-		return errors.Trace(err)
+		return nil, errors.Trace(err)
 	}
-
 	// Set SQL mode to None for avoiding SQL compatibility problem
-	err = rc.unsafeSession.Execute(ctx, "set @@sql_mode=''")
+	err = unsafeSession.Execute(ctx, "set @@sql_mode=''")
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	return unsafeSession, nil
+}
+
+func createSessions(ctx context.Context, g glue.Glue, store kv.Storage, count uint) (createdUnsafeSessions []glue.Session, createErr error) {
+	unsafeSessions := make([]glue.Session, 0, count)
+	defer func() {
+		if createErr != nil {
+			closeSessions(unsafeSessions)
+		}
+	}()
+	for range count {
+		unsafeSession, err := createSession(ctx, g, store)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+		unsafeSessions = append(unsafeSessions, unsafeSession)
+	}
+	return unsafeSessions, nil
+}
+
+func closeSessions(sessions []glue.Session) {
+	for _, session := range sessions {
+		if session != nil {
+			session.Close()
+		}
+	}
+}
+
+// Init create db connection and domain for storage.
+func (rc *LogClient) Init(ctx context.Context, g glue.Glue, store kv.Storage) error {
+	var err error
+	rc.unsafeSession, err = createSession(ctx, g, store)
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -1765,39 +1801,60 @@ func (rc *LogClient) RepairIngestIndex(ctx context.Context, ingestRecorder *inge
 
 	info := rc.dom.InfoSchema()
 	console := glue.GetConsole(g)
-NEXTSQL:
-	for _, sql := range sqls {
-		progressTitle := fmt.Sprintf("repair ingest index %s for table %s.%s", sql.IndexName, sql.SchemaName, sql.TableName)
-
+	for i, sql := range sqls {
 		tableInfo, err := info.TableByName(ctx, sql.SchemaName, sql.TableName)
 		if err != nil {
 			return errors.Trace(err)
 		}
-		oldIndexIDFound := false
+		sqls[i].OldIndexIDFound = false
+		sqls[i].IndexRepaired = false
 		if fromCheckpoint {
 			for _, idx := range tableInfo.Indices() {
 				indexInfo := idx.Meta()
 				if indexInfo.ID == sql.IndexID {
 					// the original index id is not dropped
-					oldIndexIDFound = true
+					sqls[i].OldIndexIDFound = true
 					break
 				}
 				// what if index's state is not public?
 				if indexInfo.Name.O == sql.IndexName {
+					progressTitle := fmt.Sprintf("repair ingest index %s for table %s.%s", sql.IndexName, sql.SchemaName, sql.TableName)
 					// find the same name index, but not the same index id,
 					// which means the repaired index id is created
 					if _, err := fmt.Fprintf(console.Out(), "%s ... %s\n", progressTitle, color.HiGreenString("SKIPPED DUE TO CHECKPOINT MODE")); err != nil {
 						return errors.Trace(err)
 					}
-					continue NEXTSQL
+					sqls[i].IndexRepaired = true
+					break
 				}
 			}
 		}
+	}
 
-		if err := func(sql checkpoint.CheckpointIngestIndexRepairSQL) error {
-			w := console.StartProgressBar(progressTitle, glue.OnlyOneTask)
-			defer w.Close()
+	sessionCount := defaultRepairIndexSessionCount
+	indexSessions, err := createSessions(ctx, g, rc.dom.Store(), sessionCount)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	defer func() {
+		closeSessions(indexSessions)
+	}()
+	workerpool := tidbutil.NewWorkerPool(sessionCount, "repair ingest index")
+	eg, ectx := errgroup.WithContext(ctx)
+	mp := console.StartMultiProgress()
+	for _, sql := range sqls {
+		if sql.IndexRepaired {
+			continue
+		}
+		if ectx.Err() != nil {
+			break
+		}
+		progressTitle := fmt.Sprintf("repair ingest index %s for table %s.%s", sql.IndexName, sql.SchemaName, sql.TableName)
+		w := mp.AddTextBar(progressTitle, 1)
+		workerpool.ApplyWithIDInErrorGroup(eg, func(id uint64) error {
+			defer w.Done()
 
+			indexSession := indexSessions[id%uint64(len(indexSessions))]
 			// TODO: When the TiDB supports the DROP and CREATE the same name index in one SQL,
 			//   the checkpoint for ingest recorder can be removed and directly use the SQL:
 			//      ALTER TABLE db.tbl DROP INDEX `i_1`, ADD IDNEX `i_1` ...
@@ -1808,8 +1865,8 @@ NEXTSQL:
 			// restored metakv and then skips repairing it.
 
 			// only when first execution or old index id is not dropped
-			if !fromCheckpoint || oldIndexIDFound {
-				if err := rc.unsafeSession.ExecuteInternal(ctx, alterTableDropIndexSQL, sql.SchemaName.O, sql.TableName.O, sql.IndexName); err != nil {
+			if !fromCheckpoint || sql.OldIndexIDFound {
+				if err := indexSession.ExecuteInternal(ectx, alterTableDropIndexSQL, sql.SchemaName.O, sql.TableName.O, sql.IndexName); err != nil {
 					return errors.Trace(err)
 				}
 			}
@@ -1819,17 +1876,15 @@ NEXTSQL:
 				}
 			})
 			// create the repaired index when first execution or not found it
-			if err := rc.unsafeSession.ExecuteInternal(ctx, sql.AddSQL, sql.AddArgs...); err != nil {
-				return errors.Trace(err)
-			}
-			w.Inc()
-			if err := w.Wait(ctx); err != nil {
+			if err := indexSession.ExecuteInternal(ectx, sql.AddSQL, sql.AddArgs...); err != nil {
 				return errors.Trace(err)
 			}
+			w.Increment()
 			return nil
-		}(sql); err != nil {
-			return errors.Trace(err)
-		}
+		})
+	}
+	if err := eg.Wait(); err != nil {
+		return errors.Trace(err)
 	}
 
 	return nil