storage: benchmark replicated lock acquisition

pkg/storage/bench_pebble_test.go

@@ -17,8 +17,10 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -514,6 +516,42 @@ func BenchmarkBatchApplyBatchRepr_Pebble(b *testing.B) {
 	}
 }
 
+func BenchmarkMVCCCheckForAcquireLock_Pebble(b *testing.B) {
+	defer log.Scope(b).Close(b)
+	ctx := context.Background()
+	testutils.RunTrueAndFalse(b, "batch", func(b *testing.B, batch bool) {
+		testutils.RunTrueAndFalse(b, "heldOtherTxn", func(b *testing.B, heldOtherTxn bool) {
+			testutils.RunTrueAndFalse(b, "heldSameTxn", func(b *testing.B, heldSameTxn bool) {
+				if heldOtherTxn && heldSameTxn {
+					return // not possible
+				}
+				strs := []lock.Strength{lock.Shared, lock.Exclusive}
+				testutils.RunValues(b, "strength", strs, func(b *testing.B, strength lock.Strength) {
+					runMVCCCheckForAcquireLock(ctx, b, setupMVCCInMemPebble, batch, heldOtherTxn, heldSameTxn, strength)
+				})
+			})
+		})
+	})
+}
+
+func BenchmarkMVCCAcquireLock_Pebble(b *testing.B) {
+	defer log.Scope(b).Close(b)
+	ctx := context.Background()
+	testutils.RunTrueAndFalse(b, "batch", func(b *testing.B, batch bool) {
+		testutils.RunTrueAndFalse(b, "heldOtherTxn", func(b *testing.B, heldOtherTxn bool) {
+			testutils.RunTrueAndFalse(b, "heldSameTxn", func(b *testing.B, heldSameTxn bool) {
+				if heldOtherTxn && heldSameTxn {
+					return // not possible
+				}
+				strs := []lock.Strength{lock.Shared, lock.Exclusive}
+				testutils.RunValues(b, "strength", strs, func(b *testing.B, strength lock.Strength) {
+					runMVCCAcquireLock(ctx, b, setupMVCCInMemPebble, batch, heldOtherTxn, heldSameTxn, strength)
+				})
+			})
+		})
+	})
+}
+
 func BenchmarkBatchBuilderPut(b *testing.B) {
 	defer log.Scope(b).Close(b)
 	value := make([]byte, 10)

pkg/storage/bench_test.go

@@ -24,6 +24,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/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
@@ -1580,6 +1581,107 @@ func runBatchApplyBatchRepr(
 	b.StopTimer()
 }
 
+func runMVCCCheckForAcquireLock(
+	ctx context.Context,
+	b *testing.B,
+	emk engineMaker,
+	useBatch bool,
+	heldOtherTxn bool,
+	heldSameTxn bool,
+	strength lock.Strength,
+) {
+	runMVCCAcquireLockCommon(ctx, b, emk, useBatch, heldOtherTxn, heldSameTxn, strength, true /* checkFor */)
+}
+
+func runMVCCAcquireLock(
+	ctx context.Context,
+	b *testing.B,
+	emk engineMaker,
+	useBatch bool,
+	heldOtherTxn bool,
+	heldSameTxn bool,
+	strength lock.Strength,
+) {
+	runMVCCAcquireLockCommon(ctx, b, emk, useBatch, heldOtherTxn, heldSameTxn, strength, false /* checkFor */)
+}
+
+func runMVCCAcquireLockCommon(
+	ctx context.Context,
+	b *testing.B,
+	emk engineMaker,
+	useBatch bool,
+	heldOtherTxn bool,
+	heldSameTxn bool,
+	strength lock.Strength,
+	checkFor bool,
+) {
+	if heldOtherTxn && heldSameTxn {
+		b.Fatalf("heldOtherTxn and heldSameTxn cannot both be true")
+	}
+
+	keyBuf := append(make([]byte, 0, 64), []byte("key-")...)
+	makeKey := func(i int) roachpb.Key {
+		// NOTE: we're appending to a buffer with sufficient capacity, so this does
+		// not allocate, but as a result, we need to watch out for aliasing bugs.
+		return encoding.EncodeUvarintAscending(keyBuf[:4], uint64(i))
+	}
+	makeTxn := func(name string) roachpb.Transaction {
+		return roachpb.MakeTransaction(name, keyBuf, 0, 0, hlc.Timestamp{WallTime: 1}, 0, 0, 0)
+	}
+	txn1 := makeTxn("txn1")
+	txn2 := makeTxn("txn2")
+
+	loc := "acquire_lock"
+	if checkFor {
+		loc = "check_for_acquire_lock"
+	}
+	eng := emk(b, loc)
+	defer eng.Close()
+
+	for i := 0; i < b.N; i++ {
+		key := makeKey(i)
+		if heldOtherTxn || heldSameTxn {
+			txn := &txn1
+			if heldOtherTxn {
+				txn = &txn2
+			}
+			// Acquire a shared and an exclusive lock on the key.
+			err := MVCCAcquireLock(ctx, eng, txn, lock.Shared, key, nil, 0)
+			require.NoError(b, err)
+			err = MVCCAcquireLock(ctx, eng, txn, lock.Exclusive, key, nil, 0)
+			require.NoError(b, err)
+		}
+	}
+
+	rw := ReadWriter(eng)
+	if useBatch {
+		batch := eng.NewBatch()
+		defer batch.Close()
+		rw = batch
+	}
+	ms := &enginepb.MVCCStats{}
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		key := makeKey(i)
+		txn := &txn1
+		var err error
+		if checkFor {
+			err = MVCCCheckForAcquireLock(ctx, rw, txn, strength, key, 0)
+		} else {
+			err = MVCCAcquireLock(ctx, rw, txn, strength, key, ms, 0)
+		}
+		if heldOtherTxn {
+			require.Error(b, err)
+		} else {
+			require.NoError(b, err)
+		}
+	}
+
+	b.StopTimer()
+}
+
 type mvccExportToSSTOpts struct {
 	numKeys, numRevisions, numRangeKeys     int
 	exportAllRevisions, useElasticCPUHandle bool

pkg/testutils/subtest.go

@@ -10,25 +10,29 @@
 
 package testutils
 
-import (
-	"fmt"
-	"testing"
-)
+import "fmt"
 
 // RunTrueAndFalse calls the provided function in a subtest, first with the
 // boolean argument set to false and next with the boolean argument set to true.
-func RunTrueAndFalse(t *testing.T, name string, fn func(t *testing.T, b bool)) {
+func RunTrueAndFalse[T testingTB[T]](t T, name string, fn func(t T, b bool)) {
 	t.Helper()
 	RunValues(t, name, []bool{false, true}, fn)
 }
 
 // RunValues calls the provided function in a subtest for each of the
 // provided values.
-func RunValues[T any](t *testing.T, name string, values []T, fn func(*testing.T, T)) {
+func RunValues[T testingTB[T], V any](t T, name string, values []V, fn func(T, V)) {
 	t.Helper()
 	for _, v := range values {
-		t.Run(fmt.Sprintf("%s=%v", name, v), func(t *testing.T) {
+		t.Run(fmt.Sprintf("%s=%v", name, v), func(t T) {
 			fn(t, v)
 		})
 	}
 }
+
+// testingTB is an interface that matches *testing.T and *testing.B, without
+// incurring the package dependency.
+type testingTB[T any] interface {
+	Run(name string, f func(t T)) bool
+	Helper()
+}