Skip to content

Commit 4116e19

Browse files
pav-kvpav-kv
committed
kvstorage: mv TestIterateIDPrefixKeys to kvstorage
The function that this test exercised lives in the kvstorage package. Epic: none Release note: none
1 parent cd17f4e commit 4116e19

File tree

3 files changed

+154
-128
lines changed

3 files changed

+154
-128
lines changed

‎pkg/kv/kvserver/kvstorage/BUILD.bazel‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ go_test(
4242
"cluster_version_test.go",
4343
"datadriven_test.go",
4444
"destroy_test.go",
45+
"init_test.go",
4546
"initial_test.go",
4647
"stateloader_test.go",
4748
],
@@ -66,11 +67,13 @@ go_test(
6667
"//pkg/util/hlc",
6768
"//pkg/util/leaktest",
6869
"//pkg/util/log",
70+
"//pkg/util/randutil",
6971
"//pkg/util/stop",
7072
"//pkg/util/tracing",
7173
"//pkg/util/uuid",
7274
"@com_github_cockroachdb_datadriven//:datadriven",
7375
"@com_github_cockroachdb_redact//:redact",
76+
"@com_github_kr_pretty//:pretty",
7477
"@com_github_stretchr_testify//require",
7578
],
7679
)

‎pkg/kv/kvserver/kvstorage/init_test.go‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
// Copyright 2025 The Cockroach Authors.
2+
//
3+
// Use of this software is governed by the CockroachDB Software License
4+
// included in the /LICENSE file.
5+
6+
package kvstorage
7+
8+
import (
9+
"context"
10+
"fmt"
11+
"math/rand"
12+
"sort"
13+
"testing"
14+
15+
"github.com/cockroachdb/cockroach/pkg/keys"
16+
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
17+
"github.com/cockroachdb/cockroach/pkg/roachpb"
18+
"github.com/cockroachdb/cockroach/pkg/storage"
19+
"github.com/cockroachdb/cockroach/pkg/util/hlc"
20+
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
21+
"github.com/cockroachdb/cockroach/pkg/util/log"
22+
"github.com/cockroachdb/cockroach/pkg/util/randutil"
23+
"github.com/cockroachdb/cockroach/pkg/util/stop"
24+
"github.com/kr/pretty"
25+
"github.com/stretchr/testify/require"
26+
)
27+
28+
// TestIterateIDPrefixKeys lays down a number of tombstones (at keys.RangeTombstoneKey) interspersed
29+
// with other irrelevant keys (both chosen randomly). It then verifies that IterateIDPrefixKeys
30+
// correctly returns only the relevant keys and values.
31+
func TestIterateIDPrefixKeys(t *testing.T) {
32+
defer leaktest.AfterTest(t)()
33+
defer log.Scope(t).Close(t)
34+
35+
ctx := context.Background()
36+
stopper := stop.NewStopper()
37+
defer stopper.Stop(ctx)
38+
39+
eng := storage.NewDefaultInMemForTesting()
40+
stopper.AddCloser(eng)
41+
42+
seed := randutil.NewPseudoSeed()
43+
// const seed = -1666367124291055473
44+
t.Logf("seed is %d", seed)
45+
rng := rand.New(rand.NewSource(seed))
46+
47+
ops := []func(rangeID roachpb.RangeID) roachpb.Key{
48+
keys.RaftHardStateKey, // unreplicated; sorts after tombstone
49+
// Replicated key-anchored local key (i.e. not one we should care about).
50+
// Will be written at zero timestamp, but that's ok.
51+
func(rangeID roachpb.RangeID) roachpb.Key {
52+
return keys.RangeDescriptorKey([]byte(fmt.Sprintf("fakerange%d", rangeID)))
53+
},
54+
func(rangeID roachpb.RangeID) roachpb.Key {
55+
return roachpb.Key(fmt.Sprintf("fakeuserkey%d", rangeID))
56+
},
57+
}
58+
59+
const rangeCount = 10
60+
rangeIDFn := func() roachpb.RangeID {
61+
return 1 + roachpb.RangeID(rng.Intn(10*rangeCount)) // spread rangeIDs out
62+
}
63+
64+
// Write a number of keys that should be irrelevant to the iteration in this test.
65+
for i := 0; i < rangeCount; i++ {
66+
rangeID := rangeIDFn()
67+
68+
// Grab between one and all ops, randomly.
69+
for _, opIdx := range rng.Perm(len(ops))[:rng.Intn(1+len(ops))] {
70+
key := ops[opIdx](rangeID)
71+
t.Logf("writing op=%d rangeID=%d", opIdx, rangeID)
72+
if _, err := storage.MVCCPut(
73+
ctx,
74+
eng,
75+
key,
76+
hlc.Timestamp{},
77+
roachpb.MakeValueFromString("fake value for "+key.String()),
78+
storage.MVCCWriteOptions{},
79+
); err != nil {
80+
t.Fatal(err)
81+
}
82+
}
83+
}
84+
85+
type seenT struct {
86+
rangeID roachpb.RangeID
87+
tombstone kvserverpb.RangeTombstone
88+
}
89+
90+
// Next, write the keys we're planning to see again.
91+
var wanted []seenT
92+
{
93+
used := make(map[roachpb.RangeID]struct{})
94+
for {
95+
rangeID := rangeIDFn()
96+
if _, ok := used[rangeID]; ok {
97+
// We already wrote this key, so roll the dice again.
98+
continue
99+
}
100+
101+
tombstone := kvserverpb.RangeTombstone{
102+
NextReplicaID: roachpb.ReplicaID(rng.Int31n(100)),
103+
}
104+
105+
used[rangeID] = struct{}{}
106+
wanted = append(wanted, seenT{rangeID: rangeID, tombstone: tombstone})
107+
108+
t.Logf("writing tombstone at rangeID=%d", rangeID)
109+
require.NoError(t, MakeStateLoader(rangeID).SetRangeTombstone(ctx, eng, tombstone))
110+
111+
if len(wanted) >= rangeCount {
112+
break
113+
}
114+
}
115+
}
116+
117+
sort.Slice(wanted, func(i, j int) bool {
118+
return wanted[i].rangeID < wanted[j].rangeID
119+
})
120+
121+
var seen []seenT
122+
var tombstone kvserverpb.RangeTombstone
123+
124+
handleTombstone := func(rangeID roachpb.RangeID) error {
125+
seen = append(seen, seenT{rangeID: rangeID, tombstone: tombstone})
126+
return nil
127+
}
128+
129+
if err := IterateIDPrefixKeys(ctx, eng, keys.RangeTombstoneKey, &tombstone, handleTombstone); err != nil {
130+
t.Fatal(err)
131+
}
132+
placeholder := seenT{
133+
rangeID: roachpb.RangeID(9999),
134+
}
135+
136+
if len(wanted) != len(seen) {
137+
t.Errorf("wanted %d results, got %d", len(wanted), len(seen))
138+
}
139+
140+
for len(wanted) < len(seen) {
141+
wanted = append(wanted, placeholder)
142+
}
143+
for len(seen) < len(wanted) {
144+
seen = append(seen, placeholder)
145+
}
146+
147+
if diff := pretty.Diff(wanted, seen); len(diff) > 0 {
148+
pretty.Ldiff(t, wanted, seen)
149+
t.Fatal("diff(wanted, seen) is nonempty")
150+
}
151+
}

