conn: use goccy/go-json

Author: zchee

conn.go

@@ -1,47 +1,54 @@
-// SPDX-License-Identifier: BSD-3-Clause
 // SPDX-FileCopyrightText: Copyright 2021 The Go Language Server Authors
+// SPDX-License-Identifier: BSD-3-Clause
 
 package jsonrpc2
 
 import (
-	"bytes"
 	"context"
-	"encoding/json"
 	"fmt"
 	"sync"
 	"sync/atomic"
 
+	json "github.com/goccy/go-json"
+
 	"go.lsp.dev/pkg/event"
 	"go.lsp.dev/pkg/event/label"
 	"go.lsp.dev/pkg/event/tag"
 )
 
 // Conn is the common interface to jsonrpc clients and servers.
+//
 // Conn is bidirectional; it does not have a designated server or client end.
 // It manages the jsonrpc2 protocol, connecting responses back to their calls.
 type Conn interface {
 	// Call invokes the target method and waits for a response.
+	//
 	// The params will be marshaled to JSON before sending over the wire, and will
 	// be handed to the method invoked.
+	//
 	// The response will be unmarshaled from JSON into the result.
+	//
 	// The id returned will be unique from this connection, and can be used for
 	// logging or tracking.
 	Call(ctx context.Context, method string, params, result interface{}) (ID, error)
 
 	// Notify invokes the target method but does not wait for a response.
+	//
 	// The params will be marshaled to JSON before sending over the wire, and will
 	// be handed to the method invoked.
 	Notify(ctx context.Context, method string, params interface{}) error
 
 	// Go starts a goroutine to handle the connection.
-	// It must be called exactly once for each Conn.
-	// It returns immediately.
-	// You must block on Done() to wait for the connection to shut down.
+	//
+	// It must be called exactly once for each Conn. It returns immediately.
+	// Must block on Done() to wait for the connection to shut down.
+	//
 	// This is a temporary measure, this should be started automatically in the
 	// future.
 	Go(ctx context.Context, handler Handler)
 
 	// Close closes the connection and it's underlying stream.
+	//
 	// It does not wait for the close to complete, use the Done() channel for
 	// that.
 	Close() error
@@ -52,19 +59,20 @@ type Conn interface {
 	Done() <-chan struct{}
 
 	// Err returns an error if there was one from within the processing goroutine.
+	//
 	// If err returns non nil, the connection will be already closed or closing.
 	Err() error
 }
 
 type conn struct {
-	seq       int64      // access atomically
-	writeMu   sync.Mutex // protects writes to the stream
-	stream    Stream
-	pendingMu sync.Mutex // protects the pending map
-	pending   map[ID]chan *Response
-
-	done chan struct{}
-	err  atomic.Value
+	seq       int64                 // access atomically
+	writeMu   sync.Mutex            // protects writes to the stream
+	stream    Stream                // supplied stream
+	pendingMu sync.Mutex            // protects the pending map
+	pending   map[ID]chan *Response // holds the pending response channel with the ID as the key.
+
+	done chan struct{} // closed when done
+	err  atomic.Value  // holds run error
 }
 
 // NewConn creates a new connection object around the supplied stream.
@@ -78,9 +86,9 @@ func NewConn(s Stream) Conn {
 }
 
 // Call implements Conn.
-func (c *conn) Call(ctx context.Context, method string, params, result interface{}) (_ ID, err error) {
+func (c *conn) Call(ctx context.Context, method string, params, result interface{}) (id ID, err error) {
 	// generate a new request identifier
-	id := ID{number: atomic.AddInt64(&c.seq, 1)}
+	id = NewNumberID(atomic.AddInt64(&c.seq, 1))
 	call, err := NewCall(id, method, params)
 	if err != nil {
 		return id, fmt.Errorf("marshaling call parameters: %w", err)
@@ -123,22 +131,22 @@ func (c *conn) Call(ctx context.Context, method string, params, result interface
 
 	// now wait for the response
 	select {
-	case response := <-rchan:
-		switch {
-		case response.err != nil: // is it an error response?
-			return id, response.err
+	case resp := <-rchan:
+		// is it an error response?
+		if resp.err != nil {
+			return id, resp.err
+		}
 
-		case result == nil || len(response.result) == 0:
+		if result == nil || len(resp.result) == 0 {
 			return id, nil
+		}
 
-		default:
-			dec := json.NewDecoder(bytes.NewReader(response.result))
-			if err := dec.Decode(result); err != nil {
-				return id, fmt.Errorf("unmarshaling result: %w", err)
-			}
-			return id, nil
+		if err := json.Unmarshal(resp.result, result); err != nil {
+			return id, fmt.Errorf("unmarshaling result: %w", err)
 		}
 
+		return id, nil
+
 	case <-ctx.Done():
 		return id, ctx.Err()
 	}
@@ -150,6 +158,7 @@ func (c *conn) Notify(ctx context.Context, method string, params interface{}) (e
 	if err != nil {
 		return fmt.Errorf("marshaling notify parameters: %w", err)
 	}
+
 	ctx, done := event.Start(ctx, method,
 		tag.Method.Of(method),
 		tag.RPCDirection.Of(tag.Outbound),
@@ -162,6 +171,7 @@ func (c *conn) Notify(ctx context.Context, method string, params interface{}) (e
 	event.Metric(ctx, tag.Started.Of(1))
 	n, err := c.write(ctx, notify)
 	event.Metric(ctx, tag.SentBytes.Of(n))
+
 	return err
 }
 
@@ -176,25 +186,31 @@ func (c *conn) replier(req Message, spanDone func()) Replier {
 			// request was a notify, no need to respond
 			return nil
 		}
+
 		response, err := NewResponse(call.id, result, err)
 		if err != nil {
 			return err
 		}
+
 		n, err := c.write(ctx, response)
 		event.Metric(ctx, tag.SentBytes.Of(n))
 		if err != nil {
-			// TODO(iancottrell): if a stream write fails, we really need to shut down
-			// the whole stream
+			// TODO(iancottrell): if a stream write fails, we really need to shut down the whole stream
 			return err
 		}
 		return nil
 	}
 }
 
-func (c *conn) write(ctx context.Context, msg Message) (int64, error) {
+func (c *conn) write(ctx context.Context, msg Message) (n int64, err error) {
 	c.writeMu.Lock()
-	defer c.writeMu.Unlock()
-	return c.stream.Write(ctx, msg)
+	n, err = c.stream.Write(ctx, msg)
+	c.writeMu.Unlock()
+	if err != nil {
+		return 0, fmt.Errorf("write to stream: %w", err)
+	}
+
+	return n, nil
 }
 
 // Go implements Conn.
@@ -204,6 +220,7 @@ func (c *conn) Go(ctx context.Context, handler Handler) {
 
 func (c *conn) run(ctx context.Context, handler Handler) {
 	defer close(c.done)
+
 	for {
 		// get the next message
 		msg, n, err := c.stream.Read(ctx)
@@ -212,6 +229,7 @@ func (c *conn) run(ctx context.Context, handler Handler) {
 			c.fail(err)
 			return
 		}
+
 		switch msg := msg.(type) {
 		case Request:
 			labels := []label.Label{
@@ -224,14 +242,18 @@ func (c *conn) run(ctx context.Context, handler Handler) {
 			} else {
 				labels = labels[:len(labels)-1]
 			}
+
 			reqCtx, spanDone := event.Start(ctx, msg.Method(), labels...)
 			event.Metric(reqCtx,
 				tag.Started.Of(1),
-				tag.ReceivedBytes.Of(n))
+				tag.ReceivedBytes.Of(n),
+			)
+
 			if err := handler(reqCtx, c.replier(msg, spanDone), msg); err != nil {
 				// delivery failed, not much we can do
 				event.Error(reqCtx, "jsonrpc2 message delivery failed", err)
 			}
+
 		case *Response:
 			// If method is not set, this should be a response, in which case we must
 			// have an id to send the response back to the caller.
@@ -269,6 +291,7 @@ func (c *conn) fail(err error) {
 	c.stream.Close()
 }
 
+// recordStatus records the status code based on the error.
 func recordStatus(ctx context.Context, err error) {
 	if err != nil {
 		event.Label(ctx, tag.StatusCode.Of("ERROR"))