sessionctx: support bytes format and percentage format for tidb_serve…

…r_memory_limit (#38655)

ref #37816

sessionctx/variable/sysvar.go

@@ -757,27 +757,24 @@ var defaultSysVars = []*SysVar{
 			return nil
 		},
 	},
-	{Scope: ScopeGlobal, Name: TiDBServerMemoryLimit, Value: strconv.FormatUint(DefTiDBServerMemoryLimit, 10), Type: TypeUnsigned, MinValue: 0, MaxValue: math.MaxUint64,
+	{Scope: ScopeGlobal, Name: TiDBServerMemoryLimit, Value: DefTiDBServerMemoryLimit, Type: TypeStr,
 		GetGlobal: func(_ context.Context, s *SessionVars) (string, error) {
-			return memory.ServerMemoryLimit.String(), nil
+			return memory.ServerMemoryLimitOriginText.Load(), nil
 		},
 		Validation: func(s *SessionVars, normalizedValue string, originalValue string, scope ScopeFlag) (string, error) {
-			intVal, err := strconv.ParseUint(normalizedValue, 10, 64)
+			_, str, err := parseMemoryLimit(s, normalizedValue, originalValue)
 			if err != nil {
 				return "", err
 			}
-			if intVal > 0 && intVal < (512<<20) { // 512 MB
-				s.StmtCtx.AppendWarning(ErrTruncatedWrongValue.GenWithStackByArgs(TiDBServerMemoryLimit, originalValue))
-				intVal = 512 << 20
-			}
-			return strconv.FormatUint(intVal, 10), nil
+			return str, nil
 		},
 		SetGlobal: func(_ context.Context, s *SessionVars, val string) error {
-			intVal, err := strconv.ParseUint(val, 10, 64)
+			bt, str, err := parseMemoryLimit(s, val, val)
 			if err != nil {
 				return err
 			}
-			memory.ServerMemoryLimit.Store(intVal)
+			memory.ServerMemoryLimitOriginText.Store(str)
+			memory.ServerMemoryLimit.Store(bt)
 			gctuner.GlobalMemoryLimitTuner.UpdateMemoryLimit()
 			return nil
 		},

sessionctx/variable/sysvar_test.go

@@ -23,9 +23,11 @@ import (
 	"sync/atomic"
 	"testing"
 
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/parser/terror"
+	"github.com/pingcap/tidb/util/memory"
 	"github.com/stretchr/testify/require"
 )
 
@@ -771,21 +773,21 @@ func TestTiDBServerMemoryLimit(t *testing.T) {
 	// Test tidb_server_memory_limit
 	serverMemoryLimit := GetSysVar(TiDBServerMemoryLimit)
 	// Check default value
-	require.Equal(t, serverMemoryLimit.Value, strconv.FormatUint(DefTiDBServerMemoryLimit, 10))
+	require.Equal(t, serverMemoryLimit.Value, DefTiDBServerMemoryLimit)
 
 	// MinValue is 512 MB
 	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, strconv.FormatUint(100*mb, 10))
 	require.NoError(t, err)
 	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimit)
 	require.NoError(t, err)
-	require.Equal(t, strconv.FormatUint(512*mb, 10), val)
+	require.Equal(t, "512MB", val)
 
 	// Test Close
 	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, strconv.FormatUint(0, 10))
 	require.NoError(t, err)
 	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimit)
 	require.NoError(t, err)
-	require.Equal(t, strconv.FormatUint(0, 10), val)
+	require.Equal(t, "0", val)
 
 	// Test MaxValue
 	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, strconv.FormatUint(math.MaxUint64, 10))
@@ -834,3 +836,117 @@ func TestTiDBServerMemoryLimit(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, strconv.FormatUint(200*mb, 10), val)
 }
