kv,kvcoord,sql: poison txnCoordSender after a retryable error

pkg/kv/db.go

@@ -851,6 +851,18 @@ func (db *DB) NewTxn(ctx context.Context, debugName string) *Txn {
 // from recoverable internal errors, and is automatically committed
 // otherwise. The retryable function should have no side effects which could
 // cause problems in the event it must be run more than once.
+// For example:
+// err := db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+//		if kv, err := txn.Get(ctx, key); err != nil {
+//			return err
+//		}
+//		// ...
+//		return nil
+//	})
+// Note that once the transaction encounters a retryable error, the txn object
+// is marked as poisoned and all future ops fail fast until the retry. The
+// callback may return either nil or the retryable error. Txn is responsible for
+// resetting the transaction and retrying the callback.
 func (db *DB) Txn(ctx context.Context, retryable func(context.Context, *Txn) error) error {
 	// TODO(radu): we should open a tracing Span here (we need to figure out how
 	// to use the correct tracer).

pkg/kv/db_test.go

@@ -13,6 +13,7 @@ package kv_test
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -715,3 +716,84 @@ func TestGenerateForcedRetryableError(t *testing.T) {
 	require.True(t, errors.As(err, &retryErr))
 	require.Equal(t, 1, int(retryErr.Transaction.Epoch))
 }
+
+// Get a retryable error within a db.Txn transaction and verify the retry
+// succeeds. We are verifying the behavior is the same whether the retryable
+// callback returns the retryable error or returns nil. Both implementations are
+// legal - returning early (with either nil or the error) after a retryable
+// error is optional.
+func TestDB_TxnRetry(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	s, db := setup(t)
+	defer s.Stopper().Stop(context.Background())
+
+	testutils.RunTrueAndFalse(t, "returnNil", func(t *testing.T, returnNil bool) {
+		keyA := fmt.Sprintf("a_return_nil_%t", returnNil)
+		keyB := fmt.Sprintf("b_return_nil_%t", returnNil)
+		runNumber := 0
+		err := db.Txn(context.Background(), func(ctx context.Context, txn *kv.Txn) error {
+			require.NoError(t, txn.Put(ctx, keyA, "1"))
+			require.NoError(t, txn.Put(ctx, keyB, "1"))
+
+			{
+				// High priority txn - will abort the other txn.
+				hpTxn := kv.NewTxn(ctx, db, 0)
+				require.NoError(t, hpTxn.SetUserPriority(roachpb.MaxUserPriority))
+				// Only write if we have not written before, because otherwise we will keep aborting
+				// the other txn forever.
+				r, err := hpTxn.Get(ctx, keyA)
+				require.NoError(t, err)
+				if !r.Exists() {
+					require.Zero(t, runNumber)
+					require.NoError(t, hpTxn.Put(ctx, keyA, "hp txn"))
+					require.NoError(t, hpTxn.Commit(ctx))
+				} else {
+					// We already wrote to keyA, meaning this is a retry, no need to write again.
+					require.Equal(t, 1, runNumber)
+					require.NoError(t, hpTxn.Rollback(ctx))
+				}
+			}
+
+			// Read, so that we'll get a retryable error.
+			r, err := txn.Get(ctx, keyA)
+			if runNumber == 0 {
+				// First run, we should get a retryable error.
+				require.Zero(t, runNumber)
+				require.IsType(t, &roachpb.TransactionRetryWithProtoRefreshError{}, err)
+				require.Equal(t, []byte(nil), r.ValueBytes())
+
+				// At this point txn is poisoned, and any op returns the same (poisoning) error.
+				r, err = txn.Get(ctx, keyB)
+				require.IsType(t, &roachpb.TransactionRetryWithProtoRefreshError{}, err)
+				require.Equal(t, []byte(nil), r.ValueBytes())
+			} else {
+				// The retry should succeed.
+				require.Equal(t, 1, runNumber)
+				require.NoError(t, err)
+				require.Equal(t, []byte("1"), r.ValueBytes())
+			}
+			runNumber++
+
+			if returnNil {
+				return nil
+			}
+			// Return the retryable error.
+			return err
+		})
+		require.NoError(t, err)
+		require.Equal(t, 2, runNumber)
+
+		err1 := db.Txn(context.Background(), func(ctx context.Context, txn *kv.Txn) error {
+			// The high priority txn was overwritten by the successful retry.
+			kv, e1 := txn.Get(ctx, keyA)
+			require.NoError(t, e1)
+			require.Equal(t, []byte("1"), kv.ValueBytes())
+			kv, e2 := txn.Get(ctx, keyB)
+			require.NoError(t, e2)
+			require.Equal(t, []byte("1"), kv.ValueBytes())
+			return nil
+		})
+		require.NoError(t, err1)
+	})
+}

pkg/kv/kvclient/kvcoord/testdata/savepoints

@@ -486,6 +486,11 @@ savepoint x
 abort
 ----
 (*roachpb.TransactionRetryWithProtoRefreshError)
+txn id not changed
+
+reset
+----
+txn error cleared
 txn id changed
 
 release x

pkg/kv/kvclient/kvcoord/txn_coord_sender.go

@@ -50,6 +50,12 @@ const (
 	// txnPending is the normal state for ongoing transactions.
 	txnPending txnState = iota
 
+	// txnRetryableError means that the transaction encountered a
+	// TransactionRetryWithProtoRefreshError, and calls to Send() fail in this
+	// state. It is possible to move back to txnPending by calling
+	// ClearTxnRetryableErr().
+	txnRetryableError
+
 	// txnError means that a batch encountered a non-retriable error. Further
 	// batches except EndTxn(commit=false) will be rejected.
 	txnError
@@ -105,6 +111,11 @@ type TxnCoordSender struct {
 		syncutil.Mutex
 
 		txnState txnState
+
+		// storedRetryableErr is set when txnState == txnRetryableError. This
+		// storedRetryableErr is returned to clients on Send().
+		storedRetryableErr *roachpb.TransactionRetryWithProtoRefreshError
+
 		// storedErr is set when txnState == txnError. This storedErr is returned to
 		// clients on Send().
 		storedErr *roachpb.Error
@@ -686,6 +697,8 @@ func (tc *TxnCoordSender) maybeRejectClientLocked(
 	switch tc.mu.txnState {
 	case txnPending:
 		// All good.
+	case txnRetryableError:
+		return roachpb.NewError(tc.mu.storedRetryableErr)
 	case txnError:
 		return tc.mu.storedErr
 	case txnFinalized:
@@ -712,19 +725,10 @@ func (tc *TxnCoordSender) maybeRejectClientLocked(
 		// The transaction heartbeat observed an aborted transaction record and
 		// this was not due to a synchronous transaction commit and transaction
 		// record garbage collection.
-		// See the comment on txnHeartbeater.mu.finalizedStatus for more details.
+		// See the comment on txnHeartbeater.mu.finalObservedStatus for more details.
 		abortedErr := roachpb.NewErrorWithTxn(
 			roachpb.NewTransactionAbortedError(roachpb.ABORT_REASON_CLIENT_REJECT), &tc.mu.txn)
-		if tc.typ == kv.LeafTxn {
-			// Leaf txns return raw retriable errors (which get handled by the
-			// root) rather than TransactionRetryWithProtoRefreshError.
-			return abortedErr
-		}
-		// Root txns handle retriable errors.
-		newTxn := roachpb.PrepareTransactionForRetry(
-			ctx, abortedErr, roachpb.NormalUserPriority, tc.clock)
-		return roachpb.NewError(roachpb.NewTransactionRetryWithProtoRefreshError(
-			abortedErr.String(), tc.mu.txn.ID, newTxn))
+		return roachpb.NewError(tc.handleRetryableErrLocked(ctx, abortedErr))
 	case protoStatus != roachpb.PENDING || hbObservedStatus != roachpb.PENDING:
 		// The transaction proto is in an unexpected state.
 		return roachpb.NewErrorf(
@@ -816,6 +820,11 @@ func (tc *TxnCoordSender) handleRetryableErrLocked(
 		errTxnID, // the id of the transaction that encountered the error
 		newTxn)
 
+	// Move to a retryable error state, where all Send() calls fail until the
+	// state is cleared.
+	tc.mu.txnState = txnRetryableError
+	tc.mu.storedRetryableErr = retErr
+
 	// If the ID changed, it means we had to start a new transaction and the
 	// old one is toast. This TxnCoordSender cannot be used any more - future
 	// Send() calls will be rejected; the client is supposed to create a new
@@ -1360,3 +1369,25 @@ func (tc *TxnCoordSender) DeferCommitWait(ctx context.Context) func(context.Cont
 		return tc.maybeCommitWait(ctx, true /* deferred */)
 	}
 }
+
+// GetTxnRetryableErr is part of the TxnSender interface.
+func (tc *TxnCoordSender) GetTxnRetryableErr(
+	ctx context.Context,
+) *roachpb.TransactionRetryWithProtoRefreshError {
+	tc.mu.Lock()
+	defer tc.mu.Unlock()
+	if tc.mu.txnState == txnRetryableError {
+		return tc.mu.storedRetryableErr
+	}
+	return nil
+}
+
+// ClearTxnRetryableErr is part of the TxnSender interface.
+func (tc *TxnCoordSender) ClearTxnRetryableErr(ctx context.Context) {
+	tc.mu.Lock()
+	defer tc.mu.Unlock()
+	if tc.mu.txnState == txnRetryableError {
+		tc.mu.storedRetryableErr = nil
+		tc.mu.txnState = txnPending
+	}
+}

pkg/kv/kvclient/kvcoord/txn_coord_sender_savepoints.go

@@ -126,16 +126,12 @@ func (tc *TxnCoordSender) RollbackToSavepoint(ctx context.Context, s kv.Savepoin
 			err = roachpb.NewTransactionRetryWithProtoRefreshError(
 				"cannot rollback to savepoint after a transaction restart",
 				tc.mu.txn.ID,
-				// The transaction inside this error doesn't matter.
-				roachpb.Transaction{},
+				tc.mu.txn,
 			)
 		}
 		return err
 	}
 
-	// Restore the transaction's state, in case we're rewiding after an error.
-	tc.mu.txnState = txnPending
-
 	tc.mu.active = sp.active
 
 	for _, reqInt := range tc.interceptorStack {
@@ -173,8 +169,7 @@ func (tc *TxnCoordSender) ReleaseSavepoint(ctx context.Context, s kv.SavepointTo
 		err = roachpb.NewTransactionRetryWithProtoRefreshError(
 			"cannot release savepoint after a transaction restart",
 			tc.mu.txn.ID,
-			// The transaction inside this error doesn't matter.
-			roachpb.Transaction{},
+			tc.mu.txn,
 		)
 	}
 	return err

pkg/kv/kvclient/kvcoord/txn_coord_sender_savepoints_test.go

@@ -124,6 +124,16 @@ func TestSavepoints(t *testing.T) {
 				}
 				fmt.Fprintf(&buf, "txn id %s\n", changed)
 
+			case "reset":
+				prevID := txn.ID()
+				txn.PrepareForRetry(ctx)
+				changed := "changed"
+				if prevID == txn.ID() {
+					changed = "not changed"
+				}
+				fmt.Fprintf(&buf, "txn error cleared\n")
+				fmt.Fprintf(&buf, "txn id %s\n", changed)
+
 			case "put":
 				b := txn.NewBatch()
 				b.Put(td.CmdArgs[0].Key, td.CmdArgs[1].Key)