internal/datasource: limit cache size

Put a limit on the cache size so long-lived servers don't run out of
memory.

Hashicorp's simple cache implementation is concurrency-safe, so we no
longer need a lock.

For golang/go#47780

Change-Id: I4fcf1f6caf74333001c3e6d79e08baa50318a26b
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/344590
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
TryBot-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Julie Qiu <julie@golang.org>
diff --git a/go.mod b/go.mod
index ddcaf3f..a3e9e88 100644
--- a/go.mod
+++ b/go.mod
@@ -28,6 +28,7 @@
 	github.com/google/go-replayers/httpreplay v0.1.0
 	github.com/google/licensecheck v0.3.1
 	github.com/google/safehtml v0.0.2
+	github.com/hashicorp/golang-lru v0.5.1
 	github.com/jackc/pgconn v1.9.0
 	github.com/jackc/pgx/v4 v4.12.0
 	github.com/jba/templatecheck v0.6.0
diff --git a/go.sum b/go.sum
index 5e636e6..9c9559c 100644
--- a/go.sum
+++ b/go.sum
@@ -322,6 +322,7 @@
 github.com/hashicorp/go-version v1.2.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09ZGVZPK5anwXA=
 github.com/hashicorp/go.net v0.0.1/go.mod h1:hjKkEWcCURg++eb33jQU7oqQcI9XDCnUzHA0oac0k90=
 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/logutils v1.0.0/go.mod h1:QIAnNjmIWmVIIkWDTG1z5v++HQmx9WQRO+LraFDTW64=
 github.com/hashicorp/mdns v1.0.0/go.mod h1:tL+uN++7HEJ6SQLQ2/p+z2pH24WQKWjBPkE0mNTz8vQ=
diff --git a/internal/datasource/datasource.go b/internal/datasource/datasource.go
index a8c966e..ff6b4da 100644
--- a/internal/datasource/datasource.go
+++ b/internal/datasource/datasource.go
@@ -8,8 +8,7 @@
 package datasource
 
 import (
-	"sync"
-
+	lru "github.com/hashicorp/golang-lru"
 	"golang.org/x/pkgsite/internal"
 	"golang.org/x/pkgsite/internal/fetch"
 	"golang.org/x/pkgsite/internal/source"
@@ -20,8 +19,7 @@
 type dataSource struct {
 	sourceClient *source.Client
 
-	mu    sync.Mutex
-	cache map[internal.Modver]cacheEntry
+	cache *lru.Cache
 }
 
 // cacheEntry holds a fetched module or an error, if the fetch failed.
@@ -30,31 +28,34 @@
 	err    error
 }
 
+const maxCachedModules = 100
+
 func newDataSource(sc *source.Client) *dataSource {
+	cache, err := lru.New(maxCachedModules)
+	if err != nil {
+		// Can only happen if size is bad.
+		panic(err)
+	}
 	return &dataSource{
 		sourceClient: sc,
-		cache:        map[internal.Modver]cacheEntry{},
+		cache:        cache,
 	}
 }
 
 // cacheGet returns information from the cache if it is present, and (nil, nil) otherwise.
 func (ds *dataSource) cacheGet(path, version string) (*internal.Module, error) {
-	ds.mu.Lock()
-	defer ds.mu.Unlock()
-	// Look for an exact match first.
-	if e, ok := ds.cache[internal.Modver{Path: path, Version: version}]; ok {
-		return e.module, e.err
-	}
-	// Look for the module path with LocalVersion, as for a directory-based or GOPATH-mode module.
-	if e, ok := ds.cache[internal.Modver{Path: path, Version: fetch.LocalVersion}]; ok {
-		return e.module, e.err
+	// Look for an exact match first, then use LocalVersion, as for a
+	// 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.module, e.err
+		}
 	}
 	return nil, nil
 }
 
 // cachePut puts information into the cache.
 func (ds *dataSource) cachePut(path, version string, m *internal.Module, err error) {
-	ds.mu.Lock()
-	defer ds.mu.Unlock()
-	ds.cache[internal.Modver{Path: path, Version: version}] = cacheEntry{m, err}
+	ds.cache.Add(internal.Modver{Path: path, Version: version}, cacheEntry{m, err})
 }