gopls/internal/regtest: automate counting of editor notifications to await

Using awaiting a certain number of work items in the regtest is a source
of flakes and latent bugs. This CL replaces all such assertions with
assertions based on the number of notifications sent by the fake.Editor.

While implementing this, I discovered several tests that had incorrect
counting, so this may fix some flakes.

Implementing this required pushing the asynchronous processing of file
events into the Editor, rather than the Workdir.

Change-Id: I9c3639409f2beed4a76295cbd53180c6e2ace126
Reviewed-on: https://go-review.googlesource.com/c/tools/+/285612
Run-TryBot: Robert Findley <rfindley@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Trust: Robert Findley <rfindley@google.com>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/gopls/internal/regtest/codelens_test.go b/gopls/internal/regtest/codelens_test.go
index a58de79..4d16dc5 100644
--- a/gopls/internal/regtest/codelens_test.go
+++ b/gopls/internal/regtest/codelens_test.go
@@ -9,7 +9,6 @@
 	"strings"
 	"testing"
 
-	"golang.org/x/tools/internal/lsp"
 	"golang.org/x/tools/internal/lsp/fake"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/source"
@@ -137,7 +136,7 @@
 				}); err != nil {
 					t.Fatal(err)
 				}
-				env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1))
+				env.Await(env.DoneWithChangeWatchedFiles())
 				got := env.Editor.BufferText("go.mod")
 				const wantGoMod = `module mod.com
 
@@ -199,7 +198,7 @@
 	runner.Run(t, shouldRemoveDep, func(t *testing.T, env *Env) {
 		env.OpenFile("go.mod")
 		env.ExecuteCodeLensCommand("go.mod", source.CommandTidy)
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1))
+		env.Await(env.DoneWithChangeWatchedFiles())
 		got := env.Editor.BufferText("go.mod")
 		const wantGoMod = `module mod.com
 
diff --git a/gopls/internal/regtest/completion_test.go b/gopls/internal/regtest/completion_test.go
index 412fdab..4c72768 100644
--- a/gopls/internal/regtest/completion_test.go
+++ b/gopls/internal/regtest/completion_test.go
@@ -9,7 +9,6 @@
 	"strings"
 	"testing"
 
-	"golang.org/x/tools/internal/lsp"
 	"golang.org/x/tools/internal/lsp/fake"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/testenv"
@@ -128,7 +127,7 @@
 				if tc.content != nil {
 					env.WriteWorkspaceFile(tc.filename, *tc.content)
 					env.Await(
-						CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+						env.DoneWithChangeWatchedFiles(),
 					)
 				}
 				env.OpenFile(tc.filename)
@@ -249,24 +248,22 @@
 		// unimported completions, and then remove it before proceeding.
 		env.RemoveWorkspaceFile("main2.go")
 		env.RunGoCommand("mod", "tidy")
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 2))
+		env.Await(env.DoneWithChangeWatchedFiles())
 
 		// Trigger unimported completions for the example.com/blah package.
 		env.OpenFile("main.go")
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1))
+		env.Await(env.DoneWithOpen())
 		pos := env.RegexpSearch("main.go", "ah")
 		completions := env.Completion("main.go", pos)
 		if len(completions.Items) == 0 {
 			t.Fatalf("no completion items")
 		}
 		env.AcceptCompletion("main.go", pos, completions.Items[0])
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1))
+		env.Await(env.DoneWithChange())
 
 		// Trigger completions once again for the blah.<> selector.
 		env.RegexpReplace("main.go", "_ = blah", "_ = blah.")
