gopls: fix raciness related to TestOrphanedFiles

While investigating the high rate of flakiness in TestOrphanedFiles, I
found two problems:

- The change to use Await(...) rather than AfterChange(..), which was
  intended to *avoid* flakes due to asynchronous log messages, actually
  increased flakiness. After that change, the test can sometimes proceed
  to make additional edits before change processing completes, so that
  the load succeeds but the snapshot is cloned before the resulting
  state can be written (sigh...).
- If the load succeeds, we should proceed to update state even if the
  context is cancelled. Writing the state is not expensive.

Fixes golang/go#61521

Change-Id: I4e70526a0108013c66d378da966289a7c2f5dbe2
Reviewed-on: https://go-review.googlesource.com/c/tools/+/518975
TryBot-Result: Gopher Robot <[email protected]>
Run-TryBot: Robert Findley <[email protected]>
Reviewed-by: Alan Donovan <[email protected]>
gopls-CI: kokoro <[email protected]>

Author: findleyr

gopls/internal/lsp/cache/snapshot.go

@@ -755,12 +755,22 @@ func (s *snapshot) MetadataForFile(ctx context.Context, uri span.URI) ([]*source
 		scope := fileLoadScope(uri)
 		err := s.load(ctx, false, scope)
 
-		// Guard against failed loads due to context cancellation.
 		//
 		// Return the context error here as the current operation is no longer
 		// valid.
-		if ctxErr := ctx.Err(); ctxErr != nil {
-			return nil, ctxErr
+		if err != nil {
+			// Guard against failed loads due to context cancellation. We don't want
+			// to mark loads as completed if they failed due to context cancellation.
+			if ctx.Err() != nil {
+				return nil, ctx.Err()
+			}
+
+			// Don't return an error here, as we may still return stale IDs.
+			// Furthermore, the result of MetadataForFile should be consistent upon
+			// subsequent calls, even if the file is marked as unloadable.
+			if !errors.Is(err, errNoPackages) {
+				event.Error(ctx, "MetadataForFile", err)
+			}
 		}
 
 		// We must clear scopes after loading.
@@ -769,13 +779,6 @@ func (s *snapshot) MetadataForFile(ctx context.Context, uri span.URI) ([]*source
 		// packages as loaded. We could do this from snapshot.load and avoid
 		// raciness.
 		s.clearShouldLoad(scope)
-
-		// Don't return an error here, as we may still return stale IDs.
-		// Furthermore, the result of MetadataForFile should be consistent upon
-		// subsequent calls, even if the file is marked as unloadable.
-		if err != nil && !errors.Is(err, errNoPackages) {
-			event.Error(ctx, "MetadataForFile", err)
-		}
 	}
 
 	// Retrieve the metadata.

gopls/internal/regtest/diagnostics/diagnostics_test.go

@@ -1307,7 +1307,14 @@ func _() {
 		env.OpenFile("a/a_exclude.go")
 
 		loadOnce := LogMatching(protocol.Info, "query=.*file=.*a_exclude.go", 1, false)
-		env.Await(loadOnce) // can't use OnceMet or AfterChange as logs are async
+
+		// can't use OnceMet or AfterChange as logs are async
+		env.Await(loadOnce)
+		// ...but ensure that the change has been fully processed before editing.
+		// Otherwise, there may be a race where the snapshot is cloned before all
+		// state changes resulting from the load have been processed
+		// (golang/go#61521).
+		env.AfterChange()
 
 		// Check that orphaned files are not reloaded, by making a change in
 		// a.go file and confirming that the workspace diagnosis did not reload