Skip to content

Commit 98809dc

Browse files
wenyihu6tbg
wenyihu6
authored and
tbg
committed
asim: move NewStateWith* helpers into separate file
Previously, NewStateWith* helpers were with simulator logic, even though they were only used for unit tests (not datadriven tests). This commit moves them into a separate file to make it clearer that they are purely testing utilities and to clear confusion when changing code. Epic: none Release note: none
1 parent 7fa48ae commit 98809dc

File tree

3 files changed

+116
-104
lines changed

3 files changed

+116
-104
lines changed

‎pkg/kv/kvserver/asim/state/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ go_library(
1010
"liveness.go",
1111
"load.go",
1212
"new_state.go",
13+
"new_state_test_helper.go",
1314
"parser_replica_placement.go",
1415
"split_decider.go",
1516
"state.go",

‎pkg/kv/kvserver/asim/state/new_state.go‎

Lines changed: 0 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import (
1010
"math/rand"
1111
"sort"
1212

13-
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
1413
"github.com/cockroachdb/cockroach/pkg/roachpb"
1514
)
1615

@@ -372,106 +371,3 @@ func RangesInfoRandDistribution(
372371
storeList, distribution, distribution, ranges, spanConfig,
373372
int64(MinKey), int64(keyspace), rangeSize)
374373
}
375-
376-
// NewStateWithDistribution returns a State where the stores given are
377-
// initialized with the specified % of the replicas. This is done on a best
378-
// effort basis, given the replication factor. It may be impossible to satisfy
379-
// some distributions, for example: percentOfReplicas {1: 0.40, 2: 0.20, 3:
380-
// 0.20}, replicationFactor 3, would be impossible to satisfy as the only
381-
// distribution possible is {1: 0.33, 2: 0.33, 3: 0.33} given a replication
382-
// factor of 3. A best effort distribution is applied in these cases.
383-
func NewStateWithDistribution(
384-
percentOfReplicas []float64,
385-
ranges, replicationFactor, keyspace int,
386-
settings *config.SimulationSettings,
387-
) State {
388-
numNodes := len(percentOfReplicas)
389-
// Currently multi-store is not tested for correctness. Default to a single
390-
// store per node.
391-
clusterInfo := ClusterInfoWithStoreCount(numNodes, 1 /* storesPerNode */)
392-
s := LoadClusterInfo(clusterInfo, settings)
393-
394-
stores := make([]StoreID, numNodes)
395-
for i, store := range s.Stores() {
396-
stores[i] = store.StoreID()
397-
}
398-
spanConfig := defaultSpanConfig
399-
spanConfig.NumReplicas = int32(replicationFactor)
400-
spanConfig.NumVoters = int32(replicationFactor)
401-
402-
rangesInfo := RangesInfoWithDistribution(
403-
stores,
404-
percentOfReplicas,
405-
percentOfReplicas,
406-
ranges,
407-
spanConfig,
408-
int64(MinKey),
409-
int64(keyspace),
410-
0, /* rangeSize */
411-
)
412-
LoadRangeInfo(s, rangesInfo...)
413-
return s
414-
}
415-
416-
// NewStateWithReplCounts returns a new test state where each store is
417-
// initialized the given number of replicas. The required number of ranges is
418-
// inferred from the replication factor and the replica count.
419-
func NewStateWithReplCounts(
420-
replCounts map[StoreID]int, replicationFactor, keyspace int, settings *config.SimulationSettings,
421-
) State {
422-
clusterInfo := ClusterInfoWithStoreCount(len(replCounts), 1 /* storesPerNode */)
423-
rangesInfo := RangesInfoWithReplicaCounts(replCounts, keyspace, replicationFactor, 0 /* rangeSize */)
424-
return LoadConfig(clusterInfo, rangesInfo, settings)
425-
}
426-
427-
// NewStateEvenDistribution returns a new State where the replica count per
428-
// store is equal.
429-
func NewStateEvenDistribution(
430-
stores, ranges, replicationFactor, keyspace int, settings *config.SimulationSettings,
431-
) State {
432-
clusterInfo := ClusterInfoWithStoreCount(stores, 1 /* storesPerNode*/)
433-
rangesInfo := RangesInfoEvenDistribution(stores, ranges, int64(MinKey), int64(keyspace), replicationFactor, 0 /* rangeSize */)
434-
return LoadConfig(clusterInfo, rangesInfo, settings)
435-
}
436-
437-
// NewStateSkewedDistribution returns a new State where the replica count per
438-
// store is skewed.
439-
func NewStateSkewedDistribution(
440-
stores, ranges, replicationFactor, keyspace int, settings *config.SimulationSettings,
441-
) State {
442-
clusterInfo := ClusterInfoWithStoreCount(stores, 1 /* storesPerNode */)
443-
rangesInfo := RangesInfoSkewedDistribution(stores, ranges, int64(MinKey), int64(keyspace), replicationFactor, 0 /* rangeSize */)
444-
return LoadConfig(clusterInfo, rangesInfo, settings)
445-
}
446-
447-
// NewStateRandDistribution returns a new State where the replica count per
448-
// store is randomized.
449-
func NewStateRandDistribution(
450-
seed int64,
451-
stores int,
452-
ranges int,
453-
keyspace int,
454-
replicationFactor int,
455-
settings *config.SimulationSettings,
456-
) State {
457-
randSource := rand.New(rand.NewSource(seed))
458-
clusterInfo := ClusterInfoWithStoreCount(stores, 1 /* storesPerNode */)
459-
rangesInfo := RangesInfoRandDistribution(randSource, stores, ranges, keyspace, replicationFactor, 0 /* rangeSize */)
460-
return LoadConfig(clusterInfo, rangesInfo, settings)
461-
}
462-
463-
// NewStateWeightedRandDistribution returns a new State where the replica count
464-
// per store is weighted randomized based on weightedStores.
465-
func NewStateWeightedRandDistribution(
466-
seed int64,
467-
weightedStores []float64,
468-
ranges int,
469-
keyspace int,
470-
replicationFactor int,
471-
settings *config.SimulationSettings,
472-
) State {
473-
randSource := rand.New(rand.NewSource(seed))
474-
clusterInfo := ClusterInfoWithStoreCount(len(weightedStores), 1 /* storesPerNode */)
475-
rangesInfo := RangesInfoWeightedRandDistribution(randSource, weightedStores, ranges, keyspace, replicationFactor, 0 /* rangeSize */)
476-
return LoadConfig(clusterInfo, rangesInfo, settings)
477-
}

