br: set memory limit (#53793)

close #53777
pingcap · Aug 5, 2024 · 2c47cd5 · 2c47cd5
1 parent f3e153a
commit 2c47cd5
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 1 deletion.
diff --git a/br/cmd/br/BUILD.bazel b/br/cmd/br/BUILD.bazel
@@ -37,9 +37,11 @@ go_library(
         "//pkg/util",
         "//pkg/util/gctuner",
         "//pkg/util/logutil",
+        "//pkg/util/mathutil",
         "//pkg/util/memory",
         "//pkg/util/metricsutil",
         "//pkg/util/redact",
+        "//pkg/util/size",
         "@com_github_gogo_protobuf//proto",
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_kvproto//pkg/brpb",
@@ -63,5 +65,8 @@ go_test(
     srcs = ["main_test.go"],
     embed = [":br_lib"],
     flaky = True,
-    deps = ["@org_uber_go_goleak//:goleak"],
+    deps = [
+        "@com_github_stretchr_testify//require",
+        "@org_uber_go_goleak//:goleak",
+    ],
 )
diff --git a/br/cmd/br/cmd.go b/br/cmd/br/cmd.go
@@ -5,8 +5,10 @@ package main
 import (
 	"context"
 	"fmt"
+	"math"
 	"os"
 	"path/filepath"
+	"runtime/debug"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -21,9 +23,12 @@ import (
 	"github.com/pingcap/tidb/pkg/config"
 	tidbutils "github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/logutil"
+	"github.com/pingcap/tidb/pkg/util/mathutil"
 	"github.com/pingcap/tidb/pkg/util/memory"
 	"github.com/pingcap/tidb/pkg/util/redact"
+	"github.com/pingcap/tidb/pkg/util/size"
 	"github.com/spf13/cobra"
+	"go.uber.org/zap"
 )
 
 var (
@@ -107,6 +112,25 @@ func AddFlags(cmd *cobra.Command) {
 	_ = cmd.PersistentFlags().MarkHidden(FlagRedactLog)
 }
 
+const quarterGiB uint64 = 256 * size.MB
+const halfGiB uint64 = 512 * size.MB
+const fourGiB uint64 = 4 * size.GB
+
+func calculateMemoryLimit(memleft uint64) uint64 {
+	// memreserved = f(memleft) = 512MB * memleft / (memleft + 4GB)
+	//  * f(0) = 0
+	//  * f(4GB) = 256MB
+	//  * f(+inf) -> 512MB
+	memreserved := halfGiB / (1 + fourGiB/(memleft|1))
+	// 0     memused          memtotal-memreserved  memtotal
+	// +--------+--------------------+----------------+
+	//          ^            br mem upper limit
+	//          +--------------------^
+	//             GOMEMLIMIT range
+	memlimit := memleft - memreserved
+	return memlimit
+}
+
 // Init initializes BR cli.
 func Init(cmd *cobra.Command) (err error) {
 	initOnce.Do(func() {
@@ -162,6 +186,34 @@ func Init(cmd *cobra.Command) (err error) {
 		}
 		log.ReplaceGlobals(lg, p)
 		memory.InitMemoryHook()
+		if debug.SetMemoryLimit(-1) == math.MaxInt64 {
+			memtotal, e := memory.MemTotal()
+			if e != nil {
+				err = e
+				return
+			}
+			memused, e := memory.MemUsed()
+			if e != nil {
+				err = e
+				return
+			}
+			if memused >= memtotal {
+				log.Warn("failed to obtain memory size, skip setting memory limit",
+					zap.Uint64("memused", memused), zap.Uint64("memtotal", memtotal))
+			} else {
+				memleft := memtotal - memused
+				memlimit := calculateMemoryLimit(memleft)
+				// BR command needs 256 MiB at least, if the left memory is less than 256 MiB,
+				// the memory limit cannot limit anyway and then finally OOM.
+				memlimit = mathutil.Max(memlimit, quarterGiB)
+				log.Info("calculate the rest memory",
+					zap.Uint64("memtotal", memtotal), zap.Uint64("memused", memused), zap.Uint64("memlimit", memlimit))
+				// No need to set memory limit because the left memory is sufficient.
+				if memlimit < uint64(math.MaxInt64) {
+					debug.SetMemoryLimit(int64(memlimit))
+				}
+			}
+		}
 
 		redactLog, e := cmd.Flags().GetBool(FlagRedactLog)
 		if e != nil {

diff --git a/br/cmd/br/main_test.go b/br/cmd/br/main_test.go
@@ -21,6 +21,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
 )
 
@@ -75,3 +76,18 @@ func TestRunMain(*testing.T) {
 
 	<-waitCh
 }
+
+func TestCalculateMemoryLimit(t *testing.T) {
+	// f(0 Byte) = 0 Byte
+	require.Equal(t, uint64(0), calculateMemoryLimit(0))
+	// f(100 KB) = 87.5 KB
+	require.Equal(t, uint64(89600), calculateMemoryLimit(100*1024))
+	// f(100 MB) = 87.5 MB
+	require.Equal(t, uint64(91763188), calculateMemoryLimit(100*1024*1024))
+	// f(3.99 GB) = 3.74 GB
+	require.Equal(t, uint64(4026531839), calculateMemoryLimit(4*1024*1024*1024-1))
+	// f(4 GB) = 3.5 GB
+	require.Equal(t, uint64(3758096384), calculateMemoryLimit(4*1024*1024*1024))
+	// f(32 GB) = 31.5 GB
+	require.Equal(t, uint64(33822867456), calculateMemoryLimit(32*1024*1024*1024))
+}