From 0573d37425f474b4245a971e084e5a8805aaa928 Mon Sep 17 00:00:00 2001
From: Bilal Akhtar <bilal@cockroachlabs.com>
Date: Thu, 2 Nov 2023 15:04:36 -0400
Subject: [PATCH] metamorphic: add MVCC{CheckFor,}AcquireLock operations

This change adds two new operations to the MVCC metamorphic
testing package that were added in #110323. We check for
determinism in the storage package by adding these operations
to the metamorphic tests as well.

Fixes #109650.

Epic: none

Release note: None
---
 pkg/storage/metamorphic/BUILD.bazel   |   1 +
 pkg/storage/metamorphic/operations.go | 108 ++++++++++++++++++++++++++
 2 files changed, 109 insertions(+)

diff --git a/pkg/storage/metamorphic/BUILD.bazel b/pkg/storage/metamorphic/BUILD.bazel
index c1bd89caabed..1947ab919412 100644
--- a/pkg/storage/metamorphic/BUILD.bazel
+++ b/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",
diff --git a/pkg/storage/metamorphic/operations.go b/pkg/storage/metamorphic/operations.go
index 19b9ba4bd1b3..386f41bb19b2 100644
--- a/pkg/storage/metamorphic/operations.go
+++ b/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,56 @@ 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
+}
+
+func (m mvccCheckForAcquireLockOp) run(ctx context.Context) string {
+	txn := m.m.getTxn(m.txn)
+	txn.Sequence++
+	writer := m.m.getReadWriter(m.writer)
+
+	err := storage.MVCCCheckForAcquireLock(ctx, writer, txn, m.strength, m.key, 64)
+	if err != nil {
+		return fmt.Sprintf("error: %s", err)
+	}
+
+	// Update the txn's lock spans to account for this intent being written.
+	return "ok"
+}
+
+type mvccAcquireLockOp struct {
+	m        *metaTestRunner
+	writer   readWriterID
+	key      roachpb.Key
+	txn      txnID
+	strength lock.Strength
+}
+
+func (m mvccAcquireLockOp) run(ctx context.Context) string {
+	txn := m.m.getTxn(m.txn)
+	txn.Sequence++
+	writer := m.m.getReadWriter(m.writer)
+
+	err := storage.MVCCAcquireLock(ctx, writer, txn, m.strength, m.key, nil, 64)
+	if err != nil {
+		if writeTooOldErr := (*kvpb.WriteTooOldError)(nil); errors.As(err, &writeTooOldErr) {
+			txn.WriteTimestamp.Forward(writeTooOldErr.ActualTimestamp)
+			// Update the txn's lock spans to account for this intent being written.
+			addKeyToLockSpans(txn, m.key)
+		}
+		return fmt.Sprintf("error: %s", err)
+	}
+
+	// Update the txn's lock spans to account for this intent being written.
+	addKeyToLockSpans(txn, m.key)
+	return "ok"
+}
+
 type mvccDeleteRangeOp struct {
 	m      *metaTestRunner
 	writer readWriterID
@@ -938,6 +989,63 @@ var opGenerators = []opGenerator{
 		},
 		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
+			}
+
+			// 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,
+			}
+		},
+		operands: []operandType{
+			operandReadWriter,
+			operandTransaction,
+			operandMVCCKey,
+			operandFloat,
+		},
+		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
+			}
+
+			return &mvccCheckForAcquireLockOp{
+				m:        m,
+				writer:   writer,
+				key:      key.Key,
+				txn:      txn,
+				strength: strength,
+			}
+		},
+		operands: []operandType{
+			operandReadWriter,
+			operandTransaction,
+			operandMVCCKey,
+			operandFloat,
+		},
+		weight: 50,
+	},
 	{
 		name: "mvcc_delete_range",
 		generate: func(ctx context.Context, m *metaTestRunner, args ...string) mvccOp {