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 float64, 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	return
142}
143
144// ensureSpace is used to ensure we have space in the cache
145func (c *TwoQueueCache) ensureSpace(recentEvict bool) {
146	// If we have space, nothing to do
147	recentLen := c.recent.Len()
148	freqLen := c.frequent.Len()
149	if recentLen+freqLen < c.size {
150		return
151	}
152
153	// If the recent buffer is larger than
154	// the target, evict from there
155	if recentLen > 0 && (recentLen > c.recentSize || (recentLen == c.recentSize && !recentEvict)) {
156		k, _, _ := c.recent.RemoveOldest()
157		c.recentEvict.Add(k, nil)
158		return
159	}
160
161	// Remove from the frequent list otherwise
162	c.frequent.RemoveOldest()
163}
164
165// Len returns the number of items in the cache.
166func (c *TwoQueueCache) Len() int {
167	c.lock.RLock()
168	defer c.lock.RUnlock()
169	return c.recent.Len() + c.frequent.Len()
170}
171
172// Keys returns a slice of the keys in the cache.
173// The frequently used keys are first in the returned slice.
174func (c *TwoQueueCache) Keys() []interface{} {
175	c.lock.RLock()
176	defer c.lock.RUnlock()
177	k1 := c.frequent.Keys()
178	k2 := c.recent.Keys()
179	return append(k1, k2...)
180}
181
182// Remove removes the provided key from the cache.
183func (c *TwoQueueCache) Remove(key interface{}) {
184	c.lock.Lock()
185	defer c.lock.Unlock()
186	if c.frequent.Remove(key) {
187		return
188	}
189	if c.recent.Remove(key) {
190		return
191	}
192	if c.recentEvict.Remove(key) {
193		return
194	}
195}
196
197// Purge is used to completely clear the cache.
198func (c *TwoQueueCache) Purge() {
199	c.lock.Lock()
200	defer c.lock.Unlock()
201	c.recent.Purge()
202	c.frequent.Purge()
203	c.recentEvict.Purge()
204}
205
206// Contains is used to check if the cache contains a key
207// without updating recency or frequency.
208func (c *TwoQueueCache) Contains(key interface{}) bool {
209	c.lock.RLock()
210	defer c.lock.RUnlock()
211	return c.frequent.Contains(key) || c.recent.Contains(key)
212}
213
214// Peek is used to inspect the cache value of a key
215// without updating recency or frequency.
216func (c *TwoQueueCache) Peek(key interface{}) (value interface{}, ok bool) {
217	c.lock.RLock()
218	defer c.lock.RUnlock()
219	if val, ok := c.frequent.Peek(key); ok {
220		return val, ok
221	}
222	return c.recent.Peek(key)
223}
224