mirror of
https://github.com/MetaCubeX/mihomo.git
synced 2026-03-05 05:27:31 +00:00
feat: support ARC for DNS cache
This commit is contained in:
Larvan2
wwqgtxx
262
common/lru/lrucache.go
Normal file
262
common/lru/lrucache.go
Normal file
@@ -0,0 +1,262 @@
|
||||
package lru
|
||||
|
||||
// Modified by https://github.com/die-net/lrucache
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
list "github.com/bahlo/generic-list-go"
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
// Option is part of Functional Options Pattern
|
||||
type Option[K comparable, V any] func(*LruCache[K, V])
|
||||
|
||||
// EvictCallback is used to get a callback when a cache entry is evicted
|
||||
type EvictCallback[K comparable, V any] func(key K, value V)
|
||||
|
||||
// WithEvict set the evict callback
|
||||
func WithEvict[K comparable, V any](cb EvictCallback[K, V]) Option[K, V] {
|
||||
return func(l *LruCache[K, V]) {
|
||||
l.onEvict = cb
|
||||
}
|
||||
}
|
||||
|
||||
// WithUpdateAgeOnGet update expires when Get element
|
||||
func WithUpdateAgeOnGet[K comparable, V any]() Option[K, V] {
|
||||
return func(l *LruCache[K, V]) {
|
||||
l.updateAgeOnGet = true
|
||||
}
|
||||
}
|
||||
|
||||
// WithAge defined element max age (second)
|
||||
func WithAge[K comparable, V any](maxAge int64) Option[K, V] {
|
||||
return func(l *LruCache[K, V]) {
|
||||
l.maxAge = maxAge
|
||||
}
|
||||
}
|
||||
|
||||
// WithSize defined max length of LruCache
|
||||
func WithSize[K comparable, V any](maxSize int) Option[K, V] {
|
||||
return func(l *LruCache[K, V]) {
|
||||
l.maxSize = maxSize
|
||||
}
|
||||
}
|
||||
|
||||
// WithStale decide whether Stale return is enabled.
|
||||
// If this feature is enabled, element will not get Evicted according to `WithAge`.
|
||||
func WithStale[K comparable, V any](stale bool) Option[K, V] {
|
||||
return func(l *LruCache[K, V]) {
|
||||
l.staleReturn = stale
|
||||
}
|
||||
}
|
||||
|
||||
// LruCache is a thread-safe, in-memory lru-cache that evicts the
|
||||
// least recently used entries from memory when (if set) the entries are
|
||||
// older than maxAge (in seconds). Use the New constructor to create one.
|
||||
type LruCache[K comparable, V any] struct {
|
||||
maxAge int64
|
||||
maxSize int
|
||||
mu sync.Mutex
|
||||
cache map[K]*list.Element[*entry[K, V]]
|
||||
lru *list.List[*entry[K, V]] // Front is least-recent
|
||||
updateAgeOnGet bool
|
||||
staleReturn bool
|
||||
onEvict EvictCallback[K, V]
|
||||
}
|
||||
|
||||
// New creates an LruCache
|
||||
func New[K comparable, V any](options ...Option[K, V]) *LruCache[K, V] {
|
||||
lc := &LruCache[K, V]{
|
||||
lru: list.New[*entry[K, V]](),
|
||||
cache: make(map[K]*list.Element[*entry[K, V]]),
|
||||
}
|
||||
|
||||
for _, option := range options {
|
||||
option(lc)
|
||||
}
|
||||
|
||||
return lc
|
||||
}
|
||||
|
||||
// Get returns the any representation of a cached response and a bool
|
||||
// set to true if the key was found.
|
||||
func (c *LruCache[K, V]) Get(key K) (V, bool) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
el := c.get(key)
|
||||
if el == nil {
|
||||
return lo.Empty[V](), false
|
||||
}
|
||||
value := el.value
|
||||
|
||||
return value, true
|
||||
}
|
||||
|
||||
func (c *LruCache[K, V]) GetOrStore(key K, constructor func() V) (V, bool) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
el := c.get(key)
|
||||
if el == nil {
|
||||
value := constructor()
|
||||
c.set(key, value)
|
||||
return value, false
|
||||
}
|
||||
value := el.value
|
||||
|
||||
return value, true
|
||||
}
|
||||
|
||||
// GetWithExpire returns the any representation of a cached response,
|
||||
// a time.Time Give expected expires,
|
||||
// and a bool set to true if the key was found.
|
||||
// This method will NOT check the maxAge of element and will NOT update the expires.
|
||||
func (c *LruCache[K, V]) GetWithExpire(key K) (V, time.Time, bool) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
el := c.get(key)
|
||||
if el == nil {
|
||||
return lo.Empty[V](), time.Time{}, false
|
||||
}
|
||||
|
||||
return el.value, time.Unix(el.expires, 0), true
|
||||
}
|
||||
|
||||
// Exist returns if key exist in cache but not put item to the head of linked list
|
||||
func (c *LruCache[K, V]) Exist(key K) bool {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
_, ok := c.cache[key]
|
||||
return ok
|
||||
}
|
||||
|
||||
// Set stores the any representation of a response for a given key.
|
||||
func (c *LruCache[K, V]) Set(key K, value V) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
c.set(key, value)
|
||||
}
|
||||
|
||||
func (c *LruCache[K, V]) set(key K, value V) {
|
||||
expires := int64(0)
|
||||
if c.maxAge > 0 {
|
||||
expires = time.Now().Unix() + c.maxAge
|
||||
}
|
||||
c.setWithExpire(key, value, time.Unix(expires, 0))
|
||||
}
|
||||
|
||||
// SetWithExpire stores the any representation of a response for a given key and given expires.
|
||||
// The expires time will round to second.
|
||||
func (c *LruCache[K, V]) SetWithExpire(key K, value V, expires time.Time) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
c.setWithExpire(key, value, expires)
|
||||
}
|
||||
|
||||
func (c *LruCache[K, V]) setWithExpire(key K, value V, expires time.Time) {
|
||||
if le, ok := c.cache[key]; ok {
|
||||
c.lru.MoveToBack(le)
|
||||
e := le.Value
|
||||
e.value = value
|
||||
e.expires = expires.Unix()
|
||||
} else {
|
||||
e := &entry[K, V]{key: key, value: value, expires: expires.Unix()}
|
||||
c.cache[key] = c.lru.PushBack(e)
|
||||
|
||||
if c.maxSize > 0 {
|
||||
if elLen := c.lru.Len(); elLen > c.maxSize {
|
||||
c.deleteElement(c.lru.Front())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
c.maybeDeleteOldest()
|
||||
}
|
||||
|
||||
// CloneTo clone and overwrite elements to another LruCache
|
||||
func (c *LruCache[K, V]) CloneTo(n *LruCache[K, V]) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
n.mu.Lock()
|
||||
defer n.mu.Unlock()
|
||||
|
||||
n.lru = list.New[*entry[K, V]]()
|
||||
n.cache = make(map[K]*list.Element[*entry[K, V]])
|
||||
|
||||
for e := c.lru.Front(); e != nil; e = e.Next() {
|
||||
elm := e.Value
|
||||
n.cache[elm.key] = n.lru.PushBack(elm)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *LruCache[K, V]) get(key K) *entry[K, V] {
|
||||
le, ok := c.cache[key]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
if !c.staleReturn && c.maxAge > 0 && le.Value.expires <= time.Now().Unix() {
|
||||
c.deleteElement(le)
|
||||
c.maybeDeleteOldest()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
c.lru.MoveToBack(le)
|
||||
el := le.Value
|
||||
if c.maxAge > 0 && c.updateAgeOnGet {
|
||||
el.expires = time.Now().Unix() + c.maxAge
|
||||
}
|
||||
return el
|
||||
}
|
||||
|
||||
// Delete removes the value associated with a key.
|
||||
func (c *LruCache[K, V]) Delete(key K) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
if le, ok := c.cache[key]; ok {
|
||||
c.deleteElement(le)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *LruCache[K, V]) maybeDeleteOldest() {
|
||||
if !c.staleReturn && c.maxAge > 0 {
|
||||
now := time.Now().Unix()
|
||||
for le := c.lru.Front(); le != nil && le.Value.expires <= now; le = c.lru.Front() {
|
||||
c.deleteElement(le)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *LruCache[K, V]) deleteElement(le *list.Element[*entry[K, V]]) {
|
||||
c.lru.Remove(le)
|
||||
e := le.Value
|
||||
delete(c.cache, e.key)
|
||||
if c.onEvict != nil {
|
||||
c.onEvict(e.key, e.value)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *LruCache[K, V]) Clear() error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
c.cache = make(map[K]*list.Element[*entry[K, V]])
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type entry[K comparable, V any] struct {
|
||||
key K
|
||||
value V
|
||||
expires int64
|
||||
}
|
||||
184
common/lru/lrucache_test.go
Normal file
184
common/lru/lrucache_test.go
Normal file
@@ -0,0 +1,184 @@
|
||||
package lru
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
var entries = []struct {
|
||||
key string
|
||||
value string
|
||||
}{
|
||||
{"1", "one"},
|
||||
{"2", "two"},
|
||||
{"3", "three"},
|
||||
{"4", "four"},
|
||||
{"5", "five"},
|
||||
}
|
||||
|
||||
func TestLRUCache(t *testing.T) {
|
||||
c := New[string, string]()
|
||||
|
||||
for _, e := range entries {
|
||||
c.Set(e.key, e.value)
|
||||
}
|
||||
|
||||
c.Delete("missing")
|
||||
_, ok := c.Get("missing")
|
||||
assert.False(t, ok)
|
||||
|
||||
for _, e := range entries {
|
||||
value, ok := c.Get(e.key)
|
||||
if assert.True(t, ok) {
|
||||
assert.Equal(t, e.value, value)
|
||||
}
|
||||
}
|
||||
|
||||
for _, e := range entries {
|
||||
c.Delete(e.key)
|
||||
|
||||
_, ok := c.Get(e.key)
|
||||
assert.False(t, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLRUMaxAge(t *testing.T) {
|
||||
c := New[string, string](WithAge[string, string](86400))
|
||||
|
||||
now := time.Now().Unix()
|
||||
expected := now + 86400
|
||||
|
||||
// Add one expired entry
|
||||
c.Set("foo", "bar")
|
||||
c.lru.Back().Value.expires = now
|
||||
|
||||
// Reset
|
||||
c.Set("foo", "bar")
|
||||
e := c.lru.Back().Value
|
||||
assert.True(t, e.expires >= now)
|
||||
c.lru.Back().Value.expires = now
|
||||
|
||||
// Set a few and verify expiration times
|
||||
for _, s := range entries {
|
||||
c.Set(s.key, s.value)
|
||||
e := c.lru.Back().Value
|
||||
assert.True(t, e.expires >= expected && e.expires <= expected+10)
|
||||
}
|
||||
|
||||
// Make sure we can get them all
|
||||
for _, s := range entries {
|
||||
_, ok := c.Get(s.key)
|
||||
assert.True(t, ok)
|
||||
}
|
||||
|
||||
// Expire all entries
|
||||
for _, s := range entries {
|
||||
le, ok := c.cache[s.key]
|
||||
if assert.True(t, ok) {
|
||||
le.Value.expires = now
|
||||
}
|
||||
}
|
||||
|
||||
// Get one expired entry, which should clear all expired entries
|
||||
_, ok := c.Get("3")
|
||||
assert.False(t, ok)
|
||||
assert.Equal(t, c.lru.Len(), 0)
|
||||
}
|
||||
|
||||
func TestLRUpdateOnGet(t *testing.T) {
|
||||
c := New[string, string](WithAge[string, string](86400), WithUpdateAgeOnGet[string, string]())
|
||||
|
||||
now := time.Now().Unix()
|
||||
expires := now + 86400/2
|
||||
|
||||
// Add one expired entry
|
||||
c.Set("foo", "bar")
|
||||
c.lru.Back().Value.expires = expires
|
||||
|
||||
_, ok := c.Get("foo")
|
||||
assert.True(t, ok)
|
||||
assert.True(t, c.lru.Back().Value.expires > expires)
|
||||
}
|
||||
|
||||
func TestMaxSize(t *testing.T) {
|
||||
c := New[string, string](WithSize[string, string](2))
|
||||
// Add one expired entry
|
||||
c.Set("foo", "bar")
|
||||
_, ok := c.Get("foo")
|
||||
assert.True(t, ok)
|
||||
|
||||
c.Set("bar", "foo")
|
||||
c.Set("baz", "foo")
|
||||
|
||||
_, ok = c.Get("foo")
|
||||
assert.False(t, ok)
|
||||
}
|
||||
|
||||
func TestExist(t *testing.T) {
|
||||
c := New[int, int](WithSize[int, int](1))
|
||||
c.Set(1, 2)
|
||||
assert.True(t, c.Exist(1))
|
||||
c.Set(2, 3)
|
||||
assert.False(t, c.Exist(1))
|
||||
}
|
||||
|
||||
func TestEvict(t *testing.T) {
|
||||
temp := 0
|
||||
evict := func(key int, value int) {
|
||||
temp = key + value
|
||||
}
|
||||
|
||||
c := New[int, int](WithEvict[int, int](evict), WithSize[int, int](1))
|
||||
c.Set(1, 2)
|
||||
c.Set(2, 3)
|
||||
|
||||
assert.Equal(t, temp, 3)
|
||||
}
|
||||
|
||||
func TestSetWithExpire(t *testing.T) {
|
||||
c := New[int, *struct{}](WithAge[int, *struct{}](1))
|
||||
now := time.Now().Unix()
|
||||
|
||||
tenSecBefore := time.Unix(now-10, 0)
|
||||
c.SetWithExpire(1, &struct{}{}, tenSecBefore)
|
||||
|
||||
// res is expected not to exist, and expires should be empty time.Time
|
||||
res, expires, exist := c.GetWithExpire(1)
|
||||
|
||||
assert.True(t, nil == res)
|
||||
assert.Equal(t, time.Time{}, expires)
|
||||
assert.Equal(t, false, exist)
|
||||
}
|
||||
|
||||
func TestStale(t *testing.T) {
|
||||
c := New[int, int](WithAge[int, int](1), WithStale[int, int](true))
|
||||
now := time.Now().Unix()
|
||||
|
||||
tenSecBefore := time.Unix(now-10, 0)
|
||||
c.SetWithExpire(1, 2, tenSecBefore)
|
||||
|
||||
res, expires, exist := c.GetWithExpire(1)
|
||||
assert.Equal(t, 2, res)
|
||||
assert.Equal(t, tenSecBefore, expires)
|
||||
assert.Equal(t, true, exist)
|
||||
}
|
||||
|
||||
func TestCloneTo(t *testing.T) {
|
||||
o := New[string, int](WithSize[string, int](10))
|
||||
o.Set("1", 1)
|
||||
o.Set("2", 2)
|
||||
|
||||
n := New[string, int](WithSize[string, int](2))
|
||||
n.Set("3", 3)
|
||||
n.Set("4", 4)
|
||||
|
||||
o.CloneTo(n)
|
||||
|
||||
assert.False(t, n.Exist("3"))
|
||||
assert.True(t, n.Exist("1"))
|
||||
|
||||
n.Set("5", 5)
|
||||
assert.False(t, n.Exist("1"))
|
||||
}
|
||||
Reference in New Issue
Block a user