-		env.Await(
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 2),
-		)
+		env.Await(env.DoneWithChange())
 		pos = env.RegexpSearch("main.go", "\n}")
 		completions = env.Completion("main.go", pos)
 		if len(completions.Items) != 1 {
diff --git a/gopls/internal/regtest/configuration_test.go b/gopls/internal/regtest/configuration_test.go
index b61a8a8..a197da7 100644
--- a/gopls/internal/regtest/configuration_test.go
+++ b/gopls/internal/regtest/configuration_test.go
@@ -7,7 +7,6 @@
 import (
 	"testing"
 
-	"golang.org/x/tools/internal/lsp"
 	"golang.org/x/tools/internal/lsp/fake"
 )
 
@@ -28,7 +27,7 @@
 	run(t, files, func(t *testing.T, env *Env) {
 		env.OpenFile("a/a.go")
 		env.Await(
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1),
+			env.DoneWithOpen(),
 			NoDiagnostics("a/a.go"),
 		)
 		cfg := &fake.EditorConfig{}
diff --git a/gopls/internal/regtest/diagnostics_test.go b/gopls/internal/regtest/diagnostics_test.go
index 8aa45ca..04a2b0d 100644
--- a/gopls/internal/regtest/diagnostics_test.go
+++ b/gopls/internal/regtest/diagnostics_test.go
@@ -46,7 +46,7 @@
 			// Once we have gotten diagnostics for the change above, we should
 			// satisfy the DiagnosticAtRegexp assertion.
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1),
+				env.DoneWithChange(),
 				env.DiagnosticAtRegexp("main.go", "Printl"),
 			),
 			// Assert that this test has sent no error logs to the client. This is not
@@ -225,7 +225,7 @@
 		env.Await(DiagnosticAt("a_test.go", 5, 3))
 		env.Sandbox.Workdir.RemoveFile(context.Background(), "a_test.go")
 		env.Await(OnceMet(
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+			env.DoneWithChangeWatchedFiles(),
 			DiagnosticAt("a_test.go", 5, 3)))
 	})
 }
@@ -375,7 +375,7 @@
 			// test to actually exercise the bug, we must wait until that work has
 			// completed.
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1),
+				env.DoneWithChange(),
 				NoDiagnostics("a.go"),
 			),
 		)
@@ -574,7 +574,7 @@
 				env.OpenFile("hello.txt")
 				env.Await(
 					OnceMet(
-						CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1),
+						env.DoneWithOpen(),
 						NoShowMessage(),
 					),
 				)
@@ -618,7 +618,7 @@
 		badFile := fmt.Sprintf("%s/found packages main (main.go) and x (x.go) in %s/src/x", dir, env.Sandbox.GOPATH())
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1),
+				env.DoneWithChange(),
 				EmptyDiagnostics("x/main.go"),
 			),
 			NoDiagnostics(badFile),
@@ -763,16 +763,13 @@
 		env.SaveBufferWithoutActions("a/a2.go")
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidSave), 1),
+				env.DoneWithSave(),
 				NoDiagnostics("a/a1.go"),
 			),
 		)
 		env.EditBuffer("a/a2.go", fake.NewEdit(0, 0, 0, 0, `package a`))
 		env.Await(
-			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1),
-				NoDiagnostics("a/a1.go"),
-			),
+			OnceMet(env.DoneWithChange(), NoDiagnostics("a/a1.go")),
 		)
 	})
 }
@@ -864,9 +861,7 @@
 	run(t, mod, func(t *testing.T, env *Env) {
 		env.OpenFile("foo/bar_test.go")
 		env.EditBuffer("foo/bar_test.go", fake.NewEdit(0, 0, 0, 0, "package foo"))
-		env.Await(
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1),
-		)
+		env.Await(env.DoneWithChange())
 		env.RegexpReplace("foo/bar_test.go", "package foo", `package foo_test
 
 import "testing"
@@ -900,11 +895,11 @@
 		env.SaveBuffer("foo/bar_test.go")
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidSave), 1),
+				env.DoneWithSave(),
 				NoDiagnostics("foo/bar_test.go"),
 			),
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidSave), 1),
+				env.DoneWithSave(),
 				NoDiagnostics("foo/foo.go"),
 			),
 		)
@@ -923,7 +918,7 @@
 		env.OpenFile("x_test.go")
 		env.EditBuffer("x_test.go", fake.NewEdit(0, 0, 0, 0, "pack"))
 		env.Await(
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1),
+			env.DoneWithChange(),
 			NoShowMessage(),
 		)
 	})
@@ -944,7 +939,7 @@
 		env.OpenFile("_foo/x.go")
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1),
+				env.DoneWithOpen(),
 				NoDiagnostics("_foo/x.go"),
 			))
 	})
@@ -980,7 +975,7 @@
 `
 	runner.Run(t, ws, func(t *testing.T, env *Env) {
 		env.OpenFile("b/b.go")
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1))
+		env.Await(env.DoneWithOpen())
 		// Delete c/c.go, the only file in package c.
 		env.RemoveWorkspaceFile("c/c.go")
 
@@ -1007,16 +1002,16 @@
 	// package loads.
 	writeGoVim := func(env *Env, name, content string) {
 		env.WriteWorkspaceFile(name, "")
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1))
+		env.Await(env.DoneWithChangeWatchedFiles())
 
 		env.CreateBuffer(name, "\n")
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1))
+		env.Await(env.DoneWithOpen())
 
 		env.EditBuffer(name, fake.NewEdit(1, 0, 1, 0, content))
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1))
+		env.Await(env.DoneWithChange())
 
 		env.EditBuffer(name, fake.NewEdit(0, 0, 1, 0, ""))
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1))
+		env.Await(env.DoneWithChange())
 	}
 
 	const p = `package p; func DoIt(s string) {};`
