internal/lsp: making cmd tests fast and small

This required extensive changes to the "internal" server handling and also to
the way we wait for diagnostics in the "check" verb.
It improves both memory and time by over an order of magnitude, hopefully
allowing us to renable the tests on the builders

Change-Id: I84e84ca4c449e9970ebf1d922a0a2ce0a8a49c72
Reviewed-on: https://go-review.googlesource.com/c/tools/+/175878
Run-TryBot: Ian Cottrell <iancottrell@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

internal/lsp/cmd/check.go

@@ -8,8 +8,8 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"time"
 
-	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/span"
 )
 
@@ -18,16 +18,6 @@ type check struct {
 	app *Application
 }
 
-type checkClient struct {
-	baseClient
-	diagnostics chan entry
-}
-
-type entry struct {
-	uri         span.URI
-	diagnostics []protocol.Diagnostic
-}
-
 func (c *check) Name() string      { return "check" }
 func (c *check) Usage() string     { return "<filename>" }
 func (c *check) ShortHelp() string { return "show diagnostic results for the specified file" }
@@ -49,49 +39,35 @@ func (c *check) Run(ctx context.Context, args ...string) error {
 		// no files, so no results
 		return nil
 	}
-	client := &checkClient{
-		diagnostics: make(chan entry),
-	}
-	checking := map[span.URI]*protocol.ColumnMapper{}
+	checking := map[span.URI]*cmdFile{}
 	// now we ready to kick things off
