kv: don't let pusher win when pushing STAGING txn with equal priority

Fixes cockroachdb#105330.
Fixes cockroachdb#101721.

This commit updates the transaction push logic to stop pushes from completing
successfully when the pushee is STAGING and the pusher has equal priority. Prior
to this change, the pusher would win in these cases when using a PUSH_TIMESTAMP
if at least one of the two transactions involved used a weak isolation level.

This had two undesirable effects:
- if the pushee was still possibly committable and requiring recovery
  (`!(knownHigherTimestamp || knownHigherEpoch)` in the code) then the pusher
  would immediately begin parallel commit recovery, attempting to disrupt the
  commit and abort the pushee. This is undesirable because the pushee may still
  be in progress and in cases of equal priority, we'd like to wait for the
  parallel commit to complete before kicking off recovery, deferring recovery to
  only the cases where the pushee/committers's heartbeat has expired.
- if the pushee was known to be uncommittable (`knownHigherTimestamp || knownHigherEpoch`
  in the code), then txn recovery was not kicked off but the pusher still could
  not perform the PUSH_TIMESTAMP (see e40c1b4), so it would return a
  `TransactionPushError`. This confused logic in `handleTransactionPushError`,
  allowing the error to escape to the client.

This commit resolves both issues by considering the pushee's transaction status
in `txnwait.ShouldPushImmediately`.

Release note: None

pkg/kv/kvnemesis/validator.go

