Add in-memory cache implementation.

Author: tomdeering-wf

cache/cache.go

@@ -0,0 +1,190 @@
+package cache
+
+import (
+	"container/list"
+	"sync"
+)
+
+// Cache is a bounded-size in-memory cache of sized items with a configurable eviction policy
+type Cache interface {
+	// Get retrieves items from the cache by key.
+	// If an item for a particular key is not found, its position in the result will be nil.
+	Get(keys ...string) []Item
+
+	// Put adds an item to the cache.
+	Put(key string, item Item)
+
+	// Remove clears items with the given keys from the cache
+	Remove(keys ...string)
+
+	// Size returns the size of all items currently in the cache.
+	Size() uint64
+}
+
+// Item is an item in a cache
+type Item interface {
+	// Size returns the item's size, in bytes
+	Size() uint64
+}
+
+// A tuple tracking a cached item and a reference to its node in the eviction list
+type cached struct {
+	item    Item
+	element *list.Element
+}
+
+// Sets the provided list element on the cached item if it is not nil
+func (c *cached) setElementIfNotNil(element *list.Element) {
+	if element != nil {
+		c.element = element
+	}
+}
+
+// Private cache implementation
+type cache struct {
+	sync.Mutex                                  // Lock for synchronizing Get, Put, Remove
+	cap          uint64                         // Capacity bound
+	size         uint64                         // Cumulative size
+	items        map[string]*cached             // Map from keys to cached items
+	keyList      *list.List                     // List of cached items in order of increasing evictability
+	recordAdd    func(key string) *list.Element // Function called to indicate that an item with the given key was added
+	recordAccess func(key string) *list.Element // Function called to indicate that an item with the given key was accessed
+}
+
+// CacheOption configures a cache.
+type CacheOption func(*cache)
+
+// Capacity sets the fixed capacity bound on the cache, in bytes.
+// If not provided, default is 10MB.
+func Capacity(cap uint64) CacheOption {
+	return func(c *cache) {
+		c.cap = cap
+	}
+}
+
+// Policy is a cache eviction policy for use with the EvictionPolicy CacheOption.
+type Policy uint8
+
+const (
+	// LeastRecentlyAdded indicates a least-recently-added eviction policy.
+	LeastRecentlyAdded Policy = iota
+	// LeastRecentlyUsed indicates a least-recently-used eviction policy.
+	LeastRecentlyUsed
+)
+
+// EvictionPolicy sets the eviction policy to be used to make room for new items.
+// If not provided, default is LeastRecentlyUsed.
+func EvictionPolicy(policy Policy) CacheOption {
+	return func(c *cache) {
+		switch policy {
+		case LeastRecentlyAdded:
+			c.recordAccess = c.noop
+			c.recordAdd = c.record
+		case LeastRecentlyUsed:
+			c.recordAccess = c.record
+			c.recordAdd = c.noop
+		}
+	}
+}
+
+// New returns a cache with the requested options configured.
+// The cache consumes memory bounded by a fixed capacity,
+// plus tracking overhead linear in the number of items.
+func New(options ...CacheOption) Cache {
+	c := &cache{
+		cap:     10 * 1024 * 1024, // Default capacity of 10MB
+		keyList: list.New(),
+		items:   map[string]*cached{},
+	}
+	// Default LRU eviction policy
+	EvictionPolicy(LeastRecentlyUsed)(c)
+
+	for _, option := range options {
+		option(c)
+	}
+
+	return c
+}
+
+func (c *cache) Get(keys ...string) []Item {
+	c.Lock()
+	defer c.Unlock()
+
+	items := make([]Item, len(keys))
+	for i, key := range keys {
+		cached := c.items[key]
+		if cached == nil {
+			items[i] = nil
+		} else {
+			c.recordAccess(key)
+			items[i] = cached.item
+		}
+	}
+
+	return items
+}
+
+func (c *cache) Put(key string, item Item) {
+	c.Lock()
+	defer c.Unlock()
+
+	// Remove the item currently with this key (if any)
+	c.remove(key)
+
+	// Make sure there's room to add this item
+	c.ensureCapacity(item.Size())
+
+	// Actually add the new item
+	cached := &cached{item: item}
+	cached.setElementIfNotNil(c.recordAdd(key))
+	cached.setElementIfNotNil(c.recordAccess(key))
+	c.items[key] = cached
+	c.size += item.Size()
+}
+
+func (c *cache) Remove(keys ...string) {
+	c.Lock()
+	defer c.Unlock()
+
+	for _, key := range keys {
+		c.remove(key)
+	}
+}
+
+func (c *cache) Size() uint64 {
+	return c.size
+}
+
+// Given the need to add some number of new bytes to the cache,
+// evict items according to the eviction policy until there is room.
+// The caller should hold the cache lock.
+func (c *cache) ensureCapacity(toAdd uint64) {
+	mustRemove := int64(c.size+toAdd) - int64(c.cap)
+	for mustRemove > 0 {
+		key := c.keyList.Back().Value.(string)
+		mustRemove -= int64(c.items[key].item.Size())
+		c.remove(key)
+	}
+}
+
+// Remove the item associated with the given key.
+// The caller should hold the cache lock.
+func (c *cache) remove(key string) {
+	if cached, ok := c.items[key]; ok {
+		delete(c.items, key)
+		c.size -= cached.item.Size()
+		c.keyList.Remove(cached.element)
+	}
+}
+
+// A no-op function that does nothing for the provided key
+func (c *cache) noop(string) *list.Element { return nil }
+
+// A function to record the given key and mark it as last to be evicted
+func (c *cache) record(key string) *list.Element {
+	if item, ok := c.items[key]; ok {
+		c.keyList.MoveToFront(item.element)
+		return item.element
+	}
+	return c.keyList.PushFront(key)
+}

