Skip to content

Commit

Permalink
freelru: Fix GetAndRefreshOrAdd
Browse files Browse the repository at this point in the history
  • Loading branch information
nekohasekai committed Jan 9, 2025
1 parent be9840c commit 2e66165
Show file tree
Hide file tree
Showing 5 changed files with 24 additions and 18 deletions.
10 changes: 4 additions & 6 deletions common/udpnat2/service.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand All @@ -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 {
Expand Down
2 changes: 1 addition & 1 deletion contrab/freelru/cache.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
16 changes: 9 additions & 7 deletions contrab/freelru/lru.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
7 changes: 5 additions & 2 deletions contrab/freelru/shardedlru.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down
7 changes: 5 additions & 2 deletions contrab/freelru/syncedlru.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 2e66165

Please sign in to comment.