Merge #70159 #75632

70159: kvclient: switch error in ctx cancel edge case r=andreimatei a=andreimatei

This patch changes the error returned by the DistSender on a cancelled
ctx. Depending on the exactly who detects the ctx as cancelled, there
are a number of different possibilities - too many to enumerate here.
Generally, the client is supposed to get a context.Canceled error,
wrapped in different layers. The code path that this patch changes is
about the cancellation being detected by sendPartialBatch() without it
having been previously detected by sendToReplicas(). This path is
unusual (sendToReplicas() generally detects the ctx cancelled).
Depending on the exact cancellation timing, it's possible though for
sendPartialBatch() to detect it instead. In this case, this patch makes
it so that, if sendToReplicas() returned a sendError (indicating that
a suitable replica could not be reached and that the higher layer is
expected to continue trying other replicas), the error returned to the
client is a cancellation error instead of the sendError.

Touches #69419

Release note: None

75632: execstats: fix flakey TestTraceAnalyzer r=yuzefovich a=adityamaru

Fixes: #75546

Release note: None

Co-authored-by: Andrei Matei <[email protected]>
Co-authored-by: Aditya Maru <[email protected]>

pkg/kv/kvclient/kvcoord/dist_sender.go

@@ -1590,9 +1590,6 @@ func (ds *DistSender) sendPartialBatch(
 		}
 
 		if err != nil {
-			// Set pErr so that, if we don't perform any more retries, the
-			// deduceRetryEarlyExitError() call below the loop is inhibited.
-			pErr = roachpb.NewError(err)
 			switch {
 			case errors.HasType(err, sendError{}):
 				// We've tried all the replicas without success. Either they're all
@@ -1614,6 +1611,9 @@ func (ds *DistSender) sendPartialBatch(
 				routingTok.Evict(ctx)
 				continue
 			}
+			// Set pErr so that the deduceRetryEarlyExitError() call below the loop is
+			// inhibited.
+			pErr = roachpb.NewError(err)
 			break
 		}
 
@@ -1677,7 +1677,7 @@ func (ds *DistSender) sendPartialBatch(
 	// channels were closed.
 	if pErr == nil {
 		if err := ds.deduceRetryEarlyExitError(ctx); err == nil {
-			log.Fatal(ctx, "exited retry loop without an error")
+			pErr = roachpb.NewError(errors.AssertionFailedf("exited retry loop without an error"))
 		} else {
 			pErr = roachpb.NewError(err)
 		}

pkg/kv/kvserver/client_merge_test.go

@@ -4127,7 +4127,7 @@ func TestStoreRangeMergeDuringShutdown(t *testing.T) {
 	// Send a dummy get request on the RHS to force a lease acquisition. We expect
 	// this to fail, as quiescing stores cannot acquire leases.
 	_, err = store.DB().Get(ctx, key.Next())
-	if exp := "not lease holder"; !testutils.IsError(err, exp) {
+	if exp := "node unavailable"; !testutils.IsError(err, exp) {
 		t.Fatalf("expected %q error, but got %v", exp, err)
 	}
 }

pkg/sql/execstats/traceanalyzer_test.go

@@ -107,17 +107,14 @@ func TestTraceAnalyzer(t *testing.T) {
 	for _, vectorizeMode := range []sessiondatapb.VectorizeExecMode{sessiondatapb.VectorizeOff, sessiondatapb.VectorizeOn} {
 		execCtx, finishAndCollect := tracing.ContextWithRecordingSpan(ctx, execCfg.AmbientCtx.Tracer, t.Name())
 		defer finishAndCollect()
-		ie := execCfg.InternalExecutor
-		ie.SetSessionData(
-			&sessiondata.SessionData{
-				SessionData: sessiondatapb.SessionData{
-					VectorizeMode: vectorizeMode,
-				},
-				LocalOnlySessionData: sessiondatapb.LocalOnlySessionData{
-					DistSQLMode: sessiondatapb.DistSQLOn,
-				},
+		ie := execCfg.InternalExecutorFactory(ctx, &sessiondata.SessionData{
+			SessionData: sessiondatapb.SessionData{
+				VectorizeMode: vectorizeMode,
 			},
-		)
+			LocalOnlySessionData: sessiondatapb.LocalOnlySessionData{
+				DistSQLMode: sessiondatapb.DistSQLOn,
+			},
+		})
 		_, err := ie.ExecEx(
 			execCtx,
 			t.Name(),