github.com/moqsien/xraycore@v1.8.5/common/cache/lru.go (about)

     1  package cache
     2  
     3  import (
     4  	"container/list"
     5  	"sync"
     6  )
     7  
     8  // Lru simple, fast lru cache implementation
     9  type Lru interface {
    10  	Get(key interface{}) (value interface{}, ok bool)
    11  	GetKeyFromValue(value interface{}) (key interface{}, ok bool)
    12  	PeekKeyFromValue(value interface{}) (key interface{}, ok bool) // Peek means check but NOT bring to top
    13  	Put(key, value interface{})
    14  }
    15  
    16  type lru struct {
    17  	capacity         int
    18  	doubleLinkedlist *list.List
    19  	keyToElement     *sync.Map
    20  	valueToElement   *sync.Map
    21  	mu               *sync.Mutex
    22  }
    23  
    24  type lruElement struct {
    25  	key   interface{}
    26  	value interface{}
    27  }
    28  
    29  // NewLru initializes a lru cache
    30  func NewLru(cap int) Lru {
    31  	return &lru{
    32  		capacity:         cap,
    33  		doubleLinkedlist: list.New(),
    34  		keyToElement:     new(sync.Map),
    35  		valueToElement:   new(sync.Map),
    36  		mu:               new(sync.Mutex),
    37  	}
    38  }
    39  
    40  func (l *lru) Get(key interface{}) (value interface{}, ok bool) {
    41  	l.mu.Lock()
    42  	defer l.mu.Unlock()
    43  	if v, ok := l.keyToElement.Load(key); ok {
    44  		element := v.(*list.Element)
    45  		l.doubleLinkedlist.MoveToFront(element)
    46  		return element.Value.(*lruElement).value, true
    47  	}
    48  	return nil, false
    49  }
    50  
    51  func (l *lru) GetKeyFromValue(value interface{}) (key interface{}, ok bool) {
    52  	l.mu.Lock()
    53  	defer l.mu.Unlock()
    54  	if k, ok := l.valueToElement.Load(value); ok {
    55  		element := k.(*list.Element)
    56  		l.doubleLinkedlist.MoveToFront(element)
    57  		return element.Value.(*lruElement).key, true
    58  	}
    59  	return nil, false
    60  }
    61  
    62  func (l *lru) PeekKeyFromValue(value interface{}) (key interface{}, ok bool) {
    63  	if k, ok := l.valueToElement.Load(value); ok {
    64  		element := k.(*list.Element)
    65  		return element.Value.(*lruElement).key, true
    66  	}
    67  	return nil, false
    68  }
    69  
    70  func (l *lru) Put(key, value interface{}) {
    71  	l.mu.Lock()
    72  	e := &lruElement{key, value}
    73  	if v, ok := l.keyToElement.Load(key); ok {
    74  		element := v.(*list.Element)
    75  		element.Value = e
    76  		l.doubleLinkedlist.MoveToFront(element)
    77  	} else {
    78  		element := l.doubleLinkedlist.PushFront(e)
    79  		l.keyToElement.Store(key, element)
    80  		l.valueToElement.Store(value, element)
    81  		if l.doubleLinkedlist.Len() > l.capacity {
    82  			toBeRemove := l.doubleLinkedlist.Back()
    83  			l.doubleLinkedlist.Remove(toBeRemove)
    84  			l.keyToElement.Delete(toBeRemove.Value.(*lruElement).key)
    85  			l.valueToElement.Delete(toBeRemove.Value.(*lruElement).value)
    86  		}
    87  	}
    88  	l.mu.Unlock()
    89  }