优化修复

dns/cache.go

@@ -19,6 +19,9 @@ type DNSCache struct {
 	cache      map[string]*DNSCacheItem // 缓存映射表
 	mutex      sync.RWMutex             // 读写锁，保护缓存
 	defaultTTL time.Duration            // 默认缓存TTL
+	maxSize    int                      // 最大缓存条目数
+	// 使用链表结构来跟踪缓存条目的访问顺序，用于LRU淘汰
+	accessList []string                 // 记录访问顺序，最新访问的放在最后
 }
 
 // NewDNSCache 创建新的DNS缓存实例
@@ -26,6 +29,8 @@ func NewDNSCache(defaultTTL time.Duration) *DNSCache {
 	cache := &DNSCache{
 		cache:      make(map[string]*DNSCacheItem),
 		defaultTTL: defaultTTL,
+		maxSize:    10000, // 默认最大缓存10000条记录
+		accessList: make([]string, 0, 10000),
 	}
 
 	// 启动缓存清理协程
@@ -110,32 +115,70 @@ func (c *DNSCache) Set(qName string, qType uint16, response *dns.Msg, ttl time.D
 	}
 
 	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	// 如果条目已存在，先从访问列表中移除
+	for i, k := range c.accessList {
+		if k == key {
+			// 移除旧位置
+			c.accessList = append(c.accessList[:i], c.accessList[i+1:]...)
+			break
+		}
+	}
+
+	// 将新条目添加到访问列表末尾
+	c.accessList = append(c.accessList, key)
 	c.cache[key] = item
-	c.mutex.Unlock()
+
+	// 检查是否超过最大大小限制，如果超过则移除最久未使用的条目
+	if len(c.cache) > c.maxSize {
+		// 最久未使用的条目是访问列表的第一个
+		oldestKey := c.accessList[0]
+		// 从缓存和访问列表中移除
+		delete(c.cache, oldestKey)
+		c.accessList = c.accessList[1:]
+	}
 }
 
 // Get 获取缓存项
 func (c *DNSCache) Get(qName string, qType uint16) (*dns.Msg, bool) {
 	key := cacheKey(qName, qType)
 
-	c.mutex.RLock()
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
 	item, found := c.cache[key]
 	if !found {
-		c.mutex.RUnlock()
 		return nil, false
 	}
 
 	// 检查是否过期
 	if time.Now().After(item.Expiry) {
-		c.mutex.RUnlock()
-		// 过期了，删除缓存项（在写锁中）
-		c.delete(key)
+		// 过期了，删除缓存项
+		delete(c.cache, key)
+		// 从访问列表中移除
+		for i, k := range c.accessList {
+			if k == key {
+				c.accessList = append(c.accessList[:i], c.accessList[i+1:]...)
+				break
+			}
+		}
 		return nil, false
 	}
 
+	// 将访问的条目移动到访问列表末尾（标记为最近使用）
+	for i, k := range c.accessList {
+		if k == key {
+			// 移除旧位置
+			c.accessList = append(c.accessList[:i], c.accessList[i+1:]...)
+			// 添加到末尾
+			c.accessList = append(c.accessList, key)
+			break
+		}
+	}
+
 	// 返回缓存的响应副本
 	response := item.Response.Copy()
-	c.mutex.RUnlock()
 
 	return response, true
 }
@@ -143,14 +186,24 @@ func (c *DNSCache) Get(qName string, qType uint16) (*dns.Msg, bool) {
 // delete 删除缓存项
 func (c *DNSCache) delete(key string) {
 	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	// 从缓存中删除
 	delete(c.cache, key)
-	c.mutex.Unlock()
+	// 从访问列表中移除
+	for i, k := range c.accessList {
+		if k == key {
+			c.accessList = append(c.accessList[:i], c.accessList[i+1:]...)
+			break
+		}
+	}
 }
 
 // Clear 清空缓存
 func (c *DNSCache) Clear() {
 	c.mutex.Lock()
 	c.cache = make(map[string]*DNSCacheItem)
+	c.accessList = make([]string, 0, c.maxSize) // 重置访问列表
 	c.mutex.Unlock()
 }
 
@@ -178,9 +231,23 @@ func (c *DNSCache) cleanupExpired() {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
 
+	// 收集所有过期的键
+	var expiredKeys []string
 	for key, item := range c.cache {
 		if now.After(item.Expiry) {
-			delete(c.cache, key)
+			expiredKeys = append(expiredKeys, key)
+		}
+	}
+
+	// 删除过期的缓存项
+	for _, key := range expiredKeys {
+		delete(c.cache, key)
+		// 从访问列表中移除
+		for i, k := range c.accessList {
+			if k == key {
+				c.accessList = append(c.accessList[:i], c.accessList[i+1:]...)
+				break
+			}
 		}
 	}
 }