stmtctx: add mutex to protect stmtCache(#36159) (#36351)

close #36159

sessionctx/stmtctx/stmtctx.go

@@ -239,7 +239,12 @@ type StatementContext struct {
 	TaskMapBakTS          uint64 // counter for
 
 	// stmtCache is used to store some statement-related values.
-	stmtCache map[StmtCacheKey]interface{}
+	// add mutex to protect stmtCache concurrent access
+	// https://github.com/pingcap/tidb/issues/36159
+	stmtCache struct {
+		mu   sync.Mutex
+		data map[StmtCacheKey]interface{}
+	}
 
 	// Map to store all CTE storages of current SQL.
 	// Will clean up at the end of the execution.
@@ -352,23 +357,29 @@ const (
 
 // GetOrStoreStmtCache gets the cached value of the given key if it exists, otherwise stores the value.
 func (sc *StatementContext) GetOrStoreStmtCache(key StmtCacheKey, value interface{}) interface{} {
-	if sc.stmtCache == nil {
-		sc.stmtCache = make(map[StmtCacheKey]interface{})
+	sc.stmtCache.mu.Lock()
+	defer sc.stmtCache.mu.Unlock()
+	if sc.stmtCache.data == nil {
+		sc.stmtCache.data = make(map[StmtCacheKey]interface{})
 	}
-	if _, ok := sc.stmtCache[key]; !ok {
-		sc.stmtCache[key] = value
+	if _, ok := sc.stmtCache.data[key]; !ok {
+		sc.stmtCache.data[key] = value
 	}
-	return sc.stmtCache[key]
+	return sc.stmtCache.data[key]
 }
 
 // ResetInStmtCache resets the cache of given key.
 func (sc *StatementContext) ResetInStmtCache(key StmtCacheKey) {
-	delete(sc.stmtCache, key)
+	sc.stmtCache.mu.Lock()
+	defer sc.stmtCache.mu.Unlock()
+	delete(sc.stmtCache.data, key)
 }
 
 // ResetStmtCache resets all cached values.
 func (sc *StatementContext) ResetStmtCache() {
-	sc.stmtCache = make(map[StmtCacheKey]interface{})
+	sc.stmtCache.mu.Lock()
+	defer sc.stmtCache.mu.Unlock()
+	sc.stmtCache.data = make(map[StmtCacheKey]interface{})
 }
 
 // SQLDigest gets normalized and digest for provided sql.