roachprod: add stop option to reset VM failure

When the reset VM failure is injected, we now have the option to stop the
processes before restarting the cluster.

This change also improves process management by starting the processes in the
correct order (System interface first, then tenants).

Epic: None
Release note: None

Author: herkolategan

pkg/roachprod/failureinjection/failures/reset.go

@@ -7,6 +7,7 @@ package failures
 
 import (
 	"context"
+	"slices"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/roachprod"
@@ -16,12 +17,22 @@ import (
 
 type (
 	ResetVMArgs struct {
-		Nodes install.Nodes
+		Nodes         install.Nodes
+		StopProcesses bool
 	}
 	resetVMFailure struct {
 		GenericFailure
-		Processes map[install.Node][]install.MonitorProcessRunning
+		processes processMap
 	}
+	// nodeSet is a set of nodes.
+	nodeSet map[install.Node]struct{}
+	// instanceMap is a map of SQL instances to a map of nodes that are running
+	// that instance.
+	instanceMap map[int]nodeSet
+	// processMap is a map of virtual cluster names to instance maps. It's a
+	// convenience type that allows grouping the processes that should be
+	// started and stopped together.
+	processMap map[string]instanceMap
 )
 
 var _ FailureMode = &resetVMFailure{}
@@ -47,6 +58,43 @@ func MakeResetVMFailure(
 	}, nil
 }
 
+func (m *processMap) add(virtualClusterName string, instance int, node install.Node) {
+	if virtualClusterName == "" {
+		virtualClusterName = install.SystemInterfaceName
+	}
+	if _, ok := (*m)[virtualClusterName]; !ok {
+		(*m)[virtualClusterName] = make(map[int]nodeSet, 0)
+	}
+	if _, ok := (*m)[virtualClusterName][instance]; !ok {
+		(*m)[virtualClusterName][instance] = make(nodeSet, 0)
+	}
+	(*m)[virtualClusterName][instance][node] = struct{}{}
+}
+
+// getStartOrder returns the order in which the processes should be started. It
+// ensures that the System interface is started first.
+func (m *processMap) getStartOrder() []string {
+	var order []string
+	// If the System interface is present, it should be the first to start.
+	if _, ok := (*m)[install.SystemInterfaceName]; ok {
+		order = append(order, install.SystemInterfaceName)
+	}
+	for virtualClusterName := range *m {
+		if virtualClusterName != install.SystemInterfaceName {
+			order = append(order, virtualClusterName)
+		}
+	}
+	return order
+}
+
+// getStopOrder returns the order in which the processes should be stopped. It
+// is the reverse of the start order.
+func (m *processMap) getStopOrder() []string {
+	order := m.getStartOrder()
+	slices.Reverse(order)
+	return order
+}
+
 // Description implements FailureMode.
 func (r *resetVMFailure) Description() string {
 	return ResetVMFailureName
@@ -62,10 +110,29 @@ func (r *resetVMFailure) Inject(ctx context.Context, l *logger.Logger, args Fail
 	// Capture the processes running on the nodes.
 	nodes := args.(ResetVMArgs).Nodes
 	monitorChan := r.c.WithNodes(nodes).Monitor(l, ctx, install.MonitorOpts{OneShot: true})
-	r.Processes = make(map[install.Node][]install.MonitorProcessRunning, 0)
+	r.processes = make(processMap, 0)
 	for e := range monitorChan {
 		if p, ok := e.Event.(install.MonitorProcessRunning); ok {
-			r.Processes[e.Node] = append(r.Processes[e.Node], p)
+			r.processes.add(p.VirtualClusterName, p.SQLInstance, e.Node)
+		}
+	}
+
+	// Optionally stop the processes.
+	if args.(ResetVMArgs).StopProcesses {
+		for _, virtualClusterName := range r.processes.getStopOrder() {
+			instanceMap := r.processes[virtualClusterName]
+			for _, nodeMap := range instanceMap {
+				var stopNodes install.Nodes
+				for node := range nodeMap {
+					stopNodes = append(nodes, node)
+				}
+				l.Printf("Stopping process %s on nodes %v", virtualClusterName, stopNodes)
+				stopOpts := roachprod.DefaultStopOpts()
+				err := r.c.WithNodes(stopNodes).Stop(ctx, l, stopOpts.Sig, stopOpts.Wait, stopOpts.GracePeriod, virtualClusterName)
+				if err != nil {
+					return err
+				}
+			}
 		}
 	}
 
@@ -80,12 +147,17 @@ func (r *resetVMFailure) Cleanup(ctx context.Context, l *logger.Logger, args Fai
 // Recover implements FailureMode.
 func (r *resetVMFailure) Recover(ctx context.Context, l *logger.Logger, args FailureArgs) error {
 	// Restart the processes.
-	for node, processes := range r.Processes {
-		for _, p := range processes {
-			l.Printf("Starting process %s on node %s", p.PID, p.VirtualClusterName)
-			err := r.c.WithNodes([]install.Node{node}).Start(ctx, l, install.StartOpts{
-				VirtualClusterName: p.VirtualClusterName,
-				SQLInstance:        p.SQLInstance,
+	for _, virtualClusterName := range r.processes.getStartOrder() {
+		instanceMap := r.processes[virtualClusterName]
+		for instance, nodeMap := range instanceMap {
+			var nodes install.Nodes
+			for node := range nodeMap {
+				nodes = append(nodes, node)
+			}
+			l.Printf("Starting process %s on nodes %v", virtualClusterName, nodes)
+			err := r.c.WithNodes(nodes).Start(ctx, l, install.StartOpts{
+				VirtualClusterName: virtualClusterName,
+				SQLInstance:        instance,
 				IsRestart:          true,
 			})
 			if err != nil {