Refactor LSP logging (#632)

* Refactor LSP logging

While working on #630, I needed
to be able to see requests, responses and notifications to the server
from the client and from the server to the client.

I learned that there is functionality in the jsonrpc2 lib that we use to
do this nicely. This PR implements a new logging configuration to make
use of this using an adapted version of the LogMessage function that is
found in the library and using the OnRecv and OnSend functions.

Because of this, I have removed the logging responsibility from the
server and moved it to the connection. The server now should log errors
only and that should be the default mode of operation.

Enabling verbose logging will show you all messages being sent on the
connection. ExcludeMethods and IncludeMethods can be used to filter the
stream of messages when you are interested in monitoring or ignoring
certain flows of messages. I know that this will be useful for me and I
hope that this is a sensible enough way to do this overall.

Signed-off-by: Charlie Egan <charlie@styra.com>

* Inline logging config

Signed-off-by: Charlie Egan <charlie@styra.com>

* Remove VerboseLogging server option for good

Also use an io.Writer for errors

Signed-off-by: Charlie Egan <charlie@styra.com>

---------

Signed-off-by: Charlie Egan <charlie@styra.com>

cmd/languageserver.go

@@ -7,7 +7,6 @@ import (
 	"os/signal"
 	"syscall"
 
-	"github.com/sourcegraph/jsonrpc2"
 	"github.com/spf13/cobra"
 
 	"github.com/styrainc/regal/internal/lsp"
@@ -26,19 +25,21 @@ func init() {
 			defer cancel()
 
 			opts := &lsp.LanguageServerOptions{
-				ErrorLog:       os.Stderr,
-				VerboseLogging: verboseLogging,
+				ErrorLog: os.Stderr,
 			}
 
 			ls := lsp.NewLanguageServer(opts)
 
-			conn := jsonrpc2.NewConn(
-				ctx,
-				jsonrpc2.NewBufferedStream(lsp.StdOutReadWriteCloser{}, jsonrpc2.VSCodeObjectCodec{}),
-				jsonrpc2.HandlerWithError(ls.Handle),
-			)
+			conn := lsp.NewConnectionFromLanguageServer(ctx, ls.Handle, &lsp.ConnectionOptions{
+				LoggingConfig: lsp.ConnectionLoggingConfig{
+					Writer:      os.Stderr,
+					LogInbound:  verboseLogging,
+					LogOutbound: verboseLogging,
+				},
+			})
 
 			ls.SetConn(conn)
+
 			go ls.StartDiagnosticsWorker(ctx)
 			go ls.StartHoverWorker(ctx)
 			go ls.StartCommandWorker(ctx)

internal/lsp/connection.go

@@ -0,0 +1,224 @@
+// The implementation of logMessages, is a heavily modified version of the original implementation
+// in https://github.com/sourcegraph/jsonrpc2
+// The original license for that code is as follows:
+// Copyright (c) 2016 Sourcegraph Inc
+//
+// # MIT License
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+// documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
+// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
+// permit persons to whom the Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all copies or substantial portions
+// of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
+// THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+// TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+package lsp
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"sync"
+
+	"github.com/sourcegraph/jsonrpc2"
+)
+
+type ConnectionOptions struct {
+	LoggingConfig ConnectionLoggingConfig
+}
+
+type ConnectionLoggingConfig struct {
+	Writer io.Writer
+
+	LogInbound  bool
+	LogOutbound bool
+
+	// IncludeMethods is a list of methods to include in the request log.
+	// If empty, all methods are included. IncludeMethods takes precedence
+	// over ExcludeMethods.
+	IncludeMethods []string
+	// ExcludeMethods is a list of methods to exclude from the request log.
+	ExcludeMethods []string
+}
+
+func (cfg *ConnectionLoggingConfig) ShouldLog(method string) bool {
+	if len(cfg.IncludeMethods) > 0 {
+		for _, m := range cfg.IncludeMethods {
+			if m == method {
+				return true
+			}
+		}
+
+		return false
+	}
+
+	for _, m := range cfg.ExcludeMethods {
+		if m == method {
+			return false
+		}
+	}
+
+	return true
+}
+
+type ConnectionHandlerFunc func(context.Context, *jsonrpc2.Conn, *jsonrpc2.Request) (result interface{}, err error)
+
+type connectionLogger struct {
+	writer io.Writer
+}
+
+func (c *connectionLogger) Printf(format string, v ...interface{}) {
+	fmt.Fprintf(c.writer, format, v...)
+}
+
+func NewConnectionFromLanguageServer(
+	ctx context.Context,
+	handler ConnectionHandlerFunc,
+	opts *ConnectionOptions,
+) *jsonrpc2.Conn {
+	return jsonrpc2.NewConn(
+		ctx,
+		jsonrpc2.NewBufferedStream(StdOutReadWriteCloser{}, jsonrpc2.VSCodeObjectCodec{}),
+		jsonrpc2.HandlerWithError(handler),
+		logMessages(opts.LoggingConfig),
+	)
+}
+
+func logMessages(cfg ConnectionLoggingConfig) jsonrpc2.ConnOpt {
+	logger := &connectionLogger{writer: cfg.Writer}
+
+	return func(c *jsonrpc2.Conn) {
+		// Remember reqs we have received so that we can helpfully show the
+		// request method in OnSend for responses.
+		var (
+			mu         sync.Mutex
+			reqMethods = map[jsonrpc2.ID]string{}
+		)
+
+		setMethod := func(id jsonrpc2.ID, method string) {
+			mu.Lock()
+			defer mu.Unlock()
+
+			reqMethods[id] = method
+		}
+
+		getMethod := func(id jsonrpc2.ID) string {
+			mu.Lock()
+			defer mu.Unlock()
+
+			return reqMethods[id]
+		}
+
+		deleteMethod := func(id jsonrpc2.ID) {
+			mu.Lock()
+			defer mu.Unlock()
+
+			delete(reqMethods, id)
+		}
+
+		if cfg.LogInbound {
+			jsonrpc2.OnRecv(buildRecvHandler(setMethod, logger, cfg))(c)
+		}
+
+		if cfg.LogOutbound {
+			jsonrpc2.OnSend(buildSendHandler(getMethod, deleteMethod, logger, cfg))(c)
+		}
+	}
+}
+
+func buildRecvHandler(
+	setMethod func(jsonrpc2.ID, string),
+	logger *connectionLogger,
+	cfg ConnectionLoggingConfig,
+) func(req *jsonrpc2.Request, resp *jsonrpc2.Response) {
+	return func(req *jsonrpc2.Request, resp *jsonrpc2.Response) {
+		switch {
+		case req != nil && resp == nil:
+			setMethod(req.ID, req.Method)
+
+			if !cfg.ShouldLog(req.Method) {
+				return
+			}
+
+			params, _ := json.Marshal(req.Params)
+			if req.Notif {
+				logger.Printf("--> notif: %s: %s\n", req.Method, params)
+			} else {
+				logger.Printf("--> request #%s: %s: %s\n", req.ID, req.Method, params)
+			}
+
+		case resp != nil:
+			var method string
+			if req != nil {
+				method = req.Method
+			} else {
+				method = "(no matching request)"
+			}
+
+			if !cfg.ShouldLog(method) {
+				return
+			}
+
+			switch {
+			case resp.Result != nil:
+				result, _ := json.Marshal(resp.Result)
+				logger.Printf("--> response #%s: %s: %s\n", resp.ID, method, result)
+			case resp.Error != nil:
+				errBs, _ := json.Marshal(resp.Error)
+				logger.Printf("--> response error #%s: %s: %s\n", resp.ID, method, errBs)
+			}
+		}
+	}
+}
+
+func buildSendHandler(
+	getMethod func(jsonrpc2.ID) string,
+	deleteMethod func(jsonrpc2.ID),
+	logger *connectionLogger,
+	cfg ConnectionLoggingConfig,
+) func(req *jsonrpc2.Request, resp *jsonrpc2.Response) {
+	return func(req *jsonrpc2.Request, resp *jsonrpc2.Response) {
+		switch {
+		case req != nil && resp == nil:
+			if !cfg.ShouldLog(req.Method) {
+				return
+			}
+
+			params, _ := json.Marshal(req.Params)
+			if req.Notif {
+				logger.Printf("<-- notif: %s: %s\n", req.Method, params)
+			} else {
+				logger.Printf("<-- request #%s: %s: %s\n", req.ID, req.Method, params)
+			}
+
+		case resp != nil:
+			method := getMethod(resp.ID)
+
+			deleteMethod(resp.ID)
+
+			if method == "" {
+				method = "(no previous request)"
+			}
+
+			if !cfg.ShouldLog(method) {
+				return
+			}
+
+			if resp.Result != nil {
+				result, _ := json.Marshal(resp.Result)
+				logger.Printf("<-- response #%s: %s: %s\n", resp.ID, method, result)
+			} else {
+				errBs, _ := json.Marshal(resp.Error)
+				logger.Printf("<-- response error #%s: %s: %s\n", resp.ID, method, errBs)
+			}
+		}
+	}
+}