kv: add error return from Txn.PrepareForRetry

Instead of using log.Fatal, return assertion errors.

Release note: None

Author: nvb

pkg/kv/db_test.go

@@ -872,7 +872,7 @@ func TestPreservingSteppingOnSenderReplacement(t *testing.T) {
 		require.NotEqual(t, pErr.TxnID, pErr.Transaction.ID)
 
 		// Reset the handle in order to get a new sender.
-		txn.PrepareForRetry(ctx)
+		require.NoError(t, txn.PrepareForRetry(ctx))
 
 		// Make sure we have a new txn ID.
 		require.NotEqual(t, pErr.TxnID, txn.ID())

pkg/kv/kvclient/kvcoord/txn_coord_sender_savepoints_test.go

@@ -127,7 +127,9 @@ func TestSavepoints(t *testing.T) {
 
 			case "reset":
 				prevID := txn.ID()
-				txn.PrepareForRetry(ctx)
+				if err := txn.PrepareForRetry(ctx); err != nil {
+					t.Fatal(err)
+				}
 				changed := "changed"
 				if prevID == txn.ID() {
 					changed = "not changed"

pkg/kv/kvclient/kvcoord/txn_coord_sender_test.go

@@ -912,7 +912,9 @@ func TestTxnCoordSenderTxnUpdatedOnError(t *testing.T) {
 			stopper.Stop(ctx)
 
 			if test.callPrepareForRetry {
-				txn.PrepareForRetry(ctx)
+				if err := txn.PrepareForRetry(ctx); err != nil {
+					t.Fatal(err)
+				}
 			}
 			if test.name != "nil" && err == nil {
 				t.Fatalf("expected an error")

pkg/kv/txn.go

@@ -994,15 +994,17 @@ func (txn *Txn) exec(ctx context.Context, fn func(context.Context, *Txn) error)
 			break
 		}
 
-		txn.PrepareForRetry(ctx)
+		if err := txn.PrepareForRetry(ctx); err != nil {
+			return err
+		}
 	}
 
 	return err
 }
 
 // PrepareForRetry needs to be called before a retry to perform some
 // book-keeping and clear errors when possible.
-func (txn *Txn) PrepareForRetry(ctx context.Context) {
+func (txn *Txn) PrepareForRetry(ctx context.Context) error {
 	// Reset commit triggers. These must be reconfigured by the client during the
 	// next retry.
 	txn.commitTriggers = nil
@@ -1012,11 +1014,11 @@ func (txn *Txn) PrepareForRetry(ctx context.Context) {
 
 	retryErr := txn.mu.sender.GetTxnRetryableErr(ctx)
 	if retryErr == nil {
-		return
+		return nil
 	}
 	if txn.typ != RootTxn {
-		panic(errors.WithContextTags(errors.NewAssertionErrorWithWrappedErrf(
-			retryErr, "PrepareForRetry() called on leaf txn"), ctx))
+		return errors.WithContextTags(errors.NewAssertionErrorWithWrappedErrf(
+			retryErr, "PrepareForRetry() called on leaf txn"), ctx)
 	}
 	log.VEventf(ctx, 2, "retrying transaction: %s because of a retryable error: %s",
 		txn.debugNameLocked(), retryErr)
@@ -1029,23 +1031,24 @@ func (txn *Txn) PrepareForRetry(ctx context.Context) {
 		// txn IDs. However, at no point can both the old and new incarnation of a
 		// transaction be active at the same time -- this would constitute a
 		// programming error.
-		log.Fatalf(
-			ctx,
-			"unexpected retryable error for old incarnation of the transaction %s; current incarnation %s",
-			retryErr.TxnID,
-			txn.mu.ID,
-		)
+		return errors.WithContextTags(
+			errors.NewAssertionErrorWithWrappedErrf(
+				retryErr,
+				"unexpected retryable error for old incarnation of the transaction %s; current incarnation %s",
+				retryErr.TxnID,
+				txn.mu.ID,
+			), ctx)
 	}
 
 	if !retryErr.PrevTxnAborted() {
 		// If the retryable error doesn't correspond to an aborted transaction,
 		// there's no need to switch out the transaction. We simply clear the
 		// retryable error and proceed.
 		txn.mu.sender.ClearTxnRetryableErr(ctx)
-		return
+		return nil
 	}
 
-	txn.handleTransactionAbortedErrorLocked(ctx, retryErr)
+	return txn.handleTransactionAbortedErrorLocked(ctx, retryErr)
 }
 
 // Send runs the specified calls synchronously in a single batch and
@@ -1354,10 +1357,11 @@ func (txn *Txn) UpdateStateOnRemoteRetryableErr(ctx context.Context, pErr *kvpb.
 // the current one with it.
 func (txn *Txn) handleTransactionAbortedErrorLocked(
 	ctx context.Context, retryErr *kvpb.TransactionRetryWithProtoRefreshError,
-) {
+) error {
 	if !retryErr.PrevTxnAborted() {
 		// Sanity check we're dealing with a TransactionAbortedError.
-		log.Fatalf(ctx, "cannot replace root sender if txn not aborted: %v", retryErr)
+		return errors.WithContextTags(errors.NewAssertionErrorWithWrappedErrf(
+			retryErr, "cannot replace root sender if txn not aborted"), ctx)
 	}
 
 	// The transaction we had been using thus far has been aborted. The proto
@@ -1368,6 +1372,7 @@ func (txn *Txn) handleTransactionAbortedErrorLocked(
 	prevSteppingMode := txn.mu.sender.GetSteppingMode(ctx)
 	txn.mu.sender = txn.db.factory.RootTransactionalSender(newTxn, txn.mu.userPriority)
 	txn.mu.sender.ConfigureStepping(ctx, prevSteppingMode)
+	return nil
 }
 
 // SetFixedTimestamp makes the transaction run in an unusual way, at a "fixed

pkg/sql/conn_fsm.go

@@ -552,11 +552,14 @@ func cleanupAndFinishOnError(args fsm.Args) error {
 
 func prepareTxnForRetry(args fsm.Args) error {
 	ts := args.Extended.(*txnState)
-	func() {
+	err := func() error {
 		ts.mu.Lock()
 		defer ts.mu.Unlock()
-		ts.mu.txn.PrepareForRetry(ts.Ctx)
+		return ts.mu.txn.PrepareForRetry(ts.Ctx)
 	}()
+	if err != nil {
+		return err
+	}
 	ts.setAdvanceInfo(
 		advanceOne,
 		noRewind,
@@ -587,13 +590,16 @@ func moveToCommitWaitAfterInternalCommit(args fsm.Args) error {
 func prepareTxnForRetryWithRewind(args fsm.Args) error {
 	pl := args.Payload.(eventRetriableErrPayload)
 	ts := args.Extended.(*txnState)
-	func() {
+	err := func() error {
 		ts.mu.Lock()
 		defer ts.mu.Unlock()
-		ts.mu.txn.PrepareForRetry(ts.Ctx)
 		ts.mu.autoRetryReason = pl.err
 		ts.mu.autoRetryCounter++
+		return ts.mu.txn.PrepareForRetry(ts.Ctx)
 	}()
+	if err != nil {
+		return err
+	}
 	// The caller will call rewCap.rewindAndUnlock().
 	ts.setAdvanceInfo(
 		rewind,