br: copy full backup to pitr storage

br/cmd/br/operator.go

@@ -35,6 +35,8 @@ func newOperatorCommand() *cobra.Command {
 	cmd.AddCommand(newBase64ifyCommand())
 	cmd.AddCommand(newListMigrationsCommand())
 	cmd.AddCommand(newMigrateToCommand())
+	cmd.AddCommand(newForceFlushCommand())
+	cmd.AddCommand(newChecksumCommand())
 	return cmd
 }
 
@@ -109,3 +111,43 @@ func newMigrateToCommand() *cobra.Command {
 	operator.DefineFlagsForMigrateToConfig(cmd.Flags())
 	return cmd
 }
+
+func newChecksumCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "checksum-as",
+		Short: "calculate the checksum with rewrite rules",
+		Long: "Calculate the checksum of the current cluster (specified by `-u`) " +
+			"with applying the rewrite rules generated from a backup (specified by `-s`). " +
+			"This can be used when you have the checksum of upstream elsewhere.",
+		Args: cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			cfg := operator.ChecksumWithRewriteRulesConfig{}
+			if err := cfg.ParseFromFlags(cmd.Flags()); err != nil {
+				return err
+			}
+			ctx := GetDefaultContext()
+			return operator.RunChecksumTable(ctx, tidbGlue, cfg)
+		},
+	}
+	task.DefineFilterFlags(cmd, []string{"!*.*"}, false)
+	operator.DefineFlagsForChecksumTableConfig(cmd.Flags())
+	return cmd
+}
+
+func newForceFlushCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "force-flush",
+		Short: "force a log backup task to flush",
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			cfg := operator.ForceFlushConfig{}
+			if err := cfg.ParseFromFlags(cmd.Flags()); err != nil {
+				return err
+			}
+			ctx := GetDefaultContext()
+			return operator.RunForceFlush(ctx, &cfg)
+		},
+	}
+	operator.DefineFlagsForForceFlushConfig(cmd.Flags())
+	return cmd
+}

br/pkg/checkpoint/restore.go

