Merge #145532

145532: concurrency: mark some unreplicated locks as ineligible for export  r=miraradeva a=stevendanna

Once a lock has been reported as missing via QueryIntent, it is
important that it stays missing. As described in #145458, if a lock
that has been previously reported as missing is somehow then found via
a subsequent QueryIntent request, this can result in a transaction
being erroneously committed by the transaction recovery process.

When unreplicated lock reliability is enabled, Lease and Merge use
ExportUnreplicatedLocks to move unreplicated locks from the lock table
into durable storage. However, these locks may include unreplicated
locks that "cover" a replicated lock that was previously reported as
missing, exactly what we must avoid.

Here, we solve that by extending QueryIntent's evaluation to mark any
locks it reports as missing as ineligible for export.

Fixes #145458

Release note: None

Co-authored-by: Steven Danna <danna@cockroachlabs.com>

Author: craig[bot]

pkg/kv/kvclient/kvcoord/BUILD.bazel

@@ -188,6 +188,7 @@ go_test(
         "//pkg/kv/kvpb/kvpbmock",
         "//pkg/kv/kvserver",
         "//pkg/kv/kvserver/closedts",
+        "//pkg/kv/kvserver/concurrency",
         "//pkg/kv/kvserver/concurrency/isolation",
         "//pkg/kv/kvserver/concurrency/lock",
         "//pkg/kv/kvserver/kvserverbase",

pkg/kv/kvclient/kvcoord/dist_sender_server_test.go

@@ -13,6 +13,7 @@ import (
 	"reflect"
 	"sort"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -25,7 +26,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/isolation"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc"
@@ -4643,3 +4646,173 @@ func TestProxyTracing(t *testing.T) {
 		t.Logf("found trace event; msg=%s, tag=%s, loc=%s", msg, tag, loc)
 	})
 }
