restore: Try to create tables in parallel

lightning/restore/restore.go

@@ -298,30 +298,37 @@ outside:
 
 func (rc *RestoreController) restoreSchema(ctx context.Context) error {
 	if !rc.cfg.Mydumper.NoSchema {
+		// CONFUSED: how we control concurrency via configuration?
+		concurrency := 16
 		if rc.tidbGlue.OwnsSQLExecutor() {
+			// CONFUSED: why don't use rc.tidbGlue.GetDB()?
 			db, err := DBFromConfig(rc.cfg.TiDB)
 			if err != nil {
 				return errors.Trace(err)
 			}
 			defer db.Close()
 			db.ExecContext(ctx, "SET SQL_MODE = ?", rc.cfg.TiDB.StrSQLMode)
 		}
-
+		db, err := rc.tidbGlue.GetDB()
+		if err != nil {
+			return errors.Trace(err)
+		}
+		defer db.SetMaxOpenConns(1)
+		db.SetMaxOpenConns(concurrency)
 		for _, dbMeta := range rc.dbMetas {
 			task := log.With(zap.String("db", dbMeta.Name)).Begin(zap.InfoLevel, "restore table schema")
 
 			tablesSchema := make(map[string]string)
 			for _, tblMeta := range dbMeta.Tables {
 				tablesSchema[tblMeta.Name] = tblMeta.GetSchema(ctx, rc.store)
 			}
-			err := InitSchema(ctx, rc.tidbGlue, dbMeta.Name, tablesSchema)
+			err = InitSchema(ctx, concurrency, rc.tidbGlue.GetParser(), rc.tidbGlue.GetSQLExecutor(), dbMeta.Name, tablesSchema)
 
 			task.End(zap.ErrorLevel, err)
 			if err != nil {
 				return errors.Annotatef(err, "restore table schema %s failed", dbMeta.Name)
 			}
 		}
-
 		// restore views. Since views can cross database we must restore views after all table schemas are restored.
 		for _, dbMeta := range rc.dbMetas {
 			if len(dbMeta.Views) > 0 {
@@ -330,7 +337,7 @@ func (rc *RestoreController) restoreSchema(ctx context.Context) error {
 				for _, viewMeta := range dbMeta.Views {
 					viewsSchema[viewMeta.Name] = viewMeta.GetSchema(ctx, rc.store)
 				}
-				err := InitSchema(ctx, rc.tidbGlue, dbMeta.Name, viewsSchema)
+				err := InitSchema(ctx, concurrency, rc.tidbGlue.GetParser(), rc.tidbGlue.GetSQLExecutor(), dbMeta.Name, viewsSchema)
 
 				task.End(zap.ErrorLevel, err)
 				if err != nil {

lightning/restore/tidb.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"sync"
 
 	tmysql "github.com/go-sql-driver/mysql"
 	"github.com/pingcap/errors"
@@ -38,6 +39,7 @@ import (
 	"github.com/pingcap/tidb-lightning/lightning/mydump"
 
 	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
 )
 
 type TiDBManager struct {
@@ -124,43 +126,74 @@ func (timgr *TiDBManager) Close() {
 	timgr.db.Close()
 }
 
-func InitSchema(ctx context.Context, g glue.Glue, database string, tablesSchema map[string]string) error {
-	logger := log.With(zap.String("db", database))
-	sqlExecutor := g.GetSQLExecutor()
+type schemaJob struct {
+	sql string
+	log zapcore.Field
+}
 
+func InitSchema(ctx context.Context, concurrency int, parser *parser.Parser, exec glue.SQLExecutor, database string, tablesSchema map[string]string) error {
+	logger := log.With(zap.String("db", database))
 	var createDatabase strings.Builder
 	createDatabase.WriteString("CREATE DATABASE IF NOT EXISTS ")
 	common.WriteMySQLIdentifier(&createDatabase, database)
-	err := sqlExecutor.ExecuteWithLog(ctx, createDatabase.String(), "create database", logger)
+	err := exec.ExecuteWithLog(ctx, createDatabase.String(), "create database", logger)
 	if err != nil {
 		return errors.Trace(err)
 	}
-
-	task := logger.Begin(zap.InfoLevel, "create tables")
+	var wg sync.WaitGroup
+	fnCreateTable := func(ctx context.Context, exec glue.SQLExecutor, jobCh chan *schemaJob, logger log.Logger, errCh chan error) {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case job := <-jobCh:
+				//TODO: maybe we should put these createStems into a transaction
+				err = exec.ExecuteWithLog(ctx, job.sql, "create table", logger.With(job.log))
+				wg.Done()
+				if err != nil {
+					errCh <- err
+				}
+			}
+		}
+	}
+	fnHandleError := func(ctx context.Context, err error, errCh chan error, quit context.CancelFunc) {
+		for {
+			select {
+			case err = <-errCh:
+				quit()
+				return
+			case <-ctx.Done():
+				return
+			}
+		}
+	}
 	var sqlCreateStmts []string
+	errCh := make(chan error)
+	queue := make(chan *schemaJob, concurrency)
+	task := logger.Begin(zap.InfoLevel, "create tables")
+	childCtx, cancel := context.WithCancel(ctx)
+	for i := 0; i < concurrency; i++ {
+		go fnCreateTable(childCtx, exec, queue, logger, errCh)
+	}
+	go fnHandleError(childCtx, err, errCh, cancel)
 	for tbl, sqlCreateTable := range tablesSchema {
 		task.Debug("create table", zap.String("schema", sqlCreateTable))
-
-		sqlCreateStmts, err = createTableIfNotExistsStmt(g.GetParser(), sqlCreateTable, database, tbl)
+		sqlCreateStmts, err = createTableIfNotExistsStmt(parser, sqlCreateTable, database, tbl)
 		if err != nil {
+			errCh <- err
 			break
 		}
-
-		//TODO: maybe we should put these createStems into a transaction
-		for _, s := range sqlCreateStmts {
-			err = sqlExecutor.ExecuteWithLog(
-				ctx,
-				s,
-				"create table",
-				logger.With(zap.String("table", common.UniqueTable(database, tbl))),
-			)
-			if err != nil {
-				break
+		log := zap.String("table", common.UniqueTable(database, tbl))
+		for _, stmt := range sqlCreateStmts {
+			wg.Add(1)
+			queue <- &schemaJob{
+				sql: stmt,
+				log: log,
 			}
 		}
 	}
+	wg.Wait()
 	task.End(zap.ErrorLevel, err)
-
 	return errors.Trace(err)
 }
 

lightning/restore/tidb_test.go

@@ -201,7 +201,7 @@ func (s *tidbSuite) TestInitSchema(c *C) {
 		ExpectClose()
 
 	s.mockDB.MatchExpectationsInOrder(false) // maps are unordered.
-	err := InitSchema(ctx, s.tiGlue, "db", map[string]string{
+	err := InitSchema(ctx, 1, s.tiGlue.GetParser(), s.tiGlue.GetSQLExecutor(), "db", map[string]string{
 		"t1": "create table t1 (a int primary key, b varchar(200));",
 		"t2": "/*!40014 SET FOREIGN_KEY_CHECKS=0*/;CREATE TABLE `db`.`t2` (xx TEXT) AUTO_INCREMENT=11203;",
 	})
@@ -218,7 +218,7 @@ func (s *tidbSuite) TestInitSchemaSyntaxError(c *C) {
 	s.mockDB.
 		ExpectClose()
 
-	err := InitSchema(ctx, s.tiGlue, "db", map[string]string{
+	err := InitSchema(ctx, 1, s.tiGlue.GetParser(), s.tiGlue.GetSQLExecutor(), "db", map[string]string{
 		"t1": "create table `t1` with invalid syntax;",
 	})
 	c.Assert(err, NotNil)
@@ -239,7 +239,7 @@ func (s *tidbSuite) TestInitSchemaUnsupportedSchemaError(c *C) {
 	s.mockDB.
 		ExpectClose()
 
-	err := InitSchema(ctx, s.tiGlue, "db", map[string]string{
+	err := InitSchema(ctx, 1, s.tiGlue.GetParser(), s.tiGlue.GetSQLExecutor(), "db", map[string]string{
 		"t1": "create table `t1` (a VARCHAR(999999999));",
 	})
 	c.Assert(err, ErrorMatches, ".*Column length too big.*")