Skip to content

Commit 92aa1eb

Browse files
johnstcndannykopping
johnstcn
and
dannykopping
authored
fix(cli): port-forward: update workspace last_used_at (coder#12659)
This PR updates the coder port-forward command to periodically inform coderd that the workspace is being used: - Adds workspaceusage.Tracker which periodically batch-updates workspace LastUsedAt - Adds coderd endpoint to signal workspace usage - Updates coder port-forward to periodically hit this endpoint - Modifies BatchUpdateWorkspacesLastUsedAt to avoid overwriting with stale data Co-authored-by: Danny Kopping <danny@coder.com>
1 parent d789a60 commit 92aa1eb

File tree

15 files changed

+708
-2
lines changed

15 files changed

+708
-2
lines changed

‎cli/portforward.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,8 @@ func (r *RootCmd) portForward() *serpent.Command {
136136
listeners[i] = l
137137
}
138138

139+
stopUpdating := client.UpdateWorkspaceUsageContext(ctx, workspace.ID)
140+
139141
// Wait for the context to be canceled or for a signal and close
140142
// all listeners.
141143
var closeErr error
@@ -156,6 +158,7 @@ func (r *RootCmd) portForward() *serpent.Command {
156158
}
157159

158160
cancel()
161+
stopUpdating()
159162
closeAllListeners()
160163
}()
161164

‎cli/portforward_test.go

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"github.com/coder/coder/v2/coderd/coderdtest"
2222
"github.com/coder/coder/v2/coderd/database"
2323
"github.com/coder/coder/v2/coderd/database/dbfake"
24+
"github.com/coder/coder/v2/coderd/database/dbtime"
2425
"github.com/coder/coder/v2/codersdk"
2526
"github.com/coder/coder/v2/pty/ptytest"
2627
"github.com/coder/coder/v2/testutil"
@@ -96,7 +97,12 @@ func TestPortForward(t *testing.T) {
9697
// Setup agent once to be shared between test-cases (avoid expensive
9798
// non-parallel setup).
9899
var (
99-
client, db = coderdtest.NewWithDatabase(t, nil)
100+
wuTick = make(chan time.Time)
101+
wuFlush = make(chan int, 1)
102+
client, db = coderdtest.NewWithDatabase(t, &coderdtest.Options{
103+
WorkspaceUsageTrackerTick: wuTick,
104+
WorkspaceUsageTrackerFlush: wuFlush,
105+
})
100106
admin = coderdtest.CreateFirstUser(t, client)
101107
member, memberUser = coderdtest.CreateAnotherUser(t, client, admin.OrganizationID)
102108
workspace = runAgent(t, client, memberUser.ID, db)
@@ -148,6 +154,13 @@ func TestPortForward(t *testing.T) {
148154
cancel()
149155
err = <-errC
150156
require.ErrorIs(t, err, context.Canceled)
157+
158+
flushCtx := testutil.Context(t, testutil.WaitShort)
159+
testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
160+
_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
161+
updated, err := client.Workspace(context.Background(), workspace.ID)
162+
require.NoError(t, err)
163+
require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
151164
})
152165

153166
t.Run(c.name+"_TwoPorts", func(t *testing.T) {
@@ -196,6 +209,13 @@ func TestPortForward(t *testing.T) {
196209
cancel()
197210
err = <-errC
198211
require.ErrorIs(t, err, context.Canceled)
212+
213+
flushCtx := testutil.Context(t, testutil.WaitShort)
214+
testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
215+
_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
216+
updated, err := client.Workspace(context.Background(), workspace.ID)
217+
require.NoError(t, err)
218+
require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
199219
})
200220
}
201221

@@ -257,6 +277,13 @@ func TestPortForward(t *testing.T) {
257277
cancel()
258278
err := <-errC
259279
require.ErrorIs(t, err, context.Canceled)
280+
281+
flushCtx := testutil.Context(t, testutil.WaitShort)
282+
testutil.RequireSendCtx(flushCtx, t, wuTick, dbtime.Now())
283+
_ = testutil.RequireRecvCtx(flushCtx, t, wuFlush)
284+
updated, err := client.Workspace(context.Background(), workspace.ID)
285+
require.NoError(t, err)
286+
require.Greater(t, updated.LastUsedAt, workspace.LastUsedAt)
260287
})
261288
}
262289

