failure-injection: add disk stall failure modes

pkg/cmd/roachtest/cluster_test.go

@@ -8,6 +8,7 @@ package main
 import (
 	"context"
 	"fmt"
+	"math/rand"
 	"strconv"
 	"strings"
 	"testing"
@@ -66,6 +67,35 @@ func TestClusterNodes(t *testing.T) {
 	}
 }
 
+func TestSeededRandGroups(t *testing.T) {
+	rng := rand.New(rand.NewSource(1))
+	testCases := []struct {
+		numNodes  int
+		numGroups int
+		expected  []string
+	}{
+		{numNodes: 1, numGroups: 1, expected: []string{":1"}},
+		{numNodes: 10, numGroups: 1, expected: []string{":1-10"}},
+		{numNodes: 10, numGroups: 2, expected: []string{":1,3,8,10", ":2,4-7,9"}},
+		{numNodes: 3, numGroups: 3, expected: []string{":3", ":2", ":1"}},
+		{numNodes: 5, numGroups: 3, expected: []string{":2", ":1,3-4", ":5"}},
+	}
+	for _, tc := range testCases {
+		t.Run("", func(t *testing.T) {
+			c := &clusterImpl{spec: spec.MakeClusterSpec(tc.numNodes)}
+			nodes := c.All()
+			groups, err := nodes.SeededRandGroups(rng, tc.numGroups)
+			require.NoError(t, err)
+			for i, group := range groups {
+				nodeList := c.MakeNodes(group)
+				if tc.expected[i] != nodeList {
+					t.Errorf("expected %s, but found %s", tc.expected[i], nodeList)
+				}
+			}
+		})
+	}
+}
+
 type testWrapper struct {
 	*testing.T
 	l *logger.Logger

pkg/cmd/roachtest/option/node_list_option.go

@@ -83,6 +83,29 @@ func (n NodeListOption) SeededRandList(rand *rand.Rand, size int) (NodeListOptio
 	return nodes[:size], nil
 }
 
+// SeededRandGroups splits up the NodeListOption into numGroups groups of nodes using
+// a seeded rand object. Nodes are not guaranteed to be evenly distributed among groups,
+// but groups are guaranteed to have at least one node.
+func (n NodeListOption) SeededRandGroups(rand *rand.Rand, numGroups int) ([]NodeListOption, error) {
+	if numGroups > len(n) {
+		return nil, fmt.Errorf("cannot partition nodes - numGroups: %d > len: %d", numGroups, len(n))
+	}
+
+	nodes := append([]int{}, n...)
+	rand.Shuffle(len(nodes), func(i, j int) { nodes[i], nodes[j] = nodes[j], nodes[i] })
+
+	groups := make([]NodeListOption, numGroups)
+	// Assign each group at least one node.
+	for i := 0; i < numGroups; i++ {
+		groups[i] = NodeListOption{nodes[i]}
+	}
+	for i := numGroups; i < len(nodes); i++ {
+		groupIdx := rand.Intn(numGroups)
+		groups[groupIdx] = append(groups[groupIdx], nodes[i])
+	}
+	return groups, nil
+}
+
 // NodeIDsString returns the nodes in the NodeListOption, separated by spaces.
 func (n NodeListOption) NodeIDsString() string {
 	result := ""

pkg/cmd/roachtest/tests/BUILD.bazel

@@ -226,6 +226,7 @@ go_library(
         "//pkg/clusterversion",
         "//pkg/cmd/cmpconn",
         "//pkg/cmd/roachprod-microbench/util",
+        "//pkg/cmd/roachprod/grafana",
         "//pkg/cmd/roachtest/cluster",
         "//pkg/cmd/roachtest/clusterstats",
         "//pkg/cmd/roachtest/grafana",