Use sync.Map

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
 )
 
 // 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
 }