-	_, err := c.app.connect(ctx, client)
+	conn, err := c.app.connect(ctx)
 	if err != nil {
 		return err
 	}
 	for _, arg := range args {
 		uri := span.FileURI(arg)
-		m, err := client.AddFile(ctx, uri)
-		if err != nil {
-			return err
+		file := conn.AddFile(ctx, uri)
+		if file.err != nil {
+			return file.err
 		}
-		checking[uri] = m
+		checking[uri] = file
 	}
 	// now wait for results
-	for entry := range client.diagnostics {
-		//TODO:timeout?
-		m, found := checking[entry.uri]
-		if !found {
-			continue
+	//TODO: maybe conn.ExecuteCommand(ctx, &protocol.ExecuteCommandParams{Command: "gopls-wait-idle"})
+	for _, file := range checking {
+		select {
+		case <-file.hasDiagnostics:
+		case <-time.Tick(30 * time.Second):
+			return fmt.Errorf("timed out waiting for results from %v", file.uri)
 		}
-		for _, d := range entry.diagnostics {
-			spn, err := m.RangeSpan(d.Range)
+		for _, d := range file.diagnostics {
+			spn, err := file.mapper.RangeSpan(d.Range)
 			if err != nil {
 				return fmt.Errorf("Could not convert position %v for %q", d.Range, d.Message)
 			}
 			fmt.Printf("%v: %v\n", spn, d.Message)
 		}
-		delete(checking, entry.uri)
-		if len(checking) == 0 {
-			return nil
-		}
-	}
-	return fmt.Errorf("did not get all results")
-}
-
-func (c *checkClient) PublishDiagnostics(ctx context.Context, p *protocol.PublishDiagnosticsParams) error {
-	c.diagnostics <- entry{
-		uri:         span.URI(p.URI),
-		diagnostics: p.Diagnostics,
 	}
 	return nil
 }

internal/lsp/cmd/check_test.go

@@ -7,7 +7,6 @@ package cmd_test
 import (
 	"context"
 	"fmt"
-	"runtime"
 	"strings"
 	"testing"
 
@@ -17,9 +16,6 @@ import (
 )
 
 func (r *runner) Diagnostics(t *testing.T, data tests.Diagnostics) {
-	if runtime.GOOS != "linux" || isRace {
-		t.Skip("currently uses too much memory, see issue #31611")
-	}
 	for uri, want := range data {
 		if len(want) == 1 && want[0].Message == "" {
 			continue
@@ -28,8 +24,7 @@ func (r *runner) Diagnostics(t *testing.T, data tests.Diagnostics) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		args := []string{"-remote=internal"}
-		args = append(args, "check", fname)
+		args := []string{"-remote=internal", "check", fname}
 		out := captureStdOut(t, func() {
 			tool.Main(context.Background(), r.app, args)
 		})

internal/lsp/cmd/cmd.go

@@ -19,6 +19,7 @@ import (
 	"net"
 	"os"
 	"strings"
+	"sync"
 
 	"golang.org/x/tools/go/packages"
 	"golang.org/x/tools/internal/jsonrpc2"
@@ -120,62 +121,104 @@ func (app *Application) commands() []tool.Application {
 	}
 }
 
-type cmdClient interface {
-	protocol.Client
+var (
+	internalMu          sync.Mutex
+	internalConnections = make(map[string]*connection)
+)
 
-	prepare(app *Application, server protocol.Server)
-}
-
-func (app *Application) connect(ctx context.Context, client cmdClient) (protocol.Server, error) {
-	var server protocol.Server
+func (app *Application) connect(ctx context.Context) (*connection, error) {
 	switch app.Remote {
 	case "":
-		server = lsp.NewClientServer(client)
+		connection := newConnection(app)
+		connection.Server = lsp.NewClientServer(connection.Client)
+		return connection, connection.initialize(ctx)
 	case "internal":
+		internalMu.Lock()
+		defer internalMu.Unlock()
+		if c := internalConnections[app.Config.Dir]; c != nil {
+			return c, nil
+		}
+		connection := newConnection(app)
+		ctx := context.Background() //TODO:a way of shutting down the internal server
 		cr, sw, _ := os.Pipe()
 		sr, cw, _ := os.Pipe()
 		var jc *jsonrpc2.Conn
-		jc, server, _ = protocol.NewClient(jsonrpc2.NewHeaderStream(cr, cw), client)
+		jc, connection.Server, _ = protocol.NewClient(jsonrpc2.NewHeaderStream(cr, cw), connection.Client)
 		go jc.Run(ctx)
 		go lsp.NewServer(jsonrpc2.NewHeaderStream(sr, sw)).Run(ctx)
+		if err := connection.initialize(ctx); err != nil {
+			return nil, err
+		}
+		internalConnections[app.Config.Dir] = connection
+		return connection, nil
 	default:
+		connection := newConnection(app)
 		conn, err := net.Dial("tcp", app.Remote)
 		if err != nil {
 			return nil, err
 		}
 		stream := jsonrpc2.NewHeaderStream(conn, conn)
 		var jc *jsonrpc2.Conn
-		jc, server, _ = protocol.NewClient(stream, client)
+		jc, connection.Server, _ = protocol.NewClient(stream, connection.Client)
 		go jc.Run(ctx)
+		return connection, connection.initialize(ctx)
 	}
+}
 
+func (c *connection) initialize(ctx context.Context) error {
 	params := &protocol.InitializeParams{}
-	params.RootURI = string(span.FileURI(app.Config.Dir))
+	params.RootURI = string(span.FileURI(c.Client.app.Config.Dir))
 	params.Capabilities.Workspace.Configuration = true
 	params.Capabilities.TextDocument.Hover.ContentFormat = []protocol.MarkupKind{protocol.PlainText}
-
-	client.prepare(app, server)
-	if _, err := server.Initialize(ctx, params); err != nil {
-		return nil, err
+	if _, err := c.Server.Initialize(ctx, params); err != nil {
+		return err
 	}
-	if err := server.Initialized(ctx, &protocol.InitializedParams{}); err != nil {
-		return nil, err
+	if err := c.Server.Initialized(ctx, &protocol.InitializedParams{}); err != nil {
+		return err
 	}
-	return server, nil
+	return nil
 }
 
-type baseClient struct {
+type connection struct {
+	protocol.Server
+	Client *cmdClient
+}
+
+type cmdClient struct {
 	protocol.Server
 	app  *Application
-	server protocol.Server
 	fset *token.FileSet
+
+	filesMu sync.Mutex
+	files   map[span.URI]*cmdFile
 }
 
-func (c *baseClient) ShowMessage(ctx context.Context, p *protocol.ShowMessageParams) error { return nil }
-func (c *baseClient) ShowMessageRequest(ctx context.Context, p *protocol.ShowMessageRequestParams) (*protocol.MessageActionItem, error) {
+type cmdFile struct {
+	uri            span.URI
+	mapper         *protocol.ColumnMapper
+	err            error
+	added          bool
+	hasDiagnostics chan struct{}
+	diagnostics    []protocol.Diagnostic
+}
+
+func newConnection(app *Application) *connection {
+	return &connection{
+		Client: &cmdClient{
+			app:   app,
+			fset:  token.NewFileSet(),
+			files: make(map[span.URI]*cmdFile),
+		},
+	}
+}
+
+func (c *cmdClient) ShowMessage(ctx context.Context, p *protocol.ShowMessageParams) error { return nil }
+
+func (c *cmdClient) ShowMessageRequest(ctx context.Context, p *protocol.ShowMessageRequestParams) (*protocol.MessageActionItem, error) {
 	return nil, nil
 }
-func (c *baseClient) LogMessage(ctx context.Context, p *protocol.LogMessageParams) error {
+
+func (c *cmdClient) LogMessage(ctx context.Context, p *protocol.LogMessageParams) error {
 	switch p.Type {
 	case protocol.Error:
 		log.Print("Error:", p.Message)
@@ -190,17 +233,22 @@ func (c *baseClient) LogMessage(ctx context.Context, p *protocol.LogMessageParam
 	}
 	return nil
 }
-func (c *baseClient) Event(ctx context.Context, t *interface{}) error { return nil }
-func (c *baseClient) RegisterCapability(ctx context.Context, p *protocol.RegistrationParams) error {
+
+func (c *cmdClient) Event(ctx context.Context, t *interface{}) error { return nil }
+
+func (c *cmdClient) RegisterCapability(ctx context.Context, p *protocol.RegistrationParams) error {
 	return nil
 }
-func (c *baseClient) UnregisterCapability(ctx context.Context, p *protocol.UnregistrationParams) error {
+
+func (c *cmdClient) UnregisterCapability(ctx context.Context, p *protocol.UnregistrationParams) error {
 	return nil
 }
-func (c *baseClient) WorkspaceFolders(ctx context.Context) ([]protocol.WorkspaceFolder, error) {
+
+func (c *cmdClient) WorkspaceFolders(ctx context.Context) ([]protocol.WorkspaceFolder, error) {
 	return nil, nil
 }
-func (c *baseClient) Configuration(ctx context.Context, p *protocol.ConfigurationParams) ([]interface{}, error) {
+
+func (c *cmdClient) Configuration(ctx context.Context, p *protocol.ConfigurationParams) ([]interface{}, error) {
 	results := make([]interface{}, len(p.Items))
 	for i, item := range p.Items {
 		if item.Section != "gopls" {
@@ -218,36 +266,66 @@ func (c *baseClient) Configuration(ctx context.Context, p *protocol.Configuratio
 	}
 	return results, nil
 }
-func (c *baseClient) ApplyEdit(ctx context.Context, p *protocol.ApplyWorkspaceEditParams) (*protocol.ApplyWorkspaceEditResponse, error) {
+
+func (c *cmdClient) ApplyEdit(ctx context.Context, p *protocol.ApplyWorkspaceEditParams) (*protocol.ApplyWorkspaceEditResponse, error) {
 	return &protocol.ApplyWorkspaceEditResponse{Applied: false, FailureReason: "not implemented"}, nil
 }
-func (c *baseClient) PublishDiagnostics(ctx context.Context, p *protocol.PublishDiagnosticsParams) error {
+
+func (c *cmdClient) PublishDiagnostics(ctx context.Context, p *protocol.PublishDiagnosticsParams) error {
+	c.filesMu.Lock()
+	defer c.filesMu.Unlock()
+	uri := span.URI(p.URI)
+	file := c.getFile(ctx, uri)
+	hadDiagnostics := file.diagnostics != nil
+	file.diagnostics = p.Diagnostics
+	if file.diagnostics == nil {
+		file.diagnostics = []protocol.Diagnostic{}
+	}
+	if !hadDiagnostics {
+		close(file.hasDiagnostics)
+	}
 	return nil
 }
 
-func (c *baseClient) prepare(app *Application, server protocol.Server) {
-	c.app = app
-	c.server = server
-	c.fset = token.NewFileSet()
-}
-
-func (c *baseClient) AddFile(ctx context.Context, uri span.URI) (*protocol.ColumnMapper, error) {
+func (c *cmdClient) getFile(ctx context.Context, uri span.URI) *cmdFile {
+	file, found := c.files[uri]
+	if !found {
+		file = &cmdFile{
+			uri:            uri,
+			hasDiagnostics: make(chan struct{}),
+		}
+		c.files[uri] = file
+	}
+	if file.mapper == nil {
 		fname, err := uri.Filename()
 		if err != nil {
-		return nil, fmt.Errorf("%v: %v", uri, err)
+			file.err = fmt.Errorf("%v: %v", uri, err)
+			return file
 		}
 		content, err := ioutil.ReadFile(fname)
 		if err != nil {
-		return nil, fmt.Errorf("%v: %v", uri, err)
+			file.err = fmt.Errorf("%v: %v", uri, err)
+			return file
 		}
 		f := c.fset.AddFile(fname, -1, len(content))
 		f.SetLinesForContent(content)
-	m := protocol.NewColumnMapper(uri, c.fset, f, content)
+		file.mapper = protocol.NewColumnMapper(uri, c.fset, f, content)
+	}
+	return file
+}
+
+func (c *connection) AddFile(ctx context.Context, uri span.URI) *cmdFile {
+	c.Client.filesMu.Lock()
+	defer c.Client.filesMu.Unlock()
+	file := c.Client.getFile(ctx, uri)
+	if !file.added {
+		file.added = true
 		p := &protocol.DidOpenTextDocumentParams{}
 		p.TextDocument.URI = string(uri)
-	p.TextDocument.Text = string(content)
-	if err := c.server.DidOpen(ctx, p); err != nil {
-		return nil, fmt.Errorf("%v: %v", uri, err)
+		p.TextDocument.Text = string(file.mapper.Content)
+		if err := c.Server.DidOpen(ctx, p); err != nil {
+			file.err = fmt.Errorf("%v: %v", uri, err)
 		}
-	return m, nil
+	}
+	return file
 }

internal/lsp/cmd/definition.go

@@ -59,17 +59,16 @@ func (d *definition) Run(ctx context.Context, args ...string) error {
 	if len(args) != 1 {
 		return tool.CommandLineErrorf("definition expects 1 argument")
 	}
-	client := &baseClient{}
-	server, err := d.query.app.connect(ctx, client)
+	conn, err := d.query.app.connect(ctx)
 	if err != nil {
 		return err
 	}
 	from := span.Parse(args[0])
-	m, err := client.AddFile(ctx, from.URI())
-	if err != nil {
-		return err
+	file := conn.AddFile(ctx, from.URI())
+	if file.err != nil {
+		return file.err
 	}
-	loc, err := m.Location(from)
+	loc, err := file.mapper.Location(from)
 	if err != nil {
 		return err
 	}
@@ -77,7 +76,7 @@ func (d *definition) Run(ctx context.Context, args ...string) error {
 		TextDocument: protocol.TextDocumentIdentifier{URI: loc.URI},
 		Position:     loc.Range.Start,
 	}
-	locs, err := server.Definition(ctx, &p)
+	locs, err := conn.Definition(ctx, &p)
 	if err != nil {
 		return fmt.Errorf("%v: %v", from, err)
 	}
@@ -85,18 +84,18 @@ func (d *definition) Run(ctx context.Context, args ...string) error {
 	if len(locs) == 0 {
 		return fmt.Errorf("%v: not an identifier", from)
 	}
-	hover, err := server.Hover(ctx, &p)
+	hover, err := conn.Hover(ctx, &p)
 	if err != nil {
 		return fmt.Errorf("%v: %v", from, err)
 	}
 	if hover == nil {
 		return fmt.Errorf("%v: not an identifier", from)
 	}
-	m, err = client.AddFile(ctx, span.NewURI(locs[0].URI))
-	if err != nil {
-		return fmt.Errorf("%v: %v", from, err)
+	file = conn.AddFile(ctx, span.NewURI(locs[0].URI))
+	if file.err != nil {
+		return fmt.Errorf("%v: %v", from, file.err)
 	}
-	definition, err := m.Span(locs[0])
+	definition, err := file.mapper.Span(locs[0])
 	if err != nil {
 		return fmt.Errorf("%v: %v", from, err)
 	}

internal/lsp/cmd/definition_test.go

@@ -34,9 +34,6 @@ func TestDefinitionHelpExample(t *testing.T) {
 	if runtime.GOOS == "android" {
 		t.Skip("not all source files are available on android")
 	}
-	if runtime.GOOS != "linux" || isRace {
-		t.Skip("currently uses too much memory, see issue #31611")
-	}
 	dir, err := os.Getwd()
 	if err != nil {
 		t.Errorf("could not get wd: %v", err)
@@ -64,7 +61,7 @@ func (r *runner) Definition(t *testing.T, data tests.Definitions) {
 			// TODO: support type definition queries
 			continue
 		}
-		args := []string{"query"}
+		args := []string{"-remote=internal", "query"}
 		if d.Flags != "" {
 			args = append(args, strings.Split(d.Flags, " ")...)
 		}

internal/lsp/cmd/format.go

@@ -50,20 +50,19 @@ func (f *format) Run(ctx context.Context, args ...string) error {
 		// no files, so no results
 		return nil
 	}
-	client := &baseClient{}
 	// now we ready to kick things off
-	server, err := f.app.connect(ctx, client)
+	conn, err := f.app.connect(ctx)
 	if err != nil {
 		return err
 	}
 	for _, arg := range args {
 		spn := span.Parse(arg)
-		m, err := client.AddFile(ctx, spn.URI())
-		if err != nil {
-			return err
+		file := conn.AddFile(ctx, spn.URI())
+		if file.err != nil {
+			return file.err
 		}
 		filename, _ := spn.URI().Filename() // this cannot fail, already checked in AddFile above
-		loc, err := m.Location(spn)
+		loc, err := file.mapper.Location(spn)
 		if err != nil {
 			return err
 		}
@@ -73,16 +72,16 @@ func (f *format) Run(ctx context.Context, args ...string) error {
 		p := protocol.DocumentFormattingParams{
 			TextDocument: protocol.TextDocumentIdentifier{URI: loc.URI},
 		}
-		edits, err := server.Formatting(ctx, &p)
+		edits, err := conn.Formatting(ctx, &p)
 		if err != nil {
 			return fmt.Errorf("%v: %v", spn, err)
 		}
-		sedits, err := lsp.FromProtocolEdits(m, edits)
+		sedits, err := lsp.FromProtocolEdits(file.mapper, edits)
 		if err != nil {
 			return fmt.Errorf("%v: %v", spn, err)
 		}
 		ops := source.EditsToDiff(sedits)
-		lines := diff.SplitLines(string(m.Content))
+		lines := diff.SplitLines(string(file.mapper.Content))
 		formatted := strings.Join(diff.ApplyEdits(lines, ops), "")
 		printIt := true
 		if f.List {

internal/lsp/cmd/format_test.go

@@ -44,7 +44,7 @@ func (r *runner) Format(t *testing.T, data tests.Formats) {
 			app := &cmd.Application{}
 			app.Config = r.data.Config
 			got := captureStdOut(t, func() {
-				tool.Main(context.Background(), app, append([]string{"format"}, args...))
+				tool.Main(context.Background(), app, append([]string{"-remote=internal", "format"}, args...))
 			})
 			got = r.normalizePaths(got)
 			// check the first two lines are the expected file header