kvnemesis: add network partitions

This commit adds support for symmetric, asymmetric, partial and full
partitions to kvnemesis, using the `rpc.Partitioner`.

Release note: None

Author: miraradeva

pkg/kv/kvnemesis/BUILD.bazel

@@ -32,6 +32,7 @@ go_library(
         "//pkg/kv/kvserver/liveness",
         "//pkg/kv/kvtestutils",
         "//pkg/roachpb",
+        "//pkg/rpc",
         "//pkg/settings/cluster",
         "//pkg/sql/catalog/bootstrap",
         "//pkg/storage",
@@ -77,6 +78,7 @@ go_test(
     deps = [
         "//pkg/base",
         "//pkg/config/zonepb",
+        "//pkg/gossip",
         "//pkg/kv",
         "//pkg/kv/kvclient/kvcoord",
         "//pkg/kv/kvnemesis/kvnemesisutil",
@@ -89,6 +91,7 @@ go_test(
         "//pkg/kv/kvserver/kvserverpb",
         "//pkg/kv/kvtestutils",
         "//pkg/roachpb",
+        "//pkg/rpc",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
         "//pkg/server",

pkg/kv/kvnemesis/applier.go

@@ -150,6 +150,12 @@ func applyOp(ctx context.Context, env *Env, db *kv.DB, op *Operation) {
 		o.Result = resultInit(ctx, err)
 	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))
+		o.Result = resultInit(ctx, err)
+	case *RemoveNetworkPartitionOperation:
+		err := env.Partitioner.RemovePartition(roachpb.NodeID(o.FromNode), roachpb.NodeID(o.ToNode))
+		o.Result = resultInit(ctx, err)
 	case *ClosureTxnOperation:
 		// Use a backoff loop to avoid thrashing on txn aborts. Don't wait between
 		// epochs of the same transaction to avoid waiting while holding locks.

pkg/kv/kvnemesis/env.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvtestutils"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/errors"
 )
@@ -29,9 +30,10 @@ type Logger interface {
 // Env manipulates the environment (cluster settings, zone configurations) that
 // the Applier operates in.
 type Env struct {
-	SQLDBs  []*gosql.DB
-	Tracker *SeqTracker
-	L       Logger
+	SQLDBs      []*gosql.DB
+	Tracker     *SeqTracker
+	L           Logger
+	Partitioner *rpc.Partitioner
 }
 
 func (e *Env) anyNode() *gosql.DB {

pkg/kv/kvnemesis/generator.go

@@ -60,6 +60,7 @@ type OperationConfig struct {
 	ChangeLease    ChangeLeaseConfig
 	ChangeSetting  ChangeSettingConfig
 	ChangeZone     ChangeZoneConfig
+	Fault          FaultConfig
 }
 
 // ClosureTxnConfig configures the relative probability of running some
@@ -392,6 +393,20 @@ type SavepointConfig struct {
 	SavepointRollback int
 }
 
+// FaultConfig configures the relative probabilities of generating different
+// types of faults. Network partitions can be symmetric or asymmetric, partial
+// or full, but they may need multiple operations to set up; e.g. a symmetric
+// partition between node A and node B requires to partitions: from A to B, and
+// from B to A.
+type FaultConfig struct {
+	// AddNetworkPartition is an operation that simulates a network partition.
+	AddNetworkPartition int
+	// RemoveNetworkPartition is an operation that simulates healing a network
+	// partition.
+	RemoveNetworkPartition int
+	// Disk stalls and node crashes belong here.
+}
+
 // newAllOperationsConfig returns a GeneratorConfig that exercises *all*
 // options. You probably want NewDefaultConfig. Most of the time, these will be
 // the same, but having both allows us to merge code for operations that do not
@@ -511,6 +526,10 @@ func newAllOperationsConfig() GeneratorConfig {
 		ChangeZone: ChangeZoneConfig{
 			ToggleGlobalReads: 1,
 		},
+		Fault: FaultConfig{
+			AddNetworkPartition:    1,
+			RemoveNetworkPartition: 1,
+		},
 	}}
 }
 
@@ -603,6 +622,11 @@ func NewDefaultConfig() GeneratorConfig {
 	config.Ops.ClosureTxn.CommitBatchOps.FlushLockTable = 0
 	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.
+	config.Ops.Fault.AddNetworkPartition = 0
+	config.Ops.Fault.RemoveNetworkPartition = 0
 	return config
 }
 
@@ -663,13 +687,27 @@ func MakeGenerator(config GeneratorConfig, replicasFn GetReplicasFn) (*Generator
 		return nil, errors.Errorf(`NumReplicas (%d) must <= NumNodes (%d)`,
 			config.NumReplicas, config.NumNodes)
 	}
+	p := partitions{
+		healthy:     make(map[connection]struct{}),
+		partitioned: make(map[connection]struct{}),
+	}
+	for i := 1; i <= config.NumNodes; i++ {
+		for j := 1; j <= config.NumNodes; j++ {
+			if i == j {
+				continue
+			}
+			conn := connection{from: i, to: j}
+			p.healthy[conn] = struct{}{}
+		}
+	}
 	g := &Generator{}
 	g.mu.generator = generator{
 		Config:           config,
 		replicasFn:       replicasFn,
 		keys:             make(map[string]string),
 		currentSplits:    make(map[string]struct{}),
 		historicalSplits: make(map[string]struct{}),
+		partitions:       p,
 	}
 	return g, nil
 }
@@ -703,6 +741,20 @@ type generator struct {
 	// emitted, regardless of whether the split has since been applied or been
 	// merged away again.
 	historicalSplits map[string]struct{}
+
+	// partitions contains the sets of healthy and partitioned connections
+	// between nodes.
+	partitions
+}
+
+type connection struct {
+	from int // node ID
+	to   int // node ID
+}
+
+type partitions struct {
+	healthy     map[connection]struct{}
+	partitioned map[connection]struct{}
 }
 
 // RandStep returns a single randomly generated next operation to execute.
@@ -763,6 +815,8 @@ func (g *generator) RandStep(rng *rand.Rand) Step {
 
 	addOpGen(&allowed, setLeaseType, g.Config.Ops.ChangeSetting.SetLeaseType)
 	addOpGen(&allowed, toggleGlobalReads, g.Config.Ops.ChangeZone.ToggleGlobalReads)
+	addOpGen(&allowed, addRandNetworkPartition, g.Config.Ops.Fault.AddNetworkPartition)
+	addOpGen(&allowed, removeRandNetworkPartition, g.Config.Ops.Fault.RemoveNetworkPartition)
 
 	return step(g.selectOp(rng, allowed))
 }
@@ -1643,6 +1697,34 @@ func toggleGlobalReads(_ *generator, _ *rand.Rand) Operation {
 	return changeZone(ChangeZoneType_ToggleGlobalReads)
 }
 
+func addRandNetworkPartition(g *generator, rng *rand.Rand) Operation {
+	if len(g.partitions.healthy) == 0 {
+		return addNetworkPartition(0, 0)
+	}
+	all := make([]connection, 0, len(g.partitions.healthy))
+	for conn := range g.partitions.healthy {
+		all = append(all, conn)
+	}
+	randConn := all[rng.Intn(len(all))]
+	delete(g.partitions.healthy, randConn)
+	g.partitions.partitioned[randConn] = struct{}{}
+	return addNetworkPartition(randConn.from, randConn.to)
+}
+
+func removeRandNetworkPartition(g *generator, rng *rand.Rand) Operation {
+	if len(g.partitions.partitioned) == 0 {
+		return removeNetworkPartition(0, 0)
+	}
+	all := make([]connection, 0, len(g.partitions.partitioned))
+	for conn := range g.partitions.partitioned {
+		all = append(all, conn)
+	}
+	randConn := all[rng.Intn(len(all))]
+	delete(g.partitions.partitioned, randConn)
+	g.partitions.healthy[randConn] = struct{}{}
+	return removeNetworkPartition(randConn.from, randConn.to)
+}
+
 func makeRandBatch(c *ClientOperationConfig) opGenFunc {
 	return func(g *generator, rng *rand.Rand) Operation {
 		var allowed []opGen
@@ -2252,6 +2334,18 @@ func rollbackSavepoint(id int) Operation {
 	return Operation{SavepointRollback: &SavepointRollbackOperation{ID: int32(id)}}
 }
 
+func addNetworkPartition(from int, to int) Operation {
+	return Operation{
+		AddNetworkPartition: &AddNetworkPartitionOperation{FromNode: int32(from), ToNode: int32(to)},
+	}
+}
+
+func removeNetworkPartition(from int, to int) Operation {
+	return Operation{
+		RemoveNetworkPartition: &RemoveNetworkPartitionOperation{FromNode: int32(from), ToNode: int32(to)},
+	}
+}
+
 type countingRandSource struct {
 	count atomic.Uint64
 	inner rand.Source64

pkg/kv/kvnemesis/generator_test.go

@@ -418,6 +418,10 @@ func TestRandStep(t *testing.T) {
 			case ChangeZoneType_ToggleGlobalReads:
 				counts.ChangeZone.ToggleGlobalReads++
 			}
+		case *AddNetworkPartitionOperation:
+			counts.Fault.AddNetworkPartition++
+		case *RemoveNetworkPartitionOperation:
+			counts.Fault.RemoveNetworkPartition++
 		default:
 			t.Fatalf("%T", o)
 		}

pkg/kv/kvnemesis/kvnemesis.go

@@ -142,9 +142,11 @@ func RunNemesis(
 		}
 		return nil
 	}
+	env.Partitioner.EnablePartitions(true)
 	if err := ctxgroup.GroupWorkers(ctx, concurrency, workerFn); err != nil {
 		return nil, err
 	}
+	env.Partitioner.EnablePartitions(false)
 
 	allSteps := make(steps, 0, numSteps)
 	for _, steps := range stepsByWorker {

pkg/kv/kvnemesis/kvnemesis_test.go

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/gossip"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvclient/kvcoord"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
@@ -25,6 +26,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvtestutils"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/rpc"
+	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
@@ -46,7 +49,7 @@ import (
 var defaultNumSteps = envutil.EnvOrDefaultInt("COCKROACH_KVNEMESIS_STEPS", 100)
 
 func (cfg kvnemesisTestCfg) testClusterArgs(
-	ctx context.Context, tr *SeqTracker,
+	ctx context.Context, tr *SeqTracker, partitioner *rpc.Partitioner,
 ) base.TestClusterArgs {
 	storeKnobs := &kvserver.StoreTestingKnobs{
 		DisableRaftLogQueue:                   true,
@@ -173,31 +176,46 @@ func (cfg kvnemesisTestCfg) testClusterArgs(
 		cfg.testSettings(ctx, st)
 	}
 
-	args := base.TestClusterArgs{
-		ServerArgs: base.TestServerArgs{
-			Knobs: base.TestingKnobs{
-				Store: storeKnobs,
-				KVClient: &kvcoord.ClientTestingKnobs{
-					// Don't let DistSender split DeleteRangeUsingTombstone across range boundaries.
-					// This does happen in real CRDB, but leads to separate atomic subunits, which
-					// would add complexity to kvnemesis that isn't worth it. Instead, the operation
-					// generator for the most part tries to avoid range-spanning requests, and the
-					// ones that do end up happening get a hard error.
-					OnRangeSpanningNonTxnalBatch: func(ba *kvpb.BatchRequest) *kvpb.Error {
-						for _, req := range ba.Requests {
-							if req.GetInner().Method() != kvpb.DeleteRange {
-								continue
-							}
-							if req.GetDeleteRange().UseRangeTombstone == true {
-								return kvpb.NewError(errDelRangeUsingTombstoneStraddlesRangeBoundary)
-							}
+	commonServerArgs := base.TestServerArgs{
+		Knobs: base.TestingKnobs{
+			Store: storeKnobs,
+			KVClient: &kvcoord.ClientTestingKnobs{
+				// Don't let DistSender split DeleteRangeUsingTombstone across range boundaries.
+				// This does happen in real CRDB, but leads to separate atomic subunits, which
+				// would add complexity to kvnemesis that isn't worth it. Instead, the operation
+				// generator for the most part tries to avoid range-spanning requests, and the
+				// ones that do end up happening get a hard error.
+				OnRangeSpanningNonTxnalBatch: func(ba *kvpb.BatchRequest) *kvpb.Error {
+					for _, req := range ba.Requests {
+						if req.GetInner().Method() != kvpb.DeleteRange {
+							continue
+						}
+						if req.GetDeleteRange().UseRangeTombstone == true {
+							return kvpb.NewError(errDelRangeUsingTombstoneStraddlesRangeBoundary)
 						}
-						return nil
-					},
+					}
+					return nil
 				},
 			},
-			Settings: st,
 		},
+		Settings: st,
+	}
+
+	args := base.TestClusterArgs{
+		ServerArgs: commonServerArgs,
+		ServerArgsPerNode: func() map[int]base.TestServerArgs {
+			perNode := make(map[int]base.TestServerArgs)
+			for i := 0; i < cfg.numNodes; i++ {
+				ctk := rpc.ContextTestingKnobs{}
+				partitioner.RegisterTestingKnobs(roachpb.NodeID(i+1), &ctk)
+				perNodeServerArgs := commonServerArgs
+				perNodeServerArgs.Knobs.Server = &server.TestingKnobs{
+					ContextTestingKnobs: ctk,
+				}
+				perNode[i] = perNodeServerArgs
+			}
+			return perNode
+		}(),
 	}
 
 	if cfg.testArgs != nil {
@@ -470,8 +488,15 @@ func testKVNemesisImpl(t testing.TB, cfg kvnemesisTestCfg) {
 	// 4 nodes so we have somewhere to move 3x replicated ranges to.
 	ctx := context.Background()
 	tr := &SeqTracker{}
-	tc := testcluster.StartTestCluster(t, cfg.numNodes, cfg.testClusterArgs(ctx, tr))
+	var partitioner rpc.Partitioner
+	tc := testcluster.StartTestCluster(t, cfg.numNodes, cfg.testClusterArgs(ctx, tr, &partitioner))
 	defer tc.Stopper().Stop(ctx)
+	for i := 0; i < cfg.numNodes; i++ {
+		g := tc.Servers[i].StorageLayer().GossipI().(*gossip.Gossip)
+		addr := g.GetNodeAddr().String()
+		nodeID := g.NodeID.Get()
+		partitioner.RegisterNodeAddr(addr, nodeID)
+	}
 	dbs, sqlDBs := make([]*kv.DB, cfg.numNodes), make([]*gosql.DB, cfg.numNodes)
 	for i := 0; i < cfg.numNodes; i++ {
 		dbs[i] = tc.Server(i).DB()
@@ -500,7 +525,7 @@ func testKVNemesisImpl(t testing.TB, cfg kvnemesisTestCfg) {
 
 	logger := newTBridge(t)
 	defer dumpRaftLogsOnFailure(t, logger.ll.dir, tc.Servers)
-	env := &Env{SQLDBs: sqlDBs, Tracker: tr, L: logger}
+	env := &Env{SQLDBs: sqlDBs, Tracker: tr, L: logger, Partitioner: &partitioner}
 	failures, err := RunNemesis(ctx, rng, env, config, cfg.concurrency, cfg.numSteps, dbs...)
 
 	for i := 0; i < cfg.numNodes; i++ {