restore: refactor split code for two phase split.

br/pkg/errors/errors.go

@@ -40,10 +40,12 @@ var (
 	ErrEnvNotSpecified           = errors.Normalize("environment variable not found", errors.RFCCodeText("BR:Common:ErrEnvNotSpecified"))
 	ErrUnsupportedOperation      = errors.Normalize("the operation is not supported", errors.RFCCodeText("BR:Common:ErrUnsupportedOperation"))
 
-	ErrPDUpdateFailed    = errors.Normalize("failed to update PD", errors.RFCCodeText("BR:PD:ErrPDUpdateFailed"))
-	ErrPDLeaderNotFound  = errors.Normalize("PD leader not found", errors.RFCCodeText("BR:PD:ErrPDLeaderNotFound"))
-	ErrPDInvalidResponse = errors.Normalize("PD invalid response", errors.RFCCodeText("BR:PD:ErrPDInvalidResponse"))
-	ErrPDBatchScanRegion = errors.Normalize("batch scan region", errors.RFCCodeText("BR:PD:ErrPDBatchScanRegion"))
+	ErrPDUpdateFailed         = errors.Normalize("failed to update PD", errors.RFCCodeText("BR:PD:ErrPDUpdateFailed"))
+	ErrPDLeaderNotFound       = errors.Normalize("PD leader not found", errors.RFCCodeText("BR:PD:ErrPDLeaderNotFound"))
+	ErrPDInvalidResponse      = errors.Normalize("PD invalid response", errors.RFCCodeText("BR:PD:ErrPDInvalidResponse"))
+	ErrPDBatchScanRegion      = errors.Normalize("batch scan region", errors.RFCCodeText("BR:PD:ErrPDBatchScanRegion"))
+	ErrPDUnknownScatterResult = errors.Normalize("failed to wait region scattered", errors.RFCCodeText("BR:PD:ErrPDUknownScatterResult"))
+	ErrPDSplitFailed          = errors.Normalize("failed to wait region splitted", errors.RFCCodeText("BR:PD:ErrPDUknownScatterResult"))
 
 	ErrBackupChecksumMismatch    = errors.Normalize("backup checksum mismatch", errors.RFCCodeText("BR:Backup:ErrBackupChecksumMismatch"))
 	ErrBackupInvalidRange        = errors.Normalize("backup range invalid", errors.RFCCodeText("BR:Backup:ErrBackupInvalidRange"))

br/pkg/restore/client.go

@@ -27,6 +27,7 @@ import (
 	"github.com/pingcap/log"
 	"github.com/pingcap/tidb/br/pkg/checkpoint"
 	"github.com/pingcap/tidb/br/pkg/checksum"
+	"github.com/pingcap/tidb/br/pkg/conn"
 	"github.com/pingcap/tidb/br/pkg/conn/util"
 	berrors "github.com/pingcap/tidb/br/pkg/errors"
 	"github.com/pingcap/tidb/br/pkg/glue"
@@ -119,10 +120,12 @@ type Client struct {
 	dbPool          []*DB
 	rateLimit       uint64
 	isOnline        bool
+	granularity     string
 	noSchema        bool
 	hasSpeedLimited bool
 
 	restoreStores []uint64
+	storeCount    int
 
 	cipher             *backuppb.CipherInfo
 	switchModeInterval time.Duration
@@ -457,6 +460,10 @@ func (rc *Client) GetDomain() *domain.Domain {
 	return rc.dom
 }
 
+func (rc *Client) GetStoreCount() int {
+	return rc.storeCount
+}
+
 // GetPDClient returns a pd client.
 func (rc *Client) GetPDClient() pd.Client {
 	return rc.pdClient
@@ -523,10 +530,25 @@ func (rc *Client) SetStorage(ctx context.Context, backend *backuppb.StorageBacke
 	return nil
 }
 
-func (rc *Client) InitClients(backend *backuppb.StorageBackend, isRawKvMode bool, isTxnKvMode bool) {
+func (rc *Client) InitClients(ctx context.Context, backend *backuppb.StorageBackend, isRawKvMode bool, isTxnKvMode bool) {
+	storeWorkerPoolMap := make(map[uint64]chan struct{})
+	storeStatisticMap := make(map[uint64]*int64)
+	stores, err := conn.GetAllTiKVStoresWithRetry(ctx, rc.pdClient, util.SkipTiFlash)
+	if err != nil {
+		log.Fatal("failed to get stores", zap.Error(err))
+	}
+	concurrencyPerStore := 512
+	for _, store := range stores {
+		ch := make(chan struct{}, concurrencyPerStore)
+		for i := 0; i < concurrencyPerStore; i += 1 {
+			ch <- struct{}{}
+		}
+		storeWorkerPoolMap[store.Id] = ch
+		storeStatisticMap[store.Id] = new(int64)
+	}
 	metaClient := split.NewSplitClient(rc.pdClient, rc.pdHTTPClient, rc.tlsConf, isRawKvMode)
 	importCli := NewImportClient(metaClient, rc.tlsConf, rc.keepaliveConf)
-	rc.fileImporter = NewFileImporter(metaClient, importCli, backend, isRawKvMode, isTxnKvMode, rc.rewriteMode)
+	rc.fileImporter = NewFileImporter(metaClient, importCli, backend, isRawKvMode, isTxnKvMode, storeWorkerPoolMap, storeStatisticMap, rc.rewriteMode)
 }
 
 func (rc *Client) SetRawKVClient(c *RawKVBatchClient) {
@@ -566,7 +588,7 @@ func (rc *Client) InitBackupMeta(
 	}
 	rc.backupMeta = backupMeta
 
-	rc.InitClients(backend, backupMeta.IsRawKv, backupMeta.IsTxnKv)
+	rc.InitClients(c, backend, backupMeta.IsRawKv, backupMeta.IsTxnKv)
 	log.Info("load backupmeta", zap.Int("databases", len(rc.databases)), zap.Int("jobs", len(rc.ddlJobs)))
 	return rc.fileImporter.CheckMultiIngestSupport(c, rc.pdClient)
 }
@@ -643,6 +665,16 @@ func (rc *Client) EnableOnline() {
 	rc.isOnline = true
 }
 
+// SetGranularity sets the ganularity of restore pipeline.
+func (rc *Client) SetGranularity(g string) {
+	rc.granularity = g
+}
+
+// GetGranularity sets the ganularity of restore pipeline.
+func (rc *Client) GetGranularity() string {
+	return rc.granularity
+}
+
 // GetTLSConfig returns the tls config.
 func (rc *Client) GetTLSConfig() *tls.Config {
 	return rc.tlsConf
@@ -1224,7 +1256,7 @@ func MockCallSetSpeedLimit(ctx context.Context, fakeImportClient ImporterClient,
 	rc.SetRateLimit(42)
 	rc.SetConcurrency(concurrency)
 	rc.hasSpeedLimited = false
-	rc.fileImporter = NewFileImporter(nil, fakeImportClient, nil, false, false, rc.rewriteMode)
+	rc.fileImporter = NewFileImporter(nil, fakeImportClient, nil, false, false, nil, nil, rc.rewriteMode)
 	return rc.setSpeedLimit(ctx, rc.rateLimit)
 }
 
@@ -2005,31 +2037,33 @@ const (
 
 // LoadRestoreStores loads the stores used to restore data.
 func (rc *Client) LoadRestoreStores(ctx context.Context) error {
-	if !rc.isOnline {
-		return nil
-	}
 	if span := opentracing.SpanFromContext(ctx); span != nil && span.Tracer() != nil {
 		span1 := span.Tracer().StartSpan("Client.LoadRestoreStores", opentracing.ChildOf(span.Context()))
 		defer span1.Finish()
 		ctx = opentracing.ContextWithSpan(ctx, span1)
 	}
+	stores, err := conn.GetAllTiKVStoresWithRetry(ctx, rc.pdClient, util.SkipTiFlash)
 
-	stores, err := rc.pdClient.GetAllStores(ctx)
 	if err != nil {
 		return errors.Trace(err)
 	}
 	for _, s := range stores {
 		if s.GetState() != metapb.StoreState_Up {
 			continue
 		}
-		for _, l := range s.GetLabels() {
-			if l.GetKey() == restoreLabelKey && l.GetValue() == restoreLabelValue {
-				rc.restoreStores = append(rc.restoreStores, s.GetId())
-				break
+		rc.storeCount++
+		if rc.isOnline {
+			for _, l := range s.GetLabels() {
+				if l.GetKey() == restoreLabelKey && l.GetValue() == restoreLabelValue {
+					rc.restoreStores = append(rc.restoreStores, s.GetId())
+					break
+				}
 			}
 		}
 	}
-	log.Info("load restore stores", zap.Uint64s("store-ids", rc.restoreStores))
+	if rc.isOnline {
+		log.Info("load restore stores", zap.Uint64s("store-ids", rc.restoreStores))
+	}
 	return nil
 }
 

br/pkg/restore/import.go

@@ -54,8 +54,9 @@ const (
 )
 
 const (
-	importScanRegionTime = 10 * time.Second
+	importScanRegionTime = 20 * time.Second
 	gRPCBackOffMaxDelay  = 3 * time.Second
+	gRPCTimeOut          = 10 * time.Minute
 )
 
 // RewriteMode is a mode flag that tells the TiKV how to handle the rewrite rules.
@@ -267,6 +268,10 @@ type FileImporter struct {
 	importClient ImporterClient
 	backend      *backuppb.StorageBackend
 
+	storeWorkerPoolRWLock sync.RWMutex
+	storeWorkerPoolMap    map[uint64]chan struct{}
+	storeStatisticMap     map[uint64]*int64
+
 	kvMode             KvMode
 	rawStartKey        []byte
 	rawEndKey          []byte
@@ -283,6 +288,8 @@ func NewFileImporter(
 	backend *backuppb.StorageBackend,
 	isRawKvMode bool,
 	isTxnKvMode bool,
+	storeWorkerPoolMap map[uint64]chan struct{},
+	storeStatisticMap map[uint64]*int64,
 	rewriteMode RewriteMode,
 ) FileImporter {
 	kvMode := TiDB
@@ -293,12 +300,14 @@ func NewFileImporter(
 		kvMode = Txn
 	}
 	return FileImporter{
-		metaClient:   metaClient,
-		backend:      backend,
-		importClient: importClient,
-		kvMode:       kvMode,
-		rewriteMode:  rewriteMode,
-		cacheKey:     fmt.Sprintf("BR-%s-%d", time.Now().Format("20060102150405"), rand.Int63()),
+		metaClient:         metaClient,
+		backend:            backend,
+		importClient:       importClient,
+		storeStatisticMap:  storeStatisticMap,
+		storeWorkerPoolMap: storeWorkerPoolMap,
+		kvMode:             kvMode,
+		rewriteMode:        rewriteMode,
+		cacheKey:           fmt.Sprintf("BR-%s-%d", time.Now().Format("20060102150405"), rand.Int63()),
 	}
 }
 
@@ -569,17 +578,18 @@ func (importer *FileImporter) ImportSSTFiles(
 			log.Debug("download file done",
 				zap.String("file-sample", files[0].Name), zap.Stringer("take", time.Since(start)),
 				logutil.Key("start", files[0].StartKey), logutil.Key("end", files[0].EndKey))
-			if errIngest := importer.ingest(ctx, info, downloadMetas); errIngest != nil {
+			start = time.Now()
+			if errIngest := importer.ingest(ctx, files, info, downloadMetas); errIngest != nil {
 				log.Warn("ingest file failed, retry later",
 					logutil.Files(files),
 					logutil.SSTMetas(downloadMetas),
 					logutil.Region(info.Region),
 					zap.Error(errIngest))
 				return errors.Trace(errIngest)
 			}
+			log.Debug("ingest file done", zap.String("file-sample", files[0].Name), zap.Stringer("take", time.Since(start)))
 		}
 
-		log.Debug("ingest file done", zap.String("file-sample", files[0].Name), zap.Stringer("take", time.Since(start)))
 		for _, f := range files {
 			summary.CollectSuccessUnit(summary.TotalKV, 1, f.TotalKvs)
 			summary.CollectSuccessUnit(summary.TotalBytes, 1, f.TotalBytes)
@@ -719,7 +729,32 @@ func (importer *FileImporter) downloadSST(
 	for _, p := range regionInfo.Region.GetPeers() {
 		peer := p
 		eg.Go(func() error {
-			resp, err := importer.importClient.DownloadSST(ectx, peer.GetStoreId(), req)
+			importer.storeWorkerPoolRWLock.RLock()
+			workerCh := importer.storeWorkerPoolMap[peer.GetStoreId()]
+			statis := importer.storeStatisticMap[peer.GetStoreId()]
+			defer func() {
+				atomic.AddInt64(statis, -1)
+				workerCh <- struct{}{}
+				importer.storeWorkerPoolRWLock.RUnlock()
+			}()
+			<-workerCh
+			atomic.AddInt64(statis, 1)
+
+			var err error
+			var resp *import_sstpb.DownloadResponse
+			for i := 0; i < 5; i += 1 {
+				dctx, cancel := context.WithTimeout(ectx, gRPCTimeOut)
+				resp, err = importer.importClient.DownloadSST(dctx, peer.GetStoreId(), req)
+				cancel()
+				if err != nil {
+					if strings.Contains(err.Error(), "context deadline exceeded") {
+						continue
+					}
+					return errors.Trace(err)
+				}
+
+				break
+			}
 			if err != nil {
 				return errors.Trace(err)
 			}
@@ -832,10 +867,13 @@ func (importer *FileImporter) downloadRawKVSST(
 }
 
 func (importer *FileImporter) ingest(
-	ctx context.Context,
+	c context.Context,
+	files []*backuppb.File,
 	info *split.RegionInfo,
 	downloadMetas []*import_sstpb.SSTMeta,
 ) error {
+	ctx, cancel := context.WithTimeout(c, gRPCTimeOut)
+	defer cancel()
 	for {
 		ingestResp, errIngest := importer.ingestSSTs(ctx, downloadMetas, info)
 		if errIngest != nil {
@@ -866,7 +904,10 @@ func (importer *FileImporter) ingest(
 						break
 					}
 					// do not get region info, wait a second and GetRegion() again.
-					log.Warn("get region by key return nil", logutil.Region(info.Region))
+					log.Warn("ingest get region by key return nil", logutil.Region(info.Region),
+						logutil.Files(files),
+						logutil.SSTMetas(downloadMetas),
+					)
 					time.Sleep(time.Second)
 				}
 			}
@@ -875,6 +916,8 @@ func (importer *FileImporter) ingest(
 				return errors.Trace(berrors.ErrKVEpochNotMatch)
 			}
 			log.Debug("ingest sst returns not leader error, retry it",
+				logutil.Files(files),
+				logutil.SSTMetas(downloadMetas),
 				logutil.Region(info.Region),
 				zap.Stringer("newLeader", newInfo.Leader))
 			info = newInfo

br/pkg/restore/import_retry.go

@@ -257,7 +257,7 @@ func (r *RPCResult) StrategyForRetryGoError() RetryStrategy {
 	// we should unwrap the error or we cannot get the write gRPC status.
 	if gRPCErr, ok := status.FromError(errors.Cause(r.Err)); ok {
 		switch gRPCErr.Code() {
-		case codes.Unavailable, codes.Aborted, codes.ResourceExhausted:
+		case codes.Unavailable, codes.Aborted, codes.ResourceExhausted, codes.DeadlineExceeded:
 			return StrategyFromThisRegion
 		}
 	}

br/pkg/restore/merge.go

@@ -119,10 +119,15 @@ func MergeFileRanges(
 	rangeTree := rtree.NewRangeTree()
 	for key := range filesMap {
 		files := filesMap[key]
+		rangeSize := uint64(0)
+		for _, f := range filesMap[key] {
+			rangeSize += f.Size_
+		}
 		if out := rangeTree.InsertRange(rtree.Range{
 			StartKey: files[0].GetStartKey(),
 			EndKey:   files[0].GetEndKey(),
 			Files:    files,
+			Size:     rangeSize,
 		}); out != nil {
 			return nil, nil, errors.Annotatef(berrors.ErrRestoreInvalidRange,
 				"duplicate range %s files %+v", out, files)
@@ -136,6 +141,7 @@ func MergeFileRanges(
 			continue
 		}
 		sortedRanges[i-1].EndKey = sortedRanges[i].EndKey
+		sortedRanges[i-1].Size += sortedRanges[i].Size
 		sortedRanges[i-1].Files = append(sortedRanges[i-1].Files, sortedRanges[i].Files...)
 		// TODO: this is slow when there are lots of ranges need to merge.
 		sortedRanges = append(sortedRanges[:i], sortedRanges[i+1:]...)

br/pkg/restore/pipeline_items.go

@@ -243,6 +243,7 @@ func NewTiKVSender(
 	cli TiKVRestorer,
 	updateCh glue.Progress,
 	splitConcurrency uint,
+	granularity string,
 ) (BatchSender, error) {
 	inCh := make(chan DrainResult, defaultChannelSize)
 	midCh := make(chan drainResultAndDone, defaultChannelSize)
@@ -257,7 +258,13 @@ func NewTiKVSender(
 
 	sender.wg.Add(2)
 	go sender.splitWorker(ctx, inCh, midCh, splitConcurrency)
-	go sender.restoreWorker(ctx, midCh)
+	if granularity == string(CoarseGrained) {
+		outCh := make(chan drainResultAndDone, defaultChannelSize)
+		go sender.blockPipelineWorker(ctx, midCh, outCh)
+		go sender.restoreWorker(ctx, outCh)
+	} else {
+		go sender.restoreWorker(ctx, midCh)
+	}
 	return sender, nil
 }
 
@@ -272,6 +279,26 @@ type drainResultAndDone struct {
 	done   func()
 }
 
+func (b *tikvSender) blockPipelineWorker(ctx context.Context,
+	inCh <-chan drainResultAndDone,
+	outCh chan<- drainResultAndDone,
+) {
+	defer close(outCh)
+	res := make([]drainResultAndDone, 0, defaultChannelSize)
+	for dr := range inCh {
+		res = append(res, dr)
+	}
+
+	for _, dr := range res {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+			outCh <- dr
+		}
+	}
+}
+
 func (b *tikvSender) splitWorker(ctx context.Context,
 	ranges <-chan DrainResult,
 	next chan<- drainResultAndDone,
@@ -385,6 +412,7 @@ func (b *tikvSender) restoreWorker(ctx context.Context, ranges <-chan drainResul
 			if !ok {
 				return
 			}
+
 			files := r.result.Files()
 			// There has been a worker in the `RestoreSSTFiles` procedure.
 			// Spawning a raw goroutine won't make too many requests to TiKV.