+
+// TestUnexpectedCommitOnTxnRecovery constructs a scenario where transaction
+// recovery could incorrectly determine that a transaction is committed. The
+// scenario is as follows:
+//
+// Txn1:
+// - Writes to keyA.
+// - Acquires an unreplicated exclusive lock on keyB.
+// - Acquires a replicated shared lock on keyB. This lock is pipelined, and
+// replication for it fails.
+// - Attempts to commit, but fails because of the lost replicated Shared lock.
+//
+// Lease is then transferred to n3. This causes the unreplicated exclusive lock
+// on keyB to be replicated.
+//
+// Txn2:
+// - Attempts to read keyA, which kicks off transaction recovery for Txn1.
+// - Txn2 (incorrectly) concludes that Txn1 is committed at epoch=1 because it
+// finds a (stronger than Shared) replicated lock on keyB.
+//
+// Txn1:
+// - Back here, we do a stateful retry. We should learn that someone (Txn2)
+// aborted us when we go and try to commit. At the time of writing, we
+// incorrectly learn that we've been (unexpectedly) committed.
+func TestUnexpectedCommitOnTxnRecovery(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	keyA := roachpb.Key("a")
+	keyB := roachpb.Key("b")
+
+	var (
+		targetTxnIDString atomic.Value
+		cmdID             atomic.Value
+	)
+	cmdID.Store(kvserverbase.CmdIDKey(""))
+	targetTxnIDString.Store("")
+	ctx := context.Background()
+	st := cluster.MakeTestingClusterSettings()
+	// This test relies on unreplicated locks to be replicated on lease transfers.
+	concurrency.UnreplicatedLockReliabilityLeaseTransfer.Override(ctx, &st.SV, true)
+	tc := testcluster.StartTestCluster(t, 3, base.TestClusterArgs{
+		ServerArgs: base.TestServerArgs{
+			Settings: st,
+			Knobs: base.TestingKnobs{
+				Store: &kvserver.StoreTestingKnobs{
+					TestingProposalFilter: func(fArgs kvserverbase.ProposalFilterArgs) *kvpb.Error {
+						if fArgs.Req.Header.Txn == nil ||
+							fArgs.Req.Header.Txn.ID.String() != targetTxnIDString.Load().(string) {
+							return nil // not our txn
+						}
+						if !fArgs.Req.IsSingleRequest() {
+							// Not the request we care about.
+							return nil
+						}
+						getReq, ok := fArgs.Req.Requests[0].GetInner().(*kvpb.GetRequest)
+						// Only fail replication on the first retry.
+						epoch := fArgs.Req.Header.Txn.Epoch
+						if ok && getReq.KeyLockingDurability == lock.Replicated && epoch == 0 {
+							t.Logf("will fail application for txn %s@epoch=%d; req: %+v; raft cmdID: %s",
+								fArgs.Req.Header.Txn.ID.String(), epoch, getReq, fArgs.CmdID)
+							cmdID.Store(fArgs.CmdID)
+						}
+						return nil
+					},
+					TestingApplyCalledTwiceFilter: func(fArgs kvserverbase.ApplyFilterArgs) (int, *kvpb.Error) {
+						if fArgs.CmdID == cmdID.Load().(kvserverbase.CmdIDKey) {
+							t.Logf("failing application for raft cmdID: %s", cmdID)
+
+							return 0, kvpb.NewErrorf("test injected error")
+						}
+						return 0, nil
+					},
+				},
+			},
+		},
+	})
+	defer tc.Stopper().Stop(ctx)
+
+	transferLease := func(idx int) {
+		desc := tc.LookupRangeOrFatal(t, keyB)
+		tc.TransferRangeLeaseOrFatal(t, desc, tc.Target(idx))
+	}
+	// Make a db with transaction heartbeating disabled. This ensures that we
+	// don't mark Txn1 as PENDING after its first failed parallel commit attempt,
+	// which would otherwise prevent Txn2 from recovering Txn1.
+	s := tc.Server(0)
+	ambient := s.AmbientCtx()
+	tsf := kvcoord.NewTxnCoordSenderFactory(
+		kvcoord.TxnCoordSenderFactoryConfig{
+			AmbientCtx:        ambient,
+			HeartbeatInterval: -1, // disabled
+			Settings:          s.ClusterSettings(),
+			Clock:             s.Clock(),
+			Stopper:           s.Stopper(),
+		},
+		s.DistSenderI().(*kvcoord.DistSender),
+	)
+	db := kv.NewDB(ambient, tsf, s.Clock(), s.Stopper())
+
+	startTxn2 := make(chan struct{})
+	blockCh := make(chan struct{})
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	// Write to keyB so that we can later get a lock on it.
+	txn := db.NewTxn(ctx, "txn")
+	err := txn.Put(ctx, keyB, "valueB")
+	require.NoError(t, err)
+	require.NoError(t, txn.Commit(ctx))
+
+	go func() {
+		defer wg.Done()
+
+		err := db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+			if txnID := targetTxnIDString.Load(); txnID == "" {
+				// Store the txnID for the testing knobs.
+				targetTxnIDString.Store(txn.ID().String())
+				t.Logf("txn1 ID is: %s", txn.ID())
+			} else if txnID != txn.ID() {
+				// Since txn recovery aborted us, we get retried again but with an
+				// entirely new txnID. This time we just return. Writing nothing.
+				return nil
+			}
+
+			switch txn.Epoch() {
+			case 0:
+				err := txn.Put(ctx, keyA, "value")
+				require.NoError(t, err)
+				res, err := txn.GetForUpdate(ctx, keyB, kvpb.BestEffort)
+				require.NoError(t, err)
+				require.Equal(t, res.ValueBytes(), []byte("valueB"))
+				res, err = txn.GetForShare(ctx, keyB, kvpb.GuaranteedDurability)
+				require.NoError(t, err)
+				require.Equal(t, res.ValueBytes(), []byte("valueB"))
+				err = txn.Commit(ctx)
+				require.Error(t, err)
+				require.ErrorContains(t, err, "RETRY_ASYNC_WRITE_FAILURE")
+				// Transfer the lease to n3.
+				transferLease(2)
+				close(startTxn2)
+				// Block until Txn2 recovers us.
+				<-blockCh
+				return err
+			case 1:
+				// When retrying the write failure we should discover that txn recovery
+				// has aborted this transaction.
+				err := txn.Put(ctx, keyA, "value")
+				require.Error(t, err)
+				require.ErrorContains(t, err, "ABORT_REASON_ABORT_SPAN")
+				return err
+			default:
+				t.Errorf("unexpected epoch: %d", txn.Epoch())
+			}
+			return nil
+		})
+		require.NoError(t, err)
+	}()
+	<-startTxn2
+
+	txn2 := db.NewTxn(ctx, "txn2")
+	res, err := txn2.Get(ctx, keyA)
+	require.NoError(t, err)
+	// NB: Nothing should exist on keyA, because txn1 didn't commit at epoch 1 (or
+	// any epoch, for that matter).
+	require.False(t, res.Exists())
+
+	close(blockCh)
+	wg.Wait()
+}

