kvserver: propagate ReadCategory to IterateReplicaKeySpans

Extend rditer.IterateReplicaKeySpans to accept a ReadCategory and propagate it
to the constructed engine iterator. Most significantly, this ensures that the
scan to read a Raft snapshot's data is appropriately labeled with the
fs.RangeSnapshotReadCategory category.

Informs #158017.
Epic: none
Release note: none

Author: jbowens

pkg/cli/debug.go

@@ -497,7 +497,7 @@ func runDebugRangeData(cmd *cobra.Command, args []string) error {
 	defer snapshot.Close()
 
 	var results int
-	return rditer.IterateReplicaKeySpans(cmd.Context(), &desc, snapshot, rditer.SelectOpts{
+	return rditer.IterateReplicaKeySpans(cmd.Context(), &desc, snapshot, fs.UnknownReadCategory, rditer.SelectOpts{
 		Ranged: rditer.SelectRangedOptions{
 			SystemKeys: true,
 			LockTable:  true,

pkg/kv/kvserver/client_merge_test.go

@@ -3913,7 +3913,7 @@ func TestStoreRangeMergeRaftSnapshot(t *testing.T) {
 				}
 			}
 
-			require.NoError(t, rditer.IterateReplicaKeySpans(ctx, inSnap.Desc, snapReader, rditer.SelectOpts{
+			require.NoError(t, rditer.IterateReplicaKeySpans(ctx, inSnap.Desc, snapReader, fs.ReplicationReadCategory, rditer.SelectOpts{
 				Ranged: rditer.SelectRangedOptions{
 					SystemKeys: true,
 					LockTable:  true,

pkg/kv/kvserver/kv_snapshot_strategy.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/storage/fs"
 	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
@@ -407,7 +408,7 @@ func (kvSS *kvBatchSnapshotStrategy) Send(
 		}
 		return err
 	}
-	if err := rditer.IterateReplicaKeySpans(ctx, snap.State.Desc, snap.EngineSnap, rditer.SelectOpts{
+	if err := rditer.IterateReplicaKeySpans(ctx, snap.State.Desc, snap.EngineSnap, fs.RangeSnapshotReadCategory, rditer.SelectOpts{
 		Ranged: rditer.SelectRangedOptions{
 			SystemKeys: true,
 			LockTable:  true,
@@ -539,7 +540,7 @@ func (kvSS *kvBatchSnapshotStrategy) Send(
 			//
 			// See: https://github.com/cockroachdb/cockroach/issues/142673
 			transitionFromSharedToRegularReplicate = true
-			err = rditer.IterateReplicaKeySpans(ctx, snap.State.Desc, snap.EngineSnap, rditer.SelectOpts{
+			err = rditer.IterateReplicaKeySpans(ctx, snap.State.Desc, snap.EngineSnap, fs.RangeSnapshotReadCategory, rditer.SelectOpts{
 				Ranged: rditer.SelectRangedOptions{
 					UserKeys: true,
 				},

pkg/kv/kvserver/raftlog/BUILD.bazel

@@ -22,6 +22,7 @@ go_library(
         "//pkg/roachpb",
         "//pkg/storage",
         "//pkg/storage/enginepb",
+        "//pkg/storage/fs",
         "//pkg/util/buildutil",
         "//pkg/util/encoding",
         "//pkg/util/iterutil",

pkg/kv/kvserver/rditer/replica_data_iter.go

@@ -421,6 +421,7 @@ func IterateReplicaKeySpans(
 	ctx context.Context,
 	desc *roachpb.RangeDescriptor,
 	reader storage.Reader,
+	readCategory fs.ReadCategory,
 	opts SelectOpts,
 	visitor func(storage.EngineIterator, roachpb.Span) error,
 ) error {
@@ -433,9 +434,10 @@ func IterateReplicaKeySpans(
 	for _, span := range spans {
 		err := func() error {
 			iter, err := reader.NewEngineIterator(ctx, storage.IterOptions{
-				KeyTypes:   storage.IterKeyTypePointsAndRanges,
-				LowerBound: span.Key,
-				UpperBound: span.EndKey,
+				KeyTypes:     storage.IterKeyTypePointsAndRanges,
+				LowerBound:   span.Key,
+				UpperBound:   span.EndKey,
+				ReadCategory: readCategory,
 			})
 			if err != nil {
 				return err

pkg/kv/kvserver/rditer/replica_data_iter_test.go

@@ -169,7 +169,7 @@ func verifyIterateReplicaKeySpans(
 	})
 
 	require.NoError(t, IterateReplicaKeySpans(
-		context.Background(), desc, readWriter, selOpts,
+		context.Background(), desc, readWriter, fs.UnknownReadCategory, selOpts,
 		func(iter storage.EngineIterator, span roachpb.Span) error {
 			var err error
 			for ok := true; ok && err == nil; ok, err = iter.NextEngineKey() {
@@ -497,7 +497,7 @@ func TestReplicaDataIteratorGlobalRangeKey(t *testing.T) {
 
 				var actualSpans []roachpb.Span
 				require.NoError(t, IterateReplicaKeySpans(
-					context.Background(), &desc, snapshot, selOpts,
+					context.Background(), &desc, snapshot, fs.UnknownReadCategory, selOpts,
 					func(iter storage.EngineIterator, span roachpb.Span) error {
 						// We should never see any point keys.
 						hasPoint, hasRange := iter.HasPointAndRange()
@@ -554,7 +554,7 @@ func BenchmarkReplicaEngineDataIterator(b *testing.B) {
 }
 
 func benchReplicaEngineDataIterator(b *testing.B, numRanges, numKeysPerRange, valueSize int) {
-	ctx := context.Background()
+	ctx := b.Context()
 
 	// Set up ranges.
 	var descs []roachpb.RangeDescriptor
@@ -586,7 +586,7 @@ func benchReplicaEngineDataIterator(b *testing.B, numRanges, numKeysPerRange, va
 	for _, desc := range descs {
 		var keyBuf roachpb.Key
 		keySpans := MakeAllKeySpans(&desc)
-		for i := 0; i < numKeysPerRange; i++ {
+		for i := range numKeysPerRange {
 			keyBuf = append(keyBuf[:0], keySpans[i%len(keySpans)].Key...)
 			keyBuf = append(keyBuf, 0, 0, 0, 0)
 			binary.BigEndian.PutUint32(keyBuf[len(keyBuf)-4:], uint32(i))
@@ -602,11 +602,9 @@ func benchReplicaEngineDataIterator(b *testing.B, numRanges, numKeysPerRange, va
 	snapshot := eng.NewSnapshot()
 	defer snapshot.Close()
 
-	b.ResetTimer()
-
-	for i := 0; i < b.N; i++ {
+	for b.Loop() {
 		for _, desc := range descs {
-			err := IterateReplicaKeySpans(context.Background(), &desc, snapshot,
+			err := IterateReplicaKeySpans(b.Context(), &desc, snapshot, fs.UnknownReadCategory,
 				SelectOpts{
 					Ranged: SelectRangedOptions{
 						RSpan:      desc.RSpan(),

pkg/kv/kvserver/replica_learner_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/storage/fs"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
@@ -2235,7 +2236,7 @@ func getExpectedSnapshotSizeBytes(
 		ReplicatedByRangeID:   true,
 		UnreplicatedByRangeID: false,
 	}
-	err = rditer.IterateReplicaKeySpans(ctx, snap.State.Desc, snap.EngineSnap, selOpts, func(iter storage.EngineIterator, _ roachpb.Span) error {
+	err = rditer.IterateReplicaKeySpans(ctx, snap.State.Desc, snap.EngineSnap, fs.ReplicationReadCategory, selOpts, func(iter storage.EngineIterator, _ roachpb.Span) error {
 		var err error
 		for ok := true; ok && err == nil; ok, err = iter.NextEngineKey() {
 			hasPoint, hasRange := iter.HasPointAndRange()

pkg/kv/kvserver/replica_test.go

@@ -58,6 +58,7 @@ import (
 	"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/storage/fs"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
@@ -7017,7 +7018,7 @@ func TestReplicaDestroy(t *testing.T) {
 		UnreplicatedByRangeID: true,
 	}
 	require.NoError(t, rditer.IterateReplicaKeySpans(
-		ctx, tc.repl.Desc(), engSnapshot, selOpts,
+		ctx, tc.repl.Desc(), engSnapshot, fs.ReplicationReadCategory, selOpts,
 		func(iter storage.EngineIterator, _ roachpb.Span) error {
 			var err error
 			for ok := true; ok && err == nil; ok, err = iter.NextEngineKey() {