internal/lsp/lsprpc: add a forwarder handler

Add a forwarder handler that alters messages before forwarding, for now,
it just intercepts the "exit" message.

Also, make it easier to write regression tests for a shared gopls
instance, by adding a helper that instantiates two connected
environments, and only runs in the shared execution modes.

Updates golang/go#36879
Updates golang/go#34111

Change-Id: I7673f72ab71b5c7fd6ad65d274c15132a942e06a
Reviewed-on: https://go-review.googlesource.com/c/tools/+/218778
Run-TryBot: Robert Findley <rfindley@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>

internal/lsp/general.go

@@ -272,13 +272,17 @@ func (s *Server) shutdown(ctx context.Context) error {
 	return nil
 }
 
+// ServerExitFunc is used to exit when requested by the client. It is mutable
+// for testing purposes.
+var ServerExitFunc = os.Exit
+
 func (s *Server) exit(ctx context.Context) error {
 	s.stateMu.Lock()
 	defer s.stateMu.Unlock()
 	if s.state != serverShutDown {
-		os.Exit(1)
+		ServerExitFunc(1)
 	}
-	os.Exit(0)
+	ServerExitFunc(0)
 	return nil
 }
 

internal/lsp/lsprpc/lsprpc.go

@@ -10,6 +10,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"os"
 
 	"golang.org/x/sync/errgroup"
 	"golang.org/x/tools/internal/jsonrpc2"
@@ -93,6 +94,7 @@ func (f *Forwarder) ServeStream(ctx context.Context, stream jsonrpc2.Stream) err
 	serverConn.AddHandler(protocol.Canceller{})
 	clientConn.AddHandler(protocol.ServerHandler(server))
 	clientConn.AddHandler(protocol.Canceller{})
+	clientConn.AddHandler(forwarderHandler{})
 	if f.withTelemetry {
 		clientConn.AddHandler(telemetryHandler{})
 	}
@@ -106,3 +108,27 @@ func (f *Forwarder) ServeStream(ctx context.Context, stream jsonrpc2.Stream) err
 	})
 	return g.Wait()
 }
+
+// ForwarderExitFunc is used to exit the forwarder process. It is mutable for
+// testing purposes.
+var ForwarderExitFunc = os.Exit
+
+// forwarderHandler intercepts 'exit' messages to prevent the shared gopls
+// instance from exiting. In the future it may also intercept 'shutdown' to
+// provide more graceful shutdown of the client connection.
+type forwarderHandler struct {
+	jsonrpc2.EmptyHandler
+}
+
+func (forwarderHandler) Deliver(ctx context.Context, r *jsonrpc2.Request, delivered bool) bool {
+	// TODO(golang.org/issues/34111): we should more gracefully disconnect here,
+	// once that process exists.
+	if r.Method == "exit" {
+		ForwarderExitFunc(0)
+		// Still return true here to prevent the message from being delivered: in
+		// tests, ForwarderExitFunc may be overridden to something that doesn't
+		// exit the process.
+		return true
+	}
+	return false
+}

internal/lsp/regtest/diagnostics_test.go

@@ -37,28 +37,6 @@ func TestDiagnosticErrorInEditedFile(t *testing.T) {
 	})
 }
 
