Merge pull request #2 from sixt/feature/lru-cache

in-memory lru cache of limited capacity

.travis.yml

@@ -1,7 +1,6 @@
 language: go
 go:
-  - master
-  - "1.11.x"
+  - "1.x"
 
 env: GO111MODULE=on
 

cmd/gomodproxy/main.go

@@ -55,6 +55,7 @@ func main() {
 	verbose := flag.Bool("v", false, "verbose logging")
 	json := flag.Bool("json", false, "json structured logging")
 	dir := flag.String("dir", filepath.Join(os.Getenv("HOME"), ".gomodproxy"), "cache directory")
+	memLimit := flag.Int64("mem", 256, "in-memory cache size in MB")
 	flag.Var(&gitPaths, "git", "list of git settings")
 
 	flag.Parse()
@@ -68,13 +69,15 @@ func main() {
 	fmt.Println("Listening on", ln.Addr())
 
 	options := []api.Option{}
+	var logger func(...interface{})
 	if *verbose || *json {
 		if *json {
-			options = append(options, api.Log(jsonLog))
+			logger = jsonLog
 		} else {
-			options = append(options, api.Log(prettyLog))
+			logger = prettyLog
 		}
 	}
+	options = append(options, api.Log(logger))
 
 	for _, path := range gitPaths {
 		kv := strings.SplitN(path, "=", 2)
@@ -84,7 +87,10 @@ func main() {
 		options = append(options, api.Git(kv[0], kv[1]))
 	}
 
-	options = append(options, api.Memory(), api.CacheDir(*dir))
+	options = append(options,
+		api.Memory(logger, *memLimit*1024*1024),
+		api.CacheDir(*dir),
+	)
 
 	sigc := make(chan os.Signal, 1)
 	signal.Notify(sigc, os.Interrupt)

pkg/api/api.go

@@ -73,9 +73,9 @@ func Git(prefix string, auth string) Option {
 }
 
 // Memory configures API to use in-memory cache for downloaded modules.
-func Memory() Option {
+func Memory(log logger, limit int64) Option {
 	return func(api *api) {
-		api.stores = append(api.stores, store.Memory())
+		api.stores = append(api.stores, store.Memory(log, limit))
 	}
 }
 

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 {