Skip to content

Commit a3c851c

Browse files
spikecurtisspikecurtis
authored
feat: add task status reporting load generator runner (#20538)
Adds the Runner, Config, and Metrics for the scaletest load generator for task status. Part of coder/internal#913
1 parent 5bfbb03 commit a3c851c

File tree

6 files changed

+1001
-0
lines changed

6 files changed

+1001
-0
lines changed

‎scaletest/taskstatus/client.go‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
package taskstatus
2+
3+
import (
4+
"context"
5+
"net/http"
6+
"net/url"
7+
8+
"github.com/google/uuid"
9+
"golang.org/x/xerrors"
10+
11+
"cdr.dev/slog"
12+
"github.com/coder/coder/v2/codersdk"
13+
"github.com/coder/coder/v2/codersdk/agentsdk"
14+
)
15+
16+
// createExternalWorkspaceResult contains the results from creating an external workspace.
17+
type createExternalWorkspaceResult struct {
18+
WorkspaceID uuid.UUID
19+
AgentToken string
20+
}
21+
22+
// client abstracts the details of using codersdk.Client for workspace operations.
23+
// This interface allows for easier testing by enabling mock implementations and
24+
// provides a cleaner separation of concerns.
25+
//
26+
// The interface is designed to be initialized in two phases:
27+
// 1. Create the client with newClient(coderClient)
28+
// 2. Configure logging when the io.Writer is available in Run()
29+
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)
33+
34+
// watchWorkspace watches for updates to a workspace.
35+
watchWorkspace(ctx context.Context, workspaceID uuid.UUID) (<-chan codersdk.Workspace, error)
36+
37+
// initialize sets up the client with the provided logger, which is only available after Run() is called.
38+
initialize(logger slog.Logger)
39+
}
40+
41+
// appStatusPatcher abstracts the details of using agentsdk.Client for updating app status.
42+
// This interface is separate from client because it requires an agent token which is only
43+
// available after creating an external workspace.
44+
type appStatusPatcher interface {
45+
// patchAppStatus updates the status of a workspace app.
46+
patchAppStatus(ctx context.Context, req agentsdk.PatchAppStatus) error
47+
48+
// initialize sets up the patcher with the provided logger and agent token.
49+
initialize(logger slog.Logger, agentToken string)
50+
}
51+
52+
// sdkClient is the concrete implementation of the client interface using
53+
// codersdk.Client.
54+
type sdkClient struct {
55+
coderClient *codersdk.Client
56+
}
57+
58+
// newClient creates a new client implementation using the provided codersdk.Client.
59+
func newClient(coderClient *codersdk.Client) client {
60+
return &sdkClient{
61+
coderClient: coderClient,
62+
}
63+
}
64+
65+
func (c *sdkClient) createExternalWorkspace(ctx context.Context, req codersdk.CreateWorkspaceRequest) (createExternalWorkspaceResult, error) {
66+
// Create the workspace
67+
workspace, err := c.coderClient.CreateUserWorkspace(ctx, codersdk.Me, req)
68+
if err != nil {
69+
return createExternalWorkspaceResult{}, err
70+
}
71+
72+
// Get the workspace with latest build details
73+
workspace, err = c.coderClient.WorkspaceByOwnerAndName(ctx, codersdk.Me, workspace.Name, codersdk.WorkspaceOptions{})
74+
if err != nil {
75+
return createExternalWorkspaceResult{}, err
76+
}
77+
78+
// Find external agents in resources
79+
for _, resource := range workspace.LatestBuild.Resources {
80+
if resource.Type != "coder_external_agent" || len(resource.Agents) == 0 {
81+
continue
82+
}
83+
84+
// Get credentials for the first agent
85+
agent := resource.Agents[0]
86+
credentials, err := c.coderClient.WorkspaceExternalAgentCredentials(ctx, workspace.ID, agent.Name)
87+
if err != nil {
88+
return createExternalWorkspaceResult{}, err
89+
}
90+
91+
return createExternalWorkspaceResult{
92+
WorkspaceID: workspace.ID,
93+
AgentToken: credentials.AgentToken,
94+
}, nil
95+
}
96+
97+
return createExternalWorkspaceResult{}, xerrors.Errorf("no external agent found in workspace")
98+
}
99+
100+
func (c *sdkClient) watchWorkspace(ctx context.Context, workspaceID uuid.UUID) (<-chan codersdk.Workspace, error) {
101+
return c.coderClient.WatchWorkspace(ctx, workspaceID)
102+
}
103+
104+
func (c *sdkClient) initialize(logger slog.Logger) {
105+
// Configure the coder client logging
106+
c.coderClient.SetLogger(logger)
107+
c.coderClient.SetLogBodies(true)
108+
}
109+
110+
// sdkAppStatusPatcher is the concrete implementation of the appStatusPatcher interface
111+
// using agentsdk.Client.
112+
type sdkAppStatusPatcher struct {
113+
agentClient *agentsdk.Client
114+
url *url.URL
115+
httpClient *http.Client
116+
}
117+
118+
// newAppStatusPatcher creates a new appStatusPatcher implementation.
119+
func newAppStatusPatcher(client *codersdk.Client) appStatusPatcher {
120+
return &sdkAppStatusPatcher{
121+
url: client.URL,
122+
httpClient: client.HTTPClient,
123+
}
124+
}
125+
126+
func (p *sdkAppStatusPatcher) patchAppStatus(ctx context.Context, req agentsdk.PatchAppStatus) error {
127+
if p.agentClient == nil {
128+
panic("agentClient not initialized - call initialize first")
129+
}
130+
return p.agentClient.PatchAppStatus(ctx, req)
131+
}
132+
133+
func (p *sdkAppStatusPatcher) initialize(logger slog.Logger, agentToken string) {
134+
// Create and configure the agent client with the provided token
135+
p.agentClient = agentsdk.New(
136+
p.url,
137+
agentsdk.WithFixedToken(agentToken),
138+
codersdk.WithHTTPClient(p.httpClient),
139+
codersdk.WithLogger(logger),
140+
codersdk.WithLogBodies(),
141+
)
142+
}
143+
144+
// Ensure sdkClient implements the client interface.
145+
var _ client = (*sdkClient)(nil)
146+
147+
// Ensure sdkAppStatusPatcher implements the appStatusPatcher interface.
148+
var _ appStatusPatcher = (*sdkAppStatusPatcher)(nil)

