br: Allow Skip Finegrained Scatter to Speed up Split & Scatter

br/pkg/restore/snap_client/client.go

@@ -96,6 +96,8 @@ type SnapClient struct {
 	cipher                *backuppb.CipherInfo
 	concurrencyPerStore   uint
 	regionScanConcurrency uint
+	splitRegionIndexStep  uint
+	coarseScatter         bool
 	keepaliveConf         keepalive.ClientParameters
 	rateLimit             uint64
 	tlsConf               *tls.Config
@@ -186,11 +188,12 @@ func NewRestoreClient(
 	keepaliveConf keepalive.ClientParameters,
 ) *SnapClient {
 	return &SnapClient{
-		pdClient:      pdClient,
-		pdHTTPClient:  pdHTTPCli,
-		tlsConf:       tlsConf,
-		keepaliveConf: keepaliveConf,
-		switchCh:      make(chan struct{}),
+		pdClient:             pdClient,
+		pdHTTPClient:         pdHTTPCli,
+		tlsConf:              tlsConf,
+		keepaliveConf:        keepaliveConf,
+		splitRegionIndexStep: split.DefaultRegionIndexStep,
+		switchCh:             make(chan struct{}),
 	}
 }
 
@@ -317,6 +320,28 @@ func (rc *SnapClient) GetRegionScanConcurrency() uint {
 	return rc.regionScanConcurrency
 }
 
+// SetSplitRegionIndexStep sets the rough split step during snapshot restore.
+func (rc *SnapClient) SetSplitRegionIndexStep(step uint) {
+	log.Info("split region index step", zap.Uint("step", step))
+	rc.splitRegionIndexStep = split.NormalizeRegionIndexStep(step)
+}
+
+// GetSplitRegionIndexStep returns the rough split step during snapshot restore.
+func (rc *SnapClient) GetSplitRegionIndexStep() uint {
+	return rc.splitRegionIndexStep
+}
+
+// SetCoarseScatter controls whether only rough split regions are scattered.
+func (rc *SnapClient) SetCoarseScatter(coarseScatter bool) {
+	log.Info("coarse scatter", zap.Bool("enabled", coarseScatter))
+	rc.coarseScatter = coarseScatter
+}
+
+// GetCoarseScatter returns whether only rough split regions are scattered.
+func (rc *SnapClient) GetCoarseScatter() bool {
+	return rc.coarseScatter
+}
+
 func (rc *SnapClient) SetBatchDdlSize(batchDdlsize uint) {
 	rc.batchDdlSize = batchDdlsize
 }

br/pkg/restore/snap_client/tikv_sender.go

@@ -367,14 +367,15 @@ func (rc *SnapClient) SplitPoints(
 		splitClientOpts = append(splitClientOpts, split.WithRawKV())
 	}
 
-	splitter := split.NewRegionSplitter(split.NewClient(
+	splitter := split.NewRegionSplitterWithRegionIndexStep(split.NewClient(
 		rc.pdClient,
 		rc.pdHTTPClient,
 		rc.tlsConf,
 		maxSplitKeysOnce,
 		rc.storeCount+1,
 		splitClientOpts...,
-	))
+	), rc.splitRegionIndexStep)
+	splitter.SetCoarseScatter(rc.coarseScatter)
 
 	return splitter.ExecuteSortedKeys(ctx, sortedSplitKeys)
 }

br/pkg/restore/split/client.go

