Dragonfly2

1
/*
2
 *     Copyright 2020 The Dragonfly Authors
3
 *
4
 * Licensed under the Apache License, Version 2.0 (the "License");
5
 * you may not use this file except in compliance with the License.
6
 * You may obtain a copy of the License at
7
 *
8
 *      http://www.apache.org/licenses/LICENSE-2.0
9
 *
10
 * Unless required by applicable law or agreed to in writing, software
11
 * distributed under the License is distributed on an "AS IS" BASIS,
12
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
 * See the License for the specific language governing permissions and
14
 * limitations under the License.
15
 */
16

17
//go:generate mockgen -destination cache_mock.go -source cache.go -package cache
18
package cache
19

20
import (
21
	"encoding/gob"
22
	"fmt"
23
	"io"
24
	"io/fs"
25
	"os"
26
	"path/filepath"
27
	"regexp"
28
	"runtime"
29
	"sync"
30
	"time"
31
)
32

33
type Item struct {
34
	Object     any
35
	Expiration int64
36
}
37

38
// Returns true if the item has expired.
39
func (item Item) Expired() bool {
40
	if item.Expiration == 0 {
41
		return false
42
	}
43
	return time.Now().UnixNano() > item.Expiration
44
}
45

46
const (
47
	// For use with functions that take an expiration time.
48
	NoExpiration time.Duration = -1
49

50
	// For use with functions that take an expiration time. Equivalent to
51
	// passing in the same expiration duration as was given to New() or
52
	// NewFrom() when the cache was created (e.g. 5 minutes.)
53
	DefaultExpiration time.Duration = 0
54

55
	// For use with functions that do not clean up.
56
	NoCleanup time.Duration = -1
57
)
58

59
type Cache interface {
60
	Scan(m string, n int) ([]string, error)
61
	Set(k string, x any, d time.Duration)
62
	SetDefault(k string, x any)
63
	Add(k string, x any, d time.Duration) error
64
	Get(k string) (any, bool)
65
	GetWithExpiration(k string) (any, time.Time, bool)
66
	Delete(k string)
67
	DeleteExpired()
68
	Keys() []string
69
	OnEvicted(f func(string, any))
70
	Save(w io.Writer) (err error)
71
	SaveFile(fname string) error
72
	Load(r io.Reader) error
73
	LoadFile(fname string) error
74
	Items() map[string]Item
75
	ItemCount() int
76
	Flush()
77
}
78

79
type cache struct {
80
	defaultExpiration time.Duration
81
	items             map[string]Item
82
	mu                sync.RWMutex
83
	onEvicted         func(string, any)
84
	janitor           *janitor
85
}
86

87
// Scan all items to get a specified number of matching regexp key.
88
func (c *cache) Scan(m string, n int) ([]string, error) {
89
	c.mu.RLock()
90
	defer c.mu.RUnlock()
91

92
	regex, err := regexp.Compile(m)
93
	if err != nil {
94
		return nil, err
95
	}
96

97
	keys := make([]string, 0, n)
98
	for item := range c.items {
99
		if len(keys) >= n {
100
			break
101
		}
102

103
		if regex.MatchString(item) {
104
			keys = append(keys, item)
105
		}
106
	}
107

108
	return keys, nil
109
}
110

111
// Add an item to the cache, replacing any existing item. If the duration is 0
112
// (DefaultExpiration), the cache's default expiration time is used. If it is -1
113
// (NoExpiration), the item never expires.
114
func (c *cache) Set(k string, x any, d time.Duration) {
115
	// "Inlining" of set
116
	var e int64
117
	if d == DefaultExpiration {
118
		d = c.defaultExpiration
119
	}
120
	if d > 0 {
121
		e = time.Now().Add(d).UnixNano()
122
	}
123
	c.mu.Lock()
124
	c.items[k] = Item{
125
		Object:     x,
126
		Expiration: e,
127
	}
128
	// TODO: Calls to mu.Unlock are currently not deferred because defer
129
	// adds ~200 ns (as of go1.)
130
	c.mu.Unlock()
131
}
132

133
func (c *cache) set(k string, x any, d time.Duration) {
134
	var e int64
135
	if d == DefaultExpiration {
136
		d = c.defaultExpiration
137
	}
138
	if d > 0 {
139
		e = time.Now().Add(d).UnixNano()
140
	}
141
	c.items[k] = Item{
142
		Object:     x,
143
		Expiration: e,
144
	}
145
}
146

147
// Add an item to the cache, replacing any existing item, using the default
148
// expiration.
149
func (c *cache) SetDefault(k string, x any) {
150
	c.Set(k, x, DefaultExpiration)
151
}
152

153
// Add an item to the cache only if an item doesn't already exist for the given
154
// key, or if the existing item has expired. Returns an error otherwise.
155
func (c *cache) Add(k string, x any, d time.Duration) error {
156
	c.mu.Lock()
157
	_, found := c.get(k)
158
	if found {
159
		c.mu.Unlock()
160
		return fmt.Errorf("Item %s already exists", k)
161
	}
162
	c.set(k, x, d)
163
	c.mu.Unlock()
164
	return nil
165
}
166

