github.com/fufuok/freelru@v0.13.3/lru.go (about)

     1  // Licensed to Elasticsearch B.V. under one or more contributor
     2  // license agreements. See the NOTICE file distributed with
     3  // this work for additional information regarding copyright
     4  // ownership. Elasticsearch B.V. licenses this file to you under
     5  // the Apache License, Version 2.0 (the "License"); you may
     6  // not use this file except in compliance with the License.
     7  // You may obtain a copy of the License at
     8  //
     9  //    http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing,
    12  // software distributed under the License is distributed on an
    13  // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    14  // KIND, either express or implied.  See the License for the
    15  // specific language governing permissions and limitations
    16  // under the License.
    17  
    18  package freelru
    19  
    20  import (
    21  	"errors"
    22  	"fmt"
    23  	"math"
    24  	"math/bits"
    25  	"time"
    26  )
    27  
    28  // OnEvictCallback is the function type for Config.OnEvict.
    29  type OnEvictCallback[K comparable, V any] func(K, V)
    30  
    31  // HashKeyCallback is the function that creates a hash from the passed key.
    32  type HashKeyCallback[K comparable] func(K) uint32
    33  
    34  type element[K comparable, V any] struct {
    35  	key   K
    36  	value V
    37  
    38  	// bucketNext and bucketPrev are indexes in the space-dimension doubly-linked list of elements.
    39  	// That is to add/remove items to the collision bucket without re-allocations and with O(1)
    40  	// complexity.
    41  	// To simplify the implementation, internally a list l is implemented
    42  	// as a ring, such that &l.latest.prev is last element and
    43  	// &l.last.next is the latest element.
    44  	nextBucket, prevBucket uint32
    45  
    46  	// bucketPos is the bucket that an element belongs to.
    47  	bucketPos uint32
    48  
    49  	// next and prev are indexes in the time-dimension doubly-linked list of elements.
    50  	// To simplify the implementation, internally a list l is implemented
    51  	// as a ring, such that &l.latest.prev is last element and
    52  	// &l.last.next is the latest element.
    53  	next, prev uint32
    54  
    55  	// expire is the point in time when the element expires.
    56  	// Its value is Unix milliseconds since epoch.
    57  	expire int64
    58  }
    59  
    60  const emptyBucket = math.MaxUint32
    61  
    62  // LRU implements a non-thread safe fixed size LRU cache.
    63  type LRU[K comparable, V any] struct {
    64  	buckets  []uint32 // contains positions of bucket lists or 'emptyBucket'
    65  	elements []element[K, V]
    66  	onEvict  OnEvictCallback[K, V]
    67  	hash     HashKeyCallback[K]
    68  	lifetime time.Duration
    69  	metrics  Metrics
    70  
    71  	// used for element clearing after removal or expiration
    72  	emptyKey   K
    73  	emptyValue V
    74  
    75  	head uint32 // index of the newest element in the cache
    76  	len  uint32 // current number of elements in the cache
    77  	cap  uint32 // max number of elements in the cache
    78  	size uint32 // size of the element array (X% larger than cap)
    79  	mask uint32 // bitmask to avoid the costly idiv in hashToPos() if size is a 2^n value
    80  }
    81  
    82  // Metrics contains metrics about the cache.
    83  type Metrics struct {
    84  	Inserts    uint64
    85  	Collisions uint64
    86  	Evictions  uint64
    87  	Removals   uint64
    88  	Hits       uint64
    89  	Misses     uint64
    90  	Capacity   uint32
    91  	Lifetime   string
    92  	Len        int
    93  }
    94  
    95  var _ Cache[int, int] = (*LRU[int, int])(nil)
    96  
    97  // SetLifetime sets the default lifetime of LRU elements.
    98  // Lifetime 0 means "forever".
    99  func (lru *LRU[K, V]) SetLifetime(lifetime time.Duration) {
   100  	lru.lifetime = lifetime
   101  	lru.metrics.Lifetime = lifetime.String()
   102  }
   103  
   104  // SetOnEvict sets the OnEvict callback function.
   105  // The onEvict function is called for each evicted lru entry.
   106  func (lru *LRU[K, V]) SetOnEvict(onEvict OnEvictCallback[K, V]) {
   107  	lru.onEvict = onEvict
   108  }
   109  
   110  // New constructs an LRU with the given capacity of elements.
   111  // The hash function calculates a hash value from the keys.
   112  func New[K comparable, V any](capacity uint32, hash HashKeyCallback[K]) (*LRU[K, V], error) {
   113  	return NewWithSize[K, V](capacity, capacity, hash)
   114  }
   115  
   116  // NewWithSize constructs an LRU with the given capacity and size.
   117  // The hash function calculates a hash value from the keys.
   118  // A size greater than the capacity increases memory consumption and decreases the CPU consumption
   119  // by reducing the chance of collisions.
   120  // Size must not be lower than the capacity.
   121  func NewWithSize[K comparable, V any](capacity, size uint32, hash HashKeyCallback[K]) (*LRU[K, V], error) {
   122  	if capacity == 0 {
   123  		return nil, errors.New("capacity must be positive")
   124  	}
   125  	if size == emptyBucket {
   126  		return nil, fmt.Errorf("size must not be %#X", size)
   127  	}
   128  	if size < capacity {
   129  		return nil, fmt.Errorf("size (%d) is smaller than capacity (%d)", size, capacity)
   130  	}
   131  	if hash == nil {
   132  		return nil, errors.New("hash function must be set")
   133  	}
   134  
   135  	buckets := make([]uint32, size)
   136  	elements := make([]element[K, V], size)
   137  
   138  	var lru LRU[K, V]
   139  	initLRU(&lru, capacity, size, hash, buckets, elements)
   140  
   141  	return &lru, nil
   142  }
   143  
   144  func initLRU[K comparable, V any](lru *LRU[K, V], capacity, size uint32, hash HashKeyCallback[K],
   145  	buckets []uint32, elements []element[K, V],
   146  ) {
   147  	lru.cap = capacity
   148  	lru.size = size
   149  	lru.hash = hash
   150  	lru.buckets = buckets
   151  	lru.elements = elements
   152  	lru.lifetime = 0
   153  	lru.metrics.Capacity = capacity
   154  	lru.metrics.Lifetime = lru.lifetime.String()
   155  
   156  	// If the size is 2^N, we can avoid costly divisions.
   157  	if bits.OnesCount32(lru.size) == 1 {
   158  		lru.mask = lru.size - 1
   159  	}
   160  
   161  	// Mark all slots as free.
   162  	for i := range lru.buckets {
   163  		lru.buckets[i] = emptyBucket
   164  	}
   165  }
   166  
   167  // hashToBucketPos converts a hash value into a position in the elements array.
   168  func (lru *LRU[K, V]) hashToBucketPos(hash uint32) uint32 {
   169  	if lru.mask != 0 {
   170  		return hash & lru.mask
   171  	}
   172  	return hash % lru.size
   173  }
   174  
   175  // hashToPos converts a key into a position in the elements array.
   176  func (lru *LRU[K, V]) hashToPos(hash uint32) (bucketPos, elemPos uint32) {
   177  	bucketPos = lru.hashToBucketPos(hash)
   178  	elemPos = lru.buckets[bucketPos]
   179  	return
   180  }
   181  
   182  // setHead links the element as the head into the list.
   183  func (lru *LRU[K, V]) setHead(pos uint32) {
   184  	// Both calls to setHead() check beforehand that pos != lru.head.
   185  	// So if you run into this situation, you likely use FreeLRU in a concurrent situation
   186  	// without proper locking. It requires a write lock, even around Get().
   187  	// But better use SyncedLRU or SharedLRU in such a case.
   188  	if pos == lru.head {
   189  		panic(pos)
   190  	}
   191  
   192  	lru.elements[pos].prev = lru.head
   193  	lru.elements[pos].next = lru.elements[lru.head].next
   194  	lru.elements[lru.elements[lru.head].next].prev = pos
   195  	lru.elements[lru.head].next = pos
   196  	lru.head = pos
   197  }
   198  
   199  // unlinkElement removes the element from the elements list.
   200  func (lru *LRU[K, V]) unlinkElement(pos uint32) {
   201  	lru.elements[lru.elements[pos].prev].next = lru.elements[pos].next
   202  	lru.elements[lru.elements[pos].next].prev = lru.elements[pos].prev
   203  }
   204  
   205  // unlinkBucket removes the element from the buckets list.
   206  func (lru *LRU[K, V]) unlinkBucket(pos uint32) {
   207  	prevBucket := lru.elements[pos].prevBucket
   208  	nextBucket := lru.elements[pos].nextBucket
   209  	if prevBucket == nextBucket && prevBucket == pos { //nolint:gocritic
   210  		// The element references itself, so it's the only bucket entry
   211  		lru.buckets[lru.elements[pos].bucketPos] = emptyBucket
   212  		return
   213  	}
   214  	lru.elements[prevBucket].nextBucket = nextBucket
   215  	lru.elements[nextBucket].prevBucket = prevBucket
   216  	lru.buckets[lru.elements[pos].bucketPos] = nextBucket
   217  }
   218  
   219  // evict evicts the element at the given position.
   220  func (lru *LRU[K, V]) evict(pos uint32) {
   221  	if pos == lru.head {
   222  		lru.head = lru.elements[pos].prev
   223  	}
   224  
   225  	lru.unlinkElement(pos)
   226  	lru.unlinkBucket(pos)
   227  	lru.len--
   228  
   229  	if lru.onEvict != nil {
   230  		// Save k/v for the eviction function.
   231  		key := lru.elements[pos].key
   232  		value := lru.elements[pos].value
   233  		lru.onEvict(key, value)
   234  	}
   235  }
   236  
   237  // Move element from position old to new.
   238  // That avoids 'gaps' and new elements can always be simply appended.
   239  func (lru *LRU[K, V]) move(to, from uint32) {
   240  	if to == from {
   241  		return
   242  	}
   243  	if from == lru.head {
   244  		lru.head = to
   245  	}
   246  
   247  	prev := lru.elements[from].prev
   248  	next := lru.elements[from].next
   249  	lru.elements[prev].next = to
   250  	lru.elements[next].prev = to
   251  
   252  	prev = lru.elements[from].prevBucket
   253  	next = lru.elements[from].nextBucket
   254  	lru.elements[prev].nextBucket = to
   255  	lru.elements[next].prevBucket = to
   256  
   257  	lru.elements[to] = lru.elements[from]
   258  
   259  	if lru.buckets[lru.elements[to].bucketPos] == from {
   260  		lru.buckets[lru.elements[to].bucketPos] = to
   261  	}
   262  }
   263  
   264  // insert stores the k/v at pos.
   265  // It updates the head to point to this position.
   266  func (lru *LRU[K, V]) insert(pos uint32, key K, value V, lifetime time.Duration) {
   267  	lru.elements[pos].key = key
   268  	lru.elements[pos].value = value
   269  	lru.elements[pos].expire = expire(lifetime)
   270  
   271  	if lru.len == 0 {
   272  		lru.elements[pos].prev = pos
   273  		lru.elements[pos].next = pos
   274  		lru.head = pos
   275  	} else if pos != lru.head {
   276  		lru.setHead(pos)
   277  	}
   278  	lru.len++
   279  	lru.metrics.Inserts++
   280  }
   281  
   282  func now() int64 {
   283  	return time.Now().UnixMilli()
   284  }
   285  
   286  func expire(lifetime time.Duration) int64 {
   287  	if lifetime == 0 {
   288  		return 0
   289  	}
   290  	return now() + lifetime.Milliseconds()
   291  }
   292  
   293  // clearKeyAndValue clears stale data to avoid memory leaks
   294  func (lru *LRU[K, V]) clearKeyAndValue(pos uint32) {
   295  	lru.elements[pos].key = lru.emptyKey
   296  	lru.elements[pos].value = lru.emptyValue
   297  }
   298  
   299  func (lru *LRU[K, V]) findKey(hash uint32, key K) (uint32, bool) {
   300  	_, startPos := lru.hashToPos(hash)
   301  	if startPos == emptyBucket {
   302  		return emptyBucket, false
   303  	}
   304  
   305  	pos := startPos
   306  	for {
   307  		if key == lru.elements[pos].key {
   308  			if lru.elements[pos].expire != 0 && lru.elements[pos].expire <= now() {
   309  				lru.clearKeyAndValue(pos)
   310  				return emptyBucket, false
   311  			}
   312  			return pos, true
   313  		}
   314  
   315  		pos = lru.elements[pos].nextBucket
   316  		if pos == startPos {
   317  			// Key not found
   318  			return emptyBucket, false
   319  		}
   320  	}
   321  }
   322  
   323  // Len returns the number of elements stored in the cache.
   324  func (lru *LRU[K, V]) Len() int {
   325  	return int(lru.len)
   326  }
   327  
   328  // AddWithLifetime adds a key:value to the cache with a lifetime.
   329  // Returns true, true if key was updated and eviction occurred.
   330  func (lru *LRU[K, V]) AddWithLifetime(key K, value V, lifetime time.Duration) (evicted bool) {
   331  	return lru.addWithLifetime(lru.hash(key), key, value, lifetime)
   332  }
   333  
   334  func (lru *LRU[K, V]) addWithLifetime(hash uint32, key K, value V, lifetime time.Duration) (evicted bool) {
   335  	bucketPos, startPos := lru.hashToPos(hash)
   336  	if startPos == emptyBucket {
   337  		pos := lru.len
   338  
   339  		if pos == lru.cap {
   340  			// Capacity reached, evict the oldest entry and
   341  			// store the new entry at evicted position.
   342  			pos = lru.elements[lru.head].next
   343  			lru.evict(pos)
   344  			lru.metrics.Evictions++
   345  			evicted = true
   346  		}
   347  
   348  		// insert new (first) entry into the bucket
   349  		lru.buckets[bucketPos] = pos
   350  		lru.elements[pos].bucketPos = bucketPos
   351  
   352  		lru.elements[pos].nextBucket = pos
   353  		lru.elements[pos].prevBucket = pos
   354  		lru.insert(pos, key, value, lifetime)
   355  		return evicted
   356  	}
   357  
   358  	// Walk through the bucket list to see whether key already exists.
   359  	pos := startPos
   360  	for {
   361  		if lru.elements[pos].key == key {
   362  			// Key exists, replace the value and update element to be the head element.
   363  			lru.elements[pos].value = value
   364  			lru.elements[pos].expire = expire(lifetime)
   365  
   366  			if pos != lru.head {
   367  				lru.unlinkElement(pos)
   368  				lru.setHead(pos)
   369  			}
   370  			// count as insert, even if it's just an update
   371  			lru.metrics.Inserts++
   372  			return false
   373  		}
   374  
   375  		pos = lru.elements[pos].nextBucket
   376  		if pos == startPos {
   377  			// Key not found
   378  			break
   379  		}
   380  	}
   381  
   382  	pos = lru.len
   383  	if pos == lru.cap {
   384  		// Capacity reached, evict the oldest entry and
   385  		// store the new entry at evicted position.
   386  		pos = lru.elements[lru.head].next
   387  		lru.evict(pos)
   388  		lru.metrics.Evictions++
   389  		evicted = true
   390  		startPos = lru.buckets[bucketPos]
   391  		if startPos == emptyBucket {
   392  			startPos = pos
   393  		}
   394  	}
   395  
   396  	// insert new entry into the existing bucket before startPos
   397  	lru.buckets[bucketPos] = pos
   398  	lru.elements[pos].bucketPos = bucketPos
   399  
   400  	lru.elements[pos].nextBucket = startPos
   401  	lru.elements[pos].prevBucket = lru.elements[startPos].prevBucket
   402  	lru.elements[lru.elements[startPos].prevBucket].nextBucket = pos
   403  	lru.elements[startPos].prevBucket = pos
   404  	lru.insert(pos, key, value, lifetime)
   405  
   406  	if lru.elements[pos].prevBucket != pos {
   407  		// The bucket now contains more than 1 element.
   408  		// That means we have a collision.
   409  		lru.metrics.Collisions++
   410  	}
   411  	return evicted
   412  }
   413  
   414  // Add adds a key:value to the cache.
   415  // Returns true, true if key was updated and eviction occurred.
   416  func (lru *LRU[K, V]) Add(key K, value V) (evicted bool) {
   417  	return lru.addWithLifetime(lru.hash(key), key, value, lru.lifetime)
   418  }
   419  
   420  func (lru *LRU[K, V]) add(hash uint32, key K, value V) (evicted bool) {
   421  	return lru.addWithLifetime(hash, key, value, lru.lifetime)
   422  }
   423  
   424  // Get looks up a key's value from the cache, setting it as the most
   425  // recently used item.
   426  func (lru *LRU[K, V]) Get(key K) (value V, ok bool) {
   427  	return lru.get(lru.hash(key), key)
   428  }
   429  
   430  func (lru *LRU[K, V]) get(hash uint32, key K) (value V, ok bool) {
   431  	if pos, ok := lru.findKey(hash, key); ok {
   432  		if pos != lru.head {
   433  			lru.unlinkElement(pos)
   434  			lru.setHead(pos)
   435  		}
   436  		lru.metrics.Hits++
   437  		return lru.elements[pos].value, ok
   438  	}
   439  
   440  	lru.metrics.Misses++
   441  	return
   442  }
   443  
   444  // Peek looks up a key's value from the cache, without changing its recent-ness.
   445  func (lru *LRU[K, V]) Peek(key K) (value V, ok bool) {
   446  	return lru.peek(lru.hash(key), key)
   447  }
   448  
   449  func (lru *LRU[K, V]) peek(hash uint32, key K) (value V, ok bool) {
   450  	if pos, ok := lru.findKey(hash, key); ok {
   451  		return lru.elements[pos].value, ok
   452  	}
   453  
   454  	return
   455  }
   456  
   457  // Contains checks for the existence of a key, without changing its recent-ness.
   458  func (lru *LRU[K, V]) Contains(key K) (ok bool) {
   459  	_, ok = lru.peek(lru.hash(key), key)
   460  	return
   461  }
   462  
   463  func (lru *LRU[K, V]) contains(hash uint32, key K) (ok bool) {
   464  	_, ok = lru.peek(hash, key)
   465  	return
   466  }
   467  
   468  // Remove removes the key from the cache.
   469  // The return value indicates whether the key existed or not.
   470  func (lru *LRU[K, V]) Remove(key K) (removed bool) {
   471  	return lru.remove(lru.hash(key), key)
   472  }
   473  
   474  func (lru *LRU[K, V]) remove(hash uint32, key K) (removed bool) {
   475  	if pos, ok := lru.findKey(hash, key); ok {
   476  		// Key exists, update element to be the head element.
   477  		lru.evict(pos)
   478  		lru.move(pos, lru.len)
   479  		lru.metrics.Removals++
   480  
   481  		// remove stale data to avoid memory leaks
   482  		lru.clearKeyAndValue(lru.len)
   483  		return ok
   484  	}
   485  
   486  	return
   487  }
   488  
   489  // Keys returns a slice of the keys in the cache, from oldest to newest.
   490  func (lru *LRU[K, V]) Keys() []K {
   491  	keys := make([]K, 0, lru.len)
   492  	pos := lru.elements[lru.head].next
   493  	for i := uint32(0); i < lru.len; i++ {
   494  		keys = append(keys, lru.elements[pos].key)
   495  		pos = lru.elements[pos].next
   496  	}
   497  	return keys
   498  }
   499  
   500  // Purge purges all data (key and value) from the LRU.
   501  func (lru *LRU[K, V]) Purge() {
   502  	for i := range lru.buckets {
   503  		lru.buckets[i] = emptyBucket
   504  	}
   505  
   506  	for i := range lru.elements {
   507  		lru.elements[i].key = lru.emptyKey
   508  		lru.elements[i].value = lru.emptyValue
   509  	}
   510  
   511  	lru.len = 0
   512  	lru.metrics = Metrics{
   513  		Capacity: lru.cap,
   514  		Lifetime: lru.lifetime.String(),
   515  	}
   516  }
   517  
   518  // Metrics returns the metrics of the cache.
   519  func (lru *LRU[K, V]) Metrics() Metrics {
   520  	lru.metrics.Len = lru.Len()
   521  	return lru.metrics
   522  }
   523  
   524  // ResetMetrics resets the metrics of the cache and returns the previous state.
   525  func (lru *LRU[K, V]) ResetMetrics() Metrics {
   526  	metrics := lru.metrics
   527  	lru.metrics = Metrics{
   528  		Capacity: lru.cap,
   529  		Lifetime: lru.lifetime.String(),
   530  	}
   531  	return metrics
   532  }
   533  
   534  // just used for debugging
   535  func (lru *LRU[K, V]) dump() {
   536  	fmt.Printf("head %d len %d cap %d size %d mask 0x%X\n",
   537  		lru.head, lru.len, lru.cap, lru.size, lru.mask)
   538  
   539  	for i := range lru.buckets {
   540  		if lru.buckets[i] == emptyBucket {
   541  			continue
   542  		}
   543  		fmt.Printf("  bucket[%d] -> %d\n", i, lru.buckets[i])
   544  		pos := lru.buckets[i]
   545  		for {
   546  			e := &lru.elements[pos]
   547  			fmt.Printf("    pos %d bucketPos %d prevBucket %d nextBucket %d prev %d next %d k %v v %v\n",
   548  				pos, e.bucketPos, e.prevBucket, e.nextBucket, e.prev, e.next, e.key, e.value)
   549  			pos = e.nextBucket
   550  			if pos == lru.buckets[i] {
   551  				break
   552  			}
   553  		}
   554  	}
   555  }
   556  
   557  func (lru *LRU[K, V]) PrintStats() {
   558  	m := &lru.metrics
   559  	fmt.Printf("Inserts: %d Collisions: %d (%.2f%%) Evictions: %d Removals: %d Hits: %d (%.2f%%) Misses: %d\n",
   560  		m.Inserts, m.Collisions, float64(m.Collisions)/float64(m.Inserts)*100,
   561  		m.Evictions, m.Removals,
   562  		m.Hits, float64(m.Hits)/float64(m.Hits+m.Misses)*100, m.Misses)
   563  }