gopls/internal/lsp/cache: record parse keys when they're created

When we parse a file through snapshot.ParseGo, we must record the parse
key by its URI for later invalidation. This wasn't done, resulting in a
memory leak.

Fix the leak by actually recording parse keys in the parseKeysByURI map,
which was previously always empty.

Write a test that diffs the cache before and after a change, which would
have caught this bug (and others).

Fixes golang/go#57222

Change-Id: I308812bf1030276dff08c26d359433750f44849a
Reviewed-on: https://go-review.googlesource.com/c/tools/+/456642
Reviewed-by: Alan Donovan <adonovan@google.com>
Run-TryBot: Robert Findley <rfindley@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>

Author: findleyr

gopls/internal/lsp/cache/parse.go

@@ -59,7 +59,7 @@ func (s *snapshot) ParseGo(ctx context.Context, fh source.FileHandle, mode sourc
 
 	// cache miss?
 	if !hit {
-		handle, release := s.store.Promise(key, func(ctx context.Context, arg interface{}) interface{} {
+		promise, release := s.store.Promise(key, func(ctx context.Context, arg interface{}) interface{} {
 			parsed, err := parseGoImpl(ctx, arg.(*snapshot).FileSet(), fh, mode)
 			return parseGoResult{parsed, err}
 		})
@@ -70,8 +70,30 @@ func (s *snapshot) ParseGo(ctx context.Context, fh source.FileHandle, mode sourc
 			entry = prev
 			release()
 		} else {
-			entry = handle
+			entry = promise
 			s.parsedGoFiles.Set(key, entry, func(_, _ interface{}) { release() })
+
+			// In order to correctly invalidate the key above, we must keep track of
+			// the parse key just created.
+			//
+			// TODO(rfindley): use a two-level map URI->parseKey->promise.
+			keys, _ := s.parseKeysByURI.Get(fh.URI())
+
+			// Only record the new key if it doesn't exist. This is overly cautious:
+			// we should only be setting the key if it doesn't exist. However, this
+			// logic will be replaced soon, and erring on the side of caution seemed
+			// wise.
+			foundKey := false
+			for _, existing := range keys {
+				if existing == key {
+					foundKey = true
+					break
+				}
+			}
+			if !foundKey {
+				keys = append(keys, key)
+				s.parseKeysByURI.Set(fh.URI(), keys)
+			}
 		}
 		s.mu.Unlock()
 	}

gopls/internal/regtest/misc/leak_test.go

@@ -0,0 +1,82 @@
+// Copyright 2022 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 misc
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"golang.org/x/tools/gopls/internal/hooks"
+	"golang.org/x/tools/gopls/internal/lsp/cache"
+	"golang.org/x/tools/gopls/internal/lsp/debug"
+	"golang.org/x/tools/gopls/internal/lsp/fake"
+	"golang.org/x/tools/gopls/internal/lsp/lsprpc"
+	. "golang.org/x/tools/gopls/internal/lsp/regtest"
+	"golang.org/x/tools/internal/jsonrpc2"
+	"golang.org/x/tools/internal/jsonrpc2/servertest"
+)
+
+// Test for golang/go#57222.
+func TestCacheLeak(t *testing.T) {
+	const files = `-- a.go --
+package a
+
+func _() {
+	println("1")
+}
+`
+	c := cache.New(nil, nil)
+	env := setupEnv(t, files, c)
+	env.Await(InitialWorkspaceLoad)
+	env.OpenFile("a.go")
+
+	// Make a couple edits to stabilize cache state.
+	//
+	// For some reason, after only one edit we're left with two parsed files
+	// (perhaps because something had to ParseHeader). If this test proves flaky,
+	// we'll need to investigate exactly what is causing various parse modes to
+	// be present (or rewrite the test to be more tolerant, for example make ~100
+	// modifications and assert that we're within a few of where we're started).
+	env.RegexpReplace("a.go", "1", "2")
+	env.RegexpReplace("a.go", "2", "3")
+	env.AfterChange()
+
+	// Capture cache state, make an arbitrary change, and wait for gopls to do
+	// its work. Afterward, we should have the exact same number of parsed
+	before := c.MemStats()
+	env.RegexpReplace("a.go", "3", "4")
+	env.AfterChange()
+	after := c.MemStats()
+
+	if diff := cmp.Diff(before, after); diff != "" {
+		t.Errorf("store objects differ after change (-before +after)\n%s", diff)
+	}
+}
+
+// setupEnv creates a new sandbox environment for editing the txtar encoded
+// content of files. It uses a new gopls instance backed by the Cache c.
+func setupEnv(t *testing.T, files string, c *cache.Cache) *Env {
+	ctx := debug.WithInstance(context.Background(), "", "off")
+	server := lsprpc.NewStreamServer(c, false, hooks.Options)
+	ts := servertest.NewPipeServer(server, jsonrpc2.NewRawStream)
+	s, err := fake.NewSandbox(&fake.SandboxConfig{
+		Files: fake.UnpackTxt(files),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	a := NewAwaiter(s.Workdir)
+	e, err := fake.NewEditor(s, fake.EditorConfig{}).Connect(ctx, ts, a.Hooks())
+
+	return &Env{
+		T:       t,
+		Ctx:     ctx,
+		Editor:  e,
+		Sandbox: s,
+		Awaiter: a,
+	}
+}