From 4b2ec2593b258df1ba15a3302aa998a5fe9a8fa6 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Tue, 7 Nov 2023 14:50:36 +0200 Subject: [PATCH] feat(cli): add template filter support to exp scaletest cleanup and traffic --- cli/exp_scaletest.go | 112 +++++++++++++++++++++++++++----------- cli/exp_scaletest_test.go | 50 +++++++++++++++++ 2 files changed, 130 insertions(+), 32 deletions(-) diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go index 7ecbb72483360..32f3bdd47e062 100644 --- a/cli/exp_scaletest.go +++ b/cli/exp_scaletest.go @@ -394,6 +394,8 @@ func (r *userCleanupRunner) Run(ctx context.Context, _ string, _ io.Writer) erro } func (r *RootCmd) scaletestCleanup() *clibase.Cmd { + var template string + cleanupStrategy := &scaletestStrategyFlags{cleanup: true} client := new(codersdk.Client) @@ -407,7 +409,7 @@ func (r *RootCmd) scaletestCleanup() *clibase.Cmd { Handler: func(inv *clibase.Invocation) error { ctx := inv.Context() - _, err := requireAdmin(ctx, client) + me, err := requireAdmin(ctx, client) if err != nil { return err } @@ -421,8 +423,15 @@ func (r *RootCmd) scaletestCleanup() *clibase.Cmd { }, } + if template != "" { + _, err := parseTemplate(ctx, client, me.OrganizationIDs, template) + if err != nil { + return xerrors.Errorf("parse template: %w", err) + } + } + cliui.Infof(inv.Stdout, "Fetching scaletest workspaces...") - workspaces, err := getScaletestWorkspaces(ctx, client) + workspaces, err := getScaletestWorkspaces(ctx, client, template) if err != nil { return err } @@ -494,6 +503,15 @@ func (r *RootCmd) scaletestCleanup() *clibase.Cmd { }, } + cmd.Options = clibase.OptionSet{ + { + Flag: "template", + Env: "CODER_SCALETEST_CLEANUP_TEMPLATE", + Description: "Name or ID of the template. Only delete workspaces created from the given template.", + Value: clibase.StringOf(&template), + }, + } + cleanupStrategy.attach(&cmd.Options) return cmd } @@ -564,34 +582,12 @@ func (r *RootCmd) scaletestCreateWorkspaces() *clibase.Cmd { return xerrors.Errorf("could not parse --output flags") } - var tpl codersdk.Template if template == "" { return xerrors.Errorf("--template is required") } - if id, err := uuid.Parse(template); err == nil && id != uuid.Nil { - tpl, err = client.Template(ctx, id) - if err != nil { - return xerrors.Errorf("get template by ID %q: %w", template, err) - } - } else { - // List templates in all orgs until we find a match. - orgLoop: - for _, orgID := range me.OrganizationIDs { - tpls, err := client.TemplatesByOrganization(ctx, orgID) - if err != nil { - return xerrors.Errorf("list templates in org %q: %w", orgID, err) - } - - for _, t := range tpls { - if t.Name == template { - tpl = t - break orgLoop - } - } - } - } - if tpl.ID == uuid.Nil { - return xerrors.Errorf("could not find template %q in any organization", template) + tpl, err := parseTemplate(ctx, client, me.OrganizationIDs, template) + if err != nil { + return xerrors.Errorf("parse template: %w", err) } cliRichParameters, err := asWorkspaceBuildParameters(parameterFlags.richParameters) @@ -859,6 +855,7 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd { tickInterval time.Duration bytesPerTick int64 ssh bool + template string client = &codersdk.Client{} tracingFlags = &scaletestTracingFlags{} @@ -876,6 +873,12 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd { ), Handler: func(inv *clibase.Invocation) error { ctx := inv.Context() + + me, err := requireAdmin(ctx, client) + if err != nil { + return err + } + reg := prometheus.NewRegistry() metrics := workspacetraffic.NewMetrics(reg, "username", "workspace_name", "agent_name") @@ -893,7 +896,14 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd { }, } - workspaces, err := getScaletestWorkspaces(inv.Context(), client) + if template != "" { + _, err := parseTemplate(ctx, client, me.OrganizationIDs, template) + if err != nil { + return xerrors.Errorf("parse template: %w", err) + } + } + + workspaces, err := getScaletestWorkspaces(inv.Context(), client, template) if err != nil { return err } @@ -997,6 +1007,13 @@ func (r *RootCmd) scaletestWorkspaceTraffic() *clibase.Cmd { } cmd.Options = []clibase.Option{ + { + Flag: "template", + FlagShorthand: "t", + Env: "CODER_SCALETEST_TEMPLATE", + Description: "Name or ID of the template. Traffic generation will be limited to workspaces created from this template.", + Value: clibase.StringOf(&template), + }, { Flag: "bytes-per-tick", Env: "CODER_SCALETEST_WORKSPACE_TRAFFIC_BYTES_PER_TICK", @@ -1281,7 +1298,7 @@ func isScaleTestWorkspace(workspace codersdk.Workspace) bool { strings.HasPrefix(workspace.Name, "scaletest-") } -func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client) ([]codersdk.Workspace, error) { +func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client, template string) ([]codersdk.Workspace, error) { var ( pageNumber = 0 limit = 100 @@ -1290,9 +1307,10 @@ func getScaletestWorkspaces(ctx context.Context, client *codersdk.Client) ([]cod for { page, err := client.Workspaces(ctx, codersdk.WorkspaceFilter{ - Name: "scaletest-", - Offset: pageNumber * limit, - Limit: limit, + Name: "scaletest-", + Template: template, + Offset: pageNumber * limit, + Limit: limit, }) if err != nil { return nil, xerrors.Errorf("fetch scaletest workspaces page %d: %w", pageNumber, err) @@ -1349,3 +1367,33 @@ func getScaletestUsers(ctx context.Context, client *codersdk.Client) ([]codersdk return users, nil } + +func parseTemplate(ctx context.Context, client *codersdk.Client, organizationIDs []uuid.UUID, template string) (tpl codersdk.Template, err error) { + if id, err := uuid.Parse(template); err == nil && id != uuid.Nil { + tpl, err = client.Template(ctx, id) + if err != nil { + return tpl, xerrors.Errorf("get template by ID %q: %w", template, err) + } + } else { + // List templates in all orgs until we find a match. + orgLoop: + for _, orgID := range organizationIDs { + tpls, err := client.TemplatesByOrganization(ctx, orgID) + if err != nil { + return tpl, xerrors.Errorf("list templates in org %q: %w", orgID, err) + } + + for _, t := range tpls { + if t.Name == template { + tpl = t + break orgLoop + } + } + } + } + if tpl.ID == uuid.Nil { + return tpl, xerrors.Errorf("could not find template %q in any organization", template) + } + + return tpl, nil +} diff --git a/cli/exp_scaletest_test.go b/cli/exp_scaletest_test.go index 556aed6c21a82..a96d0daaa9014 100644 --- a/cli/exp_scaletest_test.go +++ b/cli/exp_scaletest_test.go @@ -91,6 +91,56 @@ func TestScaleTestWorkspaceTraffic(t *testing.T) { require.ErrorContains(t, err, "no scaletest workspaces exist") } +// This test just validates that the CLI command accepts its known arguments. +func TestScaleTestWorkspaceTraffic_Template(t *testing.T) { + t.Parallel() + + ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitMedium) + 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", "workspace-traffic", + "--template", "doesnotexist", + ) + clitest.SetupConfig(t, client, root) + pty := ptytest.New(t) + inv.Stdout = pty.Output() + inv.Stderr = pty.Output() + + err := inv.WithContext(ctx).Run() + require.ErrorContains(t, err, "could not find template \"doesnotexist\" in any organization") +} + +// This test just validates that the CLI command accepts its known arguments. +func TestScaleTestCleanup_Template(t *testing.T) { + t.Parallel() + + ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitMedium) + 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", "cleanup", + "--template", "doesnotexist", + ) + clitest.SetupConfig(t, client, root) + pty := ptytest.New(t) + inv.Stdout = pty.Output() + inv.Stderr = pty.Output() + + err := inv.WithContext(ctx).Run() + require.ErrorContains(t, err, "could not find template \"doesnotexist\" in any organization") +} + // This test just validates that the CLI command accepts its known arguments. func TestScaleTestDashboard(t *testing.T) { t.Parallel()