@@ -1131,7 +1126,7 @@
 		env.RegexpReplace("foo/foo_test.go", `package main`, `package foo`)
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1),
+				env.DoneWithChange(),
 				NoDiagnostics("foo/foo.go"),
 			),
 		)
@@ -1152,12 +1147,12 @@
 	runner.Run(t, basic, func(t *testing.T, env *Env) {
 		env.Editor.CreateBuffer(env.Ctx, "foo.go", `package main`)
 		env.Await(
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1),
+			env.DoneWithOpen(),
 		)
 		env.CloseBuffer("foo.go")
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidClose), 1),
+				env.DoneWithClose(),
 				NoLogMatching(protocol.Info, "packages=0"),
 			),
 		)
@@ -1176,29 +1171,23 @@
 `
 	runner.Run(t, basic, func(t *testing.T, env *Env) {
 		env.CreateBuffer("main.go", "")
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1))
+		env.Await(env.DoneWithOpen())
 
 		env.SaveBufferWithoutActions("main.go")
-		env.Await(
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidSave), 1),
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
-		)
+		env.Await(env.DoneWithSave(), env.DoneWithChangeWatchedFiles())
 
 		env.EditBuffer("main.go", fake.NewEdit(0, 0, 0, 0, `package main
 
 func main() {
 }
 `))
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1))
+		env.Await(env.DoneWithChange())
 
 		env.SaveBuffer("main.go")
-		env.Await(
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidSave), 2),
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 2),
-		)
+		env.Await(env.DoneWithSave(), env.DoneWithChangeWatchedFiles())
 
 		env.EditBuffer("main.go", fake.NewEdit(0, 0, 4, 0, ""))
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 2))
+		env.Await(env.DoneWithChange())
 
 		env.EditBuffer("main.go", fake.NewEdit(0, 0, 0, 0, `package main
 
@@ -1486,11 +1475,11 @@
 	).run(t, contents, func(t *testing.T, env *Env) {
 		// Simulate typing character by character.
 		env.OpenFile("foo/foo_test.go")
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1))
+		env.Await(env.DoneWithOpen())
 		env.RegexpReplace("foo/foo_test.go", "_", "_t")
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1))
+		env.Await(env.DoneWithChange())
 		env.RegexpReplace("foo/foo_test.go", "_t", "_test")
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 2))
+		env.Await(env.DoneWithChange())
 
 		env.Await(
 			EmptyDiagnostics("foo/foo_test.go"),
@@ -1694,7 +1683,7 @@
 			env.OpenFile("a/a.go")
 			env.Await(
 				OnceMet(
-					CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1),
+					env.DoneWithOpen(),
 					NoDiagnostics("a/a.go"),
 				),
 				NoOutstandingWork(),
@@ -1753,7 +1742,7 @@
 	).run(t, nested, func(t *testing.T, env *Env) {
 		// Expect a diagnostic in a nested module.
 		env.OpenFile("nested/hello/hello.go")
-		didOpen := CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1)
+		didOpen := env.DoneWithOpen()
 		env.Await(
 			OnceMet(
 				didOpen,
@@ -1783,7 +1772,7 @@
 		env.RegexpReplace("main.go", "{}", "{ var x int; }") // simulate typing
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), 1),
+				env.DoneWithChange(),
 				NoLogMatching(protocol.Info, "packages=1"),
 			),
 		)
