restore: Try to create tables in parallel

lightning/glue/glue.go

@@ -17,7 +17,6 @@ import (
 	"context"
 	"database/sql"
 	"errors"
-	"sync"
 
 	"github.com/pingcap/parser"
 	"github.com/pingcap/parser/model"
@@ -48,127 +47,6 @@ type SQLExecutor interface {
 	Close()
 }
 
-type dbConnector struct {
-	dbConf   config.DBStore
-	lazyConn func(dbConf config.DBStore) (*sql.DB, error)
-	conn     *sql.DB
-	err      error
-}
-
-func (connector *dbConnector) connect() error {
-	if connector.conn == nil {
-		connector.conn, connector.err = connector.lazyConn(connector.dbConf)
-	}
-	return connector.err
-}
-
-func (connector *dbConnector) close() error {
-	if connector.conn != nil {
-		return connector.conn.Close()
-	}
-	return nil
-}
-
-func (connector *dbConnector) getError() error {
-	return connector.err
-}
-
-func (connector *dbConnector) getDB() *sql.DB {
-	return connector.conn
-}
-
-// DBPool is multi-connection implementation of `SQLExecutor` interface
-// It will pick one of any idle connections to handle DB events
-// and put it back when events done
-type DBPool struct {
-	SQLExecutor
-	singleton *dbConnector
-	pool      sync.Pool
-}
-
-func (dbPool *DBPool) getConn() *dbConnector {
-	conn := dbPool.pool.Get()
-	if conn != nil {
-		return conn.(*dbConnector)
-	}
-	return dbPool.singleton
-}
-
-func (dbPool *DBPool) putConn(conn *dbConnector) *DBPool {
-	dbPool.pool.Put(conn)
-	return dbPool
-}
-
-// NewDBPool is constructor of DBPool
-func NewDBPool(size int, dbConf config.DBStore, lazyConn func(dbConf config.DBStore) (*sql.DB, error)) (*DBPool, error) {
-	// TODO: `size` validation
-	// NOTE: make sure we have one db connection at least
-	singleton, err := lazyConn(dbConf)
-	if err != nil {
-		return nil, err
-	}
-	dbPool := &DBPool{}
-	dbPool.singleton = &dbConnector{
-		dbConf:   dbConf,
-		lazyConn: lazyConn,
-		conn:     singleton,
-		err:      nil,
-	}
-	dbPool.putConn(dbPool.singleton)
-	for i := 0; i < size-1; i++ {
-		dbPool.putConn(&dbConnector{
-			dbConf:   dbConf,
-			lazyConn: lazyConn,
-			conn:     nil,
-			err:      nil,
-		})
-	}
-	return dbPool, nil
-}
-
-// ExecuteWithLog implement `SQLExecutor` interface
-func (dbPool *DBPool) ExecuteWithLog(ctx context.Context, query string, purpose string, logger log.Logger) error {
-	conn := dbPool.getConn()
-	defer func() {
-		dbPool.putConn(conn)
-	}()
-	if err := conn.connect(); err != nil {
-		return err
-	}
-	sql := common.SQLWithRetry{
-		DB:     conn.getDB(),
-		Logger: logger,
-	}
-	return sql.Exec(ctx, purpose, query)
-}
-
-// ObtainStringWithLog implement `SQLExecutor` interface
-func (dbPool *DBPool) ObtainStringWithLog(ctx context.Context, query string, purpose string, logger log.Logger) (string, error) {
-	conn := dbPool.getConn()
-	defer func() {
-		dbPool.putConn(conn)
-	}()
-	var s string
-	var err error
-	if err = conn.connect(); err == nil {
-		err = common.SQLWithRetry{
-			DB:     conn.getDB(),
-			Logger: logger,
-		}.QueryRow(ctx, purpose, query, &s)
-	}
-	return s, err
-}
-
-// Close implement `SQLExecutor` interface
-func (dbPool *DBPool) Close() {
-	for conn := dbPool.getConn(); conn != nil; conn = dbPool.getConn() {
-		if conn.getError() == nil {
-			// CONFUSED: why we ignore error occurs via db connection close issue?
-			_ = conn.close()
-		}
-	}
-}
-
 type ExternalTiDBGlue struct {
 	db     *sql.DB
 	parser *parser.Parser

lightning/restore/restore.go

@@ -15,7 +15,6 @@ package restore
 
 import (
 	"context"
-	"database/sql"
 	"fmt"
 	"io"
 	"math"
@@ -299,51 +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)
 		}
-		// CONFUSED: how we control concurrency via configuration?
-		// OR do not implement glue.SQLExecutor inside ExternalTiDBGlue
-		// Maybe we can aggregate an instance of glue.SQLExecutor?
-		concurrency := 16
-		var sqlExec glue.SQLExecutor
-		var err error
-		if concurrency > 1 {
-			dbFactory := func(conf config.DBStore) (*sql.DB, error) {
-				db, err := DBFromConfig(conf)
-				if err == nil {
-					db.ExecContext(ctx, "SET SQL_MODE = ?", conf.StrSQLMode)
-				}
-				return db, err
-			}
-			sqlExec, err = glue.NewDBPool(concurrency, rc.cfg.TiDB, dbFactory)
-			defer sqlExec.Close()
-			if err != nil {
-				return errors.Trace(err)
-			}
-		} else {
-			sqlExec = rc.tidbGlue.GetSQLExecutor()
+		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, concurrency, rc.tidbGlue.GetParser(), sqlExec, 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 {
@@ -352,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, concurrency, rc.tidbGlue.GetParser(), sqlExec, dbMeta.Name, viewsSchema)
+				err := InitSchema(ctx, concurrency, rc.tidbGlue.GetParser(), rc.tidbGlue.GetSQLExecutor(), dbMeta.Name, viewsSchema)
 
 				task.End(zap.ErrorLevel, err)
 				if err != nil {