freelru: Fix GetAndRefreshOrAdd

Author: nekohasekai

common/udpnat2/service.go

@@ -29,7 +29,7 @@ func New(handler N.UDPConnectionHandlerEx, prepare PrepareFunc, timeout time.Dur
 	if !shared {
 		cache = common.Must1(freelru.New[netip.AddrPort, *natConn](1024, maphash.NewHasher[netip.AddrPort]().Hash32))
 	} else {
-		cache = common.Must1(freelru.NewSharded[netip.AddrPort, *natConn](1024, maphash.NewHasher[netip.AddrPort]().Hash32))
+		cache = common.Must1(freelru.NewSynced[netip.AddrPort, *natConn](1024, maphash.NewHasher[netip.AddrPort]().Hash32))
 	}
 	cache.SetLifetime(timeout)
 	cache.SetHealthCheck(func(port netip.AddrPort, conn *natConn) bool {
@@ -51,7 +51,7 @@ func New(handler N.UDPConnectionHandlerEx, prepare PrepareFunc, timeout time.Dur
 }
 
 func (s *Service) NewPacket(bufferSlices [][]byte, source M.Socksaddr, destination M.Socksaddr, userData any) {
-	conn, loaded := s.cache.GetAndRefreshOrAdd(source.AddrPort(), func() (*natConn, bool) {
+	conn, _, ok := s.cache.GetAndRefreshOrAdd(source.AddrPort(), func() (*natConn, bool) {
 		ok, ctx, writer, onClose := s.prepare(source, destination, userData)
 		if !ok {
 			return nil, false
@@ -67,10 +67,8 @@ func (s *Service) NewPacket(bufferSlices [][]byte, source M.Socksaddr, destinati
 		go s.handler.NewPacketConnectionEx(ctx, newConn, source, destination, onClose)
 		return newConn, true
 	})
-	if !loaded {
-		if conn == nil {
-			return
-		}
+	if !ok {
+		return
 	}
 	buffer := conn.readWaitOptions.NewPacketBuffer()
 	for _, bufferSlice := range bufferSlices {

contrab/freelru/cache.go

@@ -54,7 +54,7 @@ type Cache[K comparable, V comparable] interface {
 	// The lifetime of the found cache item is refreshed, even if it was already expired.
 	GetAndRefresh(key K) (V, bool)
 
-	GetAndRefreshOrAdd(key K, constructor func() (V, bool)) (V, bool)
+	GetAndRefreshOrAdd(key K, constructor func() (V, bool)) (V, bool, bool)
 
 	// Peek looks up a key's value from the cache, without changing its recent-ness.
 	// If the found entry is already expired, the evict function is called.

contrab/freelru/lru.go

@@ -522,29 +522,31 @@ func (lru *LRU[K, V]) getAndRefresh(hash uint32, key K) (value V, ok bool) {
 	return
 }
 
-func (lru *LRU[K, V]) GetAndRefreshOrAdd(key K, constructor func() (V, bool)) (V, bool) {
-	return lru.getAndRefreshOrAdd(lru.hash(key), key, constructor)
+func (lru *LRU[K, V]) GetAndRefreshOrAdd(key K, constructor func() (V, bool)) (V, bool, bool) {
+	value, updated, ok := lru.getAndRefreshOrAdd(lru.hash(key), key, constructor)
+	if !updated && ok {
+		lru.PurgeExpired()
+	}
+	return value, updated, ok
 }
 
-func (lru *LRU[K, V]) getAndRefreshOrAdd(hash uint32, key K, constructor func() (V, bool)) (value V, ok bool) {
+func (lru *LRU[K, V]) getAndRefreshOrAdd(hash uint32, key K, constructor func() (V, bool)) (value V, updated bool, ok bool) {
 	if pos, ok := lru.findKeyNoExpire(hash, key); ok {
 		if pos != lru.head {
 			lru.unlinkElement(pos)
 			lru.setHead(pos)
 		}
 		lru.metrics.Hits++
 		lru.elements[pos].expire = expire(lru.lifetime)
-		return lru.elements[pos].value, ok
+		return lru.elements[pos].value, true, true
 	}
-
 	lru.metrics.Misses++
 	value, ok = constructor()
 	if !ok {
 		return
 	}
 	lru.addWithLifetime(hash, key, value, lru.lifetime)
-	lru.PurgeExpired()
-	return value, false
+	return value, false, true
 }
 
 // Peek looks up a key's value from the cache, without changing its recent-ness.

contrab/freelru/shardedlru.go

@@ -201,14 +201,17 @@ func (lru *ShardedLRU[K, V]) GetAndRefresh(key K) (value V, ok bool) {
 	return
 }
 
-func (lru *ShardedLRU[K, V]) GetAndRefreshOrAdd(key K, constructor func() (V, bool)) (value V, updated bool) {
+func (lru *ShardedLRU[K, V]) GetAndRefreshOrAdd(key K, constructor func() (V, bool)) (value V, updated bool, ok bool) {
 	hash := lru.hash(key)
 	shard := (hash >> 16) & lru.mask
 
 	lru.mus[shard].Lock()
-	value, updated = lru.lrus[shard].getAndRefreshOrAdd(hash, key, constructor)
+	value, updated, ok = lru.lrus[shard].getAndRefreshOrAdd(hash, key, constructor)
 	lru.mus[shard].Unlock()
 
+	if !updated && ok {
+		lru.PurgeExpired()
+	}
 	return
 }
 

contrab/freelru/syncedlru.go

@@ -121,11 +121,14 @@ func (lru *SyncedLRU[K, V]) GetAndRefresh(key K) (value V, ok bool) {
 	return
 }
 
-func (lru *SyncedLRU[K, V]) GetAndRefreshOrAdd(key K, constructor func() (V, bool)) (value V, updated bool) {
+func (lru *SyncedLRU[K, V]) GetAndRefreshOrAdd(key K, constructor func() (V, bool)) (value V, updated bool, ok bool) {
 	hash := lru.lru.hash(key)
 
 	lru.mu.Lock()
-	value, updated = lru.lru.getAndRefreshOrAdd(hash, key, constructor)
+	value, updated, ok = lru.lru.getAndRefreshOrAdd(hash, key, constructor)
+	if !updated && ok {
+		lru.lru.PurgeExpired()
+	}
 	lru.mu.Unlock()
 
 	return