cache/cache_test.go

@@ -0,0 +1,148 @@
+package cache
+
+import (
+	"container/list"
+	"testing"
+)
+
+func TestCapacity(t *testing.T) {
+	c := &cache{}
+	Capacity(52)(c)
+	if c.cap != 52 {
+		t.Errorf("Capacity failed to set cap")
+	}
+}
+
+func TestEvictionPolicy(t *testing.T) {
+	c := &cache{keyList: list.New()}
+	EvictionPolicy(LeastRecentlyUsed)(c)
+	if accessed, added := c.recordAccess("foo"), c.recordAdd("foo"); accessed == nil || added != nil {
+		t.Errorf("EvictionPolicy failed to set LRU policy")
+	}
+
+	c = &cache{keyList: list.New()}
+	EvictionPolicy(LeastRecentlyAdded)(c)
+	if accessed, added := c.recordAccess("foo"), c.recordAdd("foo"); accessed != nil || added == nil {
+		t.Errorf("EvictionPolicy failed to set LRU policy")
+	}
+}
+
+func TestNew(t *testing.T) {
+	optionApplied := false
+	option := func(*cache) {
+		optionApplied = true
+	}
+
+	c := New(option).(*cache)
+
+	if c.cap != 10*1024*1024 {
+		t.Errorf("Expected default cache capacity of %d", 10*1024*1024)
+	}
+	if c.size != 0 {
+		t.Errorf("Expected initial size of zero")
+	}
+	if c.items == nil {
+		t.Errorf("Expected items to be initialized")
+	}
+	if c.keyList == nil {
+		t.Errorf("Expected keyList to be initialized")
+	}
+	if !optionApplied {
+		t.Errorf("New did not apply its provided option")
+	}
+	if accessed, added := c.recordAccess("foo"), c.recordAdd("foo"); accessed == nil || added != nil {
+		t.Errorf("Expected default LRU policy")
+	}
+}
+
+type testItem uint64
+
+func (ti testItem) Size() uint64 {
+	return uint64(ti)
+}
+
+func TestPutGetRemoveSize(t *testing.T) {
+	keys := []string{"foo", "bar", "baz"}
+	testCases := []struct {
+		label         string
+		cache         Cache
+		useCache      func(c Cache)
+		expectedSize  uint64
+		expectedItems []Item
+	}{{
+		label: "Items added, key doesn't exist",
+		cache: New(),
+		useCache: func(c Cache) {
+			c.Put("foo", testItem(1))
+		},
+		expectedSize:  1,
+		expectedItems: []Item{testItem(1), nil, nil},
+	}, {
+		label: "Items added, key exists",
+		cache: New(),
+		useCache: func(c Cache) {
+			c.Put("foo", testItem(1))
+			c.Put("foo", testItem(10))
+		},
+		expectedSize:  10,
+		expectedItems: []Item{testItem(10), nil, nil},
+	}, {
+		label: "Items added, LRA eviction",
+		cache: New(Capacity(2), EvictionPolicy(LeastRecentlyAdded)),
+		useCache: func(c Cache) {
+			c.Put("foo", testItem(1))
+			c.Put("bar", testItem(1))
+			c.Get("foo")
+			c.Put("baz", testItem(1))
+		},
+		expectedSize:  2,
+		expectedItems: []Item{nil, testItem(1), testItem(1)},
+	}, {
+		label: "Items added, LRU eviction",
+		cache: New(Capacity(2), EvictionPolicy(LeastRecentlyUsed)),
+		useCache: func(c Cache) {
+			c.Put("foo", testItem(1))
+			c.Put("bar", testItem(1))
+			c.Get("foo")
+			c.Put("baz", testItem(1))
+		},
+		expectedSize:  2,
+		expectedItems: []Item{testItem(1), nil, testItem(1)},
+	}, {
+		label: "Items removed, key doesn't exist",
+		cache: New(),
+		useCache: func(c Cache) {
+			c.Put("foo", testItem(1))
+			c.Remove("baz")
+		},
+		expectedSize:  1,
+		expectedItems: []Item{testItem(1), nil, nil},
+	}, {
+		label: "Items removed, key exists",
+		cache: New(),
+		useCache: func(c Cache) {
+			c.Put("foo", testItem(1))
+			c.Remove("foo")
+		},
+		expectedSize:  0,
+		expectedItems: []Item{nil, nil, nil},
+	}}
+
+	for _, testCase := range testCases {
+		t.Log(testCase.label)
+		testCase.useCache(testCase.cache)
+		if testCase.cache.Size() != testCase.expectedSize {
+			t.Errorf("Expected size of %d, got %d", testCase.expectedSize, testCase.cache.Size())
+		}
+		actual := testCase.cache.Get(keys...)
+		if len(actual) != len(testCase.expectedItems) {
+			t.Errorf("Expected to get %d items, got %d", len(testCase.expectedItems), len(actual))
+		} else {
+			for i, expectedItem := range testCase.expectedItems {
+				if actual[i] != expectedItem {
+					t.Errorf("Expected Get to return %v in position %d, got %v", expectedItem, i, actual[i])
+				}
+			}
+		}
+	}
+}