executor: fix the fail of triggering topn fallback action (#56187)

close #56185

pkg/executor/aggregate/agg_spill.go

@@ -377,7 +377,7 @@ func (p *ParallelAggSpillDiskAction) Action(t *memory.Tracker) {
 		return
 	}
 
-	p.triggerFallBackAction(t)
+	p.TriggerFallBackAction(t)
 }
 
 // Return true if we successfully set flag
@@ -386,12 +386,6 @@ func (p *ParallelAggSpillDiskAction) actionImpl(t *memory.Tracker) bool {
 	return p.spillHelper.setNeedSpill(p.e.memTracker, t)
 }
 
-func (p *ParallelAggSpillDiskAction) triggerFallBackAction(t *memory.Tracker) {
-	if fallback := p.GetFallback(); fallback != nil {
-		fallback.Action(t)
-	}
-}
-
 // GetPriority get the priority of the Action
 func (*ParallelAggSpillDiskAction) GetPriority() int64 {
 	return memory.DefSpillPriority

pkg/executor/sortexec/BUILD.bazel

@@ -42,7 +42,7 @@ go_test(
     timeout = "short",
     srcs = ["sort_test.go"],
     flaky = True,
-    shard_count = 15,
+    shard_count = 16,
     deps = [
         "//pkg/config",
         "//pkg/sessionctx/variable",

pkg/executor/sortexec/sort_spill.go

@@ -101,7 +101,7 @@ func (s *parallelSortSpillAction) Action(t *memory.Tracker) {
 	}
 
 	if t.CheckExceed() && !hasEnoughDataToSpill(s.spillHelper.sortExec.memTracker, t) {
-		s.GetFallback().Action(t)
+		s.TriggerFallBackAction(t)
 	}
 }
 

pkg/executor/sortexec/topn_spill.go

@@ -260,6 +260,6 @@ func (t *topNSpillAction) Action(tracker *memory.Tracker) {
 	}
 
 	if tracker.CheckExceed() && !hasEnoughData {
-		t.GetFallback()
+		t.TriggerFallBackAction(tracker)
 	}
 }

pkg/executor/sortexec/topn_spill_test.go

@@ -41,6 +41,7 @@ var totalRowNum = 10000
 var noSpillCaseHardLimit = hardLimit2
 var spillCase1HardLimit = hardLimit1
 var spillCase2HardLimit = hardLimit1
+var fallBackHardLimit = hardLimit1
 var inMemoryThenSpillHardLimit = hardLimit1 * 2
 
 // Test is successful if there is no hang
@@ -492,3 +493,32 @@ func TestIssue54206(t *testing.T) {
 	tk.MustExec("insert into t1 values(1, 1);")
 	tk.MustQuery("select t1.a+t1.b as result from t1 left join t2 on 1 = 0 order by result limit 1;")
 }
+
+func TestTopNFallBackAction(t *testing.T) {
+	sortexec.SetSmallSpillChunkSizeForTest()
+	ctx := mock.NewContext()
+	topNCase := &testutil.SortCase{Rows: totalRowNum, OrderByIdx: []int{0, 1}, Ndvs: []int{0, 0}, Ctx: ctx}
+	newRootExceedAction := new(testutil.MockActionOnExceed)
+
+	ctx.GetSessionVars().InitChunkSize = 32
+	ctx.GetSessionVars().MaxChunkSize = 32
+	ctx.GetSessionVars().MemTracker = memory.NewTracker(memory.LabelForSQLText, fallBackHardLimit)
+	ctx.GetSessionVars().StmtCtx.MemTracker.AttachTo(ctx.GetSessionVars().MemTracker)
+	ctx.GetSessionVars().MemTracker.SetActionOnExceed(newRootExceedAction)
+	ctx.GetSessionVars().MemTracker.Consume(int64(float64(fallBackHardLimit) * 0.99999))
+
+	schema := expression.NewSchema(topNCase.Columns()...)
+	dataSource := buildDataSource(topNCase, schema)
+
+	count := uint64(totalRowNum / 5)
+	offset := count / 5
+
+	exe := buildTopNExec(topNCase, dataSource, offset, count)
+
+	dataSource.PrepareChunks()
+	_ = executeTopNExecutor(t, exe)
+	err := exe.Close()
+	require.NoError(t, err)
+
+	require.Less(t, 0, newRootExceedAction.GetTriggeredNum())
+}

pkg/util/memory/action.go

@@ -93,6 +93,13 @@ func (b *BaseOOMAction) GetFallback() ActionOnExceed {
 	return b.fallbackAction
 }
 
+// TriggerFallBackAction triggers the fallback action of the current action
+func (b *BaseOOMAction) TriggerFallBackAction(tracker *Tracker) {
+	if fallback := b.GetFallback(); fallback != nil {
+		fallback.Action(tracker)
+	}
+}
+
 // Default OOM Action priority.
 const (
 	DefPanicPriority = iota