diff --git a/gopls/internal/regtest/env.go b/gopls/internal/regtest/env.go
index 70859fd..322799d 100644
--- a/gopls/internal/regtest/env.go
+++ b/gopls/internal/regtest/env.go
@@ -54,7 +54,7 @@
 	// be string, though the spec allows for numeric tokens as well.  When work
 	// completes, it is deleted from this map.
 	outstandingWork map[protocol.ProgressToken]*workProgress
-	completedWork   map[string]int
+	completedWork   map[string]uint64
 }
 
 type workProgress struct {
@@ -119,7 +119,7 @@
 		state: State{
 			diagnostics:     make(map[string]*protocol.PublishDiagnosticsParams),
 			outstandingWork: make(map[protocol.ProgressToken]*workProgress),
-			completedWork:   make(map[string]int),
+			completedWork:   make(map[string]uint64),
 		},
 		waiters: make(map[int]*condition),
 	}
diff --git a/gopls/internal/regtest/env_test.go b/gopls/internal/regtest/env_test.go
index 82fb17f..e476be9 100644
--- a/gopls/internal/regtest/env_test.go
+++ b/gopls/internal/regtest/env_test.go
@@ -16,7 +16,7 @@
 	e := &Env{
 		state: State{
 			outstandingWork: make(map[protocol.ProgressToken]*workProgress),
-			completedWork:   make(map[string]int),
+			completedWork:   make(map[string]uint64),
 		},
 	}
 	ctx := context.Background()
diff --git a/gopls/internal/regtest/expectation.go b/gopls/internal/regtest/expectation.go
index 6c479c3..037faa4 100644
--- a/gopls/internal/regtest/expectation.go
+++ b/gopls/internal/regtest/expectation.go
@@ -199,6 +199,13 @@
 	return CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChange), changes)
 }
 
+// DoneWithSave expects all didSave notifications currently sent by the editor
+// to be completely processed.
+func (e *Env) DoneWithSave() Expectation {
+	saves := e.Editor.Stats().DidSave
+	return CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidSave), saves)
+}
+
 // DoneWithChangeWatchedFiles expects all didChangeWatchedFiles notifications
 // currently sent by the editor to be completely processed.
 func (e *Env) DoneWithChangeWatchedFiles() Expectation {
@@ -206,11 +213,18 @@
 	return CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), changes)
 }
 
+// DoneWithClose expects all didClose notifications currently sent by the
+// editor to be completely processed.
+func (e *Env) DoneWithClose() Expectation {
+	changes := e.Editor.Stats().DidClose
+	return CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidClose), changes)
+}
+
 // CompletedWork expects a work item to have been completed >= atLeast times.
 //
 // Since the Progress API doesn't include any hidden metadata, we must use the
 // progress notification title to identify the work we expect to be completed.
-func CompletedWork(title string, atLeast int) SimpleExpectation {
+func CompletedWork(title string, atLeast uint64) SimpleExpectation {
 	check := func(s State) Verdict {
 		if s.completedWork[title] >= atLeast {
 			return Met
diff --git a/gopls/internal/regtest/formatting_test.go b/gopls/internal/regtest/formatting_test.go
index 63fa0ce1..e6dba33 100644
--- a/gopls/internal/regtest/formatting_test.go
+++ b/gopls/internal/regtest/formatting_test.go
@@ -8,7 +8,6 @@
 	"strings"
 	"testing"
 
-	"golang.org/x/tools/internal/lsp"
 	"golang.org/x/tools/internal/lsp/tests"
 )
 
@@ -228,7 +227,7 @@
 			run(t, "-- main.go --", func(t *testing.T, env *Env) {
 				crlf := strings.ReplaceAll(tt.want, "\n", "\r\n")
 				env.CreateBuffer("main.go", crlf)
-				env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1))
+				env.Await(env.DoneWithOpen())
 				env.OrganizeImports("main.go")
 				got := env.Editor.BufferText("main.go")
 				got = strings.ReplaceAll(got, "\r\n", "\n") // convert everything to LF for simplicity
diff --git a/gopls/internal/regtest/generate_test.go b/gopls/internal/regtest/generate_test.go
index 87e64df..8e6751f 100644
--- a/gopls/internal/regtest/generate_test.go
+++ b/gopls/internal/regtest/generate_test.go
@@ -10,8 +10,6 @@
 
 import (
 	"testing"
-
-	"golang.org/x/tools/internal/lsp"
 )
 
 func TestGenerateProgress(t *testing.T) {
@@ -47,7 +45,7 @@
 		env.RunGenerate("./lib")
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+				env.DoneWithChangeWatchedFiles(),
 				EmptyDiagnostics("lib/lib.go")),
 		)
 	})
