br: Adjust restore concurrency and tikv config

br/pkg/checkpoint/checkpoint_test.go

@@ -107,12 +107,13 @@ func testCheckpointMetaForRestore(
 	require.Equal(t, checkpointMetaForSnapshotRestore.RestoredTS, checkpointMetaForSnapshotRestore2.RestoredTS)
 
 	checkpointMetaForLogRestore := &checkpoint.CheckpointMetadataForLogRestore{
-		UpstreamClusterID: 123,
-		RestoredTS:        222,
-		StartTS:           111,
-		RewriteTS:         333,
-		GcRatio:           "1.0",
-		TiFlashItems:      map[int64]model.TiFlashReplicaInfo{1: {Count: 1}},
+		UpstreamClusterID:        123,
+		RestoredTS:               222,
+		StartTS:                  111,
+		RewriteTS:                333,
+		GcRatio:                  "1.0",
+		RocksDBMaxBackgroundJobs: "8",
+		TiFlashItems:             map[int64]model.TiFlashReplicaInfo{1: {Count: 1}},
 	}
 
 	err = logMetaManager.SaveCheckpointMetadata(ctx, checkpointMetaForLogRestore)
@@ -124,6 +125,7 @@ func testCheckpointMetaForRestore(
 	require.Equal(t, checkpointMetaForLogRestore.StartTS, checkpointMetaForLogRestore2.StartTS)
 	require.Equal(t, checkpointMetaForLogRestore.RewriteTS, checkpointMetaForLogRestore2.RewriteTS)
 	require.Equal(t, checkpointMetaForLogRestore.GcRatio, checkpointMetaForLogRestore2.GcRatio)
+	require.Equal(t, checkpointMetaForLogRestore.RocksDBMaxBackgroundJobs, checkpointMetaForLogRestore2.RocksDBMaxBackgroundJobs)
 	require.Equal(t, checkpointMetaForLogRestore.TiFlashItems, checkpointMetaForLogRestore2.TiFlashItems)
 
 	exists, err := logMetaManager.ExistsCheckpointProgress(ctx)
@@ -144,6 +146,7 @@ func testCheckpointMetaForRestore(
 	require.Equal(t, uint64(111), taskInfo.Metadata.StartTS)
 	require.Equal(t, uint64(333), taskInfo.Metadata.RewriteTS)
 	require.Equal(t, "1.0", taskInfo.Metadata.GcRatio)
+	require.Equal(t, "8", taskInfo.Metadata.RocksDBMaxBackgroundJobs)
 	require.Equal(t, true, taskInfo.HasSnapshotMetadata)
 	require.Equal(t, checkpoint.InLogRestoreAndIdMapPersisted, taskInfo.Progress)
 

br/pkg/checkpoint/log_restore.go

@@ -134,12 +134,14 @@ func AppendRangeForLogRestore(
 }
 
 type CheckpointMetadataForLogRestore struct {
-	UpstreamClusterID uint64 `json:"upstream-cluster-id"`
-	RestoreStartTS    uint64 `json:"restore-start-ts"`
-	RestoredTS        uint64 `json:"restored-ts"`
-	StartTS           uint64 `json:"start-ts"`
-	RewriteTS         uint64 `json:"rewrite-ts"`
-	GcRatio           string `json:"gc-ratio"`
+	UpstreamClusterID        uint64 `json:"upstream-cluster-id"`
+	RestoreStartTS           uint64 `json:"restore-start-ts"`
+	RestoredTS               uint64 `json:"restored-ts"`
+	StartTS                  uint64 `json:"start-ts"`
+	RewriteTS                uint64 `json:"rewrite-ts"`
+	GcRatio                  string `json:"gc-ratio"`
+	RocksDBMaxBackgroundJobs string `json:"rocksdb-max-background-jobs,omitempty"`
+	SnapshotRestoreDataSize  uint64 `json:"snapshot-restore-data-size,omitempty"`
 	// tiflash recorder items with snapshot restore records
 	TiFlashItems map[int64]model.TiFlashReplicaInfo `json:"tiflash-recorder,omitempty"`
 }

br/pkg/conn/conn.go

@@ -50,10 +50,10 @@ const (
 	// DefaultMergeRegionKeyCount is the default region key count, 960000.
 	DefaultMergeRegionKeyCount uint64 = 960000
 
-	// DefaultImportNumGoroutines is the default number of threads for import.
-	// use 128 as default value, which is 8 times of the default value of tidb.
-	// we think is proper for IO-bound cases.
-	DefaultImportNumGoroutines uint = 128
+	// DefaultImportNumGoroutines is the default number of goroutines for restore.
+	DefaultImportNumGoroutines uint = 36
+
+	minRestoreConcurrencyOverImportThreads uint = 4
 )
 
 type VersionCheckerType int
@@ -328,7 +328,8 @@ func (mgr *Mgr) GetCurrentTsFromPD(ctx context.Context) (uint64, error) {
 }
 
 // ProcessTiKVConfigs handle the tikv config for region split size, region split keys, and import goroutines in place.
-// It retrieves the config from all alive tikv stores and returns the minimum values.
+// It retrieves the config from all alive tikv stores, keeps conservative split values,
+// and makes restore concurrency no less than import.num-threads plus a small margin.
 // If retrieving the config fails, it returns the default config values.
 func (mgr *Mgr) ProcessTiKVConfigs(ctx context.Context, cfg *kvconfig.KVConfig, client *http.Client) {
 	mergeRegionSize := cfg.MergeRegionSize
@@ -362,9 +363,8 @@ func (mgr *Mgr) ProcessTiKVConfigs(ctx context.Context, cfg *kvconfig.KVConfig,
 				log.Warn("Failed to parse import num-threads from config", logutil.ShortError(e))
 				return e
 			}
-			// We use 8 times the default value because it's an IO-bound case.
-			if importGoroutines.Value == DefaultImportNumGoroutines || (threads > 0 && threads*8 < importGoroutines.Value) {
-				importGoroutines.Value = threads * 8
+			if threads > 0 {
+				importGoroutines.Value = max(importGoroutines.Value, threads+minRestoreConcurrencyOverImportThreads)
 			}
 		}
 		// replace the value

br/pkg/conn/conn_test.go

@@ -357,8 +357,8 @@ func TestGetMergeRegionSizeAndCount(t *testing.T) {
 			content: []string{
 				"{\"log-level\": \"debug\", \"coprocessor\": {\"region-split-keys\": 1, \"region-split-size\": \"1MiB\"}, \"import\": {\"num-threads\": 6}}",
 			},
-			// the number of import goroutines is 8 times than import.num-threads.
-			importNumGoroutines: 48,
+			// the default value already satisfies import.num-threads + 4.
+			importNumGoroutines: conn.DefaultImportNumGoroutines,
 			// one tikv detected in this case we are not update default size and keys because they are too small.
 			regionSplitSize: 1 * units.MiB,
 			regionSplitKeys: 1,
@@ -379,7 +379,7 @@ func TestGetMergeRegionSizeAndCount(t *testing.T) {
 			content: []string{
 				"{\"log-level\": \"debug\", \"coprocessor\": {\"region-split-keys\": 10000000, \"region-split-size\": \"1GiB\"}, \"import\": {\"num-threads\": 128}}",
 			},
-			importNumGoroutines: 1024,
+			importNumGoroutines: 132,
 			// one tikv detected in this case and we update with new size and keys.
 			regionSplitSize: 1 * units.GiB,
 			regionSplitKeys: 10000000,
@@ -411,8 +411,7 @@ func TestGetMergeRegionSizeAndCount(t *testing.T) {
 				"{\"log-level\": \"debug\", \"coprocessor\": {\"region-split-keys\": 10000000, \"region-split-size\": \"1GiB\"}, \"import\": {\"num-threads\": 128}}",
 				"{\"log-level\": \"debug\", \"coprocessor\": {\"region-split-keys\": 12000000, \"region-split-size\": \"900MiB\"}, \"import\": {\"num-threads\": 12}}",
 			},
-			// two tikv detected in this case and we choose the small one.
-			importNumGoroutines: 96,
+			importNumGoroutines: 132,
 			regionSplitSize:     1 * units.GiB,
 			regionSplitKeys:     10000000,
 		},

br/pkg/restore/log_client/BUILD.bazel

@@ -6,6 +6,7 @@ go_library(
         "batch_meta_processor.go",
         "client.go",
         "compacted_file_strategy.go",
+        "flow_control.go",
         "id_map.go",
         "import.go",
         "import_retry.go",
@@ -59,6 +60,7 @@ go_library(
         "//pkg/util/codec",
         "//pkg/util/redact",
         "//pkg/util/sqlescape",
+        "//pkg/util/sqlexec",
         "@com_github_docker_go_units//:go-units",
         "@com_github_fatih_color//:color",
         "@com_github_gogo_protobuf//proto",

br/pkg/restore/log_client/client.go

@@ -34,6 +34,7 @@ import (
 	"github.com/pingcap/failpoint"
 	backuppb "github.com/pingcap/kvproto/pkg/brpb"
 	"github.com/pingcap/kvproto/pkg/encryptionpb"
+	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/log"
 	"github.com/pingcap/tidb/br/pkg/checkpoint"
 	"github.com/pingcap/tidb/br/pkg/checksum"
@@ -86,6 +87,7 @@ import (
 const MetaKVBatchSize = 64 * 1024 * 1024
 const maxSplitKeysOnce = 10240
 const maxReadMetaKVFilesConcurrency uint = 128
+const defaultTiKVMaxReplicas uint = 3
 
 // rawKVBatchCount specifies the count of entries that the rawkv client puts into TiKV.
 const rawKVBatchCount = 64
@@ -143,6 +145,8 @@ func (l *LogRestoreManager) Close(ctx context.Context) {
 // including concurrency management, checkpoint handling, and file importing(splitting) for efficient log processing.
 type SstRestoreManager struct {
 	restorer         restore.SstRestorer
+	storeCount       uint
+	replicaCount     uint
 	checkpointRunner *checkpoint.CheckpointRunner[checkpoint.RestoreKeyType, checkpoint.RestoreValueType]
 }
 
@@ -366,13 +370,24 @@ func (rc *LogClient) RestoreSSTFileSets(
 	ctx context.Context,
 	backupFileSets restore.BatchBackupFileSet,
 	importModeSwitcher *restore.ImportModeSwitcher,
+	snapshotRestoreDataSize uint64,
+	checkpointCompactedSSTSize uint64,
 	onProgress func(int64),
 ) error {
 	begin := time.Now()
 	if len(backupFileSets) == 0 {
 		log.Info("[Compacted SST Restore] No SST files found for restoration.")
 		return nil
 	}
+	if err := rc.adjustTiKVFlowControlForCompactedSSTRestore(
+		ctx,
+		backupFileSets,
+		snapshotRestoreDataSize,
+		checkpointCompactedSSTSize,
+	); err != nil {
+		return errors.Trace(err)
+	}
+
 	err := importModeSwitcher.GoSwitchToImportMode(ctx)
 	if err != nil {
 		return errors.Trace(err)
@@ -565,6 +580,8 @@ func (rc *LogClient) InitClients(
 	if err != nil {
 		log.Fatal("failed to get stores", zap.Error(err))
 	}
+	liveStoreCount := liveTiKVStoreCount(stores)
+	replicaCount := rc.getMaxReplica(ctx)
 
 	metaClient := split.NewClient(rc.pdClient, rc.pdHTTPClient, rc.tlsConf, maxSplitKeysOnce, len(stores)+1)
 	importCli := importclient.NewImportClient(metaClient, rc.tlsConf, rc.keepaliveConf)
@@ -578,9 +595,10 @@ func (rc *LogClient) InitClients(
 	if err != nil {
 		return errors.Trace(err)
 	}
-	// This poolSize is similar to full restore, as both workflows are comparable.
-	// The poolSize should be greater than concurrencyPerStore multiplied by the number of stores.
-	poolSize := concurrencyPerStore * 32 * uint(len(stores))
+	// Keep the global SST restore pool large enough to avoid starving stores when
+	// queued file sets temporarily point to other TiKVs.
+	const sstRestoreWorkerPoolSizePerStore uint = 7186
+	poolSize := sstRestoreWorkerPoolSizePerStore * uint(len(stores))
 	log.Info("sst restore worker pool", zap.Uint("size", poolSize))
 	sstWorkerPool := tidbutil.NewWorkerPool(poolSize, "sst file")
 
@@ -613,7 +631,7 @@ func (rc *LogClient) InitClients(
 	if err != nil {
 		return errors.Trace(err)
 	}
-	sstRestoreManager := &SstRestoreManager{}
+	sstRestoreManager := &SstRestoreManager{storeCount: liveStoreCount, replicaCount: replicaCount}
 	if sstCheckpointMetaManager != nil {
 		var err error
 		sstRestoreManager.checkpointRunner, err = checkpoint.StartCheckpointRunnerForRestore(ctx, sstCheckpointMetaManager)
@@ -632,6 +650,55 @@ func (rc *LogClient) InitClients(
 	return nil
 }
 
+func liveTiKVStoreCount(stores []*metapb.Store) uint {
+	var count uint
+	for _, store := range stores {
+		if store.GetState() == metapb.StoreState_Up {
+			count++
+		}
+	}
+	return count
+}
+
+func (rc *LogClient) getMaxReplica(ctx context.Context) uint {
+	if rc.pdHTTPClient == nil {
+		return maxReplicaFromReplicateConfig(nil, errors.New("PD HTTP client is not initialized"))
+	}
+	var resp map[string]any
+	var err error
+	err = utils.WithRetry(ctx, func() error {
+		resp, err = rc.pdHTTPClient.GetReplicateConfig(ctx)
+		return err
+	}, utils.NewAggressivePDBackoffStrategy())
+	return maxReplicaFromReplicateConfig(resp, err)
+}
+
+func maxReplicaFromReplicateConfig(resp map[string]any, err error) uint {
+	if err != nil {
+		log.Warn("failed to get max replicas from PD replicate config, use default value",
+			zap.Uint("default-max-replicas", defaultTiKVMaxReplicas),
+			logutil.ShortError(err))
+		return defaultTiKVMaxReplicas
+	}
+
+	const key = "max-replicas"
+	val, ok := resp[key]
+	if !ok {
+		log.Warn("max replicas not found in PD replicate config, use default value",
+			zap.Uint("default-max-replicas", defaultTiKVMaxReplicas),
+			zap.Any("replicate-config", resp))
+		return defaultTiKVMaxReplicas
+	}
+	replicaCount, ok := val.(float64)
+	if !ok || replicaCount <= 0 {
+		log.Warn("invalid max replicas in PD replicate config, use default value",
+			zap.Uint("default-max-replicas", defaultTiKVMaxReplicas),
+			zap.Any("replicate-config", resp))
+		return defaultTiKVMaxReplicas
+	}
+	return uint(replicaCount)
+}
+
 func (rc *LogClient) InitCheckpointMetadataForCompactedSstRestore(
 	ctx context.Context,
 	sstCheckpointMetaManager checkpoint.SnapshotMetaManagerT,
@@ -665,51 +732,60 @@ func (rc *LogClient) LoadOrCreateCheckpointMetadataForLogRestore(
 	ctx context.Context,
 	restoreStartTS, startTS, restoredTS uint64,
 	gcRatio string,
+	rocksDBMaxBackgroundJobs string,
 	tiflashRecorder *tiflashrec.TiFlashRecorder,
 	logCheckpointMetaManager checkpoint.LogMetaManagerT,
-) (string, error) {
+	snapshotRestoreDataSize uint64,
+) (string, string, uint64, error) {
 	rc.useCheckpoint = true
 
 	// if the checkpoint metadata exists in the external storage, the restore is not
 	// for the first time.
 	exists, err := logCheckpointMetaManager.ExistsCheckpointMetadata(ctx)
 	if err != nil {
-		return "", errors.Trace(err)
+		return "", "", 0, errors.Trace(err)
 	}
 	if exists {
 		// load the checkpoint since this is not the first time to restore
 		log.Info("loading existing log restore checkpoint")
 		meta, err := logCheckpointMetaManager.LoadCheckpointMetadata(ctx)
 		if err != nil {
-			return "", errors.Trace(err)
+			return "", "", 0, errors.Trace(err)
 		}
 
-		log.Info("reuse gc ratio from checkpoint metadata", zap.String("old-gc-ratio", gcRatio),
-			zap.String("checkpoint-gc-ratio", meta.GcRatio))
-		return meta.GcRatio, nil
+		if meta.RocksDBMaxBackgroundJobs != "" {
+			rocksDBMaxBackgroundJobs = meta.RocksDBMaxBackgroundJobs
+		}
+		log.Info("reuse TiKV config from checkpoint metadata",
+			zap.String("gc-ratio", meta.GcRatio),
+			zap.String("rocksdb-max-background-jobs", rocksDBMaxBackgroundJobs))
+		return meta.GcRatio, rocksDBMaxBackgroundJobs, meta.SnapshotRestoreDataSize, nil
 	}
 
 	// initialize the checkpoint metadata since it is the first time to restore.
 	var items map[int64]model.TiFlashReplicaInfo
 	if tiflashRecorder != nil {
 		items = tiflashRecorder.GetItems()
 	}
-	log.Info("save gc ratio into checkpoint metadata",
+	log.Info("save TiKV config into checkpoint metadata",
 		zap.Uint64("start-ts", startTS), zap.Uint64("restored-ts", restoredTS), zap.Uint64("rewrite-ts", rc.currentTS),
-		zap.String("gc-ratio", gcRatio), zap.Int("tiflash-item-count", len(items)))
+		zap.String("gc-ratio", gcRatio), zap.String("rocksdb-max-background-jobs", rocksDBMaxBackgroundJobs),
+		zap.Int("tiflash-item-count", len(items)))
 	if err := logCheckpointMetaManager.SaveCheckpointMetadata(ctx, &checkpoint.CheckpointMetadataForLogRestore{
-		UpstreamClusterID: rc.upstreamClusterID,
-		RestoreStartTS:    restoreStartTS,
-		RestoredTS:        restoredTS,
-		StartTS:           startTS,
-		RewriteTS:         rc.currentTS,
-		GcRatio:           gcRatio,
-		TiFlashItems:      items,
+		UpstreamClusterID:        rc.upstreamClusterID,
+		RestoreStartTS:           restoreStartTS,
+		RestoredTS:               restoredTS,
+		StartTS:                  startTS,
+		RewriteTS:                rc.currentTS,
+		GcRatio:                  gcRatio,
+		RocksDBMaxBackgroundJobs: rocksDBMaxBackgroundJobs,
+		SnapshotRestoreDataSize:  snapshotRestoreDataSize,
+		TiFlashItems:             items,
 	}); err != nil {
-		return gcRatio, errors.Trace(err)
+		return gcRatio, rocksDBMaxBackgroundJobs, snapshotRestoreDataSize, errors.Trace(err)
 	}
 
-	return gcRatio, nil
+	return gcRatio, rocksDBMaxBackgroundJobs, snapshotRestoreDataSize, nil
 }
 
 type LockedMigrations struct {