From 570faa656ee3af8cb24364b8bb1ba6608b32c695 Mon Sep 17 00:00:00 2001
From: xuhuaiyu <391585975@qq.com>
Date: Wed, 2 Nov 2022 18:23:35 +0800
Subject: [PATCH 1/3] util, executor: remove OOMAction after a query finishing

---
 executor/executor_test.go | 20 ++++++++++++++++++++
 util/memory/tracker.go    |  9 +++++++++
 2 files changed, 29 insertions(+)

diff --git a/executor/executor_test.go b/executor/executor_test.go
index f03ea400f6ab0..e67fc31639570 100644
--- a/executor/executor_test.go
+++ b/executor/executor_test.go
@@ -6176,3 +6176,23 @@ func TestCompileOutOfMemoryQuota(t *testing.T) {
 	err := tk.ExecToErr("select t.a, t1.a from t use index(idx), t1 use index(idx) where t.a = t1.a")
 	require.Contains(t, err.Error(), "Out Of Memory Quota!")
 }
+
+func TestSessionRootTrackerDetach(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	defer tk.MustExec("set global tidb_mem_oom_action = DEFAULT")
+	tk.MustExec("set global tidb_mem_oom_action='CANCEL'")
+	tk.MustExec("use test")
+	tk.MustExec("create table t(a int, b int, index idx(a))")
+	tk.MustExec("create table t1(a int, c int, index idx(a))")
+	tk.MustExec("set tidb_mem_quota_query=10")
+	err := tk.ExecToErr("select /*+hash_join(t1)*/ t.a, t1.a from t use index(idx), t1 use index(idx) where t.a = t1.a")
+	require.Contains(t, err.Error(), "Out Of Memory Quota!")
+	tk.MustExec("set tidb_mem_quota_query=1000")
+	rs, err := tk.Exec("select /*+hash_join(t1)*/ t.a, t1.a from t use index(idx), t1 use index(idx) where t.a = t1.a")
+	require.NoError(t, err)
+	require.NotNil(t, tk.Session().GetSessionVars().MemTracker.GetFallbackForTest(false))
+	err = rs.Close()
+	require.NoError(t, err)
+	require.Nil(t, tk.Session().GetSessionVars().MemTracker.GetFallbackForTest(false))
+}
diff --git a/util/memory/tracker.go b/util/memory/tracker.go
index 83a34f6731450..4a97e4096b1f0 100644
--- a/util/memory/tracker.go
+++ b/util/memory/tracker.go
@@ -305,6 +305,15 @@ func (t *Tracker) Detach() {
 		t.DetachFromGlobalTracker()
 		return
 	}
+	if parent.IsRootTrackerOfSess {
+		parent.actionMuForHardLimit.Lock()
+		parent.actionMuForHardLimit.actionOnExceed = nil
+		parent.actionMuForHardLimit.Unlock()
+
+		parent.actionMuForSoftLimit.Lock()
+		parent.actionMuForSoftLimit.actionOnExceed = nil
+		parent.actionMuForSoftLimit.Unlock()
+	}
 	parent.remove(t)
 	t.mu.Lock()
 	defer t.mu.Unlock()

From fdf2202774245029ba8a56d9c1a52e4f11c696f9 Mon Sep 17 00:00:00 2001
From: xuhuaiyu <391585975@qq.com>
Date: Thu, 3 Nov 2022 11:05:38 +0800
Subject: [PATCH 2/3] fix unit test

---
 executor/executor_test.go            | 3 +++
 session/session_test/session_test.go | 5 ++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/executor/executor_test.go b/executor/executor_test.go
index e67fc31639570..af5608a78aefa 100644
--- a/executor/executor_test.go
+++ b/executor/executor_test.go
@@ -6072,12 +6072,15 @@ func TestGlobalMemoryControl(t *testing.T) {
 
 	tk1 := testkit.NewTestKit(t, store)
 	tracker1 := tk1.Session().GetSessionVars().MemTracker
+	tracker1.FallbackOldAndSetNewAction(&memory.PanicOnExceed{})
 
 	tk2 := testkit.NewTestKit(t, store)
 	tracker2 := tk2.Session().GetSessionVars().MemTracker
+	tracker2.FallbackOldAndSetNewAction(&memory.PanicOnExceed{})
 
 	tk3 := testkit.NewTestKit(t, store)
 	tracker3 := tk3.Session().GetSessionVars().MemTracker
+	tracker3.FallbackOldAndSetNewAction(&memory.PanicOnExceed{})
 
 	sm := &testkit.MockSessionManager{
 		PS: []*util.ProcessInfo{tk1.Session().ShowProcess(), tk2.Session().ShowProcess(), tk3.Session().ShowProcess()},
diff --git a/session/session_test/session_test.go b/session/session_test/session_test.go
index 79c1fd00d6a2c..c0c37cf3296a8 100644
--- a/session/session_test/session_test.go
+++ b/session/session_test/session_test.go
@@ -2112,7 +2112,8 @@ func TestSetEnableRateLimitAction(t *testing.T) {
 	result.Check(testkit.Rows("1"))
 	tk.MustExec("use test")
 	tk.MustExec("create table tmp123(id int)")
-	tk.MustQuery("select * from tmp123;")
+	rs, err := tk.Exec("select * from tmp123;")
+	require.NoError(t, err)
 	haveRateLimitAction := false
 	action := tk.Session().GetSessionVars().MemTracker.GetFallbackForTest(false)
 	for ; action != nil; action = action.GetFallback() {
@@ -2122,6 +2123,8 @@ func TestSetEnableRateLimitAction(t *testing.T) {
 		}
 	}
 	require.True(t, haveRateLimitAction)
+	err = rs.Close()
+	require.NoError(t, err)
 
 	// assert set sys variable
 	tk.MustExec("set global tidb_enable_rate_limit_action= '0';")

From b7fd630445f9ac98e45a842b7f751b34759c0d35 Mon Sep 17 00:00:00 2001
From: xuhuaiyu <391585975@qq.com>
Date: Thu, 3 Nov 2022 13:32:28 +0800
Subject: [PATCH 3/3] address comment

---
 util/memory/tracker.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/util/memory/tracker.go b/util/memory/tracker.go
index f4f690343f327..6f4bad7713e5f 100644
--- a/util/memory/tracker.go
+++ b/util/memory/tracker.go
@@ -314,6 +314,7 @@ func (t *Tracker) Detach() {
 		parent.actionMuForSoftLimit.Lock()
 		parent.actionMuForSoftLimit.actionOnExceed = nil
 		parent.actionMuForSoftLimit.Unlock()
+		parent.NeedKill.Store(false)
 	}
 	parent.remove(t)
 	t.mu.Lock()