kvnemesis: add follower reads

This commit adds follower reads support to kvnemesis. The target closed
timestamp is set to 1s, and select read-only non-locking operations run
in batches with a timestamp of now - 3s. Multi-request batches and
transactions are also eligible for follower reads as long as they
consist only of eligible operations.

Fixes: #59061

Release note: None

Author: miraradeva

pkg/kv/db.go

@@ -226,6 +226,15 @@ func (s *CrossRangeTxnWrapperSender) Send(
 		return br, pErr
 	}
 
+	// Before retrying the batch in a transaction, strip the header's timestamp.
+	// It may have been set to try a follower read, but it's not allowed in a txn
+	// (see comment near Header.Timestamp). Currently, this non-transactional API
+	// is used for follower reads only by KVNemesis.
+	//
+	// We can end up here if the batch contains transactional requests and spans
+	// multiple ranges.
+	ba.Header.Timestamp = hlc.Timestamp{}
+
 	err := s.db.Txn(ctx, func(ctx context.Context, txn *Txn) error {
 		txn.SetDebugName("auto-wrap")
 		b := txn.NewBatch()

pkg/kv/kvnemesis/applier.go

@@ -122,6 +122,16 @@ func exceptContextCanceled(err error) bool {
 		strings.Contains(err.Error(), "query execution canceled")
 }
 
+const followerReadsOffset = 2 * time.Second
+
+func makeFollowerReadTimestamp(ts hlc.Timestamp) hlc.Timestamp {
+	return ts.Add(-followerReadsOffset.Nanoseconds(), 0)
+}
+func configureBatchForFollowerReads(b *kv.Batch, ts hlc.Timestamp) {
+	b.Header.Timestamp = makeFollowerReadTimestamp(ts)
+	b.Header.RoutingPolicy = kvpb.RoutingPolicy_NEAREST
+}
+
 func (a *Applier) applyOp(ctx context.Context, db *kv.DB, op *Operation) {
 	switch o := op.GetValue().(type) {
 	case *GetOperation,
@@ -188,6 +198,15 @@ func (a *Applier) applyOp(ctx context.Context, db *kv.DB, op *Operation) {
 		})
 		var savedTxn *kv.Txn
 		txnErr := db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+			// Attempt to set a follower reads timestamp only on the first iteration.
+			// Setting it again and again on retries can lead to starvation.
+			if o.FollowerReadEligible && savedTxn == nil {
+				followerReadTs := makeFollowerReadTimestamp(db.Clock().Now())
+				err := txn.SetFixedTimestamp(ctx, followerReadTs)
+				if err != nil {
+					panic(err)
+				}
+			}
 			if err := txn.SetIsoLevel(o.IsoLevel); err != nil {
 				panic(err)
 			}
@@ -251,7 +270,7 @@ func (a *Applier) applyOp(ctx context.Context, db *kv.DB, op *Operation) {
 			}
 			if o.CommitInBatch != nil {
 				b := txn.NewBatch()
-				applyBatchOp(ctx, b, txn.CommitInBatch, o.CommitInBatch)
+				applyBatchOp(ctx, b, txn.CommitInBatch, o.CommitInBatch, nil /* clock */)
 				// The KV api disallows use of a txn after an operation on it errors.
 				if r := o.CommitInBatch.Result; r.Type == ResultType_Error {
 					return errors.DecodeError(ctx, *r.Err)
@@ -349,6 +368,13 @@ func applyClientOp(
 			} else {
 				b.Get(o.Key)
 			}
+			if !inTxn && o.FollowerReadEligible {
+				kvDB, ok := db.(*kv.DB)
+				if !ok {
+					panic(errors.AssertionFailedf("unexpected transactional interface"))
+				}
+				configureBatchForFollowerReads(b, kvDB.Clock().Now())
+			}
 		})
 		o.Result = resultInit(ctx, err)
 		if err != nil {
@@ -422,6 +448,13 @@ func applyClientOp(
 					b.Scan(o.Key, o.EndKey)
 				}
 			}
+			if !inTxn && o.FollowerReadEligible {
+				kvDB, ok := db.(*kv.DB)
+				if !ok {
+					panic(errors.AssertionFailedf("unexpected transactional interface"))
+				}
+				configureBatchForFollowerReads(b, kvDB.Clock().Now())
+			}
 		})
 		o.Result = resultInit(ctx, err)
 		if err != nil {
@@ -531,7 +564,15 @@ func applyClientOp(
 		o.Result = resultInit(ctx, err)
 	case *BatchOperation:
 		b := &kv.Batch{}
-		applyBatchOp(ctx, b, db.Run, o)
+		if inTxn {
+			applyBatchOp(ctx, b, db.Run, o, nil /* clock */)
+		} else {
+			kvDB, ok := db.(*kv.DB)
+			if !ok {
+				panic(errors.AssertionFailedf("unexpected transactional interface"))
+			}
+			applyBatchOp(ctx, b, db.Run, o, kvDB.Clock())
+		}
 	case *SavepointCreateOperation:
 		txn, ok := db.(*kv.Txn) // savepoints are only allowed with transactions
 		if !ok {
@@ -589,7 +630,11 @@ func setLastReqSeq(b *kv.Batch, seq kvnemesisutil.Seq) {
 }
 
 func applyBatchOp(
-	ctx context.Context, b *kv.Batch, run func(context.Context, *kv.Batch) error, o *BatchOperation,
+	ctx context.Context,
+	b *kv.Batch,
+	run func(context.Context, *kv.Batch) error,
+	o *BatchOperation,
+	clock *hlc.Clock,
 ) {
 	for i := range o.Ops {
 		switch subO := o.Ops[i].GetValue().(type) {
@@ -674,6 +719,9 @@ func applyBatchOp(
 			panic(errors.AssertionFailedf(`unknown batch operation type: %T %v`, subO, subO))
 		}
 	}
+	if clock != nil && o.FollowerReadEligible {
+		configureBatchForFollowerReads(b, clock.Now())
+	}
 	ts, err := batchRun(ctx, run, b)
 	o.Result = resultInit(ctx, err)
 	// NB: we intentionally fall through; the batch propagates the error

pkg/kv/kvnemesis/generator.go

@@ -111,6 +111,9 @@ type ClosureTxnConfig struct {
 type ClientOperationConfig struct {
 	// GetMissing is an operation that Gets a key that definitely doesn't exist.
 	GetMissing int
+	// GetMissingFollowerRead is an operation that Gets a key that definitely
+	// doesn't exist, and is marked eligible to be served as a follower read.
+	GetMissingFollowerRead int
 	// GetMissingForUpdate is an operation that Gets a key that definitely doesn't
 	// exist using a locking read with strength lock.Exclusive.
 	GetMissingForUpdate int
@@ -146,6 +149,9 @@ type ClientOperationConfig struct {
 	GetMissingForShareSkipLockedGuaranteedDurability int
 	// GetExisting is an operation that Gets a key that likely exists.
 	GetExisting int
+	// GetExistingFollowerRead is an operation that Gets a key that likely exists,
+	// and is marked eligible to be served as a follower read.
+	GetExistingFollowerRead int
 	// GetExistingForUpdate is an operation that Gets a key that likely exists
 	// using a locking read with strength lock.Exclusive.
 	GetExistingForUpdate int
@@ -203,6 +209,9 @@ type ClientOperationConfig struct {
 	CPutAllowIfDoesNotExist int
 	// Scan is an operation that Scans a key range that may contain values.
 	Scan int
+	// ScanFollowerRead is an operation that Scans a key range that may contain
+	// values, and is marked eligible to be served as a follower read.
+	ScanFollowerRead int
 	// ScanForUpdate is an operation that Scans a key range that may contain
 	// values using a per-key locking scan with strength lock.Exclusive.
 	ScanForUpdate int
@@ -240,6 +249,10 @@ type ClientOperationConfig struct {
 	// ReverseScan is an operation that Scans a key range that may contain
 	// values in reverse key order.
 	ReverseScan int
+	// ReverseScanFollowerRead is an operation that Scans a key range that may
+	// contain values in reverse key order, and is marked eligible to be served as
+	// a follower read.
+	ReverseScanFollowerRead int
 	// ReverseScanForUpdate is an operation that Scans a key range that may
 	// contain values using a per-key locking scan with strength lock.Exclusive in
 	// reverse key order.
@@ -419,16 +432,18 @@ type FaultConfig struct {
 // yet pass (for example, if the new operation finds a kv bug or edge case).
 func newAllOperationsConfig() GeneratorConfig {
 	clientOpConfig := ClientOperationConfig{
-		GetMissing:                                        1,
-		GetMissingForUpdate:                               1,
-		GetMissingForUpdateGuaranteedDurability:           1,
-		GetMissingForUpdateSkipLocked:                     1,
-		GetMissingForUpdateSkipLockedGuaranteedDurability: 1,
-		GetMissingForShare:                                1,
-		GetMissingForShareGuaranteedDurability:            1,
-		GetMissingForShareSkipLocked:                      1,
-		GetMissingForShareSkipLockedGuaranteedDurability:  1,
+		GetMissing:                                         1,
+		GetMissingFollowerRead:                             1,
+		GetMissingForUpdate:                                1,
+		GetMissingForUpdateGuaranteedDurability:            1,
+		GetMissingForUpdateSkipLocked:                      1,
+		GetMissingForUpdateSkipLockedGuaranteedDurability:  1,
+		GetMissingForShare:                                 1,
+		GetMissingForShareGuaranteedDurability:             1,
+		GetMissingForShareSkipLocked:                       1,
+		GetMissingForShareSkipLockedGuaranteedDurability:   1,
 		GetExisting:                                        1,
+		GetExistingFollowerRead:                            1,
 		GetExistingForUpdate:                               1,
 		GetExistingForUpdateGuaranteedDurability:           1,
 		GetExistingForShare:                                1,
@@ -447,6 +462,7 @@ func newAllOperationsConfig() GeneratorConfig {
 		CPutNoMatch:                         1,
 		CPutAllowIfDoesNotExist:             1,
 		Scan:                                1,
+		ScanFollowerRead:                    1,
 		ScanForUpdate:                       1,
 		ScanForUpdateGuaranteedDurability:   1,
 		ScanForShare:                        1,
@@ -457,6 +473,7 @@ func newAllOperationsConfig() GeneratorConfig {
 		ScanForShareSkipLocked:                      1,
 		ScanForShareSkipLockedGuaranteedDurability:  1,
 		ReverseScan:                                        1,
+		ReverseScanFollowerRead:                            1,
 		ReverseScanForUpdate:                               1,
 		ReverseScanForUpdateGuaranteedDurability:           1,
 		ReverseScanForShare:                                1,
@@ -953,6 +970,7 @@ func (g *generator) selectOp(rng *rand.Rand, contextuallyValid []opGen) Operatio
 
 func (g *generator) registerClientOps(allowed *[]opGen, c *ClientOperationConfig) {
 	addOpGen(allowed, randGetMissing, c.GetMissing)
+	addOpGen(allowed, randGetMissingFollowerRead, c.GetMissingFollowerRead)
 	addOpGen(allowed, randGetMissingForUpdate, c.GetMissingForUpdate)
 	addOpGen(
 		allowed, randGetMissingForUpdateGuaranteedDurability, c.GetMissingForUpdateGuaranteedDurability,
@@ -984,6 +1002,7 @@ func (g *generator) registerClientOps(allowed *[]opGen, c *ClientOperationConfig
 
 	if len(g.keys) > 0 {
 		addOpGen(allowed, randGetExisting, c.GetExisting)
+		addOpGen(allowed, randGetExistingFollowerRead, c.GetExistingFollowerRead)
 		addOpGen(allowed, randGetExistingForUpdate, c.GetExistingForUpdate)
 		addOpGen(
 			allowed,
@@ -1016,6 +1035,7 @@ func (g *generator) registerClientOps(allowed *[]opGen, c *ClientOperationConfig
 		addOpGen(allowed, randDelMustAcquireExclusiveLockExisting, c.DeleteMustAcquireExclusiveLockExisting)
 	}
 	addOpGen(allowed, randScan, c.Scan)
+	addOpGen(allowed, randScanFollowerRead, c.ScanFollowerRead)
 	addOpGen(allowed, randScanForUpdate, c.ScanForUpdate)
 	addOpGen(allowed, randScanForUpdateGuaranteedDurability, c.ScanForUpdateGuaranteedDurability)
 	addOpGen(allowed, randScanForShare, c.ScanForShare)
@@ -1034,6 +1054,7 @@ func (g *generator) registerClientOps(allowed *[]opGen, c *ClientOperationConfig
 		c.ScanForShareSkipLockedGuaranteedDurability,
 	)
 	addOpGen(allowed, randReverseScan, c.ReverseScan)
+	addOpGen(allowed, randReverseScanFollowerRead, c.ReverseScanFollowerRead)
 	addOpGen(allowed, randReverseScanForUpdate, c.ReverseScanForUpdate)
 	addOpGen(
 		allowed,
@@ -1073,6 +1094,12 @@ func randGetMissing(_ *generator, rng *rand.Rand) Operation {
 	return get(randKey(rng))
 }
 
+func randGetMissingFollowerRead(_ *generator, rng *rand.Rand) Operation {
+	op := get(randKey(rng))
+	op.Get.FollowerReadEligible = true
+	return op
+}
+
 func randGetMissingForUpdate(g *generator, rng *rand.Rand) Operation {
 	op := randGetMissing(g, rng)
 	op.Get.ForUpdate = true
@@ -1136,6 +1163,13 @@ func randGetExisting(g *generator, rng *rand.Rand) Operation {
 	return get(key)
 }
 
+func randGetExistingFollowerRead(g *generator, rng *rand.Rand) Operation {
+	key := randSliceKey(rng, maps.Keys(g.keys))
+	op := get(key)
+	op.Get.FollowerReadEligible = true
+	return op
+}
+
 func randGetExistingForUpdate(g *generator, rng *rand.Rand) Operation {
 	op := randGetExisting(g, rng)
 	op.Get.ForUpdate = true
@@ -1436,6 +1470,13 @@ func randScan(g *generator, rng *rand.Rand) Operation {
 	return scan(key, endKey)
 }
 
+func randScanFollowerRead(g *generator, rng *rand.Rand) Operation {
+	key, endKey := randSpan(rng)
+	op := scan(key, endKey)
+	op.Scan.FollowerReadEligible = true
+	return op
+}
+
 func randScanForUpdate(g *generator, rng *rand.Rand) Operation {
 	op := randScan(g, rng)
 	op.Scan.ForUpdate = true
@@ -1500,6 +1541,13 @@ func randReverseScan(g *generator, rng *rand.Rand) Operation {
 	return op
 }
 
+func randReverseScanFollowerRead(g *generator, rng *rand.Rand) Operation {
+	op := randScan(g, rng)
+	op.Scan.Reverse = true
+	op.Scan.FollowerReadEligible = true
+	return op
+}
+
 func randReverseScanForUpdate(g *generator, rng *rand.Rand) Operation {
 	op := randReverseScan(g, rng)
 	op.Scan.ForUpdate = true
@@ -1833,12 +1881,23 @@ func restartRandNode(g *generator, rng *rand.Rand) Operation {
 	return restartNode(randNode)
 }
 
+func isFollowerReadEligibleOp(op Operation) bool {
+	if op.Get != nil && op.Get.FollowerReadEligible {
+		return true
+	}
+	if op.Scan != nil && op.Scan.FollowerReadEligible {
+		return true
+	}
+	return false
+}
+
 func makeRandBatch(c *ClientOperationConfig) opGenFunc {
 	return func(g *generator, rng *rand.Rand) Operation {
 		var allowed []opGen
 		g.registerClientOps(&allowed, c)
 		numOps := rng.Intn(4)
 		ops := make([]Operation, numOps)
+		followerReadEligible := true
 		var addedForwardScan, addedReverseScan bool
 
 		// TODO(ssd): MutateBatchHeader is disallowed with Puts because of
@@ -1849,6 +1908,8 @@ func makeRandBatch(c *ClientOperationConfig) opGenFunc {
 
 		for i := 0; i < numOps; i++ {
 			ops[i] = g.selectOp(rng, allowed)
+			// The batch is eligible for follower reads only if all its ops are.
+			followerReadEligible = followerReadEligible && isFollowerReadEligibleOp(ops[i])
 			if ops[i].Scan != nil {
 				if !ops[i].Scan.Reverse {
 					if addedReverseScan {
@@ -1883,7 +1944,9 @@ func makeRandBatch(c *ClientOperationConfig) opGenFunc {
 				addedBatchHeaderMutation = true
 			}
 		}
-		return batch(ops...)
+		op := batch(ops...)
+		op.Batch.FollowerReadEligible = followerReadEligible
+		return op
 	}
 }
 
@@ -1933,6 +1996,7 @@ func makeClosureTxn(
 		// Stack of savepoint indexes/ids.
 		// The last element of the slice is the top of the stack.
 		var spIDs []int
+		followerReadEligible := true
 		for i := range ops {
 			// In each iteration, we start with the allowed non-savepoint ops,
 			// and we add the valid savepoint ops in registerSavepointOps.
@@ -1943,6 +2007,8 @@ func makeClosureTxn(
 			// allowed savepoint ops changes. See registerSavepointOps.
 			g.registerSavepointOps(&allowedIncludingSavepointOps, savepointOps, spIDs, i)
 			ops[i] = g.selectOp(rng, allowedIncludingSavepointOps)
+			// The transaction is eligible for follower reads only if all its ops are.
+			followerReadEligible = followerReadEligible && isFollowerReadEligibleOp(ops[i])
 			// Now that a random op is selected, we may need to update the stack of
 			// existing savepoints. See maybeUpdateSavepoints.
 			maybeUpdateSavepoints(&spIDs, ops[i])
@@ -1958,6 +2024,7 @@ func makeClosureTxn(
 			}
 			op.ClosureTxn.CommitInBatch = makeRandBatch(commitInBatch)(g, rng).Batch
 		}
+		op.ClosureTxn.FollowerReadEligible = followerReadEligible
 		return op
 	}
 }