+
+func TestTiDBServerMemoryLimit2(t *testing.T) {
+	vars := NewSessionVars(nil)
+	mock := NewMockGlobalAccessor4Tests()
+	mock.SessionVars = vars
+	vars.GlobalVarsAccessor = mock
+	var (
+		//mb  uint64 = 1 << 20
+		err error
+		val string
+	)
+	// Test tidb_server_memory_limit
+	serverMemoryLimit := GetSysVar(TiDBServerMemoryLimit)
+	// Check default value
+	require.Equal(t, serverMemoryLimit.Value, DefTiDBServerMemoryLimit)
+
+	total := memory.GetMemTotalIgnoreErr()
+	if total > 0 {
+		// Can use percentage format when TiDB can obtain physical memory
+		// Test Percentage Format
+		err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "1%")
+		require.NoError(t, err)
+		val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimit)
+		require.NoError(t, err)
+		if total/100 > uint64(512<<20) {
+			require.Equal(t, memory.ServerMemoryLimit.Load(), total/100)
+			require.Equal(t, "1%", val)
+		} else {
+			require.Equal(t, memory.ServerMemoryLimit.Load(), uint64(512<<20))
+			require.Equal(t, "512MB", val)
+		}
+
+		err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "0%")
+		require.Error(t, err)
+		err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "100%")
+		require.Error(t, err)
+
+		err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "75%")
+		require.NoError(t, err)
+		val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimit)
+		require.NoError(t, err)
+		require.Equal(t, "75%", val)
+		require.Equal(t, memory.ServerMemoryLimit.Load(), total/100*75)
+	}
+	// Test can't obtain physical memory
+	require.Nil(t, failpoint.Enable("github.com/pingcap/tidb/util/memory/GetMemTotalError", `return(true)`))
+	require.Error(t, mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "75%"))
+	require.Nil(t, failpoint.Disable("github.com/pingcap/tidb/util/memory/GetMemTotalError"))
+
+	// Test byteSize format
+	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "1234")
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimit)
+	require.NoError(t, err)
+	require.Equal(t, memory.ServerMemoryLimit.Load(), uint64(512<<20))
+	require.Equal(t, "512MB", val)
+
+	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "1234567890123")
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimit)
+	require.NoError(t, err)
+	require.Equal(t, memory.ServerMemoryLimit.Load(), uint64(1234567890123))
+	require.Equal(t, "1234567890123", val)
+
+	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "10KB")
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimit)
+	require.NoError(t, err)
+	require.Equal(t, memory.ServerMemoryLimit.Load(), uint64(512<<20))
+	require.Equal(t, "512MB", val)
+
+	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "12345678KB")
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimit)
+	require.NoError(t, err)
+	require.Equal(t, memory.ServerMemoryLimit.Load(), uint64(12345678<<10))
+	require.Equal(t, "12345678KB", val)
+
+	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "10MB")
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimit)
+	require.NoError(t, err)
+	require.Equal(t, memory.ServerMemoryLimit.Load(), uint64(512<<20))
+	require.Equal(t, "512MB", val)
+
+	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "700MB")
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimit)
+	require.NoError(t, err)
+	require.Equal(t, memory.ServerMemoryLimit.Load(), uint64(700<<20))
+	require.Equal(t, "700MB", val)
+
+	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "20GB")
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimit)
+	require.NoError(t, err)
+	require.Equal(t, memory.ServerMemoryLimit.Load(), uint64(20<<30))
+	require.Equal(t, "20GB", val)
+
+	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "2TB")
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimit)
+	require.NoError(t, err)
+	require.Equal(t, memory.ServerMemoryLimit.Load(), uint64(2<<40))
+	require.Equal(t, "2TB", val)
+
+	// Test error
+	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "123aaa123")
+	require.Error(t, err)
+	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "700MBaa")
+	require.Error(t, err)
+	err = mock.SetGlobalSysVar(context.Background(), TiDBServerMemoryLimit, "a700MB")
+	require.Error(t, err)
+}

sessionctx/variable/tidb_vars.go

