freelru: make cache growable

Author: nekohasekai

common/udpnat2/conn.go

@@ -30,7 +30,7 @@ var (
 )
 
 type natConn struct {
-	cache           *freelru.ConcurrentLRU[netip.AddrPort, *natConn]
+	cache           *freelru.Cache[netip.AddrPort, *natConn]
 	writer          N.PacketWriter
 	localAddr       M.Socksaddr
 	handlerAccess   sync.RWMutex

common/udpnat2/service.go

@@ -15,7 +15,7 @@ import (
 )
 
 type Service struct {
-	cache   *freelru.ConcurrentLRU[netip.AddrPort, *natConn]
+	cache   *freelru.Cache[netip.AddrPort, *natConn]
 	handler N.UDPConnectionHandlerEx
 	prepare PrepareFunc
 }
@@ -26,7 +26,7 @@ func New(handler N.UDPConnectionHandlerEx, prepare PrepareFunc, timeout time.Dur
 	if timeout == 0 {
 		panic("invalid timeout")
 	}
-	cache := common.Must1(freelru.NewConcurrent[netip.AddrPort, *natConn](1024, maphash.NewHasher[netip.AddrPort]().Hash32, shared))
+	cache := common.Must1(freelru.New[netip.AddrPort, *natConn](1024, maphash.NewHasher[netip.AddrPort]().Hash32, shared))
 	cache.SetLifetime(timeout)
 	cache.SetHealthCheck(func(port netip.AddrPort, conn *natConn) bool {
 		select {

contrab/freelru/cache.go

@@ -1,102 +1,204 @@
-// Licensed to Elasticsearch B.V. under one or more contributor
-// license agreements. See the NOTICE file distributed with
-// this work for additional information regarding copyright
-// ownership. Elasticsearch B.V. licenses this file to you under
-// the Apache License, Version 2.0 (the "License"); you may
-// not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//    http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied.  See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
 package freelru
 
 import "time"
 
-type Cache[K comparable, V comparable] interface {
-	// SetLifetime sets the default lifetime of LRU elements.
-	// Lifetime 0 means "forever".
-	SetLifetime(lifetime time.Duration)
+// Cache is a thread-safe LRU cache that can use either a single lock or
+// multiple shards.
+type Cache[K comparable, V comparable] struct {
+	synced  *SyncedLRU[K, V]
+	sharded *ShardedLRU[K, V]
+}
+
+// New creates a thread-safe, growable LRU cache with the given maximum capacity.
+// When sharded is false, the cache preserves exact LRU behavior with a single
+// lock. When sharded is true, it reduces lock contention at the cost of exact
+// global LRU behavior.
+func New[K comparable, V comparable](capacity uint32, hash HashKeyCallback[K], sharded bool) (*Cache[K, V], error) {
+	if sharded {
+		cache, err := newSharded[K, V](capacity, hash)
+		if err != nil {
+			return nil, err
+		}
+		return &Cache[K, V]{sharded: cache}, nil
+	}
+	cache, err := newSynced[K, V](capacity, hash)
+	if err != nil {
+		return nil, err
+	}
+	return &Cache[K, V]{synced: cache}, nil
+}
 
-	// SetOnEvict sets the OnEvict callback function.
-	// The onEvict function is called for each evicted lru entry.
-	SetOnEvict(onEvict OnEvictCallback[K, V])
+func (lru *Cache[K, V]) SetLifetime(lifetime time.Duration) {
+	if lru.sharded != nil {
+		lru.sharded.SetLifetime(lifetime)
+	} else {
+		lru.synced.SetLifetime(lifetime)
+	}
+}
 
-	SetHealthCheck(healthCheck HealthCheckCallback[K, V])
+func (lru *Cache[K, V]) SetOnEvict(onEvict OnEvictCallback[K, V]) {
+	if lru.sharded != nil {
+		lru.sharded.SetOnEvict(onEvict)
+	} else {
+		lru.synced.SetOnEvict(onEvict)
+	}
+}
 
-	// Len returns the number of elements stored in the cache.
-	Len() int
+func (lru *Cache[K, V]) SetHealthCheck(healthCheck HealthCheckCallback[K, V]) {
+	if lru.sharded != nil {
+		lru.sharded.SetHealthCheck(healthCheck)
+	} else {
+		lru.synced.SetHealthCheck(healthCheck)
+	}
+}
 
-	// AddWithLifetime adds a key:value to the cache with a lifetime.
-	// Returns true, true if key was updated and eviction occurred.
-	AddWithLifetime(key K, value V, lifetime time.Duration) (evicted bool)
+func (lru *Cache[K, V]) Len() int {
+	if lru.sharded != nil {
+		return lru.sharded.Len()
+	}
+	return lru.synced.Len()
+}
 
-	// Add adds a key:value to the cache.
-	// Returns true, true if key was updated and eviction occurred.
-	Add(key K, value V) (evicted bool)
+func (lru *Cache[K, V]) AddWithLifetime(key K, value V, lifetime time.Duration) bool {
+	if lru.sharded != nil {
+		return lru.sharded.AddWithLifetime(key, value, lifetime)
+	}
+	return lru.synced.AddWithLifetime(key, value, lifetime)
+}
+
+func (lru *Cache[K, V]) Add(key K, value V) bool {
+	if lru.sharded != nil {
+		return lru.sharded.Add(key, value)
+	}
+	return lru.synced.Add(key, value)
+}
+
+func (lru *Cache[K, V]) Get(key K) (V, bool) {
+	if lru.sharded != nil {
+		return lru.sharded.Get(key)
+	}
+	return lru.synced.Get(key)
+}
+
+func (lru *Cache[K, V]) GetWithLifetime(key K) (V, time.Time, bool) {
+	if lru.sharded != nil {
+		return lru.sharded.GetWithLifetime(key)
+	}
+	return lru.synced.GetWithLifetime(key)
+}
 
-	// Get returns the value associated with the key, setting it as the most
-	// recently used item.
-	// If the found cache item is already expired, the evict function is called
-	// and the return value indicates that the key was not found.
-	Get(key K) (V, bool)
+func (lru *Cache[K, V]) GetWithLifetimeNoExpire(key K) (V, time.Time, bool) {
+	if lru.sharded != nil {
+		return lru.sharded.GetWithLifetimeNoExpire(key)
+	}
+	return lru.synced.GetWithLifetimeNoExpire(key)
+}
 
-	GetWithLifetime(key K) (V, time.Time, bool)
+func (lru *Cache[K, V]) GetAndRefresh(key K) (V, bool) {
+	if lru.sharded != nil {
+		return lru.sharded.GetAndRefresh(key)
+	}
+	return lru.synced.GetAndRefresh(key)
+}
 
-	GetWithLifetimeNoExpire(key K) (V, time.Time, bool)
+func (lru *Cache[K, V]) GetAndRefreshOrAdd(key K, constructor func() (V, bool)) (V, bool, bool) {
+	if lru.sharded != nil {
+		return lru.sharded.GetAndRefreshOrAdd(key, constructor)
+	}
+	return lru.synced.GetAndRefreshOrAdd(key, constructor)
+}
 
-	// GetAndRefresh returns the value associated with the key, setting it as the most
-	// recently used item.
-	// The lifetime of the found cache item is refreshed, even if it was already expired.
-	GetAndRefresh(key K) (V, bool)
+func (lru *Cache[K, V]) Peek(key K) (V, bool) {
+	if lru.sharded != nil {
+		return lru.sharded.Peek(key)
+	}
+	return lru.synced.Peek(key)
+}
 
-	GetAndRefreshOrAdd(key K, constructor func() (V, bool)) (V, bool, bool)
+func (lru *Cache[K, V]) PeekWithLifetime(key K) (V, time.Time, bool) {
+	if lru.sharded != nil {
+		return lru.sharded.PeekWithLifetime(key)
+	}
+	return lru.synced.PeekWithLifetime(key)
+}
 
-	// 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)
+func (lru *Cache[K, V]) UpdateLifetime(key K, value V, lifetime time.Duration) bool {
+	if lru.sharded != nil {
+		return lru.sharded.UpdateLifetime(key, value, lifetime)
+	}
+	return lru.synced.UpdateLifetime(key, value, lifetime)
+}
 
-	PeekWithLifetime(key K) (V, time.Time, bool)
+func (lru *Cache[K, V]) Contains(key K) bool {
+	if lru.sharded != nil {
+		return lru.sharded.Contains(key)
+	}
+	return lru.synced.Contains(key)
+}
 
-	UpdateLifetime(key K, value V, lifetime time.Duration) bool
+func (lru *Cache[K, V]) Remove(key K) bool {
+	if lru.sharded != nil {
+		return lru.sharded.Remove(key)
+	}
+	return lru.synced.Remove(key)
+}
 
-	// Contains checks for the existence of a key, without changing its recent-ness.
-	// If the found entry is already expired, the evict function is called.
-	Contains(key K) bool
+func (lru *Cache[K, V]) RemoveOldest() (K, V, bool) {
+	if lru.sharded != nil {
+		return lru.sharded.RemoveOldest()
+	}
+	return lru.synced.RemoveOldest()
+}
 
-	// Remove removes the key from the cache.
-	// The return value indicates whether the key existed or not.
-	// The evict function is called for the removed entry.
-	Remove(key K) bool
+func (lru *Cache[K, V]) Keys() []K {
+	if lru.sharded != nil {
+		return lru.sharded.Keys()
+	}
+	return lru.synced.Keys()
+}
 
-	// RemoveOldest removes the oldest entry from the cache.
-	// Key, value and an indicator of whether the entry has been removed is returned.
-	// The evict function is called for the removed entry.
-	RemoveOldest() (key K, value V, removed bool)
+func (lru *Cache[K, V]) Purge() {
+	if lru.sharded != nil {
+		lru.sharded.Purge()
+	} else {
+		lru.synced.Purge()
+	}
+}
 
-	// Keys returns a slice of the keys in the cache, from oldest to newest.
-	// Expired entries are not included.
-	// The evict function is called for each expired item.
-	Keys() []K
+func (lru *Cache[K, V]) PurgeExpired() {
+	if lru.sharded != nil {
+		lru.sharded.PurgeExpired()
+	} else {
+		lru.synced.PurgeExpired()
+	}
+}
 
-	// Purge purges all data (key and value) from the LRU.
-	// The evict function is called for each expired item.
-	// The LRU metrics are reset.
-	Purge()
+func (lru *Cache[K, V]) Metrics() Metrics {
+	if lru.sharded != nil {
+		return lru.sharded.Metrics()
+	}
+	return lru.synced.Metrics()
+}
 
-	// PurgeExpired purges all expired items from the LRU.
-	// The evict function is called for each expired item.
-	PurgeExpired()
+func (lru *Cache[K, V]) ResetMetrics() Metrics {
+	if lru.sharded != nil {
+		return lru.sharded.ResetMetrics()
+	}
+	return lru.synced.ResetMetrics()
+}
 
-	// Metrics returns the metrics of the cache.
-	Metrics() Metrics
+func (lru *Cache[K, V]) dump() {
+	if lru.sharded != nil {
+		lru.sharded.dump()
+	} else {
+		lru.synced.dump()
+	}
+}
 
-	// ResetMetrics resets the metrics of the cache and returns the previous state.
-	ResetMetrics() Metrics
+func (lru *Cache[K, V]) PrintStats() {
+	if lru.sharded != nil {
+		lru.sharded.PrintStats()
+	} else {
+		lru.synced.PrintStats()
+	}
 }