167
// Get an item from the cache. Returns the item or nil, and a bool indicating
168
// whether the key was found.
169
func (c *cache) Get(k string) (any, bool) {
170
	c.mu.RLock()
171
	// "Inlining" of get and Expired
172
	item, found := c.items[k]
173
	if !found {
174
		c.mu.RUnlock()
175
		return nil, false
176
	}
177

178
	c.mu.RUnlock()
179
	return item.Object, true
180
}
181

182
// GetWithExpiration returns an item and its expiration time from the cache.
183
// It returns the item or nil, the expiration time if one is set (if the item
184
// never expires a zero value for time.Time is returned), and a bool indicating
185
// whether the key was found.
186
func (c *cache) GetWithExpiration(k string) (any, time.Time, bool) {
187
	c.mu.RLock()
188
	// "Inlining" of get and Expired
189
	item, found := c.items[k]
190
	if !found {
191
		c.mu.RUnlock()
192
		return nil, time.Time{}, false
193
	}
194

195
	if item.Expiration > 0 {
196
		if time.Now().UnixNano() > item.Expiration {
197
			c.mu.RUnlock()
198
			return nil, time.Time{}, false
199
		}
200

201
		// Return the item and the expiration time
202
		c.mu.RUnlock()
203
		return item.Object, time.Unix(0, item.Expiration), true
204
	}
205

206
	// If expiration <= 0 (i.e. no expiration time set) then return the item
207
	// and a zeroed time.Time
208
	c.mu.RUnlock()
209
	return item.Object, time.Time{}, true
210
}
211

212
func (c *cache) get(k string) (any, bool) {
213
	item, found := c.items[k]
214
	if !found {
215
		return nil, false
216
	}
217
	// "Inlining" of Expired
218
	if item.Expiration > 0 {
219
		if time.Now().UnixNano() > item.Expiration {
220
			return nil, false
221
		}
222
	}
223
	return item.Object, true
224
}
225

226
// Delete an item from the cache. Does nothing if the key is not in the cache.
227
func (c *cache) Delete(k string) {
228
	c.mu.Lock()
229
	v, evicted := c.delete(k)
230
	c.mu.Unlock()
231
	if evicted {
232
		c.onEvicted(k, v)
233
	}
234
}
235

236
func (c *cache) delete(k string) (any, bool) {
237
	if c.onEvicted != nil {
238
		if v, found := c.items[k]; found {
239
			delete(c.items, k)
240
			return v.Object, true
241
		}
242
	}
243
	delete(c.items, k)
244
	return nil, false
245
}
246

247
type keyAndValue struct {
248
	key   string
249
	value any
250
}
251

252
// Delete all expired items from the cache.
253
func (c *cache) DeleteExpired() {
254
	var evictedItems []keyAndValue
255
	now := time.Now().UnixNano()
256
	c.mu.Lock()
257
	for k, v := range c.items {
258
		// "Inlining" of expired
259
		if v.Expiration > 0 && now > v.Expiration {
260
			ov, evicted := c.delete(k)
261
			if evicted {
262
				evictedItems = append(evictedItems, keyAndValue{k, ov})
263
			}
264
		}
265
	}
266
	c.mu.Unlock()
267
	for _, v := range evictedItems {
268
		c.onEvicted(v.key, v.value)
269
	}
270
}
271

272
// Keys returns a sorted slice of all the keys in the cache.
273
func (c *cache) Keys() []string {
274
	c.mu.RLock()
275
	defer c.mu.RUnlock()
276
	keys := make([]string, len(c.items))
277
	var i int
278
	for k := range c.items {
279
		keys[i] = k
280
		i++
281
	}
282
	return keys
283
}
284

285
// Sets an (optional) function that is called with the key and value when an
286
// item is evicted from the cache. (Including when it is deleted manually, but
287
// not when it is overwritten.) Set to nil to disable.
288
func (c *cache) OnEvicted(f func(string, any)) {
289
	c.mu.Lock()
290
	c.onEvicted = f
291
	c.mu.Unlock()
292
}
293

294
// Write the cache's items (using Gob) to an io.Writer.
295
//
296
// NOTE: This method is deprecated in favor of c.Items() and NewFrom() (see the
297
// documentation for NewFrom().)
298
func (c *cache) Save(w io.Writer) (err error) {
299
	enc := gob.NewEncoder(w)
300
	defer func() {
301
		if x := recover(); x != nil {
302
			err = fmt.Errorf("Error registering item types with Gob library")
303
		}
304
	}()
305
	c.mu.RLock()
306
	defer c.mu.RUnlock()
307
	for _, v := range c.items {
308
		gob.Register(v.Object)
309
	}
310
	err = enc.Encode(&c.items)
311
	return
312
}
313

