kvnemesis: add MutateBatchHeaderOperation

This adds a new operation which sets the TargetBytes or
MaxSpanRequestKeys on the header of a batch operation. These header
options substantially change the how requests are processed in dist
sender and can lead to early termination of a batch.

Since they can lead to early termination of a batch, there is a
trade-off here since this operation may result in other operations being
no-ops.

Fixes #152668
Release note: None

Author: stevendanna

pkg/kv/kvnemesis/applier.go

@@ -605,6 +605,9 @@ func applyBatchOp(
 			panic(errors.AssertionFailedf(`Barrier cannot be used in batches`))
 		case *FlushLockTableOperation:
 			panic(errors.AssertionFailedf(`FlushLockOperation cannot be used in batches`))
+		case *MutateBatchHeaderOperation:
+			b.Header.MaxSpanRequestKeys = subO.MaxSpanRequestKeys
+			b.Header.TargetBytes = subO.TargetBytes
 		default:
 			panic(errors.AssertionFailedf(`unknown batch operation type: %T %v`, subO, subO))
 		}
@@ -615,25 +618,34 @@ func applyBatchOp(
 	// to each result.
 	err = nil
 	o.Result.OptionalTimestamp = ts
+	resultIdx := 0
 	for i := range o.Ops {
 		switch subO := o.Ops[i].GetValue().(type) {
 		case *GetOperation:
-			if b.Results[i].Err != nil {
-				subO.Result = resultInit(ctx, b.Results[i].Err)
+			res := b.Results[resultIdx]
+			if res.Err != nil {
+				subO.Result = resultInit(ctx, res.Err)
 			} else {
-				subO.Result.Type = ResultType_Value
-				result := b.Results[i].Rows[0]
-				if result.Value != nil {
-					subO.Result.Value = result.Value.RawBytes
+				if res.ResumeSpan != nil {
+					subO.Result.Type = ResultType_NoError
 				} else {
-					subO.Result.Value = nil
+					subO.Result.Type = ResultType_Value
+					result := res.Rows[0]
+					if result.Value != nil {
+						subO.Result.Value = result.Value.RawBytes
+					} else {
+						subO.Result.Value = nil
+					}
 				}
 			}
+			subO.Result.ResumeSpan = res.ResumeSpan
 		case *PutOperation:
-			err := b.Results[i].Err
+			err := b.Results[resultIdx].Err
 			subO.Result = resultInit(ctx, err)
+			subO.Result.ResumeSpan = b.Results[resultIdx].ResumeSpan
 		case *ScanOperation:
-			kvs, err := b.Results[i].Rows, b.Results[i].Err
+			res := b.Results[resultIdx]
+			kvs, err := res.Rows, res.Err
 			if err != nil {
 				subO.Result = resultInit(ctx, err)
 			} else {
@@ -646,11 +658,14 @@ func applyBatchOp(
 					}
 				}
 			}
+			subO.Result.ResumeSpan = res.ResumeSpan
 		case *DeleteOperation:
-			err := b.Results[i].Err
+			err := b.Results[resultIdx].Err
 			subO.Result = resultInit(ctx, err)
+			subO.Result.ResumeSpan = b.Results[resultIdx].ResumeSpan
 		case *DeleteRangeOperation:
-			keys, err := b.Results[i].Keys, b.Results[i].Err
+			res := b.Results[resultIdx]
+			keys, err := res.Keys, res.Err
 			if err != nil {
 				subO.Result = resultInit(ctx, err)
 			} else {
@@ -660,13 +675,20 @@ func applyBatchOp(
 					subO.Result.Keys[j] = key
 				}
 			}
+			subO.Result.ResumeSpan = res.ResumeSpan
 		case *DeleteRangeUsingTombstoneOperation:
 			subO.Result = resultInit(ctx, err)
+			subO.Result.ResumeSpan = b.Results[resultIdx].ResumeSpan
+		case *MutateBatchHeaderOperation:
+			subO.Result = resultInit(ctx, nil)
 		case *AddSSTableOperation:
 			panic(errors.AssertionFailedf(`AddSSTable cannot be used in batches`))
 		default:
 			panic(errors.AssertionFailedf(`unknown batch operation type: %T %v`, subO, subO))
 		}
+		if o.Ops[i].OperationHasResultInBatch() {
+			resultIdx++
+		}
 	}
 }
 

pkg/kv/kvnemesis/generator.go

@@ -283,6 +283,10 @@ type ClientOperationConfig struct {
 	// FlushLockTable is an operation that moves unreplicated locks in the
 	// in-memory lock table into the
 	FlushLockTable int
+
+	// MutateBatchHeader mutates elements of the batch Header that may influence
+	// batch evaluation. Only relevant for BatchOperations.
+	MutateBatchHeader int
 }
 
 // BatchOperationConfig configures the relative probability of generating a
@@ -433,6 +437,7 @@ func newAllOperationsConfig() GeneratorConfig {
 		AddSSTable:                                         1,
 		Barrier:                                            1,
 		FlushLockTable:                                     1,
+		MutateBatchHeader:                                  1,
 	}
 	batchOpConfig := BatchOperationConfig{
 		Batch: 4,
@@ -498,6 +503,14 @@ func newAllOperationsConfig() GeneratorConfig {
 // operations/make some operations more likely.
 func NewDefaultConfig() GeneratorConfig {
 	config := newAllOperationsConfig()
+
+	// MutateBatchHeader is only valid in batches.
+	config.Ops.DB.MutateBatchHeader = 0
+	config.Ops.ClosureTxn.TxnClientOps.MutateBatchHeader = 0
+
+	// TODO(#153446): Header mutations with EndTransaction are not currently safe.
+	config.Ops.ClosureTxn.CommitBatchOps.MutateBatchHeader = 0
+
 	// DeleteRangeUsingTombstone does not support transactions.
 	config.Ops.ClosureTxn.TxnClientOps.DeleteRangeUsingTombstone = 0
 	config.Ops.ClosureTxn.TxnBatchOps.Ops.DeleteRangeUsingTombstone = 0
@@ -876,6 +889,7 @@ func (g *generator) registerClientOps(allowed *[]opGen, c *ClientOperationConfig
 	addOpGen(allowed, randAddSSTable, c.AddSSTable)
 	addOpGen(allowed, randBarrier, c.Barrier)
 	addOpGen(allowed, randFlushLockTable, c.FlushLockTable)
+	addOpGen(allowed, randBatchMutation, c.MutateBatchHeader)
 }
 
 func (g *generator) registerBatchOps(allowed *[]opGen, c *BatchOperationConfig) {
@@ -1452,6 +1466,19 @@ func randMergeIsSplit(g *generator, rng *rand.Rand) Operation {
 	return merge(key)
 }
 
+func randBatchMutation(g *generator, rng *rand.Rand) Operation {
+	op := &MutateBatchHeaderOperation{}
+	// We currently only support two header option mutations, both of which can
+	// lead to early termination. Half the time we choose a value very likely to
+	// lead to early termination.
+	if rng.Float64() > 0.5 {
+		op.MaxSpanRequestKeys = randItem(rng, []int64{1, 100})
+	} else {
+		op.TargetBytes = randItem(rng, []int64{1, 1 << 20 /* 1MiB */})
+	}
+	return Operation{MutateBatchHeader: op}
+}
+
 func makeRemoveReplicaFn(key string, current []roachpb.ReplicationTarget, voter bool) opGenFunc {
 	return func(g *generator, rng *rand.Rand) Operation {
 		var changeType roachpb.ReplicaChangeType
@@ -1565,6 +1592,13 @@ func makeRandBatch(c *ClientOperationConfig) opGenFunc {
 		numOps := rng.Intn(4)
 		ops := make([]Operation, numOps)
 		var addedForwardScan, addedReverseScan bool
+
+		// TODO(ssd): MutateBatchHeader is disallowed with Puts because of
+		// validation in the txnWriteBuffer that disallows such requests. We could
+		// relax this restriction for many batches if we had information about the
+		// enclosing transaction here.
+		var addedPutOrCPut, addedBatchHeaderMutation bool
+
 		for i := 0; i < numOps; i++ {
 			ops[i] = g.selectOp(rng, allowed)
 			if ops[i].Scan != nil {
@@ -1585,6 +1619,20 @@ func makeRandBatch(c *ClientOperationConfig) opGenFunc {
 					}
 					addedReverseScan = true
 				}
+			} else if ops[i].Put != nil {
+				if addedBatchHeaderMutation {
+					i--
+					continue
+				}
+				addedPutOrCPut = true
+			} else if ops[i].MutateBatchHeader != nil {
+				// In addition to avoiding batch mutations when we have Puts or CPuts,
+				// we also skip adding mutations if one is already added.
+				if addedPutOrCPut || addedBatchHeaderMutation {
+					i--
+					continue
+				}
+				addedBatchHeaderMutation = true
 			}
 		}
 		return batch(ops...)
@@ -1803,6 +1851,10 @@ func keysBetween(keys map[string]struct{}, start, end string) []string {
 	return between
 }
 
+func randItem[T any](rng *rand.Rand, l []T) T {
+	return l[rng.Intn(len(l))]
+}
+
 func randKey(rng *rand.Rand) string {
 	// Avoid the endpoints because having point writes at the
 	// endpoints complicates randRangeSpan.

pkg/kv/kvnemesis/generator_test.go

@@ -266,6 +266,8 @@ func TestRandStep(t *testing.T) {
 				client.Barrier++
 			case *FlushLockTableOperation:
 				client.FlushLockTable++
+			case *MutateBatchHeaderOperation:
+				client.MutateBatchHeader++
 			case *BatchOperation:
 				batch.Batch++
 				countClientOps(&batch.Ops, nil, o.Ops...)
@@ -303,7 +305,8 @@ func TestRandStep(t *testing.T) {
 			*DeleteRangeUsingTombstoneOperation,
 			*AddSSTableOperation,
 			*BarrierOperation,
-			*FlushLockTableOperation:
+			*FlushLockTableOperation,
+			*MutateBatchHeaderOperation:
 			countClientOps(&counts.DB, &counts.Batch, step.Op)
 		case *ClosureTxnOperation:
 			countClientOps(&counts.ClosureTxn.TxnClientOps, &counts.ClosureTxn.TxnBatchOps, o.Ops...)

pkg/kv/kvnemesis/operations.go

@@ -62,6 +62,8 @@ func (op Operation) Result() *Result {
 		return &o.Result
 	case *SavepointRollbackOperation:
 		return &o.Result
+	case *MutateBatchHeaderOperation:
+		return &o.Result
 	default:
 		panic(errors.AssertionFailedf(`unknown operation: %T %v`, o, o))
 	}
@@ -115,6 +117,14 @@ func formatOps(w *strings.Builder, fctx formatCtx, ops []Operation) {
 	}
 }
 
+func (op Operation) OperationHasResultInBatch() bool {
+	if op.MutateBatchHeader != nil {
+		return false
+	} else {
+		return true
+	}
+}
+
 func (op Operation) String() string {
 	fctx := formatCtx{receiver: `x`, indent: ``}
 	var buf strings.Builder
@@ -222,6 +232,8 @@ func (op Operation) format(w *strings.Builder, fctx formatCtx) {
 		o.format(w, fctx)
 	case *SavepointRollbackOperation:
 		o.format(w, fctx)
+	case *MutateBatchHeaderOperation:
+		o.format(w, fctx)
 	default:
 		fmt.Fprintf(w, "%v", op.GetValue())
 	}
@@ -453,6 +465,15 @@ func (op SavepointRollbackOperation) format(w *strings.Builder, fctx formatCtx)
 	op.Result.format(w)
 }
 
+func (op MutateBatchHeaderOperation) format(w *strings.Builder, fctx formatCtx) {
+	if op.TargetBytes > 0 {
+		fmt.Fprintf(w, `%s.Header.TargetBytes = %d // MutateBatchHeaderOperation`, fctx.receiver, op.TargetBytes)
+	}
+	if op.MaxSpanRequestKeys > 0 {
+		fmt.Fprintf(w, `%s.Header.MaxSpanRequestKeys = %d // MutateBatchHeaderOperation`, fctx.receiver, op.MaxSpanRequestKeys)
+	}
+}
+
 func (r Result) format(w *strings.Builder) {
 	if r.Type == ResultType_Unknown {
 		return

pkg/kv/kvnemesis/operations.proto

@@ -165,6 +165,13 @@ message SavepointRollbackOperation {
   Result result = 2 [(gogoproto.nullable) = false];
 }
 
+message MutateBatchHeaderOperation {
+  int64 max_span_request_keys = 1;
+  int64 target_bytes = 2;
+
+  Result result = 3 [(gogoproto.nullable) = false];
+}
+
 message Operation {
   option (gogoproto.goproto_stringer) = false;
   option (gogoproto.onlyone) = true;
@@ -195,6 +202,7 @@ message Operation {
   SavepointRollbackOperation savepoint_rollback = 22;
   BarrierOperation barrier = 23;
   FlushLockTableOperation flush_lock_table = 24;
+  MutateBatchHeaderOperation mutate_batch_header = 25;
 }
 
 enum ResultType {
@@ -223,6 +231,7 @@ message Result {
   // Only set if Type is ResultType_Values. The RawBytes of a roachpb.Value.
   repeated KeyValue values = 5 [(gogoproto.nullable) = false];
   util.hlc.Timestamp optional_timestamp = 6 [(gogoproto.nullable) = false];
+  roachpb.Span resume_span = 7;
 }
 
 message Step {

pkg/kv/kvnemesis/validator.go

@@ -433,6 +433,9 @@ func (v *validator) processOp(op Operation) {
 		if _, isErr := v.checkError(op, t.Result); isErr {
 			break
 		}
+		if t.Result.ResumeSpan != nil {
+			break
+		}
 		read := &observedRead{
 			Key:        t.Key,
 			SkipLocked: t.SkipLocked,
@@ -464,6 +467,9 @@ func (v *validator) processOp(op Operation) {
 		if v.checkNonAmbError(op, t.Result) {
 			break
 		}
+		if t.Result.ResumeSpan != nil {
+			break
+		}
 		// Accumulate all the writes for this transaction.
 		write := &observedWrite{
 			Key:   t.Key,
@@ -482,6 +488,9 @@ func (v *validator) processOp(op Operation) {
 		if v.checkNonAmbError(op, t.Result) {
 			break
 		}
+		if t.Result.ResumeSpan != nil {
+			break
+		}
 		sv, _ := v.tryConsumeWrite(t.Key, t.Seq)
 		write := &observedWrite{
 			Key:       t.Key,
@@ -532,10 +541,14 @@ func (v *validator) processOp(op Operation) {
 		// not prevent new keys from being inserted in the deletion span between the
 		// transaction's read and write timestamps.
 		if v.observationFilter != observeLocking {
+			endKey := t.EndKey
+			if t.Result.ResumeSpan != nil {
+				endKey = t.Result.ResumeSpan.Key
+			}
 			v.curObservations = append(v.curObservations, &observedScan{
 				Span: roachpb.Span{
 					Key:    t.Key,
-					EndKey: t.EndKey,
+					EndKey: endKey,
 				},
 				IsDeleteRange: true, // just for printing
 				KVs:           nil,
@@ -747,13 +760,24 @@ func (v *validator) processOp(op Operation) {
 		if _, isErr := v.checkError(op, t.Result); isErr {
 			break
 		}
+		readSpan := roachpb.Span{Key: t.Key, EndKey: t.EndKey}
+		// If the ResumeSpan equals the original request span, the scan wasn't
+		// processed at all.
+		if t.Result.ResumeSpan != nil && t.Result.ResumeSpan.Equal(readSpan) {
+			break
+		}
+
 		switch v.observationFilter {
 		case observeAll:
+			if t.Result.ResumeSpan != nil {
+				if op.Scan.Reverse {
+					readSpan.Key = t.Result.ResumeSpan.EndKey
+				} else {
+					readSpan.EndKey = t.Result.ResumeSpan.Key
+				}
+			}
 			scan := &observedScan{
-				Span: roachpb.Span{
-					Key:    t.Key,
-					EndKey: t.EndKey,
-				},
+				Span:       readSpan,
 				Reverse:    t.Reverse,
 				SkipLocked: t.SkipLocked,
 				KVs:        make([]roachpb.KeyValue, len(t.Result.Values)),
@@ -926,6 +950,9 @@ func (v *validator) processOp(op Operation) {
 		// Don't fail on all errors because savepoints can be labeled with
 		// errOmitted if a previous op in the txn failed.
 		v.checkError(op, t.Result)
+	case *MutateBatchHeaderOperation:
+		execTimestampStrictlyOptional = true
+		v.checkError(op, t.Result)
 	default:
 		panic(errors.AssertionFailedf(`unknown operation type: %T %v`, t, t))
 	}