From ec6f60e77d2a1fb7ec48a8f38f62d4b590a94836 Mon Sep 17 00:00:00 2001
From: Shenghui Wu <793703860@qq.com>
Date: Tue, 27 Sep 2022 23:17:44 +0800
Subject: [PATCH] executor: ignore spill action when trigger global memory kill
  (#38198)

ref pingcap/tidb#37816
---
 executor/executor_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 util/memory/tracker.go    | 14 +++++++++++++-
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/executor/executor_test.go b/executor/executor_test.go
index 229b66bc229e1..db7a3abbd12ad 100644
--- a/executor/executor_test.go
+++ b/executor/executor_test.go
@@ -24,6 +24,7 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -6126,3 +6127,40 @@ func TestGlobalMemoryControl(t *testing.T) {
 		})
 	require.Equal(t, test[0], 0) // Keep 1GB HeapInUse
 }
+
+func TestGlobalMemoryControl2(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_limit = 1 << 30")
+	tk0.MustExec("set global tidb_server_memory_limit_sess_min_size = 128")
+
+	sm := &testkit.MockSessionManager{
+		PS: []*util.ProcessInfo{tk0.Session().ShowProcess()},
+	}
+	dom.ServerMemoryLimitHandle().SetSessionManager(sm)
+	go dom.ServerMemoryLimitHandle().Run()
+
+	tk0.MustExec("use test")
+	tk0.MustExec("create table t(a int)")
+	tk0.MustExec("insert into t select 1")
+	for i := 1; i <= 8; i++ {
+		tk0.MustExec("insert into t select * from t") // 256 Lines
+	}
+
+	var test []int
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		time.Sleep(100 * time.Millisecond) // Make sure the sql is running.
+		test = make([]int, 128<<20)        // Keep 1GB HeapInuse
+		wg.Done()
+	}()
+	sql := "select * from t t1 join t t2 join t t3 on t1.a=t2.a and t1.a=t3.a order by t1.a;" // Need 500MB
+	require.True(t, strings.Contains(tk0.QueryToErr(sql).Error(), "Out Of Memory Quota!"))
+	require.Equal(t, tk0.Session().GetSessionVars().StmtCtx.DiskTracker.MaxConsumed(), int64(0))
+	wg.Wait()
+	test[0] = 0
+	runtime.GC()
+}
diff --git a/util/memory/tracker.go b/util/memory/tracker.go
index 0077c26b13e5a..c967f47a05509 100644
--- a/util/memory/tracker.go
+++ b/util/memory/tracker.go
@@ -405,10 +405,22 @@ func (t *Tracker) Consume(bs int64) {
 		}
 	}
 
+	tryActionLastOne := func(mu *actionMu, tracker *Tracker) {
+		mu.Lock()
+		defer mu.Unlock()
+		if currentAction := mu.actionOnExceed; currentAction != nil {
+			for nextAction := currentAction.GetFallback(); nextAction != nil; {
+				currentAction = nextAction
+				nextAction = currentAction.GetFallback()
+			}
+			currentAction.Action(tracker)
+		}
+	}
+
 	if bs > 0 && sessionRootTracker != nil {
 		// Kill the Top1 session
 		if sessionRootTracker.NeedKill.Load() {
-			tryAction(&sessionRootTracker.actionMuForHardLimit, sessionRootTracker)
+			tryActionLastOne(&sessionRootTracker.actionMuForHardLimit, sessionRootTracker)
 		}
 		// Update the Top1 session
 		memUsage := sessionRootTracker.BytesConsumed()