pkg/kv/kvnemesis/kvnemesis_test.go

@@ -337,21 +337,14 @@ func TestKVNemesisMultiNode_BufferedWrites(t *testing.T) {
 	defer log.Scope(t).Close(t)
 
 	testKVNemesisImpl(t, kvnemesisTestCfg{
-		numNodes:     3,
-		numSteps:     defaultNumSteps,
-		concurrency:  5,
-		seedOverride: 0,
-		// TODO(#145458): Until #145458 is fixed reduce the
-		// rate of lost writes by avoiding lease transfers and
-		// merges and also turning off error injection.
-		invalidLeaseAppliedIndexProb: 0.0,
-		injectReproposalErrorProb:    0.0,
+		numNodes:                     3,
+		numSteps:                     defaultNumSteps,
+		concurrency:                  5,
+		seedOverride:                 0,
+		invalidLeaseAppliedIndexProb: 0.2,
+		injectReproposalErrorProb:    0.2,
 		assertRaftApply:              true,
 		bufferedWriteProb:            0.70,
-		testGeneratorConfig: func(g *GeneratorConfig) {
-			g.Ops.ChangeLease = ChangeLeaseConfig{}
-			g.Ops.Merge = MergeConfig{}
-		},
 		testSettings: func(ctx context.Context, st *cluster.Settings) {
 			kvcoord.BufferedWritesEnabled.Override(ctx, &st.SV, true)
 			concurrency.UnreplicatedLockReliabilityLeaseTransfer.Override(ctx, &st.SV, true)
@@ -362,9 +355,7 @@ func TestKVNemesisMultiNode_BufferedWrites(t *testing.T) {
 }
 
 // TestKVNemesisMultiNode_BufferedWritesNoPipelining turns on buffered
-// writes and turns off write pipelining. Turning off write pipelining
-// allows us to test the lock reliability features even without a fix
-// for #145458.
+// writes and turns off write pipelining.
 func TestKVNemesisMultiNode_BufferedWritesNoPipelining(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)

pkg/kv/kvserver/batcheval/cmd_query_intent.go

@@ -157,8 +157,11 @@ func QueryIntent(
 		}
 	}
 
+	res := result.Result{}
 	if !reply.FoundIntent && args.ErrorIfMissing {
-		return result.Result{}, kvpb.NewIntentMissingError(args.Key, intent)
+		l := roachpb.MakeLockAcquisition(args.Txn, args.Key, lock.Replicated, args.Strength, args.IgnoredSeqNums)
+		res.Local.ReportedMissingLocks = []roachpb.LockAcquisition{l}
+		return res, kvpb.NewIntentMissingError(args.Key, intent)
 	}
-	return result.Result{}, nil
+	return res, nil
 }

pkg/kv/kvserver/batcheval/result/result.go

@@ -38,6 +38,10 @@ type LocalResult struct {
 	// UpdatedTxns stores transaction records that have been updated by
 	// calls to EndTxn, PushTxn, and RecoverTxn.
 	UpdatedTxns []*roachpb.Transaction
+	// ReportedMissingLocks stores lock acquisition structs that represent locks
+	// that have been reported as missing via QueryIntent. Such locks must be
+	// reported to the concurrency manager.
+	ReportedMissingLocks []roachpb.LockAcquisition
 	// EndTxns stores completed transactions. If the transaction
 	// contains unresolved intents, they should be handed off for
 	// asynchronous intent resolution. A bool in each EndTxnIntents
@@ -131,6 +135,17 @@ func (lResult *LocalResult) DetachEncounteredIntents() []roachpb.Intent {
 	return r
 }
 
+// DetachMissingLocks returns (and removes) those locks that have been reported
+// missing during an QueryIntentRequest and must be handled.
+func (lResult *LocalResult) DetachMissingLocks() []roachpb.LockAcquisition {
+	if lResult == nil {
+		return nil
+	}
+	r := lResult.ReportedMissingLocks
+	lResult.ReportedMissingLocks = nil
+	return r
+}
+
 // DetachEndTxns returns (and removes) the EndTxnIntent objects from
 // the local result. If alwaysOnly is true, the slice is filtered to
 // include only those which have specified returnAlways=true, meaning
@@ -418,6 +433,13 @@ func (p *Result) MergeAndDestroy(q Result) error {
 	}
 	q.Local.ResolvedLocks = nil
 
+	if p.Local.ReportedMissingLocks == nil {
+		p.Local.ReportedMissingLocks = q.Local.ReportedMissingLocks
+	} else {
+		p.Local.ReportedMissingLocks = append(p.Local.ReportedMissingLocks, q.Local.ReportedMissingLocks...)
+	}
+	q.Local.ReportedMissingLocks = nil
+
 	if p.Local.UpdatedTxns == nil {
 		p.Local.UpdatedTxns = q.Local.UpdatedTxns
 	} else {

pkg/kv/kvserver/concurrency/concurrency_control.go

@@ -254,6 +254,11 @@ type LockManager interface {
 	// acquired a new lock or re-acquired an existing lock that it already held.
 	OnLockAcquired(context.Context, *roachpb.LockAcquisition)
 
+	// OnLockMissing informs the concurrency manager that a lock has been reported
+	// missing to a client via QueryIntent. Such locks cannot later be
+	// materialized via a lock table flush.
+	OnLockMissing(context.Context, *roachpb.LockAcquisition)
+
 	// OnLockUpdated informs the concurrency manager that a transaction has
 	// updated or released a lock or range of locks that it previously held.
 	// The Durability field of the lock update struct is ignored.
@@ -697,6 +702,11 @@ type lockTable interface {
 	// intent has been applied to the replicated state machine.
 	AcquireLock(*roachpb.LockAcquisition) error
 
+	// MarkIneligibleForExport marks any locks held by this transaction on the
+	// same key as ineligible for export from the lock table for replication since
+	// doing so could result in a transaction being erroneously committed.
+	MarkIneligibleForExport(*roachpb.LockAcquisition) error
+
 	// UpdateLocks informs the lockTable that an existing lock or range of locks
 	// was either updated or released.
 	//

pkg/kv/kvserver/concurrency/concurrency_manager.go

@@ -131,9 +131,7 @@ var UnreplicatedLockReliabilityLeaseTransfer = settings.RegisterBoolSetting(
 	settings.SystemOnly,
 	"kv.lock_table.unreplicated_lock_reliability.lease_transfer.enabled",
 	"whether the replica should attempt to keep unreplicated locks during lease transfers",
-	// TODO(#145458): We've disabled this by default to avoid flakes until the underlying bug is fixed.
-	// metamorphic.ConstantWithTestBool("kv.lock_table.unreplicated_lock_reliability.lease_transfer.enabled", true),
-	false,
+	metamorphic.ConstantWithTestBool("kv.lock_table.unreplicated_lock_reliability.lease_transfer.enabled", false),
 )
 
 // UnreplicatedLockReliabilityMerge controls whether the replica will
@@ -142,9 +140,7 @@ var UnreplicatedLockReliabilityMerge = settings.RegisterBoolSetting(
 	settings.SystemOnly,
 	"kv.lock_table.unreplicated_lock_reliability.merge.enabled",
 	"whether the replica should attempt to keep unreplicated locks during range merges",
-	// TODO(#145458): We've disabled this by default to avoid flakes until the underlying bug is fixed.
-	// metamorphic.ConstantWithTestBool("kv.lock_table.unreplicated_lock_reliability.merge.enabled", true),
-	false,
+	metamorphic.ConstantWithTestBool("kv.lock_table.unreplicated_lock_reliability.merge.enabled", false),
 )
 
 var MaxLockFlushSize = settings.RegisterByteSizeSetting(
@@ -599,6 +595,15 @@ func (m *managerImpl) OnLockAcquired(ctx context.Context, acq *roachpb.LockAcqui
 	}
 }
 
+// OnLockMissing implements the Lockmanager interface.
+func (m *managerImpl) OnLockMissing(ctx context.Context, acq *roachpb.LockAcquisition) {
+	if err := m.lt.MarkIneligibleForExport(acq); err != nil {
+		// We don't currently expect any errors other than assertion failures that represent
+		// programming errors from this method.
+		log.Fatalf(ctx, "%v", err)
+	}
+}
+
 // OnLockUpdated implements the LockManager interface.
 func (m *managerImpl) OnLockUpdated(ctx context.Context, up *roachpb.LockUpdate) {
 	if err := m.lt.UpdateLocks(up); err != nil {