‎cli/server.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ import (
8686
stringutil "github.com/coder/coder/v2/coderd/util/strings"
8787
"github.com/coder/coder/v2/coderd/workspaceapps"
8888
"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
89+
"github.com/coder/coder/v2/coderd/workspaceusage"
8990
"github.com/coder/coder/v2/codersdk"
9091
"github.com/coder/coder/v2/codersdk/drpc"
9192
"github.com/coder/coder/v2/cryptorand"
@@ -968,6 +969,13 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
968969
purger := dbpurge.New(ctx, logger, options.Database)
969970
defer purger.Close()
970971

972+
// Updates workspace usage
973+
tracker := workspaceusage.New(options.Database,
974+
workspaceusage.WithLogger(logger.Named("workspace_usage_tracker")),
975+
)
976+
options.WorkspaceUsageTracker = tracker
977+
defer tracker.Close()
978+
971979
// Wrap the server in middleware that redirects to the access URL if
972980
// the request is not to a local IP.
973981
var handler http.Handler = coderAPI.RootHandler

‎coderd/apidoc/docs.go

Lines changed: 29 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/apidoc/swagger.json

Lines changed: 27 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/coderd.go

Lines changed: 15 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,12 @@ 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.WithLogger(options.Logger.Named("workspace_usage_tracker")),
372+
)
373+
}
374+
365375
ctx, cancel := context.WithCancel(context.Background())
366376
r := chi.NewRouter()
367377

@@ -405,6 +415,7 @@ func New(options *Options) *API {
405415
options.Logger.Named("acquirer"),
406416
options.Database,
407417
options.Pubsub),
418+
workspaceUsageTracker: options.WorkspaceUsageTracker,
408419
}
409420

410421
api.AppearanceFetcher.Store(&appearance.DefaultFetcher)
@@ -972,6 +983,7 @@ func New(options *Options) *API {
972983
})
973984
r.Get("/watch", api.watchWorkspace)
974985
r.Put("/extend", api.putExtendWorkspace)
986+
r.Post("/usage", api.postWorkspaceUsage)
975987
r.Put("/dormant", api.putWorkspaceDormant)
976988
r.Put("/favorite", api.putFavoriteWorkspace)
977989
r.Delete("/favorite", api.deleteFavoriteWorkspace)
@@ -1179,6 +1191,8 @@ type API struct {
11791191
statsBatcher *batchstats.Batcher
11801192

11811193
Acquirer *provisionerdserver.Acquirer
1194+
1195+
workspaceUsageTracker *workspaceusage.Tracker
11821196
}
11831197

11841198
// Close waits for all WebSocket connections to drain before returning.
@@ -1200,6 +1214,7 @@ func (api *API) Close() error {
12001214
_ = (*coordinator).Close()
12011215
}
12021216
_ = api.agentProvider.Close()
1217+
api.workspaceUsageTracker.Close()
12031218
return nil
12041219
}
12051220

‎coderd/coderdtest/coderdtest.go

Lines changed: 34 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,8 @@ type Options struct {
146147
WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
147148
AllowWorkspaceRenames bool
148149
NewTicker func(duration time.Duration) (<-chan time.Time, func())
150+
WorkspaceUsageTrackerFlush chan int
151+
WorkspaceUsageTrackerTick chan time.Time
149152
}
150153

151154
// New constructs a codersdk client connected to an in-memory API instance.
@@ -306,6 +309,36 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
306309
hangDetector.Start()
307310
t.Cleanup(hangDetector.Close)
308311

