session: fix the memory tracker oom action for nt-dml (#34557)

ref #33485

session/nontransactional.go

@@ -94,8 +94,7 @@ func HandleNonTransactionalDelete(ctx context.Context, stmt *ast.NonTransactiona
 
 	// TODO: choose an appropriate quota.
 	// Use the mem-quota-query as a workaround. As a result, a NT-DML may consume 2x of the memory quota.
-	memTracker := memory.NewTracker(memory.LabelForNonTransactionalDML, se.GetSessionVars().MemQuotaQuery)
-	memTracker.AttachToGlobalTracker(executor.GlobalMemoryUsageTracker)
+	memTracker := setMemTracker(se)
 	defer memTracker.DetachFromGlobalTracker()
 	jobs, err := buildShardJobs(ctx, stmt, se, selectSQL, shardColumnInfo, memTracker)
 	if err != nil {
@@ -112,6 +111,25 @@ func HandleNonTransactionalDelete(ctx context.Context, stmt *ast.NonTransactiona
 	return buildExecuteResults(ctx, jobs, se.GetSessionVars().BatchSize.MaxChunkSize, se.GetSessionVars().EnableRedactLog)
 }
 
+func setMemTracker(se Session) *memory.Tracker {
+	memTracker := memory.NewTracker(memory.LabelForNonTransactionalDML, se.GetSessionVars().MemQuotaQuery)
+	globalConfig := config.GetGlobalConfig()
+	switch globalConfig.OOMAction {
+	case config.OOMActionCancel:
+		action := &memory.PanicOnExceed{ConnID: se.GetSessionVars().ConnectionID}
+		action.SetLogHook(domain.GetDomain(se).ExpensiveQueryHandle().LogOnQueryExceedMemQuota)
+		memTracker.SetActionOnExceed(action)
+	case config.OOMActionLog:
+		fallthrough
+	default:
+		action := &memory.LogOnExceed{ConnID: se.GetSessionVars().ConnectionID}
+		action.SetLogHook(domain.GetDomain(se).ExpensiveQueryHandle().LogOnQueryExceedMemQuota)
+		memTracker.SetActionOnExceed(action)
+	}
+	memTracker.AttachToGlobalTracker(executor.GlobalMemoryUsageTracker)
+	return memTracker
+}
+
 func checkConstraint(stmt *ast.NonTransactionalDeleteStmt, se Session) error {
 	sessVars := se.GetSessionVars()
 	if !(sessVars.IsAutocommit() && !sessVars.InTxn()) {