314
// Save the cache's items to the given filename, creating the file if it
315
// doesn't exist, and overwriting it if it does.
316
//
317
// NOTE: This method is deprecated in favor of c.Items() and NewFrom() (see the
318
// documentation for NewFrom().)
319
func (c *cache) SaveFile(fname string) error {
320
	dir := filepath.Dir(fname)
321
	if err := os.MkdirAll(dir, fs.FileMode(0700)); err != nil {
322
		return err
323
	}
324

325
	fp, err := os.Create(fname)
326
	if err != nil {
327
		return err
328
	}
329

330
	defer fp.Close()
331

332
	return c.Save(fp)
333
}
334

335
// Add (Gob-serialized) cache items from an io.Reader, excluding any items with
336
// keys that already exist (and haven't expired) in the current cache.
337
//
338
// NOTE: This method is deprecated in favor of c.Items() and NewFrom() (see the
339
// documentation for NewFrom().)
340
func (c *cache) Load(r io.Reader) error {
341
	dec := gob.NewDecoder(r)
342
	items := map[string]Item{}
343
	err := dec.Decode(&items)
344
	if err == nil {
345
		c.mu.Lock()
346
		defer c.mu.Unlock()
347
		for k, v := range items {
348
			ov, found := c.items[k]
349
			if !found || ov.Expired() {
350
				c.items[k] = v
351
			}
352
		}
353
	}
354
	return err
355
}
356

357
// Load and add cache items from the given filename, excluding any items with
358
// keys that already exist in the current cache.
359
//
360
// NOTE: This method is deprecated in favor of c.Items() and NewFrom() (see the
361
// documentation for NewFrom().)
362
func (c *cache) LoadFile(fname string) error {
363
	fp, err := os.Open(fname)
364
	if err != nil {
365
		return err
366
	}
367
	err = c.Load(fp)
368
	if err != nil {
369
		fp.Close()
370
		return err
371
	}
372
	return fp.Close()
373
}
374

375
// Copies all unexpired items in the cache into a new map and returns it.
376
func (c *cache) Items() map[string]Item {
377
	c.mu.RLock()
378
	defer c.mu.RUnlock()
379
	m := make(map[string]Item, len(c.items))
380
	now := time.Now().UnixNano()
381
	for k, v := range c.items {
382
		// "Inlining" of Expired
383
		if v.Expiration > 0 {
384
			if now > v.Expiration {
385
				continue
386
			}
387
		}
388
		m[k] = v
389
	}
390
	return m
391
}
392

393
// Returns the number of items in the cache. This may include items that have
394
// expired, but have not yet been cleaned up.
395
func (c *cache) ItemCount() int {
396
	c.mu.RLock()
397
	n := len(c.items)
398
	c.mu.RUnlock()
399
	return n
400
}
401

402
// Delete all items from the cache.
403
func (c *cache) Flush() {
404
	c.mu.Lock()
405
	c.items = map[string]Item{}
406
	c.mu.Unlock()
407
}
408

409
type janitor struct {
410
	Interval time.Duration
411
	stop     chan bool
412
}
413

414
func (j *janitor) Run(c *cache) {
415
	ticker := time.NewTicker(j.Interval)
416
	for {
417
		select {
418
		case <-ticker.C:
419
			c.DeleteExpired()
420
		case <-j.stop:
421
			ticker.Stop()
422
			return
423
		}
424
	}
425
}
426

427
func stopJanitor(c *cache) {
428
	c.janitor.stop <- true
429
}
430

431
func runJanitor(c *cache, ci time.Duration) {
432
	j := &janitor{
433
		Interval: ci,
434
		stop:     make(chan bool),
435
	}
436
	c.janitor = j
437
	go j.Run(c)
438
}
439

440
func newCache(de time.Duration) *cache {
441
	if de == 0 {
442
		de = -1
443
	}
444
	c := &cache{
445
		defaultExpiration: de,
446
		items:             make(map[string]Item),
447
	}
448
	return c
449
}
450

451
func newCacheWithJanitor(de time.Duration, ci time.Duration) Cache {
452
	c := newCache(de)
453
	// This trick ensures that the janitor goroutine (which--granted it
454
	// was enabled--is running DeleteExpired on c forever) does not keep
455
	// the returned C object from being garbage collected. When it is
456
	// garbage collected, the finalizer stops the janitor goroutine, after
457
	// which c can be collected.
458
	if ci > 0 {
459
		runJanitor(c, ci)
460
		runtime.SetFinalizer(c, stopJanitor)
461
	}
462
	return c
463
}
464

465
// Return a new cache with a given default expiration duration and cleanup
466
// interval. If the expiration duration is less than one (or NoExpiration),
467
// the items in the cache never expire (by default), and must be deleted
468
// manually. If the cleanup interval is less than one, expired items are not
469
// deleted from the cache before calling c.DeleteExpired().
470
func New(defaultExpiration, cleanupInterval time.Duration) Cache {
471
	return newCacheWithJanitor(defaultExpiration, cleanupInterval)
472
}
473