diff --git a/br/cmd/br/BUILD.bazel b/br/cmd/br/BUILD.bazel index c6b4e1d646e11..b82ecbd2bc8dc 100644 --- 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 index df0395fa1d719..02939c78b4676 100644 --- 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 index 5fa5b1439a2ec..bdd6f7282ee62 100644 --- 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)) +}