br: support to stop gc memory limit tuner when restore (#51082)

close pingcap/tidb#51078

br/cmd/br/BUILD.bazel

@@ -35,6 +35,7 @@ go_library(
         "//pkg/parser/model",
         "//pkg/session",
         "//pkg/util",
+        "//pkg/util/gctuner",
         "//pkg/util/logutil",
         "//pkg/util/memory",
         "//pkg/util/metricsutil",

br/cmd/br/backup.go

@@ -13,6 +13,7 @@ import (
 	"github.com/pingcap/tidb/br/pkg/version/build"
 	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/session"
+	"github.com/pingcap/tidb/pkg/util/gctuner"
 	"github.com/pingcap/tidb/pkg/util/metricsutil"
 	"github.com/spf13/cobra"
 	"go.uber.org/zap"
@@ -48,6 +49,10 @@ func runBackupCommand(command *cobra.Command, cmdName string) error {
 	// No need to cache the coproceesor result
 	config.GetGlobalConfig().TiKVClient.CoprCache.CapacityMB = 0
 
+	// Disable the memory limit tuner. That's because the server memory is get from TiDB node instead of BR node.
+	gctuner.GlobalMemoryLimitTuner.DisableAdjustMemoryLimit()
+	defer gctuner.GlobalMemoryLimitTuner.EnableAdjustMemoryLimit()
+
 	if err := task.RunBackup(ctx, tidbGlue, cmdName, &cfg); err != nil {
 		log.Error("failed to backup", zap.Error(err))
 		return errors.Trace(err)

br/cmd/br/restore.go

@@ -16,6 +16,7 @@ import (
 	"github.com/pingcap/tidb/br/pkg/version/build"
 	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/session"
+	"github.com/pingcap/tidb/pkg/util/gctuner"
 	"github.com/pingcap/tidb/pkg/util/metricsutil"
 	"github.com/spf13/cobra"
 	"go.uber.org/zap"
@@ -67,6 +68,10 @@ func runRestoreCommand(command *cobra.Command, cmdName string) error {
 	// No need to cache the coproceesor result
 	config.GetGlobalConfig().TiKVClient.CoprCache.CapacityMB = 0
 
+	// Disable the memory limit tuner. That's because the server memory is get from TiDB node instead of BR node.
+	gctuner.GlobalMemoryLimitTuner.DisableAdjustMemoryLimit()
+	defer gctuner.GlobalMemoryLimitTuner.EnableAdjustMemoryLimit()
+
 	if err := task.RunRestore(GetDefaultContext(), tidbGlue, cmdName, &cfg); err != nil {
 		log.Error("failed to restore", zap.Error(err))
 		printWorkaroundOnFullRestoreError(command, err)

br/pkg/metautil/statsfile.go

@@ -71,6 +71,12 @@ func newStatsWriter(
 	}
 }
 
+// flush temporary and clear []byte to make it garbage collected as soon as possible
+func (s *StatsWriter) flushTemporary() ([]byte, error) {
+	defer s.clearTemporary()
+	return proto.Marshal(s.statsFile)
+}
+
 func (s *StatsWriter) clearTemporary() {
 	// clear the temporary variables
 	s.totalSize = 0
@@ -81,7 +87,7 @@ func (s *StatsWriter) clearTemporary() {
 
 func (s *StatsWriter) writeStatsFileAndClear(ctx context.Context, physicalID int64) error {
 	fileName := getStatsFileName(physicalID)
-	content, err := proto.Marshal(s.statsFile)
+	content, err := s.flushTemporary()
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -92,7 +98,7 @@ func (s *StatsWriter) writeStatsFileAndClear(ctx context.Context, physicalID int
 	}
 
 	checksum := sha256.Sum256(content)
-
+	sizeOri := uint64(len(content))
 	encryptedContent, iv, err := Encrypt(content, s.cipher)
 	if err != nil {
 		return errors.Trace(err)
@@ -106,11 +112,9 @@ func (s *StatsWriter) writeStatsFileAndClear(ctx context.Context, physicalID int
 		Name:     fileName,
 		Sha256:   checksum[:],
 		SizeEnc:  uint64(len(encryptedContent)),
-		SizeOri:  uint64(len(content)),
+		SizeOri:  sizeOri,
 		CipherIv: iv,
 	})
-
-	s.clearTemporary()
 	return nil
 }
 
@@ -225,6 +229,9 @@ func downloadStats(
 				if err := json.Unmarshal(block.JsonTable, jsonTable); err != nil {
 					return errors.Trace(err)
 				}
+				// reset the block.JsonTable to nil to make it garbage collected as soon as possible
+				block.JsonTable = nil
+
 				select {
 				case <-ectx.Done():
 					return nil

pkg/util/gctuner/memory_limit_tuner.go

@@ -39,6 +39,10 @@ type memoryLimitTuner struct {
 	serverMemLimitBeforeAdjust   atomicutil.Uint64
 	percentageBeforeAdjust       atomicutil.Float64
 	nextGCTriggeredByMemoryLimit atomicutil.Bool
+
+	// The flag to disable memory limit adjust. There might be many tasks need to activate it in future,
+	// so it is integer type.
+	adjustDisabled atomicutil.Int64
 }
 
 // fallbackPercentage indicates the fallback memory limit percentage when turning.
@@ -55,6 +59,18 @@ func WaitMemoryLimitTunerExitInTest() {
 	}
 }
 
+// DisableAdjustMemoryLimit makes memoryLimitTuner directly return `initGOMemoryLimitValue` when function `calcMemoryLimit` is called.
+func (t *memoryLimitTuner) DisableAdjustMemoryLimit() {
+	t.adjustDisabled.Add(1)
+	debug.SetMemoryLimit(initGOMemoryLimitValue)
+}
+
+// EnableAdjustMemoryLimit makes memoryLimitTuner return an adjusted memory limit when function `calcMemoryLimit` is called.
+func (t *memoryLimitTuner) EnableAdjustMemoryLimit() {
+	t.adjustDisabled.Add(-1)
+	t.UpdateMemoryLimit()
+}
+
 // tuning check the memory nextGC and judge whether this GC is trigger by memory limit.
 // Go runtime ensure that it will be called serially.
 func (t *memoryLimitTuner) tuning() {
@@ -155,7 +171,10 @@ func (t *memoryLimitTuner) UpdateMemoryLimit() {
 	debug.SetMemoryLimit(memoryLimit)
 }
 
-func (*memoryLimitTuner) calcMemoryLimit(percentage float64) int64 {
+func (t *memoryLimitTuner) calcMemoryLimit(percentage float64) int64 {
+	if t.adjustDisabled.Load() > 0 {
+		return initGOMemoryLimitValue
+	}
 	memoryLimit := int64(float64(memory.ServerMemoryLimit.Load()) * percentage) // `tidb_server_memory_limit` * `tidb_server_memory_limit_gc_trigger`
 	if memoryLimit == 0 {
 		memoryLimit = math.MaxInt64

pkg/util/gctuner/memory_limit_tuner_test.go

@@ -230,3 +230,14 @@ func TestIssue48741(t *testing.T) {
 	waitingTunningFinishFn()
 	checkIfMemoryLimitIsModified()
 }
+
+func TestSetMemoryLimit(t *testing.T) {
+	GlobalMemoryLimitTuner.DisableAdjustMemoryLimit()
+	memory.ServerMemoryLimit.Store(1 << 30)   // 1GB
+	GlobalMemoryLimitTuner.SetPercentage(0.8) // 1GB * 80% = 800MB
+	GlobalMemoryLimitTuner.UpdateMemoryLimit()
+	require.Equal(t, initGOMemoryLimitValue, debug.SetMemoryLimit(-1))
+	GlobalMemoryLimitTuner.EnableAdjustMemoryLimit()
+	GlobalMemoryLimitTuner.UpdateMemoryLimit()
+	require.Equal(t, int64(1<<30*80/100), debug.SetMemoryLimit(-1))
+}