Skip to content

Commit 97e4700

Browse files
f0sself0ssel
committed
chore: accept payload on workspace usage route
1 parent 5ccf508 commit 97e4700

File tree

5 files changed

+130
-16
lines changed

5 files changed

+130
-16
lines changed

‎cli/portforward.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ func (r *RootCmd) portForward() *serpent.Command {
137137
listeners[i] = l
138138
}
139139

140-
stopUpdating := client.UpdateWorkspaceUsageContext(ctx, workspace.ID)
140+
stopUpdating := client.UpdateWorkspaceUsageContext(ctx, workspace.ID, nil)
141141

142142
// Wait for the context to be canceled or for a signal and close
143143
// all listeners.

‎coderd/workspaces.go

+90
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"errors"
88
"fmt"
99
"net/http"
10+
"slices"
1011
"strconv"
1112
"time"
1213

@@ -15,6 +16,7 @@ import (
1516
"golang.org/x/xerrors"
1617

1718
"cdr.dev/slog"
19+
"github.com/coder/coder/v2/agent/proto"
1820
"github.com/coder/coder/v2/coderd/audit"
1921
"github.com/coder/coder/v2/coderd/database"
2022
"github.com/coder/coder/v2/coderd/database/db2sdk"
@@ -1106,6 +1108,7 @@ func (api *API) putExtendWorkspace(rw http.ResponseWriter, r *http.Request) {
11061108
// @Security CoderSessionToken
11071109
// @Tags Workspaces
11081110
// @Param workspace path string true "Workspace ID" format(uuid)
1111+
// @Param request body codersdk.PostWorkspaceUsageRequest false "Post workspace usage request"
11091112
// @Success 204
11101113
// @Router /workspaces/{workspace}/usage [post]
11111114
func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
@@ -1116,6 +1119,93 @@ func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
11161119
}
11171120

11181121
api.statsReporter.TrackUsage(workspace.ID)
1122+
1123+
if !api.Experiments.Enabled(codersdk.ExperimentWorkspaceUsage) {
1124+
// Continue previous behavior if the experiment is not enabled.
1125+
rw.WriteHeader(http.StatusNoContent)
1126+
return
1127+
}
1128+
1129+
ctx := r.Context()
1130+
var req codersdk.PostWorkspaceUsageRequest
1131+
if !httpapi.Read(ctx, rw, r, &req) {
1132+
return
1133+
}
1134+
1135+
if req.AgentID == uuid.Nil && req.AppName == "" {
1136+
// Continue previous behavior if body is empty.
1137+
rw.WriteHeader(http.StatusNoContent)
1138+
return
1139+
}
1140+
if req.AgentID == uuid.Nil {
1141+
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
1142+
Message: "Invalid request",
1143+
Validations: []codersdk.ValidationError{{
1144+
Field: "agent_id",
1145+
Detail: "must be set when app_name is set",
1146+
}},
1147+
})
1148+
return
1149+
}
1150+
if req.AppName == "" {
1151+
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
1152+
Message: "Invalid request",
1153+
Validations: []codersdk.ValidationError{{
1154+
Field: "app_name",
1155+
Detail: "must be set when agent_id is set",
1156+
}},
1157+
})
1158+
return
1159+
}
1160+
if !slices.Contains(codersdk.AllowedAppNames, req.AppName) {
1161+
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
1162+
Message: "Invalid request",
1163+
Validations: []codersdk.ValidationError{{
1164+
Field: "app_name",
1165+
Detail: fmt.Sprintf("must be one of %v", codersdk.AllowedAppNames),
1166+
}},
1167+
})
1168+
return
1169+
}
1170+
1171+
stat := &proto.Stats{}
1172+
switch req.AppName {
1173+
case codersdk.UsageAppNameVscode:
1174+
stat.SessionCountVscode = 1
1175+
case codersdk.UsageAppNameJetbrains:
1176+
stat.SessionCountJetbrains = 1
1177+
case codersdk.UsageAppNameReconnectingPty:
1178+
stat.SessionCountReconnectingPty = 1
1179+
case codersdk.UsageAppNameSsh:
1180+
stat.SessionCountSsh = 1
1181+
default:
1182+
// This means the app_name is not in the list of allowed app names.
1183+
httpapi.InternalServerError(rw, xerrors.Errorf("unknown app_name %q", req.AppName))
1184+
return
1185+
}
1186+
1187+
agent, err := api.Database.GetWorkspaceAgentByID(ctx, req.AgentID)
1188+
if err != nil {
1189+
if httpapi.Is404Error(err) {
1190+
httpapi.ResourceNotFound(rw)
1191+
return
1192+
}
1193+
httpapi.InternalServerError(rw, err)
1194+
return
1195+
}
1196+
1197+
template, err := api.Database.GetTemplateByID(ctx, workspace.TemplateID)
1198+
if err != nil {
1199+
httpapi.InternalServerError(rw, err)
1200+
return
1201+
}
1202+
1203+
err = api.statsReporter.ReportAgentStats(ctx, dbtime.Now(), workspace, agent, template.Name, stat)
1204+
if err != nil {
1205+
httpapi.InternalServerError(rw, err)
1206+
return
1207+
}
1208+
11191209
rw.WriteHeader(http.StatusNoContent)
11201210
}
11211211

‎coderd/workspacestats/tracker_test.go

