diff --git a/cli/exp_scaletest_notifications.go b/cli/exp_scaletest_notifications.go index 6b765bc7d61..918729ddcd5 100644 --- a/cli/exp_scaletest_notifications.go +++ b/cli/exp_scaletest_notifications.go @@ -36,6 +36,7 @@ func (r *RootCmd) scaletestNotifications() *serpent.Command { smtpRequestTimeout time.Duration dialTimeout time.Duration noCleanup bool + reuseUsers bool smtpAPIURL string tracingFlags = &scaletestTracingFlags{} @@ -117,7 +118,19 @@ func (r *RootCmd) scaletestNotifications() *serpent.Command { <-time.After(prometheusFlags.Wait) }() - _, _ = fmt.Fprintln(inv.Stderr, "Creating users...") + var adminReuse, regularReuse []notificationReuseUser + if reuseUsers { + _, _ = fmt.Fprintln(inv.Stderr, "Reusing existing scaletest users...") + // Bound token lifetime to just beyond the run so tokens orphaned by a + // hard kill expire quickly rather than at the deployment default. + tokenLifetime := notificationTimeout + dialTimeout + time.Hour + adminReuse, regularReuse, err = selectNotificationReuseUsers(ctx, client, int(templateAdminCount), int(regularUserCount), tokenLifetime) + if err != nil { + return err + } + } else { + _, _ = fmt.Fprintln(inv.Stderr, "Creating users...") + } dialBarrier := &sync.WaitGroup{} templateAdminWatchBarrier := &sync.WaitGroup{} @@ -143,12 +156,8 @@ func (r *RootCmd) scaletestNotifications() *serpent.Command { } configs := make([]notifications.Config, 0, userCount) - for range templateAdminCount { + for i := range int(templateAdminCount) { config := notifications.Config{ - User: createusers.Config{ - OrganizationID: me.OrganizationIDs[0], - }, - Roles: []string{codersdk.RoleTemplateAdmin}, NotificationTimeout: notificationTimeout, DialTimeout: dialTimeout, DialBarrier: dialBarrier, @@ -159,23 +168,33 @@ func (r *RootCmd) scaletestNotifications() *serpent.Command { SMTPRequestTimeout: smtpRequestTimeout, SMTPHttpClient: smtpHTTPClient, } + if reuseUsers { + config.SessionToken = adminReuse[i].sessionToken + config.PreCreatedUser = adminReuse[i].user + } else { + config.User = createusers.Config{OrganizationID: me.OrganizationIDs[0]} + config.Roles = []string{codersdk.RoleTemplateAdmin} + } if err := config.Validate(); err != nil { return xerrors.Errorf("validate config: %w", err) } configs = append(configs, config) } - for range regularUserCount { + for i := range int(regularUserCount) { config := notifications.Config{ - User: createusers.Config{ - OrganizationID: me.OrganizationIDs[0], - }, - Roles: []string{}, NotificationTimeout: notificationTimeout, DialTimeout: dialTimeout, DialBarrier: dialBarrier, ReceivingWatchBarrier: templateAdminWatchBarrier, Metrics: metrics, } + if reuseUsers { + config.SessionToken = regularReuse[i].sessionToken + config.PreCreatedUser = regularReuse[i].user + } else { + config.User = createusers.Config{OrganizationID: me.OrganizationIDs[0]} + config.Roles = []string{} + } if err := config.Validate(); err != nil { return xerrors.Errorf("validate config: %w", err) } @@ -302,6 +321,12 @@ func (r *RootCmd) scaletestNotifications() *serpent.Command { Description: "Do not clean up resources after the test completes.", Value: serpent.BoolOf(&noCleanup), }, + { + Flag: "reuse-users", + Env: "CODER_SCALETEST_NOTIFICATION_REUSE_USERS", + Description: "Reuse existing scaletest users instead of creating new ones. Enough users, including template admins, must already exist (see \"coder exp scaletest create-users\") or the command errors. Run only one user-selecting scaletest command at a time to avoid overlapping selections.", + Value: serpent.BoolOf(&reuseUsers), + }, { Flag: "smtp-api-url", Env: "CODER_SCALETEST_SMTP_API_URL", @@ -475,3 +500,71 @@ func triggerNotifications( expectedNotifications[notificationsLib.TemplateTemplateDeleted] <- time.Now() close(expectedNotifications[notificationsLib.TemplateTemplateDeleted]) } + +type notificationReuseUser struct { + user codersdk.User + sessionToken string +} + +// selectNotificationReuseUsers selects existing scaletest users and mints a token +// for each, erroring if the pool lacks enough template admins or regular users. +func selectNotificationReuseUsers(ctx context.Context, client *codersdk.Client, adminCount, regularCount int, tokenLifetime time.Duration) (admins, regulars []notificationReuseUser, err error) { + users, err := getScaletestUsers(ctx, client) + if err != nil { + return nil, nil, xerrors.Errorf("list scaletest users: %w", err) + } + + var adminUsers, regularUsers []codersdk.User + for _, u := range users { + if userHasRole(u, codersdk.RoleTemplateAdmin) { + adminUsers = append(adminUsers, u) + } else { + regularUsers = append(regularUsers, u) + } + } + + if len(adminUsers) < adminCount || len(regularUsers) < regularCount { + total := adminCount + regularCount + var pct float64 + if total > 0 { + pct = float64(adminCount) / float64(total) * 100 + } + return nil, nil, xerrors.Errorf( + "not enough scaletest users to reuse: found %d template admins and %d regular users, need %d and %d. "+ + "Create them first, for example: coder exp scaletest create-users --count %d --template-admin-percentage %.2f --no-cleanup", + len(adminUsers), len(regularUsers), adminCount, regularCount, total, pct) + } + + // Token names are auto-generated by the server; we only need the returned key. + mint := func(selected []codersdk.User) ([]notificationReuseUser, error) { + reuse := make([]notificationReuseUser, 0, len(selected)) + for _, u := range selected { + res, err := client.CreateToken(ctx, u.ID.String(), codersdk.CreateTokenRequest{ + Lifetime: tokenLifetime, + }) + if err != nil { + return nil, xerrors.Errorf("mint token for user %q: %w", u.Username, err) + } + reuse = append(reuse, notificationReuseUser{user: u, sessionToken: res.Key}) + } + return reuse, nil + } + + if admins, err = mint(adminUsers[:adminCount]); err != nil { + return nil, nil, err + } + if regulars, err = mint(regularUsers[:regularCount]); err != nil { + return nil, nil, err + } + + return admins, regulars, nil +} + +func userHasRole(u codersdk.User, roleName string) bool { + for _, role := range u.Roles { + if role.Name == roleName { + return true + } + } + return false +} diff --git a/cli/exp_scaletest_test.go b/cli/exp_scaletest_test.go index fbafdf8e79d..f399f1a78a5 100644 --- a/cli/exp_scaletest_test.go +++ b/cli/exp_scaletest_test.go @@ -116,6 +116,40 @@ func TestScaleTestCreateUsers(t *testing.T) { require.Equal(t, 1, templateAdmins) } +// TestScaleTestNotifications_ReuseUsersInsufficient verifies that --reuse-users +// checks the pool up front and exits with an actionable error when not enough +// scaletest users (or template admins) exist, rather than creating any. +func TestScaleTestNotifications_ReuseUsersInsufficient(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{ + Logger: &log, + }) + _ = coderdtest.CreateFirstUser(t, client) + + inv, root := clitest.New(t, "exp", "scaletest", "notifications", + "--user-count", "2", + "--template-admin-percentage", "50", + "--reuse-users", + "--dial-timeout", "5s", + "--notification-timeout", "5s", + "--scaletest-prometheus-address", "127.0.0.1:0", + "--scaletest-prometheus-wait", "0s", + "--output", "text", + ) + clitest.SetupConfig(t, client, root) + err := inv.WithContext(ctx).Run() + require.ErrorContains(t, err, "not enough scaletest users to reuse") +} + // 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) { diff --git a/scaletest/notifications/config.go b/scaletest/notifications/config.go index 372199bf932..0bd5a56e981 100644 --- a/scaletest/notifications/config.go +++ b/scaletest/notifications/config.go @@ -8,6 +8,7 @@ import ( "github.com/google/uuid" "golang.org/x/xerrors" + "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/scaletest/createusers" ) @@ -18,6 +19,15 @@ type Config struct { // Roles are the roles to assign to the user. Roles []string `json:"roles"` + // SessionToken, when set, makes the runner reuse PreCreatedUser instead of + // creating one. User and Roles are then ignored. + SessionToken string `json:"-"` + + // PreCreatedUser is the user to run as when SessionToken is set. It must + // already hold any role needed to receive the notifications under test. Only + // ID and Email are used. + PreCreatedUser codersdk.User `json:"-"` + // NotificationTimeout is how long to wait for notifications after triggering. NotificationTimeout time.Duration `json:"notification_timeout"` @@ -46,13 +56,20 @@ type Config struct { } func (c Config) Validate() error { - // The runner always needs an org; ensure we propagate it into the user config. - if c.User.OrganizationID == uuid.Nil { - return xerrors.New("user organization_id must be set") - } + if c.SessionToken != "" { + // Reuse mode does not create a user, so only the existing user's ID is required. + if c.PreCreatedUser.ID == uuid.Nil { + return xerrors.New("pre_created_user must be set when session_token is set") + } + } else { + // The runner always needs an org; ensure we propagate it into the user config. + if c.User.OrganizationID == uuid.Nil { + return xerrors.New("user organization_id must be set") + } - if err := c.User.Validate(); err != nil { - return xerrors.Errorf("user config: %w", err) + if err := c.User.Validate(); err != nil { + return xerrors.Errorf("user config: %w", err) + } } if c.DialBarrier == nil { diff --git a/scaletest/notifications/run.go b/scaletest/notifications/run.go index bfc305d9744..b1572b3c289 100644 --- a/scaletest/notifications/run.go +++ b/scaletest/notifications/run.go @@ -88,29 +88,45 @@ func (r *Runner) Run(ctx context.Context, id string, logs io.Writer) error { r.client.SetLogger(logger) r.client.SetLogBodies(true) - r.createUserRunner = createusers.NewRunner(r.client, r.cfg.User) - newUserAndToken, err := r.createUserRunner.RunReturningUser(ctx, id, logs) - if err != nil { - r.cfg.Metrics.AddError("create_user") - return xerrors.Errorf("create user: %w", err) - } - newUser := newUserAndToken.User - newUserClient := codersdk.New(r.client.URL, - codersdk.WithSessionToken(newUserAndToken.SessionToken), - codersdk.WithLogger(logger), - codersdk.WithLogBodies()) + var ( + newUser codersdk.User + newUserClient *codersdk.Client + ) + if r.cfg.SessionToken != "" { + // Reuse mode assigns no roles; the user must already hold any role needed + // to receive the notifications under test. + newUser = r.cfg.PreCreatedUser + newUserClient = codersdk.New(r.client.URL, + codersdk.WithSessionToken(r.cfg.SessionToken), + codersdk.WithLogger(logger), + codersdk.WithLogBodies()) - logger.Info(ctx, "runner user created", slog.F("username", newUser.Username), slog.F("user_id", newUser.ID.String())) - - if len(r.cfg.Roles) > 0 { - logger.Info(ctx, "assigning roles to user", slog.F("roles", r.cfg.Roles)) - - _, err := r.client.UpdateUserRoles(ctx, newUser.ID.String(), codersdk.UpdateRoles{ - Roles: r.cfg.Roles, - }) + logger.Info(ctx, "reusing existing user", slog.F("username", newUser.Username), slog.F("user_id", newUser.ID.String())) + } else { + r.createUserRunner = createusers.NewRunner(r.client, r.cfg.User) + newUserAndToken, err := r.createUserRunner.RunReturningUser(ctx, id, logs) if err != nil { - r.cfg.Metrics.AddError("assign_roles") - return xerrors.Errorf("assign roles: %w", err) + r.cfg.Metrics.AddError("create_user") + return xerrors.Errorf("create user: %w", err) + } + newUser = newUserAndToken.User + newUserClient = codersdk.New(r.client.URL, + codersdk.WithSessionToken(newUserAndToken.SessionToken), + codersdk.WithLogger(logger), + codersdk.WithLogBodies()) + + logger.Info(ctx, "runner user created", slog.F("username", newUser.Username), slog.F("user_id", newUser.ID.String())) + + if len(r.cfg.Roles) > 0 { + logger.Info(ctx, "assigning roles to user", slog.F("roles", r.cfg.Roles)) + + _, err := r.client.UpdateUserRoles(ctx, newUser.ID.String(), codersdk.UpdateRoles{ + Roles: r.cfg.Roles, + }) + if err != nil { + r.cfg.Metrics.AddError("assign_roles") + return xerrors.Errorf("assign roles: %w", err) + } } }