‎pkg/kv/kvserver/asim/state/new_state_test_helper.go‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
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 state
7+
8+
import (
9+
"math/rand"
10+
11+
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
12+
)
13+
14+
// NewStateWithDistribution returns a State where the stores given are
15+
// initialized with the specified % of the replicas. This is done on a best
16+
// effort basis, given the replication factor. It may be impossible to satisfy
17+
// some distributions, for example: percentOfReplicas {1: 0.40, 2: 0.20, 3:
18+
// 0.20}, replicationFactor 3, would be impossible to satisfy as the only
19+
// distribution possible is {1: 0.33, 2: 0.33, 3: 0.33} given a replication
20+
// factor of 3. A best effort distribution is applied in these cases.
21+
func NewStateWithDistribution(
22+
percentOfReplicas []float64,
23+
ranges, replicationFactor, keyspace int,
24+
settings *config.SimulationSettings,
25+
) State {
26+
numNodes := len(percentOfReplicas)
27+
// Currently multi-store is not tested for correctness. Default to a single
28+
// store per node.
29+
clusterInfo := ClusterInfoWithStoreCount(numNodes, 1 /* storesPerNode */)
30+
s := LoadClusterInfo(clusterInfo, settings)
31+
32+
stores := make([]StoreID, numNodes)
33+
for i, store := range s.Stores() {
34+
stores[i] = store.StoreID()
35+
}
36+
spanConfig := defaultSpanConfig
37+
spanConfig.NumReplicas = int32(replicationFactor)
38+
spanConfig.NumVoters = int32(replicationFactor)
39+
40+
rangesInfo := RangesInfoWithDistribution(
41+
stores,
42+
percentOfReplicas,
43+
percentOfReplicas,
44+
ranges,
45+
spanConfig,
46+
int64(MinKey),
47+
int64(keyspace),
48+
0, /* rangeSize */
49+
)
50+
LoadRangeInfo(s, rangesInfo...)
51+
return s
52+
}
53+
54+
// NewStateWithReplCounts returns a new test state where each store is
55+
// initialized the given number of replicas. The required number of ranges is
56+
// inferred from the replication factor and the replica count.
57+
func NewStateWithReplCounts(
58+
replCounts map[StoreID]int, replicationFactor, keyspace int, settings *config.SimulationSettings,
59+
) State {
60+
clusterInfo := ClusterInfoWithStoreCount(len(replCounts), 1 /* storesPerNode */)
61+
rangesInfo := RangesInfoWithReplicaCounts(replCounts, keyspace, replicationFactor, 0 /* rangeSize */)
62+
return LoadConfig(clusterInfo, rangesInfo, settings)
63+
}
64+
65+
// NewStateEvenDistribution returns a new State where the replica count per
66+
// store is equal.
67+
func NewStateEvenDistribution(
68+
stores, ranges, replicationFactor, keyspace int, settings *config.SimulationSettings,
69+
) State {
70+
clusterInfo := ClusterInfoWithStoreCount(stores, 1 /* storesPerNode*/)
71+
rangesInfo := RangesInfoEvenDistribution(stores, ranges, int64(MinKey), int64(keyspace), replicationFactor, 0 /* rangeSize */)
72+
return LoadConfig(clusterInfo, rangesInfo, settings)
73+
}
74+
75+
// NewStateSkewedDistribution returns a new State where the replica count per
76+
// store is skewed.
77+
func NewStateSkewedDistribution(
78+
stores, ranges, replicationFactor, keyspace int, settings *config.SimulationSettings,
79+
) State {
80+
clusterInfo := ClusterInfoWithStoreCount(stores, 1 /* storesPerNode */)
81+
rangesInfo := RangesInfoSkewedDistribution(stores, ranges, int64(MinKey), int64(keyspace), replicationFactor, 0 /* rangeSize */)
82+
return LoadConfig(clusterInfo, rangesInfo, settings)
83+
}
84+
85+
// NewStateRandDistribution returns a new State where the replica count per
86+
// store is randomized.
87+
func NewStateRandDistribution(
88+
seed int64,
89+
stores int,
90+
ranges int,
91+
keyspace int,
92+
replicationFactor int,
93+
settings *config.SimulationSettings,
94+
) State {
95+
randSource := rand.New(rand.NewSource(seed))
96+
clusterInfo := ClusterInfoWithStoreCount(stores, 1 /* storesPerNode */)
97+
rangesInfo := RangesInfoRandDistribution(randSource, stores, ranges, keyspace, replicationFactor, 0 /* rangeSize */)
98+
return LoadConfig(clusterInfo, rangesInfo, settings)
99+
}
100+
101+
// NewStateWeightedRandDistribution returns a new State where the replica count
102+
// per store is weighted randomized based on weightedStores.
103+
func NewStateWeightedRandDistribution(
104+
seed int64,
105+
weightedStores []float64,
106+
ranges int,
107+
keyspace int,
108+
replicationFactor int,
109+
settings *config.SimulationSettings,
110+
) State {
111+
randSource := rand.New(rand.NewSource(seed))
112+
clusterInfo := ClusterInfoWithStoreCount(len(weightedStores), 1 /* storesPerNode */)
113+
rangesInfo := RangesInfoWeightedRandDistribution(randSource, weightedStores, ranges, keyspace, replicationFactor, 0 /* rangeSize */)
114+
return LoadConfig(clusterInfo, rangesInfo, settings)
115+
}

0 commit comments

Comments
 (0)