Add freelru.SetUpdateLifetimeOnGet/GetWithLifetime

Author: nekohasekai

common/udpnat2/service.go

@@ -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:

contrab/freelru/cache.go

@@ -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.
@@ -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)

contrab/freelru/lru.go

@@ -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
@@ -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
@@ -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
@@ -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)
 			}
-			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
 		}
 	}
 }
@@ -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++
@@ -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
 	}
 
@@ -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
 	}

contrab/freelru/lru_test.go

@@ -14,18 +14,21 @@ 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")
-	time.Sleep(time.Second + time.Millisecond*100)
 	_, ok := lru.Get("hello")
 	require.True(t, ok)
+	time.Sleep(time.Second + time.Millisecond*100)
+	_, ok = lru.Get("hello")
+	require.True(t, ok)
 }
 
 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")

contrab/freelru/sharedlru.go

@@ -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]) {
@@ -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) {