restore: Try to create tables in parallel

lightning/glue/glue.go

@@ -19,12 +19,15 @@ import (
 	"errors"
 
 	"github.com/pingcap/parser"
+	"github.com/pingcap/parser/ast"
 	"github.com/pingcap/parser/model"
 	"github.com/pingcap/parser/mysql"
 	"github.com/pingcap/tidb-lightning/lightning/checkpoints"
 	"github.com/pingcap/tidb-lightning/lightning/common"
 	"github.com/pingcap/tidb-lightning/lightning/config"
 	"github.com/pingcap/tidb-lightning/lightning/log"
+	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tidb/util/sqlexec"
 )
 
 type Glue interface {
@@ -33,7 +36,7 @@ type Glue interface {
 	GetDB() (*sql.DB, error)
 	GetParser() *parser.Parser
 	GetTables(context.Context, string) ([]*model.TableInfo, error)
-	GetSession() (checkpoints.Session, error)
+	GetSession(context.Context) (checkpoints.Session, error)
 	OpenCheckpointsDB(context.Context, *config.Config) (checkpoints.CheckpointsDB, error)
 	// Record is used to report some information (key, value) to host TiDB, including progress, stage currently
 	Record(string, uint64)
@@ -48,6 +51,39 @@ type SQLExecutor interface {
 	Close()
 }
 
+// sqlConnSession implement checkpoints.Session used only for lighting itself
+type sqlConnSession struct {
+	checkpoints.Session
+	conn *sql.Conn
+}
+
+func (session *sqlConnSession) Close() {
+	session.conn.Close()
+}
+
+func (session *sqlConnSession) Execute(ctx context.Context, sql string) ([]sqlexec.RecordSet, error) {
+	_, err := session.conn.ExecContext(ctx, sql)
+	return nil, err
+}
+
+func (session *sqlConnSession) CommitTxn(context.Context) error {
+	return errors.New("sqlConnSession doesn't have a valid CommitTxn implementation")
+}
+
+func (session *sqlConnSession) RollbackTxn(context.Context) {}
+
+func (session *sqlConnSession) PrepareStmt(sql string) (stmtID uint32, paramCount int, fields []*ast.ResultField, err error) {
+	return 0, 0, nil, errors.New("sqlConnSession doesn't have a valid PrepareStmt implementation")
+}
+
+func (session *sqlConnSession) ExecutePreparedStmt(ctx context.Context, stmtID uint32, param []types.Datum) (sqlexec.RecordSet, error) {
+	return nil, errors.New("sqlConnSession doesn't have a valid ExecutePreparedStmt implementation")
+}
+
+func (session *sqlConnSession) DropPreparedStmt(stmtID uint32) error {
+	return errors.New("sqlConnSession doesn't have a valid DropPreparedStmt implementation")
+}
+
 type ExternalTiDBGlue struct {
 	db     *sql.DB
 	parser *parser.Parser
@@ -125,8 +161,12 @@ func (e ExternalTiDBGlue) GetTables(context.Context, string) ([]*model.TableInfo
 	return nil, errors.New("ExternalTiDBGlue doesn't have a valid GetTables function")
 }
 
-func (e ExternalTiDBGlue) GetSession() (checkpoints.Session, error) {
-	return nil, errors.New("ExternalTiDBGlue doesn't have a valid GetSession function")
+func (e ExternalTiDBGlue) GetSession(ctx context.Context) (checkpoints.Session, error) {
+	conn, err := e.db.Conn(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return &sqlConnSession{conn: conn}, nil
 }
 
 func (e *ExternalTiDBGlue) OpenCheckpointsDB(ctx context.Context, cfg *config.Config) (checkpoints.CheckpointsDB, error) {

lightning/restore/restore.go

@@ -30,6 +30,7 @@ import (
 	"github.com/pingcap/failpoint"
 	sstpb "github.com/pingcap/kvproto/pkg/import_sstpb"
 	"github.com/pingcap/parser/model"
+	"github.com/pingcap/tidb-lightning/lightning/checkpoints"
 	"github.com/pingcap/tidb-lightning/lightning/glue"
 	tidbcfg "github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/meta/autoid"
@@ -304,48 +305,238 @@ outside:
 	return errors.Trace(err)
 }
 
-func (rc *RestoreController) restoreSchema(ctx context.Context) error {
-	if !rc.cfg.Mydumper.NoSchema {
-		if rc.tidbGlue.OwnsSQLExecutor() {
-			db, err := DBFromConfig(rc.cfg.TiDB)
-			if err != nil {
-				return errors.Annotate(err, "connect to tidb failed")
-			}
-			defer db.Close()
-			db.ExecContext(ctx, "SET SQL_MODE = ?", rc.cfg.TiDB.StrSQLMode)
-		}
+type schemaStmtType int
+
+const (
+	schemaCreateDatabase = iota
+	schemaCreateTable
+	schemaCreateView
+)
 
-		for _, dbMeta := range rc.dbMetas {
-			task := log.With(zap.String("db", dbMeta.Name)).Begin(zap.InfoLevel, "restore table schema")
+type schemaJob struct {
+	dbName   string
+	tblName  string // empty for create db jobs
+	stmtType schemaStmtType
+	stmts    []*schemaStmt
+}
 
-			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)
+type schemaStmt struct {
+	sql string
+}
 
-			task.End(zap.ErrorLevel, err)
-			if err != nil {
-				return errors.Annotatef(err, "restore table schema %s failed", dbMeta.Name)
+type restoreSchemaWorker struct {
+	ctx   context.Context
+	quit  context.CancelFunc
+	jobCh chan *schemaJob
+	errCh chan error
+	wg    sync.WaitGroup
+	glue  glue.Glue
+	store storage.ExternalStorage
+}
+
+func (worker *restoreSchemaWorker) makeJobs(dbMetas []*mydump.MDDatabaseMeta) error {
+	defer close(worker.jobCh)
+	var err error
+	// 1. restore databases, execute statements concurrency
+	for _, dbMeta := range dbMetas {
+		restoreSchemaJob := &schemaJob{
+			dbName:   dbMeta.Name,
+			stmtType: schemaCreateDatabase,
+			stmts:    make([]*schemaStmt, 0, 1),
+		}
+		restoreSchemaJob.stmts = append(restoreSchemaJob.stmts, &schemaStmt{
+			sql: createDatabaseIfNotExistStmt(dbMeta.Name),
+		})
+		err = worker.appendJob(restoreSchemaJob)
+		if err != nil {
+			return err
+		}
+	}
+	err = worker.wait()
+	if err != nil {
+		return err
+	}
+	// 2. restore tables, execute statements concurrency
+	for _, dbMeta := range dbMetas {
+		for _, tblMeta := range dbMeta.Tables {
+			sql := tblMeta.GetSchema(worker.ctx, worker.store)
+			if sql != "" {
+				stmts, err := createTableIfNotExistsStmt(worker.glue.GetParser(), sql, dbMeta.Name, tblMeta.Name)
+				if err != nil {
+					return err
+				}
+				restoreSchemaJob := &schemaJob{
+					dbName:   dbMeta.Name,
+					tblName:  tblMeta.Name,
+					stmtType: schemaCreateTable,
+					stmts:    make([]*schemaStmt, 0, len(stmts)),
+				}
+				for _, sql := range stmts {
+					restoreSchemaJob.stmts = append(restoreSchemaJob.stmts, &schemaStmt{
+						sql: sql,
+					})
+				}
+				err = worker.appendJob(restoreSchemaJob)
+				if err != nil {
+					return err
+				}
 			}
 		}
-
-		// 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 {
-				task := log.With(zap.String("db", dbMeta.Name)).Begin(zap.InfoLevel, "restore view schema")
-				viewsSchema := make(map[string]string)
-				for _, viewMeta := range dbMeta.Views {
-					viewsSchema[viewMeta.Name] = viewMeta.GetSchema(ctx, rc.store)
+	}
+	err = worker.wait()
+	if err != nil {
+		return err
+	}
+	// 3. restore views. Since views can cross database we must restore views after all table schemas are restored.
+	for _, dbMeta := range dbMetas {
+		for _, viewMeta := range dbMeta.Views {
+			sql := viewMeta.GetSchema(worker.ctx, worker.store)
+			if sql != "" {
+				stmts, err := createTableIfNotExistsStmt(worker.glue.GetParser(), sql, dbMeta.Name, viewMeta.Name)
+				if err != nil {
+					return err
+				}
+				restoreSchemaJob := &schemaJob{
+					dbName:   dbMeta.Name,
+					tblName:  viewMeta.Name,
+					stmtType: schemaCreateView,
+					stmts:    make([]*schemaStmt, 0, len(stmts)),
+				}
+				for _, sql := range stmts {
+					restoreSchemaJob.stmts = append(restoreSchemaJob.stmts, &schemaStmt{
+						sql: sql,
+					})
+				}
+				err = worker.appendJob(restoreSchemaJob)
+				if err != nil {
+					return err
 				}
-				err := InitSchema(ctx, rc.tidbGlue, dbMeta.Name, viewsSchema)
+				// we don't support restore views concurrency, cauz it maybe will raise a error
+				err = worker.wait()
+				if err != nil {
+					return err
+				}
+			}
+		}
+	}
+	return nil
+}
 
+func (worker *restoreSchemaWorker) doJob() {
+	var logger log.Logger
+	var session checkpoints.Session
+	defer func() {
+		if session != nil {
+			session.Close()
+		}
+	}()
+	for {
+		select {
+		case <-worker.ctx.Done():
+			return
+		case job := <-worker.jobCh:
+			if job == nil {
+				return
+			}
+			var err error
+			if session == nil {
+				session, err = worker.glue.GetSession(worker.ctx)
+				if err != nil {
+					worker.throw(err)
+					return
+				}
+			}
+			if job.stmtType == schemaCreateDatabase {
+				logger = log.With(zap.String("db", job.dbName))
+			} else if job.stmtType == schemaCreateTable {
+				logger = log.With(zap.String("table", common.UniqueTable(job.dbName, job.tblName)))
+			} else if job.stmtType == schemaCreateView {
+				logger = log.With(zap.String("table", common.UniqueTable(job.dbName, job.tblName)))
+			}
+			for _, stmt := range job.stmts {
+				task := logger.Begin(zap.DebugLevel, fmt.Sprintf("execute SQL: %s", stmt.sql))
+				_, err = session.Execute(worker.ctx, stmt.sql)
 				task.End(zap.ErrorLevel, err)
 				if err != nil {
-					return errors.Annotatef(err, "restore view schema %s failed", dbMeta.Name)
+					switch job.stmtType {
+					case schemaCreateDatabase:
+						err = errors.Annotatef(err, "restore database schema %s failed", job.dbName)
+					case schemaCreateTable:
+						err = errors.Annotatef(err, "restore table schema %s failed", job.tblName)
+					case schemaCreateView:
+						err = errors.Annotatef(err, "restore view schema %s failed", job.tblName)
+					}
+					worker.wg.Done()
+					worker.throw(err)
+					return
 				}
 			}
+			worker.wg.Done()
+		}
+	}
+}
 
+func (worker *restoreSchemaWorker) wait() error {
+	select {
+	case err := <-worker.errCh:
+		if err != nil {
+			worker.quit()
+		}
+		return err
+	case <-worker.ctx.Done():
+		return worker.ctx.Err()
+	default:
+		worker.wg.Wait()
+		return nil
+	}
+}
+
+func (worker *restoreSchemaWorker) throw(err error) {
+	if err != nil {
+		select {
+		case <-worker.ctx.Done():
+			err := worker.ctx.Err()
+			if err != nil {
+				worker.errCh <- err
+			}
+		case worker.errCh <- err:
+			worker.quit()
+		}
+	}
+}
+
+func (worker *restoreSchemaWorker) appendJob(job *schemaJob) error {
+	select {
+	case err := <-worker.errCh:
+		return err
+	case <-worker.ctx.Done():
+		return worker.ctx.Err()
+	case worker.jobCh <- job:
+		worker.wg.Add(1)
+		return nil
+	}
+}
+
+func (rc *RestoreController) restoreSchema(ctx context.Context) error {
+	if !rc.cfg.Mydumper.NoSchema {
+		logTask := log.L().Begin(zap.InfoLevel, "restore all schema")
+		concurrency := 16
+		childCtx, cancel := context.WithCancel(ctx)
+		worker := restoreSchemaWorker{
+			ctx:   childCtx,
+			quit:  cancel,
+			jobCh: make(chan *schemaJob, concurrency),
+			errCh: make(chan error),
+			glue:  rc.tidbGlue,
+			store: rc.store,
+		}
+		for i := 0; i < concurrency; i++ {
+			go worker.doJob()
+		}
+		err := worker.makeJobs(rc.dbMetas)
+		logTask.End(zap.ErrorLevel, err)
+		if err != nil {
+			return err
 		}
 	}
 	getTableFunc := rc.backend.FetchRemoteTableModels

lightning/restore/tidb.go

@@ -182,6 +182,13 @@ loopCreate:
 	return errors.Trace(err)
 }
 
+func createDatabaseIfNotExistStmt(dbName string) string {
+	var createDatabase strings.Builder
+	createDatabase.WriteString("CREATE DATABASE IF NOT EXISTS ")
+	common.WriteMySQLIdentifier(&createDatabase, dbName)
+	return createDatabase.String()
+}
+
 func createTableIfNotExistsStmt(p *parser.Parser, createTable, dbName, tblName string) ([]string, error) {
 	stmts, _, err := p.Parse(createTable, "", "")
 	if err != nil {