@@ -1232,16 +1232,6 @@ func (v *validator) failIfError(
 ) (ambiguous, hasError bool) {
 	exceptions = append(exceptions[:len(exceptions):len(exceptions)], func(err error) bool {
 		return errors.Is(err, errInjected)
-	}, func(err error) bool {
-		// Work-around for [1].
-		//
-		// TODO(arul): find out why we (as of [2]) sometimes leaking
-		// *TransactionPushError (wrapped in `UnhandledRetryableError`) from
-		// `db.Get`, `db.Scan`, etc.
-		//
-		// [1]: https://github.com/cockroachdb/cockroach/issues/105330
-		// [2]: https://github.com/cockroachdb/cockroach/pull/97779
-		return errors.HasType(err, (*kvpb.UnhandledRetryableError)(nil))
 	})
 	switch r.Type {
 	case ResultType_Unknown:

pkg/kv/kvserver/batcheval/cmd_push_txn.go

@@ -26,6 +26,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/must"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/redact"
 )
@@ -213,38 +214,9 @@ func PushTxn(
 	}
 	reply.PusheeTxn.UpgradePriority(args.PusheeTxn.Priority)
 
-	// If the pusher is aware that the pushee's currently recorded attempt at a
-	// parallel commit failed, either because it found intents at a higher
-	// timestamp than the parallel commit attempt or because it found intents at
-	// a higher epoch than the parallel commit attempt, it should not consider
-	// the pushee to be performing a parallel commit. Its commit status is not
-	// indeterminate.
-	if (knownHigherTimestamp || knownHigherEpoch) && reply.PusheeTxn.Status == roachpb.STAGING {
-		reply.PusheeTxn.Status = roachpb.PENDING
-		reply.PusheeTxn.InFlightWrites = nil
-		// If the pusher is aware that the pushee's currently recorded attempt
-		// at a parallel commit failed, upgrade PUSH_TIMESTAMPs to PUSH_ABORTs.
-		// We don't want to move the transaction back to PENDING, as this is not
-		// (currently) allowed by the recovery protocol. We also don't want to
-		// move the transaction to a new timestamp while retaining the STAGING
-		// status, as this could allow the transaction to enter an implicit
-		// commit state without its knowledge, leading to atomicity violations.
-		//
-		// This has no effect on pushes that fail with a TransactionPushError.
-		// Such pushes will still wait on the pushee to retry its commit and
-		// eventually commit or abort. It also has no effect on expired pushees,
-		// as they would have been aborted anyway. This only impacts pushes
-		// which would have succeeded due to priority mismatches. In these
-		// cases, the push acts the same as a short-circuited transaction
-		// recovery process, because the transaction recovery procedure always
-		// finalizes target transactions, even if initiated by a PUSH_TIMESTAMP.
-		if pushType == kvpb.PUSH_TIMESTAMP {
-			pushType = kvpb.PUSH_ABORT
-		}
-	}
-
 	pusherIso, pusheeIso := args.PusherTxn.IsoLevel, reply.PusheeTxn.IsoLevel
 	pusherPri, pusheePri := args.PusherTxn.Priority, reply.PusheeTxn.Priority
+	pusheeStatus := reply.PusheeTxn.Status
 	var pusherWins bool
 	var reason string
 	switch {
@@ -258,7 +230,7 @@ func PushTxn(
 		// If just attempting to cleanup old or already-committed txns,
 		// pusher always fails.
 		pusherWins = false
-	case txnwait.CanPushWithPriority(pushType, pusherIso, pusheeIso, pusherPri, pusheePri):
+	case txnwait.CanPushWithPriority(pushType, pusherIso, pusheeIso, pusherPri, pusheePri, pusheeStatus):
 		reason = "pusher has priority"
 		pusherWins = true
 	case args.Force:
@@ -282,13 +254,42 @@ func PushTxn(
 	// If the pushed transaction is in the staging state, we can't change its
 	// record without first going through the transaction recovery process and
 	// attempting to finalize it.
+	pusheeStaging := pusheeStatus == roachpb.STAGING
+	// However, if the pusher is aware that the pushee's currently recorded
+	// attempt at a parallel commit failed, either because it found intents at a
+	// higher timestamp than the parallel commit attempt or because it found
+	// intents at a higher epoch than the parallel commit attempt, it should not
+	// consider the pushee to be performing a parallel commit. Its commit status
+	// is not indeterminate.
+	pusheeStagingFailed := pusheeStaging && (knownHigherTimestamp || knownHigherEpoch)
 	recoverOnFailedPush := cArgs.EvalCtx.EvalKnobs().RecoverIndeterminateCommitsOnFailedPushes
-	if reply.PusheeTxn.Status == roachpb.STAGING && (pusherWins || recoverOnFailedPush) {
+	if pusheeStaging && !pusheeStagingFailed && (pusherWins || recoverOnFailedPush) {
 		err := kvpb.NewIndeterminateCommitError(reply.PusheeTxn)
 		log.VEventf(ctx, 1, "%v", err)
 		return result.Result{}, err
 	}
 
+	// If the pusher is aware that the pushee's currently recorded attempt at a
+	// parallel commit failed, upgrade PUSH_TIMESTAMPs to PUSH_ABORTs. We don't
+	// want to move the transaction back to PENDING, as this is not (currently)
+	// allowed by the recovery protocol. We also don't want to move the
+	// transaction to a new timestamp while retaining the STAGING status, as this
+	// could allow the transaction to enter an implicit commit state without its
+	// knowledge, leading to atomicity violations.
+	//
+	// This has no effect on pushes that fail with a TransactionPushError. Such
+	// pushes will still wait on the pushee to retry its commit and eventually
+	// commit or abort. It also has no effect on expired pushees, as they would
+	// have been aborted anyway. This only impacts pushes which would have
+	// succeeded due to priority mismatches. In these cases, the push acts the
+	// same as a short-circuited transaction recovery process, because the
+	// transaction recovery procedure always finalizes target transactions, even
+	// if initiated by a PUSH_TIMESTAMP.
+	if pusheeStaging && pusherWins && pushType == kvpb.PUSH_TIMESTAMP {
+		_ = must.True(ctx, pusheeStagingFailed, "parallel commit must be known to have failed for push to succeed")
+		pushType = kvpb.PUSH_ABORT
+	}
+
 	if !pusherWins {
 		err := kvpb.NewTransactionPushError(reply.PusheeTxn)
 		log.VEventf(ctx, 1, "%v", err)
@@ -306,6 +307,8 @@ func PushTxn(
 		// Forward the timestamp to accommodate AbortSpan GC. See method comment for
 		// details.
 		reply.PusheeTxn.WriteTimestamp.Forward(reply.PusheeTxn.LastActive())
+		// If the transaction was previously staging, clear its in-flight writes.
+		reply.PusheeTxn.InFlightWrites = nil
 		// If the transaction record was already present, persist the updates to it.
 		// If not, then we don't want to create it. This could allow for finalized
 		// transactions to be revived. Instead, we obey the invariant that only the

pkg/kv/kvserver/concurrency/concurrency_manager_test.go

@@ -747,18 +747,19 @@ func (c *cluster) PushTransaction(
 		pusheeTxn, pusheeRecordSig := pusheeRecord.asTxn()
 		pusheeIso := pusheeTxn.IsoLevel
 		pusheePri := pusheeTxn.Priority
+		pusheeStatus := pusheeTxn.Status
 		// NOTE: this logic is adapted from cmd_push_txn.go.
 		var pusherWins bool
 		switch {
-		case pusheeTxn.Status.IsFinalized():
+		case pusheeStatus.IsFinalized():
 			// Already finalized.
 			return pusheeTxn, nil
 		case pushType == kvpb.PUSH_TIMESTAMP && pushTo.LessEq(pusheeTxn.WriteTimestamp):
 			// Already pushed.
 			return pusheeTxn, nil
 		case pushType == kvpb.PUSH_TOUCH:
 			pusherWins = false
-		case txnwait.CanPushWithPriority(pushType, pusherIso, pusheeIso, pusherPri, pusheePri):
+		case txnwait.CanPushWithPriority(pushType, pusherIso, pusheeIso, pusherPri, pusheePri, pusheeStatus):
 			pusherWins = true
 		default:
 			pusherWins = false

pkg/kv/kvserver/concurrency/lock_table_waiter.go

@@ -1298,7 +1298,11 @@ func canPushWithPriority(req Request, s waitingState) bool {
 	}
 	pusheeIso = s.txn.IsoLevel
 	pusheePri = s.txn.Priority
-	return txnwait.CanPushWithPriority(pushType, pusherIso, pusheeIso, pusherPri, pusheePri)
+	// We assume that the pushee is in the PENDING state when deciding whether
+	// to push. A push may determine that the pushee is STAGING or has already
+	// been finalized.
+	pusheeStatus := roachpb.PENDING
+	return txnwait.CanPushWithPriority(pushType, pusherIso, pusheeIso, pusherPri, pusheePri, pusheeStatus)
 }
 
 func logResolveIntent(ctx context.Context, intent roachpb.LockUpdate) {

pkg/kv/kvserver/replica_send.go

@@ -773,7 +773,7 @@ func (r *Replica) handleTransactionPushError(
 	dontRetry := r.store.cfg.TestingKnobs.DontRetryPushTxnFailures
 	if !dontRetry && ba.IsSinglePushTxnRequest() {
 		pushReq := ba.Requests[0].GetInner().(*kvpb.PushTxnRequest)
-		dontRetry = txnwait.ShouldPushImmediately(pushReq)
+		dontRetry = txnwait.ShouldPushImmediately(pushReq, t.PusheeTxn.Status)
 	}
 	if dontRetry {
 		return g, pErr

pkg/kv/kvserver/txn_recovery_integration_test.go

@@ -372,20 +372,6 @@ func TestTxnRecoveryFromStagingWithoutHighPriority(t *testing.T) {
 			pErrC <- pErr
 		}))
 
-		// If the pushee is not serializable and the pusher is reading, we
-		// currently expect the conflicting operation to immediately succeed,
-		// sometimes with an error, sometimes not. See #105330.
-		if pusheeIsoLevel.ToleratesWriteSkew() && !pusherWriting {
-			pErr = <-pErrC
-			if pusheeCommits {
-				require.Nil(t, pErr, "error: %s", pErr)
-			} else {
-				require.NotNil(t, pErr)
-				require.Regexp(t, "failed to push", pErr)
-			}
-			return
-		}
-
 		// Wait for the conflict to push and be queued in the txn wait queue.
 		testutils.SucceedsSoon(t, func() error {
 			select {

pkg/kv/kvserver/txnwait/queue.go

@@ -70,23 +70,26 @@ func TestingOverrideTxnLivenessThreshold(t time.Duration) func() {
 // proceed without queueing. This is true for pushes which are neither
 // ABORT nor TIMESTAMP, but also for ABORT and TIMESTAMP pushes where
 // the pushee has min priority or pusher has max priority.
-func ShouldPushImmediately(req *kvpb.PushTxnRequest) bool {
+func ShouldPushImmediately(req *kvpb.PushTxnRequest, pusheeStatus roachpb.TransactionStatus) bool {
 	if req.Force {
 		return true
 	}
 	return CanPushWithPriority(
 		req.PushType,
 		req.PusherTxn.IsoLevel, req.PusheeTxn.IsoLevel,
 		req.PusherTxn.Priority, req.PusheeTxn.Priority,
+		pusheeStatus,
 	)
 }
 
 // CanPushWithPriority returns true if the pusher can perform the specified push
-// type on the pushee, based on the two txns' isolation levels and priorities.
+// type on the pushee, based on the two txns' isolation levels, their priorities,
+// and the pushee's status.
 func CanPushWithPriority(
 	pushType kvpb.PushTxnType,
 	pusherIso, pusheeIso isolation.Level,
 	pusherPri, pusheePri enginepb.TxnPriority,
+	pusheeStatus roachpb.TransactionStatus,
 ) bool {
 	// Normalize the transaction priorities so that normal user priorities are
 	// considered equal for the purposes of pushing.
@@ -103,6 +106,15 @@ func CanPushWithPriority(
 	case kvpb.PUSH_ABORT:
 		return pusherPri > pusheePri
 	case kvpb.PUSH_TIMESTAMP:
+		// If the pushee transaction is STAGING, only let the PUSH_TIMESTAMP through
+		// to disrupt the transaction commit if the pusher has a higher priority. If
+		// the priorities are equal, the PUSH_TIMESTAMP should wait for the commit
+		// to complete.
+		if pusheeStatus == roachpb.STAGING {
+			return pusherPri > pusheePri
+		}
+		// Otherwise, the pushee has not yet started committing...
+
 		// If the pushee transaction tolerates write skew, the PUSH_TIMESTAMP is
 		// harmless, so let it through.
 		return pusheeIso.ToleratesWriteSkew() ||
@@ -476,10 +488,6 @@ func (q *Queue) releaseWaitingQueriesLocked(ctx context.Context, txnID uuid.UUID
 func (q *Queue) MaybeWaitForPush(
 	ctx context.Context, req *kvpb.PushTxnRequest,
 ) (*kvpb.PushTxnResponse, *kvpb.Error) {
-	if ShouldPushImmediately(req) {
-		return nil, nil
-	}
-
 	q.mu.Lock()
 	// If the txn wait queue is not enabled or if the request is not
 	// contained within the replica, do nothing. The request can fall
@@ -501,6 +509,9 @@ func (q *Queue) MaybeWaitForPush(
 	if txn := pending.getTxn(); isPushed(req, txn) {
 		q.mu.Unlock()
 		return createPushTxnResponse(txn), nil
+	} else if ShouldPushImmediately(req, txn.Status) {
+		q.mu.Unlock()
+		return nil, nil
 	}
 
 	push := &waitingPush{