kvserver: don't generate large writes during lock durability upgrade

Since this is still a best-effort feature for the time being, we want
to protect against writing too many locks during transfer or merge
requests.

Epic: none
Release note: None

Author: stevendanna

pkg/kv/kvserver/batcheval/cmd_lease.go

@@ -9,6 +9,7 @@ import (
 	"context"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/batcheval/result"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/readsummary"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/readsummary/rspb"
@@ -69,17 +70,26 @@ func evalNewLease(
 			localReadSum := rec.GetCurrentReadSummary(ctx)
 			priorReadSum = &localReadSum
 
+			durabilityUpgradeLimit := concurrency.GetMaxLockFlushSize(&rec.ClusterSettings().SV)
 			// If this is a lease transfer, we write out all unreplicated leases to
 			// storage, so that any waiters will discover them on the new leaseholder.
-			acquisitions := rec.GetConcurrencyManager().OnRangeLeaseTransferEval()
+			acquisitions, approxSize := rec.GetConcurrencyManager().OnRangeLeaseTransferEval()
 			log.VEventf(ctx, 2, "upgrading durability of %d locks", len(acquisitions))
-			for _, acq := range acquisitions {
-				if err := storage.MVCCAcquireLock(ctx, readWriter,
-					&acq.Txn, acq.IgnoredSeqNums, acq.Strength, acq.Key, ms, 0, 0); err != nil {
-					return newFailedLeaseTrigger(isTransfer), err
+			if approxSize > durabilityUpgradeLimit {
+				log.Warningf(ctx,
+					"refusing to upgrade lock durability of %d locks since approximate lock size of %d byte exceeds %d bytes",
+					len(acquisitions),
+					approxSize,
+					durabilityUpgradeLimit)
+			} else {
+				for _, acq := range acquisitions {
+					if err := storage.MVCCAcquireLock(ctx, readWriter,
+						&acq.Txn, acq.IgnoredSeqNums, acq.Strength, acq.Key, ms, 0, 0); err != nil {
+						return newFailedLeaseTrigger(isTransfer), err
+					}
 				}
+				locksWritten = len(acquisitions)
 			}
-			locksWritten = len(acquisitions)
 		} else {
 			// If the new lease is not equivalent to the old lease (i.e. either the
 			// lease is changing hands or the leaseholder restarted), construct a

pkg/kv/kvserver/batcheval/cmd_subsume.go

@@ -14,6 +14,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/batcheval/result"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/lockspanset"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/spanset"
 	"github.com/cockroachdb/cockroach/pkg/storage"
@@ -185,19 +186,28 @@ func Subsume(
 	// (*Store).MergeRange.
 	stats := cArgs.EvalCtx.GetMVCCStats()
 	if args.PreserveUnreplicatedLocks {
-		acquisitions := cArgs.EvalCtx.GetConcurrencyManager().OnRangeSubsumeEval()
-		log.VEventf(ctx, 2, "upgrading durability of %d locks", len(acquisitions))
-		statsDelta := enginepb.MVCCStats{}
-		for _, acq := range acquisitions {
-			if err := storage.MVCCAcquireLock(ctx, readWriter,
-				&acq.Txn, acq.IgnoredSeqNums, acq.Strength, acq.Key, &statsDelta, 0, 0); err != nil {
-				return result.Result{}, err
+		durabilityUpgradeLimit := concurrency.GetMaxLockFlushSize(&cArgs.EvalCtx.ClusterSettings().SV)
+		acquisitions, approxSize := cArgs.EvalCtx.GetConcurrencyManager().OnRangeSubsumeEval()
+		if approxSize > durabilityUpgradeLimit {
+			log.Warningf(ctx,
+				"refusing to upgrade lock durability of %d locks since approximate lock size of %d byte exceeds %d bytes",
+				len(acquisitions),
+				approxSize,
+				durabilityUpgradeLimit)
+		} else {
+			log.VEventf(ctx, 2, "upgrading durability of %d locks", len(acquisitions))
+			statsDelta := enginepb.MVCCStats{}
+			for _, acq := range acquisitions {
+				if err := storage.MVCCAcquireLock(ctx, readWriter,
+					&acq.Txn, acq.IgnoredSeqNums, acq.Strength, acq.Key, &statsDelta, 0, 0); err != nil {
+					return result.Result{}, err
+				}
 			}
+			// Apply the stats delta to both the stats snapshot we are sending in the
+			// response and to the stats update we expect as part of this proposal.
+			stats.Add(statsDelta)
+			cArgs.Stats.Add(statsDelta)
 		}
-		// Apply the stats delta to both the stats snapshot we are sending in the
-		// response and to the stats update we expect as part of this proposal.
-		stats.Add(statsDelta)
-		cArgs.Stats.Add(statsDelta)
 	}
 
 	// Now that the range is frozen, collect some information to ship to the LHS

pkg/kv/kvserver/client_replica_test.go

@@ -5911,6 +5911,115 @@ func TestLeaseTransferReplicatesLocks(t *testing.T) {
 	require.NoError(t, g.Wait())
 }
 
+func TestLeaseTransferDropsLocksIfLargerThanCommandSize(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	require.NoError(t, log.SetVModule("cmd_lease=2"))
+
+	// Test Plan:
+	//
+	// - Reduce MaxRaftCommandSize
+	// - Move scratch range to known location.
+	// - Take out a large number of unreplicated locks
+	// - Transfer lease without an error
+	//
+	ctx := context.Background()
+	st := cluster.MakeClusterSettings()
+	concurrency.UnreplicatedLockReliabilityLeaseTransfer.Override(ctx, &st.SV, true)
+	kvserverbase.MaxCommandSize.Override(ctx, &st.SV, 1<<20)
+	// To see the test fail:
+	// concurrency.MaxLockFlushSize.Override(ctx, &st.SV, 2<<20)
+
+	tc := testcluster.StartTestCluster(t, 3, base.TestClusterArgs{
+		ServerArgs: base.TestServerArgs{Settings: st},
+	})
+	defer tc.Stopper().Stop(ctx)
+
+	scratch := tc.ScratchRange(t)
+	desc, err := tc.LookupRange(scratch)
+	require.NoError(t, err)
+
+	// Start with the lease on store 1.
+	t.Logf("transfering to s1")
+	require.NoError(t, tc.TransferRangeLease(desc, tc.Target(0)))
+	t.Logf("done transfering to s1")
+
+	mkRandomScratchKey := func() roachpb.Key {
+		return append(scratch.Clone(), uuid.MakeV4().String()...)
+	}
+
+	numLocks := 9000
+	txn := tc.Server(1).DB().NewTxn(ctx, "test-lots-o-locks")
+	b := txn.NewBatch()
+	for range numLocks {
+		b.GetForUpdate(mkRandomScratchKey(), kvpb.BestEffort)
+		b.Requests()[len(b.Requests())-1].GetGet().LockNonExisting = true
+	}
+	require.NoError(t, txn.Run(ctx, b))
+
+	t.Log("transfering lease from s1 -> s2")
+	require.NoError(t, tc.TransferRangeLease(desc, tc.Target(1)))
+}
+
+func TestMergeDropsLocksIfLargerThanMax(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	require.NoError(t, log.SetVModule("cmd_subsume=2"))
+
+	// Test Plan:
+	//
+	// - Reduce MaxRaftCommandSize
+	// - Move scratch range to known location.
+	// - Take out a large number of unreplicated locks
+	// - Merge range without an error
+	//
+	var (
+		splitPoint = "b"
+		ctx        = context.Background()
+		st         = cluster.MakeClusterSettings()
+	)
+
+	concurrency.UnreplicatedLockReliabilityMerge.Override(ctx, &st.SV, true)
+	kvserverbase.MaxCommandSize.Override(ctx, &st.SV, 1<<20)
+	// To see the test fail:
+	// concurrency.MaxLockFlushSize.Override(ctx, &st.SV, 2<<20)
+
+	tc := testcluster.StartTestCluster(t, 1, base.TestClusterArgs{
+		ServerArgs: base.TestServerArgs{Settings: st},
+	})
+	defer tc.Stopper().Stop(ctx)
+
+	scratch := tc.ScratchRange(t)
+
+	mkKey := func(s string) roachpb.Key {
+		prefix := scratch.Clone()
+		return append(prefix[:len(prefix):len(prefix)], s...)
+	}
+
+	splitKey := mkKey(splitPoint)
+	tc.SplitRangeOrFatal(t, splitKey)
+
+	mkRandomScratchKey := func() roachpb.Key {
+		return append(mkKey(splitPoint), uuid.MakeV4().String()...)
+	}
+
+	numLocks := 6000
+	txn := tc.Server(0).DB().NewTxn(ctx, "test-lots-o-locks")
+	b := txn.NewBatch()
+	for range numLocks {
+		b.GetForUpdate(mkRandomScratchKey(), kvpb.BestEffort)
+		b.Requests()[len(b.Requests())-1].GetGet().LockNonExisting = true
+	}
+	require.NoError(t, txn.Run(ctx, b))
+
+	// Merge Range
+	t.Logf("merging range %s", scratch)
+	_, err := tc.MergeRanges(scratch)
+	require.NoError(t, err)
+}
+
 func TestMergeReplicatesLocks(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)

pkg/kv/kvserver/concurrency/BUILD.bazel

@@ -23,13 +23,15 @@ go_library(
         "//pkg/kv/kvserver/concurrency/lock",
         "//pkg/kv/kvserver/concurrency/poison",
         "//pkg/kv/kvserver/intentresolver",
+        "//pkg/kv/kvserver/kvserverbase",
         "//pkg/kv/kvserver/lockspanset",
         "//pkg/kv/kvserver/spanlatch",
         "//pkg/kv/kvserver/spanset",
         "//pkg/kv/kvserver/txnwait",
         "//pkg/roachpb",
         "//pkg/settings",
         "//pkg/settings/cluster",
+        "//pkg/storage",
         "//pkg/storage/enginepb",
         "//pkg/util/buildutil",
         "//pkg/util/container/list",

pkg/kv/kvserver/concurrency/concurrency_control.go

@@ -290,11 +290,11 @@ type RangeStateListener interface {
 	// may want to flush to disk as replicated. Since lease transfers declare
 	// latches that conflict with all requests, the caller knows that nothing is
 	// going to modify the lock table as its evaluating.
-	OnRangeLeaseTransferEval() []*roachpb.LockAcquisition
+	OnRangeLeaseTransferEval() ([]*roachpb.LockAcquisition, int64)
 
 	// OnRangeSubsumeEval informs the concurrency manager that the range is
 	// evaluating a merge.
-	OnRangeSubsumeEval() []*roachpb.LockAcquisition
+	OnRangeSubsumeEval() ([]*roachpb.LockAcquisition, int64)
 
 	// OnRangeLeaseUpdated informs the concurrency manager that its range's
 	// lease has been updated. The argument indicates whether this manager's

pkg/kv/kvserver/concurrency/concurrency_manager.go

@@ -14,13 +14,15 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/lockspanset"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/spanlatch"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/spanset"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/txnwait"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/util/debugutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -141,6 +143,24 @@ var UnreplicatedLockReliabilityMerge = settings.RegisterBoolSetting(
 	metamorphic.ConstantWithTestBool("kv.lock_table.unreplicated_lock_reliability.merge.enabled", true),
 )
 
+var MaxLockFlushSize = settings.RegisterByteSizeSetting(
+	settings.SystemOnly,
+	"kv.lock_table.unreplicated_lock_reliability.max_flush_size",
+	"maximum size of locks that will be flushed during merge and transfer operations (if 0, defaults to half of the MaxCommandSizeDefault)",
+	0,
+)
+
+// MaxLockFlushSize is the maximum number of lock bytes that we will attempt to
+// flush during merge and transfer operations.
+func GetMaxLockFlushSize(sv *settings.Values) int64 {
+	s := MaxLockFlushSize.Get(sv)
+	if s > 0 {
+		return s
+	} else {
+		return kvserverbase.MaxCommandSize.Get(sv) / 2
+	}
+}
+
 // managerImpl implements the Manager interface.
 type managerImpl struct {
 	st *cluster.Settings
@@ -607,33 +627,23 @@ func (m *managerImpl) OnRangeDescUpdated(desc *roachpb.RangeDescriptor) {
 var allKeysSpan = roachpb.Span{Key: keys.MinKey, EndKey: keys.MaxKey}
 
 // OnRangeLeaseTransferEval implements the RangeStateListener interface.
-func (m *managerImpl) OnRangeLeaseTransferEval() []*roachpb.LockAcquisition {
+func (m *managerImpl) OnRangeLeaseTransferEval() ([]*roachpb.LockAcquisition, int64) {
 	if !UnreplicatedLockReliabilityLeaseTransfer.Get(&m.st.SV) {
-		return nil
+		return nil, 0
 	}
 
-	// TODO(ssd): Expose a function that allows us to pre-allocate this a bit better.
-	acquistions := make([]*roachpb.LockAcquisition, 0)
-	m.lt.ExportUnreplicatedLocks(allKeysSpan, func(acq *roachpb.LockAcquisition) {
-		acquistions = append(acquistions, acq)
-	})
-	return acquistions
+	return m.exportUnreplicatedLocks()
 }
 
 // OnRangeSubumeEval implements the RangeStateListener interface. It is called
 // during evalutation of Subsume. The returned LockAcquisition structs represent
 // held locks that we may want to flush to disk as replicated.
-func (m *managerImpl) OnRangeSubsumeEval() []*roachpb.LockAcquisition {
+func (m *managerImpl) OnRangeSubsumeEval() ([]*roachpb.LockAcquisition, int64) {
 	if !UnreplicatedLockReliabilityMerge.Get(&m.st.SV) {
-		return nil
+		return nil, 0
 	}
 
-	// TODO(ssd): Expose a function that allows us to pre-allocate this a bit better.
-	acquistions := make([]*roachpb.LockAcquisition, 0)
-	m.lt.ExportUnreplicatedLocks(allKeysSpan, func(acq *roachpb.LockAcquisition) {
-		acquistions = append(acquistions, acq)
-	})
-	return acquistions
+	return m.exportUnreplicatedLocks()
 }
 
 // OnRangeLeaseUpdated implements the RangeStateListener interface.
@@ -706,6 +716,17 @@ func (m *managerImpl) LockTableMetrics() LockTableMetrics {
 	return m.lt.Metrics()
 }
 
+func (m *managerImpl) exportUnreplicatedLocks() ([]*roachpb.LockAcquisition, int64) {
+	// TODO(ssd): Expose a function that allows us to pre-allocate this a bit better.
+	approximateBatchSize := int64(0)
+	acquistions := make([]*roachpb.LockAcquisition, 0)
+	m.lt.ExportUnreplicatedLocks(allKeysSpan, func(acq *roachpb.LockAcquisition) {
+		approximateBatchSize += storage.ApproximateLockTableSize(acq)
+		acquistions = append(acquistions, acq)
+	})
+	return acquistions, approximateBatchSize
+}
+
 // TestingLockTableString implements the MetricExporter interface.
 func (m *managerImpl) TestingLockTableString() string {
 	return m.lt.String()

pkg/kv/kvserver/concurrency/concurrency_manager_test.go

@@ -545,7 +545,7 @@ func TestConcurrencyManagerBasic(t *testing.T) {
 
 			case "on-lease-transfer-eval":
 				mon.runSync("eval transfer lease", func(ctx context.Context) {
-					locksToWrite := m.OnRangeLeaseTransferEval()
+					locksToWrite, _ := m.OnRangeLeaseTransferEval()
 					if len(locksToWrite) > 0 {
 						log.Eventf(ctx, "locks to propose as replicated: %d", len(locksToWrite))
 					}

pkg/storage/mvcc.go

@@ -6319,6 +6319,18 @@ func MVCCVerifyLock(
 	return false, nil
 }
 
+var didUpdate bool
+
+func ApproximateLockTableSize(acq *roachpb.LockAcquisition) int64 {
+	keySize := int64(len(acq.Key)) + engineKeyVersionLockTableLen
+	metaSize := int64((&enginepb.MVCCMetadata{
+		Txn:                 &acq.Txn,
+		Timestamp:           acq.Txn.WriteTimestamp.ToLegacyTimestamp(),
+		TxnDidNotUpdateMeta: &didUpdate,
+	}).Size())
+	return keySize + metaSize
+}
+
 // mvccReleaseLockInternal releases a lock at the specified key and strength and
 // by the specified transaction. The function accepts the instructions for how
 // to release the lock (encoded in the LockUpdate), and the current value of the