@@ -21,7 +21,6 @@ import (
 	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/sessionctx/variable/featuretag/concurrencyddl"
-	"github.com/pingcap/tidb/util/mathutil"
 	"github.com/pingcap/tidb/util/memory"
 	"github.com/pingcap/tidb/util/paging"
 	"github.com/pingcap/tidb/util/size"
@@ -1136,7 +1135,7 @@ var (
 
 	// DefTiDBServerMemoryLimit indicates the default value of TiDBServerMemoryLimit(TotalMem * 80%).
 	// It should be a const and shouldn't be modified after tidb is started.
-	DefTiDBServerMemoryLimit = mathutil.Max(memory.GetMemTotalIgnoreErr()/10*8, 512<<20)
+	DefTiDBServerMemoryLimit = serverMemoryLimitDefaultValue()
 	GOGCTunerThreshold       = atomic.NewFloat64(DefTiDBGOGCTunerThreshold)
 )
 
@@ -1162,3 +1161,11 @@ var (
 	// GetExternalTimestamp is the func registered by staleread to get externaltimestamp from pd
 	GetExternalTimestamp func(ctx context.Context) (uint64, error)
 )
+
+func serverMemoryLimitDefaultValue() string {
+	total, err := memory.MemTotal()
+	if err == nil && total != 0 {
+		return "80%"
+	}
+	return "0"
+}

sessionctx/variable/varsutil.go

@@ -16,6 +16,7 @@ package variable
 
 import (
 	"fmt"
+	"io"
 	"strconv"
 	"strings"
 	"sync/atomic"
@@ -27,6 +28,7 @@ import (
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/collate"
+	"github.com/pingcap/tidb/util/memory"
 	"github.com/pingcap/tidb/util/timeutil"
 	"github.com/tikv/client-go/v2/oracle"
 	"golang.org/x/exp/slices"
@@ -382,6 +384,66 @@ func parseTimeZone(s string) (*time.Location, error) {
 	return nil, ErrUnknownTimeZone.GenWithStackByArgs(s)
 }
 
+func parseMemoryLimit(s *SessionVars, normalizedValue string, originalValue string) (byteSize uint64, normalizedStr string, err error) {
+	defer func() {
+		if err == nil && byteSize > 0 && byteSize < (512<<20) {
+			s.StmtCtx.AppendWarning(ErrTruncatedWrongValue.GenWithStackByArgs(TiDBServerMemoryLimit, originalValue))
+			byteSize = 512 << 20
+			normalizedStr = "512MB"
+		}
+	}()
+
+	// 1. Try parse percentage format: x%
+	if total := memory.GetMemTotalIgnoreErr(); total != 0 {
+		perc, str := parsePercentage(normalizedValue)
+		if perc != 0 {
+			intVal := total / 100 * perc
+			return intVal, str, nil
+		}
+	}
+
+	// 2. Try parse byteSize format: xKB/MB/GB/TB or byte number
+	bt, str := parseByteSize(normalizedValue)
+	if str != "" {
+		return bt, str, nil
+	}
+
+	return 0, "", ErrTruncatedWrongValue.GenWithStackByArgs(TiDBServerMemoryLimit, originalValue)
+}
+
+func parsePercentage(s string) (percentage uint64, normalizedStr string) {
+	defer func() {
+		if percentage == 0 || percentage >= 100 {
+			percentage, normalizedStr = 0, ""
+		}
+	}()
+	var endString string
+	if n, err := fmt.Sscanf(s, "%d%%%s", &percentage, &endString); n == 1 && err == io.EOF {
+		return percentage, fmt.Sprintf("%d%%", percentage)
+	}
+	return 0, ""
+}
+
+func parseByteSize(s string) (byteSize uint64, normalizedStr string) {
+	var endString string
+	if n, err := fmt.Sscanf(s, "%d%s", &byteSize, &endString); n == 1 && err == io.EOF {
+		return byteSize, fmt.Sprintf("%d", byteSize)
+	}
+	if n, err := fmt.Sscanf(s, "%dKB%s", &byteSize, &endString); n == 1 && err == io.EOF {
+		return byteSize << 10, fmt.Sprintf("%dKB", byteSize)
+	}
+	if n, err := fmt.Sscanf(s, "%dMB%s", &byteSize, &endString); n == 1 && err == io.EOF {
+		return byteSize << 20, fmt.Sprintf("%dMB", byteSize)
+	}
+	if n, err := fmt.Sscanf(s, "%dGB%s", &byteSize, &endString); n == 1 && err == io.EOF {
+		return byteSize << 30, fmt.Sprintf("%dGB", byteSize)
+	}
+	if n, err := fmt.Sscanf(s, "%dTB%s", &byteSize, &endString); n == 1 && err == io.EOF {
+		return byteSize << 40, fmt.Sprintf("%dTB", byteSize)
+	}
+	return 0, ""
+}
+
 func setSnapshotTS(s *SessionVars, sVal string) error {
 	if sVal == "" {
 		s.SnapshotTS = 0

util/memory/meminfo.go

@@ -18,6 +18,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/parser/terror"
 	"github.com/pingcap/tidb/util/cgroup"
 	"github.com/pingcap/tidb/util/mathutil"
@@ -33,6 +34,11 @@ var MemUsed func() (uint64, error)
 // GetMemTotalIgnoreErr returns the total amount of RAM on this system/container. If error occurs, return 0.
 func GetMemTotalIgnoreErr() uint64 {
 	if memTotal, err := MemTotal(); err == nil {
+		failpoint.Inject("GetMemTotalError", func(val failpoint.Value) {
+			if val, ok := val.(bool); val && ok {
+				memTotal = 0
+			}
+		})
 		return memTotal
 	}
 	return 0

util/memory/tracker.go

@@ -33,6 +33,7 @@ const TrackMemWhenExceeds = 104857600 // 100MB
 
 // Process global variables for memory limit.
 var (
+	ServerMemoryLimitOriginText  = atomicutil.NewString("0")
 	ServerMemoryLimit            = atomicutil.NewUint64(0)
 	ServerMemoryLimitSessMinSize = atomicutil.NewUint64(128 << 20)