*: global memory controller should not kill session whose mem less than limit_sess_min_size (#42803)

executor/adapter.go

@@ -734,12 +734,7 @@ func (a *ExecStmt) handleNoDelay(ctx context.Context, e Executor, isPessimistic
 		// done in the `defer` function. If the rs is not nil, the detachment will be done in
 		// `rs.Close` in `handleStmt`
 		if handled && sc != nil && rs == nil {
-			if sc.MemTracker != nil {
-				sc.MemTracker.Detach()
-			}
-			if sc.DiskTracker != nil {
-				sc.DiskTracker.Detach()
-			}
+			sc.DetachMemDiskTracker()
 		}
 	}()
 
@@ -1417,15 +1412,7 @@ func (a *ExecStmt) FinishExecuteStmt(txnTS uint64, err error, hasMoreResults boo
 func (a *ExecStmt) CloseRecordSet(txnStartTS uint64, lastErr error) {
 	a.FinishExecuteStmt(txnStartTS, lastErr, false)
 	a.logAudit()
-	// Detach the Memory and disk tracker for the previous stmtCtx from GlobalMemoryUsageTracker and GlobalDiskUsageTracker
-	if stmtCtx := a.Ctx.GetSessionVars().StmtCtx; stmtCtx != nil {
-		if stmtCtx.DiskTracker != nil {
-			stmtCtx.DiskTracker.Detach()
-		}
-		if stmtCtx.MemTracker != nil {
-			stmtCtx.MemTracker.Detach()
-		}
-	}
+	a.Ctx.GetSessionVars().StmtCtx.DetachMemDiskTracker()
 }
 
 // LogSlowQuery is used to print the slow query in the log files.

executor/issuetest/BUILD.bazel

@@ -15,10 +15,12 @@ go_test(
         "//parser/auth",
         "//parser/charset",
         "//parser/mysql",
+        "//session",
         "//sessionctx/variable",
         "//statistics",
         "//testkit",
         "//util",
+        "//util/memory",
         "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_stretchr_testify//require",
         "@com_github_tikv_client_go_v2//tikv",

executor/issuetest/executor_issue_test.go

@@ -20,17 +20,20 @@ import (
 	"math/rand"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/parser/auth"
 	"github.com/pingcap/tidb/parser/charset"
 	"github.com/pingcap/tidb/parser/mysql"
+	"github.com/pingcap/tidb/session"
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/statistics"
 	"github.com/pingcap/tidb/testkit"
 	"github.com/pingcap/tidb/util"
+	"github.com/pingcap/tidb/util/memory"
 	"github.com/stretchr/testify/require"
 )
 
@@ -1348,3 +1351,46 @@ func TestIssue40158(t *testing.T) {
 	tk.MustExec("insert into t1 values (1, null);")
 	tk.MustQuery("select * from t1 where c1 is null and _id < 1;").Check(testkit.Rows())
 }
+
+func TestIssue42662(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.Session().GetSessionVars().ConnectionID = 12345
+	tk.Session().GetSessionVars().MemTracker = memory.NewTracker(memory.LabelForSession, -1)
+	tk.Session().GetSessionVars().MemTracker.SessionID = 12345
+	tk.Session().GetSessionVars().MemTracker.IsRootTrackerOfSess = true
+
+	sm := &testkit.MockSessionManager{
+		PS: []*util.ProcessInfo{tk.Session().ShowProcess()},
+	}
+	sm.Conn = make(map[uint64]session.Session)
+	sm.Conn[tk.Session().GetSessionVars().ConnectionID] = tk.Session()
+	dom.ServerMemoryLimitHandle().SetSessionManager(sm)
+	go dom.ServerMemoryLimitHandle().Run()
+
+	tk.MustExec("use test")
+	tk.MustQuery("select connection_id()").Check(testkit.Rows("12345"))
+	tk.MustExec("drop table if exists t1, t2")
+	tk.MustExec("create table t1 (a int, b int, c int)")
+	tk.MustExec("create table t2 (a int, b int, c int)")
+	tk.MustExec("insert into t1 values (1, 1, 1), (1, 2, 2), (2, 1, 3), (2, 2, 4)")
+	tk.MustExec("insert into t2 values (1, 1, 1), (1, 2, 2), (2, 1, 3), (2, 2, 4)")
+	// set tidb_server_memory_limit to 1.6GB, tidb_server_memory_limit_sess_min_size to 128MB
+	tk.MustExec("set global tidb_server_memory_limit='1600MB'")
+	tk.MustExec("set global tidb_server_memory_limit_sess_min_size=128*1024*1024")
+	tk.MustExec("set global tidb_mem_oom_action = 'cancel'")
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/executor/issue42662_1", `return(true)`))
+	// tk.Session() should be marked as MemoryTop1Tracker but not killed.
+	tk.MustQuery("select /*+ hash_join(t1)*/ * from t1 join t2 on t1.a = t2.a and t1.b = t2.b")
+
+	// try to trigger the kill top1 logic
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/util/servermemorylimit/issue42662_2", `return(true)`))
+	time.Sleep(1 * time.Second)
+
+	// no error should be returned
+	tk.MustQuery("select count(*) from t1")
+	tk.MustQuery("select count(*) from t1")
+
+	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/executor/issue42662_1"))
+	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/util/servermemorylimit/issue42662_2"))
+}

executor/join.go

@@ -310,6 +310,15 @@ func (w *buildWorker) fetchBuildSideRows(ctx context.Context, chkCh chan<- *chun
 			return
 		}
 	})
+	failpoint.Inject("issue42662_1", func(val failpoint.Value) {
+		if val.(bool) {
+			if w.hashJoinCtx.sessCtx.GetSessionVars().ConnectionID != 0 {
+				// consume 170MB memory, this sql should be tracked into MemoryTop1Tracker
+				w.hashJoinCtx.memTracker.Consume(170 * 1024 * 1024)
+			}
+			return
+		}
+	})
 	sessVars := w.hashJoinCtx.sessCtx.GetSessionVars()
 	for {
 		if w.hashJoinCtx.finished.Load() {

server/conn.go

@@ -2121,12 +2121,20 @@ func (cc *clientConn) handleStmt(ctx context.Context, stmt ast.StmtNode, warns [
 	cc.audit(plugin.Starting)
 	rs, err := cc.ctx.ExecuteStmt(ctx, stmt)
 	reg.End()
-	// The session tracker detachment from global tracker is solved in the `rs.Close` in most cases.
-	// If the rs is nil, the detachment will be done in the `handleNoDelay`.
+	// - If rs is not nil, the statement tracker detachment from session tracker
+	//   is done in the `rs.Close` in most cases.
+	// - If the rs is nil and err is not nil, the detachment will be done in
+	//   the `handleNoDelay`.
 	if rs != nil {
 		defer terror.Call(rs.Close)
 	}
 	if err != nil {
+		// If error is returned during the planner phase or the executor.Open
+		// phase, the rs will be nil, and StmtCtx.MemTracker StmtCtx.DiskTracker
+		// will not be detached. We need to detach them manually.
+		if sv := cc.ctx.GetSessionVars(); sv != nil && sv.StmtCtx != nil {
+			sv.StmtCtx.DetachMemDiskTracker()
+		}
 		return true, err
 	}
 

sessionctx/stmtctx/stmtctx.go

@@ -1114,6 +1114,19 @@ func (sc *StatementContext) UseDynamicPartitionPrune() bool {
 	return sc.UseDynamicPruneMode
 }
 
+// DetachMemDiskTracker detaches the memory and disk tracker from the sessionTracker.
+func (sc *StatementContext) DetachMemDiskTracker() {
+	if sc == nil {
+		return
+	}
+	if sc.MemTracker != nil {
+		sc.MemTracker.Detach()
+	}
+	if sc.DiskTracker != nil {
+		sc.DiskTracker.Detach()
+	}
+}
+
 // CopTasksDetails collects some useful information of cop-tasks during execution.
 type CopTasksDetails struct {
 	NumCopTasks int

testkit/mocksessionmanager.go

@@ -33,7 +33,7 @@ type MockSessionManager struct {
 	SerID   uint64
 	TxnInfo []*txninfo.TxnInfo
 	Dom     *domain.Domain
-	conn    map[uint64]session.Session
+	Conn    map[uint64]session.Session
 	mu      sync.Mutex
 }
 
@@ -44,8 +44,8 @@ func (msm *MockSessionManager) ShowTxnList() []*txninfo.TxnInfo {
 	if len(msm.TxnInfo) > 0 {
 		return msm.TxnInfo
 	}
-	rs := make([]*txninfo.TxnInfo, 0, len(msm.conn))
-	for _, se := range msm.conn {
+	rs := make([]*txninfo.TxnInfo, 0, len(msm.Conn))
+	for _, se := range msm.Conn {
 		info := se.TxnInfo()
 		if info != nil {
 			rs = append(rs, info)
@@ -66,7 +66,7 @@ func (msm *MockSessionManager) ShowProcessList() map[uint64]*util.ProcessInfo {
 		return ret
 	}
 	msm.mu.Lock()
-	for connID, pi := range msm.conn {
+	for connID, pi := range msm.Conn {
 		ret[connID] = pi.ShowProcess()
 	}
 	msm.mu.Unlock()
@@ -89,7 +89,7 @@ func (msm *MockSessionManager) GetProcessInfo(id uint64) (*util.ProcessInfo, boo
 	}
 	msm.mu.Lock()
 	defer msm.mu.Unlock()
-	if sess := msm.conn[id]; sess != nil {
+	if sess := msm.Conn[id]; sess != nil {
 		return sess.ShowProcess(), true
 	}
 	if msm.Dom != nil {
@@ -130,7 +130,7 @@ func (*MockSessionManager) GetInternalSessionStartTSList() []uint64 {
 
 // KillNonFlashbackClusterConn implement SessionManager interface.
 func (msm *MockSessionManager) KillNonFlashbackClusterConn() {
-	for _, se := range msm.conn {
+	for _, se := range msm.Conn {
 		processInfo := se.ShowProcess()
 		ddl, ok := processInfo.StmtCtx.GetPlan().(*core.DDL)
 		if !ok {
@@ -148,7 +148,7 @@ func (msm *MockSessionManager) KillNonFlashbackClusterConn() {
 // CheckOldRunningTxn implement SessionManager interface.
 func (msm *MockSessionManager) CheckOldRunningTxn(job2ver map[int64]int64, job2ids map[int64]string) {
 	msm.mu.Lock()
-	for _, se := range msm.conn {
+	for _, se := range msm.Conn {
 		session.RemoveLockDDLJobs(se, job2ver, job2ids)
 	}
 	msm.mu.Unlock()

testkit/testkit.go

@@ -69,10 +69,10 @@ func NewTestKit(t testing.TB, store kv.Storage) *TestKit {
 		mockSm, ok := sm.(*MockSessionManager)
 		if ok {
 			mockSm.mu.Lock()
-			if mockSm.conn == nil {
-				mockSm.conn = make(map[uint64]session.Session)
+			if mockSm.Conn == nil {
+				mockSm.Conn = make(map[uint64]session.Session)
 			}
-			mockSm.conn[tk.session.GetSessionVars().ConnectionID] = tk.session
+			mockSm.Conn[tk.session.GetSessionVars().ConnectionID] = tk.session
 			mockSm.mu.Unlock()
 		}
 		tk.session.SetSessionManager(sm)

util/memory/tracker.go

@@ -302,6 +302,9 @@ func (t *Tracker) AttachTo(parent *Tracker) {
 
 // Detach de-attach the tracker child from its parent, then set its parent property as nil
 func (t *Tracker) Detach() {
+	if t == nil {
+		return
+	}
 	parent := t.getParent()
 	if parent == nil {
 		return
@@ -446,6 +449,7 @@ func (t *Tracker) Consume(bs int64) {
 				currentAction = nextAction
 				nextAction = currentAction.GetFallback()
 			}
+			logutil.BgLogger().Warn("global memory controller, lastAction", zap.Any("action", currentAction))
 			currentAction.Action(tracker)
 		}
 	}
@@ -471,6 +475,7 @@ func (t *Tracker) Consume(bs int64) {
 				}
 				oldTracker = MemUsageTop1Tracker.Load()
 			}
+			logutil.BgLogger().Error("global memory controller, update the Top1 session", zap.Int64("memUsage", memUsage), zap.Uint64("conn", sessionRootTracker.SessionID), zap.Uint64("limitSessMinSize", limitSessMinSize))
 		}
 	}
 

util/servermemorylimit/BUILD.bazel

@@ -11,6 +11,7 @@ go_library(
         "//util",
         "//util/logutil",
         "//util/memory",
+        "@com_github_pingcap_failpoint//:failpoint",
         "@org_uber_go_atomic//:atomic",
         "@org_uber_go_zap//:zap",
     ],

util/servermemorylimit/servermemorylimit.go

@@ -21,6 +21,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util"
@@ -88,6 +89,15 @@ type sessionToBeKilled struct {
 	lastLogTime   time.Time
 }
 
+func (s *sessionToBeKilled) reset() {
+	s.isKilling = false
+	s.sqlStartTime = time.Time{}
+	s.sessionID = 0
+	s.sessionTracker = nil
+	s.killStartTime = time.Time{}
+	s.lastLogTime = time.Time{}
+}
+
 func killSessIfNeeded(s *sessionToBeKilled, bt uint64, sm util.SessionManager) {
 	if s.isKilling {
 		if info, ok := sm.GetProcessInfo(s.sessionID); ok {
@@ -104,7 +114,7 @@ func killSessIfNeeded(s *sessionToBeKilled, bt uint64, sm util.SessionManager) {
 				return
 			}
 		}
-		s.isKilling = false
+		s.reset()
 		IsKilling.Store(false)
 		memory.MemUsageTop1Tracker.CompareAndSwap(s.sessionTracker, nil)
 		//nolint: all_revive,revive
@@ -115,14 +125,25 @@ func killSessIfNeeded(s *sessionToBeKilled, bt uint64, sm util.SessionManager) {
 	if bt == 0 {
 		return
 	}
+	failpoint.Inject("issue42662_2", func(val failpoint.Value) {
+		if val.(bool) {
+			bt = 1
+		}
+	})
 	instanceStats := memory.ReadMemStats()
 	if instanceStats.HeapInuse > MemoryMaxUsed.Load() {
 		MemoryMaxUsed.Store(instanceStats.HeapInuse)
 	}
+	limitSessMinSize := memory.ServerMemoryLimitSessMinSize.Load()
 	if instanceStats.HeapInuse > bt {
 		t := memory.MemUsageTop1Tracker.Load()
 		if t != nil {
-			if info, ok := sm.GetProcessInfo(t.SessionID); ok {
+			memUsage := t.BytesConsumed()
+			// If the memory usage of the top1 session is less than tidb_server_memory_limit_sess_min_size, we do not need to kill it.
+			if uint64(memUsage) < limitSessMinSize {
+				memory.MemUsageTop1Tracker.CompareAndSwap(t, nil)
+				t = nil
+			} else if info, ok := sm.GetProcessInfo(t.SessionID); ok {
 				logutil.BgLogger().Warn("global memory controller tries to kill the top1 memory consumer",
 					zap.Uint64("connID", info.ID),
 					zap.String("sql digest", info.Digest),
@@ -146,6 +167,17 @@ func killSessIfNeeded(s *sessionToBeKilled, bt uint64, sm util.SessionManager) {
 				s.killStartTime = time.Now()
 			}
 		}
+		// If no one larger than tidb_server_memory_limit_sess_min_size is found, we will not kill any one.
+		if t == nil {
+			if s.lastLogTime.IsZero() {
+				s.lastLogTime = time.Now()
+			}
+			if time.Since(s.lastLogTime) < 5*time.Second {
+				return
+			}
+			logutil.BgLogger().Warn("global memory controller tries to kill the top1 memory consumer, but no one larger than tidb_server_memory_limit_sess_min_size is found", zap.Uint64("tidb_server_memory_limit_sess_min_size", limitSessMinSize))
+			s.lastLogTime = time.Now()
+		}
 	}
 }