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