From 04327e3343571072898a2f775c5b608f69d7986d Mon Sep 17 00:00:00 2001 From: jansenjiang Date: Fri, 17 Jul 2026 16:54:54 +0800 Subject: [PATCH] internal/jsonrpc2: plumb cancellation cause through request context handleAsync reported why a request was canceled by inspecting the connection-level s.writeErr, guarded by a TODO awaiting golang/go#51365. That API shipped as context.WithCancelCause/context.Cause in Go 1.21, and this module targets Go 1.25. Carry the cancellation cause on each request's context instead of inferring it from global state. Both the write-failure and read-side teardowns cancel each in-flight request with a cause wrapping ErrServerClosing; the remaining cancel sites (peer cancel, normal completion) pass nil. handleAsync now reads context.Cause(req.ctx) instead of guessing from s.writeErr. This also fixes a latent mis-attribution: s.writeErr is connection-level, so a request canceled for an unrelated reason was reported as ErrServerClosing whenever s.writeErr happened to be non-nil. ctx.Err() still returns context.Canceled in all cases, so external behavior for callers inspecting ctx.Err() is unchanged. --- internal/jsonrpc2/conn.go | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) 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)) } } })