internal/lsp/lsprpc: clean up client session on disconnection

Gopls behavior on disconnection is currently somewhat undefined, because
it hasn't mattered when there was a single gopls session per binary
invocation. With golang/go#34111, this changes.

Checks are added to ensure clients and sessions are cleaned up when an LSP
connection closes. Also, normal client disconnection is differentiated
with the jsonrpc2.ErrDisconnected value.

Updates golang/go#34111

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

5 files changed, 36 insertions, 4 deletions

diff --git a/internal/jsonrpc2/jsonrpc2.go b/internal/jsonrpc2/jsonrpc2.go
index 6cbd60762..39ac3eba7 100644
--- a/internal/jsonrpc2/jsonrpc2.go
+++ b/internal/jsonrpc2/jsonrpc2.go
@@ -316,7 +316,8 @@ func (c *Conn) Run(runCtx context.Context) error {
 		// get the data for a message
 		data, n, err := c.stream.Read(runCtx)
 		if err != nil {
-			// the stream failed, we cannot continue
+			// The stream failed, we cannot continue. If the client disconnected
+			// normally, we should get ErrDisconnected here.
 			return err
 		}
 		// read a combined message
diff --git a/internal/jsonrpc2/stream.go b/internal/jsonrpc2/stream.go
index f850c27c7..2d1e7c440 100644
--- a/internal/jsonrpc2/stream.go
+++ b/internal/jsonrpc2/stream.go
@@ -8,6 +8,7 @@ import (
 	"bufio"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"strconv"
@@ -28,6 +29,9 @@ type Stream interface {
 	Write(context.Context, []byte) (int64, error)
 }
 
+// ErrDisconnected signals that the stream or connection exited normally.
+var ErrDisconnected = errors.New("disconnected")
+
 // NewStream returns a Stream built on top of an io.Reader and io.Writer
 // The messages are sent with no wrapping, and rely on json decode consistency
 // to determine message boundaries.
@@ -52,6 +56,9 @@ func (s *plainStream) Read(ctx context.Context) ([]byte, int64, error) {
 	}
 	var raw json.RawMessage
 	if err := s.in.Decode(&raw); err != nil {
+		if err == io.EOF {
+			return nil, 0, ErrDisconnected
+		}
 		return nil, 0, err
 	}
 	return raw, int64(len(raw)), nil
@@ -96,6 +103,10 @@ func (s *headerStream) Read(ctx context.Context) ([]byte, int64, error) {
 	for {
 		line, err := s.in.ReadString('\n')
 		total += int64(len(line))
+		if err == io.EOF {
+			// A normal disconnection will terminate with EOF before the next header.
+			return nil, total, ErrDisconnected
+		}
 		if err != nil {
 			return nil, total, fmt.Errorf("failed reading header line %q", err)
 		}
diff --git a/internal/lsp/general.go b/internal/lsp/general.go
index 95ffeeaf0..9ea2f01d6 100644
--- a/internal/lsp/general.go
+++ b/internal/lsp/general.go
@@ -266,9 +266,11 @@ func (s *Server) shutdown(ctx context.Context) error {
 	if s.state < serverInitialized {
 		return jsonrpc2.NewErrorf(jsonrpc2.CodeInvalidRequest, "server not initialized")
 	}
-	// drop all the active views
-	s.session.Shutdown(ctx)
-	s.state = serverShutDown
+	if s.state != serverShutDown {
+		// drop all the active views
+		s.session.Shutdown(ctx)
+		s.state = serverShutDown
+	}
 	return nil
 }
 
diff --git a/internal/lsp/lsprpc/lsprpc.go b/internal/lsp/lsprpc/lsprpc.go
index 0bf170c32..d37f4fd1f 100644
--- a/internal/lsp/lsprpc/lsprpc.go
+++ b/internal/lsp/lsprpc/lsprpc.go
@@ -126,10 +126,16 @@ func (s *StreamServer) ServeStream(ctx context.Context, stream jsonrpc2.Stream)
 		session: session,
 	}
 	s.debug.State.AddClient(dc)
+	defer s.debug.State.DropClient(dc)
+
 	server := s.serverForTest
 	if server == nil {
 		server = lsp.NewServer(session, client)
 	}
+	// Clients may or may not send a shutdown message. Make sure the server is
+	// shut down.
+	// TODO(rFindley): this shutdown should perhaps be on a disconnected context.
+	defer server.Shutdown(ctx)
 	conn.AddHandler(protocol.ServerHandler(server))
 	conn.AddHandler(protocol.Canceller{})
 	if s.withTelemetry {
@@ -174,6 +180,7 @@ func NewForwarder(network, addr string, withTelemetry bool, debugInstance *debug
 		stdlog.Printf("error getting gopls path for forwarder: %v", err)
 		gp = ""
 	}
+
 	return &Forwarder{
 		network:       network,
 		addr:          addr,
diff --git a/internal/lsp/lsprpc/lsprpc_test.go b/internal/lsp/lsprpc/lsprpc_test.go
index 215d936b5..1bca64093 100644
--- a/internal/lsp/lsprpc/lsprpc_test.go
+++ b/internal/lsp/lsprpc/lsprpc_test.go
@@ -37,6 +37,10 @@ func (s pingServer) DidOpen(ctx context.Context, params *protocol.DidOpenTextDoc
 	return nil
 }
 
+func (s pingServer) Shutdown(ctx context.Context) error {
+	return nil
+}
+
 func TestClientLogging(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -48,6 +52,7 @@ func TestClientLogging(t *testing.T) {
 	ss := NewStreamServer(cache.New(nil, di.State), false, di)
 	ss.serverForTest = server
 	ts := servertest.NewPipeServer(ctx, ss)
+	defer ts.Close()
 	cc := ts.Connect(ctx)
 	cc.AddHandler(protocol.ClientHandler(client))
 
@@ -91,6 +96,10 @@ func (s waitableServer) Resolve(_ context.Context, item *protocol.CompletionItem
 	return item, nil
 }
 
+func (s waitableServer) Shutdown(ctx context.Context) error {
+	return nil
+}
+
 func TestRequestCancellation(t *testing.T) {
 	server := waitableServer{
 		started: make(chan struct{}),
@@ -100,9 +109,11 @@ func TestRequestCancellation(t *testing.T) {
 	ss.serverForTest = server
 	ctx := context.Background()
 	tsDirect := servertest.NewTCPServer(ctx, ss)
+	defer tsDirect.Close()
 
 	forwarder := NewForwarder("tcp", tsDirect.Addr, false, debug.NewInstance("", ""))
 	tsForwarded := servertest.NewPipeServer(ctx, forwarder)
+	defer tsForwarded.Close()
 
 	tests := []struct {
 		serverType string