From 854991099c83a7bf5bc99aa8115f9620156d0e83 Mon Sep 17 00:00:00 2001 From: Callum Styan Date: Thu, 20 Aug 2026 11:46:44 -0700 Subject: [PATCH] feat(scaletest): add create-users scaletest subcommand (#28287) --- cli/exp_scaletest.go | 1 + cli/exp_scaletest_createusers.go | 245 +++++++++++++++++++++++++++++++ cli/exp_scaletest_test.go | 57 +++++++ 3 files changed, 303 insertions(+) create mode 100644 cli/exp_scaletest_createusers.go diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go index c49a228a54d..d8979b675ab 100644 --- a/cli/exp_scaletest.go +++ b/cli/exp_scaletest.go @@ -61,6 +61,7 @@ func (r *RootCmd) scaletestCmd() *serpent.Command { r.scaletestCleanup(), r.scaletestDashboard(), r.scaletestDynamicParameters(), + r.scaletestCreateUsers(), r.scaletestCreateWorkspaces(), r.scaletestWorkspaceUpdates(), r.scaletestWorkspaceTraffic(), diff --git a/cli/exp_scaletest_createusers.go b/cli/exp_scaletest_createusers.go new file mode 100644 index 00000000000..2f5468bf46d --- /dev/null +++ b/cli/exp_scaletest_createusers.go @@ -0,0 +1,245 @@ +//go:build !slim + +package cli + +import ( + "context" + "fmt" + "io" + "os/signal" + "strconv" + + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/tracing" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/scaletest/createusers" + "github.com/coder/coder/v2/scaletest/harness" + "github.com/coder/coder/v2/scaletest/loadtestutil" + "github.com/coder/serpent" +) + +func (r *RootCmd) scaletestCreateUsers() *serpent.Command { + var ( + count int64 + templateAdminPercentage float64 + noCleanup bool + + tracingFlags = &scaletestTracingFlags{} + strategy = &scaletestStrategyFlags{} + cleanupStrategy = newScaletestCleanupStrategy() + output = &scaletestOutputFlags{} + ) + + cmd := &serpent.Command{ + Use: "create-users", + Short: "Create many users for use by other scale tests.", + Long: "Pre-provision a fixed number of users (optionally promoting a percentage to " + + "template admin) so that other load generators can target existing users instead " + + "of creating their own inline. Intended to be run against dedicated, throwaway " + + "scaletest deployments.", + Handler: func(inv *serpent.Invocation) error { + ctx := inv.Context() + client, err := r.InitClient(inv) + if err != nil { + return err + } + + notifyCtx, stop := signal.NotifyContext(ctx, StopSignals...) + defer stop() + ctx = notifyCtx + + me, err := RequireAdmin(ctx, client) + if err != nil { + return err + } + + if count <= 0 { + return xerrors.New("--count must be greater than 0") + } + if templateAdminPercentage < 0 || templateAdminPercentage > 100 { + return xerrors.New("--template-admin-percentage must be between 0 and 100") + } + + // int64 truncates toward zero, so this floors the count. Ensure at + // least one admin whenever a nonzero percentage was requested. + templateAdminCount := int64(float64(count) * templateAdminPercentage / 100) + if templateAdminPercentage > 0 { + templateAdminCount = max(templateAdminCount, 1) + } + regularUserCount := count - templateAdminCount + + _, _ = fmt.Fprintf(inv.Stderr, "Distribution plan:\n") + _, _ = fmt.Fprintf(inv.Stderr, " Total users: %d\n", count) + _, _ = fmt.Fprintf(inv.Stderr, " Template admins: %d (%.1f%%)\n", templateAdminCount, templateAdminPercentage) + _, _ = fmt.Fprintf(inv.Stderr, " Regular users: %d (%.1f%%)\n", regularUserCount, 100.0-templateAdminPercentage) + + outputs, err := output.parse() + if err != nil { + return xerrors.Errorf("could not parse --output flags") + } + + tracerProvider, closeTracing, tracingEnabled, err := tracingFlags.provider(ctx) + if err != nil { + return xerrors.Errorf("create tracer provider: %w", err) + } + defer func() { + _, _ = fmt.Fprintln(inv.Stderr, "\nUploading traces...") + if err := closeTracing(ctx); err != nil { + _, _ = fmt.Fprintf(inv.Stderr, "\nError uploading traces: %+v\n", err) + } + }() + tracer := tracerProvider.Tracer(scaletestTracerName) + + th := harness.NewTestHarness(strategy.toStrategy(), cleanupStrategy.toStrategy()) + + _, _ = fmt.Fprintln(inv.Stderr, "Creating users...") + for i := range count { + id := strconv.FormatInt(i, 10) + name := fmt.Sprintf("createusers-%s", id) + + var roles []string + if i < templateAdminCount { + roles = []string{codersdk.RoleTemplateAdmin} + } + + // Use an independent client per runner so they don't reuse TCP + // connections, which can unbalance requests among Coder replicas. + runnerClient, err := loadtestutil.DupClientCopyingHeaders(client, BypassHeader) + if err != nil { + return xerrors.Errorf("create runner client: %w", err) + } + + // Username and email are generated by createusers when left empty. + config := createusers.Config{ + OrganizationID: me.OrganizationIDs[0], + } + if err := config.Validate(); err != nil { + return xerrors.Errorf("validate config: %w", err) + } + + var runner harness.Runnable = &createUsersRunner{ + client: runnerClient, + runner: createusers.NewRunner(runnerClient, config), + roles: roles, + } + if tracingEnabled { + runner = &runnableTraceWrapper{ + tracer: tracer, + spanName: name, + runner: runner, + } + } + + th.AddRun(name, id, runner) + } + + _, _ = fmt.Fprintln(inv.Stderr, "Running user creation scaletest...") + testCtx, testCancel := strategy.toContext(ctx) + defer testCancel() + err = th.Run(testCtx) + if err != nil { + return xerrors.Errorf("run test harness (harness failure, not a test failure): %w", err) + } + + // If the command was interrupted, skip stats. + if notifyCtx.Err() != nil { + return notifyCtx.Err() + } + + res := th.Results() + for _, o := range outputs { + err = o.write(res, inv.Stdout) + if err != nil { + return xerrors.Errorf("write output %q to %q: %w", o.format, o.path, err) + } + } + + if noCleanup { + _, _ = fmt.Fprintln(inv.Stderr, "\nSkipping cleanup (users left in place for other scale tests).") + } else { + _, _ = fmt.Fprintln(inv.Stderr, "\nCleaning up...") + cleanupCtx, cleanupCancel := cleanupStrategy.toContext(ctx) + defer cleanupCancel() + err = th.Cleanup(cleanupCtx) + if err != nil { + return xerrors.Errorf("cleanup tests: %w", err) + } + } + + if res.TotalFail > 0 { + return xerrors.New("load test failed, see above for more details") + } + + return nil + }, + } + + cmd.Options = serpent.OptionSet{ + { + Flag: "count", + FlagShorthand: "c", + Env: "CODER_SCALETEST_CREATE_USERS_COUNT", + Description: "Required: Total number of users to create.", + Value: serpent.Int64Of(&count), + Required: true, + }, + { + Flag: "template-admin-percentage", + Env: "CODER_SCALETEST_CREATE_USERS_TEMPLATE_ADMIN_PERCENTAGE", + Default: "20.0", + Description: "Percentage of users to assign the Template Admin role to (0-100).", + Value: serpent.Float64Of(&templateAdminPercentage), + }, + { + Flag: "no-cleanup", + Env: "CODER_SCALETEST_NO_CLEANUP", + Description: "Do not clean up resources after the test completes. Use this when other scale tests will reuse the created users.", + Value: serpent.BoolOf(&noCleanup), + }, + } + + tracingFlags.attach(&cmd.Options) + strategy.attach(&cmd.Options) + cleanupStrategy.attach(&cmd.Options) + output.attach(&cmd.Options) + return cmd +} + +// createUsersRunner adapts createusers.Runner to harness.Runnable and assigns +// roles after creation. The client must be a site admin to assign roles. +type createUsersRunner struct { + client *codersdk.Client + runner *createusers.Runner + roles []string +} + +var ( + _ harness.Runnable = &createUsersRunner{} + _ harness.Cleanable = &createUsersRunner{} +) + +func (r *createUsersRunner) Run(ctx context.Context, id string, logs io.Writer) error { + ctx, span := tracing.StartSpan(ctx) + defer span.End() + + user, err := r.runner.RunReturningUser(ctx, id, logs) + if err != nil { + return xerrors.Errorf("create user: %w", err) + } + + if len(r.roles) > 0 { + _, _ = fmt.Fprintf(logs, "Assigning roles to user %q: %v\n", user.Username, r.roles) + if _, err := r.client.UpdateUserRoles(ctx, user.ID.String(), codersdk.UpdateRoles{ + Roles: r.roles, + }); err != nil { + return xerrors.Errorf("assign roles %v to user %q: %w", r.roles, user.Username, err) + } + } + + return nil +} + +func (r *createUsersRunner) Cleanup(ctx context.Context, id string, logs io.Writer) error { + return r.runner.Cleanup(ctx, id, logs) +} diff --git a/cli/exp_scaletest_test.go b/cli/exp_scaletest_test.go index 98d2071ad0a..fbafdf8e79d 100644 --- a/cli/exp_scaletest_test.go +++ b/cli/exp_scaletest_test.go @@ -10,6 +10,8 @@ import ( "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/cli/clitest" "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/scaletest/loadtestutil" "github.com/coder/coder/v2/testutil" ) @@ -59,6 +61,61 @@ func TestScaleTestCreateWorkspaces(t *testing.T) { require.ErrorContains(t, err, "could not find template \"doesnotexist\" in any organization") } +func TestScaleTestCreateUsers(t *testing.T) { + t.Parallel() + + if testutil.RaceEnabled() { + t.Skip("Skipping due to race detector") + } + + ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitLong) + defer cancelFunc() + + log := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + client := coderdtest.New(t, &coderdtest.Options{ + // No provisioner daemons are required because this command only creates + // users. + Logger: &log, + }) + _ = coderdtest.CreateFirstUser(t, client) + + inv, root := clitest.New(t, "exp", "scaletest", "create-users", + "--count", "3", + "--template-admin-percentage", "34", + "--no-cleanup", + "--concurrency", "2", + "--timeout", "30s", + "--job-timeout", "15s", + "--cleanup-concurrency", "1", + "--cleanup-timeout", "30s", + "--cleanup-job-timeout", "15s", + "--output", "text", + ) + clitest.SetupConfig(t, client, root) + err := inv.WithContext(ctx).Run() + require.NoError(t, err) + + // Verify the users were created and roughly the requested percentage are + // template admins (34% of 3 rounds down to 1). + res, err := client.Users(ctx, codersdk.UsersRequest{Search: loadtestutil.ScaleTestPrefix + "-"}) + require.NoError(t, err) + + var created, templateAdmins int + for _, u := range res.Users { + if !loadtestutil.IsScaleTestUser(u.Username, u.Email) { + continue + } + created++ + for _, role := range u.Roles { + if role.Name == codersdk.RoleTemplateAdmin { + templateAdmins++ + } + } + } + require.Equal(t, 3, created) + require.Equal(t, 1, templateAdmins) +} + // This test just validates that the CLI command accepts its known arguments. // A more comprehensive test is performed in workspacetraffic/run_test.go func TestScaleTestWorkspaceTraffic(t *testing.T) {