*: limit IO concurrency

lightning/config/config.go

@@ -72,6 +72,7 @@ type Lightning struct {
 	common.LogConfig
 	TableConcurrency  int  `toml:"table-concurrency" json:"table-concurrency"`
 	RegionConcurrency int  `toml:"region-concurrency" json:"region-concurrency"`
+	IOConcurrency     int  `toml:"io-concurrency" json:"region-concurrency"`
 	ProfilePort       int  `toml:"pprof-port" json:"pprof-port"`
 	CheckRequirements bool `toml:"check-requirements" json:"check-requirements"`
 }
@@ -130,6 +131,7 @@ func NewConfig() *Config {
 		App: Lightning{
 			RegionConcurrency: runtime.NumCPU(),
 			TableConcurrency:  8,
+			IOConcurrency:     5,
 			CheckRequirements: true,
 		},
 		TiDB: DBStore{

lightning/mydump/parser.go

@@ -8,6 +8,7 @@ import (
 	"github.com/pkg/errors"
 
 	"github.com/pingcap/tidb-lightning/lightning/metric"
+	"github.com/pingcap/tidb-lightning/lightning/worker"
 )
 
 // ChunkParser is a parser of the data files (the file containing only INSERT
@@ -32,6 +33,7 @@ type ChunkParser struct {
 	// cache
 	remainBuf *bytes.Buffer
 	appendBuf *bytes.Buffer
+	ioWorkers *worker.RestoreWorkerPool
 }
 
 // Chunk represents a portion of the data file.
@@ -49,12 +51,13 @@ type Row struct {
 }
 
 // NewChunkParser creates a new parser which can read chunks out of a file.
-func NewChunkParser(reader io.Reader, blockBufSize int64) *ChunkParser {
+func NewChunkParser(reader io.Reader, blockBufSize int64, ioWorkers *worker.RestoreWorkerPool) *ChunkParser {
 	return &ChunkParser{
 		reader:    reader,
 		blockBuf:  make([]byte, blockBufSize),
 		remainBuf: &bytes.Buffer{},
 		appendBuf: &bytes.Buffer{},
+		ioWorkers: ioWorkers,
 	}
 }
 

lightning/mydump/parser_test.go

@@ -1,12 +1,15 @@
 package mydump_test
 
 import (
+	"context"
 	"io"
 	"strings"
 
 	. "github.com/pingcap/check"
 	"github.com/pingcap/tidb-lightning/lightning/config"
 	"github.com/pingcap/tidb-lightning/lightning/mydump"
+	"github.com/pingcap/tidb-lightning/lightning/worker"
+
 	"github.com/pkg/errors"
 )
 
@@ -25,7 +28,8 @@ func (s *testMydumpParserSuite) TestReadRow(c *C) {
 			"insert another_table values (10, 11, 12, '(13)', '(', 14, ')');",
 	)
 
-	parser := mydump.NewChunkParser(reader, config.ReadBlockSize)
+	ioWorkers := worker.NewRestoreWorkerPool(context.Background(), 5, "test")
+	parser := mydump.NewChunkParser(reader, config.ReadBlockSize, ioWorkers)
 
 	c.Assert(parser.ReadRow(), IsNil)
 	c.Assert(parser.LastRow(), DeepEquals, mydump.Row{
@@ -73,7 +77,8 @@ func (s *testMydumpParserSuite) TestReadChunks(c *C) {
 		INSERT foo VALUES (29,30,31,32),(33,34,35,36);
 	`)
 
-	parser := mydump.NewChunkParser(reader, config.ReadBlockSize)
+	ioWorkers := worker.NewRestoreWorkerPool(context.Background(), 5, "test")
+	parser := mydump.NewChunkParser(reader, config.ReadBlockSize, ioWorkers)
 
 	chunks, err := parser.ReadChunks(32)
 	c.Assert(err, IsNil)
@@ -119,7 +124,8 @@ func (s *testMydumpParserSuite) TestNestedRow(c *C) {
 		("789",CONVERT("[]" USING UTF8MB4));
 	`)
 
-	parser := mydump.NewChunkParser(reader, config.ReadBlockSize)
+	ioWorkers := worker.NewRestoreWorkerPool(context.Background(), 5, "test")
+	parser := mydump.NewChunkParser(reader, config.ReadBlockSize, ioWorkers)
 	chunks, err := parser.ReadChunks(96)
 
 	c.Assert(err, IsNil)

lightning/restore/restore.go

@@ -24,6 +24,8 @@ import (
 	"github.com/pingcap/tidb-lightning/lightning/metric"
 	"github.com/pingcap/tidb-lightning/lightning/mydump"
 	verify "github.com/pingcap/tidb-lightning/lightning/verification"
+	"github.com/pingcap/tidb-lightning/lightning/worker"
+
 	tidbcfg "github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/meta/autoid"
 	"github.com/pingcap/tidb/util/hack"
@@ -95,8 +97,9 @@ type RestoreController struct {
 	cfg             *config.Config
 	dbMetas         []*mydump.MDDatabaseMeta
 	dbInfos         map[string]*TidbDBInfo
-	tableWorkers    *RestoreWorkerPool
-	regionWorkers   *RestoreWorkerPool
+	tableWorkers    *worker.RestoreWorkerPool
+	regionWorkers   *worker.RestoreWorkerPool
+	ioWorkers       *worker.RestoreWorkerPool
 	importer        *kv.Importer
 	tidbMgr         *TiDBManager
 	postProcessLock sync.Mutex // a simple way to ensure post-processing is not concurrent without using complicated goroutines
@@ -129,8 +132,9 @@ func NewRestoreController(ctx context.Context, dbMetas []*mydump.MDDatabaseMeta,
 	rc := &RestoreController{
 		cfg:           cfg,
 		dbMetas:       dbMetas,
-		tableWorkers:  NewRestoreWorkerPool(ctx, cfg.App.TableConcurrency, "table"),
-		regionWorkers: NewRestoreWorkerPool(ctx, cfg.App.RegionConcurrency, "region"),
+		tableWorkers:  worker.NewRestoreWorkerPool(ctx, cfg.App.TableConcurrency, "table"),
+		regionWorkers: worker.NewRestoreWorkerPool(ctx, cfg.App.RegionConcurrency, "region"),
+		ioWorkers:     worker.NewRestoreWorkerPool(ctx, cfg.App.IOConcurrency, "io"),
 		importer:      importer,
 		tidbMgr:       tidbMgr,
 
@@ -439,9 +443,9 @@ func (rc *RestoreController) restoreTables(ctx context.Context) error {
 			// Note: We still need tableWorkers to control the concurrency of tables. In the future, we will investigate more about
 			// the difference between restoring tables concurrently and restoring tables one by one.
 
-			worker := rc.tableWorkers.Apply()
+			restoreWorker := rc.tableWorkers.Apply()
 			wg.Add(1)
-			go func(w *RestoreWorker, t *TableRestore, cp *TableCheckpoint) {
+			go func(w *worker.RestoreWorker, t *TableRestore, cp *TableCheckpoint) {
 				defer wg.Done()
 
 				closedEngine, err := t.restore(ctx, rc, cp)
@@ -465,7 +469,7 @@ func (rc *RestoreController) restoreTables(ctx context.Context) error {
 				}
 
 				err = t.postProcess(ctx, closedEngine, rc, cp)
-			}(worker, tr, cp)
+			}(restoreWorker, tr, cp)
 		}
 	}
 
@@ -548,15 +552,15 @@ func (t *TableRestore) restore(ctx context.Context, rc *RestoreController, cp *T
 		// 	3. load kvs data (into kv deliver server)
 		// 	4. flush kvs data (into tikv node)
 
-		cr, err := newChunkRestore(chunkIndex, chunk, rc.cfg.Mydumper.ReadBlockSize)
+		cr, err := newChunkRestore(chunkIndex, chunk, rc.cfg.Mydumper.ReadBlockSize, rc.ioWorkers)
 		if err != nil {
 			return nil, errors.Trace(err)
 		}
 		metric.ChunkCounter.WithLabelValues(metric.ChunkStatePending).Inc()
 
-		worker := rc.regionWorkers.Apply()
+		restoreWorker := rc.regionWorkers.Apply()
 		wg.Add(1)
-		go func(w *RestoreWorker, cr *chunkRestore) {
+		go func(w *worker.RestoreWorker, cr *chunkRestore) {
 			// Restore a chunk.
 			defer func() {
 				cr.close()
@@ -581,7 +585,7 @@ func (t *TableRestore) restore(ctx context.Context, rc *RestoreController, cp *T
 
 			handled := int(atomic.AddInt32(handledChunksCount, 1))
 			common.AppLogger.Infof("[%s] handled region count = %d (%s)", t.tableName, handled, common.Percent(handled, len(cp.Chunks)))
-		}(worker, cr)
+		}(restoreWorker, cr)
 	}
 
 	wg.Wait()
@@ -860,56 +864,18 @@ func (rc *RestoreController) getTables() []string {
 	return tables
 }
 
-////////////////////////////////////////////////////////////////
-
-type RestoreWorkerPool struct {
-	limit   int
-	workers chan *RestoreWorker
-	name    string
-}
-
-type RestoreWorker struct {
-	ID int64
-}
-
-func NewRestoreWorkerPool(ctx context.Context, limit int, name string) *RestoreWorkerPool {
-	workers := make(chan *RestoreWorker, limit)
-	for i := 0; i < limit; i++ {
-		workers <- &RestoreWorker{ID: int64(i + 1)}
-	}
-
-	metric.IdleWorkersGauge.WithLabelValues(name).Set(float64(limit))
-	return &RestoreWorkerPool{
-		limit:   limit,
-		workers: workers,
-		name:    name,
-	}
-}
-
-func (pool *RestoreWorkerPool) Apply() *RestoreWorker {
-	worker := <-pool.workers
-	metric.IdleWorkersGauge.WithLabelValues(pool.name).Set(float64(len(pool.workers)))
-	return worker
-}
-func (pool *RestoreWorkerPool) Recycle(worker *RestoreWorker) {
-	pool.workers <- worker
-	metric.IdleWorkersGauge.WithLabelValues(pool.name).Set(float64(len(pool.workers)))
-}
-
-////////////////////////////////////////////////////////////////
-
 type chunkRestore struct {
 	parser *mydump.ChunkParser
 	index  int
 	chunk  *ChunkCheckpoint
 }
 
-func newChunkRestore(index int, chunk *ChunkCheckpoint, blockBufSize int64) (*chunkRestore, error) {
+func newChunkRestore(index int, chunk *ChunkCheckpoint, blockBufSize int64, ioWorkers *worker.RestoreWorkerPool) (*chunkRestore, error) {
 	reader, err := os.Open(chunk.Key.Path)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
-	parser := mydump.NewChunkParser(reader, blockBufSize)
+	parser := mydump.NewChunkParser(reader, blockBufSize, ioWorkers)
 
 	reader.Seek(chunk.Chunk.Offset, io.SeekStart)
 	parser.SetPos(chunk.Chunk.Offset, chunk.Chunk.PrevRowIDMax)

lightning/worker/worker.go

@@ -0,0 +1,45 @@
+package worker
+
+import (
+	"context"
+
+	"github.com/pingcap/tidb-lightning/lightning/metric"
+)
+
+////////////////////////////////////////////////////////////////
+
+type RestoreWorkerPool struct {
+	limit   int
+	workers chan *RestoreWorker
+	name    string
+}
+
+type RestoreWorker struct {
+	ID int64
+}
+
+func NewRestoreWorkerPool(ctx context.Context, limit int, name string) *RestoreWorkerPool {
+	workers := make(chan *RestoreWorker, limit)
+	for i := 0; i < limit; i++ {
+		workers <- &RestoreWorker{ID: int64(i + 1)}
+	}
+
+	metric.IdleWorkersGauge.WithLabelValues(name).Set(float64(limit))
+	return &RestoreWorkerPool{
+		limit:   limit,
+		workers: workers,
+		name:    name,
+	}
+}
+
+func (pool *RestoreWorkerPool) Apply() *RestoreWorker {
+	worker := <-pool.workers
+	metric.IdleWorkersGauge.WithLabelValues(pool.name).Set(float64(len(pool.workers)))
+	return worker
+}
+func (pool *RestoreWorkerPool) Recycle(worker *RestoreWorker) {
+	pool.workers <- worker
+	metric.IdleWorkersGauge.WithLabelValues(pool.name).Set(float64(len(pool.workers)))
+}
+
+////////////////////////////////////////////////////////////////