asim: add randomized range generation

pkg/kv/kvserver/asim/gen/generator.go

@@ -209,8 +209,16 @@ type PlacementType int
 const (
 	Uniform PlacementType = iota
 	Skewed
+	Random
+	WeightedRandom
 )
 
+func GetAvailablePlacementTypes() []PlacementType {
+	// WeightedRandom is enabled if and only if default setting contains a
+	// non-zero array for weighted rand.
+	return []PlacementType{Uniform, Skewed, Random}
+}
+
 // BaseRanges provides basic ranges functionality and are embedded in
 // other specialized range structs.
 type BaseRanges struct {
@@ -220,21 +228,27 @@ type BaseRanges struct {
 	Bytes             int64
 }
 
-// getRangesInfo generates RangesInfo, with its distribution defined by
+// GetRangesInfo generates RangesInfo, with its distribution defined by
 // PlacementType and other configurations determined by BaseRanges fields.
-func (b BaseRanges) getRangesInfo(pType PlacementType, numOfStores int) state.RangesInfo {
+func (b BaseRanges) GetRangesInfo(
+	pType PlacementType, numOfStores int, randSource *rand.Rand, weightedRandom []float64,
+) state.RangesInfo {
 	switch pType {
 	case Uniform:
 		return state.RangesInfoEvenDistribution(numOfStores, b.Ranges, b.KeySpace, b.ReplicationFactor, b.Bytes)
 	case Skewed:
 		return state.RangesInfoSkewedDistribution(numOfStores, b.Ranges, b.KeySpace, b.ReplicationFactor, b.Bytes)
+	case WeightedRandom:
+		return state.RangesInfoWeightedRandDistribution(randSource, weightedRandom, b.Ranges, b.KeySpace, b.ReplicationFactor, b.Bytes)
+	case Random:
+		return state.RangesInfoRandDistribution(randSource, numOfStores, b.Ranges, b.KeySpace, b.ReplicationFactor, b.Bytes)
 	default:
 		panic(fmt.Sprintf("unexpected range placement type %v", pType))
 	}
 }
 
 // LoadRangeInfo loads the given state with the specified rangesInfo.
-func (b BaseRanges) loadRangeInfo(s state.State, rangesInfo state.RangesInfo) {
+func (b BaseRanges) LoadRangeInfo(s state.State, rangesInfo state.RangesInfo) {
 	for _, rangeInfo := range rangesInfo {
 		rangeInfo.Size = b.Bytes
 	}
@@ -250,6 +264,9 @@ type BasicRanges struct {
 func NewBasicRanges(
 	ranges int, placementType PlacementType, keySpace int, replicationFactor int, bytes int64,
 ) BasicRanges {
+	if placementType == WeightedRandom || placementType == Random {
+		panic(fmt.Sprintf("basic ranges cannot use randomized type %v", placementType))
+	}
 	return BasicRanges{
 		BaseRanges: BaseRanges{
 			Ranges:            ranges,
@@ -266,8 +283,8 @@ func NewBasicRanges(
 func (br BasicRanges) Generate(
 	seed int64, settings *config.SimulationSettings, s state.State,
 ) state.State {
-	rangesInfo := br.getRangesInfo(br.PlacementType, len(s.Stores()))
-	br.loadRangeInfo(s, rangesInfo)
+	rangesInfo := br.GetRangesInfo(br.PlacementType, len(s.Stores()), nil, []float64{})
+	br.LoadRangeInfo(s, rangesInfo)
 	return s
 }
 

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

@@ -12,6 +12,7 @@ package state
 
 import (
 	"fmt"
+	"math/rand"
 	"sort"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
@@ -62,6 +63,53 @@ func exactDistribution(counts []int) []float64 {
 	return distribution
 }
 
+type Weighted struct {
+	cumulativeWeights []float64
+}
+
+func NewWeighted(weightedStores []float64) Weighted {
+	cumulativeWeights := make([]float64, len(weightedStores))
+	prefixSumWeight := float64(0)
+	for i, item := range weightedStores {
+		prefixSumWeight += item
+		cumulativeWeights[i] = prefixSumWeight
+	}
+	return Weighted{cumulativeWeights: cumulativeWeights}
+}
+
+func (w Weighted) Rand(randSource *rand.Rand) int {
+	r := randSource.Float64()
+	index := sort.Search(len(w.cumulativeWeights), func(i int) bool { return w.cumulativeWeights[i] >= r })
+	return index
+}
+
+func weightedRandDistribution(randSource *rand.Rand, weightedStores []float64) []float64 {
+	w := NewWeighted(weightedStores)
+	// Vote for 10 times and give a distribution.
+	numSamples := 10
+	votes := make([]int, len(weightedStores))
+	for i := 0; i < numSamples; i++ {
+		index := w.Rand(randSource)
+		votes[index] += 1
+	}
+	return exactDistribution(votes)
+}
+
+func randDistribution(randSource *rand.Rand, n int) []float64 {
+	total := float64(0)
+	distribution := make([]float64, n)
+	for i := 0; i < n; i++ {
+		num := float64(randSource.Intn(10))
+		distribution[i] = num
+		total += num
+	}
+
+	for i := 0; i < n; i++ {
+		distribution[i] = distribution[i] / total
+	}
+	return distribution
+}
+
 // RangesInfoWithDistribution returns a RangesInfo, where the stores given are
 // initialized with the specified % of the replicas. This is done on a best
 // effort basis, given the replication factor. It may be impossible to satisfy
@@ -247,6 +295,59 @@ func RangesInfoEvenDistribution(
 		int64(MinKey), int64(keyspace), rangeSize)
 }
 
+// Weighted distribution: vote for 10 times and see which bucket the number falls under
+func RangesInfoWeightedRandDistribution(
+	randSource *rand.Rand,
+	weightedStores []float64,
+	ranges int,
+	keyspace int,
+	replicationFactor int,
+	rangeSize int64,
+) RangesInfo {
+	if randSource == nil || len(weightedStores) == 0 {
+		panic("unexpected arguments for weighted random ranges info")
+	}
+	distribution := weightedRandDistribution(randSource, weightedStores)
+	storeList := makeStoreList(len(weightedStores))
+	spanConfig := defaultSpanConfig
+	spanConfig.NumReplicas = int32(replicationFactor)
+	spanConfig.NumVoters = int32(replicationFactor)
+	return RangesInfoWithDistribution(
+		storeList,
+		distribution,
+		distribution,
+		ranges,
+		spanConfig,
+		int64(MinKey),
+		int64(keyspace),
+		rangeSize, /* rangeSize */
+	)
+}
+
+func RangesInfoRandDistribution(
+	randSource *rand.Rand,
+	stores int,
+	ranges int,
+	keyspace int,
+	replicationFactor int,
+	rangeSize int64,
+) RangesInfo {
+	if randSource == nil {
+		// BETTER NAME HERE
+		panic("unexpected arguments for weighted random ranges info")
+	}
+	distribution := randDistribution(randSource, stores)
+	storeList := makeStoreList(stores)
+
+	spanConfig := defaultSpanConfig
+	spanConfig.NumReplicas = int32(replicationFactor)
+	spanConfig.NumVoters = int32(replicationFactor)
+
+	return RangesInfoWithDistribution(
+		storeList, distribution, distribution, ranges, spanConfig,
+		int64(MinKey), int64(keyspace), rangeSize)
+}
+
 // NewStateWithDistribution returns a State where the stores given are
 // initialized with the specified % of the replicas. This is done on a best
 // effort basis, given the replication factor. It may be impossible to satisfy

pkg/kv/kvserver/asim/tests/BUILD.bazel

@@ -33,9 +33,9 @@ go_library(
     name = "tests",
     srcs = [
         "assert.go",
-        "default_settings.go",
         "rand_framework.go",
         "rand_gen.go",
+        "settings.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/tests",
     visibility = ["//visibility:public"],
@@ -46,6 +46,7 @@ go_library(
         "//pkg/kv/kvserver/asim/gen",
         "//pkg/kv/kvserver/asim/metrics",
         "//pkg/kv/kvserver/asim/state",
+        "//pkg/kv/kvserver/asim/workload",
         "//pkg/roachpb",
         "//pkg/spanconfig/spanconfigtestutils",
         "//pkg/util/log",

pkg/kv/kvserver/asim/tests/default_settings.go

@@ -100,3 +100,24 @@ func defaultPlotSettings() plotSettings {
 		width:  defaultWidth,
 	}
 }
+
+type rangeGenSettings struct {
+	rangeKeyGenType generatorType
+	keySpaceGenType generatorType
+	weightedRand    []float64
+}
+
+const (
+	defaultRangeKeyGenType = uniformGenerator
+	defaultKeySpaceGenType = uniformGenerator
+)
+
+var defaultWeightedRand []float64
+
+func defaultRangeGenSettings() rangeGenSettings {
+	return rangeGenSettings{
+		rangeKeyGenType: defaultRangeKeyGenType,
+		keySpaceGenType: defaultKeySpaceGenType,
+		weightedRand:    defaultWeightedRand,
+	}
+}

pkg/kv/kvserver/asim/tests/rand_framework.go

@@ -13,6 +13,7 @@ package tests
 import (
 	"context"
 	"fmt"
+	"math"
 	"math/rand"
 	"strings"
 	"testing"
@@ -21,6 +22,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/gen"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/metrics"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/workload"
 	"github.com/guptarohit/asciigraph"
 )
 
@@ -30,27 +32,40 @@ type testSettings struct {
 	verbose       bool
 	randSource    *rand.Rand
 	randOptions   map[string]bool
+	rangeGen      rangeGenSettings
 }
 
 func newTestSettings(
-	numIterations int, duration time.Duration, verbose bool, seed int64, randOptions map[string]bool,
+	numIterations int,
+	duration time.Duration,
+	verbose bool,
+	seed int64,
+	randOptions map[string]bool,
+	rangeGenSettings rangeGenSettings,
 ) testSettings {
 	return testSettings{
 		numIterations: numIterations,
 		duration:      duration,
 		verbose:       verbose,
 		randSource:    rand.New(rand.NewSource(seed)),
 		randOptions:   randOptions,
+		rangeGen:      rangeGenSettings,
 	}
 }
 
 type randTestingFramework struct {
-	s testSettings
+	s                 testSettings
+	rangeGenerator    generator
+	keySpaceGenerator generator
 }
 
 func newRandTestingFramework(settings testSettings) randTestingFramework {
+	rangeGenerator := newRandomizedGenerator(settings.randSource, defaultMinRange, defaultMaxRange, settings.rangeGen.rangeKeyGenType)
+	keySpaceGenerator := newRandomizedGenerator(settings.randSource, defaultMinKeySpace, defaultMaxKeySpace, settings.rangeGen.keySpaceGenType)
 	return randTestingFramework{
-		s: settings,
+		s:                 settings,
+		rangeGenerator:    rangeGenerator,
+		keySpaceGenerator: keySpaceGenerator,
 	}
 }
 
@@ -65,7 +80,7 @@ func (f randTestingFramework) getRanges() gen.RangeGen {
 	if !f.s.randOptions["ranges"] {
 		return defaultBasicRangesGen()
 	}
-	return gen.BasicRanges{}
+	return f.randomBasicRangesGen()
 }
 
 func (f randTestingFramework) getLoad() gen.LoadGen {
@@ -183,3 +198,74 @@ func checkAssertions(
 	}
 	return false, ""
 }
+
+type generator interface {
+	Num() int64
+}
+
+type generatorType int
+
+const (
+	uniformGenerator generatorType = iota
+	zipfGenerator
+)
+
+func newGenerator(randSource *rand.Rand, iMin int64, iMax int64, gType generatorType) generator {
+	switch gType {
+	case uniformGenerator:
+		return workload.NewUniformKeyGen(iMin, iMax, randSource)
+	case zipfGenerator:
+		return workload.NewZipfianKeyGen(iMin, iMax, 1.1, 1, randSource)
+	default:
+		panic(fmt.Sprintf("unexpected generator type %v", gType))
+	}
+}
+
+func newRandomizedGenerator(
+	randSource *rand.Rand, iMin int64, iMax int64, gType generatorType,
+) generator {
+	return newGenerator(randSource, iMin, iMax, gType)
+}
+
+const (
+	defaultMinRange    = 1
+	defaultMaxRange    = 1000
+	defaultMinKeySpace = 1000
+	defaultMaxKeySpace = 200000
+)
+
+func convertInt64ToInt(num int64) int {
+	// Should be impossible since we have set imax, imin to something smaller to imax32
+	if num < math.MinInt32 {
+		return math.MinInt32
+	}
+	if num > math.MaxUint32 {
+		return math.MaxUint32
+	}
+	return int(num)
+}
+
+func (f randTestingFramework) randomBasicRangesGen() gen.RangeGen {
+	options := gen.GetAvailablePlacementTypes()
+	randIndex := f.s.randSource.Intn(len(options))
+	chosenType := options[randIndex]
+	if len(f.s.rangeGen.weightedRand) == 0 {
+		return NewRandomizedBasicRanges(
+			f.s.randSource,
+			convertInt64ToInt(f.rangeGenerator.Num()),
+			convertInt64ToInt(f.keySpaceGenerator.Num()),
+			chosenType,
+			defaultReplicationFactor,
+			defaultBytes,
+		)
+	} else {
+		return NewWeightedRandomizedBasicRanges(
+			f.s.randSource,
+			f.s.rangeGen.weightedRand,
+			convertInt64ToInt(f.rangeGenerator.Num()),
+			convertInt64ToInt(f.keySpaceGenerator.Num()),
+			defaultReplicationFactor,
+			defaultBytes,
+		)
+	}
+}