@@ -19,6 +19,7 @@ import (
 	"encoding/json"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/br/pkg/glue"
 	"github.com/pingcap/tidb/br/pkg/pdutil"
@@ -137,6 +138,8 @@ type CheckpointMetadataForSnapshotRestore struct {
 	UpstreamClusterID uint64                `json:"upstream-cluster-id"`
 	RestoredTS        uint64                `json:"restored-ts"`
 	SchedulersConfig  *pdutil.ClusterConfig `json:"schedulers-config"`
+
+	RestoreUUID uuid.UUID `json:"restore-uuid"`
 }
 
 func LoadCheckpointMetadataForSnapshotRestore(

br/pkg/checkpoint/storage.go

@@ -32,6 +32,11 @@ import (
 	"go.uber.org/zap"
 )
 
+type hookedOnFlush struct {
+	checkpointStorage
+	cb func(ctx context.Context) error
+}
+
 type checkpointStorage interface {
 	flushCheckpointData(ctx context.Context, data []byte) error
 	flushCheckpointChecksum(ctx context.Context, data []byte) error

br/pkg/errors/errors.go

@@ -40,6 +40,7 @@ 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"))
 	ErrInvalidRange                 = errors.Normalize("invalid restore range", errors.RFCCodeText("BR:Common:ErrInvalidRange"))
+	ErrMigrationNotFound            = errors.Normalize("no migrtion found", errors.RFCCodeText("BR:Common:ErrMigrationNotFound"))
 	ErrMigrationVersionNotSupported = errors.Normalize("the migration version isn't supported", errors.RFCCodeText("BR:Common:ErrMigrationVersionNotSupported"))
 
 	ErrPDUpdateFailed         = errors.Normalize("failed to update PD", errors.RFCCodeText("BR:PD:ErrPDUpdateFailed"))

br/pkg/glue/glue.go

@@ -4,6 +4,7 @@ package glue
 
 import (
 	"context"
+	"sync/atomic"
 
 	"github.com/pingcap/tidb/pkg/ddl"
 	"github.com/pingcap/tidb/pkg/domain"
@@ -82,3 +83,22 @@ type Progress interface {
 	// called.
 	Close()
 }
+
+type CounterProgress struct {
+	Counter atomic.Int64
+}
+
+func (c *CounterProgress) Inc() {
+	c.Counter.Add(1)
+}
+
+func (c *CounterProgress) IncBy(cnt int64) {
+	c.Counter.Add(cnt)
+}
+
+func (c *CounterProgress) GetCurrent() int64 {
+	return c.Counter.Load()
+}
+
+func (c *CounterProgress) Close() {
+}

br/pkg/gluetidb/glue.go

@@ -55,6 +55,10 @@ type Glue struct {
 	startDomainMu *sync.Mutex
 }
 
+func WrapSession(se sessiontypes.Session) glue.Session {
+	return &tidbSession{se: se}
+}
+
 type tidbSession struct {
 	se sessiontypes.Session
 }

br/pkg/restore/log_client/BUILD.bazel

@@ -11,6 +11,7 @@ go_library(
         "log_file_map.go",
         "log_split_strategy.go",
         "migration.go",
+        "ssts.go",
     ],
     importpath = "github.com/pingcap/tidb/br/pkg/restore/log_client",
     visibility = ["//visibility:public"],
@@ -43,11 +44,13 @@ go_library(
         "//pkg/kv",
         "//pkg/meta",
         "//pkg/meta/model",
+        "//pkg/tablecodec",
         "//pkg/util",
         "//pkg/util/codec",
         "//pkg/util/redact",
         "//pkg/util/sqlexec",
         "//pkg/util/table-filter",
+        "@com_github_docker_go_units//:go-units",
         "@com_github_fatih_color//:color",
         "@com_github_gogo_protobuf//proto",
         "@com_github_opentracing_opentracing_go//:opentracing-go",
@@ -90,7 +93,7 @@ go_test(
     ],
     embed = [":log_client"],
     flaky = True,
-    shard_count = 45,
+    shard_count = 50,
     deps = [
         "//br/pkg/errors",
         "//br/pkg/glue",
@@ -119,6 +122,7 @@ go_test(
         "//pkg/util/sqlexec",
         "//pkg/util/table-filter",
         "@com_github_docker_go_units//:go-units",
+        "@com_github_google_uuid//:uuid",
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_pingcap_kvproto//pkg/brpb",

br/pkg/restore/log_client/client.go

@@ -25,13 +25,16 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
+	"github.com/docker/go-units"
 	"github.com/fatih/color"
 	"github.com/gogo/protobuf/proto"
 	"github.com/opentracing/opentracing-go"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
+	backup "github.com/pingcap/kvproto/pkg/brpb"
 	backuppb "github.com/pingcap/kvproto/pkg/brpb"
 	"github.com/pingcap/kvproto/pkg/encryptionpb"
 	"github.com/pingcap/log"
@@ -40,6 +43,7 @@ import (
 	"github.com/pingcap/tidb/br/pkg/conn"
 	"github.com/pingcap/tidb/br/pkg/conn/util"
 	"github.com/pingcap/tidb/br/pkg/encryption"
+	berrors "github.com/pingcap/tidb/br/pkg/errors"
 	"github.com/pingcap/tidb/br/pkg/glue"
 	"github.com/pingcap/tidb/br/pkg/logutil"
 	"github.com/pingcap/tidb/br/pkg/metautil"
@@ -178,6 +182,7 @@ func NewSstRestoreManager(
 
 type LogClient struct {
 	*LogFileManager
+
 	logRestoreManager *LogRestoreManager
 	sstRestoreManager *SstRestoreManager
 
@@ -209,6 +214,16 @@ type LogClient struct {
 
 	// checkpoint information for log restore
 	useCheckpoint bool
+
+	logFilesStat logFilesStatistic
+	restoreStat  restoreStatistic
+}
+
+type restoreStatistic struct {
+	restoreSSTKVSize  uint64
+	restoreSSTKVCount uint64
+	restoreSSTPhySize uint64
+	restoreSSTTakes   uint64
 }
 
 // NewRestoreClient returns a new RestoreClient.
@@ -252,11 +267,12 @@ func (rc *LogClient) Close(ctx context.Context) {
 
 func (rc *LogClient) RestoreCompactedSstFiles(
 	ctx context.Context,
-	compactionsIter iter.TryNextor[*backuppb.LogFileSubcompaction],
+	compactionsIter iter.TryNextor[SSTs],
 	rules map[int64]*restoreutils.RewriteRules,
 	importModeSwitcher *restore.ImportModeSwitcher,
 	onProgress func(int64),
 ) error {
+	begin := time.Now()
 	backupFileSets := make([]restore.BackupFileSet, 0, 8)
 	// Collect all items from the iterator in advance to avoid blocking during restoration.
 	// This approach ensures that we have all necessary data ready for processing,
@@ -267,14 +283,36 @@ func (rc *LogClient) RestoreCompactedSstFiles(
 			return r.Err
 		}
 		i := r.Item
-		rewriteRules, ok := rules[i.Meta.TableId]
+
+		tid := i.TableID()
+		if r, ok := i.(RewrittenSST); ok && r.RewrittenTo() > 0 {
+			tid = r.RewrittenTo()
+		}
+		rewriteRules, ok := rules[tid]
 		if !ok {
-			log.Warn("[Compacted SST Restore] Skipping excluded table during restore.", zap.Int64("table_id", i.Meta.TableId))
+			log.Warn("[Compacted SST Restore] Skipping excluded table during restore.", zap.Int64("table_id", i.TableID()))
 			continue
 		}
+
+		if r, ok := i.(RewrittenSST); ok {
+			rewritten := r.RewrittenTo()
+			if rewritten > 0 && rewritten != i.TableID() {
+				rewriteRules = rewriteRules.Clone()
+				if !rewriteRules.RewriteSourceTableID(rewritten, i.TableID()) {
+					return errors.Annotatef(
+						berrors.ErrUnknown,
+						"table rewritten from a table id (%d) to (%d) which doesn't exist in the stream",
+						rewritten,
+						i.TableID(),
+					)
+				}
+				log.Info("Rewritten rewrite rules.", zap.Stringer("rules", rewriteRules), zap.Int64("table_id", i.TableID()), zap.Int64("rewritten_to", rewritten))
+			}
+		}
+
 		set := restore.BackupFileSet{
-			TableID:      i.Meta.TableId,
-			SSTFiles:     i.SstOutputs,
+			TableID:      i.TableID(),
+			SSTFiles:     i.GetSSTs(),
 			RewriteRules: rewriteRules,
 		}
 		backupFileSets = append(backupFileSets, set)
@@ -311,7 +349,30 @@ func (rc *LogClient) RestoreCompactedSstFiles(
 	if err != nil {
 		return errors.Trace(err)
 	}
-	return rc.sstRestoreManager.restorer.WaitUntilFinish()
+	err = rc.sstRestoreManager.restorer.WaitUntilFinish()
+
+	for _, files := range backupFileSets {
+		for _, f := range files.SSTFiles {
+			log.Info("Collected file.", zap.Uint64("total_kv", f.TotalKvs), zap.Uint64("total_bytes", f.TotalBytes), zap.Uint64("size", f.Size_))
+			atomic.AddUint64(&rc.restoreStat.restoreSSTKVCount, f.TotalKvs)
+			atomic.AddUint64(&rc.restoreStat.restoreSSTKVSize, f.TotalBytes)
+			atomic.AddUint64(&rc.restoreStat.restoreSSTPhySize, f.Size_)
+		}
+	}
+	atomic.AddUint64(&rc.restoreStat.restoreSSTTakes, uint64(time.Since(begin)))
+	return err
+}
+
+func (rc *LogClient) RestoreSSTStatisticFields(pushTo *[]zapcore.Field) {
+	takes := time.Duration(rc.restoreStat.restoreSSTTakes)
+	fields := []zapcore.Field{
+		zap.Uint64("restore-sst-kv-count", rc.restoreStat.restoreSSTKVCount),
+		zap.Uint64("restore-sst-kv-size", rc.restoreStat.restoreSSTKVSize),
+		zap.Uint64("restore-sst-physical-size (after compression)", rc.restoreStat.restoreSSTPhySize),
+		zap.Duration("restore-sst-total-take", takes),
+		zap.String("average-speed (sst)", units.HumanSize(float64(rc.restoreStat.restoreSSTKVSize)/takes.Seconds())+"/s"),
+	}
+	*pushTo = append(*pushTo, fields...)
 }
 
 func (rc *LogClient) SetRawKVBatchClient(
@@ -516,13 +577,29 @@ func (rc *LogClient) InitCheckpointMetadataForLogRestore(
 	return gcRatio, nil
 }
 
-func (rc *LogClient) GetMigrations(ctx context.Context) ([]*backuppb.Migration, error) {
-	ext := stream.MigerationExtension(rc.storage)
+type LockedMigrations struct {
+	Migs     []*backup.Migration
+	ReadLock storage.RemoteLock
+}
+
+func (rc *LogClient) GetMigrations(ctx context.Context) (*LockedMigrations, error) {
+	ext := stream.MigrationExtension(rc.storage)
 	migs, err := ext.Load(ctx)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
-	return migs.ListAll(), nil
+
+	ms := migs.ListAll()
+	readLock, err := ext.GetReadLock(ctx, "restore stream")
+	if err != nil {
+		return nil, err
+	}
+
+	lms := &LockedMigrations{
+		Migs:     ms,
+		ReadLock: readLock,
+	}
+	return lms, nil
 }
 
 func (rc *LogClient) InstallLogFileManager(ctx context.Context, startTS, restoreTS uint64, metadataDownloadBatchSize uint,
@@ -544,6 +621,8 @@ func (rc *LogClient) InstallLogFileManager(ctx context.Context, startTS, restore
 	if err != nil {
 		return err
 	}
+	rc.logFilesStat = logFilesStatistic{}
+	rc.LogFileManager.Stats = &rc.logFilesStat
 	return nil
 }
 
@@ -1509,15 +1588,15 @@ func (rc *LogClient) UpdateSchemaVersion(ctx context.Context) error {
 // It uses a region splitter to handle the splitting logic based on the provided rules and checkpoint sets.
 func (rc *LogClient) WrapCompactedFilesIterWithSplitHelper(
 	ctx context.Context,
-	compactedIter iter.TryNextor[*backuppb.LogFileSubcompaction],
+	compactedIter iter.TryNextor[SSTs],
 	rules map[int64]*restoreutils.RewriteRules,
 	checkpointSets map[string]struct{},
 	updateStatsFn func(uint64, uint64),
 	splitSize uint64,
 	splitKeys int64,
-) (iter.TryNextor[*backuppb.LogFileSubcompaction], error) {
+) (iter.TryNextor[SSTs], error) {
 	client := split.NewClient(rc.pdClient, rc.pdHTTPClient, rc.tlsConf, maxSplitKeysOnce, 3)
-	wrapper := restore.PipelineRestorerWrapper[*backuppb.LogFileSubcompaction]{
+	wrapper := restore.PipelineRestorerWrapper[SSTs]{
 		PipelineRegionsSplitter: split.NewPipelineRegionsSplitter(client, splitSize, splitKeys),
 	}
 	strategy := NewCompactedFileSplitStrategy(rules, checkpointSets, updateStatsFn)