github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/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  	Put(key, value interface{})
    13  }
    14  
    15  type lru struct {
    16  	capacity         int
    17  	doubleLinkedlist *list.List
    18  	keyToElement     *sync.Map
    19  	valueToElement   *sync.Map
    20  	mu               *sync.Mutex
    21  }
    22  
    23  type lruElement struct {
    24  	key   interface{}
    25  	value interface{}
    26  }
    27  
    28  // NewLru initializes a lru cache
    29  func NewLru(cap int) Lru {
    30  	return &lru{
    31  		capacity:         cap,
    32  		doubleLinkedlist: list.New(),
    33  		keyToElement:     new(sync.Map),
    34  		valueToElement:   new(sync.Map),
    35  		mu:               new(sync.Mutex),
    36  	}
    37  }
    38  
    39  func (l *lru) Get(key interface{}) (value interface{}, ok bool) {
    40  	l.mu.Lock()
    41  	defer l.mu.Unlock()
    42  	if v, ok := l.keyToElement.Load(key); ok {
    43  		element := v.(*list.Element)
    44  		l.doubleLinkedlist.MoveToFront(element)
    45  		return element.Value.(*lruElement).value, true
    46  	}
    47  	return nil, false
    48  }
    49  
    50  func (l *lru) GetKeyFromValue(value interface{}) (key interface{}, ok bool) {
    51  	l.mu.Lock()
    52  	defer l.mu.Unlock()
    53  	if k, ok := l.valueToElement.Load(value); ok {
    54  		element := k.(*list.Element)
    55  		l.doubleLinkedlist.MoveToFront(element)
    56  		return element.Value.(*lruElement).key, true
    57  	}
    58  	return nil, false
    59  }
    60  
    61  func (l *lru) Put(key, value interface{}) {
    62  	l.mu.Lock()
    63  	e := &lruElement{key, value}
    64  	if v, ok := l.keyToElement.Load(key); ok {
    65  		element := v.(*list.Element)
    66  		element.Value = e
    67  		l.doubleLinkedlist.MoveToFront(element)
    68  	} else {
    69  		element := l.doubleLinkedlist.PushFront(e)
    70  		l.keyToElement.Store(key, element)
    71  		l.valueToElement.Store(value, element)
    72  		if l.doubleLinkedlist.Len() > l.capacity {
    73  			toBeRemove := l.doubleLinkedlist.Back()
    74  			l.doubleLinkedlist.Remove(toBeRemove)
    75  			l.keyToElement.Delete(toBeRemove.Value.(*lruElement).key)
    76  			l.valueToElement.Delete(toBeRemove.Value.(*lruElement).value)
    77  		}
    78  	}
    79  	l.mu.Unlock()
    80  }