diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 4994c63b..97905045 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -172,7 +172,7 @@ func (s *inFlightState) shuttingDown(errClosing error) error { type incomingRequest struct { *Request // the request being processed ctx context.Context - cancel context.CancelFunc + cancel context.CancelCauseFunc } // Reader abstracts the transport mechanics from the JSON RPC protocol. @@ -458,7 +458,7 @@ func (c *Connection) Cancel(id ID) { req = s.incomingByID[id] }) if req != nil { - req.cancel() + req.cancel(nil) } } @@ -554,7 +554,7 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter // response either, so parked handlers have nothing useful left to do. // Mirrors the equivalent cleanup on write failure. for _, r := range s.incomingByID { - r.cancel() + r.cancel(fmt.Errorf("%w: %v", ErrServerClosing, err)) } }) } @@ -564,7 +564,7 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter func (c *Connection) acceptRequest(ctx context.Context, msg *Request, preempter Preempter) { // In theory notifications cannot be cancelled, but we build them a cancel // context anyway. - reqCtx, cancel := context.WithCancel(ctx) + reqCtx, cancel := context.WithCancelCause(ctx) req := &incomingRequest{ Request: msg, ctx: reqCtx, @@ -665,15 +665,14 @@ func (c *Connection) handleAsync() { return } - // Only deliver to the Handler if not already canceled. + // Only deliver to the Handler if not already canceled. If the request + // was canceled with a cause (e.g. a write failure cancels every + // in-flight request), report that cause rather than the bare + // context.Canceled. if err := req.ctx.Err(); err != nil { - c.updateInFlight(func(s *inFlightState) { - if s.writeErr != nil { - // Assume that req.ctx was canceled due to s.writeErr. - // TODO(#51365): use a Context API to plumb this through req.ctx. - err = fmt.Errorf("%w: %v", ErrServerClosing, s.writeErr) - } - }) + if cause := context.Cause(req.ctx); cause != nil { + err = cause + } c.processResult("handleAsync", req, nil, err) continue } @@ -734,7 +733,7 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e } // Cancel the request to free any associated resources. - req.cancel() + req.cancel(nil) c.updateInFlight(func(s *inFlightState) { if s.incoming == 0 { panic("jsonrpc2: processResult called when incoming count is already zero") @@ -777,7 +776,7 @@ func (c *Connection) write(ctx context.Context, msg Message) error { if s.writeErr == nil { s.writeErr = err for _, r := range s.incomingByID { - r.cancel() + r.cancel(fmt.Errorf("%w: %v", ErrServerClosing, err)) } } })