diff --git a/pkg/kv/kvserver/rditer/BUILD.bazel b/pkg/kv/kvserver/rditer/BUILD.bazel index 35c4fe2f1e28..d1120aae5f2a 100644 --- a/pkg/kv/kvserver/rditer/BUILD.bazel +++ b/pkg/kv/kvserver/rditer/BUILD.bazel @@ -29,6 +29,7 @@ go_test( srcs = [ "replica_data_iter_test.go", "select_test.go", + "stats_test.go", ], data = glob(["testdata/**"]), embed = [":rditer"], @@ -39,12 +40,14 @@ go_test( "//pkg/roachpb", "//pkg/settings/cluster", "//pkg/storage", + "//pkg/storage/enginepb", "//pkg/testutils", "//pkg/testutils/datapathutils", "//pkg/testutils/echotest", "//pkg/testutils/skip", "//pkg/util/hlc", "//pkg/util/leaktest", + "//pkg/util/protoutil", "//pkg/util/randutil", "//pkg/util/uuid", "@com_github_olekukonko_tablewriter//:tablewriter", diff --git a/pkg/kv/kvserver/rditer/replica_data_iter.go b/pkg/kv/kvserver/rditer/replica_data_iter.go index aafbafc028cd..aa7ba2ee24b7 100644 --- a/pkg/kv/kvserver/rditer/replica_data_iter.go +++ b/pkg/kv/kvserver/rditer/replica_data_iter.go @@ -137,6 +137,25 @@ func MakeReplicatedKeySpanSet(d *roachpb.RangeDescriptor) *spanset.SpanSet { return ss } +// MakeReplicatedKeySpansUserOnly returns all key spans corresponding to user +// keys. +func MakeReplicatedKeySpansUserOnly(d *roachpb.RangeDescriptor) []roachpb.Span { + return Select(d.RangeID, SelectOpts{ + ReplicatedBySpan: d.RSpan(), + ReplicatedSpansFilter: ReplicatedSpansUserOnly, + }) +} + +// MakeReplicatedKeySpansExcludingUser returns all key spans corresponding to +// non-user keys. +func MakeReplicatedKeySpansExcludingUser(d *roachpb.RangeDescriptor) []roachpb.Span { + return Select(d.RangeID, SelectOpts{ + ReplicatedBySpan: d.RSpan(), + ReplicatedByRangeID: true, + ReplicatedSpansFilter: ReplicatedSpansExcludeUser, + }) +} + // makeReplicatedKeySpansExceptLockTable returns all key spans that are fully Raft // replicated for the given Range, except for the lock table spans. These are // returned in the following sorted order: diff --git a/pkg/kv/kvserver/rditer/replica_data_iter_test.go b/pkg/kv/kvserver/rditer/replica_data_iter_test.go index 8c46e58d132d..33fc2bd6fea0 100644 --- a/pkg/kv/kvserver/rditer/replica_data_iter_test.go +++ b/pkg/kv/kvserver/rditer/replica_data_iter_test.go @@ -24,12 +24,14 @@ 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/enginepb" "github.com/cockroachdb/cockroach/pkg/testutils" "github.com/cockroachdb/cockroach/pkg/testutils/datapathutils" "github.com/cockroachdb/cockroach/pkg/testutils/echotest" "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/protoutil" "github.com/cockroachdb/cockroach/pkg/util/randutil" "github.com/cockroachdb/cockroach/pkg/util/uuid" "github.com/olekukonko/tablewriter" @@ -118,6 +120,13 @@ func createRangeData( TxnUUID: testTxnID, }, } + lockMeta := enginepb.MVCCMetadata{ + Txn: &enginepb.TxnMeta{ID: testTxnID}, + KeyBytes: storage.MVCCVersionTimestampSize, + ValBytes: int64(len(value.RawBytes)), + } + lockVal, err := protoutil.Marshal(&lockMeta) + require.NoError(t, err) rs = append(rs, storage.MVCCRangeKey{ // emitted last because we emit all point keys before range keys StartKey: desc.StartKey.AsRawKey().Clone(), @@ -134,7 +143,7 @@ func createRangeData( } for _, l := range locks { sl, _ := l.ToEngineKey(nil) - require.NoError(t, eng.PutEngineKey(sl, []byte("fake lock"))) + require.NoError(t, eng.PutEngineKey(sl, lockVal)) } return ps, rs diff --git a/pkg/kv/kvserver/rditer/stats.go b/pkg/kv/kvserver/rditer/stats.go index 95e5722da54b..83bb3cf4d705 100644 --- a/pkg/kv/kvserver/rditer/stats.go +++ b/pkg/kv/kvserver/rditer/stats.go @@ -27,6 +27,24 @@ func ComputeStatsForRange( ctx, d, reader, nowNanos, storage.ComputeStatsVisitors{}) } +// ComputeStatsForRangeUserOnly is like ComputeStatsForRange but iterates only +// over spans corresponding to user keys. +func ComputeStatsForRangeUserOnly( + ctx context.Context, d *roachpb.RangeDescriptor, reader storage.Reader, nowNanos int64, +) (enginepb.MVCCStats, error) { + return ComputeStatsForRangeWithVisitorsUserOnly( + ctx, d, reader, nowNanos, storage.ComputeStatsVisitors{}) +} + +// ComputeStatsForRangeExcludingUser is like ComputeStatsForRange but iterates +// only over spans *not* corresponding to user keys. +func ComputeStatsForRangeExcludingUser( + ctx context.Context, d *roachpb.RangeDescriptor, reader storage.Reader, nowNanos int64, +) (enginepb.MVCCStats, error) { + return ComputeStatsForRangeWithVisitorsExcludingUser( + ctx, d, reader, nowNanos, storage.ComputeStatsVisitors{}) +} + // ComputeStatsForRangeWithVisitors is like ComputeStatsForRange but also // calls the given callbacks for every key. func ComputeStatsForRangeWithVisitors( @@ -35,9 +53,45 @@ func ComputeStatsForRangeWithVisitors( reader storage.Reader, nowNanos int64, visitors storage.ComputeStatsVisitors, +) (enginepb.MVCCStats, error) { + return computeStatsForSpansWithVisitors(ctx, MakeReplicatedKeySpans(d), reader, nowNanos, visitors) +} + +// ComputeStatsForRangeWithVisitorsUserOnly is like +// ComputeStatsForRangeWithVisitors but computes stats only for spans +// corresponding to user keys. +func ComputeStatsForRangeWithVisitorsUserOnly( + ctx context.Context, + d *roachpb.RangeDescriptor, + reader storage.Reader, + nowNanos int64, + visitors storage.ComputeStatsVisitors, +) (enginepb.MVCCStats, error) { + return computeStatsForSpansWithVisitors(ctx, MakeReplicatedKeySpansUserOnly(d), reader, nowNanos, visitors) +} + +// ComputeStatsForRangeWithVisitorsExcludingUser is like +// ComputeStatsForRangeWithVisitors but computes stats only for spans +// corresponding to non-user keys. +func ComputeStatsForRangeWithVisitorsExcludingUser( + ctx context.Context, + d *roachpb.RangeDescriptor, + reader storage.Reader, + nowNanos int64, + visitors storage.ComputeStatsVisitors, +) (enginepb.MVCCStats, error) { + return computeStatsForSpansWithVisitors(ctx, MakeReplicatedKeySpansExcludingUser(d), reader, nowNanos, visitors) +} + +func computeStatsForSpansWithVisitors( + ctx context.Context, + spans []roachpb.Span, + reader storage.Reader, + nowNanos int64, + visitors storage.ComputeStatsVisitors, ) (enginepb.MVCCStats, error) { var ms enginepb.MVCCStats - for _, keySpan := range MakeReplicatedKeySpans(d) { + for _, keySpan := range spans { msDelta, err := storage.ComputeStatsWithVisitors( ctx, reader, keySpan.Key, keySpan.EndKey, nowNanos, visitors) if err != nil { diff --git a/pkg/kv/kvserver/rditer/stats_test.go b/pkg/kv/kvserver/rditer/stats_test.go new file mode 100644 index 000000000000..2213501316ee --- /dev/null +++ b/pkg/kv/kvserver/rditer/stats_test.go @@ -0,0 +1,54 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package rditer + +import ( + "context" + "testing" + "time" + + "github.com/cockroachdb/cockroach/pkg/roachpb" + "github.com/cockroachdb/cockroach/pkg/storage" + "github.com/cockroachdb/cockroach/pkg/util/leaktest" + "github.com/stretchr/testify/require" +) + +// TestComputeStats verifies that the sum of the stats returned by +// ComputeStatsForRangeUserOnly and ComputeStatsForRangeExcludingUser equals the +// total range stats returned by ComputeStatsForRange. +func TestComputeStats(t *testing.T) { + defer leaktest.AfterTest(t)() + + storage.DisableMetamorphicSimpleValueEncoding(t) + + eng := storage.NewDefaultInMemForTesting() + defer eng.Close() + + desc := roachpb.RangeDescriptor{ + RangeID: 1, + StartKey: roachpb.RKey("a"), + EndKey: roachpb.RKey("b"), + } + + createRangeData(t, eng, desc) + nowNanos := time.Now().UnixNano() + + userOnly, err := ComputeStatsForRangeUserOnly(context.Background(), &desc, eng, nowNanos) + require.NoError(t, err) + nonUserOnly, err := ComputeStatsForRangeExcludingUser(context.Background(), &desc, eng, nowNanos) + require.NoError(t, err) + all, err := ComputeStatsForRange(context.Background(), &desc, eng, nowNanos) + require.NoError(t, err) + + userPlusNonUser := userOnly + userPlusNonUser.Add(nonUserOnly) + require.Equal(t, all, userPlusNonUser) +}