cache.go (7288B)
1 package cache 2 3 import ( 4 "context" 5 "errors" 6 "fmt" 7 "sync" 8 "time" 9 10 "dkforest/pkg/clockwork" 11 ) 12 13 const ( 14 NoExpiration time.Duration = -1 15 DefaultExpiration time.Duration = 0 16 ) 17 18 var ErrItemAlreadyExists error = errors.New("item already exists") 19 var ErrItemNotFound error = errors.New("item does not exists") 20 21 // Item wrap the user provided value and add data to it 22 type Item[V any] struct { 23 value V 24 expiration int64 25 } 26 27 type Cache[K comparable, V any] struct { 28 ctx context.Context // Context is used to stop the auto-cleanup thread 29 cancel context.CancelFunc // Cancel the context and stop the auto-cleanup thread 30 mtx sync.RWMutex // This mutex should only be used in exported methods 31 defaultExpiration time.Duration // Default expiration for items in cache 32 items map[K]Item[V] // Hashmap that contains all items in the cache 33 clock clockwork.Clock // Clock object for time related features 34 } 35 36 // Creates a cache with K as string 37 func New[V any](defaultExpiration, cleanupInterval time.Duration) *Cache[string, V] { 38 return newCache[string, V](context.Background(), defaultExpiration, cleanupInterval) 39 } 40 41 // Creates a cache with a context provided by the user 42 func NewWithContext[V any](ctx context.Context, defaultExpiration, cleanupInterval time.Duration) *Cache[string, V] { 43 return newCache[string, V](ctx, defaultExpiration, cleanupInterval) 44 } 45 46 // Create a cache with a custom comparable K provided by the user 47 func NewWithKey[K comparable, V any](defaultExpiration, cleanupInterval time.Duration) *Cache[K, V] { 48 return newCache[K, V](context.Background(), defaultExpiration, cleanupInterval) 49 } 50 51 // Destroy the cache object, cleanup all resources 52 func (c *Cache[K, V]) Destroy() { 53 c.cancel() 54 c = nil 55 } 56 57 // Has returns either or not the key is present in the cache 58 func (c *Cache[K, V]) Has(k K) bool { 59 c.mtx.RLock() 60 found := c.has(k) 61 c.mtx.RUnlock() 62 return found 63 } 64 65 // Get an value associated to the given key 66 func (c *Cache[K, V]) Get(k K) (V, bool) { 67 c.mtx.RLock() 68 value, found := c.get(k) 69 c.mtx.RUnlock() 70 return value, found 71 } 72 73 // GetWithExpiration a value and it's expiration 74 func (c *Cache[K, V]) GetWithExpiration(k K) (V, time.Time, bool) { 75 c.mtx.RLock() 76 value, expiration, found := c.getWithExpiration(k) 77 c.mtx.RUnlock() 78 return value, expiration, found 79 } 80 81 // Set a key/value pair in the cache 82 func (c *Cache[K, V]) Set(k K, v V, d time.Duration) { 83 c.mtx.Lock() 84 c.set(k, v, d) 85 c.mtx.Unlock() 86 } 87 88 // SetD same as Set, but use the DefaultExpiration automatically 89 func (c *Cache[K, V]) SetD(k K, v V) { 90 c.mtx.Lock() 91 c.set(k, v, c.defaultExpiration) 92 c.mtx.Unlock() 93 } 94 95 // Add an item to the cache only if an item doesn't already exist for the given 96 // key, or if the existing item has expired. Returns an error otherwise. 97 func (c *Cache[K, V]) Add(k K, v V, d time.Duration) error { 98 c.mtx.Lock() 99 err := c.add(k, v, d) 100 c.mtx.Unlock() 101 return err 102 } 103 104 // Replace set a new value for the cache key only if it already exists, and the existing 105 // item hasn't expired. Returns an error otherwise. 106 func (c *Cache[K, V]) Replace(k K, v V, d time.Duration) error { 107 c.mtx.Lock() 108 err := c.replace(k, v, d) 109 c.mtx.Unlock() 110 return err 111 } 112 113 func (c *Cache[K, V]) Update(k K, v V) error { 114 c.mtx.Lock() 115 err := c.update(k, v) 116 c.mtx.Unlock() 117 return err 118 } 119 120 // Delete an item from the cache 121 func (c *Cache[K, V]) Delete(k K) { 122 c.mtx.Lock() 123 c.delete(k) 124 c.mtx.Unlock() 125 } 126 127 // DeleteExpired deletes all expired items from the cache 128 func (c *Cache[K, V]) DeleteExpired() { 129 c.mtx.Lock() 130 c.deleteExpired() 131 c.mtx.Unlock() 132 } 133 134 // DeleteAll deletes all items from the cache 135 func (c *Cache[K, V]) DeleteAll() { 136 c.mtx.Lock() 137 c.deleteAll() 138 c.mtx.Unlock() 139 } 140 141 // Len returns the number of items in the cache. This may include items that have 142 // expired, but have not yet been cleaned up. 143 func (c *Cache[K, V]) Len() int { 144 c.mtx.Lock() 145 n := len(c.items) 146 c.mtx.Unlock() 147 return n 148 } 149 150 // Items copies all unexpired items in the cache into a new map and returns it. 151 func (c *Cache[K, V]) Items() map[K]Item[V] { 152 c.mtx.RLock() 153 items := c.getItems() 154 c.mtx.RUnlock() 155 return items 156 } 157 158 // SetClock set the clock object 159 func (c *Cache[K, V]) SetClock(clock clockwork.Clock) { 160 c.clock = clock 161 } 162 163 func newCache[K comparable, V any](ctx context.Context, defaultExpiration, cleanupInterval time.Duration) *Cache[K, V] { 164 items := make(map[K]Item[V]) 165 c := new(Cache[K, V]) 166 c.ctx, c.cancel = context.WithCancel(ctx) 167 c.clock = clockwork.NewRealClock() 168 c.defaultExpiration = defaultExpiration 169 c.items = items 170 if cleanupInterval > 0 { 171 go c.autoCleanup(cleanupInterval) 172 } 173 return c 174 } 175 176 func (c *Cache[K, V]) autoCleanup(cleanupInterval time.Duration) { 177 for { 178 select { 179 case <-c.ctx.Done(): 180 return 181 case <-time.After(cleanupInterval): 182 } 183 // Important to call the exported method to lock the mutex 184 c.DeleteExpired() 185 } 186 } 187 188 func (c *Cache[K, V]) has(k K) bool { 189 _, found := c.get(k) 190 return found 191 } 192 193 func (c *Cache[K, V]) getWithExpiration(k K) (V, time.Time, bool) { 194 var zero V 195 now := c.clock.Now().UnixNano() 196 item, found := c.items[k] 197 if !found { 198 return zero, time.Time{}, false 199 } 200 e := time.Time{} 201 if item.expiration > 0 { 202 if item.expiration < now { 203 return zero, time.Time{}, false 204 } 205 e = time.Unix(0, item.expiration) 206 } 207 return item.value, e, found 208 } 209 210 func (c *Cache[K, V]) get(k K) (V, bool) { 211 value, _, found := c.getWithExpiration(k) 212 return value, found 213 } 214 215 func (c *Cache[K, V]) set(k K, v V, d time.Duration) { 216 e := int64(NoExpiration) 217 if d == DefaultExpiration { 218 d = c.defaultExpiration 219 } 220 e = c.clock.Now().Add(d).UnixNano() 221 c.items[k] = Item[V]{value: v, expiration: e} 222 } 223 224 func (c *Cache[K, V]) add(k K, v V, d time.Duration) error { 225 if _, found := c.get(k); found { 226 return ErrItemAlreadyExists 227 } 228 c.set(k, v, d) 229 return nil 230 } 231 232 func (c *Cache[K, V]) replace(k K, v V, d time.Duration) error { 233 if _, found := c.get(k); !found { 234 return ErrItemNotFound 235 } 236 c.set(k, v, d) 237 return nil 238 } 239 240 func (c *Cache[K, V]) update(k K, v V) error { 241 vv, found := c.items[k] 242 if !found || vv.IsExpired() { 243 return fmt.Errorf("item %v not found", k) 244 } 245 vv.value = v 246 c.items[k] = vv 247 return nil 248 } 249 250 func (c *Cache[K, V]) deleteAll() { 251 c.items = make(map[K]Item[V]) 252 } 253 254 func (c *Cache[K, V]) delete(k K) { 255 delete(c.items, k) 256 } 257 258 func (c *Cache[K, V]) deleteExpired() { 259 now := c.clock.Now().UnixNano() 260 for k := range c.items { 261 item := c.items[k] 262 if item.isExpired(now) { 263 c.delete(k) 264 } 265 } 266 } 267 268 func (c *Cache[K, V]) getItems() map[K]Item[V] { 269 now := c.clock.Now().UnixNano() 270 m := make(map[K]Item[V], len(c.items)) 271 for k, v := range c.items { 272 if !v.isExpired(now) { 273 m[k] = v 274 } 275 } 276 return m 277 } 278 279 // Value returns the value contained by the item 280 func (i Item[V]) Value() V { 281 return i.value 282 } 283 284 // Expiration returns the expiration time 285 func (i Item[V]) Expiration() time.Time { 286 return time.Unix(0, i.expiration) 287 } 288 289 // IsExpired returns either or not the item is expired right now 290 func (i Item[V]) IsExpired() bool { 291 now := time.Now().UnixNano() 292 return i.isExpired(now) 293 } 294 295 // Given a unix (nano) timestamp, return either or not the item is expired 296 func (i Item[V]) isExpired(ts int64) bool { 297 return i.expiration > 0 && i.expiration < ts 298 }