internal/jsonrpc2: allow asynchronous writes

This CL introduces a fundamental change to the jsonrpc2 library:
connection writes, which were previously serialized by the jsonrpc2
library itself, are now allowed to be concurrent.

The change in semantics of the jsonrpc2 library should hopefully be easy
to review, since moving the synchronization to the Writer implementation
is equivalent to the previous logic. However, this change is critical
for the streamable client transport, because it allows for concurrent
http requests to the server.

Consider that a write is a POST to the server, and we don't know that
write succeeded until we get the response header. Previously, we had the
following problem: if the client POSTs a request, and the server blocks
its response on a request made through the hanging GET, the client was
unable to respond because the initial POST is still blocked. We could
update our streamable server transport to force a flush of the response
headers, but we can't guarantee that other servers behave the same way.
Fundamentally, writes in the spec are asynchronous, and we need to
support that.

internal/jsonrpc2/conn.go

@@ -65,8 +65,7 @@ type Connection struct {
 	state   inFlightState // accessed only in updateInFlight
 	done    chan struct{} // closed (under stateMu) when state.closed is true and all goroutines have completed
 
-	writer chan Writer // 1-buffered; stores the writer when not in use
-
+	writer  Writer
 	handler Handler
 
 	onInternalError func(error)
@@ -214,12 +213,11 @@ func NewConnection(ctx context.Context, cfg ConnectionConfig) *Connection {
 	c := &Connection{
 		state:           inFlightState{closer: cfg.Closer},
 		done:            make(chan struct{}),
-		writer:          make(chan Writer, 1),
+		writer:          cfg.Writer,
 		onDone:          cfg.OnDone,
 		onInternalError: cfg.OnInternalError,
 	}
 	c.handler = cfg.Bind(c)
-	c.writer <- cfg.Writer
 	c.start(ctx, cfg.Reader, cfg.Preempter)
 	return c
 }
@@ -239,7 +237,6 @@ func bindConnection(bindCtx context.Context, rwc io.ReadWriteCloser, binder Bind
 	c := &Connection{
 		state:  inFlightState{closer: rwc},
 		done:   make(chan struct{}),
-		writer: make(chan Writer, 1),
 		onDone: onDone,
 	}
 	// It's tempting to set a finalizer on c to verify that the state has gone
@@ -259,7 +256,7 @@ func bindConnection(bindCtx context.Context, rwc io.ReadWriteCloser, binder Bind
 	}
 	c.onInternalError = options.OnInternalError
 
-	c.writer <- framer.Writer(rwc)
+	c.writer = framer.Writer(rwc)
 	reader := framer.Reader(rwc)
 	c.start(ctx, reader, options.Preempter)
 	return c
@@ -728,9 +725,7 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e
 // write is used by all things that write outgoing messages, including replies.
 // it makes sure that writes are atomic
 func (c *Connection) write(ctx context.Context, msg Message) error {
-	writer := <-c.writer
-	defer func() { c.writer <- writer }()
-	err := writer.Write(ctx, msg)
+	err := c.writer.Write(ctx, msg)
 
 	if err != nil && ctx.Err() == nil {
 		// The call to Write failed, and since ctx.Err() is nil we can't attribute

internal/jsonrpc2/frame.go

@@ -12,12 +12,14 @@ import (
 	"io"
 	"strconv"
 	"strings"
+	"sync"
 )
 
 // Reader abstracts the transport mechanics from the JSON RPC protocol.
 // A Conn reads messages from the reader it was provided on construction,
 // and assumes that each call to Read fully transfers a single message,
 // or returns an error.
+//
 // A reader is not safe for concurrent use, it is expected it will be used by
 // a single Conn in a safe manner.
 type Reader interface {
@@ -29,8 +31,9 @@ type Reader interface {
 // A Conn writes messages using the writer it was provided on construction,
 // and assumes that each call to Write fully transfers a single message,
 // or returns an error.
-// A writer is not safe for concurrent use, it is expected it will be used by
-// a single Conn in a safe manner.
+//
+// A writer must be safe for concurrent use, as writes may occur concurrently
+// in practice: libraries may make calls or respond to requests asynchronously.
 type Writer interface {
 	// Write sends a message to the stream.
 	Write(context.Context, Message) error
@@ -62,7 +65,10 @@ func RawFramer() Framer { return rawFramer{} }
 
 type rawFramer struct{}
 type rawReader struct{ in *json.Decoder }
-type rawWriter struct{ out io.Writer }
+type rawWriter struct {
+	mu  sync.Mutex
+	out io.Writer
+}
 
 func (rawFramer) Reader(rw io.Reader) Reader {
 	return &rawReader{in: json.NewDecoder(rw)}
@@ -92,10 +98,14 @@ func (w *rawWriter) Write(ctx context.Context, msg Message) error {
 		return ctx.Err()
 	default:
 	}
+
 	data, err := EncodeMessage(msg)
 	if err != nil {
 		return fmt.Errorf("marshaling message: %v", err)
 	}
+
+	w.mu.Lock()
+	defer w.mu.Unlock()
 	_, err = w.out.Write(data)
 	return err
 }
@@ -107,7 +117,10 @@ func HeaderFramer() Framer { return headerFramer{} }
 
 type headerFramer struct{}
 type headerReader struct{ in *bufio.Reader }
-type headerWriter struct{ out io.Writer }
+type headerWriter struct {
+	mu  sync.Mutex
+	out io.Writer
+}
 
 func (headerFramer) Reader(rw io.Reader) Reader {
 	return &headerReader{in: bufio.NewReader(rw)}
@@ -180,6 +193,9 @@ func (w *headerWriter) Write(ctx context.Context, msg Message) error {
 		return ctx.Err()
 	default:
 	}
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
 	data, err := EncodeMessage(msg)
 	if err != nil {
 		return fmt.Errorf("marshaling message: %v", err)