executor: fix unstable TestGlobalMemoryTrackerOnCleanUp (#37372)

close #36585

executor/insert.go

@@ -315,6 +315,7 @@ func (e *InsertExec) Next(ctx context.Context, req *chunk.Chunk) error {
 
 // Close implements the Executor Close interface.
 func (e *InsertExec) Close() error {
+	defer e.memTracker.ReplaceBytesUsed(0)
 	e.ctx.GetSessionVars().CurrInsertValues = chunk.Row{}
 	e.ctx.GetSessionVars().CurrInsertBatchExtraCols = e.ctx.GetSessionVars().CurrInsertBatchExtraCols[0:0:0]
 	e.setMessage()

executor/memtest/mem_test.go

@@ -17,30 +17,30 @@ package memtest
 import (
 	"testing"
 
-	"github.com/pingcap/tidb/executor"
 	"github.com/pingcap/tidb/testkit"
 	"github.com/stretchr/testify/require"
 )
 
-func TestGlobalMemoryTrackerOnCleanUp(t *testing.T) {
-	originConsume := executor.GlobalMemoryUsageTracker.BytesConsumed()
+func TestInsertUpdateTrackerOnCleanUp(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)
 	tk.MustExec("use test")
 	tk.MustExec("drop table if exists t")
 	tk.MustExec("create table t (id int)")
 
+	originConsume := tk.Session().GetSessionVars().StmtCtx.MemTracker.BytesConsumed()
 	// assert insert
 	tk.MustExec("insert t (id) values (1)")
 	tk.MustExec("insert t (id) values (2)")
 	tk.MustExec("insert t (id) values (3)")
-	afterConsume := executor.GlobalMemoryUsageTracker.BytesConsumed()
+	afterConsume := tk.Session().GetSessionVars().StmtCtx.MemTracker.BytesConsumed()
 	require.Equal(t, afterConsume, originConsume)
 
+	originConsume = tk.Session().GetSessionVars().StmtCtx.MemTracker.BytesConsumed()
 	// assert update
 	tk.MustExec("update t set id = 4 where id = 1")
 	tk.MustExec("update t set id = 5 where id = 2")
 	tk.MustExec("update t set id = 6 where id = 3")
-	afterConsume = executor.GlobalMemoryUsageTracker.BytesConsumed()
+	afterConsume = tk.Session().GetSessionVars().StmtCtx.MemTracker.BytesConsumed()
 	require.Equal(t, afterConsume, originConsume)
 }

executor/update.go

@@ -427,14 +427,14 @@ func (e *UpdateExec) composeGeneratedColumns(rowIdx int, newRowData []types.Datu
 
 // Close implements the Executor Close interface.
 func (e *UpdateExec) Close() error {
+	defer e.memTracker.ReplaceBytesUsed(0)
 	e.setMessage()
 	if e.runtimeStats != nil && e.stats != nil {
 		txn, err := e.ctx.Txn(false)
 		if err == nil && txn.Valid() && txn.GetSnapshot() != nil {
 			txn.GetSnapshot().SetOption(kv.CollectRuntimeStats, nil)
 		}
 	}
-	defer e.memTracker.ReplaceBytesUsed(0)
 	return e.children[0].Close()
 }