LRU implemented

Author: Shatwik7

keyVal.go

@@ -1,28 +1,92 @@
 package main
 
-import "sync"
+import (
+	"container/list"
+	"sync"
+	"time"
+)
 
 type KeyVal struct {
-	mu   sync.RWMutex
-	data map[string][]byte
+	mu           sync.RWMutex
+	data         map[string][]byte
+	lruList      *list.List
+	keyToElement map[string]*list.Element
+	memoryLimit  int64
+	currentSize  int64
 }
 
-func NewKeyVal() *KeyVal {
-	return &KeyVal{
-		data: map[string][]byte{},
+type lruItem struct {
+	key   string
+	value []byte
+}
+
+// memoryLimit should be in MB
+func NewKeyVal(memoryLimit int64) *KeyVal {
+	memoryLimit = memoryLimit * 1024 * 1024 //Size in mb
+	kv := &KeyVal{
+		data:         make(map[string][]byte),
+		lruList:      list.New(),
+		keyToElement: make(map[string]*list.Element),
+		memoryLimit:  memoryLimit,
+		currentSize:  0,
 	}
+	go kv.startEviction() // Start the background eviction goroutine
+	return kv
 }
+func (kv *KeyVal) Set(key, val []byte) error {
+	kv.mu.Lock()
+	defer kv.mu.Unlock()
+
+	keyStr := string(key)
 
-func (KeyVal *KeyVal) Set(key, val []byte) error {
-	KeyVal.mu.Lock()
-	defer KeyVal.mu.Unlock()
-	KeyVal.data[string(key)] = []byte(val)
+	if existingElem, exists := kv.keyToElement[keyStr]; exists {
+		existingItem := existingElem.Value.(*lruItem)
+		kv.currentSize -= int64(len(existingItem.value))
+		existingItem.value = val
+		kv.currentSize += int64(len(val))
+		kv.lruList.MoveToFront(existingElem)
+	} else {
+		item := &lruItem{key: keyStr, value: val}
+		element := kv.lruList.PushFront(item)
+		kv.keyToElement[keyStr] = element
+		kv.data[keyStr] = val
+		kv.currentSize += int64(len(val))
+	}
+	kv.evictIfNeeded()
 	return nil
 }
 
-func (KeyVal *KeyVal) Get(key []byte) ([]byte, bool) {
-	KeyVal.mu.RLock()
-	defer KeyVal.mu.RUnlock()
-	val, ok := KeyVal.data[string(key)]
-	return val, ok
+func (kv *KeyVal) Get(key []byte) ([]byte, bool) {
+	kv.mu.Lock()
+	defer kv.mu.Unlock()
+
+	keyStr := string(key)
+	if element, exists := kv.keyToElement[keyStr]; exists {
+		kv.lruList.MoveToFront(element)
+		return kv.data[keyStr], true
+	}
+	return nil, false
+}
+
+func (kv *KeyVal) evictIfNeeded() {
+	for kv.currentSize > kv.memoryLimit {
+		lastElem := kv.lruList.Back()
+		if lastElem == nil {
+			break
+		}
+		item := lastElem.Value.(*lruItem)
+		delete(kv.data, item.key)
+		kv.lruList.Remove(lastElem)
+		delete(kv.keyToElement, item.key)
+		kv.currentSize -= int64(len(item.value))
+	}
+}
+
+func (kv *KeyVal) startEviction() {
+	for {
+		time.Sleep(10 * time.Second)
+		kv.mu.Lock()
+		kv.evictIfNeeded()
+		kv.mu.Unlock()
+	}
 }

keyVal_test.go

@@ -5,7 +5,7 @@ import (
 )
 
 func TestKeyVal(t *testing.T) {
-	KeyVal := NewKeyVal()
+	KeyVal := NewKeyVal(2)
 	key := "mykey"
 	val := "myval"
 	err := KeyVal.Set([]byte(key), []byte(val))
@@ -14,9 +14,9 @@ func TestKeyVal(t *testing.T) {
 	}
 	val2, ok := KeyVal.Get([]byte(key))
 	if !ok {
-		t.Fatal(err)
+		t.Fatal("CAN NOT GET THE VALUE")
 	}
 	if string(val2) != val {
-		t.Fatal(err)
+		t.Fatal("VALUE DOES NOT MATCH")
 	}
 }

main.go

@@ -42,7 +42,7 @@ func NewServer(cfg Config) *Server {
 		quitCh:    make(chan struct{}),
 		msgCh:     make(chan Message),
 		delPeerCh: make(chan *Peer),
-		KeyVal:    NewKeyVal(),
+		KeyVal:    NewKeyVal(10),
 	}
 }