feat: add cleanup to task-status load test runner (#20799)

Implement Cleanup in the task status Runner, to delete the external workspaces created.

cli/exp_scaletest_taskstatus.go

@@ -206,6 +206,17 @@ After all runners connect, it waits for the baseline duration before triggering
 				}
 			}
 
+			cleanupCtx, cleanupCancel := cleanupStrategy.toContext(ctx)
+			defer cleanupCancel()
+			err = th.Cleanup(cleanupCtx)
+			if err != nil {
+				return xerrors.Errorf("cleanup tests: %w", err)
+			}
+
+			if res.TotalFail > 0 {
+				return xerrors.New("load test failed, see above for more details")
+			}
+
 			return nil
 		},
 	}

scaletest/taskstatus/client.go

@@ -34,6 +34,9 @@ type client interface {
 	// watchWorkspace watches for updates to a workspace.
 	watchWorkspace(ctx context.Context, workspaceID uuid.UUID) (<-chan codersdk.Workspace, error)
 
+	// deleteWorkspace deletes the workspace by creating a build with delete transition.
+	deleteWorkspace(ctx context.Context, workspaceID uuid.UUID) error
+
 	// initialize sets up the client with the provided logger, which is only available after Run() is called.
 	initialize(logger slog.Logger)
 }
@@ -101,6 +104,18 @@ func (c *sdkClient) watchWorkspace(ctx context.Context, workspaceID uuid.UUID) (
 	return c.coderClient.WatchWorkspace(ctx, workspaceID)
 }
 
