1package lru
2
3import (
4	"sync"
5
6	"github.com/hashicorp/golang-lru/simplelru"
7)
8
9const (
10	// DefaultEvictedBufferSize defines the default buffer size to store evicted key/val
11	DefaultEvictedBufferSize = 16
12)
13
14// Cache is a thread-safe fixed size LRU cache.
15type Cache struct {
16	lru                      *simplelru.LRU
17	evictedKeys, evictedVals []interface{}
18	onEvictedCB              func(k, v interface{})
19	lock                     sync.RWMutex
20}
21
22// New creates an LRU of the given size.
23func New(size int) (*Cache, error) {
24	return NewWithEvict(size, nil)
25}
26
27// NewWithEvict constructs a fixed size cache with the given eviction
28// callback.
29func NewWithEvict(size int, onEvicted func(key, value interface{})) (c *Cache, err error) {
30	// create a cache with default settings
31	c = &Cache{
32		onEvictedCB: onEvicted,
33	}
34	if onEvicted != nil {
35		c.initEvictBuffers()
36		onEvicted = c.onEvicted
37	}
38	c.lru, err = simplelru.NewLRU(size, onEvicted)
39	return
40}
41
42func (c *Cache) initEvictBuffers() {
43	c.evictedKeys = make([]interface{}, 0, DefaultEvictedBufferSize)
44	c.evictedVals = make([]interface{}, 0, DefaultEvictedBufferSize)
45}
46
47// onEvicted save evicted key/val and sent in externally registered callback
48// outside of critical section
49func (c *Cache) onEvicted(k, v interface{}) {
50	c.evictedKeys = append(c.evictedKeys, k)
51	c.evictedVals = append(c.evictedVals, v)
52}
53
54// Purge is used to completely clear the cache.
55func (c *Cache) Purge() {
56	var ks, vs []interface{}
57	c.lock.Lock()
58	c.lru.Purge()
59	if c.onEvictedCB != nil && len(c.evictedKeys) > 0 {
60		ks, vs = c.evictedKeys, c.evictedVals
61		c.initEvictBuffers()
62	}
63	c.lock.Unlock()
64	// invoke callback outside of critical section
65	if c.onEvictedCB != nil {
66		for i := 0; i < len(ks); i++ {
67			c.onEvictedCB(ks[i], vs[i])
68		}
69	}
70}
71
72// Add adds a value to the cache. Returns true if an eviction occurred.
73func (c *Cache) Add(key, value interface{}) (evicted bool) {
74	var k, v interface{}
75	c.lock.Lock()
76	evicted = c.lru.Add(key, value)
77	if c.onEvictedCB != nil && evicted {
78		k, v = c.evictedKeys[0], c.evictedVals[0]
79		c.evictedKeys, c.evictedVals = c.evictedKeys[:0], c.evictedVals[:0]
80	}
81	c.lock.Unlock()
82	if c.onEvictedCB != nil && evicted {
83		c.onEvictedCB(k, v)
84	}
85	return
86}
87
88// Get looks up a key's value from the cache.
89func (c *Cache) Get(key interface{}) (value interface{}, ok bool) {
90	c.lock.Lock()
91	value, ok = c.lru.Get(key)
92	c.lock.Unlock()
93	return value, ok
94}
95
96// Contains checks if a key is in the cache, without updating the
97// recent-ness or deleting it for being stale.
98func (c *Cache) Contains(key interface{}) bool {
99	c.lock.RLock()
100	containKey := c.lru.Contains(key)
101	c.lock.RUnlock()
102	return containKey
103}
104
105// Peek returns the key value (or undefined if not found) without updating
106// the "recently used"-ness of the key.
107func (c *Cache) Peek(key interface{}) (value interface{}, ok bool) {
108	c.lock.RLock()
109	value, ok = c.lru.Peek(key)
110	c.lock.RUnlock()
111	return value, ok
112}
113
114// ContainsOrAdd checks if a key is in the cache without updating the
115// recent-ness or deleting it for being stale, and if not, adds the value.
116// Returns whether found and whether an eviction occurred.
117func (c *Cache) ContainsOrAdd(key, value interface{}) (ok, evicted bool) {
118	var k, v interface{}
119	c.lock.Lock()
120	if c.lru.Contains(key) {
121		c.lock.Unlock()
122		return true, false
123	}
124	evicted = c.lru.Add(key, value)
125	if c.onEvictedCB != nil && evicted {
126		k, v = c.evictedKeys[0], c.evictedVals[0]
127		c.evictedKeys, c.evictedVals = c.evictedKeys[:0], c.evictedVals[:0]
128	}
129	c.lock.Unlock()
130	if c.onEvictedCB != nil && evicted {
131		c.onEvictedCB(k, v)
132	}
133	return false, evicted
134}
135
136// PeekOrAdd checks if a key is in the cache without updating the
137// recent-ness or deleting it for being stale, and if not, adds the value.
138// Returns whether found and whether an eviction occurred.
139func (c *Cache) PeekOrAdd(key, value interface{}) (previous interface{}, ok, evicted bool) {
140	var k, v interface{}
141	c.lock.Lock()
142	previous, ok = c.lru.Peek(key)
143	if ok {
144		c.lock.Unlock()
145		return previous, true, false
146	}
147	evicted = c.lru.Add(key, value)
148	if c.onEvictedCB != nil && evicted {
149		k, v = c.evictedKeys[0], c.evictedVals[0]
150		c.evictedKeys, c.evictedVals = c.evictedKeys[:0], c.evictedVals[:0]
151	}
152	c.lock.Unlock()
153	if c.onEvictedCB != nil && evicted {
154		c.onEvictedCB(k, v)
155	}
156	return nil, false, evicted
157}
158
159// Remove removes the provided key from the cache.
160func (c *Cache) Remove(key interface{}) (present bool) {
161	var k, v interface{}
162	c.lock.Lock()
163	present = c.lru.Remove(key)
164	if c.onEvictedCB != nil && present {
165		k, v = c.evictedKeys[0], c.evictedVals[0]
166		c.evictedKeys, c.evictedVals = c.evictedKeys[:0], c.evictedVals[:0]
167	}
168	c.lock.Unlock()
169	if c.onEvictedCB != nil && present {
170		c.onEvicted(k, v)
171	}
172	return
173}
174
175// Resize changes the cache size.
176func (c *Cache) Resize(size int) (evicted int) {
177	var ks, vs []interface{}
178	c.lock.Lock()
179	evicted = c.lru.Resize(size)
180	if c.onEvictedCB != nil && evicted > 0 {
181		ks, vs = c.evictedKeys, c.evictedVals
182		c.initEvictBuffers()
183	}
184	c.lock.Unlock()
185	if c.onEvictedCB != nil && evicted > 0 {
186		for i := 0; i < len(ks); i++ {
187			c.onEvictedCB(ks[i], vs[i])
188		}
189	}
190	return evicted
191}
192
193// RemoveOldest removes the oldest item from the cache.
194func (c *Cache) RemoveOldest() (key, value interface{}, ok bool) {
195	var k, v interface{}
196	c.lock.Lock()
197	key, value, ok = c.lru.RemoveOldest()
198	if c.onEvictedCB != nil && ok {
199		k, v = c.evictedKeys[0], c.evictedVals[0]
200		c.evictedKeys, c.evictedVals = c.evictedKeys[:0], c.evictedVals[:0]
201	}
202	c.lock.Unlock()
203	if c.onEvictedCB != nil && ok {
204		c.onEvictedCB(k, v)
205	}
206	return
207}
208
209// GetOldest returns the oldest entry
210func (c *Cache) GetOldest() (key, value interface{}, ok bool) {
211	c.lock.RLock()
212	key, value, ok = c.lru.GetOldest()
213	c.lock.RUnlock()
214	return
215}
216
217// Keys returns a slice of the keys in the cache, from oldest to newest.
218func (c *Cache) Keys() []interface{} {
219	c.lock.RLock()
220	keys := c.lru.Keys()
221	c.lock.RUnlock()
222	return keys
223}
224
225// Len returns the number of items in the cache.
226func (c *Cache) Len() int {
227	c.lock.RLock()
228	length := c.lru.Len()
229	c.lock.RUnlock()
230	return length
231}
232