*: support global memory control for tidb

executor/executor.go

@@ -1942,6 +1942,7 @@ func ResetContextOfStmt(ctx sessionctx.Context, s ast.StmtNode) (err error) {
 	} else {
 		sc.InitMemTracker(memory.LabelForSQLText, vars.MemQuotaQuery)
 		sc.MemTracker.AttachToGlobalTracker(GlobalMemoryUsageTracker)
+		sc.MemTracker.IsSession, sc.MemTracker.SessionID = true, vars.ConnectionID
 	}
 
 	sc.InitDiskTracker(memory.LabelForSQLText, -1)

executor/executor_test.go

@@ -6001,3 +6001,66 @@ func TestBinaryStrNumericOperator(t *testing.T) {
 	// there should be no warning.
 	tk.MustQuery("show warnings").Check(testkit.Rows())
 }
+
+func TestGlobalMemoryControl(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+
+	tk0 := testkit.NewTestKit(t, store)
+	tk0.MustExec("set global tidb_mem_oom_action = 'cancel'")
+	tk0.MustExec("set global tidb_server_memory_quota = 512 << 20")
+	tk0.MustExec("set global tidb_server_memory_limit_sess_min_size = 128")
+
+	tk1 := testkit.NewTestKit(t, store)
+	tracker1 := tk1.Session().GetSessionVars().StmtCtx.MemTracker
+
+	tk2 := testkit.NewTestKit(t, store)
+	tracker2 := tk2.Session().GetSessionVars().StmtCtx.MemTracker
+
+	tk3 := testkit.NewTestKit(t, store)
+	tracker3 := tk3.Session().GetSessionVars().StmtCtx.MemTracker
+
+	sm := &testkit.MockSessionManager{
+		PS: []*util.ProcessInfo{tk1.Session().ShowProcess(), tk2.Session().ShowProcess(), tk3.Session().ShowProcess()},
+	}
+	dom.ExpensiveQueryHandle().SetSessionManager(sm)
+	go dom.ExpensiveQueryHandle().Run()
+
+	tracker1.Consume(100 << 20) // 100 MB
+	tracker2.Consume(200 << 20) // 200 MB
+	tracker3.Consume(300 << 20) // 300 MB
+
+	test := make([]int, 128<<20) // Keep 1GB HeapInUse
+	time.Sleep(500 * time.Millisecond)
+
+	// Kill Top1
+	require.False(t, tracker1.IsKilled.Load())
+	require.False(t, tracker2.IsKilled.Load())
+	require.True(t, tracker3.IsKilled.Load())
+	require.Equal(t, memory.MemUsageTop1Tracker.Load(), tracker3)
+	util.WithRecovery( // Next Consume() will panic and cancel the SQL
+		func() {
+			tracker3.Consume(1)
+		}, func(r interface{}) {
+			require.True(t, strings.Contains(r.(string), "Out Of Memory Quota!"))
+		})
+	tracker2.Consume(300 << 20) // Sum 500MB, Not Panic, Waiting t3 cancel finish.
+	time.Sleep(500 * time.Millisecond)
+	require.False(t, tracker2.IsKilled.Load())
+	// Kill Finished
+	tracker3.Consume(-(300 << 20))
+	// Simulated SQL is Canceled and the time is updated
+	sm.PSMu.Lock()
+	ps := *sm.PS[2]
+	ps.Time = time.Now()
+	sm.PS[2] = &ps
+	sm.PSMu.Unlock()
+	time.Sleep(500 * time.Millisecond)
+	// Kill the Next SQL
+	util.WithRecovery( // Next Consume() will panic and cancel the SQL
+		func() {
+			tracker2.Consume(1)
+		}, func(r interface{}) {
+			require.True(t, strings.Contains(r.(string), "Out Of Memory Quota!"))
+		})
+	require.Equal(t, test[0], 0) // Keep 1GB HeapInUse
+}

sessionctx/variable/sysvar.go

@@ -697,6 +697,54 @@ var defaultSysVars = []*SysVar{
 			return nil
 		},
 	},
