expirable_lru.go

Documentation: github.com/hashicorp/golang-lru/v2/expirable

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: MPL-2.0
     3  
     4  package expirable
     5  
     6  import (
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/hashicorp/golang-lru/v2/internal"
    11  )
    12  
    13  // EvictCallback is used to get a callback when a cache entry is evicted
    14  type EvictCallback[K comparable, V any] func(key K, value V)
    15  
    16  // LRU implements a thread-safe LRU with expirable entries.
    17  type LRU[K comparable, V any] struct {
    18  	size      int
    19  	evictList *internal.LruList[K, V]
    20  	items     map[K]*internal.Entry[K, V]
    21  	onEvict   EvictCallback[K, V]
    22  
    23  	// expirable options
    24  	mu   sync.Mutex
    25  	ttl  time.Duration
    26  	done chan struct{}
    27  
    28  	// buckets for expiration
    29  	buckets []bucket[K, V]
    30  	// uint8 because it's number between 0 and numBuckets
    31  	nextCleanupBucket uint8
    32  }
    33  
    34  // bucket is a container for holding entries to be expired
    35  type bucket[K comparable, V any] struct {
    36  	entries     map[K]*internal.Entry[K, V]
    37  	newestEntry time.Time
    38  }
    39  
    40  // noEvictionTTL - very long ttl to prevent eviction
    41  const noEvictionTTL = time.Hour * 24 * 365 * 10
    42  
    43  // because of uint8 usage for nextCleanupBucket, should not exceed 256.
    44  // casting it as uint8 explicitly requires type conversions in multiple places
    45  const numBuckets = 100
    46  
    47  // NewLRU returns a new thread-safe cache with expirable entries.
    48  //
    49  // Size parameter set to 0 makes cache of unlimited size, e.g. turns LRU mechanism off.
    50  //
    51  // Providing 0 TTL turns expiring off.
    52  //
    53  // Delete expired entries every 1/100th of ttl value. Goroutine which deletes expired entries runs indefinitely.
    54  func NewLRU[K comparable, V any](size int, onEvict EvictCallback[K, V], ttl time.Duration) *LRU[K, V] {
    55  	if size < 0 {
    56  		size = 0
    57  	}
    58  	if ttl <= 0 {
    59  		ttl = noEvictionTTL
    60  	}
    61  
    62  	res := LRU[K, V]{
    63  		ttl:       ttl,
    64  		size:      size,
    65  		evictList: internal.NewList[K, V](),
    66  		items:     make(map[K]*internal.Entry[K, V]),
    67  		onEvict:   onEvict,
    68  		done:      make(chan struct{}),
    69  	}
    70  
    71  	// initialize the buckets
    72  	res.buckets = make([]bucket[K, V], numBuckets)
    73  	for i := 0; i < numBuckets; i++ {
    74  		res.buckets[i] = bucket[K, V]{entries: make(map[K]*internal.Entry[K, V])}
    75  	}
    76  
    77  	// enable deleteExpired() running in separate goroutine for cache with non-zero TTL
    78  	//
    79  	// Important: done channel is never closed, so deleteExpired() goroutine will never exit,
    80  	// it's decided to add functionality to close it in the version later than v2.
    81  	if res.ttl != noEvictionTTL {
    82  		go func(done <-chan struct{}) {
    83  			ticker := time.NewTicker(res.ttl / numBuckets)
    84  			defer ticker.Stop()
    85  			for {
    86  				select {
    87  				case <-done:
    88  					return
    89  				case <-ticker.C:
    90  					res.deleteExpired()
    91  				}
    92  			}
    93  		}(res.done)
    94  	}
    95  	return &res
    96  }
    97  
    98  // Purge clears the cache completely.
    99  // onEvict is called for each evicted key.
   100  func (c *LRU[K, V]) Purge() {
   101  	c.mu.Lock()
   102  	defer c.mu.Unlock()
   103  	for k, v := range c.items {
   104  		if c.onEvict != nil {
   105  			c.onEvict(k, v.Value)
   106  		}
   107  		delete(c.items, k)
   108  	}
   109  	for _, b := range c.buckets {
   110  		for _, ent := range b.entries {
   111  			delete(b.entries, ent.Key)
   112  		}
   113  	}
   114  	c.evictList.Init()
   115  }
   116  
   117  // Add adds a value to the cache. Returns true if an eviction occurred.
   118  // Returns false if there was no eviction: the item was already in the cache,
   119  // or the size was not exceeded.
   120  func (c *LRU[K, V]) Add(key K, value V) (evicted bool) {
   121  	c.mu.Lock()
   122  	defer c.mu.Unlock()
   123  	now := time.Now()
   124  
   125  	// Check for existing item
   126  	if ent, ok := c.items[key]; ok {
   127  		c.evictList.MoveToFront(ent)
   128  		c.removeFromBucket(ent) // remove the entry from its current bucket as expiresAt is renewed
   129  		ent.Value = value
   130  		ent.ExpiresAt = now.Add(c.ttl)
   131  		c.addToBucket(ent)
   132  		return false
   133  	}
   134  
   135  	// Add new item
   136  	ent := c.evictList.PushFrontExpirable(key, value, now.Add(c.ttl))
   137  	c.items[key] = ent
   138  	c.addToBucket(ent) // adds the entry to the appropriate bucket and sets entry.expireBucket
   139  
   140  	evict := c.size > 0 && c.evictList.Length() > c.size
   141  	// Verify size not exceeded
   142  	if evict {
   143  		c.removeOldest()
   144  	}
   145  	return evict
   146  }
   147  
   148  // Get looks up a key's value from the cache.
   149  func (c *LRU[K, V]) Get(key K) (value V, ok bool) {
   150  	c.mu.Lock()
   151  	defer c.mu.Unlock()
   152  	var ent *internal.Entry[K, V]
   153  	if ent, ok = c.items[key]; ok {
   154  		// Expired item check
   155  		if time.Now().After(ent.ExpiresAt) {
   156  			return value, false
   157  		}
   158  		c.evictList.MoveToFront(ent)
   159  		return ent.Value, true
   160  	}
   161  	return
   162  }
   163  
   164  // Contains checks if a key is in the cache, without updating the recent-ness
   165  // or deleting it for being stale.
   166  func (c *LRU[K, V]) Contains(key K) (ok bool) {
   167  	c.mu.Lock()
   168  	defer c.mu.Unlock()
   169  	_, ok = c.items[key]
   170  	return ok
   171  }
   172  
   173  // Peek returns the key value (or undefined if not found) without updating
   174  // the "recently used"-ness of the key.
   175  func (c *LRU[K, V]) Peek(key K) (value V, ok bool) {
   176  	c.mu.Lock()
   177  	defer c.mu.Unlock()
   178  	var ent *internal.Entry[K, V]
   179  	if ent, ok = c.items[key]; ok {
   180  		// Expired item check
   181  		if time.Now().After(ent.ExpiresAt) {
   182  			return value, false
   183  		}
   184  		return ent.Value, true
   185  	}
   186  	return
   187  }
   188  
   189  // Remove removes the provided key from the cache, returning if the
   190  // key was contained.
   191  func (c *LRU[K, V]) Remove(key K) bool {
   192  	c.mu.Lock()
   193  	defer c.mu.Unlock()
   194  	if ent, ok := c.items[key]; ok {
   195  		c.removeElement(ent)
   196  		return true
   197  	}
   198  	return false
   199  }
   200  
   201  // RemoveOldest removes the oldest item from the cache.
   202  func (c *LRU[K, V]) RemoveOldest() (key K, value V, ok bool) {
   203  	c.mu.Lock()
   204  	defer c.mu.Unlock()
   205  	if ent := c.evictList.Back(); ent != nil {
   206  		c.removeElement(ent)
   207  		return ent.Key, ent.Value, true
   208  	}
   209  	return
   210  }
   211  
   212  // GetOldest returns the oldest entry
   213  func (c *LRU[K, V]) GetOldest() (key K, value V, ok bool) {
   214  	c.mu.Lock()
   215  	defer c.mu.Unlock()
   216  	if ent := c.evictList.Back(); ent != nil {
   217  		return ent.Key, ent.Value, true
   218  	}
   219  	return
   220  }
   221  
   222  // Keys returns a slice of the keys in the cache, from oldest to newest.
   223  func (c *LRU[K, V]) Keys() []K {
   224  	c.mu.Lock()
   225  	defer c.mu.Unlock()
   226  	keys := make([]K, 0, len(c.items))
   227  	for ent := c.evictList.Back(); ent != nil; ent = ent.PrevEntry() {
   228  		keys = append(keys, ent.Key)
   229  	}
   230  	return keys
   231  }
   232  
   233  // Values returns a slice of the values in the cache, from oldest to newest.
   234  // Expired entries are filtered out.
   235  func (c *LRU[K, V]) Values() []V {
   236  	c.mu.Lock()
   237  	defer c.mu.Unlock()
   238  	values := make([]V, len(c.items))
   239  	i := 0
   240  	now := time.Now()
   241  	for ent := c.evictList.Back(); ent != nil; ent = ent.PrevEntry() {
   242  		if now.After(ent.ExpiresAt) {
   243  			continue
   244  		}
   245  		values[i] = ent.Value
   246  		i++
   247  	}
   248  	return values
   249  }
   250  
   251  // Len returns the number of items in the cache.
   252  func (c *LRU[K, V]) Len() int {
   253  	c.mu.Lock()
   254  	defer c.mu.Unlock()
   255  	return c.evictList.Length()
   256  }
   257  
   258  // Resize changes the cache size. Size of 0 means unlimited.
   259  func (c *LRU[K, V]) Resize(size int) (evicted int) {
   260  	c.mu.Lock()
   261  	defer c.mu.Unlock()
   262  	if size <= 0 {
   263  		c.size = 0
   264  		return 0
   265  	}
   266  	diff := c.evictList.Length() - size
   267  	if diff < 0 {
   268  		diff = 0
   269  	}
   270  	for i := 0; i < diff; i++ {
   271  		c.removeOldest()
   272  	}
   273  	c.size = size
   274  	return diff
   275  }
   276  
   277  // Close destroys cleanup goroutine. To clean up the cache, run Purge() before Close().
   278  // func (c *LRU[K, V]) Close() {
   279  //	c.mu.Lock()
   280  //	defer c.mu.Unlock()
   281  //	select {
   282  //	case <-c.done:
   283  //		return
   284  //	default:
   285  //	}
   286  //	close(c.done)
   287  // }
   288  
   289  // removeOldest removes the oldest item from the cache. Has to be called with lock!
   290  func (c *LRU[K, V]) removeOldest() {
   291  	if ent := c.evictList.Back(); ent != nil {
   292  		c.removeElement(ent)
   293  	}
   294  }
   295  
   296  // removeElement is used to remove a given list element from the cache. Has to be called with lock!
   297  func (c *LRU[K, V]) removeElement(e *internal.Entry[K, V]) {
   298  	c.evictList.Remove(e)
   299  	delete(c.items, e.Key)
   300  	c.removeFromBucket(e)
   301  	if c.onEvict != nil {
   302  		c.onEvict(e.Key, e.Value)
   303  	}
   304  }
   305  
   306  // deleteExpired deletes expired records from the oldest bucket, waiting for the newest entry
   307  // in it to expire first.
   308  func (c *LRU[K, V]) deleteExpired() {
   309  	c.mu.Lock()
   310  	bucketIdx := c.nextCleanupBucket
   311  	timeToExpire := time.Until(c.buckets[bucketIdx].newestEntry)
   312  	// wait for newest entry to expire before cleanup without holding lock
   313  	if timeToExpire > 0 {
   314  		c.mu.Unlock()
   315  		time.Sleep(timeToExpire)
   316  		c.mu.Lock()
   317  	}
   318  	for _, ent := range c.buckets[bucketIdx].entries {
   319  		c.removeElement(ent)
   320  	}
   321  	c.nextCleanupBucket = (c.nextCleanupBucket + 1) % numBuckets
   322  	c.mu.Unlock()
   323  }
   324  
   325  // addToBucket adds entry to expire bucket so that it will be cleaned up when the time comes. Has to be called with lock!
   326  func (c *LRU[K, V]) addToBucket(e *internal.Entry[K, V]) {
   327  	bucketID := (numBuckets + c.nextCleanupBucket - 1) % numBuckets
   328  	e.ExpireBucket = bucketID
   329  	c.buckets[bucketID].entries[e.Key] = e
   330  	if c.buckets[bucketID].newestEntry.Before(e.ExpiresAt) {
   331  		c.buckets[bucketID].newestEntry = e.ExpiresAt
   332  	}
   333  }
   334  
   335  // removeFromBucket removes the entry from its corresponding bucket. Has to be called with lock!
   336  func (c *LRU[K, V]) removeFromBucket(e *internal.Entry[K, V]) {
   337  	delete(c.buckets[e.ExpireBucket].entries, e.Key)
   338  }
   339
View as plain text