From 2e661659f7afc2bcc370269dce6a33cee9f126f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= Date: Thu, 9 Jan 2025 15:57:55 +0800 Subject: [PATCH] freelru: Fix GetAndRefreshOrAdd --- common/udpnat2/service.go | 10 ++++------ contrab/freelru/cache.go | 2 +- contrab/freelru/lru.go | 16 +++++++++------- contrab/freelru/shardedlru.go | 7 +++++-- contrab/freelru/syncedlru.go | 7 +++++-- 5 files changed, 24 insertions(+), 18 deletions(-) diff --git a/common/udpnat2/service.go b/common/udpnat2/service.go index 07e786d8..97f753cd 100644 --- a/common/udpnat2/service.go +++ b/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 { diff --git a/contrab/freelru/cache.go b/contrab/freelru/cache.go index 319e6629..e1877fbf 100644 --- a/contrab/freelru/cache.go +++ b/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. diff --git a/contrab/freelru/lru.go b/contrab/freelru/lru.go index 55e89371..6c318570 100644 --- a/contrab/freelru/lru.go +++ b/contrab/freelru/lru.go @@ -522,11 +522,15 @@ 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) @@ -534,17 +538,15 @@ func (lru *LRU[K, V]) getAndRefreshOrAdd(hash uint32, key K, constructor func() } 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. diff --git a/contrab/freelru/shardedlru.go b/contrab/freelru/shardedlru.go index e6aca654..db97efa5 100644 --- a/contrab/freelru/shardedlru.go +++ b/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 } diff --git a/contrab/freelru/syncedlru.go b/contrab/freelru/syncedlru.go index d42da709..38a854a8 100644 --- a/contrab/freelru/syncedlru.go +++ b/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