br: support batch create table for restore

br/pkg/glue/glue.go

@@ -39,6 +39,11 @@ type Session interface {
 	Close()
 }
 
+// BatchCreateTableSession is an interface to batch create table parallelly
+type BatchCreateTableSession interface {
+	CreateTables(ctx context.Context, tables map[string][]*model.TableInfo) error
+}
+
 // Progress is an interface recording the current execution progress.
 type Progress interface {
 	// Inc increases the progress. This method must be goroutine-safe, and can

br/pkg/gluetidb/glue.go

@@ -5,6 +5,7 @@ package gluetidb
 import (
 	"bytes"
 	"context"
+	"strings"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/log"
@@ -21,6 +22,7 @@ import (
 	"github.com/pingcap/tidb/session"
 	"github.com/pingcap/tidb/sessionctx"
 	pd "github.com/tikv/pd/client"
+	"go.uber.org/zap"
 )
 
 const (
@@ -126,6 +128,48 @@ func (gs *tidbSession) CreateDatabase(ctx context.Context, schema *model.DBInfo)
 		schema.Charset = mysql.DefaultCharset
 	}
 	return d.CreateSchemaWithInfo(gs.se, schema, ddl.OnExistIgnore)
+
+}
+
+// CreateTables implements glue.BatchCreateTableSession.
+func (gs *tidbSession) CreateTables(ctx context.Context, tables map[string][]*model.TableInfo) error {
+	d := domain.GetDomain(gs.se).DDL()
+	var dbName model.CIStr
+
+	for db, tablesInDB := range tables {
+		dbName = model.NewCIStr(db)
+		queryBuilder := strings.Builder{}
+		cloneTables := make([]*model.TableInfo, 0, len(tablesInDB))
+		for _, table := range tablesInDB {
+			query, err := gs.showCreateTable(table)
+			if err != nil {
+				return errors.Trace(err)
+			}
+
+			queryBuilder.WriteString(query)
+			queryBuilder.WriteString(";")
+
+			table = table.Clone()
+			// Clone() does not clone partitions yet :(
+			if table.Partition != nil {
+				newPartition := *table.Partition
+				newPartition.Definitions = append([]model.PartitionDefinition{}, table.Partition.Definitions...)
+				table.Partition = &newPartition
+			}
+			cloneTables = append(cloneTables, table)
+		}
+		gs.se.SetValue(sessionctx.QueryString, queryBuilder.String())
+		err := d.BatchCreateTableWithInfo(gs.se, dbName, cloneTables, ddl.OnExistIgnore)
+		if err != nil {
+			//It is possible to failure when TiDB does not support model.ActionCreateTables.
+			//In this circumstance, BatchCreateTableWithInfo returns errno.ErrInvalidDDLJob,
+			//we fall back to old way that creating table one by one
+			log.Warn("batch create table from tidb failure", zap.Error(err))
+			return err
+		}
+	}
+
+	return nil
 }
 
 // CreateTable implements glue.Session.

br/pkg/restore/client.go

@@ -33,8 +33,10 @@ import (
 	"github.com/pingcap/tidb/br/pkg/summary"
 	"github.com/pingcap/tidb/br/pkg/utils"
 	"github.com/pingcap/tidb/domain"
+	"github.com/pingcap/tidb/errno"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/parser/model"
+	"github.com/pingcap/tidb/parser/terror"
 	"github.com/pingcap/tidb/statistics/handle"
 	"github.com/pingcap/tidb/tablecodec"
 	"github.com/pingcap/tidb/util/codec"
@@ -52,6 +54,7 @@ import (
 // defaultChecksumConcurrency is the default number of the concurrent
 // checksum tasks.
 const defaultChecksumConcurrency = 64
+const minBatchDdlSize = 1
 
 // Client sends requests to restore files.
 type Client struct {
@@ -93,6 +96,8 @@ type Client struct {
 	// and restore stats with #dump.LoadStatsFromJSON
 	statsHandler *handle.Handle
 	dom          *domain.Domain
+
+	batchDdlSize uint
 }
 
 // NewRestoreClient returns a new RestoreClient.
@@ -165,6 +170,14 @@ func (rc *Client) SetSwitchModeInterval(interval time.Duration) {
 	rc.switchModeInterval = interval
 }
 
+func (rc *Client) SetBatchDdlSize(batchDdlsize uint) {
+	rc.batchDdlSize = batchDdlsize
+}
+
+func (rc *Client) GetBatchDdlSize() uint {
+	return rc.batchDdlSize
+}
+
 // Close a client.
 func (rc *Client) Close() {
 	// rc.db can be nil in raw kv mode.
@@ -400,6 +413,46 @@ func (rc *Client) CreateTables(
 	}
 	return rewriteRules, newTables, nil
 }
+func (rc *Client) createTables(
+	ctx context.Context,
+	db *DB,
+	dom *domain.Domain,
+	tables []*metautil.Table,
+	newTS uint64,
+) ([]CreatedTable, error) {
+	log.Info("client to create tables")
+	if rc.IsSkipCreateSQL() {
+		log.Info("skip create table and alter autoIncID")
+	} else {
+		err := db.CreateTables(ctx, tables)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+	}
+	cts := make([]CreatedTable, 0, len(tables))
+	for _, table := range tables {
+		newTableInfo, err := rc.GetTableSchema(dom, table.DB.Name, table.Info.Name)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+		if newTableInfo.IsCommonHandle != table.Info.IsCommonHandle {
+			return nil, errors.Annotatef(berrors.ErrRestoreModeMismatch,
+				"Clustered index option mismatch. Restored cluster's @@tidb_enable_clustered_index should be %v (backup table = %v, created table = %v).",
+				transferBoolToValue(table.Info.IsCommonHandle),
+				table.Info.IsCommonHandle,
+				newTableInfo.IsCommonHandle)
+		}
+		rules := GetRewriteRules(newTableInfo, table.Info, newTS)
+		ct := CreatedTable{
+			RewriteRule: rules,
+			Table:       newTableInfo,
+			OldTable:    table,
+		}
+		log.Debug("new created tables", zap.Any("table", ct))
+		cts = append(cts, ct)
+	}
+	return cts, nil
+}
 
 func (rc *Client) createTable(
 	ctx context.Context,
@@ -459,11 +512,33 @@ func (rc *Client) GoCreateTables(
 	}
 	outCh := make(chan CreatedTable, len(tables))
 	rater := logutil.TraceRateOver(logutil.MetricTableCreatedCounter)
+
+	var err error
+
+	if rc.batchDdlSize > minBatchDdlSize && len(dbPool) > 0 {
+
+		err = rc.createTablesInWorkerPool(ctx, dom, tables, dbPool, newTS, outCh)
+
+		if err == nil {
+			defer log.Debug("all tables are created")
+			close(outCh)
+			return outCh
+			// fall back to old create table (sequential create table)
+		} else if errors.Cause(err).(*terror.Error).Code() == errno.ErrInvalidDDLJob {
+			log.Info("fall back to the sequential create table")
+		} else {
+			errCh <- err
+			close(outCh)
+			return outCh
+		}
+	}
+
 	createOneTable := func(c context.Context, db *DB, t *metautil.Table) error {
 		select {
 		case <-c.Done():
 			return c.Err()
 		default:
+
 		}
 		rt, err := rc.createTable(c, db, dom, t, newTS, ddlTables)
 		if err != nil {
@@ -497,6 +572,7 @@ func (rc *Client) GoCreateTables(
 			errCh <- err
 		}
 	}()
+
 	return outCh
 }
 
@@ -526,6 +602,41 @@ func (rc *Client) createTablesWithDBPool(ctx context.Context,
 	return eg.Wait()
 }
 
+func (rc *Client) createTablesInWorkerPool(ctx context.Context, dom *domain.Domain, tables []*metautil.Table, dbPool []*DB, newTS uint64, outCh chan<- CreatedTable) error {
+	eg, ectx := errgroup.WithContext(ctx)
+	rater := logutil.TraceRateOver(logutil.MetricTableCreatedCounter)
+	workers := utils.NewWorkerPool(uint(len(dbPool)), "Create Tables Worker")
+	numOfTables := len(tables)
+
+	for lastSent := 0; lastSent < numOfTables; lastSent += int(rc.batchDdlSize) {
+		end := utils.MinInt(lastSent+int(rc.batchDdlSize), len(tables))
+		log.Info("create tables", zap.Int("table start", lastSent), zap.Int("table end", end))
+
+		tableSlice := tables[lastSent:end]
+		workers.ApplyWithIDInErrorGroup(eg, func(id uint64) error {
+			db := dbPool[id%uint64(len(dbPool))]
+			cts, err := rc.createTables(ectx, db, dom, tableSlice, newTS) // ddl job for [lastSent:i)
+			if err != nil {
+				log.Error("create tables fail")
+				return err
+			}
+			for _, ct := range cts {
+				log.Debug("table created and send to next",
+					zap.Int("output chan size", len(outCh)),
+					zap.Stringer("table", ct.OldTable.Info.Name),
+					zap.Stringer("database", ct.OldTable.DB.Name))
+				outCh <- ct
+				rater.Inc()
+				rater.L().Info("table created",
+					zap.Stringer("table", ct.OldTable.Info.Name),
+					zap.Stringer("database", ct.OldTable.DB.Name))
+			}
+			return err
+		})
+	}
+	return eg.Wait()
+}
+
 // ExecDDLs executes the queries of the ddl jobs.
 func (rc *Client) ExecDDLs(ctx context.Context, ddlJobs []*model.Job) error {
 	// Sort the ddl jobs by schema version in ascending order.

br/pkg/restore/client_test.go

@@ -40,6 +40,7 @@ func TestCreateTables(t *testing.T) {
 	dbSchema, isExist := info.SchemaByName(model.NewCIStr("test"))
 	require.True(t, isExist)
 
+	client.SetBatchDdlSize(1)
 	tables := make([]*metautil.Table, 4)
 	intField := types.NewFieldType(mysql.TypeLong)
 	intField.Charset = "binary"

br/pkg/restore/db.go

@@ -124,22 +124,11 @@ func (db *DB) CreateDatabase(ctx context.Context, schema *model.DBInfo) error {
 	return errors.Trace(err)
 }
 
-// CreateTable executes a CREATE TABLE SQL.
-func (db *DB) CreateTable(ctx context.Context, table *metautil.Table, ddlTables map[UniqueTableName]bool) error {
-	err := db.se.CreateTable(ctx, table.DB.Name, table.Info)
-	if err != nil {
-		log.Error("create table failed",
-			zap.Stringer("db", table.DB.Name),
-			zap.Stringer("table", table.Info.Name),
-			zap.Error(err))
-		return errors.Trace(err)
-	}
-
+//
+func (db *DB) restoreSequence(ctx context.Context, table *metautil.Table) error {
 	var restoreMetaSQL string
-	switch {
-	case table.Info.IsView():
-		return nil
-	case table.Info.IsSequence():
+	var err error
+	if table.Info.IsSequence() {
 		setValFormat := fmt.Sprintf("do setval(%s.%s, %%d);",
 			utils.EncloseName(table.DB.Name.O),
 			utils.EncloseName(table.Info.Name.O))
@@ -168,7 +157,6 @@ func (db *DB) CreateTable(ctx context.Context, table *metautil.Table, ddlTables
 					zap.Error(err))
 				return errors.Trace(err)
 			}
-
 			// trigger cycle round > 0
 			err = db.se.Execute(ctx, nextSeqSQL)
 			if err != nil {
@@ -182,12 +170,56 @@ func (db *DB) CreateTable(ctx context.Context, table *metautil.Table, ddlTables
 		}
 		restoreMetaSQL = fmt.Sprintf(setValFormat, table.Info.AutoIncID)
 		err = db.se.Execute(ctx, restoreMetaSQL)
+	}
+	if err != nil {
+		log.Error("restore meta sql failed",
+			zap.String("query", restoreMetaSQL),
+			zap.Stringer("db", table.DB.Name),
+			zap.Stringer("table", table.Info.Name),
+			zap.Error(err))
+		return errors.Trace(err)
+	}
+	return errors.Trace(err)
+}
+
+// CreateTables execute a internal CREATE TABLES.
+func (db *DB) CreateTables(ctx context.Context, tables []*metautil.Table) error {
+	if batchSession, ok := db.se.(glue.BatchCreateTableSession); ok {
+		m := map[string][]*model.TableInfo{}
+		for _, table := range tables {
+			m[table.DB.Name.L] = append(m[table.DB.Name.L], table.Info)
+		}
+		if err := batchSession.CreateTables(ctx, m); err != nil {
+			return err
+		}
+
+		for _, table := range tables {
+			err := db.restoreSequence(ctx, table)
+			return err
+		}
+	}
+
+	return nil
+}
+
+// CreateTable executes a CREATE TABLE SQL.
+func (db *DB) CreateTable(ctx context.Context, table *metautil.Table, ddlTables map[UniqueTableName]bool) error {
+	err := db.se.CreateTable(ctx, table.DB.Name, table.Info)
+	if err != nil {
+		log.Error("create table failed",
+			zap.Stringer("db", table.DB.Name),
+			zap.Stringer("table", table.Info.Name),
+			zap.Error(err))
+		return errors.Trace(err)
+	}
+
+	var restoreMetaSQL string
+	switch {
+	case table.Info.IsView():
+		return nil
+	case table.Info.IsSequence():
+		err = db.restoreSequence(ctx, table)
 		if err != nil {
-			log.Error("restore meta sql failed",
-				zap.String("query", restoreMetaSQL),
-				zap.Stringer("db", table.DB.Name),
-				zap.Stringer("table", table.Info.Name),
-				zap.Error(err))
 			return errors.Trace(err)
 		}
 	// only table exists in ddlJobs during incremental restoration should do alter after creation.