[close #421] br: Split for smaller batch of keys (#422)

Signed-off-by: Ping Yu <yuping@pingcap.com>

Author: pingyu

.github/config/br_pd.toml

@@ -2,3 +2,8 @@
 [replication]
 enable-placement-rules = true
 max-replicas = 1
+
+[schedule]
+max-merge-region-size = 2
+max-merge-region-keys = 40000
+split-merge-interval = "10s"

.github/config/br_rawkv.toml

@@ -7,11 +7,18 @@ pd-heartbeat-tick-interval = "2s"
 pd-store-heartbeat-tick-interval = "5s"
 split-region-check-tick-interval = "1s"
 
+# 10000000 keys & ~480MB data size in integration tests
+[coprocessor]
+region-split-size = "10MiB"
+region-split-keys = 200000
+batch-split-limit = 100
+
 [rocksdb]
 max-open-files = 10000
 
 [raftdb]
 max-open-files = 10000
 
 [storage]
+reserve-space = "0MiB"
 enable-ttl = true

br/Makefile

@@ -22,13 +22,14 @@ PACKAGES    := go list ./...
 DIRECTORIES := $(PACKAGES) | sed 's|github.com/tikv/migration/br/||'
 
 # build & test
-BR_BIN_PATH    ?= bin/tikv-br
-TEST_BIN_PATH  ?= bin/tikv-br.test
-COVERAGE_DIR   ?= build
-TEST_PARALLEL  ?= 8
-PD_ADDR        ?= 127.0.0.1:2379
-BR_LOCAL_STORE ?= /tmp/backup_restore_test
-API_VERSION    ?= 1
+BR_BIN_PATH           ?= bin/tikv-br
+TEST_BIN_PATH         ?= bin/tikv-br.test
+COVERAGE_DIR          ?= build
+TEST_PARALLEL         ?= 8
+PD_ADDR               ?= 127.0.0.1:2379
+SPLIT_REGION_MAX_KEYS ?= 4
+BR_LOCAL_STORE        ?= /tmp/backup_restore_test
+API_VERSION           ?= 1
 
 LDFLAGS += -X "github.com/tikv/migration/br/pkg/version/build.ReleaseVersion=$(shell git describe --tags --dirty --always)"
 LDFLAGS += -X "github.com/tikv/migration/br/pkg/version/build.BuildTS=$(shell date -u '+%Y-%m-%d %H:%M:%S')"
@@ -61,6 +62,7 @@ test: tools/bin/gocov tools/bin/gocov-xml
 test/integration: build/br-test build/rawkv-integration-test
 	./bin/rawkv_test --pd=${PD_ADDR} \
 		--br='${TEST_BIN_PATH}' \
+		--split-region-max-keys=${SPLIT_REGION_MAX_KEYS} \
 		--br-storage=${BR_LOCAL_STORE} \
 		--api-version=${API_VERSION}
 

br/pkg/restore/client.go

@@ -43,6 +43,7 @@ type Client struct {
 	workerPool    *utils.WorkerPool
 	tlsConf       *tls.Config
 	keepaliveConf keepalive.ClientParameters
+	spliterConf   SpliterConfig
 
 	backupMeta    *backuppb.BackupMeta
 	dstAPIVersion kvrpcpb.APIVersion
@@ -64,6 +65,7 @@ func NewRestoreClient(
 	pdClient pd.Client,
 	tlsConf *tls.Config,
 	keepaliveConf keepalive.ClientParameters,
+	spliterConf SpliterConfig,
 	isRawKv bool,
 ) (*Client, error) {
 	apiVerion, err := conn.GetTiKVApiVersion(context.Background(), pdClient, tlsConf)
@@ -72,9 +74,10 @@ func NewRestoreClient(
 	}
 	return &Client{
 		pdClient:      pdClient,
-		toolClient:    NewSplitClient(pdClient, tlsConf, isRawKv),
+		toolClient:    NewSplitClient(pdClient, tlsConf, spliterConf, isRawKv),
 		tlsConf:       tlsConf,
 		keepaliveConf: keepaliveConf,
+		spliterConf:   spliterConf,
 		switchCh:      make(chan struct{}),
 		dstAPIVersion: apiVerion,
 	}, nil
@@ -132,7 +135,7 @@ func (rc *Client) InitBackupMeta(
 	}
 	rc.backupMeta = backupMeta
 
-	metaClient := NewSplitClient(rc.pdClient, rc.tlsConf, rc.backupMeta.IsRawKv)
+	metaClient := NewSplitClient(rc.pdClient, rc.tlsConf, rc.spliterConf, rc.backupMeta.IsRawKv)
 	importCli := NewImportClient(metaClient, rc.tlsConf, rc.keepaliveConf)
 	rc.fileImporter = NewFileImporter(metaClient, importCli, backend, rc.backupMeta.IsRawKv,
 		rc.backupMeta.ApiVersion, rc.rateLimit)

br/pkg/restore/split.go

@@ -53,15 +53,23 @@ var (
 	ScanRegionAttemptTimes = 30
 )
 
+// SpliterConfig is the config for region split.
+type SpliterConfig struct {
+	GRPCMaxRecvMsgSize int
+	SplitRegionMaxKeys int
+}
+
 // RegionSplitter is a executor of region split by rules.
 type RegionSplitter struct {
 	client SplitClient
+	conf   SpliterConfig
 }
 
 // NewRegionSplitter returns a new RegionSplitter.
-func NewRegionSplitter(client SplitClient) *RegionSplitter {
+func NewRegionSplitter(client SplitClient, conf SpliterConfig) *RegionSplitter {
 	return &RegionSplitter{
 		client: client,
+		conf:   conf,
 	}
 }
 
@@ -122,46 +130,68 @@ SplitRegions:
 		for _, region := range regions {
 			regionMap[region.Region.GetId()] = region
 		}
-		for regionID, keys := range splitKeyMap {
-			log.Info("get split keys for region", zap.Int("len", len(keys)), zap.Uint64("region", regionID))
-			var newRegions []*RegionInfo
-			region := regionMap[regionID]
-			log.Info("split regions",
-				logutil.Region(region.Region), logutil.Keys(keys), rtree.ZapRanges(ranges))
-			newRegions, errSplit = rs.splitAndScatterRegions(ctx, region, keys)
-			if errSplit != nil {
-				if strings.Contains(errSplit.Error(), "no valid key") {
-					for _, key := range keys {
-						// Region start/end keys are encoded. split_region RPC
-						// requires raw keys (without encoding).
-						log.Error("split regions no valid key",
-							logutil.Key("startKey", region.Region.StartKey),
-							logutil.Key("endKey", region.Region.EndKey),
-							logutil.Key("key", codec.EncodeBytes(nil, key)),
-							rtree.ZapRanges(ranges))
+		for regionID, regionKeys := range splitKeyMap {
+			log.Info("get split keys for region", zap.Int("len", len(regionKeys)), zap.Uint64("region", regionID))
+
+			for i := 0; i < len(regionKeys); i += rs.conf.SplitRegionMaxKeys {
+				end := i + rs.conf.SplitRegionMaxKeys
+				if end > len(regionKeys) {
+					end = len(regionKeys)
+				}
+				keys := regionKeys[i:end]
+
+				var region *RegionInfo
+				if i == 0 {
+					region = regionMap[regionID]
+				} else {
+					encodedKey := codec.EncodeBytes(nil, keys[0])
+					var errGetRegion error
+					region, errGetRegion = rs.client.GetRegion(ctx, encodedKey)
+					if errGetRegion != nil {
+						time.Sleep(interval)
+						log.Warn("get region failed, retry", logutil.Key("encodedKey", encodedKey), zap.Error(errGetRegion))
+						continue SplitRegions
 					}
-					return errors.Trace(errSplit)
 				}
-				interval = 2 * interval
-				if interval > SplitMaxRetryInterval {
-					interval = SplitMaxRetryInterval
+
+				log.Info("split regions",
+					logutil.Region(region.Region), logutil.Keys(keys), rtree.ZapRanges(ranges))
+				var newRegions []*RegionInfo
+				newRegions, errSplit = rs.splitAndScatterRegions(ctx, region, keys)
+				if errSplit != nil {
+					if strings.Contains(errSplit.Error(), "no valid key") {
+						for _, key := range keys {
+							// Region start/end keys are encoded. split_region RPC
+							// requires raw keys (without encoding).
+							log.Error("split regions no valid key",
+								logutil.Key("startKey", region.Region.StartKey),
+								logutil.Key("endKey", region.Region.EndKey),
+								logutil.Key("key", codec.EncodeBytes(nil, key)),
+								rtree.ZapRanges(ranges))
+						}
+						return errors.Trace(errSplit)
+					}
+					interval = 2 * interval
+					if interval > SplitMaxRetryInterval {
+						interval = SplitMaxRetryInterval
+					}
+					time.Sleep(interval)
+					log.Warn("split regions failed, retry",
+						zap.Error(errSplit),
+						logutil.Region(region.Region),
+						logutil.Leader(region.Leader),
+						logutil.Keys(keys), rtree.ZapRanges(ranges))
+					continue SplitRegions
 				}
-				time.Sleep(interval)
-				log.Warn("split regions failed, retry",
-					zap.Error(errSplit),
-					logutil.Region(region.Region),
-					logutil.Leader(region.Leader),
-					logutil.Keys(keys), rtree.ZapRanges(ranges))
-				continue SplitRegions
-			}
-			log.Info("scattered regions", zap.Int("count", len(newRegions)))
-			if len(newRegions) != len(keys) {
-				log.Warn("split key count and new region count mismatch",
-					zap.Int("new region count", len(newRegions)),
-					zap.Int("split key count", len(keys)))
+				log.Info("scattered regions", zap.Int("count", len(newRegions)))
+				if len(newRegions) != len(keys) {
+					log.Warn("split key count and new region count mismatch",
+						zap.Int("new region count", len(newRegions)),
+						zap.Int("split key count", len(keys)))
+				}
+				scatterRegions = append(scatterRegions, newRegions...)
+				onSplit(keys)
 			}
-			scatterRegions = append(scatterRegions, newRegions...)
-			onSplit(keys)
 		}
 		break
 	}

br/pkg/restore/split_client.go

@@ -79,10 +79,11 @@ type SplitClient interface {
 
 // pdClient is a wrapper of pd client, can be used by RegionSplitter.
 type pdClient struct {
-	mu         sync.Mutex
-	client     pd.Client
-	tlsConf    *tls.Config
-	storeCache map[uint64]*metapb.Store
+	mu          sync.Mutex
+	client      pd.Client
+	tlsConf     *tls.Config
+	spliterConf SpliterConfig
+	storeCache  map[uint64]*metapb.Store
 
 	// FIXME when config changed during the lifetime of pdClient,
 	// 	this may mislead the scatter.
@@ -93,12 +94,13 @@ type pdClient struct {
 }
 
 // NewSplitClient returns a client used by RegionSplitter.
-func NewSplitClient(client pd.Client, tlsConf *tls.Config, isRawKv bool) SplitClient {
+func NewSplitClient(client pd.Client, tlsConf *tls.Config, spliterConf SpliterConfig, isRawKv bool) SplitClient {
 	cli := &pdClient{
-		client:     client,
-		tlsConf:    tlsConf,
-		storeCache: make(map[uint64]*metapb.Store),
-		isRawKv:    isRawKv,
+		client:      client,
+		tlsConf:     tlsConf,
+		spliterConf: spliterConf,
+		storeCache:  make(map[uint64]*metapb.Store),
+		isRawKv:     isRawKv,
 	}
 	return cli
 }
@@ -194,7 +196,12 @@ func (c *pdClient) SplitRegion(ctx context.Context, regionInfo *RegionInfo, key
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
-	conn, err := grpc.Dial(store.GetAddress(), grpc.WithInsecure())
+	opt := grpc.WithInsecure()
+	if c.tlsConf != nil {
+		opt = grpc.WithTransportCredentials(credentials.NewTLS(c.tlsConf))
+	}
+	conn, err := grpc.Dial(store.GetAddress(), opt,
+		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(c.spliterConf.GRPCMaxRecvMsgSize)))
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -318,7 +325,8 @@ func (c *pdClient) sendSplitRegionRequest(
 		if c.tlsConf != nil {
 			opt = grpc.WithTransportCredentials(credentials.NewTLS(c.tlsConf))
 		}
-		conn, err := grpc.Dial(store.GetAddress(), opt)
+		conn, err := grpc.Dial(store.GetAddress(), opt,
+			grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(c.spliterConf.GRPCMaxRecvMsgSize)))
 		if err != nil {
 			return nil, multierr.Append(splitErrors, err)
 		}

br/pkg/restore/split_test.go

@@ -276,7 +276,7 @@ func ScatterFinishInTimeImpl(t *testing.T, needEncodeKey bool) {
 	client := initTestClient()
 	ranges := initRanges()
 	rewriteRules := initRewriteRules()
-	regionSplitter := restore.NewRegionSplitter(client)
+	regionSplitter := restore.NewRegionSplitter(client, defaultSpliterCfg)
 
 	ctx := context.Background()
 	err := regionSplitter.Split(ctx, ranges, rewriteRules, needEncodeKey, func(key [][]byte) {}) // TODO: add test case for "isRawKV=true"
@@ -335,7 +335,7 @@ func TestSplitAndScatter(t *testing.T) {
 func runTestSplitAndScatterWith(t *testing.T, client *TestClient, needEncodeKey bool) {
 	ranges := initRanges()
 	rewriteRules := initRewriteRules()
-	regionSplitter := restore.NewRegionSplitter(client)
+	regionSplitter := restore.NewRegionSplitter(client, defaultSpliterCfg)
 
 	ctx := context.Background()
 	err := regionSplitter.Split(ctx, ranges, rewriteRules, needEncodeKey, func(key [][]byte) {}) // TODO: add test case for "isRawKV=true"

br/pkg/restore/util.go

@@ -110,13 +110,14 @@ func matchOldPrefix(key []byte, rewriteRules *RewriteRules) *import_sstpb.Rewrit
 func SplitRanges(
 	ctx context.Context,
 	client *Client,
+	spliterConf SpliterConfig,
 	ranges []rtree.Range,
 	rewriteRules *RewriteRules,
 	updateCh glue.Progress,
 	isRawKv bool,
 	needEncodeKey bool,
 ) error {
-	splitter := NewRegionSplitter(NewSplitClient(client.GetPDClient(), client.GetTLSConfig(), isRawKv))
+	splitter := NewRegionSplitter(NewSplitClient(client.GetPDClient(), client.GetTLSConfig(), spliterConf, isRawKv), spliterConf)
 
 	return splitter.Split(ctx, ranges, rewriteRules, needEncodeKey, func(keys [][]byte) {
 		updateCh.Inc()

br/pkg/task/common.go

@@ -27,6 +27,7 @@ import (
 	berrors "github.com/tikv/migration/br/pkg/errors"
 	"github.com/tikv/migration/br/pkg/glue"
 	"github.com/tikv/migration/br/pkg/metautil"
+	"github.com/tikv/migration/br/pkg/restore"
 	"github.com/tikv/migration/br/pkg/storage"
 	pd "github.com/tikv/pd/client"
 	"go.uber.org/zap"
@@ -60,14 +61,19 @@ const (
 	flagGrpcKeepaliveTime = "grpc-keepalive-time"
 	// flagGrpcKeepaliveTimeout is the max time a grpc conn can keep idel before killed.
 	flagGrpcKeepaliveTimeout = "grpc-keepalive-timeout"
+	// flagGrpcMaxRecvMsgSize is the max allowed size of a message received from tikv.
+	flagGrpcMaxRecvMsgSize = "grpc-max-recv-msg-size"
 	// flagEnableOpenTracing is whether to enable opentracing
-	flagEnableOpenTracing = "enable-opentracing"
-	flagSkipCheckPath     = "skip-check-path"
+	flagEnableOpenTracing  = "enable-opentracing"
+	flagSkipCheckPath      = "skip-check-path"
+	flagSplitRegionMaxKeys = "split-region-max-keys"
 
 	defaultSwitchInterval       = 5 * time.Minute
 	defaultGRPCKeepaliveTime    = 10 * time.Second
 	defaultGRPCKeepaliveTimeout = 3 * time.Second
+	defaultGRPCMaxRecvMsgSize   = 32 * 1024 * 1024 // 32MB
 	defaultChecksumConcurrency  = 512
+	defaultSplitRegionMaxKeys   = 1024
 
 	flagCipherType    = "crypter.method"
 	flagCipherKey     = "crypter.key"
@@ -115,8 +121,11 @@ func DefineCommonFlags(flags *pflag.FlagSet) {
 		"the interval of pinging gRPC peer, must keep the same value with TiKV and PD")
 	flags.Duration(flagGrpcKeepaliveTimeout, defaultGRPCKeepaliveTimeout,
 		"the max time a gRPC connection can keep idle before killed, must keep the same value with TiKV and PD")
+	flags.Uint(flagGrpcMaxRecvMsgSize, defaultGRPCMaxRecvMsgSize,
+		"the max allowed size of a message received from TiKV")
 	_ = flags.MarkHidden(flagGrpcKeepaliveTime)
 	_ = flags.MarkHidden(flagGrpcKeepaliveTimeout)
+	_ = flags.MarkHidden(flagGrpcMaxRecvMsgSize)
 
 	flags.Bool(flagEnableOpenTracing, false,
 		"Set whether to enable opentracing during the backup/restore process")
@@ -125,6 +134,9 @@ func DefineCommonFlags(flags *pflag.FlagSet) {
 	_ = flags.MarkHidden(flagNoCreds)
 	flags.BoolP(flagSkipCheckPath, "", false, "Skip path verification")
 	_ = flags.MarkHidden(flagSkipCheckPath)
+	flags.Uint(flagSplitRegionMaxKeys, defaultSplitRegionMaxKeys,
+		"The max number of keys in a single split region request")
+	_ = flags.MarkHidden(flagSplitRegionMaxKeys)
 
 	flags.String(flagCipherType, "plaintext", "Encrypt/decrypt method, "+
 		"be one of plaintext|aes128-ctr|aes192-ctr|aes256-ctr case-insensitively, "+
@@ -343,6 +355,14 @@ func GetKeepalive(cfg *Config) keepalive.ClientParameters {
 	}
 }
 
+// GetSpliterConfig get the spliter config.
+func GetSpliterConfig(cfg *Config) restore.SpliterConfig {
+	return restore.SpliterConfig{
+		GRPCMaxRecvMsgSize: int(cfg.GRPCMaxRecvMsgSize),
+		SplitRegionMaxKeys: int(cfg.SplitRegionMaxKeys),
+	}
+}
+
 func normalizePDURL(pd string, useTLS bool) (string, error) {
 	if strings.HasPrefix(pd, "http://") {
 		if useTLS {