br: retain latest mvcc version when merge into one SST

br/pkg/restore/internal/import_client/import_client.go

@@ -62,6 +62,13 @@ type ImporterClient interface {
 		req *import_sstpb.DownloadRequest,
 	) (*import_sstpb.DownloadResponse, error)
 
+	// BatchDownloadLatestMVCC downloads SSTs keeping only the latest MVCC version per key (compacted SST restore).
+	BatchDownloadLatestMVCC(
+		ctx context.Context,
+		storeID uint64,
+		req *import_sstpb.DownloadRequest,
+	) (*import_sstpb.DownloadResponse, error)
+
 	MultiIngest(
 		ctx context.Context,
 		storeID uint64,
@@ -83,6 +90,9 @@ type ImporterClient interface {
 
 	CheckBatchDownloadSupport(ctx context.Context, stores []uint64) (bool, error)
 
+	// CheckBatchDownloadLatestMVCCSupport returns an error if any store returns Unimplemented for BatchDownloadLatestMVCC.
+	CheckBatchDownloadLatestMVCCSupport(ctx context.Context, stores []uint64) error
+
 	CheckMultiIngestSupport(ctx context.Context, stores []uint64) error
 
 	AddForcePartitionRange(ctx context.Context, storeID uint64, req *import_sstpb.AddPartitionRangeRequest) error
@@ -163,6 +173,18 @@ func (ic *importClient) BatchDownloadSST(
 	return client.BatchDownload(ctx, req)
 }
 
+func (ic *importClient) BatchDownloadLatestMVCC(
+	ctx context.Context,
+	storeID uint64,
+	req *import_sstpb.DownloadRequest,
+) (*import_sstpb.DownloadResponse, error) {
+	client, err := ic.GetImportClient(ctx, storeID)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	return client.BatchDownloadLatestMVCC(ctx, req)
+}
+
 func (ic *importClient) SetDownloadSpeedLimit(
 	ctx context.Context,
 	storeID uint64,
@@ -283,6 +305,24 @@ func (ic *importClient) CheckBatchDownloadSupport(ctx context.Context, stores []
 	return true, nil
 }
 
+func (ic *importClient) CheckBatchDownloadLatestMVCCSupport(ctx context.Context, stores []uint64) error {
+	for _, storeID := range stores {
+		_, err := ic.BatchDownloadLatestMVCC(ctx, storeID, &import_sstpb.DownloadRequest{})
+		if err != nil {
+			if s, ok := status.FromError(err); ok {
+				if s.Code() == codes.Unimplemented {
+					return errors.Errorf(
+						"tikv node doesn't support BatchDownloadLatestMVCC; upgrade TiKV or disable --retain-latest-mvcc-version (store id %d)",
+						storeID,
+					)
+				}
+			}
+			return errors.Annotatef(err, "failed to check BatchDownloadLatestMVCC support. (store id %d)", storeID)
+		}
+	}
+	return nil
+}
+
 func (ic *importClient) CheckMultiIngestSupport(ctx context.Context, stores []uint64) error {
 	for _, storeID := range stores {
 		_, err := ic.MultiIngest(ctx, storeID, &import_sstpb.MultiIngestRequest{})

br/pkg/restore/log_client/client.go

@@ -293,22 +293,20 @@ func (rc *LogClient) rewriteRulesFor(sst SSTs, rules *restoreutils.RewriteRules)
 	return rules, nil
 }
 
-func (rc *LogClient) RestoreSSTFiles(
+// CollectSSTFileSets collects all items from the iterator in advance to avoid blocking during restoration.
+// This approach ensures that we have all necessary data ready for processing,
+// preventing any potential delays caused by waiting for the iterator to yield more items.
+func (rc *LogClient) CollectSSTFileSets(
 	ctx context.Context,
 	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,
-	// preventing any potential delays caused by waiting for the iterator to yield more items.
+) (restore.BatchBackupFileSet, int64, error) {
+	backupFileSets := make(restore.BatchBackupFileSet, 0, 8)
+	totalKVs := int64(0)
 	start := time.Now()
 	for r := compactionsIter.TryNext(ctx); !r.Finished; r = compactionsIter.TryNext(ctx) {
 		if r.Err != nil {
-			return r.Err
+			return nil, 0, r.Err
 		}
 		i := r.Item
 
@@ -323,16 +321,35 @@ func (rc *LogClient) RestoreSSTFiles(
 		}
 		newRules, err := rc.rewriteRulesFor(i, rewriteRules)
 		if err != nil {
-			return err
+			return nil, 0, err
 		}
 
+		sstFiles := i.GetSSTs()
+		for _, f := range sstFiles {
+			totalKVs += int64(f.TotalKvs)
+		}
 		set := restore.BackupFileSet{
 			TableID:      i.TableID(),
-			SSTFiles:     i.GetSSTs(),
+			SSTFiles:     sstFiles,
 			RewriteRules: newRules,
 		}
 		backupFileSets = append(backupFileSets, set)
 	}
+	log.Info("[Compacted SST Restore] Collected SST files",
+		zap.Int("sst-file-count", len(backupFileSets)),
+		zap.Int64("total-kvs", totalKVs),
+		zap.Duration("iterate-take", time.Since(start)))
+	return backupFileSets, totalKVs, nil
+}
+
+// RestoreSSTFileSets restores pre-collected SST file sets.
+func (rc *LogClient) RestoreSSTFileSets(
+	ctx context.Context,
+	backupFileSets restore.BatchBackupFileSet,
+	importModeSwitcher *restore.ImportModeSwitcher,
+	onProgress func(int64),
+) error {
+	begin := time.Now()
 	if len(backupFileSets) == 0 {
 		log.Info("[Compacted SST Restore] No SST files found for restoration.")
 		return nil
@@ -349,10 +366,9 @@ func (rc *LogClient) RestoreSSTFiles(
 	}()
 
 	log.Info("[Compacted SST Restore] Start to restore SST files",
-		zap.Int("sst-file-count", len(backupFileSets)), zap.Duration("iterate-take", time.Since(start)))
-	start = time.Now()
+		zap.Int("sst-file-count", len(backupFileSets)))
 	defer func() {
-		log.Info("[Compacted SST Restore] Restore SST files finished", zap.Duration("restore-take", time.Since(start)))
+		log.Info("[Compacted SST Restore] Restore SST files finished", zap.Duration("restore-take", time.Since(begin)))
 	}()
 
 	// To optimize performance and minimize cross-region downloads,
@@ -524,6 +540,7 @@ func (rc *LogClient) InitClients(
 	sstCheckpointMetaManager checkpoint.SnapshotMetaManagerT,
 	concurrency uint,
 	concurrencyPerStore uint,
+	retainLatestMVCCVersion bool,
 ) error {
 	stores, err := conn.GetAllTiKVStoresWithRetry(ctx, rc.pdClient, util.SkipTiFlash)
 	if err != nil {
@@ -556,6 +573,11 @@ func (rc *LogClient) InitClients(
 	createCallBacks = append(createCallBacks, func(importer *snapclient.SnapFileImporter) error {
 		return importer.CheckBatchDownloadSupport(ctx, stores)
 	})
+	if retainLatestMVCCVersion {
+		createCallBacks = append(createCallBacks, func(importer *snapclient.SnapFileImporter) error {
+			return importer.CheckBatchDownloadLatestMVCCSupport(ctx, stores)
+		})
+	}
 	if rc.rateLimit != 0 {
 		createCallBack, closeCallBack := snapclient.SetSpeedLimitCallbacks(ctx, rc.pdClient, sstWorkerPool, rc.rateLimit)
 		createCallBacks = append(createCallBacks, createCallBack)
@@ -564,7 +586,8 @@ func (rc *LogClient) InitClients(
 
 	opt := snapclient.NewSnapFileImporterOptions(
 		rc.cipher, metaClient, importCli, backend,
-		snapclient.RewriteModeKeyspace, stores, concurrencyPerStore, rc.regionScanConcurrency, createCallBacks, closeCallBacks,
+		snapclient.RewriteModeKeyspace, stores, concurrencyPerStore, rc.regionScanConcurrency,
+		retainLatestMVCCVersion, createCallBacks, closeCallBacks,
 	)
 	snapFileImporter, err := snapclient.NewSnapFileImporter(
 		ctx, rc.dom.Store().GetCodec().GetAPIVersion(), snapclient.TiDBCompacted, opt)

br/pkg/restore/log_client/client_test.go

@@ -1798,6 +1798,52 @@ func (m *mockCompactedStrategy) ShouldSplit() bool {
 	return m.AccumulateCount%m.expectSplitCount == 0
 }
 
+func TestCollectSSTFileSets(t *testing.T) {
+	ctx := context.Background()
+	client := logclient.TEST_NewLogClient(123, 1, 2, 1, nil, nil)
+	rules := map[int64]*utils.RewriteRules{
+		1: {
+			Data: []*import_sstpb.RewriteRule{
+				{
+					OldKeyPrefix: tablecodec.GenTableRecordPrefix(1),
+					NewKeyPrefix: tablecodec.GenTableRecordPrefix(100),
+				},
+			},
+		},
+	}
+
+	t.Run("skips ssts without rewrite rule", func(t *testing.T) {
+		fileSets, totalKVs, err := client.CollectSSTFileSets(ctx, iter.FromSlice([]logclient.SSTs{
+			fakeSubCompactionWithOneSst(1, 100, 16*units.MiB, 100),
+			fakeSubCompactionWithOneSst(3, 100, 16*units.MiB, 300),
+		}), rules)
+		require.NoError(t, err)
+		require.Equal(t, int64(100), totalKVs)
+		require.Len(t, fileSets, 1)
+		require.Equal(t, int64(1), fileSets[0].TableID)
+		require.Len(t, fileSets[0].SSTFiles, 1)
+	})
+
+	t.Run("counts only ssts that still need restore", func(t *testing.T) {
+		var skippedKVs uint64
+		strategy := logclient.NewCompactedFileSplitStrategy(rules, map[string]struct{}{
+			"1:200": {},
+		}, func(kvCount, _ uint64) {
+			skippedKVs += kvCount
+		})
+		ssts := fakeSubCompactionWithMultiSsts(1, 200, 32*units.MiB, 200)
+		require.False(t, strategy.ShouldSkip(ssts))
+
+		fileSets, totalKVs, err := client.CollectSSTFileSets(ctx, iter.FromSlice([]logclient.SSTs{ssts}), rules)
+		require.NoError(t, err)
+		require.Equal(t, uint64(200), skippedKVs)
+		require.Equal(t, int64(200), totalKVs)
+		require.Len(t, fileSets, 1)
+		require.Len(t, fileSets[0].SSTFiles, 1)
+		require.Equal(t, "1:201", fileSets[0].SSTFiles[0].Name)
+	})
+}
+
 func TestCompactedSplitStrategy(t *testing.T) {
 	ctx := context.Background()