1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
|
// Package lru provides a weighted least-recently-used cache.
package lru
import "container/list"
// WeightFunc reports one entry's weight used for eviction budgeting.
//
// Returned weights MUST be non-negative.
type WeightFunc[K comparable, V any] func(key K, value V) int64
// OnEvictFunc runs when an entry leaves the cache.
//
// It is called for evictions, explicit removals, Clear, and replacement by Add.
type OnEvictFunc[K comparable, V any] func(key K, value V)
// Cache is a non-concurrent weighted LRU cache.
//
// Methods on Cache are not safe for concurrent use.
type Cache[K comparable, V any] struct {
maxWeight int64
weightFn WeightFunc[K, V]
onEvict OnEvictFunc[K, V]
weight int64
items map[K]*list.Element
lru list.List
}
type entry[K comparable, V any] struct {
key K
value V
weight int64
}
// New creates a cache with a maximum total weight.
//
// New panics if maxWeight is negative or weightFn is nil.
func New[K comparable, V any](maxWeight int64, weightFn WeightFunc[K, V], onEvict OnEvictFunc[K, V]) *Cache[K, V] {
if maxWeight < 0 {
panic("lru: negative max weight")
}
if weightFn == nil {
panic("lru: nil weight function")
}
return &Cache[K, V]{
maxWeight: maxWeight,
weightFn: weightFn,
onEvict: onEvict,
items: make(map[K]*list.Element),
}
}
// Add inserts or replaces key and marks it most-recently-used.
//
// Add returns false when the entry's weight exceeds MaxWeight even for an empty
// cache. In that case the cache is unchanged.
//
// Add panics if weightFn returns a negative weight.
func (cache *Cache[K, V]) Add(key K, value V) bool {
w := cache.weightFn(key, value)
if w < 0 {
panic("lru: negative entry weight")
}
if w > cache.maxWeight {
return false
}
if elem, ok := cache.items[key]; ok {
cache.removeElem(elem)
}
ent := &entry[K, V]{
key: key,
value: value,
weight: w,
}
elem := cache.lru.PushBack(ent)
cache.items[key] = elem
cache.weight += w
cache.evictOverBudget()
return true
}
// Get returns value for key and marks it most-recently-used.
func (cache *Cache[K, V]) Get(key K) (V, bool) {
elem, ok := cache.items[key]
if !ok {
var zero V
return zero, false
}
cache.lru.MoveToBack(elem)
return elem.Value.(*entry[K, V]).value, true
}
// Peek returns value for key without changing recency.
func (cache *Cache[K, V]) Peek(key K) (V, bool) {
elem, ok := cache.items[key]
if !ok {
var zero V
return zero, false
}
return elem.Value.(*entry[K, V]).value, true
}
// Remove deletes key from the cache.
func (cache *Cache[K, V]) Remove(key K) (V, bool) {
elem, ok := cache.items[key]
if !ok {
var zero V
return zero, false
}
ent := cache.removeElem(elem)
return ent.value, true
}
// Clear removes all entries from the cache.
func (cache *Cache[K, V]) Clear() {
for elem := cache.lru.Front(); elem != nil; {
next := elem.Next()
cache.removeElem(elem)
elem = next
}
}
// Len returns the number of cached entries.
func (cache *Cache[K, V]) Len() int {
return len(cache.items)
}
// Weight returns the current total weight.
func (cache *Cache[K, V]) Weight() int64 {
return cache.weight
}
// MaxWeight returns the configured total weight budget.
func (cache *Cache[K, V]) MaxWeight() int64 {
return cache.maxWeight
}
// SetMaxWeight updates the total weight budget and evicts LRU entries as
// needed.
//
// SetMaxWeight panics if maxWeight is negative.
func (cache *Cache[K, V]) SetMaxWeight(maxWeight int64) {
if maxWeight < 0 {
panic("lru: negative max weight")
}
cache.maxWeight = maxWeight
cache.evictOverBudget()
}
func (cache *Cache[K, V]) evictOverBudget() {
for cache.weight > cache.maxWeight {
elem := cache.lru.Front()
if elem == nil {
return
}
cache.removeElem(elem)
}
}
func (cache *Cache[K, V]) removeElem(elem *list.Element) *entry[K, V] {
ent := elem.Value.(*entry[K, V])
cache.lru.Remove(elem)
delete(cache.items, ent.key)
cache.weight -= ent.weight
if cache.onEvict != nil {
cache.onEvict(ent.key, ent.value)
}
return ent
}
|