312+
// Did last_used_at not update? Scratching your noggin? Here's why.
313+
// Workspace usage tracking must be triggered manually in tests.
314+
// The vast majority of existing tests do not depend on last_used_at
315+
// and adding an extra time-based background goroutine to all existing
316+
// tests may lead to future flakes and goleak complaints.
317+
// Instead, pass in your own flush and ticker like so:
318+
//
319+
// tickCh = make(chan time.Time)
320+
// flushCh = make(chan int, 1)
321+
// client = coderdtest.New(t, &coderdtest.Options{
322+
// WorkspaceUsageTrackerFlush: flushCh,
323+
// WorkspaceUsageTrackerTick: tickCh
324+
// })
325+
//
326+
// Now to trigger a tick, just write to `tickCh`.
327+
// Reading from `flushCh` will ensure that workspaceusage.Tracker flushed.
328+
// See TestPortForward or TestTracker_MultipleInstances for how this works in practice.
329+
if options.WorkspaceUsageTrackerFlush == nil {
330+
options.WorkspaceUsageTrackerFlush = make(chan int, 1) // buffering just in case
331+
}
332+
if options.WorkspaceUsageTrackerTick == nil {
333+
options.WorkspaceUsageTrackerTick = make(chan time.Time, 1) // buffering just in case
334+
}
335+
// Close is called by API.Close()
336+
wuTracker := workspaceusage.New(
337+
options.Database,
338+
workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
339+
workspaceusage.WithTickFlush(options.WorkspaceUsageTrackerTick, options.WorkspaceUsageTrackerFlush),
340+
)
341+
309342
var mutex sync.RWMutex
310343
var handler http.Handler
311344
srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -454,6 +487,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
454487
WorkspaceAppsStatsCollectorOptions: options.WorkspaceAppsStatsCollectorOptions,
455488
AllowWorkspaceRenames: options.AllowWorkspaceRenames,
456489
NewTicker: options.NewTicker,
490+
WorkspaceUsageTracker: wuTracker,
457491
}
458492
}
459493

‎coderd/database/dbmem/dbmem.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1046,6 +1046,10 @@ func (q *FakeQuerier) BatchUpdateWorkspaceLastUsedAt(_ context.Context, arg data
10461046
if _, found := m[q.workspaces[i].ID]; !found {
10471047
continue
10481048
}
1049+
// WHERE last_used_at < @last_used_at
1050+
if !q.workspaces[i].LastUsedAt.Before(arg.LastUsedAt) {
1051+
continue
1052+
}
10491053
q.workspaces[i].LastUsedAt = arg.LastUsedAt
10501054
n++
10511055
}

‎coderd/database/queries.sql.go

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/workspaces.sql

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -433,7 +433,10 @@ UPDATE
433433
SET
434434
last_used_at = @last_used_at
435435
WHERE
436-
id = ANY(@ids :: uuid[]);
436+
id = ANY(@ids :: uuid[])
437+
AND
438+
-- Do not overwrite with older data
439+
last_used_at < @last_used_at;
437440

438441
-- name: GetDeploymentWorkspaceStats :one
439442
WITH workspaces_with_jobs AS (

‎coderd/workspaces.go

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

1087+
// @Summary Post Workspace Usage by ID
1088+
// @ID post-workspace-usage-by-id
1089+
// @Security CoderSessionToken
1090+
// @Tags Workspaces
1091+
// @Param workspace path string true "Workspace ID" format(uuid)
1092+
// @Success 204
1093+
// @Router /workspaces/{workspace}/usage [post]
1094+
func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
1095+
workspace := httpmw.WorkspaceParam(r)
1096+
if !api.Authorize(r, rbac.ActionUpdate, workspace) {
1097+
httpapi.Forbidden(rw)
1098+
return
1099+
}
1100+
1101+
api.workspaceUsageTracker.Add(workspace.ID)
1102+
rw.WriteHeader(http.StatusNoContent)
1103+
}
1104+
10871105
// @Summary Favorite workspace by ID.
10881106
// @ID favorite-workspace-by-id
10891107
// @Security CoderSessionToken

0 commit comments

Comments
 (0)