kvnemesis: add lease transfers

This commit adds lease transfers to KV nemesis. This will be useful in
validating that I don't mess anything up when addressing cockroachdb#57688 or when
performing the precursor refactor to pull lease checks below latching.

I stressed this for 50,000 iterations on a 20 node roachprod cluster
without failure, so it looks like lease transfers aren't already broken!

pkg/kv/kvnemesis/applier.go

@@ -82,6 +82,9 @@ func applyOp(ctx context.Context, db *kv.DB, op *Operation) {
 		_, err := db.AdminChangeReplicas(ctx, o.Key, desc, o.Changes)
 		// TODO(dan): Save returned desc?
 		o.Result = resultError(ctx, err)
+	case *TransferLeaseOperation:
+		err := db.AdminTransferLease(ctx, o.Key, o.Target)
+		o.Result = resultError(ctx, err)
 	case *ClosureTxnOperation:
 		var savedTxn *kv.Txn
 		txnErr := db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {

pkg/kv/kvnemesis/applier_test.go

@@ -140,4 +140,10 @@ db0.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
 		`db1.AdminSplit(ctx, "foo") // context canceled`)
 	checkErr(t, step(merge(`foo`)),
 		`db0.AdminMerge(ctx, "foo") // context canceled`)
+
+	// Lease transfers
+	check(t, step(transferLease(`foo`, 1)),
+		`db1.TransferLeaseOperation(ctx, "foo", 1) // nil`)
+	checkErr(t, step(transferLease(`foo`, 1)),
+		`db0.TransferLeaseOperation(ctx, "foo", 1) // context canceled`)
 }

pkg/kv/kvnemesis/doc.go

@@ -25,7 +25,10 @@
 // TODO
 // - CPut/InitPut/Increment/Delete
 // - DeleteRange/ClearRange/RevertRange/ReverseScan
-// - TransferLease
+// - AdminRelocateRange
+// - AdminUnsplit
+// - AdminScatter
+// - CheckConsistency
 // - ExportRequest
 // - AddSSTable
 // - Root and leaf transactions

pkg/kv/kvnemesis/generator.go

@@ -42,6 +42,7 @@ type OperationConfig struct {
 	Split          SplitConfig
 	Merge          MergeConfig
 	ChangeReplicas ChangeReplicasConfig
+	ChangeLease    ChangeLeaseConfig
 }
 
 // ClosureTxnConfig configures the relative probability of running some
@@ -129,6 +130,13 @@ type ChangeReplicasConfig struct {
 	AtomicSwapReplica int
 }
 
+// ChangeLeaseConfig configures the relative probability of generating an
+// operation that causes a leaseholder change.
+type ChangeLeaseConfig struct {
+	// Transfer the lease to a random replica.
+	TransferLease int
+}
+
 // 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
@@ -170,6 +178,9 @@ func newAllOperationsConfig() GeneratorConfig {
 			RemoveReplica:     1,
 			AtomicSwapReplica: 1,
 		},
+		ChangeLease: ChangeLeaseConfig{
+			TransferLease: 1,
+		},
 	}}
 }
 
@@ -324,6 +335,8 @@ func (g *generator) RandStep(rng *rand.Rand) Step {
 		removeReplicaFn := makeRemoveReplicaFn(key, current)
 		addOpGen(&allowed, removeReplicaFn, g.Config.Ops.ChangeReplicas.RemoveReplica)
 	}
+	transferLeaseFn := makeTransferLeaseFn(key, current)
+	addOpGen(&allowed, transferLeaseFn, g.Config.Ops.ChangeLease.TransferLease)
 
 	return step(g.selectOp(rng, allowed))
 }
@@ -473,6 +486,13 @@ func makeAddReplicaFn(key string, current []roachpb.ReplicationTarget, atomicSwa
 	}
 }
 
+func makeTransferLeaseFn(key string, current []roachpb.ReplicationTarget) opGenFunc {
+	return func(g *generator, rng *rand.Rand) Operation {
+		target := current[rng.Intn(len(current))]
+		return transferLease(key, target.StoreID)
+	}
+}
+
 func makeRandBatch(c *ClientOperationConfig) opGenFunc {
 	return func(g *generator, rng *rand.Rand) Operation {
 		var allowed []opGen
@@ -600,3 +620,7 @@ func merge(key string) Operation {
 func changeReplicas(key string, changes ...roachpb.ReplicationChange) Operation {
 	return Operation{ChangeReplicas: &ChangeReplicasOperation{Key: []byte(key), Changes: changes}}
 }
+
+func transferLease(key string, target roachpb.StoreID) Operation {
+	return Operation{TransferLease: &TransferLeaseOperation{Key: []byte(key), Target: target}}
+}

pkg/kv/kvnemesis/generator_test.go

@@ -169,6 +169,8 @@ func TestRandStep(t *testing.T) {
 			} else if adds == 1 && removes == 1 {
 				counts.ChangeReplicas.AtomicSwapReplica++
 			}
+		case *TransferLeaseOperation:
+			counts.ChangeLease.TransferLease++
 		}
 		updateKeys(step.Op)
 

pkg/kv/kvnemesis/operations.go

@@ -33,6 +33,10 @@ func (op Operation) Result() *Result {
 		return &o.Result
 	case *MergeOperation:
 		return &o.Result
+	case *ChangeReplicasOperation:
+		return &o.Result
+	case *TransferLeaseOperation:
+		return &o.Result
 	case *BatchOperation:
 		return &o.Result
 	case *ClosureTxnOperation:
@@ -104,6 +108,8 @@ func (op Operation) format(w *strings.Builder, fctx formatCtx) {
 		o.format(w, fctx)
 	case *ChangeReplicasOperation:
 		o.format(w, fctx)
+	case *TransferLeaseOperation:
+		o.format(w, fctx)
 	case *BatchOperation:
 		newFctx := fctx
 		newFctx.indent = fctx.indent + `  `
@@ -231,6 +237,11 @@ func (op ChangeReplicasOperation) format(w *strings.Builder, fctx formatCtx) {
 	op.Result.format(w)
 }
 
+func (op TransferLeaseOperation) format(w *strings.Builder, fctx formatCtx) {
+	fmt.Fprintf(w, `%s.TransferLeaseOperation(ctx, %s, %d)`, fctx.receiver, roachpb.Key(op.Key), op.Target)
+	op.Result.format(w)
+}
+
 func (r Result) format(w *strings.Builder) {
 	switch r.Type {
 	case ResultType_NoError: