gopls/internal/test: synchronize notifications during commands

It turns out that in the jsonrcp2 package, call responses are
asynchronous to other notifications. Therefore, we must synchronize
tests using progress notifications.

Introduce a DelayMessages test option to reproduce these types of races.
(It worked for reproducing golang/go#70342.)

Fixes golang/go#70342

Change-Id: I4cfcd7675335694a47eaf1a2547be0301fc244c9
Reviewed-on: https://go-review.googlesource.com/c/tools/+/627696
Reviewed-by: Alan Donovan <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: findleyr

gopls/internal/server/command.go

@@ -50,6 +50,15 @@ func (s *server) ExecuteCommand(ctx context.Context, params *protocol.ExecuteCom
 	ctx, done := event.Start(ctx, "lsp.Server.executeCommand")
 	defer done()
 
+	// For test synchronization, always create a progress notification.
+	//
+	// This may be in addition to user-facing progress notifications created in
+	// the course of command execution.
+	if s.Options().VerboseWorkDoneProgress {
+		work := s.progress.Start(ctx, params.Command, "Verbose: running command...", nil, nil)
+		defer work.End(ctx, "Done.")
+	}
+
 	var found bool
 	for _, name := range s.Options().SupportedCommands {
 		if name == params.Command {

gopls/internal/test/integration/fake/editor.go

@@ -10,13 +10,15 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math/rand/v2"
 	"os"
 	"path"
 	"path/filepath"
 	"regexp"
 	"slices"
 	"strings"
 	"sync"
+	"time"
 
 	"golang.org/x/tools/gopls/internal/protocol"
 	"golang.org/x/tools/gopls/internal/protocol/command"
@@ -136,6 +138,10 @@ type EditorConfig struct {
 	// If non-nil, MessageResponder is used to respond to ShowMessageRequest
 	// messages.
 	MessageResponder func(params *protocol.ShowMessageRequestParams) (*protocol.MessageActionItem, error)
+
+	// MaxMessageDelay is used for fuzzing message delivery to reproduce test
+	// flakes.
+	MaxMessageDelay time.Duration
 }
 
 // NewEditor creates a new Editor.
@@ -162,10 +168,11 @@ func (e *Editor) Connect(ctx context.Context, connector servertest.Connector, ho
 	e.serverConn = conn
 	e.Server = protocol.ServerDispatcher(conn)
 	e.client = &Client{editor: e, hooks: hooks}
-	conn.Go(bgCtx,
-		protocol.Handlers(
-			protocol.ClientHandler(e.client,
-				jsonrpc2.MethodNotFound)))
+	handler := protocol.ClientHandler(e.client, jsonrpc2.MethodNotFound)
+	if e.config.MaxMessageDelay > 0 {
+		handler = DelayedHandler(e.config.MaxMessageDelay, handler)
+	}
+	conn.Go(bgCtx, protocol.Handlers(handler))
 
 	if err := e.initialize(ctx); err != nil {
 		return nil, err
@@ -174,6 +181,18 @@ func (e *Editor) Connect(ctx context.Context, connector servertest.Connector, ho
 	return e, nil
 }
 
+// DelayedHandler waits [0, maxDelay) before handling each message.
+func DelayedHandler(maxDelay time.Duration, handler jsonrpc2.Handler) jsonrpc2.Handler {
+	return func(ctx context.Context, reply jsonrpc2.Replier, req jsonrpc2.Request) error {
+		delay := time.Duration(rand.Int64N(int64(maxDelay)))
+		select {
+		case <-ctx.Done():
+		case <-time.After(delay):
+		}
+		return handler(ctx, reply, req)
+	}
+}
+
 func (e *Editor) Stats() CallCounts {
 	e.callsMu.Lock()
 	defer e.callsMu.Unlock()

gopls/internal/test/integration/misc/webserver_test.go

@@ -135,6 +135,11 @@ const A = 1
 				collectMessages := env.Awaiter.ListenToShownMessages()
 				env.ExecuteCommand(params, &result)
 
+				// golang/go#70342: just because the command has finished does not mean
+				// that we will have received the necessary notifications. Synchronize
+				// using progress reports.
+				env.Await(CompletedWork(params.Command, 1, false))
+
 				wantDocs, wantMessages := 0, 1
 				if supported {
 					wantDocs, wantMessages = 1, 0
@@ -144,10 +149,10 @@ const A = 1
 				messages := collectMessages()
 
 				if gotDocs := len(docs); gotDocs != wantDocs {
-					t.Errorf("GoDoc: got %d showDocument requests, want %d", gotDocs, wantDocs)
+					t.Errorf("gopls.doc: got %d showDocument requests, want %d", gotDocs, wantDocs)
 				}
 				if gotMessages := len(messages); gotMessages != wantMessages {
-					t.Errorf("GoDoc: got %d showMessage requests, want %d", gotMessages, wantMessages)
+					t.Errorf("gopls.doc: got %d showMessage requests, want %d", gotMessages, wantMessages)
 				}
 			})
 		})

gopls/internal/test/integration/options.go

@@ -7,6 +7,7 @@ package integration
 import (
 	"strings"
 	"testing"
+	"time"
 
 	"golang.org/x/tools/gopls/internal/protocol"
 	"golang.org/x/tools/gopls/internal/test/integration/fake"
@@ -192,3 +193,14 @@ func MessageResponder(f func(*protocol.ShowMessageRequestParams) (*protocol.Mess
 		opts.editor.MessageResponder = f
 	})
 }
+
+// DelayMessages can be used to fuzz message delivery delays for the purpose of
+// reproducing test flakes.
+//
+// (Even though this option may be unused, keep it around to aid in debugging
+// future flakes.)
+func DelayMessages(upto time.Duration) RunOption {
+	return optionSetter(func(opts *runConfig) {
+		opts.editor.MaxMessageDelay = upto
+	})
+}