+func (c *sdkClient) deleteWorkspace(ctx context.Context, workspaceID uuid.UUID) error {
+	// Create a build with delete transition to delete the workspace
+	_, err := c.coderClient.CreateWorkspaceBuild(ctx, workspaceID, codersdk.CreateWorkspaceBuildRequest{
+		Transition: codersdk.WorkspaceTransitionDelete,
+		Reason:     codersdk.CreateWorkspaceBuildReasonCLI,
+	})
+	if err != nil {
+		return xerrors.Errorf("create delete build: %w", err)
+	}
+	return nil
+}
+
 func (c *sdkClient) initialize(logger slog.Logger) {
 	// Configure the coder client logging
 	c.coderClient.SetLogger(logger)

scaletest/taskstatus/run.go

@@ -41,7 +41,10 @@ type Runner struct {
 	clock quartz.Clock
 }
 
-var _ harness.Runnable = &Runner{}
+var (
+	_ harness.Runnable  = &Runner{}
+	_ harness.Cleanable = &Runner{}
+)
 
 // NewRunner creates a new Runner with the provided codersdk.Client and configuration.
 func NewRunner(coderClient *codersdk.Client, cfg Config) *Runner {
@@ -111,6 +114,30 @@ func (r *Runner) Run(ctx context.Context, name string, logs io.Writer) error {
 	return nil
 }
 
+// Cleanup deletes the external workspace created by this runner.
+func (r *Runner) Cleanup(ctx context.Context, id string, logs io.Writer) error {
+	if r.workspaceID == uuid.Nil {
+		// No workspace was created, nothing to cleanup
+		return nil
+	}
+
+	logs = loadtestutil.NewSyncWriter(logs)
+	logger := slog.Make(sloghuman.Sink(logs)).Leveled(slog.LevelDebug).Named(id)
+
+	logger.Info(ctx, "deleting external workspace", slog.F("workspace_id", r.workspaceID))
+
+	err := r.client.deleteWorkspace(ctx, r.workspaceID)
+	if err != nil {
+		logger.Error(ctx, "failed to delete external workspace",
+			slog.F("workspace_id", r.workspaceID),
+			slog.Error(err))
+		return xerrors.Errorf("delete external workspace: %w", err)
+	}
+
+	logger.Info(ctx, "successfully deleted external workspace", slog.F("workspace_id", r.workspaceID))
+	return nil
+}
+
 func (r *Runner) watchWorkspaceUpdates(ctx context.Context) error {
 	shouldMarkConnectedDone := true
 	defer func() {

scaletest/taskstatus/run_internal_test.go

@@ -14,6 +14,7 @@ import (
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/sloghuman"
 	"github.com/coder/quartz"
 
 	"github.com/coder/coder/v2/codersdk"
@@ -57,6 +58,12 @@ func (m *fakeClient) createExternalWorkspace(ctx context.Context, req codersdk.C
 	}, nil
 }
 
+func (m *fakeClient) deleteWorkspace(ctx context.Context, workspaceID uuid.UUID) error {
+	m.logger.Debug(ctx, "called fake DeleteWorkspace", slog.F("workspace_id", workspaceID.String()))
+	// Simulate successful deletion in tests
+	return nil
+}
+
 // fakeAppStatusPatcher implements the appStatusPatcher interface for testing
 type fakeAppStatusPatcher struct {
 	t          *testing.T
@@ -480,3 +487,68 @@ func TestParseStatusMessage(t *testing.T) {
 		})
 	}
 }
+
+func TestRunner_Cleanup(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitMedium)
+
+	fakeClient := &fakeClientWithCleanupTracking{
+		fakeClient:           newFakeClient(t),
+		deleteWorkspaceCalls: make([]uuid.UUID, 0),
+	}
+	fakeClient.initialize(slog.Make(sloghuman.Sink(testutil.NewTestLogWriter(t))).Leveled(slog.LevelDebug))
+
+	cfg := Config{
+		AppSlug:              "test-app",
+		TemplateID:           uuid.UUID{5, 6, 7, 8},
+		WorkspaceName:        "test-workspace",
+		MetricLabelValues:    []string{"test"},
+		Metrics:              NewMetrics(prometheus.NewRegistry(), "test"),
+		ReportStatusPeriod:   100 * time.Millisecond,
+		ReportStatusDuration: 200 * time.Millisecond,
+		StartReporting:       make(chan struct{}),
+		ConnectedWaitGroup:   &sync.WaitGroup{},
+	}
+
+	runner := &Runner{
+		client:  fakeClient,
+		patcher: newFakeAppStatusPatcher(t),
+		cfg:     cfg,
+		clock:   quartz.NewMock(t),
+	}
+
+	logWriter := testutil.NewTestLogWriter(t)
+
+	// Case 1: No workspace created - Cleanup should do nothing
+	err := runner.Cleanup(ctx, "test-runner", logWriter)
+	require.NoError(t, err)
+	require.Len(t, fakeClient.deleteWorkspaceCalls, 0, "deleteWorkspace should not be called when no workspace was created")
+
+	// Case 2: Workspace created - Cleanup should delete it
+	runner.workspaceID = uuid.UUID{1, 2, 3, 4}
+	err = runner.Cleanup(ctx, "test-runner", logWriter)
+	require.NoError(t, err)
+	require.Len(t, fakeClient.deleteWorkspaceCalls, 1, "deleteWorkspace should be called once")
+	require.Equal(t, runner.workspaceID, fakeClient.deleteWorkspaceCalls[0], "deleteWorkspace should be called with correct workspace ID")
+
+	// Case 3: Cleanup with error
+	fakeClient.deleteError = xerrors.New("delete failed")
+	runner.workspaceID = uuid.UUID{5, 6, 7, 8}
+	err = runner.Cleanup(ctx, "test-runner", logWriter)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "delete external workspace")
+}
+
+// fakeClientWithCleanupTracking extends fakeClient to track deleteWorkspace calls
+type fakeClientWithCleanupTracking struct {
+	*fakeClient
+	deleteWorkspaceCalls []uuid.UUID
+	deleteError          error
+}
+
+func (c *fakeClientWithCleanupTracking) deleteWorkspace(ctx context.Context, workspaceID uuid.UUID) error {
+	c.deleteWorkspaceCalls = append(c.deleteWorkspaceCalls, workspaceID)
+	c.logger.Debug(ctx, "called fake DeleteWorkspace with tracking", slog.F("workspace_id", workspaceID.String()))
+	return c.deleteError
+}