-func TestSimultaneousEdits(t *testing.T) {
-	t.Parallel()
-	runner.Run(t, exampleProgram, func(ctx context.Context, t *testing.T, env1 *Env) {
-		// Create a second test session connected to the same workspace and server
-		// as the first.
-		env2 := NewEnv(ctx, t, env1.W, env1.Server)
-
-		// In editor #1, break fmt.Println as before.
-		edit1 := fake.NewEdit(5, 11, 5, 12, "")
-		env1.OpenFile("main.go")
-		env1.EditBuffer("main.go", edit1)
-		// In editor #2 remove the closing brace.
-		edit2 := fake.NewEdit(6, 0, 6, 1, "")
-		env2.OpenFile("main.go")
-		env2.EditBuffer("main.go", edit2)
-
-		// Now check that we got different diagnostics in each environment.
-		env1.Await(DiagnosticAt("main.go", 5, 5))
-		env2.Await(DiagnosticAt("main.go", 7, 0))
-	})
-}
-
 const brokenFile = `package main
 
 const Foo = "abc

internal/lsp/regtest/env.go

@@ -40,9 +40,9 @@ const (
 // remote), any tests that execute on the same Runner will share the same
 // state.
 type Runner struct {
-	ts      *servertest.TCPServer
-	modes   EnvMode
-	timeout time.Duration
+	ts           *servertest.TCPServer
+	defaultModes EnvMode
+	timeout      time.Duration
 }
 
 // NewTestRunner creates a Runner with its shared state initialized, ready to
@@ -51,9 +51,9 @@ func NewTestRunner(modes EnvMode, testTimeout time.Duration) *Runner {
 	ss := lsprpc.NewStreamServer(cache.New(nil), false)
 	ts := servertest.NewTCPServer(context.Background(), ss)
 	return &Runner{
-		ts:      ts,
-		modes:   modes,
-		timeout: testTimeout,
+		ts:           ts,
+		defaultModes: modes,
+		timeout:      testTimeout,
 	}
 }
 
@@ -62,12 +62,17 @@ func (r *Runner) Close() error {
 	return r.ts.Close()
 }
 
-// Run executes the test function in in all configured gopls execution modes.
-// For each a test run, a new workspace is created containing the un-txtared
-// files specified by filedata.
+// Run executes the test function in the default configured gopls execution
+// modes. For each a test run, a new workspace is created containing the
+// un-txtared files specified by filedata.
 func (r *Runner) Run(t *testing.T, filedata string, test func(context.Context, *testing.T, *Env)) {
 	t.Helper()
+	r.RunInMode(r.defaultModes, t, filedata, test)
+}
 
+// RunInMode runs the test in the execution modes specified by the modes bitmask.
+func (r *Runner) RunInMode(modes EnvMode, t *testing.T, filedata string, test func(ctx context.Context, t *testing.T, e *Env)) {
+	t.Helper()
 	tests := []struct {
 		name         string
 		mode         EnvMode
@@ -80,7 +85,7 @@ func (r *Runner) Run(t *testing.T, filedata string, test func(context.Context, *
 
 	for _, tc := range tests {
 		tc := tc
-		if r.modes&tc.mode == 0 {
+		if modes&tc.mode == 0 {
 			continue
 		}
 		t.Run(tc.name, func(t *testing.T) {
@@ -231,10 +236,17 @@ func (e *Env) onDiagnostics(_ context.Context, d *protocol.PublishDiagnosticsPar
 			close(condition.met)
 		}
 	}
-
 	return nil
 }
 
+// CloseEditor shuts down the editor, calling t.Fatal on any error.
+func (e *Env) CloseEditor() {
+	e.t.Helper()
+	if err := e.E.ShutdownAndExit(e.ctx); err != nil {
+		e.t.Fatal(err)
+	}
+}
+
 func meetsCondition(m map[string]*protocol.PublishDiagnosticsParams, expectations []DiagnosticExpectation) bool {
 	for _, e := range expectations {
 		if !e.IsMet(m) {

internal/lsp/regtest/reg_test.go

@@ -5,14 +5,29 @@
 package regtest
 
 import (
+	"fmt"
 	"os"
 	"testing"
 	"time"
+
+	"golang.org/x/tools/internal/lsp"
+	"golang.org/x/tools/internal/lsp/lsprpc"
 )
 
 var runner *Runner
 
 func TestMain(m *testing.M) {
+	// Override functions that would shut down the test process
+	defer func(lspExit, forwarderExit func(code int)) {
+		lsp.ServerExitFunc = lspExit
+		lsprpc.ForwarderExitFunc = forwarderExit
+	}(lsp.ServerExitFunc, lsprpc.ForwarderExitFunc)
+	// None of these regtests should be able to shut down a server process.
+	lsp.ServerExitFunc = func(code int) {
+		panic(fmt.Sprintf("LSP server exited with code %d", code))
+	}
+	// We don't want our forwarders to exit, but it's OK if they would have.
+	lsprpc.ForwarderExitFunc = func(code int) {}
 	runner = NewTestRunner(AllModes, 30*time.Second)
 	defer runner.Close()
 	os.Exit(m.Run())

internal/lsp/regtest/shared_test.go

@@ -0,0 +1,69 @@
+// Copyright 2020 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 regtest
+
+import (
+	"context"
+	"testing"
+
+	"golang.org/x/tools/internal/lsp/fake"
+)
+
+const sharedProgram = `
+-- go.mod --
+module mod
+
+go 1.12
+-- main.go --
+package main
+
+import "fmt"
+
+func main() {
+	fmt.Println("Hello World.")
+}`
+
+func runShared(t *testing.T, program string, testFunc func(ctx context.Context, t *testing.T, env1 *Env, env2 *Env)) {
+	runner.RunInMode(Forwarded, t, sharedProgram, func(ctx context.Context, t *testing.T, env1 *Env) {
+		// Create a second test session connected to the same workspace and server
+		// as the first.
+		env2 := NewEnv(ctx, t, env1.W, env1.Server)
+		testFunc(ctx, t, env1, env2)
+	})
+}
+
+func TestSimultaneousEdits(t *testing.T) {
+	t.Parallel()
+	runner.Run(t, exampleProgram, func(ctx context.Context, t *testing.T, env1 *Env) {
+		// Create a second test session connected to the same workspace and server
+		// as the first.
+		env2 := NewEnv(ctx, t, env1.W, env1.Server)
+
+		// In editor #1, break fmt.Println as before.
+		edit1 := fake.NewEdit(5, 11, 5, 12, "")
+		env1.OpenFile("main.go")
+		env1.EditBuffer("main.go", edit1)
+		// In editor #2 remove the closing brace.
+		edit2 := fake.NewEdit(6, 0, 6, 1, "")
+		env2.OpenFile("main.go")
+		env2.EditBuffer("main.go", edit2)
+
+		// Now check that we got different diagnostics in each environment.
+		env1.Await(DiagnosticAt("main.go", 5, 5))
+		env2.Await(DiagnosticAt("main.go", 7, 0))
+	})
+}
+
+func TestShutdown(t *testing.T) {
+	t.Parallel()
+	runShared(t, sharedProgram, func(ctx context.Context, t *testing.T, env1 *Env, env2 *Env) {
+		env1.CloseEditor()
+		// Now make an edit in editor #2 to trigger diagnostics.
+		edit2 := fake.NewEdit(6, 0, 6, 1, "")
+		env2.OpenFile("main.go")
+		env2.EditBuffer("main.go", edit2)
+		env2.Await(DiagnosticAt("main.go", 7, 0))
+	})
+}