asim: extend datadriven test to support recovery

Previously, only rebalancing was supported in the data driven simulation
test. This commit extends the syntax to support recovery scenarios. As
part of the extension, the state generator is split into range
generation and cluster generation.

Examples are added for each command, along with common testing scenarios
such as decommissioning, IO overload, disk fullness and adding a node
(with a store).

The newly supported commands are listed below:

```
- "load_cluster": config=<name>
  Load a defined cluster configuration to be the generated cluster in the
  simulation. The available confiurations are: single_region: 15 nodes in
  region=US, 5 in each zone US_1/US_2/US_3. single_region_multi_store: 3
  nodes, 5 stores per node with the same zone/region configuration as
  above. multi_region: 36 nodes, 12 in each region and 4 in each zone,
  regions having 3 zones. complex: 28 nodes, 3 regions with a skewed
  number of nodes per region.

- "gen_ranges" [ranges=<int>] [placement_skew=<bool>] [repl_factor=<int>]
  [keyspace=<int>] [range_bytes=<int>]
  Initialize the range generator parameters. On the next call to eval, the
  range generator is called to assign an ranges and their replica
  placement. The default values are ranges=1 repl_factor=3
  placement_skew=false keyspace=10000.

- set_liveness node=<int> [delay=<duration>]
  status=(dead|decommissioning|draining|unavailable)
  Set the liveness status of the node with ID NodeID. This applies at the
  start of the simulation or with some delay after the simulation starts,
  if specified.

- add_node: [stores=<int>] [locality=<string>] [delay=<duration>]
  Add a node to the cluster after initial generation with some delay,
  locality and number of stores on the node. The default values are
  stores=0 locality=none delay=0.

- set_span_config [delay=<duration>]
  [startKey, endKey): <span_config> Provide a new line separated list
  of spans and span configurations e.g.
  [0,100): num_replicas=5 num_voters=3 constraints={'+region=US_East'}
  [100, 500): num_replicas=3
  ...
  This will update the span config for the span [0,100) to specify 3
  voting replicas and 2 non-voting replicas, with a constraint that all
  replicas are in the region US_East.

- assertion extended to support two new assertion types:
  For type=stat assertions, if the stat (e.g. stat=replicas) value of the
  last ticks (e.g. ticks=5) duration is not exactly equal to threshold,
  the assertion fails. This applies for a specified store which must be
  provided with store=storeID.

  For type=conformance assertions, you may assert on the number of
  replicas that you expect to be under-replicated(under),
  over-replicated(over), unavailable(unavailable) and violating
  constraints(violating) at the end of the evaluation.

- "topology" [sample=<int>]
  Print the cluster locality topology of the sample given (default=last).
  e.g. for the load_cluster config=single_region
  US
  ..US_1
  ....└── [1 2 3 4 5]
  ..US_2
  ....└── [6 7 8 9 10]
  ..US_3
  ....└── [11 12 13 14 15]
```

Release note: None

pkg/kv/kvserver/asim/asim.go

@@ -68,6 +68,7 @@ type Simulator struct {
 // TODO(kvoli): Add a range log like structure to the history.
 type History struct {
 	Recorded [][]metrics.StoreMetrics
+	S        state.State
 }
 
 // Listen implements the metrics.StoreMetricListener interface.
@@ -109,7 +110,7 @@ func NewSimulator(
 		shuffler:       state.NewShuffler(settings.Seed),
 		// TODO(kvoli): Keeping the state around is a bit hacky, find a better
 		// method of reporting the ranges.
-		history:  History{Recorded: [][]metrics.StoreMetrics{}},
+		history:  History{Recorded: [][]metrics.StoreMetrics{}, S: initialState},
 		events:   events,
 		settings: settings,
 	}
@@ -121,6 +122,7 @@ func NewSimulator(
 	s.state.RegisterConfigChangeListener(s)
 
 	m.Register(&s.history)
+	s.AddLogTag("asim", nil)
 	return s
 }
 

pkg/kv/kvserver/asim/asim_test.go

@@ -87,6 +87,6 @@ func TestAllocatorSimulatorDeterministic(t *testing.T) {
 			refRun = history
 			continue
 		}
-		require.Equal(t, refRun, history)
+		require.Equal(t, refRun.Recorded, history.Recorded)
 	}
 }

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

