address comment

Signed-off-by: lance6716 <lance6716@gmail.com>

br/pkg/restore/split/client.go

@@ -539,12 +539,15 @@ func (c *pdClient) SplitKeysAndScatter(ctx context.Context, sortedSplitKeys [][]
 	}
 	// we need to find the regions that contain the split keys. However, the scan
 	// region API accepts a key range [start, end) where end key is exclusive, and if
-	// sortedSplitKeys length is 1, we scan region may return empty result. So we
+	// sortedSplitKeys length is 1, scan region may return empty result. So we
 	// increase the end key a bit. If the end key is on the region boundaries, it
 	// will be skipped by getSplitKeysOfRegions.
 	scanStart := codec.EncodeBytesExt(nil, sortedSplitKeys[0], c.isRawKv)
 	lastKey := kv.Key(sortedSplitKeys[len(sortedSplitKeys)-1])
-	scanEnd := codec.EncodeBytesExt(nil, lastKey.Next(), c.isRawKv)
+	if len(lastKey) > 0 {
+		lastKey = lastKey.Next()
+	}
+	scanEnd := codec.EncodeBytesExt(nil, lastKey, c.isRawKv)
 
 	// mu protects ret, retrySplitKeys, lastSplitErr
 	mu := sync.Mutex{}

br/pkg/restore/split/client_test.go

@@ -181,6 +181,47 @@ func TestSplitScatterRawKV(t *testing.T) {
 	require.Equal(t, len(result)-3, mockPDClient.scatterRegions.regionCount)
 }
 
+func TestSplitScatterEmptyEndKey(t *testing.T) {
+	backup := maxBatchSplitSize
+	maxBatchSplitSize = 7
+	t.Cleanup(func() {
+		maxBatchSplitSize = backup
+	})
+
+	mockPDClient := NewMockPDClientForSplit()
+	keys := [][]byte{[]byte(""), []byte("aay"), []byte("bba"), []byte("bbh"), []byte("cca"), []byte("")}
+	mockPDClient.SetRegions(keys)
+	mockClient := &pdClient{
+		client:           mockPDClient,
+		splitConcurrency: 10,
+		splitBatchKeyCnt: 100,
+		isRawKv:          true, // make tests more readable
+	}
+	ctx := context.Background()
+
+	splitKeys := [][]byte{{'b'}, {'c'}, {}}
+
+	_, err := mockClient.SplitKeysAndScatter(ctx, splitKeys)
+	require.NoError(t, err)
+
+	// check split ranges
+	regions, err := PaginateScanRegion(ctx, mockClient, []byte{}, []byte{}, 5)
+	require.NoError(t, err)
+	result := [][]byte{
+		[]byte(""), []byte("aay"),
+		[]byte("b"), []byte("bba"), []byte("bbh"),
+		[]byte("c"), []byte("cca"), []byte(""),
+	}
+	checkRegionsBoundaries(t, regions, result)
+
+	// test only one split key which is empty
+	_, err = mockClient.SplitKeysAndScatter(ctx, [][]byte{{}})
+	require.NoError(t, err)
+	regions, err = PaginateScanRegion(ctx, mockClient, []byte{}, []byte{}, 5)
+	require.NoError(t, err)
+	checkRegionsBoundaries(t, regions, result)
+}
+
 func TestScanRegionEmptyResult(t *testing.T) {
 	backup := WaitRegionOnlineAttemptTimes
 	WaitRegionOnlineAttemptTimes = 2

br/pkg/restore/split_test.go

@@ -285,40 +285,6 @@ func initRewriteRules() *RewriteRules {
 	}
 }
 
-// expected regions after split:
-//
-//	[, aay), [aay, bba), [bba, bbf), [bbf, bbh), [bbh, bbj),
-//	[bbj, cca), [cca, xxe), [xxe, xxz), [xxz, )
-func validateRegions(regions map[uint64]*split.RegionInfo) bool {
-	keys := [...]string{"", "aay", "bba", "bbf", "bbh", "bbj", "cca", "xxe", "xxz", ""}
-	return validateRegionsExt(regions, keys[:], false)
-}
-
-func validateRegionsExt(regions map[uint64]*split.RegionInfo, expectedKeys []string, isRawKv bool) bool {
-	if len(regions) != len(expectedKeys)-1 {
-		return false
-	}
-FindRegion:
-	for i := 1; i < len(expectedKeys); i++ {
-		for _, region := range regions {
-			startKey := []byte(expectedKeys[i-1])
-			if len(startKey) != 0 {
-				startKey = codec.EncodeBytesExt([]byte{}, startKey, isRawKv)
-			}
-			endKey := []byte(expectedKeys[i])
-			if len(endKey) != 0 {
-				endKey = codec.EncodeBytesExt([]byte{}, endKey, isRawKv)
-			}
-			if bytes.Equal(region.Region.GetStartKey(), startKey) &&
-				bytes.Equal(region.Region.GetEndKey(), endKey) {
-				continue FindRegion
-			}
-		}
-		return false
-	}
-	return true
-}
-
 type fakeRestorer struct {
 	mu sync.Mutex