From 2a6765a79a6148444c409916babd02a063b024ce Mon Sep 17 00:00:00 2001
From: Andrew Baptist <baptist@cockroachlabs.com>
Date: Wed, 25 Oct 2023 10:57:32 -0400
Subject: [PATCH] workload: add follower read support

The performance of the system can be different if there are reads as of
a system time. This PR adds the ability to issue follower reads.

Epic: none

Release note (cli change): Add --follower-read-percent to the workload kv run.
---
 pkg/workload/kv/kv.go | 62 ++++++++++++++++++++++++++-----------------
 1 file changed, 37 insertions(+), 25 deletions(-)

diff --git a/pkg/workload/kv/kv.go b/pkg/workload/kv/kv.go
index 49058229bea8..889b7d9d5969 100644
--- a/pkg/workload/kv/kv.go
+++ b/pkg/workload/kv/kv.go
@@ -79,6 +79,7 @@ type kv struct {
 	minBlockSizeBytes, maxBlockSizeBytes int
 	cycleLength                          int64
 	readPercent                          int
+	followerReadPercent                  int
 	spanPercent                          int
 	delPercent                           int
 	spanLimit                            int
@@ -142,6 +143,8 @@ var kvMeta = workload.Meta{
 			`Number of keys repeatedly accessed by each writer through upserts.`)
 		g.flags.IntVar(&g.readPercent, `read-percent`, 0,
 			`Percent (0-100) of operations that are reads of existing keys.`)
+		g.flags.IntVar(&g.followerReadPercent, `follower-read-percent`, 0,
+			`Percent (0-100) of read operations that are follower reads.`)
 		g.flags.IntVar(&g.spanPercent, `span-percent`, 0,
 			`Percent (0-100) of operations that are spanning queries of all ranges.`)
 		g.flags.IntVar(&g.delPercent, `del-percent`, 0,
@@ -449,25 +452,27 @@ func (w *kv) Ops(
 
 	// Read statement
 	var buf strings.Builder
+	var folBuf strings.Builder
 	if w.enum {
 		buf.WriteString(`SELECT k, v, e FROM kv WHERE k IN (`)
-		for i := 0; i < w.batchSize; i++ {
-			if i > 0 {
-				buf.WriteString(", ")
-			}
-			fmt.Fprintf(&buf, `$%d`, i+1)
-		}
+		folBuf.WriteString(`SELECT k, v, e FROM kv AS OF SYSTEM TIME follower_read_timestamp() WHERE k IN (`)
 	} else {
 		buf.WriteString(`SELECT k, v FROM kv WHERE k IN (`)
-		for i := 0; i < w.batchSize; i++ {
-			if i > 0 {
-				buf.WriteString(", ")
-			}
-			fmt.Fprintf(&buf, `$%d`, i+1)
+		folBuf.WriteString(`SELECT k, v FROM kv AS OF SYSTEM TIME follower_read_timestamp() WHERE k IN (`)
+	}
+	for i := 0; i < w.batchSize; i++ {
+		if i > 0 {
+			buf.WriteString(", ")
+			folBuf.WriteString(", ")
 		}
+		fmt.Fprintf(&buf, `$%d`, i+1)
+		fmt.Fprintf(&folBuf, `$%d`, i+1)
 	}
 	buf.WriteString(`)`)
+	folBuf.WriteString(`)`)
+
 	readStmtStr := buf.String()
+	followerReadStmtStr := folBuf.String()
 
 	// Write statement
 	buf.Reset()
@@ -531,6 +536,7 @@ func (w *kv) Ops(
 			numEmptyResults: &numEmptyResults,
 		}
 		op.readStmt = op.sr.Define(readStmtStr)
+		op.followerReadStmt = op.sr.Define(followerReadStmtStr)
 		op.writeStmt = op.sr.Define(writeStmtStr)
 		if len(sfuStmtStr) > 0 {
 			op.sfuStmt = op.sr.Define(sfuStmtStr)
@@ -554,19 +560,20 @@ func (w *kv) Ops(
 }
 
 type kvOp struct {
-	config          *kv
-	hists           *histogram.Histograms
-	sr              workload.SQLRunner
-	mcp             *workload.MultiConnPool
-	qosStmt         *workload.StmtHandle
-	readStmt        workload.StmtHandle
-	writeStmt       workload.StmtHandle
-	spanStmt        workload.StmtHandle
-	sfuStmt         workload.StmtHandle
-	delStmt         workload.StmtHandle
-	g               keyGenerator
-	t               keyTransformer
-	numEmptyResults *atomic.Int64
+	config           *kv
+	hists            *histogram.Histograms
+	sr               workload.SQLRunner
+	mcp              *workload.MultiConnPool
+	qosStmt          *workload.StmtHandle
+	readStmt         workload.StmtHandle
+	followerReadStmt workload.StmtHandle
+	writeStmt        workload.StmtHandle
+	spanStmt         workload.StmtHandle
+	sfuStmt          workload.StmtHandle
+	delStmt          workload.StmtHandle
+	g                keyGenerator
+	t                keyTransformer
+	numEmptyResults  *atomic.Int64
 }
 
 func (o *kvOp) run(ctx context.Context) (retErr error) {
@@ -589,7 +596,12 @@ func (o *kvOp) run(ctx context.Context) (retErr error) {
 			args[i] = o.t.getKey(o.g.readKey())
 		}
 		start := timeutil.Now()
-		rows, err := o.readStmt.Query(ctx, args...)
+		readStmt := o.readStmt
+
+		if o.g.rand().Intn(100) < o.config.followerReadPercent {
+			readStmt = o.followerReadStmt
+		}
+		rows, err := readStmt.Query(ctx, args...)
 		if err != nil {
 			return err
 		}