From d50e74dc421abd7005215cfefe9a2a75450c125f Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Wed, 8 Nov 2023 11:05:12 +0800
Subject: [PATCH] executor: fix panic when to enable tidb_mem_quota_analyze 
 (#48358) (#48374)

close pingcap/tidb#48357
---
 pkg/executor/executor.go                             | 12 ++++++++++++
 .../test/analyzetest/memorycontrol/BUILD.bazel       |  2 +-
 .../analyzetest/memorycontrol/memory_control_test.go | 10 ++++++++++
 pkg/util/util.go                                     |  8 ++++++++
 4 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/pkg/executor/executor.go b/pkg/executor/executor.go
index 0f742760ac641..861511f804462 100644
--- a/pkg/executor/executor.go
+++ b/pkg/executor/executor.go
@@ -17,6 +17,7 @@ package executor
 import (
 	"cmp"
 	"context"
+	stderrors "errors"
 	"fmt"
 	"math"
 	"runtime/pprof"
@@ -199,6 +200,7 @@ func (a *globalPanicOnExceed) Action(t *memory.Tracker) {
 	default:
 		msg = "Out of Unknown Resource Quota!"
 	}
+	// TODO(hawkingrei): should return error instead.
 	panic(msg)
 }
 
@@ -1940,6 +1942,16 @@ func (e *UnionExec) Close() error {
 // ResetContextOfStmt resets the StmtContext and session variables.
 // Before every execution, we must clear statement context.
 func ResetContextOfStmt(ctx sessionctx.Context, s ast.StmtNode) (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			logutil.BgLogger().Warn("ResetContextOfStmt panicked", zap.Stack("stack"), zap.Any("recover", r), zap.Error(err))
+			if err != nil {
+				err = stderrors.Join(err, util.GetRecoverError(r))
+			} else {
+				err = util.GetRecoverError(r)
+			}
+		}
+	}()
 	vars := ctx.GetSessionVars()
 	for name, val := range vars.StmtCtx.SetVarHintRestore {
 		err := vars.SetSystemVar(name, val)
diff --git a/pkg/executor/test/analyzetest/memorycontrol/BUILD.bazel b/pkg/executor/test/analyzetest/memorycontrol/BUILD.bazel
index 758c006a15981..bf8c263bb6016 100644
--- a/pkg/executor/test/analyzetest/memorycontrol/BUILD.bazel
+++ b/pkg/executor/test/analyzetest/memorycontrol/BUILD.bazel
@@ -8,7 +8,7 @@ go_test(
         "memory_control_test.go",
     ],
     flaky = True,
-    shard_count = 4,
+    shard_count = 5,
     deps = [
         "//pkg/config",
         "//pkg/executor",
diff --git a/pkg/executor/test/analyzetest/memorycontrol/memory_control_test.go b/pkg/executor/test/analyzetest/memorycontrol/memory_control_test.go
index d39e0a27f459f..a64e517e60d95 100644
--- a/pkg/executor/test/analyzetest/memorycontrol/memory_control_test.go
+++ b/pkg/executor/test/analyzetest/memorycontrol/memory_control_test.go
@@ -202,3 +202,13 @@ func TestMemQuotaAnalyze(t *testing.T) {
 	tk.MustExec("set global tidb_mem_quota_analyze=128;")
 	tk.MustExecToErr("analyze table tbl_2;")
 }
+
+func TestMemQuotaAnalyze2(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table tbl_2 ( col_20 decimal default 84232 , col_21 tinyint not null , col_22 int default 80814394 , col_23 mediumint default -8036687 not null , col_24 smallint default 9185 not null , col_25 tinyint unsigned default 65 , col_26 char(115) default 'ZyfroRODMbNDRZnPNRW' not null , col_27 bigint not null , col_28 tinyint not null , col_29 char(130) default 'UMApsVgzHblmY' , primary key idx_14 ( col_28,col_22 ) , unique key idx_15 ( col_24,col_22 ) , key idx_16 ( col_21,col_20,col_24,col_25,col_27,col_28,col_26,col_29 ) , key idx_17 ( col_24,col_25 ) , unique key idx_18 ( col_25,col_23,col_29,col_27,col_26,col_22 ) , key idx_19 ( col_25,col_22,col_26,col_23 ) , unique key idx_20 ( col_22,col_24,col_28,col_29,col_26,col_20 ) , key idx_21 ( col_25,col_24,col_26,col_29,col_27,col_22,col_28 ) );")
+	tk.MustExec("insert ignore into tbl_2 values ( 942,33,-1915007317,3408149,-3699,193,'Trywdis',1876334369465184864,115,null );")
+	tk.MustExec("set global tidb_mem_quota_analyze=128;")
+	tk.MustExecToErr("analyze table tbl_2;")
+}
diff --git a/pkg/util/util.go b/pkg/util/util.go
index eb2f9d5e584dd..ced69cfbd0bb7 100644
--- a/pkg/util/util.go
+++ b/pkg/util/util.go
@@ -296,3 +296,11 @@ func IsInCorrectIdentifierName(name string) bool {
 	}
 	return false
 }
+
+// GetRecoverError gets the error from recover.
+func GetRecoverError(r interface{}) error {
+	if err, ok := r.(error); ok {
+		return err
+	}
+	return errors.Errorf("%v", r)
+}