From 258256402d4091eb3851b7307a86ebb959425411 Mon Sep 17 00:00:00 2001
From: Jernej Kos <jernej@kos.mx>
Date: Wed, 25 Mar 2020 11:00:17 +0100
Subject: [PATCH] txsource/workload/queries: Add num last kept versions
 argument

---
 .changelog/2674.internal.md                   |  1 +
 go/oasis-node/cmd/debug/txsource/txsource.go  |  2 +-
 .../cmd/debug/txsource/workload/queries.go    | 27 ++++++++++++++++---
 .../cmd/debug/txsource/workload/workload.go   |  9 +++++++
 4 files changed, 34 insertions(+), 5 deletions(-)
 create mode 100644 .changelog/2674.internal.md

diff --git a/.changelog/2674.internal.md b/.changelog/2674.internal.md
new file mode 100644
index 00000000000..8ea9feca89e
--- /dev/null
+++ b/.changelog/2674.internal.md
@@ -0,0 +1 @@
+txsource/workload/queries: Add num last kept versions argument
diff --git a/go/oasis-node/cmd/debug/txsource/txsource.go b/go/oasis-node/cmd/debug/txsource/txsource.go
index a34ba724f96..7fe443e9ad4 100644
--- a/go/oasis-node/cmd/debug/txsource/txsource.go
+++ b/go/oasis-node/cmd/debug/txsource/txsource.go
@@ -138,7 +138,7 @@ func init() {
 	_ = viper.BindPFlags(fs)
 	txsourceCmd.Flags().AddFlagSet(fs)
 
-	txsourceCmd.Flags().AddFlagSet(workload.RuntimeFlags)
+	txsourceCmd.Flags().AddFlagSet(workload.Flags)
 	txsourceCmd.Flags().AddFlagSet(cmdGrpc.ClientFlags)
 	txsourceCmd.Flags().AddFlagSet(cmdFlags.DebugTestEntityFlags)
 	txsourceCmd.Flags().AddFlagSet(cmdFlags.GenesisFileFlags)
diff --git a/go/oasis-node/cmd/debug/txsource/workload/queries.go b/go/oasis-node/cmd/debug/txsource/workload/queries.go
index e09d4ba44ea..b90780816a2 100644
--- a/go/oasis-node/cmd/debug/txsource/workload/queries.go
+++ b/go/oasis-node/cmd/debug/txsource/workload/queries.go
@@ -6,6 +6,7 @@ import (
 	"math/rand"
 	"time"
 
+	flag "github.com/spf13/pflag"
 	"github.com/spf13/viper"
 	"google.golang.org/grpc"
 
@@ -27,10 +28,17 @@ const (
 	// NameQueries is the name of the queries workload.
 	NameQueries = "queries"
 
+	// CfgNumKeptVersions is the number of last versions that the nodes are keeping (e.g., due to
+	// a configured pruning policy). Versions older than that will not be queried.
+	CfgNumKeptVersions = "queries.num_kept_versions"
+
 	// Ratio of queries that should query height 1.
 	queriesEarliestHeightRatio = 0.1
 )
 
+// QueriesFlags are the queries workload flags.
+var QueriesFlags = flag.NewFlagSet("", flag.ContinueOnError)
+
 type queries struct {
 	logger *logging.Logger
 
@@ -408,16 +416,22 @@ func (q *queries) Run(gracefulExit context.Context, rng *rand.Rand, conn *grpc.C
 			return fmt.Errorf("consensus.GetBlock error: %w", err)
 		}
 
-		// Select height at which queries should be done. Earliest (height=1)
+		// Determine the earliest height that we can query.
+		earliestHeight := int64(1)
+		if numKept := viper.GetInt64(CfgNumKeptVersions); numKept < block.Height {
+			earliestHeight = block.Height - numKept
+		}
+
+		// Select height at which queries should be done. Earliest
 		// is special cased with increased probability.
 		var height int64
 		p := rng.Float32()
 		switch {
 		case p < queriesEarliestHeightRatio:
-			height = 1
+			height = earliestHeight
 		default:
-			// [1, block.Height]
-			height = rng.Int63n(block.Height) + 1
+			// [earliestHeight, block.Height]
+			height = rng.Int63n(block.Height-earliestHeight+1) + earliestHeight
 		}
 
 		q.logger.Debug("Doing queries",
@@ -454,3 +468,8 @@ func (q *queries) Run(gracefulExit context.Context, rng *rand.Rand, conn *grpc.C
 		}
 	}
 }
+
+func init() {
+	QueriesFlags.Int64(CfgNumKeptVersions, 0, "Number of last versions kept by nodes")
+	_ = viper.BindPFlags(QueriesFlags)
+}
diff --git a/go/oasis-node/cmd/debug/txsource/workload/workload.go b/go/oasis-node/cmd/debug/txsource/workload/workload.go
index 7f69af5dd4c..87ecb2d2ce9 100644
--- a/go/oasis-node/cmd/debug/txsource/workload/workload.go
+++ b/go/oasis-node/cmd/debug/txsource/workload/workload.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/cenkalti/backoff/v4"
+	flag "github.com/spf13/pflag"
 	"google.golang.org/grpc"
 
 	"github.com/oasislabs/oasis-core/go/common/crypto/signature"
@@ -183,3 +184,11 @@ var ByName = map[string]Workload{
 	NameRuntime:      &runtime{},
 	NameTransfer:     &transfer{},
 }
+
+// Flags has the workload flags.
+var Flags = flag.NewFlagSet("", flag.ContinueOnError)
+
+func init() {
+	Flags.AddFlagSet(QueriesFlags)
+	Flags.AddFlagSet(RuntimeFlags)
+}