From 950660e392e560a84abf7c61a7190e21f87ce73c Mon Sep 17 00:00:00 2001 From: Callum Styan Date: Wed, 29 Apr 2026 10:28:27 -0700 Subject: [PATCH] fix(cli): extend `exp scaletest cleanup` to properly clean up prebuilds (#23628) Signed-off-by: Callum Styan Co-authored-by: Claude Sonnet 4.6 --- cli/exp_scaletest.go | 162 +++++++++++++++++++ cli/exp_scaletest_prebuilds_internal_test.go | 82 ++++++++++ scaletest/prebuilds/run.go | 6 +- 3 files changed, 249 insertions(+), 1 deletion(-) create mode 100644 cli/exp_scaletest_prebuilds_internal_test.go diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go index 8173d76c3d..1a6922c747 100644 --- a/cli/exp_scaletest.go +++ b/cli/exp_scaletest.go @@ -38,6 +38,7 @@ import ( "github.com/coder/coder/v2/scaletest/dashboard" "github.com/coder/coder/v2/scaletest/harness" "github.com/coder/coder/v2/scaletest/loadtestutil" + "github.com/coder/coder/v2/scaletest/prebuilds" "github.com/coder/coder/v2/scaletest/reconnectingpty" "github.com/coder/coder/v2/scaletest/workspacebuild" "github.com/coder/coder/v2/scaletest/workspacetraffic" @@ -519,6 +520,88 @@ func (r *userCleanupRunner) Run(ctx context.Context, _ string, _ io.Writer) erro return nil } +// prebuildTemplateCleanupRunner deletes a single scaletest prebuilds template. +// All prebuild workspaces must be deleted before this runs. +type prebuildTemplateCleanupRunner struct { + client *codersdk.Client + template codersdk.Template +} + +var _ harness.Runnable = &prebuildTemplateCleanupRunner{} + +// Run implements Runnable. +func (r *prebuildTemplateCleanupRunner) Run(ctx context.Context, _ string, _ io.Writer) error { + ctx, span := tracing.StartSpan(ctx) + defer span.End() + + if err := r.client.DeleteTemplate(ctx, r.template.ID); err != nil { + return xerrors.Errorf("delete template %q: %w", r.template.Name, err) + } + return nil +} + +// getScaletestPrebuildWorkspaces returns all prebuild workspaces that belong +// to scaletest templates. It uses getScaletestPrebuildsTemplates to scope the +// query so that legitimate (non-scaletest) prebuilds on the deployment are not +// caught in the cleanup. If template is non-empty only workspaces for that +// template are returned. +func getScaletestPrebuildWorkspaces(ctx context.Context, client *codersdk.Client, template string) ([]codersdk.Workspace, error) { + const pageSize = 100 + + templates, err := getScaletestPrebuildsTemplates(ctx, client, template) + if err != nil { + return nil, xerrors.Errorf("list scaletest prebuild templates: %w", err) + } + + seen := make(map[uuid.UUID]struct{}) + var result []codersdk.Workspace + + for _, tmpl := range templates { + for page := 0; ; page++ { + resp, err := client.Workspaces(ctx, codersdk.WorkspaceFilter{ + Template: tmpl.Name, + Offset: page * pageSize, + Limit: pageSize, + }) + if err != nil { + return nil, xerrors.Errorf("list workspaces for template %q (page %d): %w", tmpl.Name, page, err) + } + for _, ws := range resp.Workspaces { + if _, ok := seen[ws.ID]; !ok { + seen[ws.ID] = struct{}{} + result = append(result, ws) + } + } + if len(resp.Workspaces) < pageSize { + break + } + } + } + + return result, nil +} + +// getScaletestPrebuildsTemplates returns all templates created by the scaletest +// prebuilds runner (identified by prebuilds.TemplatePrefix). If template is +// non-empty only that named template is returned; it must start with +// prebuilds.TemplatePrefix or an error is returned. +func getScaletestPrebuildsTemplates(ctx context.Context, client *codersdk.Client, template string) ([]codersdk.Template, error) { + var filter codersdk.TemplateFilter + if template != "" { + if !strings.HasPrefix(template, prebuilds.TemplatePrefix) { + return nil, xerrors.Errorf("template %q is not a scaletest prebuilds template (expected prefix %q)", template, prebuilds.TemplatePrefix) + } + filter = codersdk.TemplateFilter{ExactName: template} + } else { + filter = codersdk.TemplateFilter{FuzzyName: prebuilds.TemplatePrefix} + } + templates, err := client.Templates(ctx, filter) + if err != nil { + return nil, xerrors.Errorf("list templates: %w", err) + } + return templates, nil +} + func (r *RootCmd) scaletestCleanup() *serpent.Command { var template string cleanupStrategy := newScaletestCleanupStrategy() @@ -555,6 +638,85 @@ func (r *RootCmd) scaletestCleanup() *serpent.Command { } } + cliui.Infof(inv.Stdout, "Pausing prebuilds reconciler...") + setPrebuild := func(val bool) error { + return client.PutPrebuildsSettings(ctx, codersdk.PrebuildsSettings{ReconciliationPaused: val}) + } + if err = setPrebuild(true); err != nil { + return xerrors.Errorf("pause prebuilds reconciler: %w", err) + } + defer func() { + cliui.Infof(inv.Stdout, "Resuming prebuilds reconciler...") + if resumeErr := setPrebuild(false); resumeErr != nil { + cliui.Errorf(inv.Stderr, "Failed to resume prebuilds reconciler: %+v\n", resumeErr) + } + }() + + cliui.Infof(inv.Stdout, "Fetching scaletest prebuild workspaces...") + prebuildWorkspaces, err := getScaletestPrebuildWorkspaces(ctx, client, template) + if err != nil { + return err + } + + cliui.Errorf(inv.Stderr, "Found %d scaletest prebuild workspaces\n", len(prebuildWorkspaces)) + if len(prebuildWorkspaces) != 0 { + cliui.Infof(inv.Stdout, "Deleting scaletest prebuild workspaces...") + prebuildWsHarness := harness.NewTestHarness(cleanupStrategy.toStrategy(), harness.ConcurrentExecutionStrategy{}) + + for i, ws := range prebuildWorkspaces { + const testName = "cleanup-prebuild-workspace" + prebuildWsHarness.AddRun(testName, strconv.Itoa(i), workspacebuild.NewCleanupRunner(client, ws.ID)) + } + + prebuildWsCtx, prebuildWsCancel := cleanupStrategy.toContext(ctx) + defer prebuildWsCancel() + if err := prebuildWsHarness.Run(prebuildWsCtx); err != nil { + return xerrors.Errorf("run test harness to delete prebuild workspaces (harness failure, not a test failure): %w", err) + } + + cliui.Infof(inv.Stdout, "Done deleting scaletest prebuild workspaces:") + prebuildWsRes := prebuildWsHarness.Results() + prebuildWsRes.PrintText(inv.Stderr) + + if prebuildWsRes.TotalFail > 0 { + return xerrors.Errorf("failed to delete %d scaletest prebuild workspace(s)", prebuildWsRes.TotalFail) + } + } + + cliui.Infof(inv.Stdout, "Fetching scaletest prebuilds templates...") + prebuildTemplates, err := getScaletestPrebuildsTemplates(ctx, client, template) + if err != nil { + return err + } + + cliui.Errorf(inv.Stderr, "Found %d scaletest prebuilds templates\n", len(prebuildTemplates)) + if len(prebuildTemplates) != 0 { + cliui.Infof(inv.Stdout, "Deleting scaletest prebuilds templates...") + prebuildTplHarness := harness.NewTestHarness(cleanupStrategy.toStrategy(), harness.ConcurrentExecutionStrategy{}) + + for i, t := range prebuildTemplates { + const testName = "cleanup-prebuilds-template" + prebuildTplHarness.AddRun(testName, strconv.Itoa(i), &prebuildTemplateCleanupRunner{ + client: client, + template: t, + }) + } + + prebuildTplCtx, prebuildTplCancel := cleanupStrategy.toContext(ctx) + defer prebuildTplCancel() + if err := prebuildTplHarness.Run(prebuildTplCtx); err != nil { + return xerrors.Errorf("run test harness to delete prebuilds templates (harness failure, not a test failure): %w", err) + } + + cliui.Infof(inv.Stdout, "Done deleting scaletest prebuilds templates:") + prebuildTplRes := prebuildTplHarness.Results() + prebuildTplRes.PrintText(inv.Stderr) + + if prebuildTplRes.TotalFail > 0 { + return xerrors.Errorf("failed to delete %d scaletest prebuilds template(s)", prebuildTplRes.TotalFail) + } + } + cliui.Infof(inv.Stdout, "Fetching scaletest workspaces...") workspaces, _, err := getScaletestWorkspaces(ctx, client, "", template) if err != nil { diff --git a/cli/exp_scaletest_prebuilds_internal_test.go b/cli/exp_scaletest_prebuilds_internal_test.go new file mode 100644 index 0000000000..fd3acfc5fc --- /dev/null +++ b/cli/exp_scaletest_prebuilds_internal_test.go @@ -0,0 +1,82 @@ +//go:build !slim + +package cli + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/scaletest/prebuilds" + "github.com/coder/coder/v2/testutil" +) + +func Test_getScaletestPrebuildsTemplates(t *testing.T) { + t.Parallel() + + client, _, _ := coderdtest.NewWithAPI(t, &coderdtest.Options{ + IncludeProvisionerDaemon: true, + }) + user := coderdtest.CreateFirstUser(t, client) + + makeTemplate := func(t *testing.T, name string) { + t.Helper() + version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil) + coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) + coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID, func(r *codersdk.CreateTemplateRequest) { + r.Name = name + }) + } + + // The real runner uses a small integer suffix (e.g. "0", "1"), keeping the + // total name within the 32-character limit enforced by NameValid. + const ( + scaletestPrebuildName = prebuilds.TemplatePrefix + "0" + prebuildNoScaletest = "prebuild-other" + scaletestNoPrebuild = "scaletest-other" + unrelatedTemplate = "unrelated-template" + ) + + makeTemplate(t, scaletestPrebuildName) + makeTemplate(t, prebuildNoScaletest) + makeTemplate(t, scaletestNoPrebuild) + makeTemplate(t, unrelatedTemplate) + + t.Run("NoFilter", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitMedium) + got, err := getScaletestPrebuildsTemplates(ctx, client, "") + require.NoError(t, err) + require.Len(t, got, 1) + assert.Equal(t, scaletestPrebuildName, got[0].Name) + }) + + t.Run("MatchingTemplate", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitMedium) + got, err := getScaletestPrebuildsTemplates(ctx, client, scaletestPrebuildName) + require.NoError(t, err) + require.Len(t, got, 1) + assert.Equal(t, scaletestPrebuildName, got[0].Name) + }) + + t.Run("NonExistentScaletestTemplate", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitMedium) + got, err := getScaletestPrebuildsTemplates(ctx, client, prebuilds.TemplatePrefix+"99") + require.NoError(t, err) + assert.Empty(t, got) + }) + + t.Run("NonScaletestTemplateReturnsError", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitMedium) + for _, name := range []string{prebuildNoScaletest, scaletestNoPrebuild, unrelatedTemplate} { + _, err := getScaletestPrebuildsTemplates(ctx, client, name) + require.Error(t, err, "expected error for template %q", name) + } + }) +} diff --git a/scaletest/prebuilds/run.go b/scaletest/prebuilds/run.go index 246dd5ee68..612f93e1fe 100644 --- a/scaletest/prebuilds/run.go +++ b/scaletest/prebuilds/run.go @@ -28,6 +28,10 @@ type Runner struct { template codersdk.Template } +// TemplatePrefix is the name prefix applied to all templates created by the +// scaletest prebuilds runner. +const TemplatePrefix = "scaletest-prebuilds-template-" + var ( _ harness.Runnable = &Runner{} _ harness.Cleanable = &Runner{} @@ -64,7 +68,7 @@ func (r *Runner) Run(ctx context.Context, id string, logs io.Writer) error { r.client.SetLogger(logger) r.client.SetLogBodies(true) - templateName := "scaletest-prebuilds-template-" + id + templateName := TemplatePrefix + id version, err := r.createTemplateVersion(ctx, uuid.Nil, r.cfg.NumPresets, r.cfg.NumPresetPrebuilds) if err != nil {