gopls/internal/regtest/bench: support benchmarking multiple repos

Extract benchmark state into a new repo type, so that we may run
benchmarks in multiple shared workspaces. Also, add missing cleanup
code.

Additionally, simplify to always run gopls in a separate process. This
means that the normal test profiling flags won't be useful, so add
support for threading through profiling flags to the external gopls
process.

For golang/go#53538

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

12 files changed, 304 insertions, 214 deletions

diff --git a/gopls/internal/regtest/bench/bench_test.go b/gopls/internal/regtest/bench/bench_test.go
index e285aa52c..fa06e274e 100644
--- a/gopls/internal/regtest/bench/bench_test.go
+++ b/gopls/internal/regtest/bench/bench_test.go
@@ -18,10 +18,8 @@ import (
 	"time"
 
 	"golang.org/x/tools/gopls/internal/hooks"
-	"golang.org/x/tools/gopls/internal/lsp/cache"
 	"golang.org/x/tools/gopls/internal/lsp/cmd"
 	"golang.org/x/tools/gopls/internal/lsp/fake"
-	"golang.org/x/tools/gopls/internal/lsp/lsprpc"
 	"golang.org/x/tools/internal/bug"
 	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/fakenet"
@@ -32,15 +30,21 @@ import (
 	. "golang.org/x/tools/gopls/internal/lsp/regtest"
 )
 
-// This package implements benchmarks that share a common editor session.
-//
-// It is a work-in-progress.
-//
-// Remaining TODO(rfindley):
-//   - add detailed documentation for how to write a benchmark, as a package doc
-//   - add benchmarks for more features
-//   - eliminate flags, and just run benchmarks on with a predefined set of
-//     arguments
+var (
+	goplsPath = flag.String("gopls_path", "", "if set, use this gopls for testing; incompatible with -gopls_commit")
+
+	installGoplsOnce sync.Once // guards installing gopls at -gopls_commit
+	goplsCommit      = flag.String("gopls_commit", "", "if set, install and use gopls at this commit for testing; incompatible with -gopls_path")
+
+	cpuProfile = flag.String("gopls_cpuprofile", "", "if set, the cpu profile file suffix; see \"Profiling\" in the package doc")
+	memProfile = flag.String("gopls_memprofile", "", "if set, the mem profile file suffix; see \"Profiling\" in the package doc")
+	trace      = flag.String("gopls_trace", "", "if set, the trace file suffix; see \"Profiling\" in the package doc")
+
+	// If non-empty, tempDir is a temporary working dir that was created by this
+	// test suite.
+	makeTempDirOnce sync.Once // guards creation of the temp dir
+	tempDir         string
+)
 
 // if runAsGopls is "true", run the gopls command instead of the testing.M.
 const runAsGopls = "_GOPLS_BENCH_RUN_AS_GOPLS"
@@ -52,56 +56,16 @@ func TestMain(m *testing.M) {
 		os.Exit(0)
 	}
 	event.SetExporter(nil) // don't log to stderr
-	code := doMain(m)
-	os.Exit(code)
-}
-
-func doMain(m *testing.M) (code int) {
-	defer func() {
-		if editor != nil {
-			if err := editor.Close(context.Background()); err != nil {
-				fmt.Fprintf(os.Stderr, "closing editor: %v", err)
-				if code == 0 {
-					code = 1
-				}
-			}
-		}
-		if tempDir != "" {
-			if err := os.RemoveAll(tempDir); err != nil {
-				fmt.Fprintf(os.Stderr, "cleaning temp dir: %v", err)
-				if code == 0 {
-					code = 1
-				}
-			}
+	code := m.Run()
+	if err := cleanup(); err != nil {
+		fmt.Fprintf(os.Stderr, "cleaning up after benchmarks: %v\n", err)
+		if code == 0 {
+			code = 1
 		}
-	}()
-	return m.Run()
+	}
+	os.Exit(code)
 }
 
-var (
-	workdir   = flag.String("workdir", "", "if set, working directory to use for benchmarks; overrides -repo and -commit")
-	repo      = flag.String("repo", "https://go.googlesource.com/tools", "if set (and -workdir is unset), run benchmarks in this repo")
-	file      = flag.String("file", "go/ast/astutil/util.go", "active file, for benchmarks that operate on a file")
-	commitish = flag.String("commit", "gopls/v0.9.0", "if set (and -workdir is unset), run benchmarks at this commit")
-
-	goplsPath   = flag.String("gopls_path", "", "if set, use this gopls for testing; incompatible with -gopls_commit")
-	goplsCommit = flag.String("gopls_commit", "", "if set, install and use gopls at this commit for testing; incompatible with -gopls_path")
-
-	// If non-empty, tempDir is a temporary working dir that was created by this
-	// test suite.
-	//
-	// The sync.Once variables guard various modifications of the temp directory.
-	makeTempDirOnce  sync.Once
-	checkoutRepoOnce sync.Once
-	installGoplsOnce sync.Once
-	tempDir          string
-
-	setupEditorOnce sync.Once
-	sandbox         *fake.Sandbox
-	editor          *fake.Editor
-	awaiter         *Awaiter
-)
-
 // getTempDir returns the temporary directory to use for benchmark files,
 // creating it if necessary.
 func getTempDir() string {
@@ -115,31 +79,6 @@ func getTempDir() string {
 	return tempDir
 }
 
-// benchmarkDir returns the directory to use for benchmarks.
-//
-// If -workdir is set, just use that directory. Otherwise, check out a shallow
-// copy of -repo at the given -commit, and clean up when the test suite exits.
-func benchmarkDir() string {
-	if *workdir != "" {
-		return *workdir
-	}
-	if *repo == "" {
-		log.Fatal("-repo must be provided if -workdir is unset")
-	}
-	if *commitish == "" {
-		log.Fatal("-commit must be provided if -workdir is unset")
-	}
-
-	dir := filepath.Join(getTempDir(), "repo")
-	checkoutRepoOnce.Do(func() {
-		log.Printf("creating working dir: checking out %s@%s to %s\n", *repo, *commitish, dir)
-		if err := shallowClone(dir, *repo, *commitish); err != nil {
-			log.Fatal(err)
-		}
-	})
-	return dir
-}
-
 // shallowClone performs a shallow clone of repo into dir at the given
 // 'commitish' ref (any commit reference understood by git).
 //
@@ -163,70 +102,6 @@ func shallowClone(dir, repo, commitish string) error {
 	return nil
 }
 
-// sharedEnv returns a shared benchmark environment.
-//
-// Every call to sharedEnv uses the same editor and sandbox. If -gopls_path and
-// -gopls_commit are unset, this environment will run gopls in-process.
-func sharedEnv(tb testing.TB) *Env {
-	setupEditorOnce.Do(func() {
-		dir := benchmarkDir()
-
-		var err error
-		ts := getServer()
-		sandbox, editor, awaiter, err = connectEditor(dir, fake.EditorConfig{}, ts)
-		if err != nil {
-			log.Fatalf("connecting editor: %v", err)
-		}
-
-		if err := awaiter.Await(context.Background(), InitialWorkspaceLoad); err != nil {
-			panic(err)
-		}
-	})
-
-	return &Env{
-		T:       tb,
-		Ctx:     context.Background(),
-		Editor:  editor,
-		Sandbox: sandbox,
-		Awaiter: awaiter,
-	}
-}
-
-// newEnv returns a new Env connected to separate gopls process communicating
-// over stdin/stdout.
-//
-// Every call to newEnv returns a different Env connected to a distinct gopls
-// process.
-//
-// TODO(rfindley): consolidate gopls server construction: always use a sidecar,
-// and make it easy to collect profiles.
-func newEnv(dir string, tb testing.TB) *Env {
-	goplsPath := getGoplsPath()
-	if goplsPath == "" {
-		var err error
-		goplsPath, err = os.Executable()
-		if err != nil {
-			tb.Fatal(err)
-		}
-	}
-	ts := &SidecarServer{
-		goplsPath: goplsPath,
-		env:       []string{fmt.Sprintf("%s=true", runAsGopls)},
-	}
-	server, editor, awaiter, err := connectEditor(dir, fake.EditorConfig{}, ts)
-	if err != nil {
-		tb.Fatalf("connecting editor: %v", err)
-	}
-
-	return &Env{
-		T:       tb,
-		Ctx:     context.Background(),
-		Editor:  editor,
-		Sandbox: server,
-		Awaiter: awaiter,
-	}
-}
-
 // connectEditor connects a fake editor session in the given dir, using the
 // given editor config.
 func connectEditor(dir string, config fake.EditorConfig, ts servertest.Connector) (*fake.Sandbox, *fake.Editor, *Awaiter, error) {
@@ -246,30 +121,41 @@ func connectEditor(dir string, config fake.EditorConfig, ts servertest.Connector
 	return s, e, a, nil
 }
 
-// getServer returns a server connector that either starts a new in-process
-// server, or starts a separate gopls process.
-func getServer() servertest.Connector {
+// newGoplsServer returns a connector that connects to a new gopls process.
+func newGoplsServer(name string) (servertest.Connector, error) {
 	if *goplsPath != "" && *goplsCommit != "" {
 		panic("can't set both -gopls_path and -gopls_commit")
 	}
-	if path := getGoplsPath(); path != "" {
-		return &SidecarServer{goplsPath: *goplsPath}
+	var (
+		goplsPath = *goplsPath
+		env       []string
+	)
+	if *goplsCommit != "" {
+		goplsPath = getInstalledGopls()
 	}
-	server := lsprpc.NewStreamServer(cache.New(nil, nil), false, hooks.Options)
-	return servertest.NewPipeServer(server, jsonrpc2.NewRawStream)
-}
-
-// getGoplsPath returns the path to the external gopls binary to use for
-// benchmarks, or the empty string if no external gopls is configured via
-// -gopls_path or -gopls_commit.
-func getGoplsPath() string {
-	if *goplsPath != "" {
-		return *goplsPath
+	if goplsPath == "" {
+		var err error
+		goplsPath, err = os.Executable()
+		if err != nil {
+			return nil, err
+		}
+		env = []string{fmt.Sprintf("%s=true", runAsGopls)}
 	}
-	if *goplsCommit != "" {
-		return getInstalledGopls()
+	var args []string
+	if *cpuProfile != "" {
+		args = append(args, fmt.Sprintf("-profile.cpu=%s", name+"."+*cpuProfile))
+	}
+	if *memProfile != "" {
+		args = append(args, fmt.Sprintf("-profile.mem=%s", name+"."+*memProfile))
 	}
-	return ""
+	if *trace != "" {
+		args = append(args, fmt.Sprintf("-profile.trace=%s", name+"."+*trace))
+	}
+	return &SidecarServer{
+		goplsPath: goplsPath,
+		env:       env,
+		args:      args,
+	}, nil
 }
 
 // getInstalledGopls builds gopls at the given -gopls_commit, returning the
@@ -307,11 +193,18 @@ func getInstalledGopls() string {
 type SidecarServer struct {
 	goplsPath string
 	env       []string // additional environment bindings
+	args      []string // command-line arguments
 }
 
 // Connect creates new io.Pipes and binds them to the underlying StreamServer.
+//
+// It implements the servertest.Connector interface.
 func (s *SidecarServer) Connect(ctx context.Context) jsonrpc2.Conn {
-	cmd := exec.CommandContext(ctx, s.goplsPath, "serve")
+	// Note: don't use CommandContext here, as we want gopls to exit gracefully
+	// in order to write out profile data.
+	//
+	// We close the connection on context cancelation below.
+	cmd := exec.Command(s.goplsPath, s.args...)
 
 	stdin, err := cmd.StdinPipe()
 	if err != nil {
@@ -321,15 +214,34 @@ func (s *SidecarServer) Connect(ctx context.Context) jsonrpc2.Conn {
 	if err != nil {
 		log.Fatal(err)
 	}
-	cmd.Stderr = os.Stdout
+	cmd.Stderr = os.Stderr
 	cmd.Env = append(os.Environ(), s.env...)
 	if err := cmd.Start(); err != nil {
 		log.Fatalf("starting gopls: %v", err)
 	}
 
-	go cmd.Wait() // to free resources; error is ignored
+	go func() {
+		// If we don't log.Fatal here, benchmarks may hang indefinitely if gopls
+		// exits abnormally.
+		//
+		// TODO(rfindley): ideally we would shut down the connection gracefully,
+		// but that doesn't currently work.
+		if err := cmd.Wait(); err != nil {
+			log.Fatalf("gopls invocation failed with error: %v", err)
+		}
+	}()
 
 	clientStream := jsonrpc2.NewHeaderStream(fakenet.NewConn("stdio", stdout, stdin))
 	clientConn := jsonrpc2.NewConn(clientStream)
+
+	go func() {
+		select {
+		case <-ctx.Done():
+			clientConn.Close()
+			clientStream.Close()
+		case <-clientConn.Done():
+		}
+	}()
+
 	return clientConn
 }
diff --git a/gopls/internal/regtest/bench/completion_test.go b/gopls/internal/regtest/bench/completion_test.go
index f597ab9d7..a89a4ff68 100644
--- a/gopls/internal/regtest/bench/completion_test.go
+++ b/gopls/internal/regtest/bench/completion_test.go
@@ -5,14 +5,11 @@
 package bench
 
 import (
-	"context"
 	"fmt"
 	"testing"
 
 	"golang.org/x/tools/gopls/internal/lsp/protocol"
 	. "golang.org/x/tools/gopls/internal/lsp/regtest"
-
-	"golang.org/x/tools/gopls/internal/lsp/fake"
 )
 
 type completionBenchOptions struct {
@@ -24,32 +21,9 @@ type completionBenchOptions struct {
 }
 
 func benchmarkCompletion(options completionBenchOptions, b *testing.B) {
-	dir := benchmarkDir()
-
-	// Use a new environment for each test, to avoid any existing state from the
-	// previous session.
-	sandbox, editor, awaiter, err := connectEditor(dir, fake.EditorConfig{
-		Settings: map[string]interface{}{
-			"completionBudget": "1m", // arbitrary long completion budget
-		},
-	}, getServer())
-	if err != nil {
-		b.Fatal(err)
-	}
-	ctx := context.Background()
-	defer func() {
-		if err := editor.Close(ctx); err != nil {
-			b.Errorf("closing editor: %v", err)
-		}
-	}()
-
-	env := &Env{
-		T:       b,
-		Ctx:     ctx,
-		Editor:  editor,
-		Sandbox: sandbox,
-		Awaiter: awaiter,
-	}
+	repo := repos["tools"]
+	env := repo.newEnv(b)
+	defer env.Close()
 
 	// Run edits required for this completion.
 	if options.setup != nil {
diff --git a/gopls/internal/regtest/bench/definition_test.go b/gopls/internal/regtest/bench/definition_test.go
index 20b75de73..cdffcf654 100644
--- a/gopls/internal/regtest/bench/definition_test.go
+++ b/gopls/internal/regtest/bench/definition_test.go
@@ -4,10 +4,12 @@
 
 package bench
 
-import "testing"
+import (
+	"testing"
+)
 
-func BenchmarkGoToDefinition(b *testing.B) {
-	env := sharedEnv(b)
+func BenchmarkDefinition(b *testing.B) {
+	env := repos["tools"].sharedEnv(b)
 
 	env.OpenFile("internal/imports/mod.go")
 	loc := env.RegexpSearch("internal/imports/mod.go", "ModuleJSON")
diff --git a/gopls/internal/regtest/bench/didchange_test.go b/gopls/internal/regtest/bench/didchange_test.go
index 4e6bd232b..e18ad4e9a 100644
--- a/gopls/internal/regtest/bench/didchange_test.go
+++ b/gopls/internal/regtest/bench/didchange_test.go
@@ -18,16 +18,17 @@ import (
 //
 // Uses -workdir and -file to control where the edits occur.
 func BenchmarkDidChange(b *testing.B) {
-	env := sharedEnv(b)
-	env.OpenFile(*file)
-	env.Await(env.DoneWithOpen())
+	env := repos["tools"].sharedEnv(b)
+	const filename = "go/ast/astutil/util.go"
+	env.OpenFile(filename)
+	env.AfterChange()
 
 	// Insert the text we'll be modifying at the top of the file.
-	env.EditBuffer(*file, protocol.TextEdit{NewText: "// __REGTEST_PLACEHOLDER_0__\n"})
+	env.EditBuffer(filename, protocol.TextEdit{NewText: "// __REGTEST_PLACEHOLDER_0__\n"})
 
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		env.EditBuffer(*file, protocol.TextEdit{
+		env.EditBuffer(filename, protocol.TextEdit{
 			Range: protocol.Range{
 				Start: protocol.Position{Line: 0, Character: 0},
 				End:   protocol.Position{Line: 1, Character: 0},
diff --git a/gopls/internal/regtest/bench/doc.go b/gopls/internal/regtest/bench/doc.go
new file mode 100644
index 000000000..a9f2fbffa
--- /dev/null
+++ b/gopls/internal/regtest/bench/doc.go
@@ -0,0 +1,33 @@
+// 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.
+
+// The bench package implements benchmarks for various LSP operations.
+//
+// Benchmarks check out specific commits of popular and/or exemplary
+// repositories, and script an external gopls process via a fake text editor.
+// By default, benchmarks run the test executable as gopls (using a special
+// "gopls mode" environment variable). A different gopls binary may be used by
+// setting the -gopls_path or -gopls_commit flags.
+//
+// This package is a work in progress.
+//
+// # Profiling
+//
+// As benchmark functions run gopls in a separate process, the normal test
+// flags for profiling are not useful. Instead the -gopls_cpuprofile,
+// -gopls_memprofile, and -gopls_trace flags may be used to pass through
+// profiling flags to the gopls process. Each of these flags sets a suffix
+// for the respective gopls profiling flag, which is prefixed with a name
+// corresponding to the shared repository or (in some cases) benchmark name.
+// For example, settings -gopls_cpuprofile=cpu.out will result in profiles
+// named tools.cpu.out, BenchmarkInitialWorkspaceLoad.cpu.out, etc. Here,
+// tools.cpu.out is the cpu profile for the shared x/tools session, which may
+// be used by multiple benchmark functions, and BenchmarkInitialWorkspaceLoad
+// is the cpu profile for the last iteration of the initial workspace load
+// test, which starts a new editor session for each iteration.
+//
+// # TODO
+//   - add more benchmarks, and more repositories
+//   - improve this documentation
+package bench
diff --git a/gopls/internal/regtest/bench/hover_test.go b/gopls/internal/regtest/bench/hover_test.go
index 78fdc930a..ebd89b87d 100644
--- a/gopls/internal/regtest/bench/hover_test.go
+++ b/gopls/internal/regtest/bench/hover_test.go
@@ -9,7 +9,7 @@ import (
 )
 
 func BenchmarkHover(b *testing.B) {
-	env := sharedEnv(b)
+	env := repos["tools"].sharedEnv(b)
 
 	env.OpenFile("internal/imports/mod.go")
 	loc := env.RegexpSearch("internal/imports/mod.go", "bytes")
diff --git a/gopls/internal/regtest/bench/implementations_test.go b/gopls/internal/regtest/bench/implementations_test.go
index 610a2d28c..7f83987d0 100644
--- a/gopls/internal/regtest/bench/implementations_test.go
+++ b/gopls/internal/regtest/bench/implementations_test.go
@@ -6,8 +6,8 @@ package bench
 
 import "testing"
 
-func BenchmarkFindAllImplementations(b *testing.B) {
-	env := sharedEnv(b)
+func BenchmarkImplementations(b *testing.B) {
+	env := repos["tools"].sharedEnv(b)
 
 	env.OpenFile("internal/imports/mod.go")
 	loc := env.RegexpSearch("internal/imports/mod.go", "initAllMods")
diff --git a/gopls/internal/regtest/bench/iwl_test.go b/gopls/internal/regtest/bench/iwl_test.go
index 87df19974..44cdc785f 100644
--- a/gopls/internal/regtest/bench/iwl_test.go
+++ b/gopls/internal/regtest/bench/iwl_test.go
@@ -15,12 +15,16 @@ import (
 // BenchmarkInitialWorkspaceLoad benchmarks the initial workspace load time for
 // a new editing session.
 func BenchmarkInitialWorkspaceLoad(b *testing.B) {
-	dir := benchmarkDir()
+	repo := repos["tools"]
 	b.ResetTimer()
 
 	for i := 0; i < b.N; i++ {
-		env := newEnv(dir, b)
-		// TODO(rfindley): this depends on the repository being x/tools. Fix this.
+		// Exclude the time to set up the env from the benchmark time, as this may
+		// involve installing gopls and/or checking out the repo dir.
+		b.StopTimer()
+		env := repo.newEnv(b)
+		b.StartTimer()
+
 		env.OpenFile("internal/lsp/diagnostics.go")
 		env.Await(InitialWorkspaceLoad)
 		b.StopTimer()
diff --git a/gopls/internal/regtest/bench/references_test.go b/gopls/internal/regtest/bench/references_test.go
index e5f1f63df..782275053 100644
--- a/gopls/internal/regtest/bench/references_test.go
+++ b/gopls/internal/regtest/bench/references_test.go
@@ -7,7 +7,7 @@ package bench
 import "testing"
 
 func BenchmarkReferences(b *testing.B) {
-	env := sharedEnv(b)
+	env := repos["tools"].sharedEnv(b)
 
 	env.OpenFile("internal/imports/mod.go")
 	loc := env.RegexpSearch("internal/imports/mod.go", "gopathwalk")
diff --git a/gopls/internal/regtest/bench/rename_test.go b/gopls/internal/regtest/bench/rename_test.go
index e6db663a4..7339c7625 100644
--- a/gopls/internal/regtest/bench/rename_test.go
+++ b/gopls/internal/regtest/bench/rename_test.go
@@ -10,7 +10,7 @@ import (
 )
 
 func BenchmarkRename(b *testing.B) {
-	env := sharedEnv(b)
+	env := repos["tools"].sharedEnv(b)
 
 	env.OpenFile("internal/imports/mod.go")
 	env.Await(env.DoneWithOpen())
diff --git a/gopls/internal/regtest/bench/repo_test.go b/gopls/internal/regtest/bench/repo_test.go
new file mode 100644
index 000000000..9b7ce72c6
--- /dev/null
+++ b/gopls/internal/regtest/bench/repo_test.go
@@ -0,0 +1,164 @@
+// 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 bench
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"log"
+	"os"
+	"path/filepath"
+	"sync"
+	"testing"
+
+	"golang.org/x/tools/gopls/internal/lsp/fake"
+	. "golang.org/x/tools/gopls/internal/lsp/regtest"
+)
+
+// repos holds shared repositories for use in benchmarks.
+var repos = map[string]*repo{
+	"tools": {name: "tools", url: "https://go.googlesource.com/tools", commit: "gopls/v0.9.0"},
+}
+
+// A repo represents a working directory for a repository checked out at a
+// specific commit.
+//
+// Repos are used for sharing state across benchmarks that operate on the same
+// codebase.
+type repo struct {
+	// static configuration
+	name   string // must be unique, used for subdirectory
+	url    string // repo url
+	commit string // commitish, e.g. tag or short commit hash
+
+	dirOnce sync.Once
+	dir     string // directory contaning source code checked out to url@commit
+
+	// shared editor state
+	editorOnce sync.Once
+	editor     *fake.Editor
+	sandbox    *fake.Sandbox
+	awaiter    *Awaiter
+}
+
+// getDir returns directory containing repo source code, creating it if
+// necessary. It is safe for concurrent use.
+func (r *repo) getDir() string {
+	r.dirOnce.Do(func() {
+		r.dir = filepath.Join(getTempDir(), r.name)
+		log.Printf("cloning %s@%s into %s", r.url, r.commit, r.dir)
+		if err := shallowClone(r.dir, r.url, r.commit); err != nil {
+			log.Fatal(err)
+		}
+	})
+	return r.dir
+}
+
+// sharedEnv returns a shared benchmark environment. It is safe for concurrent
+// use.
+//
+// Every call to sharedEnv uses the same editor and sandbox, as a means to
+// avoid reinitializing the editor for large repos. Calling repo.Close cleans
+// up the shared environment.
+//
+// Repos in the package-local Repos var are closed at the end of the test main
+// function.
+func (r *repo) sharedEnv(tb testing.TB) *Env {
+	r.editorOnce.Do(func() {
+		dir := r.getDir()
+
+		ts, err := newGoplsServer(r.name)
+		if err != nil {
+			log.Fatal(err)
+		}
+		r.sandbox, r.editor, r.awaiter, err = connectEditor(dir, fake.EditorConfig{}, ts)
+		if err != nil {
+			log.Fatalf("connecting editor: %v", err)
+		}
+
+		if err := r.awaiter.Await(context.Background(), InitialWorkspaceLoad); err != nil {
+			log.Fatal(err)
+		}
+	})
+
+	return &Env{
+		T:       tb,
+		Ctx:     context.Background(),
+		Editor:  r.editor,
+		Sandbox: r.sandbox,
+		Awaiter: r.awaiter,
+	}
+}
+
+// newEnv returns a new Env connected to a new gopls process communicating
+// over stdin/stdout. It is safe for concurrent use.
+//
+// It is the caller's responsibility to call Close on the resulting Env when it
+// is no longer needed.
+func (r *repo) newEnv(tb testing.TB) *Env {
+	dir := r.getDir()
+
+	ts, err := newGoplsServer(tb.Name())
+	if err != nil {
+		tb.Fatal(err)
+	}
+	sandbox, editor, awaiter, err := connectEditor(dir, fake.EditorConfig{}, ts)
+	if err != nil {
+		log.Fatalf("connecting editor: %v", err)
+	}
+
+	return &Env{
+		T:       tb,
+		Ctx:     context.Background(),
+		Editor:  editor,
+		Sandbox: sandbox,
+		Awaiter: awaiter,
+	}
+}
+
+// Close cleans up shared state referenced by the repo.
+func (r *repo) Close() error {
+	var errBuf bytes.Buffer
+	if r.editor != nil {
+		if err := r.editor.Close(context.Background()); err != nil {
+			fmt.Fprintf(&errBuf, "closing editor: %v", err)
+		}
+	}
+	if r.sandbox != nil {
+		if err := r.sandbox.Close(); err != nil {
+			fmt.Fprintf(&errBuf, "closing sandbox: %v", err)
+		}
+	}
+	if r.dir != "" {
+		if err := os.RemoveAll(r.dir); err != nil {
+			fmt.Fprintf(&errBuf, "cleaning dir: %v", err)
+		}
+	}
+	if errBuf.Len() > 0 {
+		return errors.New(errBuf.String())
+	}
+	return nil
+}
+
+// cleanup cleans up state that is shared across benchmark functions.
+func cleanup() error {
+	var errBuf bytes.Buffer
+	for _, repo := range repos {
+		if err := repo.Close(); err != nil {
+			fmt.Fprintf(&errBuf, "closing %q: %v", repo.name, err)
+		}
+	}
+	if tempDir != "" {
+		if err := os.RemoveAll(tempDir); err != nil {
+			fmt.Fprintf(&errBuf, "cleaning tempDir: %v", err)
+		}
+	}
+	if errBuf.Len() > 0 {
+		return errors.New(errBuf.String())
+	}
+	return nil
+}
diff --git a/gopls/internal/regtest/bench/workspace_symbols_test.go b/gopls/internal/regtest/bench/workspace_symbols_test.go
index 482425c38..ac9ad531b 100644
--- a/gopls/internal/regtest/bench/workspace_symbols_test.go
+++ b/gopls/internal/regtest/bench/workspace_symbols_test.go
@@ -15,7 +15,7 @@ var symbolQuery = flag.String("symbol_query", "test", "symbol query to use in be
 // BenchmarkWorkspaceSymbols benchmarks the time to execute a workspace symbols
 // request (controlled by the -symbol_query flag).
 func BenchmarkWorkspaceSymbols(b *testing.B) {
-	env := sharedEnv(b)
+	env := repos["tools"].sharedEnv(b)
 
 	// Make an initial symbol query to warm the cache.
 	symbols := env.Symbol(*symbolQuery)