feat(scaletest/notifications): add opt-in --reuse-users (#28289)

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
+}

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) {

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 {

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)
+			}
 		}
 	}