cache.go
| 1 | package util |
| 2 | |
| 3 | import ( |
| 4 | "sync" |
| 5 | "time" |
| 6 | ) |
| 7 | |
| 8 | type cacheEntry[V any] struct { |
| 9 | value V |
| 10 | size int64 |
| 11 | // expiresAt is the zero time when the cache has no TTL. |
| 12 | expiresAt time.Time |
| 13 | } |
| 14 | |
| 15 | // Cache is a bounded, concurrency-safe cache that evicts in insertion order. |
| 16 | // |
| 17 | // ponytail: FIFO eviction, swap for LRU if hit rate matters. |
| 18 | type Cache[K comparable, V any] struct { |
| 19 | mu sync.Mutex |
| 20 | max int |
| 21 | ttl time.Duration |
| 22 | maxBytes int64 |
| 23 | sizeOf func(V) int64 |
| 24 | bytes int64 |
| 25 | items map[K]cacheEntry[V] |
| 26 | order []K |
| 27 | } |
| 28 | |
| 29 | // NewCache returns a cache holding at most max entries. A zero ttl means |
| 30 | // entries never expire. |
| 31 | func NewCache[K comparable, V any](max int, ttl time.Duration) *Cache[K, V] { |
| 32 | return &Cache[K, V]{max: max, ttl: ttl, items: make(map[K]cacheEntry[V], max)} |
| 33 | } |
| 34 | |
| 35 | // NewSizedCache is NewCache with a second bound: the sum of sizeOf over all |
| 36 | // entries stays at or below maxBytes. Values larger than maxBytes are not cached. |
| 37 | func NewSizedCache[K comparable, V any](max int, maxBytes int64, ttl time.Duration, sizeOf func(V) int64) *Cache[K, V] { |
| 38 | c := NewCache[K, V](max, ttl) |
| 39 | c.maxBytes, c.sizeOf = maxBytes, sizeOf |
| 40 | return c |
| 41 | } |
| 42 | |
| 43 | // Get returns the value for key. Expired entries are dropped and report false. |
| 44 | func (c *Cache[K, V]) Get(key K) (V, bool) { |
| 45 | c.mu.Lock() |
| 46 | defer c.mu.Unlock() |
| 47 | e, ok := c.items[key] |
| 48 | if !ok { |
| 49 | return *new(V), false |
| 50 | } |
| 51 | if !e.expiresAt.IsZero() && !time.Now().Before(e.expiresAt) { |
| 52 | c.remove(key) |
| 53 | return *new(V), false |
| 54 | } |
| 55 | return e.value, true |
| 56 | } |
| 57 | |
| 58 | // Set stores a value, evicting the oldest entries when the cache is full. |
| 59 | func (c *Cache[K, V]) Set(key K, v V) { c.SetTTL(key, v, c.ttl) } |
| 60 | |
| 61 | // SetTTL is Set with a per-entry ttl. A zero ttl means the entry never expires. |
| 62 | func (c *Cache[K, V]) SetTTL(key K, v V, ttl time.Duration) { |
| 63 | c.mu.Lock() |
| 64 | defer c.mu.Unlock() |
| 65 | e := cacheEntry[V]{value: v} |
| 66 | if ttl > 0 { |
| 67 | e.expiresAt = time.Now().Add(ttl) |
| 68 | } |
| 69 | c.remove(key) |
| 70 | if c.sizeOf != nil { |
| 71 | e.size = c.sizeOf(v) |
| 72 | if e.size > c.maxBytes { |
| 73 | return |
| 74 | } |
| 75 | } |
| 76 | for len(c.order) > 0 && (len(c.order) >= c.max || c.bytes+e.size > c.maxBytes && c.sizeOf != nil) { |
| 77 | c.bytes -= c.items[c.order[0]].size |
| 78 | delete(c.items, c.order[0]) |
| 79 | c.order = c.order[1:] |
| 80 | } |
| 81 | c.items[key] = e |
| 82 | c.order = append(c.order, key) |
| 83 | c.bytes += e.size |
| 84 | } |
| 85 | |
| 86 | // Delete drops one entry. |
| 87 | func (c *Cache[K, V]) Delete(key K) { |
| 88 | c.mu.Lock() |
| 89 | defer c.mu.Unlock() |
| 90 | c.remove(key) |
| 91 | } |
| 92 | |
| 93 | // Len reports how many entries are held, expired ones included. |
| 94 | func (c *Cache[K, V]) Len() int { |
| 95 | c.mu.Lock() |
| 96 | defer c.mu.Unlock() |
| 97 | return len(c.items) |
| 98 | } |
| 99 | |
| 100 | // remove deletes key from both the map and the order. The caller holds the lock. |
| 101 | func (c *Cache[K, V]) remove(key K) { |
| 102 | e, ok := c.items[key] |
| 103 | if !ok { |
| 104 | return |
| 105 | } |
| 106 | c.bytes -= e.size |
| 107 | delete(c.items, key) |
| 108 | for i, k := range c.order { |
| 109 | if k == key { |
| 110 | c.order = append(c.order[:i], c.order[i+1:]...) |
| 111 | return |
| 112 | } |
| 113 | } |
| 114 | } |
| 115 |