br: split regions before pitr restore (#39941)

close #37207
pingcap · Jan 22, 2023 · 6ebbe64 · 6ebbe64
1 parent f1647c8
commit 6ebbe64
Show file tree

Hide file tree

Showing 12 changed files with 1,295 additions and 2 deletions.
diff --git a/br/pkg/restore/BUILD.bazel b/br/pkg/restore/BUILD.bazel
@@ -61,6 +61,7 @@ go_library(
         "//util/collate",
         "//util/hack",
         "//util/mathutil",
+        "//util/sqlexec",
         "//util/table-filter",
         "@com_github_emirpasic_gods//maps/treemap",
         "@com_github_go_sql_driver_mysql//:mysql",

diff --git a/br/pkg/restore/client.go b/br/pkg/restore/client.go
@@ -53,6 +53,7 @@ import (
 	"github.com/pingcap/tidb/util/codec"
 	"github.com/pingcap/tidb/util/collate"
 	"github.com/pingcap/tidb/util/mathutil"
+	"github.com/pingcap/tidb/util/sqlexec"
 	filter "github.com/pingcap/tidb/util/table-filter"
 	"github.com/tikv/client-go/v2/oracle"
 	pd "github.com/tikv/pd/client"
@@ -1126,6 +1127,18 @@ func (rc *Client) SplitRanges(ctx context.Context,
 	return SplitRanges(ctx, rc, ranges, rewriteRules, updateCh, isRawKv)
 }
 
+func (rc *Client) WrapLogFilesIterWithSplitHelper(iter LogIter, rules map[int64]*RewriteRules, g glue.Glue, store kv.Storage) (LogIter, error) {
+	se, err := g.CreateSession(store)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	execCtx := se.GetSessionCtx().(sqlexec.RestrictedSQLExecutor)
+	splitSize, splitKeys := utils.GetRegionSplitInfo(execCtx)
+	log.Info("get split threshold from tikv config", zap.Uint64("split-size", splitSize), zap.Int64("split-keys", splitKeys))
+	client := split.NewSplitClient(rc.GetPDClient(), rc.GetTLSConfig(), false)
+	return NewLogFilesIterWithSplitHelper(iter, rules, client, splitSize, splitKeys), nil
+}
+
 // RestoreSSTFiles tries to restore the files.
 func (rc *Client) RestoreSSTFiles(
 	ctx context.Context,