workload: add --scatter flag to kv workload

msbutler · msbutler · commit d16c0114baf1 · 2023-05-10T15:43:41.000-04:00
The user can now run `./workload init kv --scatter ....` which scatters the kv
table across the cluster after the initial data load. This flag is best used
with `--splits`, `--max-block-bytes`, and `--insert-count`.

Epic: none

Release note: none
diff --git a/pkg/workload/kv/kv.go b/pkg/workload/kv/kv.go
@@ -88,6 +88,7 @@ type kv struct {
 	zipfian                              bool
 	sfuDelay                             time.Duration
 	splits                               int
+	scatter                              bool
 	secondaryIndex                       bool
 	shards                               int
 	targetCompressionRatio               float64
@@ -127,6 +128,7 @@ var kvMeta = workload.Meta{
 			`span-limit`:     {RuntimeOnly: true},
 			`del-percent`:    {RuntimeOnly: true},
 			`splits`:         {RuntimeOnly: true},
+			`scatter`:        {RuntimeOnly: true},
 			`timeout`:        {RuntimeOnly: true},
 		}
 		g.flags.IntVar(&g.batchSize, `batch`, 1,
@@ -157,6 +159,8 @@ var kvMeta = workload.Meta{
 				`previous --sequential run and R implies a previous random run.`)
 		g.flags.IntVar(&g.splits, `splits`, 0,
 			`Number of splits to perform before starting normal operations.`)
+		g.flags.BoolVar(&g.scatter, `scatter`, false,
+			`Scatter ranges before starting normal operations.`)
 		g.flags.BoolVar(&g.secondaryIndex, `secondary-index`, false,
 			`Add a secondary index to the schema.`)
 		g.flags.IntVar(&g.shards, `num-shards`, 0,
@@ -189,13 +193,20 @@ func (w *kv) Flags() workload.Flags { return w.flags }
 func (w *kv) Hooks() workload.Hooks {
 	return workload.Hooks{
 		PostLoad: func(_ context.Context, db *gosql.DB) error {
-			if !w.enum {
-				return nil
-			}
-			_, err := db.Exec(`
+			if w.enum {
+				_, err := db.Exec(`
 CREATE TYPE enum_type AS ENUM ('v');
 ALTER TABLE kv ADD COLUMN e enum_type NOT NULL AS ('v') STORED;`)
-			return err
+				if err != nil {
+					return err
+				}
+			}
+			if w.scatter {
+				if _, err := db.Exec(`ALTER TABLE kv SCATTER`); err != nil {
+					return err
+				}
+			}
+			return nil
 		},
 		Validate: w.validateConfig,
 	}
