drpcmanager: fix context cancellation error for unary RPCs

Fix a race condition where a unary RPC with a cancelled context could
return io.EOF instead of codes.Canceled. Two changes, mirroring how
gRPC handles this:

1. Early ctx.Err() check in NewClientStream before creating the stream.
2. Deferred stream.CheckCancelError in doInvoke to convert io.EOF to
   the cancel error if the stream was cancelled mid-operation.

The problem:

With multiplexing, each stream gets a manageStream goroutine that
watches ctx.Done() and calls SendCancel + Cancel when the context is
cancelled. This races with doInvoke, which writes invoke and message
frames through the same stream. The race has three outcomes depending
on who acquires the stream's write lock first:

  1. doInvoke wins the lock, completes all writes, and the RPC succeeds
     even though it should have been cancelled.
  2. SendCancel wins, sets send=io.EOF before doInvoke runs.
     rawWriteLocked sees send.IsSet() and returns io.EOF. Invoke's
     ToRPCErr passes io.EOF through unchanged, so the caller gets the
     wrong error code.
  3. doInvoke finishes writes, then MsgRecv sees the cancellation and
     returns codes.Canceled. This is the correct outcome but only
     happens by luck of timing.

Why this didn't happen before multiplexing:

The old single-stream manager used a non-blocking SendCancel that
returned (busy=true) when the write lock was held by an in-progress
write. With SoftCancel=false (the default), the fallback path was:

  manageStream calls stream.Cancel(ctx.Err()). The stream is not
  finished because doInvoke holds the write lock, so the manager calls
  m.terminate(), which closes the entire transport. The in-flight
  Writer.Write() fails with an IO error, and checkCancelError sees
  cancel.IsSet() and returns context.Canceled.

The correct error surfaced, but through connection termination. This was
fine in single-stream mode where one stream is one connection. With
multiplexing, we cannot terminate the entire connection for one stream's
cancellation. The new SendCancel blocks on the write lock to guarantee
the cancel frame is sent, and that introduced this race.

How gRPC handles this (verified against grpc-go source):

gRPC uses two mechanisms. First, newAttemptLocked (stream.go:408) checks
cs.ctx.Err() before creating the transport stream. This catches the
already-cancelled case without allocating resources. Second, for unary
RPCs, csAttempt.sendMsg (stream.go:1092) swallows write errors and
returns nil when !cs.desc.ClientStreams. The real error always surfaces
from RecvMsg, which detects context cancellation via
recvBufferReader.readClient (transport.go:239) and returns
status.Error(codes.Canceled, ...). This means gRPC never returns io.EOF
from a unary RPC because it never short-circuits on a send error.

For streaming RPCs, gRPC returns io.EOF from Send() after cancel (the
stream is done for writing) and codes.Canceled from Recv() (the actual
reason). Our grpccompat tests confirm this by comparing gRPC and DRPC
error results for identical cancel scenarios.

Our fix:

Rather than restructuring doInvoke to swallow send errors like gRPC, we
use the stream's existing CheckCancelError mechanism.

NewClientStream checks ctx.Err() before creating the stream. This
mirrors gRPC's newAttemptLocked check and avoids wasting a stream ID,
spawning a goroutine, and allocating stream resources.

doInvoke defers stream.CheckCancelError on its return value. If any
operation in doInvoke fails because SendCancel won the write lock race
(returning io.EOF via the send signal), CheckCancelError replaces it
with the cancel signal's error (context.Canceled). This is the same
function the stream already uses internally for transport write
failures. CheckCancelError is exported (was checkCancelError) so that
doInvoke in the drpcconn package can call it.

On TOCTOU:

The NewClientStream check is technically TOCTOU: the context could be
cancelled immediately after the check passes. This is acceptable because
Go's context cancellation model is cooperative, not preemptive. The
context package provides Done() "for use in select statements," and
operations check at natural boundaries rather than continuously. The
standard library follows this pattern: http.Client.Do checks between
redirect hops, database/sql checks before query execution, and gRPC
checks in newAttemptLocked before creating the transport stream. If the
context is cancelled mid-operation, manageStream handles cleanup and the
deferred CheckCancelError corrects the error code.

