Skip to content

Commit 3dbe206

Browse files
johnstcnjohnstcn
committed
add workspace usager tracking to coderd, add endpoint
1 parent 05c51a1 commit 3dbe206

File tree

3 files changed

+44
-0
lines changed

3 files changed

+44
-0
lines changed

‎coderd/coderd.go

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ import (
6666
"github.com/coder/coder/v2/coderd/updatecheck"
6767
"github.com/coder/coder/v2/coderd/util/slice"
6868
"github.com/coder/coder/v2/coderd/workspaceapps"
69+
"github.com/coder/coder/v2/coderd/workspaceusage"
6970
"github.com/coder/coder/v2/codersdk"
7071
"github.com/coder/coder/v2/codersdk/drpc"
7172
"github.com/coder/coder/v2/provisionerd/proto"
@@ -190,6 +191,9 @@ type Options struct {
190191

191192
// NewTicker is used for unit tests to replace "time.NewTicker".
192193
NewTicker func(duration time.Duration) (tick <-chan time.Time, done func())
194+
195+
// WorkspaceUsageTracker tracks workspace usage by the CLI.
196+
WorkspaceUsageTracker *workspaceusage.Tracker
193197
}
194198

195199
// @title Coder API
@@ -362,6 +366,14 @@ func New(options *Options) *API {
362366
OIDC: options.OIDCConfig,
363367
}
364368

369+
if options.WorkspaceUsageTracker == nil {
370+
options.WorkspaceUsageTracker = workspaceusage.New(options.Database,
371+
workspaceusage.WithFlushInterval(workspaceusage.DefaultFlushInterval),
372+
workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
373+
)
374+
}
375+
go options.WorkspaceUsageTracker.Loop()
376+
365377
ctx, cancel := context.WithCancel(context.Background())
366378
r := chi.NewRouter()
367379

@@ -405,6 +417,7 @@ func New(options *Options) *API {
405417
options.Logger.Named("acquirer"),
406418
options.Database,
407419
options.Pubsub),
420+
workspaceUsageTracker: options.WorkspaceUsageTracker,
408421
}
409422

410423
api.AppearanceFetcher.Store(&appearance.DefaultFetcher)
@@ -972,6 +985,7 @@ func New(options *Options) *API {
972985
})
973986
r.Get("/watch", api.watchWorkspace)
974987
r.Put("/extend", api.putExtendWorkspace)
988+
r.Post("/usage", api.postWorkspaceUsage)
975989
r.Put("/dormant", api.putWorkspaceDormant)
976990
r.Put("/favorite", api.putFavoriteWorkspace)
977991
r.Delete("/favorite", api.deleteFavoriteWorkspace)
@@ -1179,6 +1193,8 @@ type API struct {
11791193
statsBatcher *batchstats.Batcher
11801194

11811195
Acquirer *provisionerdserver.Acquirer
1196+
1197+
workspaceUsageTracker *workspaceusage.Tracker
11821198
}
11831199

11841200
// Close waits for all WebSocket connections to drain before returning.
@@ -1200,6 +1216,7 @@ func (api *API) Close() error {
12001216
_ = (*coordinator).Close()
12011217
}
12021218
_ = api.agentProvider.Close()
1219+
api.workspaceUsageTracker.Close()
12031220
return nil
12041221
}
12051222

‎coderd/coderdtest/coderdtest.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ import (
7070
"github.com/coder/coder/v2/coderd/util/ptr"
7171
"github.com/coder/coder/v2/coderd/workspaceapps"
7272
"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
73+
"github.com/coder/coder/v2/coderd/workspaceusage"
7374
"github.com/coder/coder/v2/codersdk"
7475
"github.com/coder/coder/v2/codersdk/agentsdk"
7576
"github.com/coder/coder/v2/codersdk/drpc"
@@ -146,6 +147,7 @@ type Options struct {
146147
WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
147148
AllowWorkspaceRenames bool
148149
NewTicker func(duration time.Duration) (<-chan time.Time, func())
150+
WorkspaceUsageTracker *workspaceusage.Tracker
149151
}
150152

151153
// New constructs a codersdk client connected to an in-memory API instance.
@@ -306,6 +308,12 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
306308
hangDetector.Start()
307309
t.Cleanup(hangDetector.Close)
308310

311+
if options.WorkspaceUsageTracker == nil {
312+
// Workspace usage tracking must be triggered manually in tests.
313+
options.WorkspaceUsageTracker = workspaceusage.New(options.Database)
314+
}
315+
t.Cleanup(options.WorkspaceUsageTracker.Close)
316+
309317
var mutex sync.RWMutex
310318
var handler http.Handler
311319
srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -454,6 +462,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
454462
WorkspaceAppsStatsCollectorOptions: options.WorkspaceAppsStatsCollectorOptions,
455463
AllowWorkspaceRenames: options.AllowWorkspaceRenames,
456464
NewTicker: options.NewTicker,
465+
WorkspaceUsageTracker: options.WorkspaceUsageTracker,
457466
}
458467
}
459468

‎coderd/workspaces.go

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1096,6 +1096,24 @@ func (api *API) putExtendWorkspace(rw http.ResponseWriter, r *http.Request) {
10961096
httpapi.Write(ctx, rw, code, resp)
10971097
}
10981098

1099+
// @Summary Post Workspace Usage by ID
1100+
// @ID post-workspace-usage-by-id
1101+
// @Security CoderSessionToken
1102+
// @Tags Workspaces
1103+
// @Param workspace path string true "Workspace ID" format(uuid)
1104+
// @Success 204
1105+
// @Router /workspaces/{workspace}/usage [post]
1106+
func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
1107+
workspace := httpmw.WorkspaceParam(r)
1108+
if !api.Authorize(r, rbac.ActionUpdate, workspace) {
1109+
httpapi.Forbidden(rw)
1110+
return
1111+
}
1112+
1113+
api.workspaceUsageTracker.Add(workspace.ID)
1114+
rw.WriteHeader(http.StatusNoContent)
1115+
}
1116+
10991117
// @Summary Favorite workspace by ID.
11001118
// @ID favorite-workspace-by-id
11011119
// @Security CoderSessionToken

0 commit comments

Comments
 (0)