From a9e2b767265a04feca4699bbcae6560ecd1a8edf Mon Sep 17 00:00:00 2001
From: ti-srebot <66930949+ti-srebot@users.noreply.github.com>
Date: Thu, 4 Aug 2022 17:58:07 +0800
Subject: [PATCH] session: fix multiple domain creation racing in domap
 (#36792) (#36810)

close pingcap/tidb#36791
---
 session/tidb.go | 24 ++++++------------------
 1 file changed, 6 insertions(+), 18 deletions(-)

diff --git a/session/tidb.go b/session/tidb.go
index 12ee40da2d4be..c7d93f71d7081 100644
--- a/session/tidb.go
+++ b/session/tidb.go
@@ -44,24 +44,17 @@ import (
 )
 
 type domainMap struct {
+	mu      sync.Mutex
 	domains map[string]*domain.Domain
-	mu      sync.RWMutex
 }
 
 func (dm *domainMap) Get(store kv.Storage) (d *domain.Domain, err error) {
-	dm.mu.RLock()
+	key := store.UUID()
 
-	// If this is the only domain instance, and the caller doesn't provide store.
-	if len(dm.domains) == 1 && store == nil {
-		for _, r := range dm.domains {
-			dm.mu.RUnlock()
-			return r, nil
-		}
-	}
+	dm.mu.Lock()
+	defer dm.mu.Unlock()
 
-	key := store.UUID()
 	d = dm.domains[key]
-	dm.mu.RUnlock()
 	if d != nil {
 		return
 	}
@@ -94,7 +87,8 @@ func (dm *domainMap) Get(store kv.Storage) (d *domain.Domain, err error) {
 	if err != nil {
 		return nil, err
 	}
-	dm.Set(store, d)
+
+	dm.domains[key] = d
 
 	return
 }
@@ -105,12 +99,6 @@ func (dm *domainMap) Delete(store kv.Storage) {
 	dm.mu.Unlock()
 }
 
-func (dm *domainMap) Set(store kv.Storage, domain *domain.Domain) {
-	dm.mu.Lock()
-	dm.domains[store.UUID()] = domain
-	dm.mu.Unlock()
-}
-
 var (
 	domap = &domainMap{
 		domains: map[string]*domain.Domain{},