kv/kvnemesis: add Txn.Prepare

Informs cockroachdb#22329.

This commit adds support for Txn.Prepare in kvnemesis.

Release note: None

pkg/kv/kvnemesis/applier.go

@@ -209,6 +209,11 @@ func applyOp(ctx context.Context, env *Env, db *kv.DB, op *Operation) {
 				}
 				return err
 			}
+			if o.Prepare {
+				if err := txn.Prepare(ctx); err != nil {
+					return err
+				}
+			}
 			if o.CommitInBatch != nil {
 				b := txn.NewBatch()
 				applyBatchOp(ctx, b, txn.CommitInBatch, o.CommitInBatch)

pkg/kv/kvnemesis/applier_test.go

@@ -280,6 +280,18 @@ func TestApplier(t *testing.T) {
 		{
 			"txn-si-rollback", step(closureTxn(ClosureTxnType_Rollback, isolation.Snapshot, put(k5, 5))),
 		},
+		{
+			"txn-ssi-prepare-commit", step(closureTxnPrepare(ClosureTxnType_Rollback, isolation.Serializable, put(k5, 5))),
+		},
+		{
+			"txn-si-prepare-commit", step(closureTxnPrepare(ClosureTxnType_Rollback, isolation.Snapshot, put(k5, 5))),
+		},
+		{
+			"txn-ssi-prepare-rollback", step(closureTxnPrepare(ClosureTxnType_Rollback, isolation.Serializable, put(k5, 5))),
+		},
+		{
+			"txn-si-prepare-rollback", step(closureTxnPrepare(ClosureTxnType_Rollback, isolation.Snapshot, put(k5, 5))),
+		},
 		{
 			"split", step(split(k2)),
 		},

pkg/kv/kvnemesis/generator.go

@@ -90,6 +90,24 @@ type ClosureTxnConfig struct {
 	// via the CommitInBatchMethod. This is an important part of the 1pc txn
 	// fastpath.
 	CommitReadCommittedInBatch int
+	// PrepareCommitSerializable is a serializable transaction that prepares and
+	// then commits normally.
+	PrepareCommitSerializable int
+	// PrepareCommitSnapshot is a snapshot transaction that prepares and then
+	// commits normally.
+	PrepareCommitSnapshot int
+	// PrepareCommitReadCommitted is a read committed transaction that prepares
+	// and then commits normally.
+	PrepareCommitReadCommitted int
+	// PrepareRollbackSerializable is a serializable transaction that prepares and
+	// encounters an (external) error at the end and has to roll back.
+	PrepareRollbackSerializable int
+	// PrepareRollbackSnapshot is a snapshot transaction that prepares and
+	// encounters an (external) error at the end and has to roll back.
+	PrepareRollbackSnapshot int
+	// PrepareRollbackReadCommitted is a read committed transaction that prepares
+	// and encounters an (external) error at the end and has to roll back.
+	PrepareRollbackReadCommitted int
 
 	TxnClientOps ClientOperationConfig
 	TxnBatchOps  BatchOperationConfig
@@ -413,19 +431,25 @@ func newAllOperationsConfig() GeneratorConfig {
 		DB:    clientOpConfig,
 		Batch: batchOpConfig,
 		ClosureTxn: ClosureTxnConfig{
-			CommitSerializable:         2,
-			CommitSnapshot:             2,
-			CommitReadCommitted:        2,
-			RollbackSerializable:       2,
-			RollbackSnapshot:           2,
-			RollbackReadCommitted:      2,
-			CommitSerializableInBatch:  2,
-			CommitSnapshotInBatch:      2,
-			CommitReadCommittedInBatch: 2,
-			TxnClientOps:               clientOpConfig,
-			TxnBatchOps:                batchOpConfig,
-			CommitBatchOps:             clientOpConfig,
-			SavepointOps:               savepointConfig,
+			CommitSerializable:           2,
+			CommitSnapshot:               2,
+			CommitReadCommitted:          2,
+			RollbackSerializable:         2,
+			RollbackSnapshot:             2,
+			RollbackReadCommitted:        2,
+			CommitSerializableInBatch:    2,
+			CommitSnapshotInBatch:        2,
+			CommitReadCommittedInBatch:   2,
+			PrepareCommitSerializable:    1,
+			PrepareCommitSnapshot:        1,
+			PrepareCommitReadCommitted:   1,
+			PrepareRollbackSerializable:  1,
+			PrepareRollbackSnapshot:      1,
+			PrepareRollbackReadCommitted: 1,
+			TxnClientOps:                 clientOpConfig,
+			TxnBatchOps:                  batchOpConfig,
+			CommitBatchOps:               clientOpConfig,
+			SavepointOps:                 savepointConfig,
 		},
 		Split: SplitConfig{
 			SplitNew:   1,
@@ -1477,29 +1501,43 @@ func makeRandBatch(c *ClientOperationConfig) opGenFunc {
 func (g *generator) registerClosureTxnOps(allowed *[]opGen, c *ClosureTxnConfig) {
 	const Commit, Rollback = ClosureTxnType_Commit, ClosureTxnType_Rollback
 	const SSI, SI, RC = isolation.Serializable, isolation.Snapshot, isolation.ReadCommitted
+	const Prepare, NoPrepare = true, false
 	addOpGen(allowed,
-		makeClosureTxn(Commit, SSI, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.CommitSerializable)
+		makeClosureTxn(Commit, SSI, NoPrepare, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.CommitSerializable)
 	addOpGen(allowed,
-		makeClosureTxn(Commit, SI, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.CommitSnapshot)
+		makeClosureTxn(Commit, SI, NoPrepare, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.CommitSnapshot)
 	addOpGen(allowed,
-		makeClosureTxn(Commit, RC, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.CommitReadCommitted)
+		makeClosureTxn(Commit, RC, NoPrepare, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.CommitReadCommitted)
 	addOpGen(allowed,
-		makeClosureTxn(Rollback, SSI, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.RollbackSerializable)
+		makeClosureTxn(Rollback, SSI, NoPrepare, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.RollbackSerializable)
 	addOpGen(allowed,
-		makeClosureTxn(Rollback, SI, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.RollbackSnapshot)
+		makeClosureTxn(Rollback, SI, NoPrepare, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.RollbackSnapshot)
 	addOpGen(allowed,
-		makeClosureTxn(Rollback, RC, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.RollbackReadCommitted)
+		makeClosureTxn(Rollback, RC, NoPrepare, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.RollbackReadCommitted)
 	addOpGen(allowed,
-		makeClosureTxn(Commit, SSI, &c.TxnClientOps, &c.TxnBatchOps, &c.CommitBatchOps, &c.SavepointOps), c.CommitSerializableInBatch)
+		makeClosureTxn(Commit, SSI, NoPrepare, &c.TxnClientOps, &c.TxnBatchOps, &c.CommitBatchOps, &c.SavepointOps), c.CommitSerializableInBatch)
 	addOpGen(allowed,
-		makeClosureTxn(Commit, SI, &c.TxnClientOps, &c.TxnBatchOps, &c.CommitBatchOps, &c.SavepointOps), c.CommitSnapshotInBatch)
+		makeClosureTxn(Commit, SI, NoPrepare, &c.TxnClientOps, &c.TxnBatchOps, &c.CommitBatchOps, &c.SavepointOps), c.CommitSnapshotInBatch)
 	addOpGen(allowed,
-		makeClosureTxn(Commit, RC, &c.TxnClientOps, &c.TxnBatchOps, &c.CommitBatchOps, &c.SavepointOps), c.CommitReadCommittedInBatch)
+		makeClosureTxn(Commit, RC, NoPrepare, &c.TxnClientOps, &c.TxnBatchOps, &c.CommitBatchOps, &c.SavepointOps), c.CommitReadCommittedInBatch)
+	addOpGen(allowed,
+		makeClosureTxn(Commit, SSI, Prepare, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.PrepareCommitSerializable)
+	addOpGen(allowed,
+		makeClosureTxn(Commit, SI, Prepare, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.PrepareCommitSnapshot)
+	addOpGen(allowed,
+		makeClosureTxn(Commit, RC, Prepare, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.PrepareCommitReadCommitted)
+	addOpGen(allowed,
+		makeClosureTxn(Rollback, SSI, Prepare, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.PrepareRollbackSerializable)
+	addOpGen(allowed,
+		makeClosureTxn(Rollback, SI, Prepare, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.PrepareRollbackSnapshot)
+	addOpGen(allowed,
+		makeClosureTxn(Rollback, RC, Prepare, &c.TxnClientOps, &c.TxnBatchOps, nil /* commitInBatch*/, &c.SavepointOps), c.PrepareRollbackReadCommitted)
 }
 
 func makeClosureTxn(
 	txnType ClosureTxnType,
 	iso isolation.Level,
+	prepare bool,
 	txnClientOps *ClientOperationConfig,
 	txnBatchOps *BatchOperationConfig,
 	commitInBatch *ClientOperationConfig,
@@ -1532,10 +1570,14 @@ func makeClosureTxn(
 			maybeUpdateSavepoints(&spIDs, ops[i])
 		}
 		op := closureTxn(txnType, iso, ops...)
+		op.ClosureTxn.Prepare = prepare
 		if commitInBatch != nil {
 			if txnType != ClosureTxnType_Commit {
 				panic(errors.AssertionFailedf(`CommitInBatch must commit got: %s`, txnType))
 			}
+			if prepare {
+				panic(errors.AssertionFailedf(`CommitInBatch cannot prepare`))
+			}
 			op.ClosureTxn.CommitInBatch = makeRandBatch(commitInBatch)(g, rng).Batch
 		}
 		return op
@@ -1768,6 +1810,12 @@ func closureTxnSSI(typ ClosureTxnType, ops ...Operation) Operation {
 	return closureTxn(typ, isolation.Serializable, ops...)
 }
 
+func closureTxnPrepare(typ ClosureTxnType, iso isolation.Level, ops ...Operation) Operation {
+	o := closureTxn(typ, iso, ops...)
+	o.ClosureTxn.Prepare = true
+	return o
+}
+
 func closureTxnCommitInBatch(
 	iso isolation.Level, commitInBatch []Operation, ops ...Operation,
 ) Operation {

pkg/kv/kvnemesis/generator_test.go

@@ -309,22 +309,46 @@ func TestRandStep(t *testing.T) {
 			} else if o.Type == ClosureTxnType_Commit {
 				switch o.IsoLevel {
 				case isolation.Serializable:
-					counts.ClosureTxn.CommitSerializable++
+					if o.Prepare {
+						counts.ClosureTxn.PrepareCommitSerializable++
+					} else {
+						counts.ClosureTxn.CommitSerializable++
+					}
 				case isolation.Snapshot:
-					counts.ClosureTxn.CommitSnapshot++
+					if o.Prepare {
+						counts.ClosureTxn.PrepareCommitSnapshot++
+					} else {
+						counts.ClosureTxn.CommitSnapshot++
+					}
 				case isolation.ReadCommitted:
-					counts.ClosureTxn.CommitReadCommitted++
+					if o.Prepare {
+						counts.ClosureTxn.PrepareCommitReadCommitted++
+					} else {
+						counts.ClosureTxn.CommitReadCommitted++
+					}
 				default:
 					t.Fatalf("unexpected isolation level %s", o.IsoLevel)
 				}
 			} else if o.Type == ClosureTxnType_Rollback {
 				switch o.IsoLevel {
 				case isolation.Serializable:
-					counts.ClosureTxn.RollbackSerializable++
+					if o.Prepare {
+						counts.ClosureTxn.PrepareRollbackSerializable++
+					} else {
+						counts.ClosureTxn.RollbackSerializable++
+					}
 				case isolation.Snapshot:
-					counts.ClosureTxn.RollbackSnapshot++
+					if o.Prepare {
+						counts.ClosureTxn.PrepareRollbackSnapshot++
+					} else {
+						counts.ClosureTxn.RollbackSnapshot++
+					}
 				case isolation.ReadCommitted:
-					counts.ClosureTxn.RollbackReadCommitted++
+					if o.Prepare {
+						counts.ClosureTxn.PrepareRollbackReadCommitted++
+					} else {
+						counts.ClosureTxn.RollbackReadCommitted++
+					}
 				default:
 					t.Fatalf("unexpected isolation level %s", o.IsoLevel)
 				}

pkg/kv/kvnemesis/operations.go

@@ -177,6 +177,12 @@ func (op Operation) format(w *strings.Builder, fctx formatCtx) {
 		w.WriteString(newFctx.receiver)
 		fmt.Fprintf(w, `.SetIsoLevel(isolation.%s)`, o.IsoLevel)
 		formatOps(w, newFctx, o.Ops)
+		if o.Prepare {
+			w.WriteString("\n")
+			w.WriteString(newFctx.indent)
+			w.WriteString(newFctx.receiver)
+			w.WriteString(`.Prepare(ctx)`)
+		}
 		if o.CommitInBatch != nil {
 			newFctx.receiver = `b`
 			o.CommitInBatch.format(w, newFctx)

pkg/kv/kvnemesis/operations.proto

@@ -36,6 +36,7 @@ message ClosureTxnOperation {
   BatchOperation commit_in_batch = 3;
   ClosureTxnType type = 4;
   cockroach.kv.kvserver.concurrency.isolation.Level iso_level = 7;
+  bool prepare = 8;
   Result result = 5 [(gogoproto.nullable) = false];
   roachpb.Transaction txn = 6;
 }

pkg/kv/kvnemesis/operations_test.go

@@ -97,6 +97,13 @@ func TestOperationsFormat(t *testing.T) {
 		},
 		{step: step(barrier(k1, k2, false /* withLAI */))},
 		{step: step(barrier(k3, k4, true /* withLAI */))},
+		{
+			step: step(
+				closureTxnPrepare(ClosureTxnType_Rollback,
+					isolation.ReadCommitted,
+					get(k8),
+				)),
+		},
 	}
 
 	w := echotest.NewWalker(t, datapathutils.TestDataPath(t, t.Name()))

pkg/kv/kvnemesis/testdata/TestApplier/txn-si-prepare-commit

@@ -0,0 +1,8 @@
+echo
+----
+db0.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+  txn.SetIsoLevel(isolation.Snapshot)
+  txn.Put(ctx, tk(5), sv(5)) // @<ts> <nil>
+  txn.Prepare(ctx)
+  return errors.New("rollback")
+}) // rollback

pkg/kv/kvnemesis/testdata/TestApplier/txn-si-prepare-rollback

@@ -0,0 +1,8 @@
+echo
+----
+db0.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+  txn.SetIsoLevel(isolation.Snapshot)
+  txn.Put(ctx, tk(5), sv(5)) // @<ts> <nil>
+  txn.Prepare(ctx)
+  return errors.New("rollback")
+}) // rollback

pkg/kv/kvnemesis/testdata/TestApplier/txn-ssi-prepare-commit

@@ -0,0 +1,8 @@
+echo
+----
+db0.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+  txn.SetIsoLevel(isolation.Serializable)
+  txn.Put(ctx, tk(5), sv(5)) // @<ts> <nil>
+  txn.Prepare(ctx)
+  return errors.New("rollback")
+}) // rollback

pkg/kv/kvnemesis/testdata/TestApplier/txn-ssi-prepare-rollback

@@ -0,0 +1,8 @@
+echo
+----
+db0.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+  txn.SetIsoLevel(isolation.Serializable)
+  txn.Put(ctx, tk(5), sv(5)) // @<ts> <nil>
+  txn.Prepare(ctx)
+  return errors.New("rollback")
+}) // rollback

pkg/kv/kvnemesis/testdata/TestOperationsFormat/8

@@ -0,0 +1,8 @@
+echo
+----
+···db0.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+···  txn.SetIsoLevel(isolation.ReadCommitted)
+···  txn.Get(ctx, tk(8))
+···  txn.Prepare(ctx)
+···  return errors.New("rollback")
+···})