internal/lsp: factor out progress reporting to a new WorkDone handle

Our current usage of WorkDone progress reporting (new in v3.15 of the
LSP spec) is in reporting progress on `go generate` commands. In
preparation for using this API more widely, factor out the reporting API
from the current io.WriteCloser wrapper (workDoneWriter).

Change-Id: Ib528093d81d4fc065528df90e100859e850b10df
Reviewed-on: https://go-review.googlesource.com/c/tools/+/229459
Run-TryBot: Robert Findley <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>
Reviewed-by: Heschi Kreinick <[email protected]>
Reviewed-by: Rebecca Stambler <[email protected]>

Author: findleyr

internal/lsp/general.go

@@ -31,7 +31,7 @@ func (s *Server) initialize(ctx context.Context, params *protocol.ParamInitializ
 	s.stateMu.Unlock()
 
 	s.supportsWorkDoneProgress = params.Capabilities.Window.WorkDoneProgress
-	s.inProgress = map[string]func(){}
+	s.inProgress = map[string]*WorkDone{}
 
 	options := s.session.Options()
 	defer func() { s.session.SetOptions(options) }()

internal/lsp/generate.go

@@ -7,8 +7,6 @@ package lsp
 import (
 	"context"
 	"io"
-	"math/rand"
-	"strconv"
 
 	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/gocommand"
@@ -21,14 +19,8 @@ func (s *Server) runGenerate(ctx context.Context, dir string, recursive bool) {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	token := strconv.FormatInt(rand.Int63(), 10)
-	s.inProgressMu.Lock()
-	s.inProgress[token] = cancel
-	s.inProgressMu.Unlock()
-	defer s.clearInProgress(token)
-
 	er := &eventWriter{ctx: ctx}
-	wc := s.newProgressWriter(ctx, cancel, token)
+	wc := s.newProgressWriter(ctx, cancel)
 	defer wc.Close()
 	args := []string{"-x"}
 	if recursive {
@@ -68,18 +60,14 @@ func (ew *eventWriter) Write(p []byte) (n int, err error) {
 // newProgressWriter returns an io.WriterCloser that can be used
 // to report progress on the "go generate" command based on the
 // client capabilities.
-func (s *Server) newProgressWriter(ctx context.Context, cancel func(), token string) io.WriteCloser {
-	var wc interface {
-		io.WriteCloser
-		start()
-	}
+func (s *Server) newProgressWriter(ctx context.Context, cancel func()) io.WriteCloser {
 	if s.supportsWorkDoneProgress {
-		wc = &workDoneWriter{ctx, token, s.client}
-	} else {
-		wc = &messageWriter{ctx, cancel, s.client}
+		wd := s.StartWork(ctx, "generate", "running go generate", cancel)
+		return &workDoneWriter{ctx, wd}
 	}
-	wc.start()
-	return wc
+	mw := &messageWriter{ctx, cancel, s.client}
+	mw.start()
+	return mw
 }
 
 // messageWriter implements progressWriter
@@ -126,58 +114,19 @@ func (lw *messageWriter) Close() error {
 	})
 }
 
-// workDoneWriter implements progressWriter
-// that will send $/progress notifications
-// to the client. Request cancellations
-// happens separately through the
-// window/workDoneProgress/cancel request
-// in which case the given context will be rendered
-// done.
+// workDoneWriter implements progressWriter by sending $/progress notifications
+// to the client. Request cancellations happens separately through the
+// window/workDoneProgress/cancel request, in which case the given context will
+// be rendered done.
 type workDoneWriter struct {
-	ctx    context.Context
-	token  string
-	client protocol.Client
+	ctx context.Context
+	wd  *WorkDone
 }
 
 func (wdw *workDoneWriter) Write(p []byte) (n int, err error) {
-	return len(p), wdw.client.Progress(wdw.ctx, &protocol.ProgressParams{
-		Token: wdw.token,
-		Value: &protocol.WorkDoneProgressReport{
-			Kind:        "report",
-			Cancellable: true,
-			Message:     string(p),
-		},
-	})
-}
-
-func (wdw *workDoneWriter) start() {
-	err := wdw.client.WorkDoneProgressCreate(wdw.ctx, &protocol.WorkDoneProgressCreateParams{
-		Token: wdw.token,
-	})
-	if err != nil {
-		event.Error(wdw.ctx, "generate progress create", err)
-		return
-	}
-	err = wdw.client.Progress(wdw.ctx, &protocol.ProgressParams{
-		Token: wdw.token,
-		Value: &protocol.WorkDoneProgressBegin{
-			Kind:        "begin",
-			Cancellable: true,
-			Message:     "running go generate",
-			Title:       "generate",
-		},
-	})
-	if err != nil {
-		event.Error(wdw.ctx, "generate progress begin", err)
-	}
+	return len(p), wdw.wd.Progress(wdw.ctx, string(p), 0)
 }
 
 func (wdw *workDoneWriter) Close() error {
-	return wdw.client.Progress(wdw.ctx, &protocol.ProgressParams{
-		Token: wdw.token,
-		Value: protocol.WorkDoneProgressEnd{
-			Kind:    "end",
-			Message: "finished",
-		},
-	})
+	return wdw.wd.End(wdw.ctx, "finished")
 }

internal/lsp/progress.go

@@ -0,0 +1,123 @@
+// 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 lsp
+
+import (
+	"context"
+	"errors"
+	"math/rand"
+	"strconv"
+
+	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/lsp/protocol"
+)
+
+// WorkDone represents a unit of work that is reported to the client via the
+// progress API.
+type WorkDone struct {
+	client   protocol.Client
+	startErr error
+	token    string
+	cancel   func()
+	cleanup  func()
+}
+
+// StartWork creates a unique token and issues a $/progress notification to
+// begin a unit of work on the server. The returned WorkDone handle may be used
+// to report incremental progress, and to report work completion. In
+// particular, it is an error to call StartWork and not call End(...) on the
+// returned WorkDone handle.
+//
+// The progress item is considered cancellable if the given cancel func is
+// non-nil.
+//
+// Example:
+//  func Generate(ctx) (err error) {
+//    ctx, cancel := context.WithCancel(ctx)
+//    defer cancel()
+//    work := s.StartWork(ctx, "generate", "running go generate", cancel)
+//    defer func() {
+//      if err != nil {
+//        work.End(ctx, fmt.Sprintf("generate failed: %v", err))
+//      } else {
+//        work.End(ctx, "done")
+//      }
+//    }()
+//    // Do the work...
+//  }
+//
+func (s *Server) StartWork(ctx context.Context, title, message string, cancel func()) *WorkDone {
+	wd := &WorkDone{
+		client: s.client,
+		token:  strconv.FormatInt(rand.Int63(), 10),
+		cancel: cancel,
+	}
+	if !s.supportsWorkDoneProgress {
+		wd.startErr = errors.New("workdone reporting is not supported")
+		return wd
+	}
+	err := wd.client.WorkDoneProgressCreate(ctx, &protocol.WorkDoneProgressCreateParams{
+		Token: wd.token,
+	})
+	if err != nil {
+		wd.startErr = err
+		event.Error(ctx, "starting work for "+title, err)
+		return wd
+	}
+	s.addInProgress(wd)
+	wd.cleanup = func() {
+		s.removeInProgress(wd.token)
+	}
+	err = wd.client.Progress(ctx, &protocol.ProgressParams{
+		Token: wd.token,
+		Value: &protocol.WorkDoneProgressBegin{
+			Kind:        "begin",
+			Cancellable: wd.cancel != nil,
+			Message:     message,
+			Title:       title,
+		},
+	})
+	if err != nil {
+		event.Error(ctx, "generate progress begin", err)
+	}
+	return wd
+}
+
+// Progress reports an update on WorkDone progress back to the client.
+func (wd *WorkDone) Progress(ctx context.Context, message string, percentage float64) error {
+	if wd.startErr != nil {
+		return wd.startErr
+	}
+	return wd.client.Progress(ctx, &protocol.ProgressParams{
+		Token: wd.token,
+		Value: &protocol.WorkDoneProgressReport{
+			Kind: "report",
+			// Note that in the LSP spec, the value of Cancellable may be changed to
+			// control whether the cancel button in the UI is enabled. Since we don't
+			// yet use this feature, the value is kept constant here.
+			Cancellable: wd.cancel != nil,
+			Message:     message,
+			Percentage:  percentage,
+		},
+	})
+}
+
+// End reports a workdone completion back to the client.
+func (wd *WorkDone) End(ctx context.Context, message string) error {
+	if wd.startErr != nil {
+		return wd.startErr
+	}
+	err := wd.client.Progress(ctx, &protocol.ProgressParams{
+		Token: wd.token,
+		Value: protocol.WorkDoneProgressEnd{
+			Kind:    "end",
+			Message: message,
+		},
+	})
+	if wd.cleanup != nil {
+		wd.cleanup()
+	}
+	return err
+}

internal/lsp/server.go

@@ -81,7 +81,7 @@ type Server struct {
 	// to determine if the client can support progress notifications
 	supportsWorkDoneProgress bool
 	inProgressMu             sync.Mutex
-	inProgress               map[string]func()
+	inProgress               map[string]*WorkDone
 }
 
 // sentDiagnostics is used to cache diagnostics that have been sent for a given file.
@@ -146,15 +146,24 @@ func (s *Server) workDoneProgressCancel(ctx context.Context, params *protocol.Wo
 	}
 	s.inProgressMu.Lock()
 	defer s.inProgressMu.Unlock()
-	cancel, ok := s.inProgress[token]
+	wd, ok := s.inProgress[token]
 	if !ok {
 		return errors.Errorf("token %q not found in progress", token)
 	}
-	cancel()
+	if wd.cancel == nil {
+		return errors.Errorf("work %q is not cancellable", token)
+	}
+	wd.cancel()
 	return nil
 }
 
-func (s *Server) clearInProgress(token string) {
+func (s *Server) addInProgress(wd *WorkDone) {
+	s.inProgressMu.Lock()
+	s.inProgress[wd.token] = wd
+	s.inProgressMu.Unlock()
+}
+
+func (s *Server) removeInProgress(token string) {
 	s.inProgressMu.Lock()
 	delete(s.inProgress, token)
 	s.inProgressMu.Unlock()