Skip to content

Commit da5eec8

Browse files
spikecurtisspikecurtis
committed
fix: wait for build in task status load generator
1 parent 0bbb7dd commit da5eec8

File tree

3 files changed

+283
-61
lines changed

3 files changed

+283
-61
lines changed

‎scaletest/taskstatus/client.go‎

Lines changed: 21 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,9 @@ import (
1111
"cdr.dev/slog"
1212
"github.com/coder/coder/v2/codersdk"
1313
"github.com/coder/coder/v2/codersdk/agentsdk"
14+
"github.com/coder/quartz"
1415
)
1516

16-
// createExternalWorkspaceResult contains the results from creating an external workspace.
17-
type createExternalWorkspaceResult struct {
18-
WorkspaceID uuid.UUID
19-
AgentToken string
20-
}
21-
2217
// client abstracts the details of using codersdk.Client for workspace operations.
2318
// This interface allows for easier testing by enabling mock implementations and
2419
// provides a cleaner separation of concerns.
@@ -27,9 +22,14 @@ type createExternalWorkspaceResult struct {
2722
// 1. Create the client with newClient(coderClient)
2823
// 2. Configure logging when the io.Writer is available in Run()
2924
type client interface {
30-
// createExternalWorkspace creates an external workspace and returns the workspace ID
31-
// and agent token for the first external agent found in the workspace resources.
32-
createExternalWorkspace(ctx context.Context, req codersdk.CreateWorkspaceRequest) (createExternalWorkspaceResult, error)
25+
// CreateUserWorkspace creates a workspace for a user.
26+
CreateUserWorkspace(ctx context.Context, userID string, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error)
27+
28+
// WorkspaceByOwnerAndName retrieves a workspace by owner and name.
29+
WorkspaceByOwnerAndName(ctx context.Context, owner string, name string, params codersdk.WorkspaceOptions) (codersdk.Workspace, error)
30+
31+
// WorkspaceExternalAgentCredentials retrieves credentials for an external agent.
32+
WorkspaceExternalAgentCredentials(ctx context.Context, workspaceID uuid.UUID, agentName string) (codersdk.ExternalAgentCredentials, error)
3333

3434
// watchWorkspace watches for updates to a workspace.
3535
watchWorkspace(ctx context.Context, workspaceID uuid.UUID) (<-chan codersdk.Workspace, error)
@@ -56,48 +56,28 @@ type appStatusPatcher interface {
5656
// codersdk.Client.
5757
type sdkClient struct {
5858
coderClient *codersdk.Client
59+
clock quartz.Clock
60+
logger slog.Logger
5961
}
6062

6163
// newClient creates a new client implementation using the provided codersdk.Client.
6264
func newClient(coderClient *codersdk.Client) client {
6365
return &sdkClient{
6466
coderClient: coderClient,
67+
clock: quartz.NewReal(),
6568
}
6669
}
6770

68-
func (c *sdkClient) createExternalWorkspace(ctx context.Context, req codersdk.CreateWorkspaceRequest) (createExternalWorkspaceResult, error) {
69-
// Create the workspace
70-
workspace, err := c.coderClient.CreateUserWorkspace(ctx, codersdk.Me, req)
71-
if err != nil {
72-
return createExternalWorkspaceResult{}, err
73-
}
74-
75-
// Get the workspace with latest build details
76-
workspace, err = c.coderClient.WorkspaceByOwnerAndName(ctx, codersdk.Me, workspace.Name, codersdk.WorkspaceOptions{})
77-
if err != nil {
78-
return createExternalWorkspaceResult{}, err
79-
}
71+
func (c *sdkClient) CreateUserWorkspace(ctx context.Context, userID string, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
72+
return c.coderClient.CreateUserWorkspace(ctx, userID, req)
73+
}
8074

81-
// Find external agents in resources
82-
for _, resource := range workspace.LatestBuild.Resources {
83-
if resource.Type != "coder_external_agent" || len(resource.Agents) == 0 {
84-
continue
85-
}
86-
87-
// Get credentials for the first agent
88-
agent := resource.Agents[0]
89-
credentials, err := c.coderClient.WorkspaceExternalAgentCredentials(ctx, workspace.ID, agent.Name)
90-
if err != nil {
91-
return createExternalWorkspaceResult{}, err
92-
}
93-
94-
return createExternalWorkspaceResult{
95-
WorkspaceID: workspace.ID,
96-
AgentToken: credentials.AgentToken,
97-
}, nil
98-
}
75+
func (c *sdkClient) WorkspaceByOwnerAndName(ctx context.Context, owner string, name string, params codersdk.WorkspaceOptions) (codersdk.Workspace, error) {
76+
return c.coderClient.WorkspaceByOwnerAndName(ctx, owner, name, params)
77+
}
9978

100-
return createExternalWorkspaceResult{}, xerrors.Errorf("no external agent found in workspace")
79+
func (c *sdkClient) WorkspaceExternalAgentCredentials(ctx context.Context, workspaceID uuid.UUID, agentName string) (codersdk.ExternalAgentCredentials, error) {
80+
return c.coderClient.WorkspaceExternalAgentCredentials(ctx, workspaceID, agentName)
10181
}
10282

10383
func (c *sdkClient) watchWorkspace(ctx context.Context, workspaceID uuid.UUID) (<-chan codersdk.Workspace, error) {
@@ -118,6 +98,7 @@ func (c *sdkClient) deleteWorkspace(ctx context.Context, workspaceID uuid.UUID)
11898

11999
func (c *sdkClient) initialize(logger slog.Logger) {
120100
// Configure the coder client logging
101+
c.logger = logger
121102
c.coderClient.SetLogger(logger)
122103
c.coderClient.SetLogBodies(true)
123104
}

‎scaletest/taskstatus/run.go‎

Lines changed: 88 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,12 @@ import (
2323

2424
const statusUpdatePrefix = "scaletest status update:"
2525

26+
// createExternalWorkspaceResult contains the results from creating an external workspace.
27+
type createExternalWorkspaceResult struct {
28+
workspaceID uuid.UUID
29+
agentToken string
30+
}
31+
2632
type Runner struct {
2733
client client
2834
patcher appStatusPatcher
@@ -65,6 +71,10 @@ func (r *Runner) Run(ctx context.Context, name string, logs io.Writer) error {
6571
}
6672
}()
6773

74+
// ensure these labels are initialized, so we see the time series right away in prometheus.
75+
r.cfg.Metrics.MissingStatusUpdatesTotal.WithLabelValues(r.cfg.MetricLabelValues...).Add(0)
76+
r.cfg.Metrics.ReportTaskStatusErrorsTotal.WithLabelValues(r.cfg.MetricLabelValues...).Add(0)
77+
6878
logs = loadtestutil.NewSyncWriter(logs)
6979
r.logger = slog.Make(sloghuman.Sink(logs)).Leveled(slog.LevelDebug).Named(name)
7080
r.client.initialize(r.logger)
@@ -74,26 +84,23 @@ func (r *Runner) Run(ctx context.Context, name string, logs io.Writer) error {
7484
slog.F("template_id", r.cfg.TemplateID),
7585
slog.F("workspace_name", r.cfg.WorkspaceName))
7686

77-
result, err := r.client.createExternalWorkspace(ctx, codersdk.CreateWorkspaceRequest{
87+
result, err := r.createExternalWorkspace(ctx, codersdk.CreateWorkspaceRequest{
7888
TemplateID: r.cfg.TemplateID,
7989
Name: r.cfg.WorkspaceName,
8090
})
8191
if err != nil {
92+
r.cfg.Metrics.ReportTaskStatusErrorsTotal.WithLabelValues(r.cfg.MetricLabelValues...).Inc()
8293
return xerrors.Errorf("create external workspace: %w", err)
8394
}
8495

8596
// Set the workspace ID
86-
r.workspaceID = result.WorkspaceID
97+
r.workspaceID = result.workspaceID
8798
r.logger.Info(ctx, "created external workspace", slog.F("workspace_id", r.workspaceID))
8899

89100
// Initialize the patcher with the agent token
90-
r.patcher.initialize(r.logger, result.AgentToken)
101+
r.patcher.initialize(r.logger, result.agentToken)
91102
r.logger.Info(ctx, "initialized app status patcher with agent token")
92103

93-
// ensure these labels are initialized, so we see the time series right away in prometheus.
94-
r.cfg.Metrics.MissingStatusUpdatesTotal.WithLabelValues(r.cfg.MetricLabelValues...).Add(0)
95-
r.cfg.Metrics.ReportTaskStatusErrorsTotal.WithLabelValues(r.cfg.MetricLabelValues...).Add(0)
96-
97104
workspaceUpdatesCtx, cancelWorkspaceUpdates := context.WithCancel(ctx)
98105
defer cancelWorkspaceUpdates()
99106
workspaceUpdatesResult := make(chan error, 1)
@@ -257,3 +264,77 @@ func parseStatusMessage(message string) (int, bool) {
257264
}
258265
return msgNo, true
259266
}
267+
268+
// createExternalWorkspace creates an external workspace and returns the workspace ID
269+
// and agent token for the first external agent found in the workspace resources.
270+
func (r *Runner) createExternalWorkspace(ctx context.Context, req codersdk.CreateWorkspaceRequest) (createExternalWorkspaceResult, error) {
271+
// Create the workspace
272+
workspace, err := r.client.CreateUserWorkspace(ctx, codersdk.Me, req)
273+
if err != nil {
274+
return createExternalWorkspaceResult{}, err
275+
}
276+
277+
r.logger.Info(ctx, "waiting for workspace build to complete",
278+
slog.F("workspace_name", workspace.Name),
279+
slog.F("workspace_id", workspace.ID))
280+
281+
// Poll the workspace until the build is complete
282+
var finalWorkspace codersdk.Workspace
283+
buildComplete := xerrors.New("build complete") // sentinel error
284+
waiter := r.clock.TickerFunc(ctx, 30*time.Second, func() error {
285+
// Get the workspace with latest build details
286+
workspace, err := r.client.WorkspaceByOwnerAndName(ctx, codersdk.Me, workspace.Name, codersdk.WorkspaceOptions{})
287+
if err != nil {
288+
r.logger.Error(ctx, "failed to poll workspace while waiting for build to complete", slog.Error(err))
289+
return nil
290+
}
291+
292+
jobStatus := workspace.LatestBuild.Job.Status
293+
r.logger.Debug(ctx, "checking workspace build status",
294+
slog.F("status", jobStatus),
295+
slog.F("build_id", workspace.LatestBuild.ID))
296+
297+
switch jobStatus {
298+
case codersdk.ProvisionerJobSucceeded:
299+
// Build succeeded
300+
r.logger.Info(ctx, "workspace build succeeded")
301+
finalWorkspace = workspace
302+
return buildComplete
303+
case codersdk.ProvisionerJobFailed:
304+
return xerrors.Errorf("workspace build failed: %s", workspace.LatestBuild.Job.Error)
305+
case codersdk.ProvisionerJobCanceled:
306+
return xerrors.Errorf("workspace build was canceled")
307+
case codersdk.ProvisionerJobPending, codersdk.ProvisionerJobRunning, codersdk.ProvisionerJobCanceling:
308+
// Still in progress, continue polling
309+
return nil
310+
default:
311+
return xerrors.Errorf("unexpected job status: %s", jobStatus)
312+
}
313+
}, "createExternalWorkspace")
314+
315+
err = waiter.Wait()
316+
if err != nil && !xerrors.Is(err, buildComplete) {
317+
return createExternalWorkspaceResult{}, xerrors.Errorf("wait for build completion: %w", err)
318+
}
319+
320+
// Find external agents in resources
321+
for _, resource := range finalWorkspace.LatestBuild.Resources {
322+
if resource.Type != "coder_external_agent" || len(resource.Agents) == 0 {
323+
continue
324+
}
325+
326+
// Get credentials for the first agent
327+
agent := resource.Agents[0]
328+
credentials, err := r.client.WorkspaceExternalAgentCredentials(ctx, finalWorkspace.ID, agent.Name)
329+
if err != nil {
330+
return createExternalWorkspaceResult{}, err
331+
}
332+
333+
return createExternalWorkspaceResult{
334+
workspaceID: finalWorkspace.ID,
335+
agentToken: credentials.AgentToken,
336+
}, nil
337+
}
338+
339+
return createExternalWorkspaceResult{}, xerrors.Errorf("no external agent found in workspace")
340+
}

0 commit comments

Comments
 (0)