diff --git a/gopls/internal/regtest/modfile_test.go b/gopls/internal/regtest/modfile_test.go
index faef053..4cfc063 100644
--- a/gopls/internal/regtest/modfile_test.go
+++ b/gopls/internal/regtest/modfile_test.go
@@ -9,7 +9,6 @@
 	"strings"
 	"testing"
 
-	"golang.org/x/tools/internal/lsp"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/tests"
 	"golang.org/x/tools/internal/testenv"
@@ -92,9 +91,9 @@
 
 			env.RemoveWorkspaceFile("a/main.go")
 			env.Await(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1),
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidSave), 1),
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 2),
+				env.DoneWithOpen(),
+				env.DoneWithSave(),
+				env.DoneWithChangeWatchedFiles(),
 			)
 
 			env.WriteWorkspaceFile("main.go", mainContent)
@@ -828,7 +827,7 @@
 		Modes(Singleton),
 	).run(t, mod, func(t *testing.T, env *Env) {
 		env.OpenFile("go.mod")
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1))
+		env.Await(env.DoneWithOpen())
 		env.RegexpReplace("go.mod", "module", "modul")
 		// Confirm that we still have metadata with only on-disk edits.
 		env.OpenFile("main.go")
diff --git a/gopls/internal/regtest/vendor_test.go b/gopls/internal/regtest/vendor_test.go
index f9d43ee..dd318ff 100644
--- a/gopls/internal/regtest/vendor_test.go
+++ b/gopls/internal/regtest/vendor_test.go
@@ -7,7 +7,6 @@
 import (
 	"testing"
 
-	"golang.org/x/tools/internal/lsp"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/source"
 	"golang.org/x/tools/internal/testenv"
@@ -61,7 +60,7 @@
 			// so once we see the request, we can assume that `go mod vendor`
 			// will be executed.
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1),
+				env.DoneWithOpen(),
 				env.DiagnosticAtRegexp("go.mod", "module mod.com"),
 			),
 		)
diff --git a/gopls/internal/regtest/watch_test.go b/gopls/internal/regtest/watch_test.go
index 3802765..5b5321a 100644
--- a/gopls/internal/regtest/watch_test.go
+++ b/gopls/internal/regtest/watch_test.go
@@ -7,7 +7,6 @@
 import (
 	"testing"
 
-	"golang.org/x/tools/internal/lsp"
 	"golang.org/x/tools/internal/lsp/fake"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/testenv"
@@ -48,11 +47,11 @@
 			// Insert a trivial edit so that we don't automatically update the buffer
 			// (see CL 267577).
 			env.EditBuffer("a/a.go", fake.NewEdit(0, 0, 0, 0, " "))
-			env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1))
+			env.Await(env.DoneWithOpen())
 			env.WriteWorkspaceFile("a/a.go", `package a; func _() {};`)
 			env.Await(
 				OnceMet(
-					CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+					env.DoneWithChangeWatchedFiles(),
 					env.DiagnosticAtRegexp("a/a.go", "x"),
 				))
 		})
