storage: improve BenchmarkUpdateSSTTimestamps

This patch improves `BenchmarkUpdateSSTTimestamps`, by running separate
benchmarks varying the number of keys and concurrency. In particular,
this exercises the `sstable.RewriteKeySuffixes` fast path, while the old
benchmark only used the naïve read/write slow path.

Release note: None

pkg/storage/sst_test.go

@@ -15,12 +15,12 @@ import (
 	"encoding/binary"
 	"fmt"
 	"math/rand"
-	"os"
-	"runtime/pprof"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
+	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -104,74 +104,62 @@ func TestCheckSSTConflictsMaxIntents(t *testing.T) {
 }
 
 func BenchmarkUpdateSSTTimestamps(b *testing.B) {
-	const (
-		modeZero    = iota + 1 // all zeroes
-		modeCounter            // uint64 counter in first 8 bytes
-		modeRandom             // random values
-
-		concurrency = 0 // 0 uses naïve replacement
-		sstSize     = 0
-		keyCount    = 500000
-		valueSize   = 8
-		valueMode   = modeRandom
-		profile     = false // cpuprofile.pprof
-	)
-
-	if sstSize > 0 && keyCount > 0 {
-		b.Fatal("Can't set both sstSize and keyCount")
+	defer log.Scope(b).Close(b)
+	skip.UnderShort(b)
+
+	ctx := context.Background()
+
+	for _, numKeys := range []int{1, 10, 100, 1000, 10000, 100000} {
+		b.Run(fmt.Sprintf("numKeys=%d", numKeys), func(b *testing.B) {
+			for _, concurrency := range []int{0, 1, 2, 4, 8} { // 0 uses naïve read/write loop
+				b.Run(fmt.Sprintf("concurrency=%d", concurrency), func(b *testing.B) {
+					runUpdateSSTTimestamps(ctx, b, numKeys, concurrency)
+				})
+			}
+		})
 	}
+}
 
-	b.StopTimer()
+func runUpdateSSTTimestamps(ctx context.Context, b *testing.B, numKeys int, concurrency int) {
+	const valueSize = 8
 
 	r := rand.New(rand.NewSource(7))
-
-	ctx := context.Background()
 	st := cluster.MakeTestingClusterSettings()
 	sstFile := &MemFile{}
 	writer := MakeIngestionSSTWriter(ctx, st, sstFile)
 	defer writer.Close()
 
+	sstTimestamp := hlc.MinTimestamp
+	reqTimestamp := hlc.Timestamp{WallTime: 1634899098417970999, Logical: 9}
+
 	key := make([]byte, 8)
 	value := make([]byte, valueSize)
-	sstTimestamp := hlc.Timestamp{WallTime: 1}
-	var i uint64
-	for i = 0; (keyCount > 0 && i < keyCount) || (sstSize > 0 && sstFile.Len() < sstSize); i++ {
-		binary.BigEndian.PutUint64(key, i)
-
-		switch valueMode {
-		case modeZero:
-		case modeCounter:
-			binary.BigEndian.PutUint64(value, i)
-		case modeRandom:
-			r.Read(value)
-		default:
-			b.Fatalf("unknown value mode %d", valueMode)
-		}
-
-		var v MVCCValue
-		v.Value.SetBytes(value)
-		v.Value.InitChecksum(key)
+	for i := 0; i < numKeys; i++ {
+		binary.BigEndian.PutUint64(key, uint64(i))
+		r.Read(value)
 
-		require.NoError(b, writer.PutMVCC(MVCCKey{Key: key, Timestamp: sstTimestamp}, v))
-	}
-	writer.Close()
-	b.Logf("%vMB %v keys", sstFile.Len()/1e6, i)
-
-	if profile {
-		f, err := os.Create("cpuprofile.pprof")
-		require.NoError(b, err)
-		defer f.Close()
+		var mvccValue MVCCValue
+		mvccValue.Value.SetBytes(value)
+		mvccValue.Value.InitChecksum(key)
 
-		require.NoError(b, pprof.StartCPUProfile(f))
-		defer pprof.StopCPUProfile()
+		if err := writer.PutMVCC(MVCCKey{Key: key, Timestamp: sstTimestamp}, mvccValue); err != nil {
+			require.NoError(b, err) // for performance
+		}
 	}
+	require.NoError(b, writer.Finish())
 
-	requestTimestamp := hlc.Timestamp{WallTime: 1634899098417970999, Logical: 9}
+	b.SetBytes(int64(numKeys * (len(key) + len(value))))
+	b.ResetTimer()
 
-	b.StartTimer()
+	var res []byte
 	for i := 0; i < b.N; i++ {
-		_, _, err := UpdateSSTTimestamps(
-			ctx, st, sstFile.Bytes(), sstTimestamp, requestTimestamp, concurrency, nil /* stats */)
-		require.NoError(b, err)
+		var ms enginepb.MVCCStats
+		var err error
+		res, _, err = UpdateSSTTimestamps(
+			ctx, st, sstFile.Bytes(), sstTimestamp, reqTimestamp, concurrency, &ms)
+		if err != nil {
+			require.NoError(b, err) // for performance
+		}
 	}
+	_ = res
 }