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

     1  package bitmap
     2  
     3  import (
     4  	"container/list"
     5  	"sync"
     6  )
     7  
     8  type Cache struct {
     9  	maxWeight int64
    10  	curWeight int64
    11  
    12  	ll    *list.List
    13  	cache map[string]*list.Element
    14  
    15  	sync.Mutex
    16  }
    17  
    18  type entry struct {
    19  	key    string
    20  	value  *Range
    21  	weight int64
    22  }
    23  
    24  func NewLRUCache(maxWeight int64) *Cache {
    25  	return &Cache{
    26  		maxWeight: maxWeight,
    27  		ll:        list.New(),
    28  		cache:     make(map[string]*list.Element),
    29  	}
    30  }
    31  
    32  func (c *Cache) Add(key string, value *Range) {
    33  	if c.maxWeight <= 0 {
    34  		return
    35  	}
    36  	weight := value.RoughSizeBytes()
    37  	if weight > c.maxWeight {
    38  		weight = c.maxWeight
    39  	}
    40  
    41  	c.Lock()
    42  	defer c.Unlock()
    43  
    44  	if ee, ok := c.cache[key]; ok {
    45  		e := ee.Value.(*entry)
    46  		c.ll.MoveToFront(ee)
    47  		diff := weight - e.weight
    48  		e.weight = weight
    49  		e.value = value
    50  		c.curWeight += diff
    51  	} else {
    52  		c.curWeight += weight
    53  		ele := c.ll.PushFront(&entry{key, value, weight})
    54  		c.cache[key] = ele
    55  	}
    56  
    57  	for c.maxWeight > 0 && c.curWeight > c.maxWeight {
    58  		last := c.ll.Back()
    59  		kv := last.Value.(*entry)
    60  		c.ll.Remove(last)
    61  		c.curWeight -= last.Value.(*entry).weight
    62  		delete(c.cache, kv.key)
    63  	}
    64  }
    65  
    66  func (c *Cache) Get(key string) *Range {
    67  	c.Lock()
    68  	defer c.Unlock()
    69  
    70  	if ele, hit := c.cache[key]; hit {
    71  		e := ele.Value.(*entry)
    72  		c.ll.MoveToFront(ele)
    73  		return e.value
    74  	}
    75  	return nil
    76  }
    77  
    78  func (c *Cache) Len() (len int) {
    79  	c.Lock()
    80  	len = c.ll.Len()
    81  	c.Unlock()
    82  	return
    83  }