@@ -83,7 +82,7 @@
 `
 	runner.Run(t, pkg, func(t *testing.T, env *Env) {
 		env.OpenFile("a/a.go")
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1))
+		env.Await(env.DoneWithOpen())
 		env.WriteWorkspaceFile("b/b.go", `package b; func B() {};`)
 		env.Await(
 			env.DiagnosticAtRegexp("a/a.go", "b.B"),
@@ -160,7 +159,7 @@
 `
 	runner.Run(t, pkg, func(t *testing.T, env *Env) {
 		env.OpenFile("a/a.go")
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1))
+		env.Await(env.DoneWithOpen())
 		env.RemoveWorkspaceFile("b/b.go")
 		env.Await(
 			env.DiagnosticAtRegexp("a/a.go", "\"mod.com/b\""),
@@ -318,7 +317,7 @@
 			env.WriteWorkspaceFile("b/b.go", newMethod)
 			env.Await(
 				OnceMet(
-					CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+					env.DoneWithChangeWatchedFiles(),
 					DiagnosticAt("a/a.go", 12, 12),
 				),
 			)
@@ -334,7 +333,7 @@
 			env.WriteWorkspaceFile("a/a.go", implementation)
 			env.Await(
 				OnceMet(
-					CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+					env.DoneWithChangeWatchedFiles(),
 					NoDiagnostics("a/a.go"),
 				),
 			)
@@ -353,7 +352,7 @@
 			})
 			env.Await(
 				OnceMet(
-					CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+					env.DoneWithChangeWatchedFiles(),
 					NoDiagnostics("a/a.go"),
 				),
 				NoDiagnostics("b/b.go"),
@@ -387,7 +386,7 @@
 			env.OpenFile("a/a_unneeded.go")
 			env.Await(
 				OnceMet(
-					CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 2),
+					env.DoneWithOpen(),
 					LogMatching(protocol.Info, "a_unneeded.go", 1),
 				),
 			)
@@ -402,7 +401,7 @@
 			env.SaveBuffer("a/a.go")
 			env.Await(
 				OnceMet(
-					CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidSave), 1),
+					env.DoneWithSave(),
 					// There should only be one log message containing
 					// a_unneeded.go, from the initial workspace load, which we
 					// check for earlier. If there are more, there's a bug.
@@ -421,7 +420,7 @@
 			env.OpenFile("a/a_unneeded.go")
 			env.Await(
 				OnceMet(
-					CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 2),
+					env.DoneWithOpen(),
 					LogMatching(protocol.Info, "a_unneeded.go", 1),
 				),
 			)
@@ -436,7 +435,7 @@
 			env.SaveBuffer("a/a.go")
 			env.Await(
 				OnceMet(
-					CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidSave), 1),
+					env.DoneWithSave(),
 					// There should only be one log message containing
 					// a_unneeded.go, from the initial workspace load, which we
 					// check for earlier. If there are more, there's a bug.
@@ -510,7 +509,7 @@
 		})
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+				env.DoneWithChangeWatchedFiles(),
 				NoDiagnostics("main.go"),
 			),
 		)
@@ -585,7 +584,7 @@
 `,
 		})
 		env.Await(
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+			env.DoneWithChangeWatchedFiles(),
 			NoDiagnostics("main.go"),
 		)
 	})
@@ -620,7 +619,7 @@
 		}
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+				env.DoneWithChangeWatchedFiles(),
 				env.DiagnosticAtRegexp("foo/main.go", `"blah"`),
 			),
 		)
@@ -660,7 +659,7 @@
 		env.RemoveWorkspaceFile("foo/go.mod")
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+				env.DoneWithChangeWatchedFiles(),
 				env.DiagnosticAtRegexp("foo/main.go", `"mod.com/blah"`),
 			),
 		)
@@ -711,11 +710,11 @@
 		})
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+				env.DoneWithChangeWatchedFiles(),
 				NoDiagnostics("a/a.go"),
 			),
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+				env.DoneWithChangeWatchedFiles(),
 				NoDiagnostics("a/a_test.go"),
 			),
 		)
@@ -743,11 +742,11 @@
 		})
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 2),
+				env.DoneWithChangeWatchedFiles(),
 				NoDiagnostics("a/a_test.go"),
 			),
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 2),
+				env.DoneWithChangeWatchedFiles(),
 				NoDiagnostics("a/a2_test.go"),
 			),
 		)
diff --git a/gopls/internal/regtest/workspace_test.go b/gopls/internal/regtest/workspace_test.go
index f4c21f7..ae89b8d 100644
--- a/gopls/internal/regtest/workspace_test.go
+++ b/gopls/internal/regtest/workspace_test.go
@@ -12,7 +12,6 @@
 	"strings"
 	"testing"
 
-	"golang.org/x/tools/internal/lsp"
 	"golang.org/x/tools/internal/lsp/fake"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/testenv"
@@ -165,7 +164,7 @@
 `, env.ReadWorkspaceFile("pkg/go.mod"), dir)
 		env.WriteWorkspaceFile("pkg/go.mod", goModWithReplace)
 		env.Await(
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+			env.DoneWithChangeWatchedFiles(),
 			UnregistrationMatching("didChangeWatchedFiles"),
 			RegistrationMatching("didChangeWatchedFiles"),
 		)