+	{Scope: ScopeGlobal, Name: TiDBServerMemoryQuota, Value: strconv.FormatUint(DefTiDBServerMemoryQuota, 10), Type: TypeUnsigned, MinValue: 0, MaxValue: math.MaxUint64,
+		GetGlobal: func(s *SessionVars) (string, error) {
+			return memory.ServerMemoryQuota.String(), nil
+		},
+		Validation: func(s *SessionVars, normalizedValue string, originalValue string, scope ScopeFlag) (string, error) {
+			intVal, err := strconv.ParseUint(normalizedValue, 10, 64)
+			if err != nil {
+				return "", err
+			}
+			if intVal > 0 && intVal < (512<<20) { // 512 MB
+				s.StmtCtx.AppendWarning(ErrTruncatedWrongValue.GenWithStackByArgs(TiDBServerMemoryQuota, originalValue))
+				intVal = 512 << 20
+			}
+			return strconv.FormatUint(intVal, 10), nil
+		},
+		SetGlobal: func(s *SessionVars, val string) error {
+			intVal, err := strconv.ParseUint(val, 10, 64)
+			if err != nil {
+				return err
+			}
+			memory.ServerMemoryQuota.Store(intVal)
+			return nil
+		},
+	},
+	{Scope: ScopeGlobal, Name: TiDBServerMemoryLimitSessMinSize, Value: strconv.FormatUint(DefTiDBServerMemoryLimitSessMinSize, 10), Type: TypeUnsigned, MinValue: 0, MaxValue: math.MaxUint64,
+		GetGlobal: func(s *SessionVars) (string, error) {
+			return memory.ServerMemoryLimitSessMinSize.String(), nil
+		},
+		Validation: func(s *SessionVars, normalizedValue string, originalValue string, scope ScopeFlag) (string, error) {
+			intVal, err := strconv.ParseUint(normalizedValue, 10, 64)
+			if err != nil {
+				return "", err
+			}
+			if intVal > 0 && intVal < 128 { // 128 Bytes
+				s.StmtCtx.AppendWarning(ErrTruncatedWrongValue.GenWithStackByArgs(TiDBServerMemoryLimitSessMinSize, originalValue))
+				intVal = 128
+			}
+			return strconv.FormatUint(intVal, 10), nil
+		},
+		SetGlobal: func(s *SessionVars, val string) error {
+			intVal, err := strconv.ParseUint(val, 10, 64)
+			if err != nil {
+				return err
+			}
+			memory.ServerMemoryLimitSessMinSize.Store(intVal)
+			return nil
+		},
+	},
 	{Scope: ScopeGlobal, Name: TiDBEnableColumnTracking, Value: BoolToOnOff(DefTiDBEnableColumnTracking), Type: TypeBool, GetGlobal: func(s *SessionVars) (string, error) {
 		return BoolToOnOff(EnableColumnTracking.Load()), nil
 	}, SetGlobal: func(s *SessionVars, val string) error {

sessionctx/variable/sysvar_test.go

@@ -17,6 +17,7 @@ package variable
 import (
 	"encoding/json"
 	"fmt"
+	"math"
 	"strconv"
 	"sync/atomic"
 	"testing"
@@ -742,3 +743,80 @@ func TestSetTIDBDiskQuota(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, strconv.FormatInt(pb, 10), val)
 }
+
+func TestTiDBServerMemoryQuota(t *testing.T) {
+	vars := NewSessionVars()
+	mock := NewMockGlobalAccessor4Tests()
+	mock.SessionVars = vars
+	vars.GlobalVarsAccessor = mock
+	var (
+		mb  uint64 = 1 << 20
+		err error
+		val string
+	)
+	// Test tidb_server_memory_quota
+	serverMemoryQuota := GetSysVar(TiDBServerMemoryQuota)
+	// Check default value
+	require.Equal(t, serverMemoryQuota.Value, strconv.FormatUint(DefTiDBServerMemoryQuota, 10))
+
+	// MinValue is 512 MB
+	err = mock.SetGlobalSysVar(TiDBServerMemoryQuota, strconv.FormatUint(100*mb, 10))
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryQuota)
+	require.NoError(t, err)
+	require.Equal(t, strconv.FormatUint(512*mb, 10), val)
+
+	// Test Close
+	err = mock.SetGlobalSysVar(TiDBServerMemoryQuota, strconv.FormatUint(0, 10))
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryQuota)
+	require.NoError(t, err)
+	require.Equal(t, strconv.FormatUint(0, 10), val)
+
+	// Test MaxValue
+	err = mock.SetGlobalSysVar(TiDBServerMemoryQuota, strconv.FormatUint(math.MaxUint64, 10))
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryQuota)
+	require.NoError(t, err)
+	require.Equal(t, strconv.FormatUint(math.MaxUint64, 10), val)
+
+	// Test Normal Value
+	err = mock.SetGlobalSysVar(TiDBServerMemoryQuota, strconv.FormatUint(1024*mb, 10))
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryQuota)
+	require.NoError(t, err)
+	require.Equal(t, strconv.FormatUint(1024*mb, 10), val)
+
+	// Test tidb_server_memory_limit_sess_min_size
+	serverMemoryLimitSessMinSize := GetSysVar(TiDBServerMemoryLimitSessMinSize)
+	// Check default value
+	require.Equal(t, serverMemoryLimitSessMinSize.Value, strconv.FormatUint(DefTiDBServerMemoryLimitSessMinSize, 10))
+
+	// MinValue is 128 Bytes
+	err = mock.SetGlobalSysVar(TiDBServerMemoryLimitSessMinSize, strconv.FormatUint(100, 10))
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimitSessMinSize)
+	require.NoError(t, err)
+	require.Equal(t, strconv.FormatUint(128, 10), val)
+
+	// Test Close
+	err = mock.SetGlobalSysVar(TiDBServerMemoryLimitSessMinSize, strconv.FormatUint(0, 10))
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimitSessMinSize)
+	require.NoError(t, err)
+	require.Equal(t, strconv.FormatUint(0, 10), val)
+
+	// Test MaxValue
+	err = mock.SetGlobalSysVar(TiDBServerMemoryLimitSessMinSize, strconv.FormatUint(math.MaxUint64, 10))
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimitSessMinSize)
+	require.NoError(t, err)
+	require.Equal(t, strconv.FormatUint(math.MaxUint64, 10), val)
+
+	// Test Normal Value
+	err = mock.SetGlobalSysVar(TiDBServerMemoryLimitSessMinSize, strconv.FormatUint(200*mb, 10))
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(TiDBServerMemoryLimitSessMinSize)
+	require.NoError(t, err)
+	require.Equal(t, strconv.FormatUint(200*mb, 10), val)
+}

