internal/lsp/protocol: send responses for cancelled requests

LSP https://microsoft.github.io/language-server-protocol/specifications/specification-current/#cancelRequest
expects the server to send back the response even when the request is cancelled.

Gopls LSP protocol implements the cancellation using the context cancellation.
That is, upon a cancellation request from the client, the server calls the
corresponding canceller that cancels the context passed to the handler.
Reusing this cancelled context for the replyer is not safe because code in any layer
can decide to shortcircuit and skip sending the data back to the client.
E.g. https://cs.opensource.google/go/tools/+/master:internal/jsonrpc2/stream.go;l=63

This CL make sure to pass the detached context to the replier.

Alternative, or a better approach to avoid any unexpected side-effect of
using a detached context is to send out the response at the point of cancellation
with a separate context. But that requires more significant code change.

Testing is currently hard, but maybe doable once the current refactoring is
done. Test is left as a TODO.

Change-Id: I4611af00ad913e96b62c6b7180c6673b0465daf9
Reviewed-on: https://go-review.googlesource.com/c/tools/+/232300
Reviewed-by: Ian Cottrell <[email protected]>

Author: hyangah

internal/lsp/protocol/protocol.go

@@ -33,16 +33,29 @@ func ServerDispatcher(conn *jsonrpc2.Conn) Server {
 
 func Handlers(handler jsonrpc2.Handler) jsonrpc2.Handler {
 	return CancelHandler(
-		CancelHandler(
-			jsonrpc2.AsyncHandler(
-				jsonrpc2.MustReplyHandler(handler))))
+		jsonrpc2.AsyncHandler(
+			jsonrpc2.MustReplyHandler(handler)))
 }
 
 func CancelHandler(handler jsonrpc2.Handler) jsonrpc2.Handler {
 	handler, canceller := jsonrpc2.CancelHandler(handler)
 	return func(ctx context.Context, reply jsonrpc2.Replier, req jsonrpc2.Request) error {
 		if req.Method() != "$/cancelRequest" {
-			return handler(ctx, reply, req)
+			// TODO(iancottrell): See if we can generate a reply for the request to be cancelled
+			// at the point of cancellation rather than waiting for gopls to naturally reply.
+			// To do that, we need to keep track of whether a reply has been sent already and
+			// be careful about racing between the two paths.
+			// TODO(iancottrell): Add a test that watches the stream and verifies the response
+			// for the cancelled request flows.
+			replyWithDetachedContext := func(ctx context.Context, resp interface{}, err error) error {
+				// https://microsoft.github.io/language-server-protocol/specifications/specification-current/#cancelRequest
+				if ctx.Err() != nil && err == nil {
+					err = RequestCancelledError
+				}
+				ctx = xcontext.Detach(ctx)
+				return reply(ctx, resp, err)
+			}
+			return handler(ctx, replyWithDetachedContext, req)
 		}
 		var params CancelParams
 		if err := json.Unmarshal(req.Params(), &params); err != nil {