‎pkg/kv/kvserver/store_test.go‎

Lines changed: 0 additions & 128 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import (
1010
"context"
1111
"fmt"
1212
"math"
13-
"math/rand"
1413
"reflect"
1514
"sort"
1615
"strings"
@@ -62,15 +61,13 @@ import (
6261
"github.com/cockroachdb/cockroach/pkg/util/log"
6362
"github.com/cockroachdb/cockroach/pkg/util/metric"
6463
"github.com/cockroachdb/cockroach/pkg/util/protoutil"
65-
"github.com/cockroachdb/cockroach/pkg/util/randutil"
6664
"github.com/cockroachdb/cockroach/pkg/util/retry"
6765
"github.com/cockroachdb/cockroach/pkg/util/stop"
6866
"github.com/cockroachdb/cockroach/pkg/util/timeutil"
6967
"github.com/cockroachdb/cockroach/pkg/util/tracing"
7068
"github.com/cockroachdb/cockroach/pkg/util/uuid"
7169
"github.com/cockroachdb/errors"
7270
"github.com/cockroachdb/redact"
73-
"github.com/kr/pretty"
7471
"github.com/stretchr/testify/assert"
7572
"github.com/stretchr/testify/require"
7673
"golang.org/x/exp/slices"
@@ -343,131 +340,6 @@ func createTestStoreWithConfig(
343340
return store
344341
}
345342

346-
// TestIterateIDPrefixKeys lays down a number of tombstones (at keys.RangeTombstoneKey) interspersed
347-
// with other irrelevant keys (both chosen randomly). It then verifies that IterateIDPrefixKeys
348-
// correctly returns only the relevant keys and values.
349-
func TestIterateIDPrefixKeys(t *testing.T) {
350-
defer leaktest.AfterTest(t)()
351-
defer log.Scope(t).Close(t)
352-
353-
ctx := context.Background()
354-
stopper := stop.NewStopper()
355-
defer stopper.Stop(ctx)
356-
357-
eng := storage.NewDefaultInMemForTesting()
358-
stopper.AddCloser(eng)
359-
360-
seed := randutil.NewPseudoSeed()
361-
// const seed = -1666367124291055473
362-
t.Logf("seed is %d", seed)
363-
rng := rand.New(rand.NewSource(seed))
364-
365-
ops := []func(rangeID roachpb.RangeID) roachpb.Key{
366-
keys.RaftHardStateKey, // unreplicated; sorts after tombstone
367-
// Replicated key-anchored local key (i.e. not one we should care about).
368-
// Will be written at zero timestamp, but that's ok.
369-
func(rangeID roachpb.RangeID) roachpb.Key {
370-
return keys.RangeDescriptorKey([]byte(fmt.Sprintf("fakerange%d", rangeID)))
371-
},
372-
func(rangeID roachpb.RangeID) roachpb.Key {
373-
return roachpb.Key(fmt.Sprintf("fakeuserkey%d", rangeID))
374-
},
375-
}
376-
377-
const rangeCount = 10
378-
rangeIDFn := func() roachpb.RangeID {
379-
return 1 + roachpb.RangeID(rng.Intn(10*rangeCount)) // spread rangeIDs out
380-
}
381-
382-
// Write a number of keys that should be irrelevant to the iteration in this test.
383-
for i := 0; i < rangeCount; i++ {
384-
rangeID := rangeIDFn()
385-
386-
// Grab between one and all ops, randomly.
387-
for _, opIdx := range rng.Perm(len(ops))[:rng.Intn(1+len(ops))] {
388-
key := ops[opIdx](rangeID)
389-
t.Logf("writing op=%d rangeID=%d", opIdx, rangeID)
390-
if _, err := storage.MVCCPut(
391-
ctx,
392-
eng,
393-
key,
394-
hlc.Timestamp{},
395-
roachpb.MakeValueFromString("fake value for "+key.String()),
396-
storage.MVCCWriteOptions{},
397-
); err != nil {
398-
t.Fatal(err)
399-
}
400-
}
401-
}
402-
403-
type seenT struct {
404-
rangeID roachpb.RangeID
405-
tombstone kvserverpb.RangeTombstone
406-
}
407-
408-
// Next, write the keys we're planning to see again.
409-
var wanted []seenT
410-
{
411-
used := make(map[roachpb.RangeID]struct{})
412-
for {
413-
rangeID := rangeIDFn()
414-
if _, ok := used[rangeID]; ok {
415-
// We already wrote this key, so roll the dice again.
416-
continue
417-
}
418-
419-
tombstone := kvserverpb.RangeTombstone{
420-
NextReplicaID: roachpb.ReplicaID(rng.Int31n(100)),
421-
}
422-
423-
used[rangeID] = struct{}{}
424-
wanted = append(wanted, seenT{rangeID: rangeID, tombstone: tombstone})
425-
426-
t.Logf("writing tombstone at rangeID=%d", rangeID)
427-
require.NoError(t, kvstorage.MakeStateLoader(rangeID).SetRangeTombstone(ctx, eng, tombstone))
428-
429-
if len(wanted) >= rangeCount {
430-
break
431-
}
432-
}
433-
}
434-
435-
sort.Slice(wanted, func(i, j int) bool {
436-
return wanted[i].rangeID < wanted[j].rangeID
437-
})
438-
439-
var seen []seenT
440-
var tombstone kvserverpb.RangeTombstone
441-
442-
handleTombstone := func(rangeID roachpb.RangeID) error {
443-
seen = append(seen, seenT{rangeID: rangeID, tombstone: tombstone})
444-
return nil
445-
}
446-
447-
if err := kvstorage.IterateIDPrefixKeys(ctx, eng, keys.RangeTombstoneKey, &tombstone, handleTombstone); err != nil {
448-
t.Fatal(err)
449-
}
450-
placeholder := seenT{
451-
rangeID: roachpb.RangeID(9999),
452-
}
453-
454-
if len(wanted) != len(seen) {
455-
t.Errorf("wanted %d results, got %d", len(wanted), len(seen))
456-
}
457-
458-
for len(wanted) < len(seen) {
459-
wanted = append(wanted, placeholder)
460-
}
461-
for len(seen) < len(wanted) {
462-
seen = append(seen, placeholder)
463-
}
464-
465-
if diff := pretty.Diff(wanted, seen); len(diff) > 0 {
466-
pretty.Ldiff(t, wanted, seen)
467-
t.Fatal("diff(wanted, seen) is nonempty")
468-
}
469-
}
470-
471343
// TestStoreConfigSetDefaults checks that StoreConfig.SetDefaults() sets proper
472344
// defaults based on numStores.
473345
func TestStoreConfigSetDefaultsNumStores(t *testing.T) {

0 commit comments

Comments
 (0)