| 1 | // SPDX-License-Identifier: AGPL-3.0-or-later |
| 2 | |
| 3 | package lru |
| 4 | |
| 5 | import ( |
| 6 | "container/list" |
| 7 | "sync" |
| 8 | "sync/atomic" |
| 9 | ) |
| 10 | |
| 11 | // SizedCache is a byte-bounded LRU. Each entry contributes a caller- |
| 12 | // supplied size; eviction runs while the running total exceeds the |
| 13 | // cap. Use this for caches whose values vary widely in memory cost |
| 14 | // (rendered HTML, diff blobs, response bodies). |
| 15 | type SizedCache[K comparable] struct { |
| 16 | mu sync.Mutex |
| 17 | maxBytes int64 |
| 18 | cur int64 |
| 19 | ll *list.List |
| 20 | items map[K]*list.Element |
| 21 | |
| 22 | hits atomic.Uint64 |
| 23 | misses atomic.Uint64 |
| 24 | evictions atomic.Uint64 |
| 25 | } |
| 26 | |
| 27 | type sizedEntry[K comparable] struct { |
| 28 | key K |
| 29 | val []byte |
| 30 | size int64 |
| 31 | } |
| 32 | |
| 33 | // NewSized constructs a byte-bounded LRU. maxBytes must be positive. |
| 34 | func NewSized[K comparable](maxBytes int64) *SizedCache[K] { |
| 35 | if maxBytes <= 0 { |
| 36 | panic("lru: maxBytes must be positive") |
| 37 | } |
| 38 | return &SizedCache[K]{ |
| 39 | maxBytes: maxBytes, |
| 40 | ll: list.New(), |
| 41 | items: make(map[K]*list.Element), |
| 42 | } |
| 43 | } |
| 44 | |
| 45 | // Get returns the cached bytes + true on hit. The returned slice is |
| 46 | // the cached buffer (zero-copy) — callers MUST NOT mutate it. Use |
| 47 | // `append([]byte(nil), v...)` if mutation is needed. |
| 48 | func (c *SizedCache[K]) Get(key K) ([]byte, bool) { |
| 49 | c.mu.Lock() |
| 50 | defer c.mu.Unlock() |
| 51 | el, ok := c.items[key] |
| 52 | if !ok { |
| 53 | c.misses.Add(1) |
| 54 | return nil, false |
| 55 | } |
| 56 | c.ll.MoveToFront(el) |
| 57 | c.hits.Add(1) |
| 58 | return el.Value.(*sizedEntry[K]).val, true |
| 59 | } |
| 60 | |
| 61 | // Set stores key→val. Replacing an existing key updates its size. |
| 62 | // Eviction runs after insertion until total bytes ≤ maxBytes. |
| 63 | func (c *SizedCache[K]) Set(key K, val []byte) { |
| 64 | c.mu.Lock() |
| 65 | defer c.mu.Unlock() |
| 66 | size := int64(len(val)) |
| 67 | if el, ok := c.items[key]; ok { |
| 68 | e := el.Value.(*sizedEntry[K]) |
| 69 | c.cur += size - e.size |
| 70 | e.val = val |
| 71 | e.size = size |
| 72 | c.ll.MoveToFront(el) |
| 73 | } else { |
| 74 | e := &sizedEntry[K]{key: key, val: val, size: size} |
| 75 | el := c.ll.PushFront(e) |
| 76 | c.items[key] = el |
| 77 | c.cur += size |
| 78 | } |
| 79 | for c.cur > c.maxBytes && c.ll.Len() > 1 { |
| 80 | c.removeElement(c.ll.Back()) |
| 81 | c.evictions.Add(1) |
| 82 | } |
| 83 | } |
| 84 | |
| 85 | // Delete removes key. No-op when absent. |
| 86 | func (c *SizedCache[K]) Delete(key K) { |
| 87 | c.mu.Lock() |
| 88 | defer c.mu.Unlock() |
| 89 | if el, ok := c.items[key]; ok { |
| 90 | c.removeElement(el) |
| 91 | } |
| 92 | } |
| 93 | |
| 94 | // Bytes reports the current total payload size. |
| 95 | func (c *SizedCache[K]) Bytes() int64 { |
| 96 | c.mu.Lock() |
| 97 | defer c.mu.Unlock() |
| 98 | return c.cur |
| 99 | } |
| 100 | |
| 101 | // Len reports the entry count. |
| 102 | func (c *SizedCache[K]) Len() int { |
| 103 | c.mu.Lock() |
| 104 | defer c.mu.Unlock() |
| 105 | return c.ll.Len() |
| 106 | } |
| 107 | |
| 108 | // Stats returns a snapshot of the hit / miss / eviction counters. |
| 109 | func (c *SizedCache[K]) Stats() Stats { |
| 110 | return Stats{ |
| 111 | Hits: c.hits.Load(), |
| 112 | Misses: c.misses.Load(), |
| 113 | Evictions: c.evictions.Load(), |
| 114 | } |
| 115 | } |
| 116 | |
| 117 | func (c *SizedCache[K]) removeElement(el *list.Element) { |
| 118 | e := el.Value.(*sizedEntry[K]) |
| 119 | c.ll.Remove(el) |
| 120 | delete(c.items, e.key) |
| 121 | c.cur -= e.size |
| 122 | } |
| 123 |