Merge #33242

33242: workload/tpcc: split, partition, and scatter during initialization r=nvanbenschoten a=nvanbenschoten

Fixes #32135.

This PR pushes splitting, partitioning, and scattering logic into `tpcc`'s `PostLoad` hook, removing it from its `Ops` method.

Co-authored-by: Nathan VanBenschoten <nvanbenschoten@gmail.com>

Author: craig[bot]

pkg/ccl/workloadccl/fixture.go

@@ -473,8 +473,9 @@ func ImportFixture(
 	if err := g.Wait(); err != nil {
 		return 0, err
 	}
-	// TODO(dan): This needs to do splits and run PostLoad hooks. Unify this,
-	// RestoreFixture, and workload.Setup.
+	if err := runPostLoadSteps(ctx, sqlDB, gen); err != nil {
+		return 0, err
+	}
 	return atomic.LoadInt64(&bytesAtomic), nil
 }
 
@@ -607,20 +608,27 @@ func RestoreFixture(
 	if err := g.Wait(); err != nil {
 		return 0, err
 	}
-	const splitConcurrency = 384 // TODO(dan): Don't hardcode this.
-	for _, table := range fixture.Generator.Tables() {
-		if err := workload.Split(ctx, sqlDB, table, splitConcurrency); err != nil {
-			return 0, errors.Wrapf(err, `splitting %s`, table.Name)
-		}
+	if err := runPostLoadSteps(ctx, sqlDB, fixture.Generator); err != nil {
+		return 0, err
 	}
-	if h, ok := fixture.Generator.(workload.Hookser); ok {
+	return atomic.LoadInt64(&bytesAtomic), nil
+}
+
+func runPostLoadSteps(ctx context.Context, sqlDB *gosql.DB, gen workload.Generator) error {
+	if h, ok := gen.(workload.Hookser); ok {
 		if hooks := h.Hooks(); hooks.PostLoad != nil {
 			if err := hooks.PostLoad(sqlDB); err != nil {
-				return 0, errors.Wrap(err, `PostLoad hook`)
+				return errors.Wrap(err, `PostLoad hook`)
 			}
 		}
 	}
-	return atomic.LoadInt64(&bytesAtomic), nil
+	const splitConcurrency = 384 // TODO(dan): Don't hardcode this.
+	for _, table := range gen.Tables() {
+		if err := workload.Split(ctx, sqlDB, table, splitConcurrency); err != nil {
+			return errors.Wrapf(err, `splitting %s`, table.Name)
+		}
+	}
+	return nil
 }
 
 // ListFixtures returns the object paths to all fixtures stored in a FixtureConfig.

pkg/cmd/roachtest/tpcc.go

@@ -63,7 +63,7 @@ type tpccOptions struct {
 // tpccFixturesCmd generates the command string to load tpcc data for the
 // specified warehouse count into a cluster using either `fixtures import`
 // or `fixtures load` depending on the cloud.
-func tpccFixturesCmd(t *test, cloud string, warehouses int, checks bool) string {
+func tpccFixturesCmd(t *test, cloud string, warehouses int, extraArgs string) string {
 	var action string
 	switch cloud {
 	case "gce":
@@ -87,8 +87,8 @@ func tpccFixturesCmd(t *test, cloud string, warehouses int, checks bool) string
 	default:
 		t.Fatalf("unknown cloud: %q", cloud)
 	}
-	return fmt.Sprintf("./workload fixtures %s tpcc --checks=%v --warehouses=%d {pgurl:1}",
-		action, checks, warehouses)
+	return fmt.Sprintf("./workload fixtures %s tpcc --warehouses=%d %s {pgurl:1}",
+		action, warehouses, extraArgs)
 }
 
 func runTPCC(ctx context.Context, t *test, c *cluster, opts tpccOptions) {
@@ -145,7 +145,7 @@ func runTPCC(ctx context.Context, t *test, c *cluster, opts tpccOptions) {
 				t.Status("loading dataset")
 				c.Start(ctx, t, crdbNodes)
 
-				c.Run(ctx, workloadNode, tpccFixturesCmd(t, cloud, opts.Warehouses, true /* checks */))
+				c.Run(ctx, workloadNode, tpccFixturesCmd(t, cloud, opts.Warehouses, ""))
 				c.Stop(ctx, crdbNodes)
 
 				c.Run(ctx, crdbNodes, "test -e /sbin/zfs && sudo zfs snapshot data1@pristine")
@@ -155,7 +155,7 @@ func runTPCC(ctx context.Context, t *test, c *cluster, opts tpccOptions) {
 			c.Start(ctx, t, crdbNodes)
 		} else {
 			c.Start(ctx, t, crdbNodes)
-			c.Run(ctx, workloadNode, tpccFixturesCmd(t, cloud, opts.Warehouses, true /* checks */))
+			c.Run(ctx, workloadNode, tpccFixturesCmd(t, cloud, opts.Warehouses, ""))
 		}
 	}()
 	t.Status("waiting")
@@ -583,31 +583,33 @@ func loadTPCCBench(
 		t.Fatal(err)
 	}
 
-	// Load the corresponding fixture.
-	t.l.Printf("restoring tpcc fixture\n")
-	cmd := tpccFixturesCmd(t, cloud, b.LoadWarehouses, false /* checks */)
-	if err := c.RunE(ctx, loadNode, cmd); err != nil {
-		return err
-	}
-
-	partArgs := ""
-	rebalanceWait := time.Duration(b.LoadWarehouses/100) * time.Minute
+	var loadArgs string
+	var rebalanceWait time.Duration
 	switch b.LoadConfig {
 	case singleLoadgen:
-		t.l.Printf("splitting and scattering\n")
+		loadArgs = `--split --scatter --checks=false`
+		rebalanceWait = time.Duration(b.LoadWarehouses/100) * time.Minute
 	case singlePartitionedLoadgen:
-		t.l.Printf("splitting, scattering, and partitioning\n")
-		partArgs = fmt.Sprintf(`--partitions=%d`, b.partitions())
+		loadArgs = fmt.Sprintf(`--split --scatter --checks=false --partitions=%d`, b.partitions())
 		rebalanceWait = time.Duration(b.LoadWarehouses/50) * time.Minute
 	case multiLoadgen:
-		t.l.Printf("splitting, scattering, and partitioning\n")
-		partArgs = fmt.Sprintf(`--partitions=%d --zones="%s" --partition-affinity=0`,
+		loadArgs = fmt.Sprintf(`--split --scatter --checks=false --partitions=%d --zones="%s"`,
 			b.partitions(), strings.Join(b.Distribution.zones(), ","))
 		rebalanceWait = time.Duration(b.LoadWarehouses/20) * time.Minute
 	default:
 		panic("unexpected")
 	}
 
+	// Load the corresponding fixture.
+	t.l.Printf("restoring tpcc fixture\n")
+	cmd := tpccFixturesCmd(t, cloud, b.LoadWarehouses, loadArgs)
+	if err := c.RunE(ctx, loadNode, cmd); err != nil {
+		return err
+	}
+	if rebalanceWait == 0 {
+		return nil
+	}
+
 	t.l.Printf("waiting %v for rebalancing\n", rebalanceWait)
 	_, err := db.ExecContext(ctx, `SET CLUSTER SETTING kv.snapshot_rebalance.max_rate='64MiB'`)
 	if err != nil {
@@ -617,9 +619,9 @@ func loadTPCCBench(
 	// Split and scatter the tables. Ramp up to the expected load in the desired
 	// distribution. This should allow for load-based rebalancing to help
 	// distribute load. Optionally pass some load configuration-specific flags.
-	cmd = fmt.Sprintf("./workload run tpcc --warehouses=%d --workers=%d --split --scatter "+
-		"--wait=false --duration=%s --tolerate-errors %s {pgurl%s}",
-		b.LoadWarehouses, b.LoadWarehouses, rebalanceWait, partArgs, roachNodes)
+	cmd = fmt.Sprintf("./workload run tpcc --warehouses=%d --workers=%d "+
+		"--wait=false --duration=%s --tolerate-errors {pgurl%s}",
+		b.LoadWarehouses, b.LoadWarehouses, rebalanceWait, roachNodes)
 	if out, err := c.RunWithBuffer(ctx, c.l, loadNode, cmd); err != nil {
 		return errors.Wrapf(err, "failed with output %q", string(out))
 	}
@@ -770,9 +772,9 @@ func runTPCCBench(ctx context.Context, t *test, c *cluster, b tpccBenchSpec) {
 					case singleLoadgen:
 						// Nothing.
 					case singlePartitionedLoadgen:
-						extraFlags = fmt.Sprintf(` --partitions=%d --split`, b.partitions())
+						extraFlags = fmt.Sprintf(` --partitions=%d`, b.partitions())
 					case multiLoadgen:
-						extraFlags = fmt.Sprintf(" --partitions=%d --partition-affinity=%d --split",
+						extraFlags = fmt.Sprintf(` --partitions=%d --partition-affinity=%d`,
 							b.partitions(), groupIdx)
 						activeWarehouses = warehouses / numLoadGroups
 					default:

pkg/workload/cli/run.go

@@ -309,6 +309,7 @@ func runRun(gen workload.Generator, urls []string, dbName string) error {
 			if !*tolerateErrors {
 				return err
 			}
+			log.Infof(ctx, "retrying after error during init: %v", err)
 		}
 	}
 
@@ -324,9 +325,16 @@ func runRun(gen workload.Generator, urls []string, dbName string) error {
 		return errors.Errorf(`no operations defined for %s`, gen.Meta().Name)
 	}
 	reg := histogram.NewRegistry()
-	ops, err := o.Ops(urls, reg)
-	if err != nil {
-		return err
+	var ops workload.QueryLoad
+	for {
+		ops, err = o.Ops(urls, reg)
+		if err == nil {
+			break
+		}
+		if !*tolerateErrors {
+			return err
+		}
+		log.Infof(ctx, "retrying after error while creating load: %v", err)
 	}
 
 	const splitConcurrency = 384 // TODO(dan): Don't hardcode this.

pkg/workload/tpcc/ddls.go

@@ -15,9 +15,9 @@
 package tpcc
 
 import (
+	gosql "database/sql"
 	"fmt"
 
-	"github.com/jackc/pgx"
 	"github.com/pkg/errors"
 	"golang.org/x/sync/errgroup"
 )
@@ -157,7 +157,7 @@ const (
 	tpccOrderLineSchemaInterleave = ` interleave in parent "order" (ol_w_id, ol_d_id, ol_o_id)`
 )
 
-func scatterRanges(db *pgx.ConnPool) {
+func scatterRanges(db *gosql.DB) error {
 	tables := []string{
 		`customer`,
 		`district`,
@@ -180,7 +180,5 @@ func scatterRanges(db *pgx.ConnPool) {
 			return nil
 		})
 	}
-	if err := g.Wait(); err != nil {
-		panic(err)
-	}
+	return g.Wait()
 }