Author: shubhamdhama

drpcconn/conn.go

@@ -89,6 +89,7 @@ func (c *Conn) Invoke(ctx context.Context, rpc string, enc drpc.Encoding, in, ou
 }
 
 func (c *Conn) doInvoke(stream *drpcstream.Stream, enc drpc.Encoding, rpc string, data []byte, metadata []byte, out drpc.Message) (err error) {
+	defer func() { err = stream.CheckCancelError(err) }()
 	if err := stream.WriteInvoke(rpc, metadata); err != nil {
 		return err
 	}

drpcmanager/manager.go

@@ -328,6 +328,9 @@ func (m *Manager) Close() error {
 
 // NewClientStream starts a stream on the managed transport for use by a client.
 func (m *Manager) NewClientStream(ctx context.Context, rpc string) (stream *drpcstream.Stream, err error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
 	return m.newStream(ctx, m.lastStreamID.Add(1), drpc.StreamKindClient, rpc)
 }
 

drpcstream/stream.go

@@ -297,10 +297,10 @@ func (s *Stream) checkFinished() {
 	}
 }
 
-// checkCancelError will replace the error with one from the cancel signal if it
+// CheckCancelError will replace the error with one from the cancel signal if it
 // is set. This is to prevent errors from reads/writes to a transport after it
 // has been asynchronously closed due to context cancelation.
-func (s *Stream) checkCancelError(err error) error {
+func (s *Stream) CheckCancelError(err error) error {
 	if s.sigs.cancel.IsSet() {
 		return s.sigs.cancel.Err()
 	}
@@ -401,7 +401,7 @@ func (s *Stream) rawWriteLocked(kind drpcwire.Kind, data []byte) (err error) {
 		s.log("SEND", fr.String)
 
 		if err := s.wr.WriteFrame(fr); err != nil {
-			return s.checkCancelError(errs.Wrap(err))
+			return s.CheckCancelError(errs.Wrap(err))
 		} else if fr.Done {
 			return nil
 		}
@@ -496,7 +496,7 @@ func (s *Stream) SendError(serr error) (err error) {
 	s.terminate(termError)
 	s.mu.Unlock()
 
-	return s.checkCancelError(s.sendPacketLocked(drpcwire.KindError, false, drpcwire.MarshalError(serr)))
+	return s.CheckCancelError(s.sendPacketLocked(drpcwire.KindError, false, drpcwire.MarshalError(serr)))
 }
 
 // SendCancel terminates the stream and sends a cancel to the remote side. It
@@ -519,7 +519,7 @@ func (s *Stream) SendCancel(err error) error {
 	s.terminate(err)
 	s.mu.Unlock()
 
-	return s.checkCancelError(s.sendPacketLocked(drpcwire.KindCancel, true, nil))
+	return s.CheckCancelError(s.sendPacketLocked(drpcwire.KindCancel, true, nil))
 }
 
 // Close terminates the stream and sends that the stream has been closed to the
@@ -540,7 +540,7 @@ func (s *Stream) Close() (err error) {
 	s.terminate(termClosed)
 	s.mu.Unlock()
 
-	return s.checkCancelError(s.sendPacketLocked(drpcwire.KindClose, false, nil))
+	return s.CheckCancelError(s.sendPacketLocked(drpcwire.KindClose, false, nil))
 }
 
 // CloseSend informs the remote that no more messages will be sent. If the remote has
@@ -563,7 +563,7 @@ func (s *Stream) CloseSend() (err error) {
 	s.terminateIfBothClosed()
 	s.mu.Unlock()
 
-	return s.checkCancelError(s.sendPacketLocked(drpcwire.KindCloseSend, false, nil))
+	return s.CheckCancelError(s.sendPacketLocked(drpcwire.KindCloseSend, false, nil))
 }
 
 // Cancel transitions the stream into a state where all writes to the transport will return