fix(cli): extend exp scaletest cleanup to properly clean up prebuilds (#23628)

Signed-off-by: Callum Styan <callumstyan@gmail.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-02 20:48:20 +00:00 · 2026-04-29 10:28:27 -07:00
parent 54d650ea79
commit 950660e392
3 changed files with 249 additions and 1 deletions
@@ -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 {
@@ -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)
+		}
+	})
+}
@@ -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 {