Go简易LRU缓存

45 阅读1分钟

Go 简易lru缓存

import (
	"container/list"
	"fmt"
)

type Entry struct {
	Key   string
	Value string
}

type Cache struct {
	maxBytes   int64
	totalBytes int64
	ll         *list.List
	data       map[string]*list.Element
	OnEvicted  func(key string, value string)
}

func NewCache(maxBytes int64) *Cache {
	return &Cache{
		maxBytes:   maxBytes,
		totalBytes: 0,
		ll:         list.New(),
		data:       make(map[string]*list.Element),
	}
}

func (c *Cache) Get(key string) (string, bool) {
	if el, ok := c.data[key]; ok {
		// 移动到链表头
		c.ll.MoveToFront(el)

		e := el.Value.(*Entry)
		return e.Value, true
	}
	return "", false
}

func (c *Cache) Add(key string, value string) {
	if el, ok := c.data[key]; ok {
		c.ll.MoveToFront(el)
		entry := el.Value.(*Entry)
		entry.Value = value
		c.totalBytes += int64(len(value)) - int64(len(entry.Value))
	} else {
		el := c.ll.PushFront(&Entry{Key: key, Value: value})
		c.data[key] = el
		c.totalBytes += int64(len(value))
	}

	if c.maxBytes > 0 && c.totalBytes > c.maxBytes {
		c.clearOldest()
	}
}

func (c *Cache) clearOldest() {
	el := c.ll.Back()
	if el != nil {
		c.ll.Remove(el)
		entry := el.Value.(*Entry)
		delete(c.data, entry.Key)
		c.totalBytes -= int64(len(entry.Value))
		if c.OnEvicted != nil {
			c.OnEvicted(entry.Key, entry.Value)
		}
	}
}

func main() {
	c := NewCache(10)
	c.OnEvicted = func(key, value string) {
		fmt.Printf("Entry(key: %s, value: %s) has been deleted\n", key, value)
	}
	c.Add("key1", "value1")
	c.Add("key2", "value2")

	fmt.Printf("total bytes: %d\n", c.totalBytes)

	if value, success := c.Get("key2"); success {
		fmt.Printf("GET Key: %s\n", value)
	}
}