+10-10
Original file line numberDiff line numberDiff line change
@@ -158,18 +158,18 @@ func TestTracker_MultipleInstances(t *testing.T) {
158158
}
159159

160160
// Use client A to update LastUsedAt of the first three
161-
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[0].Workspace.ID))
162-
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[1].Workspace.ID))
163-
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[2].Workspace.ID))
161+
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[0].Workspace.ID, nil))
162+
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[1].Workspace.ID, nil))
163+
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[2].Workspace.ID, nil))
164164
// Use client B to update LastUsedAt of the next three
165-
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[3].Workspace.ID))
166-
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[4].Workspace.ID))
167-
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[5].Workspace.ID))
165+
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[3].Workspace.ID, nil))
166+
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[4].Workspace.ID, nil))
167+
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[5].Workspace.ID, nil))
168168
// The next two will have updated from both instances
169-
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[6].Workspace.ID))
170-
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[6].Workspace.ID))
171-
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[7].Workspace.ID))
172-
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[7].Workspace.ID))
169+
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[6].Workspace.ID, nil))
170+
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[6].Workspace.ID, nil))
171+
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[7].Workspace.ID, nil))
172+
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[7].Workspace.ID, nil))
173173
// The last two will not report any usage.
174174

175175
// Tick both with different times and wait for both flushes to complete

‎codersdk/deployment.go

+1
Original file line numberDiff line numberDiff line change
@@ -2233,6 +2233,7 @@ const (
22332233
ExperimentAutoFillParameters Experiment = "auto-fill-parameters" // This should not be taken out of experiments until we have redesigned the feature.
22342234
ExperimentMultiOrganization Experiment = "multi-organization" // Requires organization context for interactions, default org is assumed.
22352235
ExperimentCustomRoles Experiment = "custom-roles" // Allows creating runtime custom roles
2236+
ExperimentWorkspaceUsage Experiment = "workspace-usage" // Enables the new workspace usage tracking
22362237
)
22372238

22382239
// ExperimentsAll should include all experiments that are safe for

‎codersdk/workspaces.go

+28-5
Original file line numberDiff line numberDiff line change
@@ -316,10 +316,31 @@ func (c *Client) PutExtendWorkspace(ctx context.Context, id uuid.UUID, req PutEx
316316
return nil
317317
}
318318

319+
type PostWorkspaceUsageRequest struct {
320+
AgentID uuid.UUID `json:"agent_id"`
321+
AppName UsageAppName `json:"app_name"`
322+
}
323+
324+
type UsageAppName string
325+
326+
const (
327+
UsageAppNameVscode UsageAppName = "vscode"
328+
UsageAppNameJetbrains UsageAppName = "jetbrains"
329+
UsageAppNameReconnectingPty UsageAppName = "reconnecting-pty"
330+
UsageAppNameSsh UsageAppName = "ssh"
331+
)
332+
333+
var AllowedAppNames = []UsageAppName{
334+
UsageAppNameVscode,
335+
UsageAppNameJetbrains,
336+
UsageAppNameReconnectingPty,
337+
UsageAppNameSsh,
338+
}
339+
319340
// PostWorkspaceUsage marks the workspace as having been used recently.
320-
func (c *Client) PostWorkspaceUsage(ctx context.Context, id uuid.UUID) error {
341+
func (c *Client) PostWorkspaceUsage(ctx context.Context, id uuid.UUID, req *PostWorkspaceUsageRequest) error {
321342
path := fmt.Sprintf("/api/v2/workspaces/%s/usage", id.String())
322-
res, err := c.Request(ctx, http.MethodPost, path, nil)
343+
res, err := c.Request(ctx, http.MethodPost, path, req)
323344
if err != nil {
324345
return xerrors.Errorf("post workspace usage: %w", err)
325346
}
@@ -334,10 +355,11 @@ func (c *Client) PostWorkspaceUsage(ctx context.Context, id uuid.UUID) error {
334355
// with the given id in the background.
335356
// The caller is responsible for calling the returned function to stop the background
336357
// process.
337-
func (c *Client) UpdateWorkspaceUsageContext(ctx context.Context, id uuid.UUID) func() {
358+
func (c *Client) UpdateWorkspaceUsageContext(ctx context.Context, workspaceID uuid.UUID, req *PostWorkspaceUsageRequest) func() {
338359
hbCtx, hbCancel := context.WithCancel(ctx)
339360
// Perform one initial update
340-
if err := c.PostWorkspaceUsage(hbCtx, id); err != nil {
361+
err := c.PostWorkspaceUsage(hbCtx, workspaceID, req)
362+
if err != nil {
341363
c.logger.Warn(ctx, "failed to post workspace usage", slog.Error(err))
342364
}
343365
ticker := time.NewTicker(time.Minute)
@@ -350,7 +372,8 @@ func (c *Client) UpdateWorkspaceUsageContext(ctx context.Context, id uuid.UUID)
350372
for {
351373
select {
352374
case <-ticker.C:
353-
if err := c.PostWorkspaceUsage(hbCtx, id); err != nil {
375+
err := c.PostWorkspaceUsage(hbCtx, workspaceID, req)
376+
if err != nil {
354377
c.logger.Warn(ctx, "failed to post workspace usage in background", slog.Error(err))
355378
}
356379
case <-hbCtx.Done():

0 commit comments

Comments
 (0)