From 10a7c9b8b73b12d992cc3ccafc5bf8776af9a0c6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=A3=8E=E6=89=87=E6=BB=91=E7=BF=94=E7=BF=BC?=
 <Fangliding.fshxy@outlook.com>
Date: Sat, 15 Feb 2025 15:19:18 +0000
Subject: [PATCH] Use sync.Map

---
 transport/internet/tls/ech.go | 23 +++++++++--------------
 1 file changed, 9 insertions(+), 14 deletions(-)

diff --git a/transport/internet/tls/ech.go b/transport/internet/tls/ech.go
index c56d0f2c..a5a6a305 100644
--- a/transport/internet/tls/ech.go
+++ b/transport/internet/tls/ech.go
@@ -57,19 +57,17 @@ type record struct {
 }
 
 var (
-	dnsCache = make(map[string]record)
-	// global Lock? I'm not sure if I need finer grained locks.
+	dnsCache sync.Map
+	// global Lock? I'm not sure if this needs finer grained locks.
 	// If we do this, we will need to nest another layer of struct
-	dnsCacheLock sync.RWMutex
-	updating     sync.Mutex
+	updating sync.Mutex
 )
 
 // QueryRecord returns the ECH config for given domain.
 // If the record is not in cache or expired, it will query the DOH server and update the cache.
 func QueryRecord(domain string, server string) ([]byte, error) {
-	dnsCacheLock.RLock()
-	rec, found := dnsCache[domain]
-	dnsCacheLock.RUnlock()
+	val, found := dnsCache.Load(domain)
+	rec, _ := val.(record)
 	if found && rec.expire.After(time.Now()) {
 		errors.LogDebug(context.Background(), "Cache hit for domain: ", domain)
 		return rec.echConfig, nil
@@ -79,9 +77,8 @@ func QueryRecord(domain string, server string) ([]byte, error) {
 	defer updating.Unlock()
 	// Try to get cache again after lock, in case another goroutine has updated it
 	// This might happen when the core tring is just stared and multiple goroutines are trying to query the same domain
-	dnsCacheLock.RLock()
-	rec, found = dnsCache[domain]
-	dnsCacheLock.RUnlock()
+	val, found = dnsCache.Load(domain)
+	rec, _ = val.(record)
 	if found && rec.expire.After(time.Now()) {
 		errors.LogDebug(context.Background(), "ECH Config cache hit for domain: ", domain, " after trying to get update lock")
 		return rec.echConfig, nil
@@ -99,14 +96,12 @@ func QueryRecord(domain string, server string) ([]byte, error) {
 		ttl = 600
 	}
 
-	// Get write lock and update cache
-	dnsCacheLock.Lock()
-	defer dnsCacheLock.Unlock()
+	// Update cache
 	newRecored := record{
 		echConfig: echConfig,
 		expire:    time.Now().Add(time.Second * time.Duration(ttl)),
 	}
-	dnsCache[domain] = newRecored
+	dnsCache.Store(domain, newRecored)
 	return echConfig, nil
 }