kvserver: do not cancel while in HandleRaftResponse

If the context is canceled in the middle of HandleRaftResponse, a
Replica can be left in a broken state, e.g. its destruction can be
partially completed. This can brick a RangeID, e.g. cause a subsequent
call to getOrCreateReplica to be stuck in an infinite loop.

Epic: none
Release note: none

Author: pav-kv

pkg/kv/kvserver/raft_transport.go

@@ -599,25 +599,24 @@ func (t *RaftTransport) StopOutgoingMessage(storeID roachpb.StoreID) {
 func (t *RaftTransport) processQueue(
 	ctx context.Context, q *raftSendQueue, client RPCMultiRaftClient, _ rpcbase.ConnectionClass,
 ) error {
-	batchCtx, cancel := context.WithCancel(ctx)
+	streamCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
-	stream, err := client.RaftMessageBatch(batchCtx) // closed via cancellation
+	stream, err := client.RaftMessageBatch(streamCtx) // closed via cancellation
 	if err != nil {
 		return errors.Wrapf(err, "creating batch client")
 	}
 
 	errCh := make(chan error, 1)
 
 	// NB: the stream context is canceled when this func returns, and causes the
-	// response handling loop to terminate asynchronously.
-	//
-	// TODO(#140958): the context cancellation in the middle of HandleRaftResponse
-	// can lead to broken state, such as a replica marked as destroyed but this
-	// not being reflected in storage.
+	// response handling loop to terminate asynchronously. Note though that the
+	// task uses the parent context for HandleRaftResponse calls, to let it finish
+	// gracefully. Previously, context cancellation inside HandleRaftResponse
+	// could cause incorrect / half-done Replica destruction (see #140958).
 	//
 	// TODO(pav-kv): wait for the task termination to prevent subsequent
 	// processQueue calls from piling up concurrent tasks.
-	goCtx, hdl, err := t.stopper.GetHandle(stream.Context(), stop.TaskOpts{
+	goCtx, hdl, err := t.stopper.GetHandle(ctx, stop.TaskOpts{
 		TaskName: "storage.RaftTransport: processing queue",
 	})
 	if err != nil {