metamorphic: add MVCC{CheckFor,}AcquireLock operations

pkg/storage/metamorphic/BUILD.bazel

@@ -16,6 +16,7 @@ go_library(
         "//pkg/base",
         "//pkg/keys",
         "//pkg/kv/kvpb",
+        "//pkg/kv/kvserver/concurrency/lock",
         "//pkg/roachpb",
         "//pkg/settings/cluster",
         "//pkg/storage",

pkg/storage/metamorphic/operations.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
@@ -280,6 +281,50 @@ func (m mvccInitPutOp) run(ctx context.Context) string {
 	return "ok"
 }
 
+type mvccCheckForAcquireLockOp struct {
+	m                *metaTestRunner
+	writer           readWriterID
+	key              roachpb.Key
+	txn              txnID
+	strength         lock.Strength
+	maxLockConflicts int
+}
+
+func (m mvccCheckForAcquireLockOp) run(ctx context.Context) string {
+	txn := m.m.getTxn(m.txn)
+	writer := m.m.getReadWriter(m.writer)
+
+	err := storage.MVCCCheckForAcquireLock(ctx, writer, txn, m.strength, m.key, int64(m.maxLockConflicts))
+	if err != nil {
+		return fmt.Sprintf("error: %s", err)
+	}
+
+	return "ok"
+}
+
+type mvccAcquireLockOp struct {
+	m                *metaTestRunner
+	writer           readWriterID
+	key              roachpb.Key
+	txn              txnID
+	strength         lock.Strength
+	maxLockConflicts int
+}
+
+func (m mvccAcquireLockOp) run(ctx context.Context) string {
+	txn := m.m.getTxn(m.txn)
+	writer := m.m.getReadWriter(m.writer)
+
+	err := storage.MVCCAcquireLock(ctx, writer, txn, m.strength, m.key, nil, int64(m.maxLockConflicts))
+	if err != nil {
+		return fmt.Sprintf("error: %s", err)
+	}
+
+	// Update the txn's lock spans to account for this lock being acquired.
+	addKeyToLockSpans(txn, m.key)
+	return "ok"
+}
+
 type mvccDeleteRangeOp struct {
 	m      *metaTestRunner
 	writer readWriterID
@@ -938,6 +983,69 @@ var opGenerators = []opGenerator{
 		},
 		weight: 50,
 	},
+	{
+		name: "mvcc_check_for_acquire_lock",
+		generate: func(ctx context.Context, m *metaTestRunner, args ...string) mvccOp {
+			writer := readWriterID(args[0])
+			txn := txnID(args[1])
+			key := m.txnKeyGenerator.parse(args[2])
+			strength := lock.Shared
+			if m.floatGenerator.parse(args[3]) < 0.5 {
+				strength = lock.Exclusive
+			}
+			maxLockConflicts := int(m.floatGenerator.parse(args[4]) * 1000)
+
+			return &mvccCheckForAcquireLockOp{
+				m:                m,
+				writer:           writer,
+				key:              key.Key,
+				txn:              txn,
+				strength:         strength,
+				maxLockConflicts: maxLockConflicts,
+			}
+		},
+		operands: []operandType{
+			operandReadWriter,
+			operandTransaction,
+			operandMVCCKey,
+			operandFloat,
+			operandFloat,
+		},
+		weight: 50,
+	},
+	{
+		name: "mvcc_acquire_lock",
+		generate: func(ctx context.Context, m *metaTestRunner, args ...string) mvccOp {
+			writer := readWriterID(args[0])
+			txn := txnID(args[1])
+			key := m.txnKeyGenerator.parse(args[2])
+			strength := lock.Shared
+			if m.floatGenerator.parse(args[3]) < 0.5 {
+				strength = lock.Exclusive
+			}
+			maxLockConflicts := int(m.floatGenerator.parse(args[4]) * 1000)
+
+			// Track this write in the txn generator. This ensures the batch will be
+			// committed before the transaction is committed.
+			m.txnGenerator.trackTransactionalWrite(writer, txn, key.Key, nil)
+			return &mvccAcquireLockOp{
+				m:                m,
+				writer:           writer,
+				key:              key.Key,
+				txn:              txn,
+				strength:         strength,
+				maxLockConflicts: maxLockConflicts,
+			}
+		},
+		operands: []operandType{
+			operandReadWriter,
+			operandTransaction,
+			operandMVCCKey,
+			operandFloat,
+			operandFloat,
+		},
+		weight: 50,
+	},
 	{
 		name: "mvcc_delete_range",
 		generate: func(ctx context.Context, m *metaTestRunner, args ...string) mvccOp {