internal/jsonrpc2_v2: error out in-flight client calls when the reader breaks

Otherwise, the Await method on the corresponding AsyncCall will never
unblock, leading to a deadlock (detected by the test changes in
CL 388597).

For golang/go#46047
For golang/go#46520
For golang/go#49387

Change-Id: I7954f80786059772ddd7f8c98d8752d56d074919
Reviewed-on: https://go-review.googlesource.com/c/tools/+/388775
Run-TryBot: Bryan Mills <bcmills@google.com>
Reviewed-by: Ian Cottrell <iancottrell@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
Auto-Submit: Bryan Mills <bcmills@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>

2 files changed, 42 insertions, 11 deletions

diff --git a/internal/jsonrpc2_v2/conn.go b/internal/jsonrpc2_v2/conn.go
index 15d40b4a4..0215c67af 100644
--- a/internal/jsonrpc2_v2/conn.go
+++ b/internal/jsonrpc2_v2/conn.go
@@ -175,9 +175,25 @@ func (c *Connection) Call(ctx context.Context, method string, params interface{}
 	// are racing the response.
 	// rchan is buffered in case the response arrives without a listener.
 	result.response = make(chan *Response, 1)
-	pending := <-c.outgoingBox
-	pending[result.id] = result.response
-	c.outgoingBox <- pending
+	outgoing, ok := <-c.outgoingBox
+	if !ok {
+		// If the call failed due to (say) an I/O error or broken pipe, attribute it
+		// as such. (If the error is nil, then the connection must have been shut
+		// down cleanly.)
+		err := c.async.wait()
+		if err == nil {
+			err = ErrClientClosing
+		}
+
+		resp, respErr := NewResponse(result.id, nil, err)
+		if respErr != nil {
+			panic(fmt.Errorf("unexpected error from NewResponse: %w", respErr))
+		}
+		result.response <- resp
+		return result
+	}
+	outgoing[result.id] = result.response
+	c.outgoingBox <- outgoing
 	// now we are ready to send
 	if err := c.write(ctx, call); err != nil {
 		// sending failed, we will never get a response, so deliver a fake one
@@ -290,8 +306,19 @@ func (c *Connection) Close() error {
 
 // readIncoming collects inbound messages from the reader and delivers them, either responding
 // to outgoing calls or feeding requests to the queue.
-func (c *Connection) readIncoming(ctx context.Context, reader Reader, toQueue chan<- *incoming) {
-	defer close(toQueue)
+func (c *Connection) readIncoming(ctx context.Context, reader Reader, toQueue chan<- *incoming) (err error) {
+	defer func() {
+		// Retire any outgoing requests that were still in flight.
+		// With the Reader no longer being processed, they necessarily cannot receive a response.
+		outgoing := <-c.outgoingBox
+		close(c.outgoingBox) // Prevent new outgoing requests, which would deadlock.
+		for id, responseBox := range outgoing {
+			responseBox <- &Response{ID: id, Error: err}
+		}
+
+		close(toQueue)
+	}()
+
 	for {
 		// get the next message
 		// no lock is needed, this is the only reader
@@ -301,7 +328,7 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, toQueue ch
 			if !isClosingError(err) {
 				c.async.setError(err)
 			}
-			return
+			return err
 		}
 		switch msg := msg.(type) {
 		case *Request:
@@ -340,12 +367,12 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, toQueue ch
 }
 
 func (c *Connection) incomingResponse(msg *Response) {
-	pending := <-c.outgoingBox
-	response, ok := pending[msg.ID]
-	if ok {
-		delete(pending, msg.ID)
+	var response chan<- *Response
+	if outgoing, ok := <-c.outgoingBox; ok {
+		response = outgoing[msg.ID]
+		delete(outgoing, msg.ID)
+		c.outgoingBox <- outgoing
 	}
-	c.outgoingBox <- pending
 	if response != nil {
 		response <- msg
 	}
diff --git a/internal/jsonrpc2_v2/wire.go b/internal/jsonrpc2_v2/wire.go
index 4da129ae6..c99e497e3 100644
--- a/internal/jsonrpc2_v2/wire.go
+++ b/internal/jsonrpc2_v2/wire.go
@@ -33,6 +33,10 @@ var (
 	ErrServerOverloaded = NewError(-32000, "JSON RPC overloaded")
 	// ErrUnknown should be used for all non coded errors.
 	ErrUnknown = NewError(-32001, "JSON RPC unknown error")
+	// ErrServerClosing is returned for calls that arrive while the server is closing.
+	ErrServerClosing = NewError(-32002, "JSON RPC server is closing")
+	// ErrClientClosing is a dummy error returned for calls initiated while the client is closing.
+	ErrClientClosing = NewError(-32003, "JSON RPC client is closing")
 )
 
 const wireVersion = "2.0"