sessionctx/variable/tidb_vars.go

@@ -20,6 +20,8 @@ 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"
 	"go.uber.org/atomic"
 )
@@ -822,6 +824,10 @@ const (
 	TiDBDDLEnableFastReorg = "tidb_ddl_enable_fast_reorg"
 	// TiDBDDLDiskQuota used to set disk quota for lightning add index.
 	TiDBDDLDiskQuota = "tidb_ddl_disk_quota"
+	// TiDBServerMemoryQuota indicates the memory limit of the tidb-server instance.
+	TiDBServerMemoryQuota = "tidb_server_memory_quota"
+	// TiDBServerMemoryLimitSessMinSize indicates the minimum memory usage of the session that can be controlled.
+	TiDBServerMemoryLimitSessMinSize = "tidb_server_memory_limit_sess_min_size"
 )
 
 // TiDB intentional limits
@@ -1054,6 +1060,7 @@ const (
 	DefTiDBForeignKeyChecks                         = false
 	DefTiDBOptRangeMaxSize                          = 0
 	DefTiDBCostModelVer                             = 1
+	DefTiDBServerMemoryLimitSessMinSize             = 128 << 20
 )
 
 // Process global variables.
@@ -1103,6 +1110,10 @@ var (
 	DDLDiskQuota = atomic.NewUint64(DefTiDBDDLDiskQuota)
 	// EnableForeignKey indicates whether to enable foreign key feature.
 	EnableForeignKey = atomic.NewBool(false)
+
+	// DefTiDBServerMemoryQuota indicates the default value of TiDBServerMemoryQuota(TotalMem * 80%).
+	// It should be a const and shouldn't be modified after tidb is started.
+	DefTiDBServerMemoryQuota = mathutil.Max(memory.GetMemTotalIgnoreErr()/10*8, 512<<20)
 )
 
 var (

testkit/mocksessionmanager.go

@@ -26,6 +26,7 @@ import (
 // MockSessionManager is a mocked session manager which is used for test.
 type MockSessionManager struct {
 	PS      []*util.ProcessInfo
+	PSMu    sync.RWMutex
 	SerID   uint64
 	TxnInfo []*txninfo.TxnInfo
 	conn    map[uint64]session.Session
@@ -39,6 +40,8 @@ func (msm *MockSessionManager) ShowTxnList() []*txninfo.TxnInfo {
 
 // ShowProcessList implements the SessionManager.ShowProcessList interface.
 func (msm *MockSessionManager) ShowProcessList() map[uint64]*util.ProcessInfo {
+	msm.PSMu.RLock()
+	defer msm.PSMu.RUnlock()
 	ret := make(map[uint64]*util.ProcessInfo)
 	for _, item := range msm.PS {
 		ret[item.ID] = item
@@ -48,6 +51,8 @@ func (msm *MockSessionManager) ShowProcessList() map[uint64]*util.ProcessInfo {
 
 // GetProcessInfo implements the SessionManager.GetProcessInfo interface.
 func (msm *MockSessionManager) GetProcessInfo(id uint64) (*util.ProcessInfo, bool) {
+	msm.PSMu.RLock()
+	defer msm.PSMu.RUnlock()
 	for _, item := range msm.PS {
 		if item.ID == id {
 			return item, true

util/expensivequery/expensivequery.go

@@ -26,6 +26,7 @@ import (
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/logutil"
+	"github.com/pingcap/tidb/util/memory"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 )
@@ -57,6 +58,7 @@ func (eqh *Handle) Run() {
 	defer ticker.Stop()
 	sm := eqh.sm.Load().(util.SessionManager)
 	record := &memoryUsageAlarm{}
+	serverMemoryQuota := &serverMemoryQuota{}
 	for {
 		select {
 		case <-ticker.C:
@@ -91,6 +93,7 @@ func (eqh *Handle) Run() {
 			if record.err == nil {
 				record.alarm4ExcessiveMemUsage(sm)
 			}
+			serverMemoryQuota.CheckQuotaAndKill(memory.ServerMemoryQuota.Load(), sm)
 		case <-eqh.exitCh:
 			return
 		}

util/expensivequery/server_memory_quota.go

@@ -0,0 +1,58 @@
+// Copyright 2022 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package expensivequery
+
+import (
+	"runtime"
+	"time"
+
+	"github.com/pingcap/tidb/util"
+	"github.com/pingcap/tidb/util/memory"
+)
+
+type serverMemoryQuota struct {
+	sqlStartTime time.Time
+	sessionID    uint64
+}
+
+func (s *serverMemoryQuota) CheckQuotaAndKill(bt uint64, sm util.SessionManager) {
+	if s.sessionID != 0 {
+		if info, ok := sm.GetProcessInfo(s.sessionID); ok {
+			if info.Time == s.sqlStartTime {
+				return // Wait killing finished.
+			}
+		}
+		s.sessionID = 0
+		s.sqlStartTime = time.Time{}
+		//nolint: all_revive,revive
+		runtime.GC()
+	}
+
+	if bt == 0 {
+		return
+	}
+	instanceStats := &runtime.MemStats{}
+	runtime.ReadMemStats(instanceStats)
+	if instanceStats.HeapInuse > bt {
+		t := memory.MemUsageTop1Tracker.Load()
+		if t != nil && uint64(t.BytesConsumed()) > memory.ServerMemoryLimitSessMinSize.Load() {
+			if info, ok := sm.GetProcessInfo(t.SessionID); ok {
+				s.sessionID = t.SessionID
+				s.sqlStartTime = info.Time
+				t.IsKilled.Store(true)
+			}
+		}
+	}
+}

util/memory/meminfo.go

@@ -32,6 +32,14 @@ var MemTotal func() (uint64, error)
 // MemUsed returns the total used amount of RAM on this system
 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 {
+		return memTotal
+	}
+	return 0
+}
+
 // MemTotalNormal returns the total amount of RAM on this system in non-container environment.
 func MemTotalNormal() (uint64, error) {
 	total, t := memLimit.get()