executor: change the behavior of mem-quota-query

executor/aggregate.go

@@ -1942,14 +1942,6 @@ func (a *AggSpillDiskAction) Action(t *memory.Tracker) {
 		atomic.StoreUint32(&a.e.inSpillMode, 1)
 		return
 	}
-	if fallback := a.GetFallback(); fallback != nil {
-		fallback.Action(t)
-	}
-}
-
-// GetPriority get the priority of the Action
-func (*AggSpillDiskAction) GetPriority() int64 {
-	return memory.DefSpillPriority
 }
 
 // SetLogHook sets the hook, it does nothing just to form the memory.ActionOnExceed interface.

executor/cte.go

@@ -438,7 +438,7 @@ func setupCTEStorageTracker(tbl cteutil.Storage, ctx sessionctx.Context, parentM
 				actionSpill = tbl.(*cteutil.StorageRC).ActionSpillForTest()
 			}
 		})
-		ctx.GetSessionVars().StmtCtx.MemTracker.FallbackOldAndSetNewAction(actionSpill)
+		ctx.GetSessionVars().StmtCtx.MemTracker.FallbackOldAndSetNewActionForSoftLimit(actionSpill)
 	}
 	return actionSpill
 }

executor/executor.go

@@ -188,11 +188,6 @@ func (a *globalPanicOnExceed) Action(t *memory.Tracker) {
 	panic(msg)
 }
 
-// GetPriority get the priority of the Action
-func (a *globalPanicOnExceed) GetPriority() int64 {
-	return memory.DefPanicPriority
-}
-
 // base returns the baseExecutor of an executor, don't override this method!
 func (e *baseExecutor) base() *baseExecutor {
 	return e

executor/executor_pkg_test.go

@@ -372,7 +372,7 @@ func TestSortSpillDisk(t *testing.T) {
 	err = exec.Close()
 	require.NoError(t, err)
 
-	ctx.GetSessionVars().StmtCtx.MemTracker = memory.NewTracker(-1, 28000)
+	ctx.GetSessionVars().StmtCtx.MemTracker = memory.NewTracker(-1, 24000/4*5)
 	dataSource.prepareChunks()
 	err = exec.Open(tmpCtx)
 	require.NoError(t, err)

executor/executor_test.go

@@ -63,7 +63,6 @@ import (
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/dbterror"
-	"github.com/pingcap/tidb/util/memory"
 	"github.com/pingcap/tidb/util/mock"
 	"github.com/pingcap/tidb/util/rowcodec"
 	"github.com/pingcap/tidb/util/sqlexec"
@@ -3172,7 +3171,7 @@ func TestInvalidDateValueInCreateTable(t *testing.T) {
 	tk.MustExec("drop table if exists t;")
 }
 
-func TestOOMActionPriority(t *testing.T) {
+func TestOOMActionFinishedAndRemoved(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 
 	tk := testkit.NewTestKit(t, store)
@@ -3193,9 +3192,9 @@ func TestOOMActionPriority(t *testing.T) {
 	tk.MustExec("create table t4(a int)")
 	tk.MustExec("insert into t4 values(1)")
 	tk.MustQuery("select * from t0 join t1 join t2 join t3 join t4 order by t0.a").Check(testkit.Rows("1 1 1 1 1"))
-	action := tk.Session().GetSessionVars().StmtCtx.MemTracker.GetFallbackForTest(true)
+	action := tk.Session().GetSessionVars().StmtCtx.MemTracker.GetFallbackForSoftLimitForTest(true)
 	// All actions are finished and removed.
-	require.Equal(t, action.GetPriority(), int64(memory.DefLogPriority))
+	require.Nil(t, action)
 }
 
 func TestTrackAggMemoryUsage(t *testing.T) {

executor/join.go

@@ -1204,7 +1204,7 @@ func (e *HashJoinExec) buildHashTableForList(buildSideResultCh <-chan *chunk.Chu
 				defer actionSpill.(*chunk.SpillDiskAction).WaitForTest()
 			}
 		})
-		e.ctx.GetSessionVars().StmtCtx.MemTracker.FallbackOldAndSetNewAction(actionSpill)
+		e.ctx.GetSessionVars().StmtCtx.MemTracker.FallbackOldAndSetNewActionForSoftLimit(actionSpill)
 	}
 	for chk := range buildSideResultCh {
 		if e.finished.Load().(bool) {

executor/merge_join.go

@@ -100,7 +100,7 @@ func (t *mergeJoinTable) init(exec *MergeJoinExec) {
 					actionSpill = t.rowContainer.ActionSpillForTest()
 				}
 			})
-			exec.ctx.GetSessionVars().StmtCtx.MemTracker.FallbackOldAndSetNewAction(actionSpill)
+			exec.ctx.GetSessionVars().StmtCtx.MemTracker.FallbackOldAndSetNewActionForSoftLimit(actionSpill)
 		}
 		t.memTracker = memory.NewTracker(memory.LabelForInnerTable, -1)
 	} else {

executor/sort.go

@@ -189,7 +189,7 @@ func (e *SortExec) fetchRowChunks(ctx context.Context) error {
 				defer e.spillAction.WaitForTest()
 			}
 		})
