kvnemesis: add node stopping and restarting

This commit adds two new kvnemesis operations: stop a running node, and
restart a stopped node. Similarly to previous operations that inject
faults, we need to be careful about either preserving quorum
(liveness mode) or allowing unavailability (safety mode). This commit
adds only a liveness mode variant of kvenemesis due to the current
behavior of `TestCluster` which could lead to hanging goroutines if
nodes are stopped sequentially and quorum is lost
(see `TestCluster.stopServers` for more details).

Part of: #64828

Release note: None

Author: miraradeva

pkg/kv/kvnemesis/BUILD.bazel

@@ -99,9 +99,11 @@ go_test(
         "//pkg/spanconfig",
         "//pkg/storage",
         "//pkg/storage/enginepb",
+        "//pkg/storage/fs",
         "//pkg/testutils",
         "//pkg/testutils/datapathutils",
         "//pkg/testutils/echotest",
+        "//pkg/testutils/listenerutil",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/skip",
         "//pkg/testutils/sqlutils",
@@ -112,6 +114,7 @@ go_test(
         "//pkg/util/leaktest",
         "//pkg/util/log",
         "//pkg/util/randutil",
+        "//pkg/util/stop",
         "//pkg/util/syncutil",
         "//pkg/util/uuid",
         "@com_github_cockroachdb_errors//:errors",

pkg/kv/kvnemesis/applier.go

@@ -28,20 +28,22 @@ import (
 
 // Applier executes Steps.
 type Applier struct {
-	env *Env
-	dbs []*kv.DB
-	mu  struct {
+	env   *Env
+	dbs   []*kv.DB
+	nodes *nodes
+	mu    struct {
 		dbIdx int
 		syncutil.Mutex
 		txns map[string]*kv.Txn
 	}
 }
 
 // MakeApplier constructs an Applier that executes against the given DBs.
-func MakeApplier(env *Env, dbs ...*kv.DB) *Applier {
+func MakeApplier(env *Env, n *nodes, dbs ...*kv.DB) *Applier {
 	a := &Applier{
-		env: env,
-		dbs: dbs,
+		env:   env,
+		dbs:   dbs,
+		nodes: n,
 	}
 	a.mu.txns = make(map[string]*kv.Txn)
 	return a
@@ -63,7 +65,7 @@ func (a *Applier) Apply(ctx context.Context, step *Step) (trace tracingpb.Record
 		}
 		trace = collectAndFinish()
 	}()
-	applyOp(recCtx, a.env, db, &step.Op)
+	a.applyOp(recCtx, db, &step.Op)
 	return collectAndFinish(), nil
 }
 
@@ -120,7 +122,7 @@ func exceptContextCanceled(err error) bool {
 		strings.Contains(err.Error(), "query execution canceled")
 }
 
-func applyOp(ctx context.Context, env *Env, db *kv.DB, op *Operation) {
+func (a *Applier) applyOp(ctx context.Context, db *kv.DB, op *Operation) {
 	switch o := op.GetValue().(type) {
 	case *GetOperation,
 		*PutOperation,
@@ -146,10 +148,10 @@ func applyOp(ctx context.Context, env *Env, db *kv.DB, op *Operation) {
 		err := db.AdminTransferLease(ctx, o.Key, o.Target)
 		o.Result = resultInit(ctx, err)
 	case *ChangeSettingOperation:
-		err := changeClusterSettingInEnv(ctx, env, o)
+		err := changeClusterSettingInEnv(ctx, a.env, o)
 		o.Result = resultInit(ctx, err)
 	case *ChangeZoneOperation:
-		err := updateZoneConfigInEnv(ctx, env, o.Type)
+		err := updateZoneConfigInEnv(ctx, a.env, o.Type)
 		o.Result = resultInit(ctx, err)
 	case *BarrierOperation:
 		var err error
@@ -162,10 +164,20 @@ func applyOp(ctx context.Context, env *Env, db *kv.DB, op *Operation) {
 	case *FlushLockTableOperation:
 		o.Result = resultInit(ctx, db.FlushLockTable(ctx, o.Key, o.EndKey))
 	case *AddNetworkPartitionOperation:
-		err := env.Partitioner.AddPartition(roachpb.NodeID(o.FromNode), roachpb.NodeID(o.ToNode))
+		err := a.env.Partitioner.AddPartition(roachpb.NodeID(o.FromNode), roachpb.NodeID(o.ToNode))
 		o.Result = resultInit(ctx, err)
 	case *RemoveNetworkPartitionOperation:
-		err := env.Partitioner.RemovePartition(roachpb.NodeID(o.FromNode), roachpb.NodeID(o.ToNode))
+		err := a.env.Partitioner.RemovePartition(roachpb.NodeID(o.FromNode), roachpb.NodeID(o.ToNode))
+		o.Result = resultInit(ctx, err)
+	case *StopNodeOperation:
+		serverID := int(o.NodeId) - 1
+		a.env.Restarter.StopServer(serverID)
+		a.nodes.setStopped(int(o.NodeId))
+		o.Result = resultInit(ctx, nil)
+	case *RestartNodeOperation:
+		serverID := int(o.NodeId) - 1
+		err := a.env.Restarter.RestartServer(serverID)
+		a.nodes.setRunning(int(o.NodeId))
 		o.Result = resultInit(ctx, err)
 	case *ClosureTxnOperation:
 		// Use a backoff loop to avoid thrashing on txn aborts. Don't wait between

pkg/kv/kvnemesis/applier_test.go

@@ -70,7 +70,7 @@ func TestApplier(t *testing.T) {
 		require.NoError(t, w.Finish())
 	}
 
-	a := MakeApplier(env, db)
+	a := MakeApplier(env, &nodes{}, db)
 
 	tests := []testCase{
 		{

pkg/kv/kvnemesis/env.go

@@ -27,13 +27,19 @@ type Logger interface {
 	WriteFile(basename string, contents string) string
 }
 
+type Restarter interface {
+	StopServer(idx int)
+	RestartServer(idx int) error
+}
+
 // Env manipulates the environment (cluster settings, zone configurations) that
 // the Applier operates in.
 type Env struct {
 	SQLDBs      []*gosql.DB
 	Tracker     *SeqTracker
 	L           Logger
 	Partitioner *rpc.Partitioner
+	Restarter   Restarter
 }
 
 func (e *Env) anyNode() *gosql.DB {

pkg/kv/kvnemesis/generator.go

@@ -406,7 +406,11 @@ type FaultConfig struct {
 	// RemoveNetworkPartition is an operation that simulates healing a network
 	// partition.
 	RemoveNetworkPartition int
-	// Disk stalls and node crashes belong here.
+	// StopNode is an operation that stops a randomly chosen node.
+	StopNode int
+	// RestartNode is an operation that restarts a randomly chosen node.
+	RestartNode int
+	// Disk stalls and other faults belong here.
 }
 
 // newAllOperationsConfig returns a GeneratorConfig that exercises *all*
@@ -531,6 +535,8 @@ func newAllOperationsConfig() GeneratorConfig {
 		Fault: FaultConfig{
 			AddNetworkPartition:    1,
 			RemoveNetworkPartition: 1,
+			StopNode:               1,
+			RestartNode:            1,
 		},
 	}}
 }
@@ -625,10 +631,12 @@ func NewDefaultConfig() GeneratorConfig {
 	config.Ops.ClosureTxn.TxnClientOps.FlushLockTable = 0
 	config.Ops.ClosureTxn.TxnBatchOps.Ops.FlushLockTable = 0
 
-	// Network partitions can result in unavailability and need to be enabled with
-	// care by specific test variants.
+	// Network partitions and node restarts can result in unavailability and need
+	// to be enabled with care by specific test variants.
 	config.Ops.Fault.AddNetworkPartition = 0
 	config.Ops.Fault.RemoveNetworkPartition = 0
+	config.Ops.Fault.StopNode = 0
+	config.Ops.Fault.RestartNode = 0
 	return config
 }
 
@@ -679,7 +687,7 @@ type Generator struct {
 
 // MakeGenerator constructs a Generator.
 func MakeGenerator(
-	config GeneratorConfig, replicasFn GetReplicasFn, mode TestMode,
+	config GeneratorConfig, replicasFn GetReplicasFn, mode TestMode, n *nodes,
 ) (*Generator, error) {
 	if config.NumNodes <= 0 {
 		return nil, errors.Errorf(`NumNodes must be positive got: %d`, config.NumNodes)
@@ -716,6 +724,7 @@ func MakeGenerator(
 		currentSplits:    make(map[string]struct{}),
 		historicalSplits: make(map[string]struct{}),
 		partitions:       p,
+		nodes:            n,
 		mode:             mode,
 	}
 	return g, nil
@@ -755,6 +764,9 @@ type generator struct {
 	// between nodes.
 	partitions
 
+	// nodes contains the sets of running and stopped nodes.
+	nodes *nodes
+
 	// mode is the test mode (e.g. Liveness or Safety). The generator needs it in
 	// order to set a timeout for range lookups under safety mode.
 	mode TestMode
@@ -770,6 +782,52 @@ type partitions struct {
 	partitioned map[connection]struct{}
 }
 
+// nodes contains the sets of running and stopped nodes. This struct is shared
+// between the generator and the applier to make sure nodes are promptly marked
+// as running/stopped when operations are generated/applied. The generator uses
+// removeRandRunning and removeRandStopped to pick nodes to stop/restart, and
+// the applier uses setRunning and setStopped to update the sets when operations
+// are actually applied. This is important because there could be a gap of
+// multiple seconds between generating a stop/restart operation and a node fully
+// stopping/restarting.
+type nodes struct {
+	mu      syncutil.RWMutex
+	running map[int]struct{}
+	stopped map[int]struct{}
+}
+
+func randNodeFromMap(m map[int]struct{}, rng *rand.Rand) int {
+	return maps.Keys(m)[rng.Intn(len(m))]
+}
+
+func (n *nodes) removeRandRunning(rng *rand.Rand) int {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	nodeID := randNodeFromMap(n.running, rng)
+	delete(n.running, nodeID)
+	return nodeID
+}
+
+func (n *nodes) removeRandStopped(rng *rand.Rand) int {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	nodeID := randNodeFromMap(n.stopped, rng)
+	delete(n.stopped, nodeID)
+	return nodeID
+}
+
+func (n *nodes) setRunning(nodeID int) {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	n.running[nodeID] = struct{}{}
+}
+
+func (n *nodes) setStopped(nodeID int) {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	n.stopped[nodeID] = struct{}{}
+}
+
 // RandStep returns a single randomly generated next operation to execute.
 //
 // RandStep is not concurrency safe.
@@ -851,6 +909,12 @@ func (g *generator) RandStep(rng *rand.Rand) Step {
 	addOpGen(&allowed, toggleGlobalReads, g.Config.Ops.ChangeZone.ToggleGlobalReads)
 	addOpGen(&allowed, addRandNetworkPartition, g.Config.Ops.Fault.AddNetworkPartition)
 	addOpGen(&allowed, removeRandNetworkPartition, g.Config.Ops.Fault.RemoveNetworkPartition)
+	if len(g.nodes.running) > 0 {
+		addOpGen(&allowed, stopRandNode, g.Config.Ops.Fault.StopNode)
+	}
+	if len(g.nodes.stopped) > 0 {
+		addOpGen(&allowed, restartRandNode, g.Config.Ops.Fault.RestartNode)
+	}
 
 	return step(g.selectOp(rng, allowed))
 }
@@ -1759,6 +1823,16 @@ func removeRandNetworkPartition(g *generator, rng *rand.Rand) Operation {
 	return removeNetworkPartition(randConn.from, randConn.to)
 }
 
+func stopRandNode(g *generator, rng *rand.Rand) Operation {
+	randNode := g.nodes.removeRandRunning(rng)
+	return stopNode(randNode)
+}
+
+func restartRandNode(g *generator, rng *rand.Rand) Operation {
+	randNode := g.nodes.removeRandStopped(rng)
+	return restartNode(randNode)
+}
+
 func makeRandBatch(c *ClientOperationConfig) opGenFunc {
 	return func(g *generator, rng *rand.Rand) Operation {
 		var allowed []opGen
@@ -2380,6 +2454,14 @@ func removeNetworkPartition(from int, to int) Operation {
 	}
 }
 
+func stopNode(nodeID int) Operation {
+	return Operation{StopNode: &StopNodeOperation{NodeId: int32(nodeID)}}
+}
+
+func restartNode(nodeID int) Operation {
+	return Operation{RestartNode: &RestartNodeOperation{NodeId: int32(nodeID)}}
+}
+
 type countingRandSource struct {
 	count atomic.Uint64
 	inner rand.Source64

pkg/kv/kvnemesis/generator_test.go

@@ -77,7 +77,11 @@ func TestRandStep(t *testing.T) {
 		return make([]roachpb.ReplicationTarget, rng.Intn(config.NumNodes)+1),
 			make([]roachpb.ReplicationTarget, rng.Intn(config.NumNodes)+1)
 	}
-	g, err := MakeGenerator(config, getReplicasFn, 0)
+	n := nodes{
+		running: map[int]struct{}{1: {}, 2: {}, 3: {}},
+		stopped: make(map[int]struct{}),
+	}
+	g, err := MakeGenerator(config, getReplicasFn, 0, &n)
 	require.NoError(t, err)
 
 	keys := make(map[string]struct{})
@@ -422,6 +426,16 @@ func TestRandStep(t *testing.T) {
 			counts.Fault.AddNetworkPartition++
 		case *RemoveNetworkPartitionOperation:
 			counts.Fault.RemoveNetworkPartition++
+		case *StopNodeOperation:
+			counts.Fault.StopNode++
+			n.mu.Lock()
+			n.stopped[int(o.NodeId)] = struct{}{}
+			n.mu.Unlock()
+		case *RestartNodeOperation:
+			counts.Fault.RestartNode++
+			n.mu.Lock()
+			n.running[int(o.NodeId)] = struct{}{}
+			n.mu.Unlock()
 		default:
 			t.Fatalf("%T", o)
 		}

pkg/kv/kvnemesis/kvnemesis.go

@@ -106,9 +106,23 @@ func RunNemesis(
 		return nil, fmt.Errorf("numSteps must be >0, got %v", numSteps)
 	}
 
-	dataSpan := GeneratorDataSpan()
+	n := nodes{
+		running: make(map[int]struct{}),
+		stopped: make(map[int]struct{}),
+	}
+	for i := 1; i <= config.NumNodes; i++ {
+		// In liveness mode, we don't allow stopping and restarting the two
+		// protected nodes (node 1 and node 2), so we don't include them in the set
+		// of running nodes at all.
+		protectedNode := i == 1 || i == 2
+		if mode == Liveness && protectedNode {
+			continue
+		}
+		n.running[i] = struct{}{}
+	}
 
-	g, err := MakeGenerator(config, newGetReplicasFn(dbs...), mode)
+	dataSpan := GeneratorDataSpan()
+	g, err := MakeGenerator(config, newGetReplicasFn(dbs...), mode, &n)
 	if err != nil {
 		return nil, err
 	}
@@ -118,7 +132,7 @@ func RunNemesis(
 	if mode == Liveness && len(applierDBs) >= 2 {
 		applierDBs = applierDBs[:2]
 	}
-	a := MakeApplier(env, applierDBs...)
+	a := MakeApplier(env, &n, applierDBs...)
 	w, err := Watch(ctx, env, dbs, dataSpan)
 	if err != nil {
 		return nil, err
@@ -182,6 +196,9 @@ func RunNemesis(
 		return nil, err
 	}
 	env.Partitioner.EnablePartitions(false)
+	for i := 0; i < config.NumNodes; i++ {
+		_ = env.Restarter.RestartServer(i)
+	}
 
 	allSteps := make(steps, 0, numSteps)
 	for _, steps := range stepsByWorker {