1package lru
2
3import (
4	"fmt"
5	"sync"
6
7	"github.com/hashicorp/golang-lru/simplelru"
8)
9
10const (
11	// Default2QRecentRatio is the ratio of the 2Q cache dedicated
12	// to recently added entries that have only been accessed once.
13	Default2QRecentRatio = 0.25
14
15	// Default2QGhostEntries is the default ratio of ghost
16	// entries kept to track entries recently evicted
17	Default2QGhostEntries = 0.50
18)
19
20// TwoQueueCache is a thread-safe fixed size 2Q cache.
21// 2Q is an enhancement over the standard LRU cache
22// in that it tracks both frequently and recently used
23// entries separately. This avoids a burst in access to new
24// entries from evicting frequently used entries. It adds some
25// additional tracking overhead to the standard LRU cache, and is
26// computationally about 2x the cost, and adds some metadata over
27// head. The ARCCache is similar, but does not require setting any
28// parameters.
29type TwoQueueCache struct {
30	size       int
31	recentSize int
32
33	recent      simplelru.LRUCache
34	frequent    simplelru.LRUCache
35	recentEvict simplelru.LRUCache
36	lock        sync.RWMutex
37}
38
39// New2Q creates a new TwoQueueCache using the default
40// values for the parameters.
41func New2Q(size int) (*TwoQueueCache, error) {
42	return New2QParams(size, Default2QRecentRatio, Default2QGhostEntries)
43}
44
45// New2QParams creates a new TwoQueueCache using the provided
46// parameter values.
47func New2QParams(size int, recentRatio, ghostRatio float64) (*TwoQueueCache, error) {
48	if size <= 0 {
49		return nil, fmt.Errorf("invalid size")
50	}
51	if recentRatio < 0.0 || recentRatio > 1.0 {
52		return nil, fmt.Errorf("invalid recent ratio")
53	}
54	if ghostRatio < 0.0 || ghostRatio > 1.0 {
55		return nil, fmt.Errorf("invalid ghost ratio")
56	}
57
58	// Determine the sub-sizes
59	recentSize := int(float64(size) * recentRatio)
60	evictSize := int(float64(size) * ghostRatio)
61
62	// Allocate the LRUs
63	recent, err := simplelru.NewLRU(size, nil)
64	if err != nil {
65		return nil, err
66	}
67	frequent, err := simplelru.NewLRU(size, nil)
68	if err != nil {
69		return nil, err
70	}
71	recentEvict, err := simplelru.NewLRU(evictSize, nil)
72	if err != nil {
73		return nil, err
74	}
75
76	// Initialize the cache
77	c := &TwoQueueCache{
78		size:        size,
79		recentSize:  recentSize,
80		recent:      recent,
81		frequent:    frequent,
82		recentEvict: recentEvict,
83	}
84	return c, nil
85}
86
87// Get looks up a key's value from the cache.
88func (c *TwoQueueCache) Get(key interface{}) (value interface{}, ok bool) {
89	c.lock.Lock()
90	defer c.lock.Unlock()
91
92	// Check if this is a frequent value
93	if val, ok := c.frequent.Get(key); ok {
94		return val, ok
95	}
96
97	// If the value is contained in recent, then we
98	// promote it to frequent
99	if val, ok := c.recent.Peek(key); ok {
100		c.recent.Remove(key)
101		c.frequent.Add(key, val)
102		return val, ok
103	}
104
105	// No hit
106	return nil, false
107}
108
109// Add adds a value to the cache.
110func (c *TwoQueueCache) Add(key, value interface{}) {
111	c.lock.Lock()
112	defer c.lock.Unlock()
113
114	// Check if the value is frequently used already,
115	// and just update the value
116	if c.frequent.Contains(key) {
117		c.frequent.Add(key, value)
118		return
119	}
120
121	// Check if the value is recently used, and promote
122	// the value into the frequent list
123	if c.recent.Contains(key) {
124		c.recent.Remove(key)
125		c.frequent.Add(key, value)
126		return
127	}
128
129	// If the value was recently evicted, add it to the
130	// frequently used list
131	if c.recentEvict.Contains(key) {
132		c.ensureSpace(true)
133		c.recentEvict.Remove(key)
134		c.frequent.Add(key, value)
135		return
136	}
137
138	// Add to the recently seen list
139	c.ensureSpace(false)
140	c.recent.Add(key, value)
141}
142
143// ensureSpace is used to ensure we have space in the cache
144func (c *TwoQueueCache) ensureSpace(recentEvict bool) {
145	// If we have space, nothing to do
146	recentLen := c.recent.Len()
147	freqLen := c.frequent.Len()
148	if recentLen+freqLen < c.size {
149		return
150	}
151
152	// If the recent buffer is larger than
153	// the target, evict from there
154	if recentLen > 0 && (recentLen > c.recentSize || (recentLen == c.recentSize && !recentEvict)) {
155		k, _, _ := c.recent.RemoveOldest()
156		c.recentEvict.Add(k, nil)
157		return
158	}
159
160	// Remove from the frequent list otherwise
161	c.frequent.RemoveOldest()
162}
163
164// Len returns the number of items in the cache.
165func (c *TwoQueueCache) Len() int {
166	c.lock.RLock()
167	defer c.lock.RUnlock()
168	return c.recent.Len() + c.frequent.Len()
169}
170
171// Keys returns a slice of the keys in the cache.
172// The frequently used keys are first in the returned slice.
173func (c *TwoQueueCache) Keys() []interface{} {
174	c.lock.RLock()
175	defer c.lock.RUnlock()
176	k1 := c.frequent.Keys()
177	k2 := c.recent.Keys()
178	return append(k1, k2...)
179}
180
181// Remove removes the provided key from the cache.
182func (c *TwoQueueCache) Remove(key interface{}) {
183	c.lock.Lock()
184	defer c.lock.Unlock()
185	if c.frequent.Remove(key) {
186		return
187	}
188	if c.recent.Remove(key) {
189		return
190	}
191	if c.recentEvict.Remove(key) {
192		return
193	}
194}
195
196// Purge is used to completely clear the cache.
197func (c *TwoQueueCache) Purge() {
198	c.lock.Lock()
199	defer c.lock.Unlock()
200	c.recent.Purge()
201	c.frequent.Purge()
202	c.recentEvict.Purge()
203}
204
205// Contains is used to check if the cache contains a key
206// without updating recency or frequency.
207func (c *TwoQueueCache) Contains(key interface{}) bool {
208	c.lock.RLock()
209	defer c.lock.RUnlock()
210	return c.frequent.Contains(key) || c.recent.Contains(key)
211}
212
213// Peek is used to inspect the cache value of a key
214// without updating recency or frequency.
215func (c *TwoQueueCache) Peek(key interface{}) (value interface{}, ok bool) {
216	c.lock.RLock()
217	defer c.lock.RUnlock()
218	if val, ok := c.frequent.Peek(key); ok {
219		return val, ok
220	}
221	return c.recent.Peek(key)
222}
223