github.com/coyove/sdss@v0.0.0-20231129015646-c2ec58cca6a2/contrib/plru/lru.go (about)

     1  package plru
     2  
     3  import (
     4  	"math"
     5  	"sync"
     6  
     7  	"github.com/coyove/sdss/future"
     8  )
     9  
    10  const Lookahead = 2
    11  
    12  type lruValue[V any] struct {
    13  	Time  int64
    14  	Value V
    15  }
    16  
    17  type Cache[K comparable, V any] struct {
    18  	mu       sync.RWMutex
    19  	onEvict  func(K, V)
    20  	storeCap int
    21  	ptr      int
    22  	store    Map[K, lruValue[V]]
    23  }
    24  
    25  func New[K comparable, V any](cap int, hash func(K) uint64, onEvict func(K, V)) *Cache[K, V] {
    26  	if cap < Lookahead {
    27  		cap = Lookahead
    28  	}
    29  	if onEvict == nil {
    30  		onEvict = func(K, V) {}
    31  	}
    32  	c := &Cache[K, V]{
    33  		onEvict:  onEvict,
    34  		storeCap: cap,
    35  		store:    *NewMap[K, lruValue[V]](cap/3*2, hash),
    36  	}
    37  	c.store.Fixed = true
    38  	return c
    39  }
    40  
    41  func (m *Cache[K, V]) Len() int {
    42  	m.mu.RLock()
    43  	defer m.mu.RUnlock()
    44  	return m.store.Len()
    45  }
    46  
    47  func (m *Cache[K, V]) Cap() int {
    48  	return m.storeCap
    49  }
    50  
    51  func (m *Cache[K, V]) Update(key K, f func(V) V) {
    52  	m.mu.Lock()
    53  	old, ok := m.store.Find(key)
    54  	if ok {
    55  		old.Value = f(old.Value)
    56  		old.Time = future.UnixNano()
    57  		m.store.Set(key, old)
    58  	}
    59  	m.mu.Unlock()
    60  	if ok {
    61  		return
    62  	}
    63  	var null V
    64  	m.Add(key, f(null))
    65  }
    66  
    67  func (m *Cache[K, V]) Add(key K, value V) {
    68  	m.mu.Lock()
    69  	defer m.mu.Unlock()
    70  
    71  	m.store.Set(key, lruValue[V]{
    72  		Time:  future.UnixNano(),
    73  		Value: value,
    74  	})
    75  
    76  	if m.store.Len() <= m.storeCap {
    77  		return
    78  	}
    79  
    80  	var k0 K
    81  	var v0 = lruValue[V]{Time: math.MaxInt64}
    82  	var e *hashItem[K, lruValue[V]]
    83  	for i := 0; i < Lookahead; i++ {
    84  		m.ptr, e = m.store.nextItem(m.ptr)
    85  		if e == nil {
    86  			m.ptr, e = m.store.nextItem(0)
    87  		}
    88  		if e.Value.Time < v0.Time {
    89  			k0, v0 = e.Key, e.Value
    90  		}
    91  		m.ptr = (m.ptr + 1) % len(m.store.items)
    92  	}
    93  	m.store.Delete(k0)
    94  	m.onEvict(k0, v0.Value)
    95  }
    96  
    97  func (m *Cache[K, V]) Get(k K) (V, bool) {
    98  	m.mu.Lock()
    99  	defer m.mu.Unlock()
   100  	v, ok := m.store.Find(k)
   101  	if ok {
   102  		v.Time = future.UnixNano()
   103  		m.store.Set(k, v)
   104  	}
   105  	return v.Value, ok
   106  }
   107  
   108  func (m *Cache[K, V]) Clear() {
   109  	m.mu.Lock()
   110  	defer m.mu.Unlock()
   111  	m.store.Clear()
   112  }
   113  
   114  func (m *Cache[K, V]) Delete(key K) V {
   115  	m.mu.Lock()
   116  	defer m.mu.Unlock()
   117  	old, ok := m.store.Delete(key)
   118  	if ok {
   119  		m.onEvict(key, old.Value)
   120  	}
   121  	return old.Value
   122  }
   123  
   124  func (m *Cache[K, V]) Range(f func(K, V) bool) {
   125  	m.mu.RLock()
   126  	defer m.mu.RUnlock()
   127  	m.store.Foreach(func(k K, v *lruValue[V]) bool {
   128  		return f(k, v.Value)
   129  	})
   130  }