From 36a118767dc12ef4c7115d0652744e13c7949d55 Mon Sep 17 00:00:00 2001 From: Michael Butler Date: Mon, 27 Mar 2023 18:16:03 -0400 Subject: [PATCH] roachtest: create simple tpce workload driver Currently, the tpce workload is used in 2 roachtests in tpce.go. This patch makes it easier for other roachtests to init and run a tpce workload by creating a tpce driver that abstracts away the actual cmds required to init and run a tpce workload. This will make it easier to address #99787, for example. Informs #99787 Release note: none --- pkg/cmd/roachtest/tests/tpce.go | 105 +++++++++++++++++++++++++------- 1 file changed, 83 insertions(+), 22 deletions(-) diff --git a/pkg/cmd/roachtest/tests/tpce.go b/pkg/cmd/roachtest/tests/tpce.go index 1e54c7668ff5..392f70f719f9 100644 --- a/pkg/cmd/roachtest/tests/tpce.go +++ b/pkg/cmd/roachtest/tests/tpce.go @@ -19,12 +19,82 @@ import ( "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry" + "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/spec" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test" "github.com/cockroachdb/cockroach/pkg/roachprod/install" + "github.com/cockroachdb/cockroach/pkg/roachprod/logger" "github.com/cockroachdb/errors" + "github.com/stretchr/testify/require" ) +type tpceSpec struct { + loadNode int + roachNodes option.NodeListOption + roachNodeIPFlags []string +} + +type tpceCmdOptions struct { + customers int + racks int + duration time.Duration + threads int +} + +func (to tpceCmdOptions) AddCommandOptions(cmd *roachtestutil.Command) { + cmd.MaybeFlag(to.customers != 0, "customers", to.customers) + cmd.MaybeFlag(to.racks != 0, "racks", to.racks) + cmd.MaybeFlag(to.duration != 0, "duration", to.duration) + cmd.MaybeFlag(to.threads != 0, "threads", to.threads) +} + +func initTPCESpec( + ctx context.Context, + l *logger.Logger, + c cluster.Cluster, + loadNode int, + roachNodes option.NodeListOption, +) (*tpceSpec, error) { + l.Printf("Installing docker") + if err := c.Install(ctx, l, c.Nodes(loadNode), "docker"); err != nil { + return nil, err + } + roachNodeIPs, err := c.InternalIP(ctx, l, roachNodes) + if err != nil { + return nil, err + } + roachNodeIPFlags := make([]string, len(roachNodeIPs)) + for i, ip := range roachNodeIPs { + roachNodeIPFlags[i] = fmt.Sprintf("--hosts=%s", ip) + } + return &tpceSpec{ + loadNode: loadNode, + roachNodes: roachNodes, + roachNodeIPFlags: roachNodeIPFlags, + }, nil +} + +func (ts *tpceSpec) newCmd(o tpceCmdOptions) *roachtestutil.Command { + cmd := roachtestutil.NewCommand(`sudo docker run cockroachdb/tpc-e:latest`) + o.AddCommandOptions(cmd) + return cmd +} + +// init initializes an empty cluster with a tpce database. This includes a bulk +// import of the data and schema creation. +func (ts *tpceSpec) init(ctx context.Context, t test.Test, c cluster.Cluster, o tpceCmdOptions) { + cmd := ts.newCmd(o).Option("init") + c.Run(ctx, c.Node(ts.loadNode), fmt.Sprintf("%s %s", cmd, ts.roachNodeIPFlags[0])) +} + +// run runs the tpce workload on cluster that has been initialized with the tpce schema. +func (ts *tpceSpec) run( + ctx context.Context, t test.Test, c cluster.Cluster, o tpceCmdOptions, +) (install.RunResultDetails, error) { + cmd := fmt.Sprintf("%s %s", ts.newCmd(o), strings.Join(ts.roachNodeIPFlags, " ")) + return c.RunWithDetailsSingleNode(ctx, t.L(), c.Node(ts.loadNode), cmd) +} + func registerTPCE(r registry.Registry) { type tpceOptions struct { owner registry.Owner // defaults to test-eng @@ -39,7 +109,7 @@ func registerTPCE(r registry.Registry) { runTPCE := func(ctx context.Context, t test.Test, c cluster.Cluster, opts tpceOptions) { roachNodes := c.Range(1, opts.nodes) - loadNode := c.Node(opts.nodes + 1) + loadNode := opts.nodes + 1 racks := opts.nodes t.Status("installing cockroach") @@ -50,10 +120,8 @@ func registerTPCE(r registry.Registry) { settings := install.MakeClusterSettings(install.NumRacksOption(racks)) c.Start(ctx, t.L(), startOpts, settings, roachNodes) - t.Status("installing docker") - if err := c.Install(ctx, t.L(), loadNode, "docker"); err != nil { - t.Fatal(err) - } + tpceSpec, err := initTPCESpec(ctx, t.L(), c, loadNode, roachNodes) + require.NoError(t, err) // Configure to increase the speed of the import. func() { @@ -73,27 +141,20 @@ func registerTPCE(r registry.Registry) { m := c.NewMonitor(ctx, roachNodes) m.Go(func(ctx context.Context) error { - const dockerRun = `sudo docker run cockroachdb/tpc-e:latest` - - roachNodeIPs, err := c.InternalIP(ctx, t.L(), roachNodes) - if err != nil { - return err - } - roachNodeIPFlags := make([]string, len(roachNodeIPs)) - for i, ip := range roachNodeIPs { - roachNodeIPFlags[i] = fmt.Sprintf("--hosts=%s", ip) - } t.Status("preparing workload") - c.Run(ctx, loadNode, fmt.Sprintf("%s --customers=%d --racks=%d --init %s", - dockerRun, opts.customers, racks, roachNodeIPFlags[0])) + tpceSpec.init(ctx, t, c, tpceCmdOptions{ + customers: opts.customers, + racks: racks, + }) t.Status("running workload") - duration := 2 * time.Hour - threads := opts.nodes * opts.cpus - result, err := c.RunWithDetailsSingleNode(ctx, t.L(), loadNode, - fmt.Sprintf("%s --customers=%d --racks=%d --duration=%s --threads=%d %s", - dockerRun, opts.customers, racks, duration, threads, strings.Join(roachNodeIPFlags, " "))) + result, err := tpceSpec.run(ctx, t, c, tpceCmdOptions{ + customers: opts.customers, + racks: racks, + duration: 2 * time.Hour, + threads: opts.nodes * opts.cpus, + }) if err != nil { t.Fatal(err.Error()) }