in-memory lru cache of limited capacity

fix memLimit flag usage

use megabytes as a unit for memory store capacity

fix incorrect item rearrangement in lru cache

use latest 1.x Go on trais

pkg/store/mem.go

@@ -10,33 +10,98 @@ import (
 
 type memory struct {
 	sync.Mutex
-	cache []Snapshot
+	log   logger
+	limit int64
+	size  int64
+	head  *lruItem
+	tail  *lruItem
 }
 
-// Memory creates an in-memory cache.
-func Memory() Store { return &memory{} }
+type lruItem struct {
+	Snapshot
+	prev *lruItem
+	next *lruItem
+}
+
+// Memory creates an in-memory LRU cache.
+func Memory(log logger, limit int64) Store { return &memory{log: log, limit: limit} }
 
 func (m *memory) Put(ctx context.Context, snapshot Snapshot) error {
 	m.Lock()
 	defer m.Unlock()
-	for _, item := range m.cache {
-		if item.Module == snapshot.Module && item.Version == snapshot.Version {
-			return nil
-		}
+	if _, err := m.lookup(snapshot.Module, snapshot.Version); err == nil {
+		return nil
+	}
+
+	item := &lruItem{Snapshot: snapshot, next: m.head}
+	m.insert(item)
+
+	for m.limit >= 0 && m.size > m.limit {
+		m.evict()
 	}
-	m.cache = append(m.cache, snapshot)
 	return nil
 }
 
 func (m *memory) Get(ctx context.Context, module string, version vcs.Version) (Snapshot, error) {
 	m.Lock()
 	defer m.Unlock()
-	for _, snapshot := range m.cache {
-		if snapshot.Module == module && snapshot.Version == version {
-			return snapshot, nil
+	return m.lookup(module, version)
+}
+
+func (m *memory) lookup(module string, version vcs.Version) (Snapshot, error) {
+	for item := m.head; item != nil; item = item.next {
+		if item.Module == module && item.Version == version {
+			m.update(item)
+			return item.Snapshot, nil
 		}
 	}
 	return Snapshot{}, errors.New("not found")
 }
 
+func (m *memory) insert(item *lruItem) {
+	m.log("mem.insert",
+		"module", item.Module, "version", item.Version, "size", len(item.Data),
+		"cachesize", m.size, "cachelimit", m.limit)
+	m.size = m.size + int64(len(item.Data))
+	if m.head == nil {
+		m.head = item
+		m.tail = item
+		return
+	}
+	item.next = m.head
+	m.head.prev = item
+	m.head = item
+}
+
+func (m *memory) update(item *lruItem) {
+	m.log("mem.update", "module", item.Module, "version", item.Version, "size", len(item.Data),
+		"cachesize", m.size, "cachelimit", m.limit)
+	if item.prev == nil {
+		return
+	}
+	item.prev.next = item.next
+	if item.next == nil {
+		m.tail = item.prev
+	} else {
+		item.next.prev = item.prev
+	}
+	item.prev = nil
+	item.next = m.head
+	m.head.prev = item
+	m.head = item
+}
+
+func (m *memory) evict() {
+	m.log("mem.evict", "module", m.tail.Module, "version", m.tail.Version, "size", len(m.tail.Data),
+		"cachesize", m.size, "cachelimit", m.limit)
+	m.size = m.size - int64(len(m.tail.Data))
+	if m.tail.prev == nil {
+		m.head = nil
+		m.tail = nil
+		return
+	}
+	m.tail.prev.next = nil
+	m.tail = m.tail.prev
+}
+
 func (m *memory) Close() error { return nil }

pkg/store/mem_test.go

@@ -0,0 +1,74 @@
+package store
+
+import (
+	"context"
+	"math/rand"
+	"testing"
+)
+
+func TestMemoryStore(t *testing.T) {
+	ctx := context.Background()
+	m := Memory(t.Log, -1)
+	m.Put(ctx, Snapshot{Module: "foo", Version: "v1.0.0", Data: []byte("hello")})
+	m.Put(ctx, Snapshot{Module: "bar", Version: "v1.0.0", Data: []byte{}})
+	m.Put(ctx, Snapshot{Module: "baz", Version: "v1.0.0", Data: []byte("world")})
+	if res, err := m.Get(ctx, "foo", "v1.0.0"); err != nil {
+		t.Fatal(err)
+	} else if res.Module != "foo" || res.Version != "v1.0.0" || string(res.Data) != "hello" {
+		t.Fatal(res)
+	}
+}
+
+func TestMemoryStoreOverflow(t *testing.T) {
+	ctx := context.Background()
+	m := Memory(t.Log, 10)
+	m.Put(ctx, Snapshot{Module: "foo", Version: "v1.0.0", Data: make([]byte, 4)})
+	m.Put(ctx, Snapshot{Module: "bar", Version: "v1.0.0", Data: make([]byte, 7)})
+
+	// "foo" should be removed, because adding "bar" exceeds the capacity
+	if res, err := m.Get(ctx, "foo", "v1.0.0"); err == nil {
+		t.Fatal(res)
+	} else if _, err := m.Get(ctx, "bar", "v1.0.0"); err != nil {
+		t.Fatal(err)
+	}
+
+	m.Put(ctx, Snapshot{Module: "baz", Version: "v1.0.0", Data: make([]byte, 3)})
+
+	// both "bar" and "baz" should be in store
+	if _, err := m.Get(ctx, "bar", "v1.0.0"); err != nil {
+		t.Fatal(err)
+	} else if _, err := m.Get(ctx, "baz", "v1.0.0"); err != nil {
+		t.Fatal(err)
+	}
+
+	m.Get(ctx, "bar", "v1.0.0")
+	m.Put(ctx, Snapshot{Module: "qux", Version: "v1.0.0", Data: make([]byte, 3)})
+
+	// "bar" should remain in store, since it was accessed recently, "baz" should be removed
+	if _, err := m.Get(ctx, "bar", "v1.0.0"); err != nil {
+		t.Fatal(err)
+	} else if res, err := m.Get(ctx, "baz", "v1.0.0"); err == nil {
+		t.Fatal(res)
+	}
+}
+
+func TestMemoryStoreRandom(t *testing.T) {
+	snaphots := []Snapshot{
+		Snapshot{Module: "a", Version: "v1.0.0", Data: make([]byte, 1)},
+		Snapshot{Module: "b", Version: "v1.0.0", Data: make([]byte, 3)},
+		Snapshot{Module: "c", Version: "v1.0.0", Data: make([]byte, 5)},
+		Snapshot{Module: "d", Version: "v1.0.0", Data: make([]byte, 7)},
+		Snapshot{Module: "e", Version: "v1.0.0", Data: make([]byte, 11)},
+		Snapshot{Module: "f", Version: "v1.0.0", Data: make([]byte, 13)},
+	}
+
+	m := Memory(t.Log, 12)
+	for i := 0; i < 100; i++ {
+		ctx := context.Background()
+		if rand.Int()%5 > 2 {
+			m.Put(ctx, snaphots[rand.Intn(len(snaphots))])
+		} else {
+			m.Get(ctx, snaphots[rand.Intn(len(snaphots))].Module, "v1.0.0")
+		}
+	}
+}

pkg/store/store.go

@@ -7,6 +7,8 @@ import (
 	"github.com/sixt/gomodproxy/pkg/vcs"
 )
 
+type logger = func(...interface{})
+
 // Store is an interface for a typical cache. It allows to put a snapshot and
 // to get snapshot of the specific version.
 type Store interface {