internal/lsp/regtest: don't run the connection on the test context

When test awaiting fails, we often fail to shut down the server because
the pipe is closed. Fix this by using a detached context for running the
connection.

Also clean up some unnecessary context arguments.

Change-Id: I535c1cc1606e44df5f8e2177c92293d57836f992
Reviewed-on: https://go-review.googlesource.com/c/tools/+/340850
TryBot-Result: Gopher Robot <gobot@golang.org>
Run-TryBot: Robert Findley <rfindley@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
diff --git a/gopls/internal/regtest/misc/shared_test.go b/gopls/internal/regtest/misc/shared_test.go
index 6861743..a6b0cd8 100644
--- a/gopls/internal/regtest/misc/shared_test.go
+++ b/gopls/internal/regtest/misc/shared_test.go
@@ -30,7 +30,8 @@
 	WithOptions(Modes(modes)).Run(t, sharedProgram, func(t *testing.T, env1 *Env) {
 		// Create a second test session connected to the same workspace and server
 		// as the first.
-		env2 := NewEnv(env1.Ctx, t, env1.Sandbox, env1.Server, env1.Editor.Config, true)
+		env2, cleanup := NewEnv(env1.Ctx, t, env1.Sandbox, env1.Server, env1.Editor.Config, true)
+		defer cleanup()
 		env2.Await(InitialWorkspaceLoad)
 		testFunc(env1, env2)
 	})
diff --git a/internal/jsonrpc2/servertest/servertest.go b/internal/jsonrpc2/servertest/servertest.go
index 392e084..b879ebd 100644
--- a/internal/jsonrpc2/servertest/servertest.go
+++ b/internal/jsonrpc2/servertest/servertest.go
@@ -68,7 +68,7 @@
 }
 
 // NewPipeServer returns a test server that can be connected to via io.Pipes.
