internal/lru: add a really simple LRU cache implementation

To replace github.com/hashicorp/golang-lru. The size used with the
cache in fetchdatasource is 100, so the efficiency of the cache is not
super important.

For golang/go#61399

Change-Id: I48383a4d1c00c4d153c0bad7b0a1a44c026b2314
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/524458
TryBot-Result: Gopher Robot <gobot@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Run-TryBot: Michael Matloob <matloob@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>
kokoro-CI: kokoro <noreply+kokoro@google.com>

go.mod

@@ -23,7 +23,6 @@ require (
 	github.com/google/go-replayers/httpreplay v1.0.0
 	github.com/google/licensecheck v0.3.1
 	github.com/google/safehtml v0.0.3-0.20211026203422-d6f0e11a5516
-	github.com/hashicorp/golang-lru v0.5.1
 	github.com/jackc/pgconn v1.10.1
 	github.com/jackc/pgx/v4 v4.14.1
 	github.com/jba/templatecheck v0.6.0

go.sum

@@ -617,7 +617,6 @@ github.com/hashicorp/go-multierror v1.0.0/go.mod h1:dHtQlpGsu+cZNNAkkCN/P3hoUDHh
 github.com/hashicorp/go-multierror v1.1.0 h1:B9UzwGQJehnUY1yNrnwREHc3fGbC2xefo8g4TbElacI=
 github.com/hashicorp/go-multierror v1.1.0/go.mod h1:spPvp8C1qA32ftKqdAHm4hHTbPw+vmowP0z+KUhOZdA=
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
-github.com/hashicorp/golang-lru v0.5.1 h1:0hERBMJE1eitiLkihrMvRVBYAkpHzc/J3QdDN+dAcgU=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=

internal/fetchdatasource/fetchdatasource.go

@@ -16,12 +16,12 @@ import (
 	"strings"
 	"time"
 
-	lru "github.com/hashicorp/golang-lru"
 	"golang.org/x/mod/semver"
 	"golang.org/x/pkgsite/internal"
 	"golang.org/x/pkgsite/internal/derrors"
 	"golang.org/x/pkgsite/internal/fetch"
 	"golang.org/x/pkgsite/internal/log"
+	"golang.org/x/pkgsite/internal/lru"
 	"golang.org/x/pkgsite/internal/proxy"
 	"golang.org/x/pkgsite/internal/version"
 )
@@ -30,7 +30,7 @@ import (
 // fetch.ModuleGetters to fetch modules and caching the results.
 type FetchDataSource struct {
 	opts  Options
-	cache *lru.Cache
+	cache *lru.Cache[internal.Modver, cacheEntry]
 }
 
 // Options are parameters for creating a new FetchDataSource.
@@ -45,11 +45,8 @@ type Options struct {
 
 // New creates a new FetchDataSource from the options.
 func (o Options) New() *FetchDataSource {
-	cache, err := lru.New(maxCachedModules)
-	if err != nil {
-		// Can only happen if size is bad, and we control it.
-		panic(err)
-	}
+	cache := lru.New[internal.Modver, cacheEntry](maxCachedModules)
+
 	opts := o
 	// Copy getters slice so caller doesn't modify us.
 	opts.Getters = make([]fetch.ModuleGetter, len(opts.Getters))
@@ -75,7 +72,6 @@ func (ds *FetchDataSource) cacheGet(path, version string) (fetch.ModuleGetter, *
 	// directory-based or GOPATH-mode module.
 	for _, v := range []string{version, fetch.LocalVersion} {
 		if e, ok := ds.cache.Get(internal.Modver{Path: path, Version: v}); ok {
-			e := e.(cacheEntry)
 			return e.g, e.module, e.err
 		}
 	}
@@ -84,7 +80,7 @@ func (ds *FetchDataSource) cacheGet(path, version string) (fetch.ModuleGetter, *
 
 // cachePut puts information into the cache.
 func (ds *FetchDataSource) cachePut(g fetch.ModuleGetter, path, version string, m *internal.Module, err error) {
-	ds.cache.Add(internal.Modver{Path: path, Version: version}, cacheEntry{g, m, err})
+	ds.cache.Put(internal.Modver{Path: path, Version: version}, cacheEntry{g, m, err})
 }
 
 // getModule gets the module at the given path and version. It first checks the

internal/lru/lru.go

@@ -0,0 +1,82 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package lru provides an LRU cache.
+package lru
+
+import (
+	"fmt"
+	"math"
+	"sync"
+)
+
+// Cache is an LRU cache.
+type Cache[K comparable, V any] struct {
+	mu      sync.Mutex
+	size    int
+	entries map[K]*entry[V]
+	tick    uint // increases every time an entry is used
+}
+
+type entry[V any] struct {
+	lastUsed uint // the tick of the last operation
+	v        V
+}
+
+// New returns a new Cache. Size must be positive or it will panic.
+func New[K comparable, V any](size int) *Cache[K, V] {
+	if size < 1 {
+		panic(fmt.Errorf("lru.New called with non-positive size %v", size))
+	}
+	return &Cache[K, V]{
+		size:    size,
+		entries: map[K]*entry[V]{},
+	}
+}
+
+// Get gets the entry for k in the Cache.
+func (c *Cache[K, V]) Get(k K) (V, bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	entry, ok := c.entries[k]
+	if !ok {
+		var zero V
+		return zero, false
+	}
+	c.tick++
+	entry.lastUsed = c.tick
+	return entry.v, true
+}
+
+// Put puts in an entry for k, v in Cache, evicting
+// the least recently used entry if necessary.
+func (c *Cache[K, V]) Put(k K, v V) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	_, ok := c.entries[k]
+	if !ok {
+		// k not already in c.entries. We need to evict the least recently
+		// used entry.
+		if c.size < 1 {
+			panic("attempting to insert into an uninitialized cache.")
+		}
+		if len(c.entries) > c.size {
+			panic(fmt.Errorf("size of cache, %d, has grown beyond size limit %d", len(c.entries), c.size))
+		}
+		if len(c.entries) == c.size {
+			// evict least recently used element.
+			var oldestTick uint = math.MaxUint
+			var oldestKey K
+			for k, e := range c.entries {
+				if e.lastUsed <= oldestTick {
+					oldestTick = e.lastUsed
+					oldestKey = k
+				}
+			}
+			delete(c.entries, oldestKey)
+		}
+	}
+	c.tick++
+	c.entries[k] = &entry[V]{lastUsed: c.tick, v: v}
+}

internal/lru/lru_test.go

@@ -0,0 +1,96 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package lru
+
+import "testing"
+
+func TestSizeOne(t *testing.T) {
+	c := New[int, int](1)
+	c.Put(1, 1)
+	got, gotOK := c.Get(1)
+	if got != 1 && gotOK != true {
+		t.Errorf("c.Get(1): got %v, %v, want %v, %v", got, gotOK, 1, true)
+	}
+	c.Put(2, 2)
+	got, gotOK = c.Get(1)
+	if got != 0 || gotOK != false {
+		t.Errorf("c.Get(1): got %v, %v, want %v, %v", got, gotOK, 0, false)
+	}
+	got, gotOK = c.Get(2)
+	if got != 2 && gotOK != true {
+		t.Errorf("c.Get(2): got %v, %v, want %v, %v", got, gotOK, 2, true)
+	}
+}
+
+func TestSizeFive(t *testing.T) {
+	c := New[int, int](5)
+	c.Put(1, 1)
+	c.Put(2, 2)
+	c.Put(3, 3)
+	c.Put(4, 4)
+	c.Put(5, 5)
+
+	getHasKey := func(k int, has bool) {
+		t.Helper()
+
+		got, ok := c.Get(k)
+		if has == false {
+			if ok == true || got != 0 {
+				t.Errorf("c.Get(%v): got %v, %v, want %v, %v", k, got, ok, 0, false)
+			}
+		} else if got != k || ok != true {
+			t.Errorf("c.Get(%v): got %v, %v, want %v, %v", k, got, ok, k, true)
+		}
+	}
+
+	getHasKey(3, true)
+	getHasKey(2, true)
+	getHasKey(1, true)
+	getHasKey(5, true)
+	getHasKey(4, true)
+	c.Put(6, 6) // 3 gets evicted
+
+	getHasKey(3, false)
+	getHasKey(1, true)
+	getHasKey(2, true)
+	getHasKey(4, true)
+	getHasKey(5, true)
+	getHasKey(6, true)
+	c.Put(7, 7)
+	c.Put(8, 8) // 1 and 2 get evicted
+
+	getHasKey(1, false)
+	getHasKey(2, false)
+	getHasKey(8, true)
+	getHasKey(7, true)
+	getHasKey(4, true)
+	getHasKey(5, true)
+	getHasKey(6, true)
+	c.Put(9, 9)   // 8 gets evicted
+	c.Put(10, 10) // 7 gets evicted
+	c.Put(11, 11) // 4 gets evicted
+	c.Put(12, 12) // 5 gets evicted
+	c.Put(13, 13) // 6 gets evicted
+	c.Put(14, 14) // 9 gets evicted
+
+	getHasKey(4, false)
+	getHasKey(5, false)
+	getHasKey(6, false)
+	getHasKey(7, false)
+	getHasKey(8, false)
+	getHasKey(9, false)
+	getHasKey(10, true)
+	getHasKey(11, true)
+	getHasKey(12, true)
+	getHasKey(13, true)
+	getHasKey(14, true)
+	c.Put(12, 12)
+
+	getHasKey(10, true)
+	getHasKey(11, true)
+	getHasKey(12, true)
+	getHasKey(13, true)
+	getHasKey(14, true)
+}