1 package ttlcache
2
3 import (
4 "container/list"
5 "context"
6 "fmt"
7 "sync"
8 "time"
9
10 "golang.org/x/sync/singleflight"
11 )
12
13
14 const (
15 EvictionReasonDeleted EvictionReason = iota + 1
16 EvictionReasonCapacityReached
17 EvictionReasonExpired
18 )
19
20
21
22 type EvictionReason int
23
24
25
26 type Cache[K comparable, V any] struct {
27 items struct {
28 mu sync.RWMutex
29 values map[K]*list.Element
30
31
32
33
34 lru *list.List
35 expQueue expirationQueue[K, V]
36
37 timerCh chan time.Duration
38 }
39
40 metricsMu sync.RWMutex
41 metrics Metrics
42
43 events struct {
44 insertion struct {
45 mu sync.RWMutex
46 nextID uint64
47 fns map[uint64]func(*Item[K, V])
48 }
49 eviction struct {
50 mu sync.RWMutex
51 nextID uint64
52 fns map[uint64]func(EvictionReason, *Item[K, V])
53 }
54 }
55
56 stopCh chan struct{}
57 options options[K, V]
58 }
59
60
61 func New[K comparable, V any](opts ...Option[K, V]) *Cache[K, V] {
62 c := &Cache[K, V]{
63 stopCh: make(chan struct{}),
64 }
65 c.items.values = make(map[K]*list.Element)
66 c.items.lru = list.New()
67 c.items.expQueue = newExpirationQueue[K, V]()
68 c.items.timerCh = make(chan time.Duration, 1)
69 c.events.insertion.fns = make(map[uint64]func(*Item[K, V]))
70 c.events.eviction.fns = make(map[uint64]func(EvictionReason, *Item[K, V]))
71
72 applyOptions(&c.options, opts...)
73
74 return c
75 }
76
77
78
79
80
81 func (c *Cache[K, V]) updateExpirations(fresh bool, elem *list.Element) {
82 var oldExpiresAt time.Time
83
84 if !c.items.expQueue.isEmpty() {
85 oldExpiresAt = c.items.expQueue[0].Value.(*Item[K, V]).expiresAt
86 }
87
88 if fresh {
89 c.items.expQueue.push(elem)
90 } else {
91 c.items.expQueue.update(elem)
92 }
93
94 newExpiresAt := c.items.expQueue[0].Value.(*Item[K, V]).expiresAt
95
96
97 if newExpiresAt.IsZero() || (!oldExpiresAt.IsZero() && !newExpiresAt.Before(oldExpiresAt)) {
98 return
99 }
100
101 d := time.Until(newExpiresAt)
102
103
104
105
106
107
108
109 if len(c.items.timerCh) > 0 {
110
111
112
113 select {
114 case d1 := <-c.items.timerCh:
115 if d1 < d {
116 d = d1
117 }
118 default:
119 }
120 }
121
122
123
124
125 c.items.timerCh <- d
126 }
127
128
129
130
131 func (c *Cache[K, V]) set(key K, value V, ttl time.Duration) *Item[K, V] {
132 if ttl == DefaultTTL {
133 ttl = c.options.ttl
134 }
135
136 elem := c.get(key, false)
137 if elem != nil {
138
139 item := elem.Value.(*Item[K, V])
140 item.update(value, ttl)
141 c.updateExpirations(false, elem)
142
143 return item
144 }
145
146 if c.options.capacity != 0 && uint64(len(c.items.values)) >= c.options.capacity {
147
148 c.evict(EvictionReasonCapacityReached, c.items.lru.Back())
149 }
150
151 if ttl == PreviousOrDefaultTTL {
152 ttl = c.options.ttl
153 }
154
155
156 item := newItem(key, value, ttl, c.options.enableVersionTracking)
157 elem = c.items.lru.PushFront(item)
158 c.items.values[key] = elem
159 c.updateExpirations(true, elem)
160
161 c.metricsMu.Lock()
162 c.metrics.Insertions++
163 c.metricsMu.Unlock()
164
165 c.events.insertion.mu.RLock()
166 for _, fn := range c.events.insertion.fns {
167 fn(item)
168 }
169 c.events.insertion.mu.RUnlock()
170
171 return item
172 }
173
174
175
176
177
178
179 func (c *Cache[K, V]) get(key K, touch bool) *list.Element {
180 elem := c.items.values[key]
181 if elem == nil {
182 return nil
183 }
184
185 item := elem.Value.(*Item[K, V])
186 if item.isExpiredUnsafe() {
187 return nil
188 }
189
190 c.items.lru.MoveToFront(elem)
191
192 if touch && item.ttl > 0 {
193 item.touch()
194 c.updateExpirations(false, elem)
195 }
196
197 return elem
198 }
199
200
201
202
203
204
205
206
207
208
209 func (c *Cache[K, V]) getWithOpts(key K, lockAndLoad bool, opts ...Option[K, V]) *Item[K, V] {
210 getOpts := options[K, V]{
211 loader: c.options.loader,
212 disableTouchOnHit: c.options.disableTouchOnHit,
213 }
214
215 applyOptions(&getOpts, opts...)
216
217 if lockAndLoad {
218 c.items.mu.Lock()
219 }
220
221 elem := c.get(key, !getOpts.disableTouchOnHit)
222
223 if lockAndLoad {
224 c.items.mu.Unlock()
225 }
226
227 if elem == nil {
228 c.metricsMu.Lock()
229 c.metrics.Misses++
230 c.metricsMu.Unlock()
231
232 if lockAndLoad && getOpts.loader != nil {
233 return getOpts.loader.Load(c, key)
234 }
235
236 return nil
237 }
238
239 c.metricsMu.Lock()
240 c.metrics.Hits++
241 c.metricsMu.Unlock()
242
243 return elem.Value.(*Item[K, V])
244 }
245
246
247
248
249
250
251 func (c *Cache[K, V]) evict(reason EvictionReason, elems ...*list.Element) {
252 if len(elems) > 0 {
253 c.metricsMu.Lock()
254 c.metrics.Evictions += uint64(len(elems))
255 c.metricsMu.Unlock()
256
257 c.events.eviction.mu.RLock()
258 for i := range elems {
259 item := elems[i].Value.(*Item[K, V])
260 delete(c.items.values, item.key)
261 c.items.lru.Remove(elems[i])
262 c.items.expQueue.remove(elems[i])
263
264 for _, fn := range c.events.eviction.fns {
265 fn(reason, item)
266 }
267 }
268 c.events.eviction.mu.RUnlock()
269
270 return
271 }
272
273 c.metricsMu.Lock()
274 c.metrics.Evictions += uint64(len(c.items.values))
275 c.metricsMu.Unlock()
276
277 c.events.eviction.mu.RLock()
278 for _, elem := range c.items.values {
279 item := elem.Value.(*Item[K, V])
280
281 for _, fn := range c.events.eviction.fns {
282 fn(reason, item)
283 }
284 }
285 c.events.eviction.mu.RUnlock()
286
287 c.items.values = make(map[K]*list.Element)
288 c.items.lru.Init()
289 c.items.expQueue = newExpirationQueue[K, V]()
290 }
291
292
293
294
295
296 func (c *Cache[K, V]) delete(key K) {
297 elem := c.items.values[key]
298 if elem == nil {
299 return
300 }
301
302 c.evict(EvictionReasonDeleted, elem)
303 }
304
305
306
307
308
309
310
311
312
313 func (c *Cache[K, V]) Set(key K, value V, ttl time.Duration) *Item[K, V] {
314 c.items.mu.Lock()
315 defer c.items.mu.Unlock()
316
317 return c.set(key, value, ttl)
318 }
319
320
321
322
323
324 func (c *Cache[K, V]) Get(key K, opts ...Option[K, V]) *Item[K, V] {
325 return c.getWithOpts(key, true, opts...)
326 }
327
328
329
330 func (c *Cache[K, V]) Delete(key K) {
331 c.items.mu.Lock()
332 defer c.items.mu.Unlock()
333
334 c.delete(key)
335 }
336
337
338 func (c *Cache[K, V]) Has(key K) bool {
339 c.items.mu.RLock()
340 defer c.items.mu.RUnlock()
341
342 _, ok := c.items.values[key]
343 return ok
344 }
345
346
347
348
349
350
351
352
353 func (c *Cache[K, V]) GetOrSet(key K, value V, opts ...Option[K, V]) (*Item[K, V], bool) {
354 c.items.mu.Lock()
355 defer c.items.mu.Unlock()
356
357 elem := c.getWithOpts(key, false, opts...)
358 if elem != nil {
359 return elem, true
360 }
361
362 setOpts := options[K, V]{
363 ttl: c.options.ttl,
364 }
365 applyOptions(&setOpts, opts...)
366
367 item := c.set(key, value, setOpts.ttl)
368
369 return item, false
370 }
371
372
373
374
375
376
377
378
379 func (c *Cache[K, V]) GetAndDelete(key K, opts ...Option[K, V]) (*Item[K, V], bool) {
380 c.items.mu.Lock()
381
382 elem := c.getWithOpts(key, false, opts...)
383 if elem == nil {
384 c.items.mu.Unlock()
385
386 getOpts := options[K, V]{
387 loader: c.options.loader,
388 }
389 applyOptions(&getOpts, opts...)
390
391 if getOpts.loader != nil {
392 item := getOpts.loader.Load(c, key)
393 return item, item != nil
394 }
395
396 return nil, false
397 }
398
399 c.delete(key)
400 c.items.mu.Unlock()
401
402 return elem, true
403 }
404
405
406 func (c *Cache[K, V]) DeleteAll() {
407 c.items.mu.Lock()
408 c.evict(EvictionReasonDeleted)
409 c.items.mu.Unlock()
410 }
411
412
413 func (c *Cache[K, V]) DeleteExpired() {
414 c.items.mu.Lock()
415 defer c.items.mu.Unlock()
416
417 if c.items.expQueue.isEmpty() {
418 return
419 }
420
421 e := c.items.expQueue[0]
422 for e.Value.(*Item[K, V]).isExpiredUnsafe() {
423 c.evict(EvictionReasonExpired, e)
424
425 if c.items.expQueue.isEmpty() {
426 break
427 }
428
429
430 e = c.items.expQueue[0]
431 }
432 }
433
434
435
436
437 func (c *Cache[K, V]) Touch(key K) {
438 c.items.mu.Lock()
439 c.get(key, true)
440 c.items.mu.Unlock()
441 }
442
443
444 func (c *Cache[K, V]) Len() int {
445 c.items.mu.RLock()
446 defer c.items.mu.RUnlock()
447
448 return len(c.items.values)
449 }
450
451
452 func (c *Cache[K, V]) Keys() []K {
453 c.items.mu.RLock()
454 defer c.items.mu.RUnlock()
455
456 res := make([]K, 0, len(c.items.values))
457 for k := range c.items.values {
458 res = append(res, k)
459 }
460
461 return res
462 }
463
464
465
466 func (c *Cache[K, V]) Items() map[K]*Item[K, V] {
467 c.items.mu.RLock()
468 defer c.items.mu.RUnlock()
469
470 items := make(map[K]*Item[K, V], len(c.items.values))
471 for k := range c.items.values {
472 item := c.get(k, false)
473 if item != nil {
474 items[k] = item.Value.(*Item[K, V])
475 }
476 }
477
478 return items
479 }
480
481
482
483 func (c *Cache[K, V]) Range(fn func(item *Item[K, V]) bool) {
484 c.items.mu.RLock()
485
486
487 if c.items.lru.Len() == 0 {
488 c.items.mu.RUnlock()
489 return
490 }
491
492 for item := c.items.lru.Front(); item != c.items.lru.Back().Next(); item = item.Next() {
493 i := item.Value.(*Item[K, V])
494 c.items.mu.RUnlock()
495
496 if !fn(i) {
497 return
498 }
499
500 if item.Next() != nil {
501 c.items.mu.RLock()
502 }
503 }
504 }
505
506
507 func (c *Cache[K, V]) Metrics() Metrics {
508 c.metricsMu.RLock()
509 defer c.metricsMu.RUnlock()
510
511 return c.metrics
512 }
513
514
515
516
517 func (c *Cache[K, V]) Start() {
518 waitDur := func() time.Duration {
519 c.items.mu.RLock()
520 defer c.items.mu.RUnlock()
521
522 if !c.items.expQueue.isEmpty() &&
523 !c.items.expQueue[0].Value.(*Item[K, V]).expiresAt.IsZero() {
524 d := time.Until(c.items.expQueue[0].Value.(*Item[K, V]).expiresAt)
525 if d <= 0 {
526
527 return time.Microsecond
528 }
529
530 return d
531 }
532
533 if c.options.ttl > 0 {
534 return c.options.ttl
535 }
536
537 return time.Hour
538 }
539
540 timer := time.NewTimer(waitDur())
541 stop := func() {
542 if !timer.Stop() {
543
544 select {
545 case <-timer.C:
546 default:
547 }
548 }
549 }
550
551 defer stop()
552
553 for {
554 select {
555 case <-c.stopCh:
556 return
557 case d := <-c.items.timerCh:
558 stop()
559 timer.Reset(d)
560 case <-timer.C:
561 c.DeleteExpired()
562 stop()
563 timer.Reset(waitDur())
564 }
565 }
566 }
567
568
569
570 func (c *Cache[K, V]) Stop() {
571 c.stopCh <- struct{}{}
572 }
573
574
575
576
577
578
579
580
581
582
583 func (c *Cache[K, V]) OnInsertion(fn func(context.Context, *Item[K, V])) func() {
584 var (
585 wg sync.WaitGroup
586 ctx, cancel = context.WithCancel(context.Background())
587 )
588
589 c.events.insertion.mu.Lock()
590 id := c.events.insertion.nextID
591 c.events.insertion.fns[id] = func(item *Item[K, V]) {
592 wg.Add(1)
593 go func() {
594 fn(ctx, item)
595 wg.Done()
596 }()
597 }
598 c.events.insertion.nextID++
599 c.events.insertion.mu.Unlock()
600
601 return func() {
602 cancel()
603
604 c.events.insertion.mu.Lock()
605 delete(c.events.insertion.fns, id)
606 c.events.insertion.mu.Unlock()
607
608 wg.Wait()
609 }
610 }
611
612
613
614
615
616
617
618
619
620
621 func (c *Cache[K, V]) OnEviction(fn func(context.Context, EvictionReason, *Item[K, V])) func() {
622 var (
623 wg sync.WaitGroup
624 ctx, cancel = context.WithCancel(context.Background())
625 )
626
627 c.events.eviction.mu.Lock()
628 id := c.events.eviction.nextID
629 c.events.eviction.fns[id] = func(r EvictionReason, item *Item[K, V]) {
630 wg.Add(1)
631 go func() {
632 fn(ctx, r, item)
633 wg.Done()
634 }()
635 }
636 c.events.eviction.nextID++
637 c.events.eviction.mu.Unlock()
638
639 return func() {
640 cancel()
641
642 c.events.eviction.mu.Lock()
643 delete(c.events.eviction.fns, id)
644 c.events.eviction.mu.Unlock()
645
646 wg.Wait()
647 }
648 }
649
650
651 type Loader[K comparable, V any] interface {
652
653
654
655
656
657 Load(c *Cache[K, V], key K) *Item[K, V]
658 }
659
660
661
662 type LoaderFunc[K comparable, V any] func(*Cache[K, V], K) *Item[K, V]
663
664
665
666
667 func (l LoaderFunc[K, V]) Load(c *Cache[K, V], key K) *Item[K, V] {
668 return l(c, key)
669 }
670
671
672
673 type SuppressedLoader[K comparable, V any] struct {
674 loader Loader[K, V]
675 group *singleflight.Group
676 }
677
678
679
680
681 func NewSuppressedLoader[K comparable, V any](loader Loader[K, V], group *singleflight.Group) *SuppressedLoader[K, V] {
682 if group == nil {
683 group = &singleflight.Group{}
684 }
685
686 return &SuppressedLoader[K, V]{
687 loader: loader,
688 group: group,
689 }
690 }
691
692
693
694
695
696
697 func (l *SuppressedLoader[K, V]) Load(c *Cache[K, V], key K) *Item[K, V] {
698
699
700
701 strKey := fmt.Sprint(key)
702
703
704
705
706
707 res, _, _ := l.group.Do(strKey, func() (interface{}, error) {
708 item := l.loader.Load(c, key)
709 if item == nil {
710 return nil, nil
711 }
712
713 return item, nil
714 })
715 if res == nil {
716 return nil
717 }
718
719 return res.(*Item[K, V])
720 }
721
View as plain text