‎scaletest/taskstatus/config.go‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
package taskstatus
2+
3+
import (
4+
"sync"
5+
"time"
6+
7+
"github.com/google/uuid"
8+
"golang.org/x/xerrors"
9+
)
10+
11+
type Config struct {
12+
// TemplateID is the template ID to use for creating the external workspace.
13+
TemplateID uuid.UUID `json:"template_id"`
14+
15+
// WorkspaceName is the name for the external workspace to create.
16+
WorkspaceName string `json:"workspace_name"`
17+
18+
// AppSlug is the slug of the app designated as the AI Agent.
19+
AppSlug string `json:"app_slug"`
20+
21+
// When the runner has connected to the watch-ws endpoint, it will call Done once on this wait group. Used to
22+
// coordinate multiple runners from the higher layer.
23+
ConnectedWaitGroup *sync.WaitGroup `json:"-"`
24+
25+
// We read on this channel before starting to report task statuses. Used to coordinate multiple runners from the
26+
// higher layer.
27+
StartReporting chan struct{} `json:"-"`
28+
29+
// Time between reporting task statuses.
30+
ReportStatusPeriod time.Duration `json:"report_status_period"`
31+
32+
// Total time to report task statuses, starting from when we successfully read from the StartReporting channel.
33+
ReportStatusDuration time.Duration `json:"report_status_duration"`
34+
35+
Metrics *Metrics `json:"-"`
36+
MetricLabelValues []string `json:"metric_label_values"`
37+
}
38+
39+
func (c *Config) Validate() error {
40+
if c.TemplateID == uuid.Nil {
41+
return xerrors.Errorf("validate template_id: must not be nil")
42+
}
43+
44+
if c.WorkspaceName == "" {
45+
return xerrors.Errorf("validate workspace_name: must not be empty")
46+
}
47+
48+
if c.AppSlug == "" {
49+
return xerrors.Errorf("validate app_slug: must not be empty")
50+
}
51+
52+
if c.ConnectedWaitGroup == nil {
53+
return xerrors.Errorf("validate connected_wait_group: must not be nil")
54+
}
55+
56+
if c.StartReporting == nil {
57+
return xerrors.Errorf("validate start_reporting: must not be nil")
58+
}
59+
60+
if c.ReportStatusPeriod <= 0 {
61+
return xerrors.Errorf("validate report_status_period: must be greater than zero")
62+
}
63+
64+
if c.ReportStatusDuration <= 0 {
65+
return xerrors.Errorf("validate report_status_duration: must be greater than zero")
66+
}
67+
68+
if c.Metrics == nil {
69+
return xerrors.Errorf("validate metrics: must not be nil")
70+
}
71+
72+
return nil
73+
}

‎scaletest/taskstatus/metrics.go‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package taskstatus
2+
3+
import "github.com/prometheus/client_golang/prometheus"
4+
5+
type Metrics struct {
6+
TaskStatusToWorkspaceUpdateLatencySeconds prometheus.HistogramVec
7+
MissingStatusUpdatesTotal prometheus.CounterVec
8+
ReportTaskStatusErrorsTotal prometheus.CounterVec
9+
}
10+
11+
func NewMetrics(reg prometheus.Registerer, labelNames ...string) *Metrics {
12+
m := &Metrics{
13+
TaskStatusToWorkspaceUpdateLatencySeconds: *prometheus.NewHistogramVec(prometheus.HistogramOpts{
14+
Namespace: "coderd",
15+
Subsystem: "scaletest",
16+
Name: "task_status_to_workspace_update_latency_seconds",
17+
Help: "Time in seconds between reporting a task status and receiving the workspace update.",
18+
}, labelNames),
19+
MissingStatusUpdatesTotal: *prometheus.NewCounterVec(prometheus.CounterOpts{
20+
Namespace: "coderd",
21+
Subsystem: "scaletest",
22+
Name: "missing_status_updates_total",
23+
Help: "Total number of missing status updates.",
24+
}, labelNames),
25+
ReportTaskStatusErrorsTotal: *prometheus.NewCounterVec(prometheus.CounterOpts{
26+
Namespace: "coderd",
27+
Subsystem: "scaletest",
28+
Name: "report_task_status_errors_total",
29+
Help: "Total number of errors when reporting task status.",
30+
}, labelNames),
31+
}
32+
reg.MustRegister(m.TaskStatusToWorkspaceUpdateLatencySeconds)
33+
reg.MustRegister(m.MissingStatusUpdatesTotal)
34+
reg.MustRegister(m.ReportTaskStatusErrorsTotal)
35+
return m
36+
}

0 commit comments

Comments
 (0)