@@ -277,7 +276,7 @@
 		env.RemoveWorkspaceFile("modb/b/b.go")
 		env.RemoveWorkspaceFile("modb/go.mod")
 		env.Await(
-			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 2),
+			env.DoneWithChangeWatchedFiles(),
 		)
 		if testenv.Go1Point() < 14 {
 			// On 1.14 and above, the go mod tidy diagnostics accidentally
@@ -340,7 +339,7 @@
 		})
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+				env.DoneWithChangeWatchedFiles(),
 				env.DiagnosticAtRegexp("modb/b/b.go", "x"),
 			),
 		)
@@ -388,7 +387,7 @@
 		env.OpenFile("modb/go.mod")
 		env.Await(
 			OnceMet(
-				CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidOpen), 1),
+				env.DoneWithOpen(),
 				DiagnosticAt("modb/go.mod", 0, 0),
 			),
 		)
@@ -659,7 +658,7 @@
 
 				replace a.com => %s/moda/a
 				`, workdir))
-		env.Await(CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1))
+		env.Await(env.DoneWithChangeWatchedFiles())
 		gotb, err = ioutil.ReadFile(modPath)
 		if err != nil {
 			t.Fatalf("reading expected workspace modfile: %v", err)
diff --git a/internal/lsp/fake/editor.go b/internal/lsp/fake/editor.go
index 758b64b..3df5b24 100644
--- a/internal/lsp/fake/editor.go
+++ b/internal/lsp/fake/editor.go
@@ -41,14 +41,17 @@
 	buffers map[string]buffer
 	// Capabilities / Options
 	serverCapabilities protocol.ServerCapabilities
+
 	// Call metrics for the purpose of expectations. This is done in an ad-hoc
 	// manner for now. Perhaps in the future we should do something more
-	// systematic.
-	calls CallCounts
+	// systematic. Guarded with a separate mutex as calls may need to be accessed
+	// asynchronously via callbacks into the Editor.
+	callsMu sync.Mutex
+	calls   CallCounts
 }
 
 type CallCounts struct {
-	DidOpen, DidChange, DidChangeWatchedFiles int
+	DidOpen, DidChange, DidSave, DidChangeWatchedFiles, DidClose uint64
 }
 
 type buffer struct {
@@ -140,8 +143,8 @@
 }
 
 func (e *Editor) Stats() CallCounts {
-	e.mu.Lock()
-	defer e.mu.Unlock()
+	e.callsMu.Lock()
+	defer e.callsMu.Unlock()
 	return e.calls
 }
 
@@ -291,36 +294,48 @@
 	return nil
 }
 
+// onFileChanges is registered to be called by the Workdir on any writes that
+// go through the Workdir API. It is called synchronously by the Workdir.
 func (e *Editor) onFileChanges(ctx context.Context, evts []FileEvent) {
 	if e.Server == nil {
 		return
 	}
-	e.mu.Lock()
-	defer e.mu.Unlock()
-	var lspevts []protocol.FileEvent
-	for _, evt := range evts {
-		// Always send an on-disk change, even for events that seem useless
-		// because they're shadowed by an open buffer.
-		lspevts = append(lspevts, evt.ProtocolEvent)
 
-		if buf, ok := e.buffers[evt.Path]; ok {
-			// Following VS Code, don't honor deletions or changes to dirty buffers.
-			if buf.dirty || evt.ProtocolEvent.Type == protocol.Deleted {
-				continue
-			}
-
-			content, err := e.sandbox.Workdir.ReadFile(evt.Path)
-			if err != nil {
-				continue // A race with some other operation.
-			}
-			// During shutdown, this call will fail. Ignore the error.
-			_ = e.setBufferContentLocked(ctx, evt.Path, false, strings.Split(content, "\n"), nil)
-		}
-	}
-	e.Server.DidChangeWatchedFiles(ctx, &protocol.DidChangeWatchedFilesParams{
-		Changes: lspevts,
-	})
+	// e may be locked when onFileChanges is called, but it is important that we
+	// synchronously increment this counter so that we can subsequently assert on
+	// the number of expected DidChangeWatchedFiles calls.
+	e.callsMu.Lock()
 	e.calls.DidChangeWatchedFiles++
+	e.callsMu.Unlock()
+
+	// Since e may be locked, we must run this mutation asynchronously.
+	go func() {
+		e.mu.Lock()
+		defer e.mu.Unlock()
+		var lspevts []protocol.FileEvent
+		for _, evt := range evts {
+			// Always send an on-disk change, even for events that seem useless
+			// because they're shadowed by an open buffer.
+			lspevts = append(lspevts, evt.ProtocolEvent)
+
+			if buf, ok := e.buffers[evt.Path]; ok {
+				// Following VS Code, don't honor deletions or changes to dirty buffers.
+				if buf.dirty || evt.ProtocolEvent.Type == protocol.Deleted {
+					continue
+				}
+
+				content, err := e.sandbox.Workdir.ReadFile(evt.Path)
+				if err != nil {
+					continue // A race with some other operation.
+				}
+				// During shutdown, this call will fail. Ignore the error.
+				_ = e.setBufferContentLocked(ctx, evt.Path, false, strings.Split(content, "\n"), nil)
+			}
+		}
+		e.Server.DidChangeWatchedFiles(ctx, &protocol.DidChangeWatchedFilesParams{
+			Changes: lspevts,
+		})
+	}()
 }
 
 // OpenFile creates a buffer for the given workdir-relative file.
@@ -371,7 +386,9 @@
 		}); err != nil {
 			return errors.Errorf("DidOpen: %w", err)
 		}
+		e.callsMu.Lock()
 		e.calls.DidOpen++
+		e.callsMu.Unlock()
 	}
 	return nil
 }
@@ -393,6 +410,9 @@
 		}); err != nil {
 			return errors.Errorf("DidClose: %w", err)
 		}
+		e.callsMu.Lock()
+		e.calls.DidClose++
+		e.callsMu.Unlock()
 	}
 	return nil
 }
@@ -458,6 +478,9 @@
 		if err := e.Server.DidSave(ctx, params); err != nil {
 			return errors.Errorf("DidSave: %w", err)
 		}
+		e.callsMu.Lock()
+		e.calls.DidSave++
+		e.callsMu.Unlock()
 	}
 	return nil
 }
@@ -653,7 +676,9 @@
 		if err := e.Server.DidChange(ctx, params); err != nil {
 			return errors.Errorf("DidChange: %w", err)
 		}
+		e.callsMu.Lock()
 		e.calls.DidChange++
+		e.callsMu.Unlock()
 	}
 	return nil
 }
diff --git a/internal/lsp/fake/workdir.go b/internal/lsp/fake/workdir.go
index 3cc6f73..5103bdb 100644
--- a/internal/lsp/fake/workdir.go
+++ b/internal/lsp/fake/workdir.go
@@ -211,7 +211,7 @@
 	copy(watchers, w.watchers)
 	w.watcherMu.Unlock()
 	for _, w := range watchers {
-		go w(ctx, evts)
+		w(ctx, evts)
 	}
 }
 
diff --git a/internal/lsp/fake/workdir_test.go b/internal/lsp/fake/workdir_test.go
index 5c9a36c..f57ea37 100644
--- a/internal/lsp/fake/workdir_test.go
+++ b/internal/lsp/fake/workdir_test.go
@@ -41,7 +41,9 @@
 
 	fileEvents := make(chan []FileEvent)
 	watch := func(_ context.Context, events []FileEvent) {
-		fileEvents <- events
+		go func() {
+			fileEvents <- events
+		}()
 	}
 	wd.AddWatcher(watch)
 	return wd, fileEvents, cleanup