-		e.ctx.GetSessionVars().StmtCtx.MemTracker.FallbackOldAndSetNewAction(e.spillAction)
+		e.ctx.GetSessionVars().StmtCtx.MemTracker.FallbackOldAndSetNewActionForSoftLimit(e.spillAction)
 		e.rowChunks.GetDiskTracker().AttachTo(e.diskTracker)
 		e.rowChunks.GetDiskTracker().SetLabel(memory.LabelForRowChunks)
 	}
@@ -218,7 +218,7 @@ func (e *SortExec) fetchRowChunks(ctx context.Context) error {
 						defer e.spillAction.WaitForTest()
 					}
 				})
-				e.ctx.GetSessionVars().StmtCtx.MemTracker.FallbackOldAndSetNewAction(e.spillAction)
+				e.ctx.GetSessionVars().StmtCtx.MemTracker.FallbackOldAndSetNewActionForSoftLimit(e.spillAction)
 				err = e.rowChunks.Add(chk)
 			}
 			if err != nil {

session/session_test/session_test.go

@@ -49,7 +49,6 @@ import (
 	"github.com/pingcap/tidb/testkit"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util"
-	"github.com/pingcap/tidb/util/memory"
 	"github.com/pingcap/tidb/util/sqlexec"
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/client-go/v2/tikv"
@@ -1719,6 +1718,7 @@ func TestDoDDLJobQuit(t *testing.T) {
 }
 
 func TestCoprocessorOOMAction(t *testing.T) {
+	t.Skip("rate limit action can't control the memory usage in time, skip now")
 	// Assert Coprocessor OOMAction
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)
@@ -2105,15 +2105,8 @@ func TestSetEnableRateLimitAction(t *testing.T) {
 	tk.MustExec("use test")
 	tk.MustExec("create table tmp123(id int)")
 	tk.MustQuery("select * from tmp123;")
-	haveRateLimitAction := false
-	action := tk.Session().GetSessionVars().StmtCtx.MemTracker.GetFallbackForTest(false)
-	for ; action != nil; action = action.GetFallback() {
-		if action.GetPriority() == memory.DefRateLimitPriority {
-			haveRateLimitAction = true
-			break
-		}
-	}
-	require.True(t, haveRateLimitAction)
+	action := tk.Session().GetSessionVars().StmtCtx.MemTracker.GetFallbackForSoftLimitForTest(false)
+	require.NotNil(t, action)
 
 	// assert set sys variable
 	tk.MustExec("set global tidb_enable_rate_limit_action= '0';")
@@ -2123,15 +2116,8 @@ func TestSetEnableRateLimitAction(t *testing.T) {
 	result = tk.MustQuery("select @@tidb_enable_rate_limit_action;")
 	result.Check(testkit.Rows("0"))
 
-	haveRateLimitAction = false
-	action = tk.Session().GetSessionVars().StmtCtx.MemTracker.GetFallbackForTest(false)
-	for ; action != nil; action = action.GetFallback() {
-		if action.GetPriority() == memory.DefRateLimitPriority {
-			haveRateLimitAction = true
-			break
-		}
-	}
-	require.False(t, haveRateLimitAction)
+	action = tk.Session().GetSessionVars().StmtCtx.MemTracker.GetFallbackForSoftLimitForTest(false)
+	require.Nil(t, action)
 }
 
 func TestStmtHints(t *testing.T) {

store/copr/coprocessor.go

@@ -172,7 +172,7 @@ func (c *CopClient) Send(ctx context.Context, req *kv.Request, variables interfa
 	}
 	it.actionOnExceed = newRateLimitAction(uint(it.sendRate.GetCapacity()))
 	if sessionMemTracker != nil && enabledRateLimitAction {
-		sessionMemTracker.FallbackOldAndSetNewAction(it.actionOnExceed)
+		sessionMemTracker.FallbackOldAndSetNewActionForSoftLimit(it.actionOnExceed)
 	}
 
 	ctx = context.WithValue(ctx, tikv.RPCCancellerCtxKey{}, it.rpcCancel)
@@ -1274,9 +1274,6 @@ func newRateLimitAction(totalTokenNumber uint) *rateLimitAction {
 // Action implements ActionOnExceed.Action
 func (e *rateLimitAction) Action(t *memory.Tracker) {
 	if !e.isEnabled() {
-		if fallback := e.GetFallback(); fallback != nil {
-			fallback.Action(t)
-		}
 		return
 	}
 	e.conditionLock()
@@ -1313,11 +1310,6 @@ func (e *rateLimitAction) SetLogHook(hook func(uint64)) {
 
 }
 
-// GetPriority get the priority of the Action.
-func (e *rateLimitAction) GetPriority() int64 {
-	return memory.DefRateLimitPriority
-}
-
 // destroyTokenIfNeeded will check the `exceed` flag after copWorker finished one task.
 // If the exceed flag is true and there is no token been destroyed before, one token will be destroyed,
 // or the token would be return back.

util/chunk/row_container.go

@@ -414,11 +414,6 @@ func (a *SpillDiskAction) Reset() {
 // SetLogHook sets the hook, it does nothing just to form the memory.ActionOnExceed interface.
 func (*SpillDiskAction) SetLogHook(_ func(uint64)) {}
 
-// GetPriority get the priority of the Action.
-func (*SpillDiskAction) GetPriority() int64 {
-	return memory.DefSpillPriority
-}
-
 // WaitForTest waits all goroutine have gone.
 func (a *SpillDiskAction) WaitForTest() {
 	a.testWg.Wait()
@@ -598,13 +593,6 @@ func (a *SortAndSpillDiskAction) Action(t *memory.Tracker) {
 		a.cond.Wait()
 	}
 	a.cond.L.Unlock()
-
-	if !t.CheckExceed() {
-		return
-	}
-	if fallback := a.GetFallback(); fallback != nil {
-		fallback.Action(t)
-	}
 }
 
 // SetLogHook sets the hook, it does nothing just to form the memory.ActionOnExceed interface.

util/chunk/row_container_test.go

@@ -94,8 +94,8 @@ func TestSpillAction(t *testing.T) {
 	var tracker *memory.Tracker
 	var err error
 	tracker = rc.GetMemTracker()
-	tracker.SetBytesLimit(chk.MemoryUsage() + 1)
-	tracker.FallbackOldAndSetNewAction(rc.ActionSpillForTest())
+	tracker.SetBytesLimit(chk.MemoryUsage()/4*5 + 10)
+	tracker.FallbackOldAndSetNewActionForSoftLimit(rc.ActionSpillForTest())
 	require.False(t, rc.AlreadySpilledSafeForTest())
 	err = rc.Add(chk)
 	rc.actionSpill.WaitForTest()
@@ -156,8 +156,8 @@ func TestSortedRowContainerSortSpillAction(t *testing.T) {
 	var tracker *memory.Tracker
 	var err error
 	tracker = rc.GetMemTracker()
-	tracker.SetBytesLimit(chk.MemoryUsage() + int64(8*chk.NumRows()) + 1)
-	tracker.FallbackOldAndSetNewAction(rc.ActionSpillForTest())
+	tracker.SetBytesLimit((chk.MemoryUsage()+int64(8*chk.NumRows()))/4*5 + 10)
+	tracker.FallbackOldAndSetNewActionForSoftLimit(rc.ActionSpillForTest())
 	require.False(t, rc.AlreadySpilledSafeForTest())
 	err = rc.Add(chk)
 	rc.actionSpill.WaitForTest()
@@ -196,8 +196,8 @@ func TestRowContainerResetAndAction(t *testing.T) {
 	var tracker *memory.Tracker
 	var err error
 	tracker = rc.GetMemTracker()
-	tracker.SetBytesLimit(chk.MemoryUsage() + 1)
-	tracker.FallbackOldAndSetNewAction(rc.ActionSpillForTest())
+	tracker.SetBytesLimit(chk.MemoryUsage()/4*5 + 10)
+	tracker.FallbackOldAndSetNewActionForSoftLimit(rc.ActionSpillForTest())
 	require.False(t, rc.AlreadySpilledSafeForTest())
 	err = rc.Add(chk)
 	require.NoError(t, err)
@@ -243,7 +243,7 @@ func TestSpillActionDeadLock(t *testing.T) {
 	tracker = rc.GetMemTracker()
 	tracker.SetBytesLimit(1)
 	ac := rc.ActionSpillForTest()
-	tracker.FallbackOldAndSetNewAction(ac)
+	tracker.FallbackOldAndSetNewActionForSoftLimit(ac)
 	require.False(t, rc.AlreadySpilledSafeForTest())
 	go func() {
 		time.Sleep(200 * time.Millisecond)
@@ -270,7 +270,7 @@ func TestActionBlocked(t *testing.T) {
 	tracker = rc.GetMemTracker()
 	tracker.SetBytesLimit(1450)
 	ac := rc.ActionSpill()
-	tracker.FallbackOldAndSetNewAction(ac)
+	tracker.FallbackOldAndSetNewActionForSoftLimit(ac)
 	for i := 0; i < 10; i++ {
 		err = rc.Add(chk)
 		require.NoError(t, err)

util/cteutil/storage_test.go

@@ -122,7 +122,7 @@ func TestSpillToDisk(t *testing.T) {
 	memTracker := storage.GetMemTracker()
 	memTracker.SetBytesLimit(inChk.MemoryUsage() + 1)
 	action := tmp.(*StorageRC).ActionSpillForTest()
-	memTracker.FallbackOldAndSetNewAction(action)
+	memTracker.FallbackOldAndSetNewActionForSoftLimit(action)
 	diskTracker := storage.GetDiskTracker()
 
 	// All data is in memory.

util/memory/action.go

@@ -39,8 +39,6 @@ type ActionOnExceed interface {
 	SetFallback(a ActionOnExceed)
 	// GetFallback get the fallback action of the Action.
 	GetFallback() ActionOnExceed
-	// GetPriority get the priority of the Action.
-	GetPriority() int64
 	// SetFinished sets the finished state of the Action.
 	SetFinished()
 	// IsFinished returns the finished state of the Action.
@@ -77,14 +75,6 @@ func (b *BaseOOMAction) GetFallback() ActionOnExceed {
 	return b.fallbackAction
 }
 
-// Default OOM Action priority.
-const (
-	DefPanicPriority = iota
-	DefLogPriority
-	DefSpillPriority
-	DefRateLimitPriority
-)
-
 // LogOnExceed logs a warning only once when memory usage exceeds memory quota.
 type LogOnExceed struct {
 	logHook func(uint64)
@@ -114,11 +104,6 @@ func (a *LogOnExceed) Action(t *Tracker) {
 	}
 }
 
-// GetPriority get the priority of the Action
-func (*LogOnExceed) GetPriority() int64 {
-	return DefLogPriority
-}
-
 // PanicOnExceed panics when memory usage exceeds memory quota.
 type PanicOnExceed struct {
 	logHook func(uint64)
@@ -146,11 +131,6 @@ func (a *PanicOnExceed) Action(_ *Tracker) {
 	panic(PanicMemoryExceed + fmt.Sprintf("[conn_id=%d]", a.ConnID))
 }
 
-// GetPriority get the priority of the Action
-func (*PanicOnExceed) GetPriority() int64 {
-	return DefPanicPriority
-}
-
 var (
 	errMemExceedThreshold = dbterror.ClassUtil.NewStd(errno.ErrMemExceedThreshold)
 )