Skip to content

Commit

Permalink
Add freelru.SetUpdateLifetimeOnGet/GetWithLifetime
Browse files Browse the repository at this point in the history
  • Loading branch information
nekohasekai committed Nov 14, 2024
1 parent c432bef commit 4d3f04d
Show file tree
Hide file tree
Showing 5 changed files with 58 additions and 20 deletions.
1 change: 1 addition & 0 deletions common/udpnat2/service.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ func New(handler N.UDPConnectionHandlerEx, prepare PrepareFunc, timeout time.Dur
cache = common.Must1(freelru.NewSharded[netip.AddrPort, *Conn](1024, maphash.NewHasher[netip.AddrPort]().Hash32))
}
cache.SetLifetime(timeout)
cache.SetUpdateLifetimeOnGet(true)
cache.SetHealthCheck(func(port netip.AddrPort, conn *Conn) bool {
select {
case <-conn.doneChan:
Expand Down
4 changes: 4 additions & 0 deletions contrab/freelru/cache.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ type Cache[K comparable, V any] interface {
// Lifetime 0 means "forever".
SetLifetime(lifetime time.Duration)

SetUpdateLifetimeOnGet(update bool)

SetHealthCheck(healthCheck HealthCheckCallback[K, V])

// SetOnEvict sets the OnEvict callback function.
Expand All @@ -47,6 +49,8 @@ type Cache[K comparable, V any] interface {
// and the return value indicates that the key was not found.
Get(key K) (V, bool)

GetWithLifetime(key K) (value V, lifetime time.Time, ok 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.
Peek(key K) (V, bool)
Expand Down
50 changes: 31 additions & 19 deletions contrab/freelru/lru.go
Original file line number Diff line number Diff line change
Expand Up @@ -63,13 +63,14 @@ const emptyBucket = math.MaxUint32

// LRU implements a non-thread safe fixed size LRU cache.
type LRU[K comparable, V any] struct {
buckets []uint32 // contains positions of bucket lists or 'emptyBucket'
elements []element[K, V]
onEvict OnEvictCallback[K, V]
hash HashKeyCallback[K]
healthCheck HealthCheckCallback[K, V]
lifetime time.Duration
metrics Metrics
buckets []uint32 // contains positions of bucket lists or 'emptyBucket'
elements []element[K, V]
onEvict OnEvictCallback[K, V]
hash HashKeyCallback[K]
healthCheck HealthCheckCallback[K, V]
lifetime time.Duration
updateLifetimeOnGet bool
metrics Metrics

// used for element clearing after removal or expiration
emptyKey K
Expand Down Expand Up @@ -100,6 +101,10 @@ func (lru *LRU[K, V]) SetLifetime(lifetime time.Duration) {
lru.lifetime = lifetime
}

func (lru *LRU[K, V]) SetUpdateLifetimeOnGet(update bool) {
lru.updateLifetimeOnGet = update
}

// SetOnEvict sets the OnEvict callback function.
// The onEvict function is called for each evicted lru entry.
// Eviction happens
Expand Down Expand Up @@ -303,10 +308,10 @@ func (lru *LRU[K, V]) clearKeyAndValue(pos uint32) {
lru.elements[pos].value = lru.emptyValue
}

func (lru *LRU[K, V]) findKey(hash uint32, key K, updateLifetimeOnGet bool) (uint32, bool) {
func (lru *LRU[K, V]) findKey(hash uint32, key K, updateLifetimeOnGet bool) (uint32, int64, bool) {
_, startPos := lru.hashToPos(hash)
if startPos == emptyBucket {
return emptyBucket, false
return emptyBucket, 0, false
}

pos := startPos
Expand All @@ -315,18 +320,18 @@ func (lru *LRU[K, V]) findKey(hash uint32, key K, updateLifetimeOnGet bool) (uin
elem := lru.elements[pos]
if (elem.expire != 0 && elem.expire <= now()) || (lru.healthCheck != nil && !lru.healthCheck(key, elem.value)) {
lru.removeAt(pos)
return emptyBucket, false
return emptyBucket, elem.expire, false
}
if updateLifetimeOnGet {
lru.elements[pos].expire = expire(lru.lifetime)
elem.expire = expire(lru.lifetime)
}
return pos, true
return pos, elem.expire, true
}

pos = lru.elements[pos].nextBucket
if pos == startPos {
// Key not found
return emptyBucket, false
return emptyBucket, 0, false
}
}
}
Expand Down Expand Up @@ -439,17 +444,24 @@ func (lru *LRU[K, V]) add(hash uint32, key K, value V) (evicted bool) {
// If the found cache item is already expired, the evict function is called
// and the return value indicates that the key was not found.
func (lru *LRU[K, V]) Get(key K) (value V, ok bool) {
return lru.get(lru.hash(key), key)
value, _, ok = lru.get(lru.hash(key), key)
return
}

func (lru *LRU[K, V]) get(hash uint32, key K) (value V, ok bool) {
if pos, ok := lru.findKey(hash, key, true); ok {
func (lru *LRU[K, V]) GetWithLifetime(key K) (value V, lifetime time.Time, ok bool) {
value, expireMills, ok := lru.get(lru.hash(key), key)
lifetime = time.UnixMilli(expireMills)
return
}

func (lru *LRU[K, V]) get(hash uint32, key K) (value V, expire int64, ok bool) {
if pos, expire, ok := lru.findKey(hash, key, lru.updateLifetimeOnGet); ok {
if pos != lru.head {
lru.unlinkElement(pos)
lru.setHead(pos)
}
lru.metrics.Hits++
return lru.elements[pos].value, ok
return lru.elements[pos].value, expire, ok
}

lru.metrics.Misses++
Expand All @@ -463,7 +475,7 @@ func (lru *LRU[K, V]) Peek(key K) (value V, ok bool) {
}

func (lru *LRU[K, V]) peek(hash uint32, key K) (value V, ok bool) {
if pos, ok := lru.findKey(hash, key, false); ok {
if pos, _, ok := lru.findKey(hash, key, false); ok {
return lru.elements[pos].value, ok
}

Expand All @@ -490,7 +502,7 @@ func (lru *LRU[K, V]) Remove(key K) (removed bool) {
}

func (lru *LRU[K, V]) remove(hash uint32, key K) (removed bool) {
if pos, ok := lru.findKey(hash, key, false); ok {
if pos, _, ok := lru.findKey(hash, key, false); ok {
lru.removeAt(pos)
return ok
}
Expand Down
2 changes: 2 additions & 0 deletions contrab/freelru/lru_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ func TestMyChange0(t *testing.T) {
t.Parallel()
lru, err := freelru.New[string, string](1024, maphash.NewHasher[string]().Hash32)
require.NoError(t, err)
lru.SetUpdateLifetimeOnGet(true)
lru.AddWithLifetime("hello", "world", 2*time.Second)
time.Sleep(time.Second)
lru.Get("hello")
Expand All @@ -26,6 +27,7 @@ func TestMyChange1(t *testing.T) {
t.Parallel()
lru, err := freelru.New[string, string](1024, maphash.NewHasher[string]().Hash32)
require.NoError(t, err)
lru.SetUpdateLifetimeOnGet(true)
lru.AddWithLifetime("hello", "world", 2*time.Second)
time.Sleep(time.Second)
lru.Peek("hello")
Expand Down
21 changes: 20 additions & 1 deletion contrab/freelru/sharedlru.go
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,14 @@ func (lru *ShardedLRU[K, V]) SetLifetime(lifetime time.Duration) {
}
}

func (lru *ShardedLRU[K, V]) SetUpdateLifetimeOnGet(update bool) {
for shard := range lru.lrus {
lru.mus[shard].Lock()
lru.lrus[shard].SetUpdateLifetimeOnGet(update)
lru.mus[shard].Unlock()
}
}

// SetOnEvict sets the OnEvict callback function.
// The onEvict function is called for each evicted lru entry.
func (lru *ShardedLRU[K, V]) SetOnEvict(onEvict OnEvictCallback[K, V]) {
Expand Down Expand Up @@ -170,12 +178,23 @@ func (lru *ShardedLRU[K, V]) Get(key K) (value V, ok bool) {
shard := (hash >> 16) & lru.mask

lru.mus[shard].Lock()
value, ok = lru.lrus[shard].get(hash, key)
value, _, ok = lru.lrus[shard].get(hash, key)
lru.mus[shard].Unlock()

return
}

func (lru *ShardedLRU[K, V]) GetWithLifetime(key K) (value V, lifetime time.Time, ok bool) {
hash := lru.hash(key)
shard := (hash >> 16) & lru.mask

lru.mus[shard].Lock()
value, expireMills, ok := lru.lrus[shard].get(hash, key)
lru.mus[shard].Unlock()
lifetime = time.UnixMilli(expireMills)
return
}

// 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.
func (lru *ShardedLRU[K, V]) Peek(key K) (value V, ok bool) {
Expand Down

0 comments on commit 4d3f04d

Please sign in to comment.