@@ -8,6 +8,7 @@ go_library(
     deps = [
         "//pkg/kv/kvserver/asim",
         "//pkg/kv/kvserver/asim/config",
+        "//pkg/kv/kvserver/asim/event",
         "//pkg/kv/kvserver/asim/metrics",
         "//pkg/kv/kvserver/asim/state",
         "//pkg/kv/kvserver/asim/workload",

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

@@ -12,10 +12,12 @@ package gen
 
 import (
 	"math/rand"
+	"sort"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/event"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/metrics"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/state"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/workload"
@@ -36,26 +38,53 @@ type LoadGen interface {
 	Generate(seed int64, settings *config.SimulationSettings) []workload.Generator
 }
 
-// StateGen provides a method to generate a state given a seed and simulation
-// settings.
-type StateGen interface {
-	// Generate returns a state that is parameterized randomly by the seed and
-	// simulation settings provided.
+// ClusterGen provides a method to generate the initial cluster state,  given a
+// seed and simulation settings. The initial cluster state includes: nodes
+// (including locality) and stores.
+type ClusterGen interface {
+	// Generate returns a new State that is parameterized randomly by the seed
+	// and simulation settings provided.
 	Generate(seed int64, settings *config.SimulationSettings) state.State
 }
 
+// RangeGen provides a method to generate the initial range splits, range
+// replica and lease placement within a cluster.
+type RangeGen interface {
+	// Generate returns an updated state, given the initial state, seed and
+	// simulation settings provided. In the updated state, ranges will have been
+	// created, replicas and leases assigned to stores in the cluster.
+	Generate(seed int64, settings *config.SimulationSettings, s state.State) state.State
+}
+
+// EventGen provides a  method to generate a list of events that will apply to
+// the simulated cluster. Currently, only delayed (fixed time) events are
+// supported.
+type EventGen interface {
+	// Generate returns a list of events, which should be exectued at the delay specified.
+	Generate(seed int64) event.DelayedEventList
+}
+
 // GenerateSimulation is a utility function that creates a new allocation
 // simulation using the provided state, workload, settings generators and seed.
 func GenerateSimulation(
-	duration time.Duration, stateGen StateGen, loadGen LoadGen, settingsGen SettingsGen, seed int64,
+	duration time.Duration,
+	clusterGen ClusterGen,
+	rangeGen RangeGen,
+	loadGen LoadGen,
+	settingsGen SettingsGen,
+	eventGen EventGen,
+	seed int64,
 ) *asim.Simulator {
 	settings := settingsGen.Generate(seed)
+	s := clusterGen.Generate(seed, &settings)
+	s = rangeGen.Generate(seed, &settings, s)
 	return asim.NewSimulator(
 		duration,
 		loadGen.Generate(seed, &settings),
-		stateGen.Generate(seed, &settings),
+		s,
 		&settings,
 		metrics.NewTracker(settings.MetricsInterval),
+		eventGen.Generate(seed)...,
 	)
 }
 
@@ -89,6 +118,10 @@ type BasicLoad struct {
 // SkewedAccess is true. The returned workload generators are seeded with the
 // provided seed.
 func (bl BasicLoad) Generate(seed int64, settings *config.SimulationSettings) []workload.Generator {
+	if bl.Rate == 0 {
+		return []workload.Generator{}
+	}
+
 	var keyGen workload.KeyGenerator
 	rand := rand.New(rand.NewSource(seed))
 	if bl.SkewedAccess {
@@ -110,25 +143,94 @@ func (bl BasicLoad) Generate(seed int64, settings *config.SimulationSettings) []
 	}
 }
 
-// BasicState implements the StateGen interface.
-type BasicState struct {
-	Stores            int
+// LoadedCluster implements the ClusterGen interface.
+type LoadedCluster struct {
+	Info state.ClusterInfo
+}
+
+// Generate returns a new simulator state, where the cluster is loaded based on
+// the cluster info the loaded cluster generator is created with. There is no
+// randomness in this cluster generation.
+func (lc LoadedCluster) Generate(seed int64, settings *config.SimulationSettings) state.State {
+	return state.LoadClusterInfo(lc.Info, settings)
+}
+
+// BasicCluster implements the ClusterGen interace.
+type BasicCluster struct {
+	Nodes         int
+	StoresPerNode int
+}
+
+// Generate returns a new simulator state, where the cluster is created with all
+// nodes having the same locality and with the specified number of stores/nodes
+// created. The cluster is created based on the stores and stores-per-node
+// values the basic cluster generator is created with.
+func (lc BasicCluster) Generate(seed int64, settings *config.SimulationSettings) state.State {
+	info := state.ClusterInfoWithStoreCount(lc.Nodes, lc.StoresPerNode)
+	return state.LoadClusterInfo(info, settings)
+}
+
+// LoadedRanges implements the RangeGen interface.
+type LoadedRanges struct {
+	Info state.RangesInfo
+}
+
+// Generate returns an updated simulator state, where the cluster is loaded
+// with the range info that the generator was created with. There is no
+// randomness in this cluster generation.
+func (lr LoadedRanges) Generate(
+	seed int64, settings *config.SimulationSettings, s state.State,
+) state.State {
+	state.LoadRangeInfo(s, lr.Info...)
+	return s
+}
+
+// PlacementType represents a type of placement distribution.
+type PlacementType int
+
+const (
+	Uniform PlacementType = iota
+	Skewed
+)
+
+// BasicRanges implements the RangeGen interface.
+type BasicRanges struct {
 	Ranges            int
-	SkewedPlacement   bool
+	PlacementType     PlacementType
 	KeySpace          int
 	ReplicationFactor int
+	Bytes             int64
 }
 
-// Generate returns a new state that is created with the number of stores,
-// ranges, keyspace and replication factor from the basic state fields. The
-// initial assignment of replicas and leases for ranges follows either a
-// uniform or powerlaw distribution depending on if SkewedPlacement is true.
-func (bs BasicState) Generate(seed int64, settings *config.SimulationSettings) state.State {
-	var s state.State
-	if bs.SkewedPlacement {
-		s = state.NewStateSkewedDistribution(bs.Stores, bs.Ranges, bs.ReplicationFactor, bs.KeySpace, settings)
-	} else {
-		s = state.NewStateEvenDistribution(bs.Stores, bs.Ranges, bs.ReplicationFactor, bs.KeySpace, settings)
+// Generate returns an updated simulator state, where the cluster is loaded
+// with ranges based on the parameters of basic ranges.
+func (br BasicRanges) Generate(
+	seed int64, settings *config.SimulationSettings, s state.State,
+) state.State {
+	stores := len(s.Stores())
+	var rangesInfo state.RangesInfo
+	switch br.PlacementType {
+	case Uniform:
+		rangesInfo = state.RangesInfoEvenDistribution(stores, br.Ranges, br.KeySpace, br.ReplicationFactor, br.Bytes)
+	case Skewed:
+		rangesInfo = state.RangesInfoSkewedDistribution(stores, br.Ranges, br.KeySpace, br.ReplicationFactor, br.Bytes)
 	}
+	for _, rangeInfo := range rangesInfo {
+		rangeInfo.Size = br.Bytes
+	}
+	state.LoadRangeInfo(s, rangesInfo...)
 	return s
 }
+
+// StaticEvents implements the EventGen interface.
+// TODO(kvoli): introduce conditional events.
+type StaticEvents struct {
+	DelayedEvents event.DelayedEventList
+}
+
+// Generate returns a list of events, exactly the same as the events
+// StaticEvents was created with.
+func (se StaticEvents) Generate(seed int64) event.DelayedEventList {
+	sort.Sort(se.DelayedEvents)
+	return se.DelayedEvents
+}

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

@@ -153,14 +153,14 @@ func RangesInfoWithDistribution(
 // their weights, a best effort apporach is taken so that the total number of
 // aggregate matches numNodes.
 func ClusterInfoWithDistribution(
-	numNodes int, storesPerNode int, regions []string, regionNodeWeights []float64,
+	nodeCount int, storesPerNode int, regions []string, regionNodeWeights []float64,
 ) ClusterInfo {
 	ret := ClusterInfo{}
 
 	ret.Regions = make([]Region, len(regions))
-	availableNodes := numNodes
+	availableNodes := nodeCount
 	for i, name := range regions {
-		allocatedNodes := int(float64(numNodes) * (regionNodeWeights[i]))
+		allocatedNodes := int(float64(nodeCount) * (regionNodeWeights[i]))
 		if allocatedNodes > availableNodes {
 			allocatedNodes = availableNodes
 		}
@@ -174,11 +174,11 @@ func ClusterInfoWithDistribution(
 	return ret
 }
 
-// ClusterInfoWithStores returns a new ClusterInfo with the specified number of
-// stores. There will be only one store per node and a single region and zone.
-func ClusterInfoWithStoreCount(stores int, storesPerNode int) ClusterInfo {
+// ClusterInfoWithStoreCount returns a new ClusterInfo with the specified number of
+// stores. There will be storesPerNode stores per node and a single region and zone.
+func ClusterInfoWithStoreCount(nodeCount int, storesPerNode int) ClusterInfo {
 	return ClusterInfoWithDistribution(
-		stores,
+		nodeCount,
 		storesPerNode,
 		[]string{"AU_EAST"}, /* regions */
 		[]float64{1},        /* regionNodeWeights */

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

@@ -7,11 +7,18 @@ go_test(
     data = glob(["testdata/**"]),
     embed = [":tests"],
     deps = [
+        "//pkg/kv/kvserver/allocator/allocatorimpl",
         "//pkg/kv/kvserver/asim",
         "//pkg/kv/kvserver/asim/config",
+        "//pkg/kv/kvserver/asim/event",
         "//pkg/kv/kvserver/asim/gen",
         "//pkg/kv/kvserver/asim/metrics",
+        "//pkg/kv/kvserver/asim/state",
+        "//pkg/kv/kvserver/liveness/livenesspb",
+        "//pkg/roachpb",
+        "//pkg/spanconfig/spanconfigtestutils",
         "//pkg/testutils/datapathutils",
+        "//pkg/util/log",
         "@com_github_cockroachdb_datadriven//:datadriven",
         "@com_github_guptarohit_asciigraph//:asciigraph",
         "@com_github_stretchr_testify//require",
@@ -26,6 +33,8 @@ go_library(
     deps = [
         "//pkg/kv/kvserver/asim",
         "//pkg/kv/kvserver/asim/metrics",
+        "//pkg/roachpb",
+        "//pkg/spanconfig/spanconfigtestutils",
         "//pkg/util/log",
         "@com_github_montanaflynn_stats//:stats",
     ],