1// Copyright 2020 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package lsp
6
7import (
8	"context"
9	"math/rand"
10	"strconv"
11	"strings"
12	"sync"
13
14	"golang.org/x/tools/internal/event"
15	"golang.org/x/tools/internal/lsp/debug/tag"
16	"golang.org/x/tools/internal/lsp/protocol"
17	"golang.org/x/tools/internal/xcontext"
18	errors "golang.org/x/xerrors"
19)
20
21type progressTracker struct {
22	client                   protocol.Client
23	supportsWorkDoneProgress bool
24
25	mu         sync.Mutex
26	inProgress map[protocol.ProgressToken]*workDone
27}
28
29func newProgressTracker(client protocol.Client) *progressTracker {
30	return &progressTracker{
31		client:     client,
32		inProgress: make(map[protocol.ProgressToken]*workDone),
33	}
34}
35
36// start notifies the client of work being done on the server. It uses either
37// ShowMessage RPCs or $/progress messages, depending on the capabilities of
38// the client.  The returned WorkDone handle may be used to report incremental
39// progress, and to report work completion. In particular, it is an error to
40// call start and not call end(...) on the returned WorkDone handle.
41//
42// If token is empty, a token will be randomly generated.
43//
44// The progress item is considered cancellable if the given cancel func is
45// non-nil. In this case, cancel is called when the work done
46//
47// Example:
48//  func Generate(ctx) (err error) {
49//    ctx, cancel := context.WithCancel(ctx)
50//    defer cancel()
51//    work := s.progress.start(ctx, "generate", "running go generate", cancel)
52//    defer func() {
53//      if err != nil {
54//        work.end(ctx, fmt.Sprintf("generate failed: %v", err))
55//      } else {
56//        work.end(ctx, "done")
57//      }
58//    }()
59//    // Do the work...
60//  }
61//
62func (t *progressTracker) start(ctx context.Context, title, message string, token protocol.ProgressToken, cancel func()) *workDone {
63	wd := &workDone{
64		ctx:    xcontext.Detach(ctx),
65		client: t.client,
66		token:  token,
67		cancel: cancel,
68	}
69	if !t.supportsWorkDoneProgress {
70		// Previous iterations of this fallback attempted to retain cancellation
71		// support by using ShowMessageCommand with a 'Cancel' button, but this is
72		// not ideal as the 'Cancel' dialog stays open even after the command
73		// completes.
74		//
75		// Just show a simple message. Clients can implement workDone progress
76		// reporting to get cancellation support.
77		if err := wd.client.ShowMessage(wd.ctx, &protocol.ShowMessageParams{
78			Type:    protocol.Log,
79			Message: message,
80		}); err != nil {
81			event.Error(ctx, "showing start message for "+title, err)
82		}
83		return wd
84	}
85	if wd.token == nil {
86		token = strconv.FormatInt(rand.Int63(), 10)
87		err := wd.client.WorkDoneProgressCreate(ctx, &protocol.WorkDoneProgressCreateParams{
88			Token: token,
89		})
90		if err != nil {
91			wd.err = err
92			event.Error(ctx, "starting work for "+title, err)
93			return wd
94		}
95		wd.token = token
96	}
97	// At this point we have a token that the client knows about. Store the token
98	// before starting work.
99	t.mu.Lock()
100	t.inProgress[wd.token] = wd
101	t.mu.Unlock()
102	wd.cleanup = func() {
103		t.mu.Lock()
104		delete(t.inProgress, token)
105		t.mu.Unlock()
106	}
107	err := wd.client.Progress(ctx, &protocol.ProgressParams{
108		Token: wd.token,
109		Value: &protocol.WorkDoneProgressBegin{
110			Kind:        "begin",
111			Cancellable: wd.cancel != nil,
112			Message:     message,
113			Title:       title,
114		},
115	})
116	if err != nil {
117		event.Error(ctx, "generate progress begin", err)
118	}
119	return wd
120}
121
122func (t *progressTracker) cancel(ctx context.Context, token protocol.ProgressToken) error {
123	t.mu.Lock()
124	defer t.mu.Unlock()
125	wd, ok := t.inProgress[token]
126	if !ok {
127		return errors.Errorf("token %q not found in progress", token)
128	}
129	if wd.cancel == nil {
130		return errors.Errorf("work %q is not cancellable", token)
131	}
132	wd.doCancel()
133	return nil
134}
135
136// workDone represents a unit of work that is reported to the client via the
137// progress API.
138type workDone struct {
139	// ctx is detached, for sending $/progress updates.
140	ctx    context.Context
141	client protocol.Client
142	// If token is nil, this workDone object uses the ShowMessage API, rather
143	// than $/progress.
144	token protocol.ProgressToken
145	// err is set if progress reporting is broken for some reason (for example,
146	// if there was an initial error creating a token).
147	err error
148
149	cancelMu  sync.Mutex
150	cancelled bool
151	cancel    func()
152
153	cleanup func()
154}
155
156func (wd *workDone) doCancel() {
157	wd.cancelMu.Lock()
158	defer wd.cancelMu.Unlock()
159	if !wd.cancelled {
160		wd.cancel()
161	}
162}
163
164// report reports an update on WorkDone report back to the client.
165func (wd *workDone) report(message string, percentage float64) {
166	if wd == nil {
167		return
168	}
169	wd.cancelMu.Lock()
170	cancelled := wd.cancelled
171	wd.cancelMu.Unlock()
172	if cancelled {
173		return
174	}
175	if wd.err != nil || wd.token == nil {
176		// Not using the workDone API, so we do nothing. It would be far too spammy
177		// to send incremental messages.
178		return
179	}
180	message = strings.TrimSuffix(message, "\n")
181	err := wd.client.Progress(wd.ctx, &protocol.ProgressParams{
182		Token: wd.token,
183		Value: &protocol.WorkDoneProgressReport{
184			Kind: "report",
185			// Note that in the LSP spec, the value of Cancellable may be changed to
186			// control whether the cancel button in the UI is enabled. Since we don't
187			// yet use this feature, the value is kept constant here.
188			Cancellable: wd.cancel != nil,
189			Message:     message,
190			Percentage:  uint32(percentage),
191		},
192	})
193	if err != nil {
194		event.Error(wd.ctx, "reporting progress", err)
195	}
196}
197
198// end reports a workdone completion back to the client.
199func (wd *workDone) end(message string) {
200	if wd == nil {
201		return
202	}
203	var err error
204	switch {
205	case wd.err != nil:
206		// There is a prior error.
207	case wd.token == nil:
208		// We're falling back to message-based reporting.
209		err = wd.client.ShowMessage(wd.ctx, &protocol.ShowMessageParams{
210			Type:    protocol.Info,
211			Message: message,
212		})
213	default:
214		err = wd.client.Progress(wd.ctx, &protocol.ProgressParams{
215			Token: wd.token,
216			Value: &protocol.WorkDoneProgressEnd{
217				Kind:    "end",
218				Message: message,
219			},
220		})
221	}
222	if err != nil {
223		event.Error(wd.ctx, "ending work", err)
224	}
225	if wd.cleanup != nil {
226		wd.cleanup()
227	}
228}
229
230// eventWriter writes every incoming []byte to
231// event.Print with the operation=generate tag
232// to distinguish its logs from others.
233type eventWriter struct {
234	ctx       context.Context
235	operation string
236}
237
238func (ew *eventWriter) Write(p []byte) (n int, err error) {
239	event.Log(ew.ctx, string(p), tag.Operation.Of(ew.operation))
240	return len(p), nil
241}
242
243// workDoneWriter wraps a workDone handle to provide a Writer interface,
244// so that workDone reporting can more easily be hooked into commands.
245type workDoneWriter struct {
246	wd *workDone
247}
248
249func (wdw workDoneWriter) Write(p []byte) (n int, err error) {
250	wdw.wd.report(string(p), 0)
251	// Don't fail just because of a failure to report progress.
252	return len(p), nil
253}
254