-func NewPipeServer(ctx context.Context, server jsonrpc2.StreamServer, framer jsonrpc2.Framer) *PipeServer {
+func NewPipeServer(server jsonrpc2.StreamServer, framer jsonrpc2.Framer) *PipeServer {
 	if framer == nil {
 		framer = jsonrpc2.NewRawStream
 	}
diff --git a/internal/jsonrpc2/servertest/servertest_test.go b/internal/jsonrpc2/servertest/servertest_test.go
index 38fa21a..1780d4f 100644
--- a/internal/jsonrpc2/servertest/servertest_test.go
+++ b/internal/jsonrpc2/servertest/servertest_test.go
@@ -26,7 +26,7 @@
 	server := jsonrpc2.HandlerServer(fakeHandler)
 	tcpTS := NewTCPServer(ctx, server, nil)
 	defer tcpTS.Close()
-	pipeTS := NewPipeServer(ctx, server, nil)
+	pipeTS := NewPipeServer(server, nil)
 	defer pipeTS.Close()
 
 	tests := []struct {
diff --git a/internal/lsp/lsprpc/lsprpc_test.go b/internal/lsp/lsprpc/lsprpc_test.go
index 795c887..cde641c 100644
--- a/internal/lsp/lsprpc/lsprpc_test.go
+++ b/internal/lsp/lsprpc/lsprpc_test.go
@@ -60,7 +60,7 @@
 	ctx = debug.WithInstance(ctx, "", "")
 	ss := NewStreamServer(cache.New(nil), false)
 	ss.serverForTest = server
-	ts := servertest.NewPipeServer(ctx, ss, nil)
+	ts := servertest.NewPipeServer(ss, nil)
 	defer checkClose(t, ts.Close)
 	cc := ts.Connect(ctx)
 	cc.Go(ctx, protocol.ClientHandler(client, jsonrpc2.MethodNotFound))
@@ -125,12 +125,11 @@
 	ss.serverForTest = s
 	tsDirect := servertest.NewTCPServer(serveCtx, ss, nil)
 
-	forwarderCtx := debug.WithInstance(ctx, "", "")
 	forwarder, err := NewForwarder("tcp;"+tsDirect.Addr, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
-	tsForwarded := servertest.NewPipeServer(forwarderCtx, forwarder, nil)
+	tsForwarded := servertest.NewPipeServer(forwarder, nil)
 	return tsDirect, tsForwarded, func() {
 		checkClose(t, tsDirect.Close)
 		checkClose(t, tsForwarded.Close)
@@ -225,7 +224,7 @@
 	if err != nil {
 		t.Fatal(err)
 	}
-	tsForwarder := servertest.NewPipeServer(clientCtx, forwarder, nil)
+	tsForwarder := servertest.NewPipeServer(forwarder, nil)
 
 	conn1 := tsForwarder.Connect(clientCtx)
 	ed1, err := fake.NewEditor(sb, fake.EditorConfig{}).Connect(clientCtx, conn1, fake.ClientHooks{})
diff --git a/internal/lsp/regtest/env.go b/internal/lsp/regtest/env.go
index f095c38..a37cbf6 100644
--- a/internal/lsp/regtest/env.go
+++ b/internal/lsp/regtest/env.go
@@ -14,6 +14,7 @@
 	"golang.org/x/tools/internal/jsonrpc2/servertest"
 	"golang.org/x/tools/internal/lsp/fake"
 	"golang.org/x/tools/internal/lsp/protocol"
+	"golang.org/x/tools/internal/xcontext"
 )
 
 // Env holds an initialized fake Editor, Workspace, and Server, which may be
@@ -109,9 +110,14 @@
 
 // NewEnv creates a new test environment using the given scratch environment
 // and gopls server.
-func NewEnv(ctx context.Context, tb testing.TB, sandbox *fake.Sandbox, ts servertest.Connector, editorConfig fake.EditorConfig, withHooks bool) *Env {
+//
+// The resulting func must be called to close the jsonrpc2 connection.
+func NewEnv(ctx context.Context, tb testing.TB, sandbox *fake.Sandbox, ts servertest.Connector, editorConfig fake.EditorConfig, withHooks bool) (_ *Env, cleanup func()) {
 	tb.Helper()
-	conn := ts.Connect(ctx)
+
+	bgCtx, cleanupConn := context.WithCancel(xcontext.Detach(ctx))
+	conn := ts.Connect(bgCtx)
+
 	env := &Env{
 		T:       tb,
 		Ctx:     ctx,
@@ -138,12 +144,12 @@
 			OnUnregistration:         env.onUnregistration,
 		}
 	}
-	editor, err := fake.NewEditor(sandbox, editorConfig).Connect(ctx, conn, hooks)
+	editor, err := fake.NewEditor(sandbox, editorConfig).Connect(bgCtx, conn, hooks)
 	if err != nil {
 		tb.Fatal(err)
 	}
 	env.Editor = editor
-	return env
+	return env, cleanupConn
 }
 
 func (e *Env) onDiagnostics(_ context.Context, d *protocol.PublishDiagnosticsParams) error {
diff --git a/internal/lsp/regtest/runner.go b/internal/lsp/regtest/runner.go
index 3cfeb77..bebec53 100644
--- a/internal/lsp/regtest/runner.go
+++ b/internal/lsp/regtest/runner.go
@@ -234,7 +234,7 @@
 	tests := []struct {
 		name      string
 		mode      Mode
-		getServer func(context.Context, *testing.T, func(*source.Options)) jsonrpc2.StreamServer
+		getServer func(*testing.T, func(*source.Options)) jsonrpc2.StreamServer
 	}{
 		{"singleton", Singleton, singletonServer},
 		{"forwarded", Forwarded, r.forwardedServer},
@@ -301,14 +301,15 @@
 			// better solution to ensure that all Go processes started by gopls have
 			// exited before we clean up.
 			r.AddCloser(sandbox)
-			ss := tc.getServer(ctx, t, config.optionsHook)
+			ss := tc.getServer(t, config.optionsHook)
 			framer := jsonrpc2.NewRawStream
 			ls := &loggingFramer{}
 			if !config.skipLogs {
 				framer = ls.framer(jsonrpc2.NewRawStream)
 			}
-			ts := servertest.NewPipeServer(ctx, ss, framer)
-			env := NewEnv(ctx, t, sandbox, ts, config.editor, !config.skipHooks)
+			ts := servertest.NewPipeServer(ss, framer)
+			env, cleanup := NewEnv(ctx, t, sandbox, ts, config.editor, !config.skipHooks)
+			defer cleanup()
 			defer func() {
 				if t.Failed() && r.PrintGoroutinesOnFailure {
 					pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
@@ -406,11 +407,11 @@
 	fmt.Fprintf(os.Stderr, "#### End Gopls Test Logs for %q\n", testname)
 }
 
-func singletonServer(ctx context.Context, t *testing.T, optsHook func(*source.Options)) jsonrpc2.StreamServer {
+func singletonServer(t *testing.T, optsHook func(*source.Options)) jsonrpc2.StreamServer {
 	return lsprpc.NewStreamServer(cache.New(optsHook), false)
 }
 
-func experimentalServer(_ context.Context, t *testing.T, optsHook func(*source.Options)) jsonrpc2.StreamServer {
+func experimentalServer(t *testing.T, optsHook func(*source.Options)) jsonrpc2.StreamServer {
 	options := func(o *source.Options) {
 		optsHook(o)
 		o.EnableAllExperiments()
@@ -421,7 +422,7 @@
 	return lsprpc.NewStreamServer(cache.New(options), false)
 }
 
-func (r *Runner) forwardedServer(ctx context.Context, t *testing.T, optsHook func(*source.Options)) jsonrpc2.StreamServer {
+func (r *Runner) forwardedServer(t *testing.T, optsHook func(*source.Options)) jsonrpc2.StreamServer {
 	ts := r.getTestServer(optsHook)
 	return newForwarder("tcp", ts.Addr)
 }
@@ -440,7 +441,7 @@
 	return r.ts
 }
 
-func (r *Runner) separateProcessServer(ctx context.Context, t *testing.T, optsHook func(*source.Options)) jsonrpc2.StreamServer {
+func (r *Runner) separateProcessServer(t *testing.T, optsHook func(*source.Options)) jsonrpc2.StreamServer {
 	// TODO(rfindley): can we use the autostart behavior here, instead of
 	// pre-starting the remote?
 	socket := r.getRemoteSocket(t)