From d77be80b40d30357440c62641d88e21bcbf95b3d Mon Sep 17 00:00:00 2001
From: yuhan6665 <1588741+yuhan6665@users.noreply.github.com>
Date: Tue, 19 Oct 2021 12:53:29 -0400
Subject: [PATCH] Refinement: LRU (#1142) (#775)

Co-authored-by: Loyalsoldier <10487845+Loyalsoldier@users.noreply.github.com>
---
 common/cache/lru.go | 38 +++++++++++++++++++++-----------------
 1 file changed, 21 insertions(+), 17 deletions(-)

diff --git a/common/cache/lru.go b/common/cache/lru.go
index 65f10a3e..7c427ace 100644
--- a/common/cache/lru.go
+++ b/common/cache/lru.go
@@ -2,7 +2,7 @@ package cache
 
 import (
 	"container/list"
-	sync "sync"
+	"sync"
 )
 
 // Lru simple, fast lru cache implementation
@@ -28,7 +28,7 @@ type lruElement struct {
 
 // NewLru init a lru cache
 func NewLru(cap int) Lru {
-	return lru{
+	return &lru{
 		capacity:         cap,
 		doubleLinkedlist: list.New(),
 		keyToElement:     new(sync.Map),
@@ -37,49 +37,53 @@ func NewLru(cap int) Lru {
 	}
 }
 
-func (l lru) Get(key interface{}) (value interface{}, ok bool) {
+func (l *lru) Get(key interface{}) (value interface{}, ok bool) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
 	if v, ok := l.keyToElement.Load(key); ok {
 		element := v.(*list.Element)
-		l.doubleLinkedlist.MoveBefore(element, l.doubleLinkedlist.Front())
-		return element.Value.(lruElement).value, true
+		l.doubleLinkedlist.MoveToFront(element)
+		return element.Value.(*lruElement).value, true
 	}
 	return nil, false
 }
 
-func (l lru) GetKeyFromValue(value interface{}) (key interface{}, ok bool) {
+func (l *lru) GetKeyFromValue(value interface{}) (key interface{}, ok bool) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
 	if k, ok := l.valueToElement.Load(value); ok {
 		element := k.(*list.Element)
-		l.doubleLinkedlist.MoveBefore(element, l.doubleLinkedlist.Front())
-		return element.Value.(lruElement).key, true
+		l.doubleLinkedlist.MoveToFront(element)
+		return element.Value.(*lruElement).key, true
 	}
 	return nil, false
 }
 
-func (l lru) PeekKeyFromValue(value interface{}) (key interface{}, ok bool) {
+func (l *lru) PeekKeyFromValue(value interface{}) (key interface{}, ok bool) {
 	if k, ok := l.valueToElement.Load(value); ok {
 		element := k.(*list.Element)
-		return element.Value.(lruElement).key, true
+		return element.Value.(*lruElement).key, true
 	}
 	return nil, false
 }
 
-func (l lru) Put(key, value interface{}) {
-	e := lruElement{key, value}
+func (l *lru) Put(key, value interface{}) {
+	l.mu.Lock()
+	e := &lruElement{key, value}
 	if v, ok := l.keyToElement.Load(key); ok {
 		element := v.(*list.Element)
 		element.Value = e
-		l.doubleLinkedlist.MoveBefore(element, l.doubleLinkedlist.Front())
+		l.doubleLinkedlist.MoveToFront(element)
 	} else {
-		l.mu.Lock()
 		element := l.doubleLinkedlist.PushFront(e)
 		l.keyToElement.Store(key, element)
 		l.valueToElement.Store(value, element)
 		if l.doubleLinkedlist.Len() > l.capacity {
 			toBeRemove := l.doubleLinkedlist.Back()
 			l.doubleLinkedlist.Remove(toBeRemove)
-			l.keyToElement.Delete(toBeRemove.Value.(lruElement).key)
-			l.valueToElement.Delete(toBeRemove.Value.(lruElement).value)
+			l.keyToElement.Delete(toBeRemove.Value.(*lruElement).key)
+			l.valueToElement.Delete(toBeRemove.Value.(*lruElement).value)
 		}
-		l.mu.Unlock()
 	}
+	l.mu.Unlock()
 }