@@ -69,6 +69,9 @@ type SplitClient interface {
 	// new regions. It returns the new regions that need to be called with
 	// WaitRegionsScattered.
 	SplitKeysAndScatter(ctx context.Context, sortedSplitKeys [][]byte) ([]*RegionInfo, error)
+	// SplitKeys splits the related regions of the keys without scattering the new
+	// regions.
+	SplitKeys(ctx context.Context, sortedSplitKeys [][]byte) ([]*RegionInfo, error)
 
 	// SplitWaitAndScatter splits a region from a batch of keys, waits for the split
 	// is finished, and scatters the new regions. It will return the original region,
@@ -627,6 +630,14 @@ func (c *pdClient) getEncodedKeys(start, end []byte) (encodedStart, encodedEnd [
 }
 
 func (c *pdClient) SplitKeysAndScatter(ctx context.Context, sortedSplitKeys [][]byte) ([]*RegionInfo, error) {
+	return c.splitKeys(ctx, sortedSplitKeys, true)
+}
+
+func (c *pdClient) SplitKeys(ctx context.Context, sortedSplitKeys [][]byte) ([]*RegionInfo, error) {
+	return c.splitKeys(ctx, sortedSplitKeys, false)
+}
+
+func (c *pdClient) splitKeys(ctx context.Context, sortedSplitKeys [][]byte, scatter bool) ([]*RegionInfo, error) {
 	if len(sortedSplitKeys) == 0 {
 		return nil, nil
 	}
@@ -686,7 +697,7 @@ func (c *pdClient) SplitKeysAndScatter(ctx context.Context, sortedSplitKeys [][]
 		for region, splitKeys := range splitKeyMap {
 			workerPool.ApplyOnErrorGroup(eg, func() error {
 				// TODO(lance6716): add error handling to retry from scan or retry from split
-				newRegions, err2 := c.SplitWaitAndScatter(eCtx, region, splitKeys)
+				newRegions, err2 := c.splitWaitAndMaybeScatter(eCtx, region, splitKeys, scatter)
 				if err2 != nil {
 					if common.IsContextCanceledError(err2) {
 						return err2
@@ -701,14 +712,16 @@ func (c *pdClient) SplitKeysAndScatter(ctx context.Context, sortedSplitKeys [][]
 					return nil
 				}
 
-				if len(newRegions) != len(splitKeys) {
+				if scatter && len(newRegions) != len(splitKeys) {
 					log.Warn("split key count and new region count mismatch",
 						zap.Int("new region count", len(newRegions)),
 						zap.Int("split key count", len(splitKeys)))
 				}
-				mu.Lock()
-				ret = append(ret, newRegions...)
-				mu.Unlock()
+				if scatter {
+					mu.Lock()
+					ret = append(ret, newRegions...)
+					mu.Unlock()
+				}
 				return nil
 			})
 		}
@@ -729,6 +742,10 @@ func isNonRetryErrForSplit(err error) bool {
 }
 
 func (c *pdClient) SplitWaitAndScatter(ctx context.Context, region *RegionInfo, keys [][]byte) ([]*RegionInfo, error) {
+	return c.splitWaitAndMaybeScatter(ctx, region, keys, true)
+}
+
+func (c *pdClient) splitWaitAndMaybeScatter(ctx context.Context, region *RegionInfo, keys [][]byte, scatter bool) ([]*RegionInfo, error) {
 	failpoint.Inject("failToSplit", func(_ failpoint.Value) {
 		failpoint.Return(nil, errors.New("retryable error"))
 	})
@@ -761,12 +778,14 @@ func (c *pdClient) SplitWaitAndScatter(ctx context.Context, region *RegionInfo,
 			if err = ctx.Err(); err != nil {
 				return nil, errors.Trace(err)
 			}
-			err = c.scatterRegions(ctx, newRegionsOfBatch)
-			if err != nil {
-				tidblogutil.Logger(ctx).Warn(
-					"scatter regions failed, will continue anyway",
-					zap.Error(err),
-				)
+			if scatter {
+				err = c.scatterRegions(ctx, newRegionsOfBatch)
+				if err != nil {
+					tidblogutil.Logger(ctx).Warn(
+						"scatter regions failed, will continue anyway",
+						zap.Error(err),
+					)
+				}
 			}
 			if c.onSplit != nil {
 				c.onSplit(keys[start:end])

br/pkg/restore/split/client_test.go

@@ -119,7 +119,23 @@ func TestSplitScatter(t *testing.T) {
 		}
 	}
 
-	_, err := mockClient.SplitKeysAndScatter(ctx, splitKeys)
+	noScatterPDClient := NewMockPDClientForSplit()
+	noScatterPDClient.SetRegions(keys)
+	noScatterClient := &pdClient{
+		client:           noScatterPDClient,
+		splitConcurrency: 20,
+		splitBatchKeyCnt: 100,
+	}
+	noScatterSplitKeys := make([][]byte, 0, len(splitKeys))
+	for _, key := range splitKeys {
+		noScatterSplitKeys = append(noScatterSplitKeys, append([]byte(nil), key...))
+	}
+	_, err := noScatterClient.SplitKeys(ctx, noScatterSplitKeys)
+	require.NoError(t, err)
+	require.Positive(t, noScatterPDClient.splitRegions.count)
+	require.Zero(t, noScatterPDClient.scatterRegions.regionCount)
+
+	_, err = mockClient.SplitKeysAndScatter(ctx, splitKeys)
 	require.NoError(t, err)
 
 	// check split ranges

br/pkg/restore/split/mock_pd_client.go

@@ -110,6 +110,14 @@ func (c *TestClient) GetRegionByID(ctx context.Context, regionID uint64) (*Regio
 	return region, nil
 }
 
+func (c *TestClient) SplitKeysAndScatter(ctx context.Context, keys [][]byte) ([]*RegionInfo, error) {
+	return c.SplitWaitAndScatter(ctx, nil, keys)
+}
+
+func (c *TestClient) SplitKeys(ctx context.Context, keys [][]byte) ([]*RegionInfo, error) {
+	return c.SplitWaitAndScatter(ctx, nil, keys)
+}
+
 func (c *TestClient) SplitWaitAndScatter(_ context.Context, _ *RegionInfo, keys [][]byte) ([]*RegionInfo, error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -690,6 +698,14 @@ func (f *FakeSplitClient) WaitRegionsScattered(context.Context, []*RegionInfo) (
 	return 0, nil
 }
 
+func (f *FakeSplitClient) SplitKeysAndScatter(context.Context, [][]byte) ([]*RegionInfo, error) {
+	return nil, nil
+}
+
+func (f *FakeSplitClient) SplitKeys(context.Context, [][]byte) ([]*RegionInfo, error) {
+	return nil, nil
+}
+
 func (*FakeSplitClient) GetCodecPDClient() *tikvclient.CodecPDClient {
 	return nil
 }

br/pkg/restore/split/split.go

@@ -36,20 +36,44 @@ const (
 	ScatterWaitUpperInterval = 30 * time.Minute
 
 	ScanRegionPaginationLimit = 128
+
+	// DefaultRegionIndexStep is the default number of split-key indexes between two rough split keys.
+	DefaultRegionIndexStep uint = 128
 )
 
+// NormalizeRegionIndexStep returns the default rough split step when the configured value is unset.
+func NormalizeRegionIndexStep(regionIndexStep uint) uint {
+	if regionIndexStep == 0 {
+		return DefaultRegionIndexStep
+	}
+	return regionIndexStep
+}
+
 // RegionSplitter is a executor of region split by rules.
 type RegionSplitter struct {
-	client SplitClient
+	client          SplitClient
+	regionIndexStep uint
+	coarseScatter   bool
 }
 
 // NewRegionSplitter returns a new RegionSplitter.
 func NewRegionSplitter(client SplitClient) *RegionSplitter {
+	return NewRegionSplitterWithRegionIndexStep(client, DefaultRegionIndexStep)
+}
+
+// NewRegionSplitterWithRegionIndexStep returns a new RegionSplitter with the configured rough split step.
+func NewRegionSplitterWithRegionIndexStep(client SplitClient, regionIndexStep uint) *RegionSplitter {
 	return &RegionSplitter{
-		client: client,
+		client:          client,
+		regionIndexStep: NormalizeRegionIndexStep(regionIndexStep),
 	}
 }
 
+// SetCoarseScatter makes only rough split regions scattered.
+func (rs *RegionSplitter) SetCoarseScatter(coarseScatter bool) {
+	rs.coarseScatter = coarseScatter
+}
+
 // ExecuteSortedKeysOnRegion expose the function `SplitWaitAndScatter` of split client.
 func (rs *RegionSplitter) ExecuteSortedKeysOnRegion(ctx context.Context, region *RegionInfo, keys [][]byte) ([]*RegionInfo, error) {
 	return rs.client.SplitWaitAndScatter(ctx, region, keys)
@@ -80,22 +104,24 @@ func (rs *RegionSplitter) executeSplitByRanges(
 ) error {
 	startTime := time.Now()
 	// Choose the rough region split keys,
-	// each splited region contains 128 regions to be splitted.
-	const regionIndexStep = 128
-
-	roughSortedSplitKeys := make([][]byte, 0, len(sortedKeys)/regionIndexStep+1)
-	for curRegionIndex := regionIndexStep; curRegionIndex < len(sortedKeys); curRegionIndex += regionIndexStep {
-		roughSortedSplitKeys = append(roughSortedSplitKeys, sortedKeys[curRegionIndex])
+	// each split region contains regionIndexStep regions to be split.
+	var roughSortedSplitKeys [][]byte
+	if rs.regionIndexStep < uint(len(sortedKeys)) {
+		regionIndexStep := int(rs.regionIndexStep)
+		roughSortedSplitKeys = make([][]byte, 0, len(sortedKeys)/regionIndexStep+1)
+		for curRegionIndex := regionIndexStep; curRegionIndex < len(sortedKeys); curRegionIndex += regionIndexStep {
+			roughSortedSplitKeys = append(roughSortedSplitKeys, sortedKeys[curRegionIndex])
+		}
 	}
 	if len(roughSortedSplitKeys) > 0 {
-		if err := rs.executeSplitByKeys(ctx, roughSortedSplitKeys); err != nil {
+		if err := rs.executeSplitByKeys(ctx, roughSortedSplitKeys, true); err != nil {
 			return errors.Trace(err)
 		}
 	}
 	log.Info("finish spliting regions roughly", zap.Duration("take", time.Since(startTime)))
 
 	// Then send split requests to each TiKV.
-	if err := rs.executeSplitByKeys(ctx, sortedKeys); err != nil {
+	if err := rs.executeSplitByKeys(ctx, sortedKeys, !rs.coarseScatter); err != nil {
 		return errors.Trace(err)
 	}
 
@@ -110,9 +136,10 @@ func (rs *RegionSplitter) executeSplitByRanges(
 func (rs *RegionSplitter) executeSplitByKeys(
 	ctx context.Context,
 	sortedKeys [][]byte,
+	scatter bool,
 ) error {
 	startTime := time.Now()
-	scatterRegions, err := rs.client.SplitKeysAndScatter(ctx, sortedKeys)
+	scatterRegions, err := rs.splitKeys(ctx, sortedKeys, scatter)
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -126,6 +153,13 @@ func (rs *RegionSplitter) executeSplitByKeys(
 	return nil
 }
 
+func (rs *RegionSplitter) splitKeys(ctx context.Context, sortedKeys [][]byte, scatter bool) ([]*RegionInfo, error) {
+	if scatter {
+		return rs.client.SplitKeysAndScatter(ctx, sortedKeys)
+	}
+	return rs.client.SplitKeys(ctx, sortedKeys)
+}
+
 // waitRegionsScattered try to wait mutilple regions scatterd in 3 minutes.
 // this could timeout, but if many regions scatterd the restore could continue
 // so we don't wait long time here.