internal/jsonrpc2: Add Close method to Stream.

Also switched the internals of the stream implementations to using
net.Conn to enable asynchronous closing, not yet exposed int the API.

Change-Id: I57f1c36e7a46729d24f4339ba2fecc3f868e823f
Reviewed-on: https://go-review.googlesource.com/c/tools/+/231698
Run-TryBot: Ian Cottrell <iancottrell@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Jonathan Amsterdam <jba@google.com>

internal/jsonrpc2/jsonrpc2.go

@@ -28,6 +28,7 @@ const (
 // Conn is bidirectional; it does not have a designated server or client end.
 type Conn struct {
 	seq       int64      // must only be accessed using atomic operations
+	writeMu   sync.Mutex // protects writes to the stream
 	stream    Stream
 	pendingMu sync.Mutex // protects the pending map
 	pending   map[ID]chan *Response
@@ -65,7 +66,7 @@ func (c *Conn) Notify(ctx context.Context, method string, params interface{}) (e
 	}()
 
 	event.Metric(ctx, tag.Started.Of(1))
-	n, err := c.stream.Write(ctx, notify)
+	n, err := c.write(ctx, notify)
 	event.Metric(ctx, tag.SentBytes.Of(n))
 	return err
 }
@@ -104,7 +105,7 @@ func (c *Conn) Call(ctx context.Context, method string, params, result interface
 		c.pendingMu.Unlock()
 	}()
 	// now we are ready to send
-	n, err := c.stream.Write(ctx, call)
+	n, err := c.write(ctx, call)
 	event.Metric(ctx, tag.SentBytes.Of(n))
 	if err != nil {
 		// sending failed, we will never get a response, so don't leave it pending
@@ -144,7 +145,7 @@ func replier(conn *Conn, req Request, spanDone func()) Replier {
 		if err != nil {
 			return err
 		}
-		n, err := conn.stream.Write(ctx, response)
+		n, err := conn.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
@@ -155,6 +156,12 @@ func replier(conn *Conn, req Request, spanDone func()) Replier {
 	}
 }
 
+func (c *Conn) write(ctx context.Context, msg Message) (int64, error) {
+	c.writeMu.Lock()
+	defer c.writeMu.Unlock()
+	return c.stream.Write(ctx, msg)
+}
+
 // Run blocks until the connection is terminated, and returns any error that
 // caused the termination.
 // It must be called exactly once for each Conn.

internal/jsonrpc2/jsonrpc2_test.go

@@ -119,11 +119,7 @@ func run(ctx context.Context, t *testing.T, withHeaders bool, r io.ReadCloser, w
 	wg.Add(1)
 	go func() {
 		defer func() {
-			// this will happen when Run returns, which means at least one of the
-			// streams has already been closed
-			// we close both streams anyway, this may be redundant but is safe
-			r.Close()
-			w.Close()
+			stream.Close()
 			// and then signal that this connection is done
 			wg.Done()
 		}()

internal/jsonrpc2/serve.go

@@ -98,6 +98,7 @@ func Serve(ctx context.Context, ln net.Listener, server StreamServer, idleTimeou
 			stream := NewHeaderStream(netConn, netConn)
 			go func() {
 				closedConns <- server.ServeStream(ctx, stream)
+				stream.Close()
 			}()
 		case err := <-doneListening:
 			return err

internal/jsonrpc2/stream.go

@@ -10,38 +10,43 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"net"
 	"strconv"
 	"strings"
-	"sync"
+
+	"golang.org/x/tools/internal/fakenet"
 )
 
 // Stream abstracts the transport mechanics from the JSON RPC protocol.
 // A Conn reads and writes messages using the stream it was provided on
 // construction, and assumes that each call to Read or Write fully transfers
 // a single message, or returns an error.
+// A stream is not safe for concurrent use, it is expected it will be used by
+// a single Conn in a safe manner.
 type Stream interface {
 	// Read gets the next message from the stream.
-	// It is never called concurrently.
 	Read(context.Context) (Message, int64, error)
 	// Write sends a message to the stream.
-	// It must be safe for concurrent use.
 	Write(context.Context, Message) (int64, error)
+	// Close closes the connection.
+	// Any blocked Read or Write operations will be unblocked and return errors.
+	Close() error
 }
 
 // NewRawStream returns a Stream built on top of an io.Reader and io.Writer.
 // The messages are sent with no wrapping, and rely on json decode consistency
 // to determine message boundaries.
-func NewRawStream(in io.Reader, out io.Writer) Stream {
+func NewRawStream(in io.ReadCloser, out io.WriteCloser) Stream {
+	conn := fakenet.NewConn("jsonrpc2.NewRawStream", in, out)
 	return &rawStream{
-		in:  json.NewDecoder(in),
-		out: out,
+		conn: conn,
+		in:   json.NewDecoder(conn),
 	}
 }
 
 type rawStream struct {
+	conn net.Conn
 	in   *json.Decoder
-	outMu sync.Mutex
-	out   io.Writer
 }
 
 func (s *rawStream) Read(ctx context.Context) (Message, int64, error) {
@@ -68,26 +73,28 @@ func (s *rawStream) Write(ctx context.Context, msg Message) (int64, error) {
 	if err != nil {
 		return 0, fmt.Errorf("marshaling message: %v", err)
 	}
-	s.outMu.Lock()
-	n, err := s.out.Write(data)
-	s.outMu.Unlock()
+	n, err := s.conn.Write(data)
 	return int64(n), err
 }
 
+func (s *rawStream) Close() error {
+	return s.conn.Close()
+}
+
 // NewHeaderStream returns a Stream built on top of an io.Reader and io.Writer.
 // The messages are sent with HTTP content length and MIME type headers.
 // This is the format used by LSP and others.
-func NewHeaderStream(in io.Reader, out io.Writer) Stream {
+func NewHeaderStream(in io.ReadCloser, out io.WriteCloser) Stream {
+	conn := fakenet.NewConn("jsonrpc2.NewHeaderStream", in, out)
 	return &headerStream{
-		in:  bufio.NewReader(in),
-		out: out,
+		conn: conn,
+		in:   bufio.NewReader(conn),
 	}
 }
 
 type headerStream struct {
+	conn net.Conn
 	in   *bufio.Reader
-	outMu sync.Mutex
-	out   io.Writer
 }
 
 func (s *headerStream) Read(ctx context.Context) (Message, int64, error) {
@@ -148,13 +155,15 @@ func (s *headerStream) Write(ctx context.Context, msg Message) (int64, error) {
 	if err != nil {
 		return 0, fmt.Errorf("marshaling message: %v", err)
 	}
-	s.outMu.Lock()
-	defer s.outMu.Unlock()
-	n, err := fmt.Fprintf(s.out, "Content-Length: %v\r\n\r\n", len(data))
+	n, err := fmt.Fprintf(s.conn, "Content-Length: %v\r\n\r\n", len(data))
 	total := int64(n)
 	if err == nil {
-		n, err = s.out.Write(data)
+		n, err = s.conn.Write(data)
 		total += int64(n)
 	}
 	return total, err
 }
+
+func (s *headerStream) Close() error {
+	return s.conn.Close()
+}

internal/lsp/protocol/log.go

@@ -36,6 +36,10 @@ func (s *loggingStream) Write(ctx context.Context, msg jsonrpc2.Message) (int64,
 	return count, err
 }
 
+func (s *loggingStream) Close() error {
+	return s.stream.Close()
+}
+
 type req struct {
 	method string
 	start  time.Time