internal/lsp: move the progress tracker to the session

Sometimes, we may want to report progress from functions inside of the
cache package, so move the progress tracker to the session to allow for
that.

Change-Id: I15409577a7a5080e7f0224a95d159de42856ffa7
Reviewed-on: https://go-review.googlesource.com/c/tools/+/319330
Trust: Rebecca Stambler <rstambler@golang.org>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>

Author: stamblerre

internal/lsp/cache/load.go

@@ -55,6 +55,15 @@ type metadata struct {
 // load calls packages.Load for the given scopes, updating package metadata,
 // import graph, and mapped files with the result.
 func (s *snapshot) load(ctx context.Context, allowNetwork bool, scopes ...interface{}) error {
+	if s.view.Options().VerboseWorkDoneProgress {
+		work := s.view.session.progress.Start(ctx, "Load", fmt.Sprintf("Loading scopes %s", scopes), nil, nil)
+		defer func() {
+			go func() {
+				work.End("Done.")
+			}()
+		}()
+	}
+
 	var query []string
 	var containsDir bool // for logging
 	for _, scope := range scopes {

internal/lsp/cache/session.go

@@ -14,6 +14,7 @@ import (
 	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/gocommand"
 	"golang.org/x/tools/internal/imports"
+	"golang.org/x/tools/internal/lsp/progress"
 	"golang.org/x/tools/internal/lsp/source"
 	"golang.org/x/tools/internal/span"
 	"golang.org/x/tools/internal/xcontext"
@@ -36,6 +37,8 @@ type Session struct {
 
 	// gocmdRunner guards go command calls from concurrency errors.
 	gocmdRunner *gocommand.Runner
+
+	progress *progress.Tracker
 }
 
 type overlay struct {
@@ -131,6 +134,11 @@ func (s *Session) SetOptions(options *source.Options) {
 	s.options = options
 }
 
+func (s *Session) SetProgressTracker(tracker *progress.Tracker) {
+	// The progress tracker should be set before any view is initialized.
+	s.progress = tracker
+}
+
 func (s *Session) Shutdown(ctx context.Context) {
 	s.viewMu.Lock()
 	defer s.viewMu.Unlock()

internal/lsp/cmd/info.go

@@ -178,6 +178,6 @@ func (l *licenses) Run(ctx context.Context, args ...string) error {
 	} else {
 		txt += opts.LicensesText
 	}
-	fmt.Fprintf(os.Stdout, txt)
+	fmt.Fprint(os.Stdout, txt)
 	return nil
 }

internal/lsp/command.go

@@ -20,6 +20,7 @@ import (
 	"golang.org/x/tools/internal/gocommand"
 	"golang.org/x/tools/internal/lsp/command"
 	"golang.org/x/tools/internal/lsp/debug"
+	"golang.org/x/tools/internal/lsp/progress"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/source"
 	"golang.org/x/tools/internal/span"
@@ -65,7 +66,7 @@ type commandConfig struct {
 type commandDeps struct {
 	snapshot source.Snapshot            // present if cfg.forURI was set
 	fh       source.VersionedFileHandle // present if cfg.forURI was set
-	work     *workDone                  // present cfg.progress was set
+	work     *progress.WorkDone         // present cfg.progress was set
 }
 
 type commandFunc func(context.Context, commandDeps) error
@@ -90,19 +91,21 @@ func (c *commandHandler) run(ctx context.Context, cfg commandConfig, run command
 	}
 	ctx, cancel := context.WithCancel(xcontext.Detach(ctx))
 	if cfg.progress != "" {
-		deps.work = c.s.progress.start(ctx, cfg.progress, "Running...", c.params.WorkDoneToken, cancel)
+		deps.work = c.s.progress.Start(ctx, cfg.progress, "Running...", c.params.WorkDoneToken, cancel)
 	}
 	runcmd := func() error {
 		defer cancel()
 		err := run(ctx, deps)
-		switch {
-		case errors.Is(err, context.Canceled):
-			deps.work.end("canceled")
-		case err != nil:
-			event.Error(ctx, "command error", err)
-			deps.work.end("failed")
-		default:
-			deps.work.end("completed")
+		if deps.work != nil {
+			switch {
+			case errors.Is(err, context.Canceled):
+				deps.work.End("canceled")
+			case err != nil:
+				event.Error(ctx, "command error", err)
+				deps.work.End("failed")
+			default:
+				deps.work.End("completed")
+			}
 		}
 		return err
 	}
@@ -349,7 +352,7 @@ func (c *commandHandler) RunTests(ctx context.Context, args command.RunTestsArgs
 	})
 }
 
-func (c *commandHandler) runTests(ctx context.Context, snapshot source.Snapshot, work *workDone, uri protocol.DocumentURI, tests, benchmarks []string) error {
+func (c *commandHandler) runTests(ctx context.Context, snapshot source.Snapshot, work *progress.WorkDone, uri protocol.DocumentURI, tests, benchmarks []string) error {
 	// TODO: fix the error reporting when this runs async.
 	pkgs, err := snapshot.PackagesForFile(ctx, uri.SpanURI(), source.TypecheckWorkspace)
 	if err != nil {
@@ -362,8 +365,8 @@ func (c *commandHandler) runTests(ctx context.Context, snapshot source.Snapshot,
 
 	// create output
 	buf := &bytes.Buffer{}
-	ew := &eventWriter{ctx: ctx, operation: "test"}
-	out := io.MultiWriter(ew, workDoneWriter{work}, buf)
+	ew := progress.NewEventWriter(ctx, "test")
+	out := io.MultiWriter(ew, progress.NewWorkDoneWriter(work), buf)
 
 	// Run `go test -run Func` on each test.
 	var failedTests int
@@ -435,7 +438,7 @@ func (c *commandHandler) Generate(ctx context.Context, args command.GenerateArgs
 		progress:    title,
 		forURI:      args.Dir,
 	}, func(ctx context.Context, deps commandDeps) error {
-		er := &eventWriter{ctx: ctx, operation: "generate"}
+		er := progress.NewEventWriter(ctx, "generate")
 
 		pattern := "."
 		if args.Recursive {
@@ -446,7 +449,7 @@ func (c *commandHandler) Generate(ctx context.Context, args command.GenerateArgs
 			Args:       []string{"-x", pattern},
 			WorkingDir: args.Dir.SpanURI().Filename(),
 		}
-		stderr := io.MultiWriter(er, workDoneWriter{deps.work})
+		stderr := io.MultiWriter(er, progress.NewWorkDoneWriter(deps.work))
 		if err := deps.snapshot.RunGoCommandPiped(ctx, source.Normal, inv, er, stderr); err != nil {
 			return err
 		}

internal/lsp/diagnostics.go

@@ -384,18 +384,18 @@ func (s *Server) showCriticalErrorStatus(ctx context.Context, snapshot source.Sn
 
 	if s.criticalErrorStatus == nil {
 		if errMsg != "" {
-			s.criticalErrorStatus = s.progress.start(ctx, WorkspaceLoadFailure, errMsg, nil, nil)
+			s.criticalErrorStatus = s.progress.Start(ctx, WorkspaceLoadFailure, errMsg, nil, nil)
 		}
 		return
 	}
 
 	// If an error is already shown to the user, update it or mark it as
 	// resolved.
 	if errMsg == "" {
-		s.criticalErrorStatus.end("Done.")
+		s.criticalErrorStatus.End("Done.")
 		s.criticalErrorStatus = nil
 	} else {
-		s.criticalErrorStatus.report(errMsg, 0)
+		s.criticalErrorStatus.Report(errMsg, 0)
 	}
 }
 

internal/lsp/general.go

@@ -46,7 +46,7 @@ func (s *Server) initialize(ctx context.Context, params *protocol.ParamInitializ
 		event.Error(ctx, "creating temp dir", err)
 		s.tempDir = ""
 	}
-	s.progress.supportsWorkDoneProgress = params.Capabilities.Window.WorkDoneProgress
+	s.progress.SetSupportsWorkDoneProgress(params.Capabilities.Window.WorkDoneProgress)
 
 	options := s.session.Options()
 	defer func() { s.session.SetOptions(options) }()
@@ -217,11 +217,11 @@ func (s *Server) addFolders(ctx context.Context, folders []protocol.WorkspaceFol
 
 	var wg sync.WaitGroup
 	if s.session.Options().VerboseWorkDoneProgress {
-		work := s.progress.start(ctx, DiagnosticWorkTitle(FromInitialWorkspaceLoad), "Calculating diagnostics for initial workspace load...", nil, nil)
+		work := s.progress.Start(ctx, DiagnosticWorkTitle(FromInitialWorkspaceLoad), "Calculating diagnostics for initial workspace load...", nil, nil)
 		defer func() {
 			go func() {
 				wg.Wait()
-				work.end("Done.")
+				work.End("Done.")
 			}()
 		}()
 	}
@@ -233,11 +233,11 @@ func (s *Server) addFolders(ctx context.Context, folders []protocol.WorkspaceFol
 		if !uri.IsFile() {
 			continue
 		}
-		work := s.progress.start(ctx, "Setting up workspace", "Loading packages...", nil, nil)
+		work := s.progress.Start(ctx, "Setting up workspace", "Loading packages...", nil, nil)
 		snapshot, release, err := s.addView(ctx, folder.Name, uri)
 		if err != nil {
 			viewErrors[uri] = err
-			work.end(fmt.Sprintf("Error loading packages: %s", err))
+			work.End(fmt.Sprintf("Error loading packages: %s", err))
 			continue
 		}
 		var swg sync.WaitGroup
@@ -247,7 +247,7 @@ func (s *Server) addFolders(ctx context.Context, folders []protocol.WorkspaceFol
 			defer swg.Done()
 			defer allFoldersWg.Done()
 			snapshot.AwaitInitialized(ctx)
-			work.end("Finished loading packages.")
+			work.End("Finished loading packages.")
 		}()
 
 		// Print each view's environment.

internal/lsp/lsp_test.go

@@ -92,6 +92,7 @@ func testLSP(t *testing.T, datum *tests.Data) {
 		normalizers: tests.CollectNormalizers(datum.Exported),
 		editRecv:    make(chan map[span.URI]string, 1),
 	}
+
 	r.server = NewServer(session, testClient{runner: r})
 	tests.Run(t, r, datum)
 }

internal/lsp/progress.go renamed to internal/lsp/progress/progress.go

@@ -2,7 +2,7 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-package lsp
+package progress
 
 import (
 	"context"
@@ -18,22 +18,26 @@ import (
 	errors "golang.org/x/xerrors"
 )
 
-type progressTracker struct {
+type Tracker struct {
 	client                   protocol.Client
 	supportsWorkDoneProgress bool
 
 	mu         sync.Mutex
-	inProgress map[protocol.ProgressToken]*workDone
+	inProgress map[protocol.ProgressToken]*WorkDone
 }
 
-func newProgressTracker(client protocol.Client) *progressTracker {
-	return &progressTracker{
+func NewTracker(client protocol.Client) *Tracker {
+	return &Tracker{
 		client:     client,
-		inProgress: make(map[protocol.ProgressToken]*workDone),
+		inProgress: make(map[protocol.ProgressToken]*WorkDone),
 	}
 }
 
-// start notifies the client of work being done on the server. It uses either
+func (tracker *Tracker) SetSupportsWorkDoneProgress(b bool) {
+	tracker.supportsWorkDoneProgress = b
+}
+
+// Start notifies the client of work being done on the server. It uses either
 // ShowMessage RPCs or $/progress messages, depending on the capabilities of
 // the client.  The returned WorkDone handle may be used to report incremental
 // progress, and to report work completion. In particular, it is an error to
@@ -59,8 +63,8 @@ func newProgressTracker(client protocol.Client) *progressTracker {
 //    // Do the work...
 //  }
 //
-func (t *progressTracker) start(ctx context.Context, title, message string, token protocol.ProgressToken, cancel func()) *workDone {
-	wd := &workDone{
+func (t *Tracker) Start(ctx context.Context, title, message string, token protocol.ProgressToken, cancel func()) *WorkDone {
+	wd := &WorkDone{
 		ctx:    xcontext.Detach(ctx),
 		client: t.client,
 		token:  token,
@@ -119,7 +123,7 @@ func (t *progressTracker) start(ctx context.Context, title, message string, toke
 	return wd
 }
 
-func (t *progressTracker) cancel(ctx context.Context, token protocol.ProgressToken) error {
+func (t *Tracker) Cancel(ctx context.Context, token protocol.ProgressToken) error {
 	t.mu.Lock()
 	defer t.mu.Unlock()
 	wd, ok := t.inProgress[token]
@@ -133,9 +137,9 @@ func (t *progressTracker) cancel(ctx context.Context, token protocol.ProgressTok
 	return nil
 }
 
-// workDone represents a unit of work that is reported to the client via the
+// WorkDone represents a unit of work that is reported to the client via the
 // progress API.
-type workDone struct {
+type WorkDone struct {
 	// ctx is detached, for sending $/progress updates.
 	ctx    context.Context
 	client protocol.Client
@@ -153,7 +157,11 @@ type workDone struct {
 	cleanup func()
 }
 
-func (wd *workDone) doCancel() {
+func (wd *WorkDone) Token() protocol.ProgressToken {
+	return wd.token
+}
+
+func (wd *WorkDone) doCancel() {
 	wd.cancelMu.Lock()
 	defer wd.cancelMu.Unlock()
 	if !wd.cancelled {
@@ -162,7 +170,7 @@ func (wd *workDone) doCancel() {
 }
 
 // report reports an update on WorkDone report back to the client.
-func (wd *workDone) report(message string, percentage float64) {
+func (wd *WorkDone) Report(message string, percentage float64) {
 	if wd == nil {
 		return
 	}
@@ -196,7 +204,7 @@ func (wd *workDone) report(message string, percentage float64) {
 }
 
 // end reports a workdone completion back to the client.
-func (wd *workDone) end(message string) {
+func (wd *WorkDone) End(message string) {
 	if wd == nil {
 		return
 	}
@@ -227,27 +235,35 @@ func (wd *workDone) end(message string) {
 	}
 }
 
-// eventWriter writes every incoming []byte to
+// EventWriter writes every incoming []byte to
 // event.Print with the operation=generate tag
 // to distinguish its logs from others.
-type eventWriter struct {
+type EventWriter struct {
 	ctx       context.Context
 	operation string
 }
 
-func (ew *eventWriter) Write(p []byte) (n int, err error) {
+func NewEventWriter(ctx context.Context, operation string) *EventWriter {
+	return &EventWriter{ctx: ctx, operation: operation}
+}
+
+func (ew *EventWriter) Write(p []byte) (n int, err error) {
 	event.Log(ew.ctx, string(p), tag.Operation.Of(ew.operation))
 	return len(p), nil
 }
 
-// workDoneWriter wraps a workDone handle to provide a Writer interface,
+// WorkDoneWriter wraps a workDone handle to provide a Writer interface,
 // so that workDone reporting can more easily be hooked into commands.
-type workDoneWriter struct {
-	wd *workDone
+type WorkDoneWriter struct {
+	wd *WorkDone
+}
+
+func NewWorkDoneWriter(wd *WorkDone) *WorkDoneWriter {
+	return &WorkDoneWriter{wd: wd}
 }
 
-func (wdw workDoneWriter) Write(p []byte) (n int, err error) {
-	wdw.wd.report(string(p), 0)
+func (wdw WorkDoneWriter) Write(p []byte) (n int, err error) {
+	wdw.wd.Report(string(p), 0)
 	// Don't fail just because of a failure to report progress.
 	return len(p), nil
 }