Skip to content

Commit d583aca

Browse files
johnstcnjohnstcn
authored
fix(coderd): workspaceapps: update last_used_at when workspace app reports stats (coder#11603)
- Adds a new query BatchUpdateLastUsedAt - Adds calls to BatchUpdateLastUsedAt in app stats handler upon flush - Passes a stats flush channel to apptest setup scaffolding and updates unit tests to assert modifications to LastUsedAt.
1 parent 5bfbf9f commit d583aca

File tree

16 files changed

+186
-15
lines changed

16 files changed

+186
-15
lines changed

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -695,6 +695,15 @@ func (q *querier) ArchiveUnusedTemplateVersions(ctx context.Context, arg databas
695695
return q.db.ArchiveUnusedTemplateVersions(ctx, arg)
696696
}
697697

698+
func (q *querier) BatchUpdateWorkspaceLastUsedAt(ctx context.Context, arg database.BatchUpdateWorkspaceLastUsedAtParams) error {
699+
// Could be any workspace and checking auth to each workspace is overkill for the purpose
700+
// of this function.
701+
if err := q.authorizeContext(ctx, rbac.ActionUpdate, rbac.ResourceWorkspace.All()); err != nil {
702+
return err
703+
}
704+
return q.db.BatchUpdateWorkspaceLastUsedAt(ctx, arg)
705+
}
706+
698707
func (q *querier) CleanTailnetCoordinators(ctx context.Context) error {
699708
if err := q.authorizeContext(ctx, rbac.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil {
700709
return err

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1549,6 +1549,13 @@ func (s *MethodTestSuite) TestWorkspace() {
15491549
ID: ws.ID,
15501550
}).Asserts(ws, rbac.ActionUpdate).Returns()
15511551
}))
1552+
s.Run("BatchUpdateWorkspaceLastUsedAt", s.Subtest(func(db database.Store, check *expects) {
1553+
ws1 := dbgen.Workspace(s.T(), db, database.Workspace{})
1554+
ws2 := dbgen.Workspace(s.T(), db, database.Workspace{})
1555+
check.Args(database.BatchUpdateWorkspaceLastUsedAtParams{
1556+
IDs: []uuid.UUID{ws1.ID, ws2.ID},
1557+
}).Asserts(rbac.ResourceWorkspace.All(), rbac.ActionUpdate).Returns()
1558+
}))
15521559
s.Run("UpdateWorkspaceTTL", s.Subtest(func(db database.Store, check *expects) {
15531560
ws := dbgen.Workspace(s.T(), db, database.Workspace{})
15541561
check.Args(database.UpdateWorkspaceTTLParams{

‎coderd/database/dbmem/dbmem.go

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -963,6 +963,31 @@ func (q *FakeQuerier) ArchiveUnusedTemplateVersions(_ context.Context, arg datab
963963
return archived, nil
964964
}
965965

966+
func (q *FakeQuerier) BatchUpdateWorkspaceLastUsedAt(_ context.Context, arg database.BatchUpdateWorkspaceLastUsedAtParams) error {
967+
err := validateDatabaseType(arg)
968+
if err != nil {
969+
return err
970+
}
971+
972+
q.mutex.Lock()
973+
defer q.mutex.Unlock()
974+
975+
// temporary map to avoid O(q.workspaces*arg.workspaceIds)
976+
m := make(map[uuid.UUID]struct{})
977+
for _, id := range arg.IDs {
978+
m[id] = struct{}{}
979+
}
980+
n := 0
981+
for i := 0; i < len(q.workspaces); i++ {
982+
if _, found := m[q.workspaces[i].ID]; !found {
983+
continue
984+
}
985+
q.workspaces[i].LastUsedAt = arg.LastUsedAt
986+
n++
987+
}
988+
return nil
989+
}
990+
966991
func (*FakeQuerier) CleanTailnetCoordinators(_ context.Context) error {
967992
return ErrUnimplemented
968993
}

‎coderd/database/dbmetrics/dbmetrics.go

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

‎coderd/database/dbmock/dbmock.go

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

‎coderd/database/querier.go

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

‎coderd/database/queries.sql.go

Lines changed: 19 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: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -357,6 +357,14 @@ SET
357357
WHERE
358358
id = $1;
359359

360+
-- name: BatchUpdateWorkspaceLastUsedAt :exec
361+
UPDATE
362+
workspaces
363+
SET
364+
last_used_at = @last_used_at
365+
WHERE
366+
id = ANY(@ids :: uuid[]);
367+
360368
-- name: GetDeploymentWorkspaceStats :one
361369
WITH workspaces_with_jobs AS (
362370
SELECT

‎coderd/httpapi/websocket.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,9 @@ import (
44
"context"
55
"time"
66

7-
"cdr.dev/slog"
87
"nhooyr.io/websocket"
8+
9+
"cdr.dev/slog"
910
)
1011

1112
// Heartbeat loops to ping a WebSocket to keep it alive.

‎coderd/workspaceagents_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1626,7 +1626,7 @@ func TestWorkspaceAgentExternalAuthListen(t *testing.T) {
16261626
cancel()
16271627
// We expect only 1
16281628
// In a failed test, you will likely see 9, as the last one
1629-
// gets cancelled.
1629+
// gets canceled.
16301630
require.Equal(t, 1, validateCalls, "validate calls duplicated on same token")
16311631
})
16321632
}

‎coderd/workspaceapps/apptest/apptest.go

Lines changed: 48 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
6464
// reconnecting-pty proxy server we want to test is mounted.
6565
client := appDetails.AppClient(t)
6666
testReconnectingPTY(ctx, t, client, appDetails.Agent.ID, "")
67+
assertWorkspaceLastUsedAtUpdated(t, appDetails)
6768
})
6869

6970
t.Run("SignedTokenQueryParameter", func(t *testing.T) {
@@ -92,6 +93,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
9293
// Make an unauthenticated client.
9394
unauthedAppClient := codersdk.New(appDetails.AppClient(t).URL)
9495
testReconnectingPTY(ctx, t, unauthedAppClient, appDetails.Agent.ID, issueRes.SignedToken)
96+
assertWorkspaceLastUsedAtUpdated(t, appDetails)
9597
})
9698
})
9799

@@ -117,6 +119,9 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
117119
body, err := io.ReadAll(resp.Body)
118120
require.NoError(t, err)
119121
require.Contains(t, string(body), "Path-based applications are disabled")
122+
// Even though path-based apps are disabled, the request should indicate
123+
// that the workspace was used.
124+
assertWorkspaceLastUsedAtNotUpdated(t, appDetails)
120125
})
121126

122127
t.Run("LoginWithoutAuthOnPrimary", func(t *testing.T) {
@@ -142,6 +147,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
142147
require.NoError(t, err)
143148
require.True(t, loc.Query().Has("message"))
144149
require.True(t, loc.Query().Has("redirect"))
150+
assertWorkspaceLastUsedAtUpdated(t, appDetails)
145151
})
146152

147153
t.Run("LoginWithoutAuthOnProxy", func(t *testing.T) {
@@ -179,6 +185,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
179185
// request is getting stripped.
180186
require.Equal(t, u.Path, redirectURI.Path+"/")
181187
require.Equal(t, u.RawQuery, redirectURI.RawQuery)
188+
assertWorkspaceLastUsedAtUpdated(t, appDetails)
182189
})
183190

184191
t.Run("NoAccessShould404", func(t *testing.T) {
@@ -195,6 +202,8 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
195202
require.NoError(t, err)
196203
defer resp.Body.Close()
197204
require.Equal(t, http.StatusNotFound, resp.StatusCode)
205+
// TODO(cian): A blocked request should not count as workspace usage.
206+
// assertWorkspaceLastUsedAtNotUpdated(t, appDetails.AppClient(t), appDetails)
198207
})
199208

200209
t.Run("RedirectsWithSlash", func(t *testing.T) {
@@ -209,6 +218,8 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
209218
require.NoError(t, err)
210219
defer resp.Body.Close()
211220
require.Equal(t, http.StatusTemporaryRedirect, resp.StatusCode)
221+
// TODO(cian): The initial redirect should not count as workspace usage.
222+
// assertWorkspaceLastUsedAtNotUpdated(t, appDetails.AppClient(t), appDetails)
212223
})
213224

214225
t.Run("RedirectsWithQuery", func(t *testing.T) {
@@ -226,6 +237,8 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
226237
loc, err := resp.Location()
227238
require.NoError(t, err)
228239
require.Equal(t, proxyTestAppQuery, loc.RawQuery)
240+
// TODO(cian): The initial redirect should not count as workspace usage.
241+
// assertWorkspaceLastUsedAtNotUpdated(t, appDetails.AppClient(t), appDetails)
229242
})
230243

231244
t.Run("Proxies", func(t *testing.T) {
@@ -267,6 +280,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
267280
require.NoError(t, err)
268281
require.Equal(t, proxyTestAppBody, string(body))
269282
require.Equal(t, http.StatusOK, resp.StatusCode)
283+
assertWorkspaceLastUsedAtUpdated(t, appDetails)
270284
})
271285

272286
t.Run("ProxiesHTTPS", func(t *testing.T) {
@@ -312,6 +326,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
312326
require.NoError(t, err)
313327
require.Equal(t, proxyTestAppBody, string(body))
314328
require.Equal(t, http.StatusOK, resp.StatusCode)
329+
assertWorkspaceLastUsedAtUpdated(t, appDetails)
315330
})
316331

317332
t.Run("BlocksMe", func(t *testing.T) {
@@ -331,6 +346,8 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
331346
body, err := io.ReadAll(resp.Body)
332347
require.NoError(t, err)
333348
require.Contains(t, string(body), "must be accessed with the full username, not @me")
349+
// TODO(cian): A blocked request should not count as workspace usage.
350+
// assertWorkspaceLastUsedAtNotUpdated(t, appDetails.AppClient(t), appDetails)
334351
})
335352

336353
t.Run("ForwardsIP", func(t *testing.T) {
@@ -349,6 +366,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
349366
require.Equal(t, proxyTestAppBody, string(body))
350367
require.Equal(t, http.StatusOK, resp.StatusCode)
351368
require.Equal(t, "1.1.1.1,127.0.0.1", resp.Header.Get("X-Forwarded-For"))
369+
assertWorkspaceLastUsedAtUpdated(t, appDetails)
352370
})
353371

354372
t.Run("ProxyError", func(t *testing.T) {
@@ -361,6 +379,9 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
361379
require.NoError(t, err)
362380
defer resp.Body.Close()
363381
require.Equal(t, http.StatusBadGateway, resp.StatusCode)
382+
// An valid authenticated attempt to access a workspace app
383+
// should count as usage regardless of success.
384+
assertWorkspaceLastUsedAtUpdated(t, appDetails)
364385
})
365386

366387
t.Run("NoProxyPort", func(t *testing.T) {
@@ -375,6 +396,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
375396
// TODO(@deansheather): This should be 400. There's a todo in the
376397
// resolve request code to fix this.
377398
require.Equal(t, http.StatusInternalServerError, resp.StatusCode)
399+
assertWorkspaceLastUsedAtUpdated(t, appDetails)
378400
})
379401
})
380402

@@ -1430,16 +1452,12 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
14301452
t.Run("ReportStats", func(t *testing.T) {
14311453
t.Parallel()
14321454

1433-
flush := make(chan chan<- struct{}, 1)
1434-
14351455
reporter := &fakeStatsReporter{}
14361456
appDetails := setupProxyTest(t, &DeploymentOptions{
14371457
StatsCollectorOptions: workspaceapps.StatsCollectorOptions{
14381458
Reporter: reporter,
14391459
ReportInterval: time.Hour,
14401460
RollupWindow: time.Minute,
1441-
1442-
Flush: flush,
14431461
},
14441462
})
14451463

@@ -1457,10 +1475,7 @@ func Run(t *testing.T, appHostIsPrimary bool, factory DeploymentFactory) {
14571475
var stats []workspaceapps.StatsReport
14581476
require.Eventually(t, func() bool {
14591477
// Keep flushing until we get a non-empty stats report.
1460-
flushDone := make(chan struct{}, 1)
1461-
flush <- flushDone
1462-
<-flushDone
1463-
1478+
appDetails.FlushStats()
14641479
stats = reporter.stats()
14651480
return len(stats) > 0
14661481
}, testutil.WaitLong, testutil.IntervalFast, "stats not reported")
@@ -1549,3 +1564,28 @@ func testReconnectingPTY(ctx context.Context, t *testing.T, client *codersdk.Cli
15491564
// Ensure the connection closes.
15501565
require.ErrorIs(t, tr.ReadUntil(ctx, nil), io.EOF)
15511566
}
1567+
1568+
// Accessing an app should update the workspace's LastUsedAt.
1569+
// NOTE: Despite our efforts with the flush channel, this is inherently racy.
1570+
func assertWorkspaceLastUsedAtUpdated(t testing.TB, details *Details) {
1571+
t.Helper()
1572+
1573+
// Wait for stats to fully flush.
1574+
require.Eventually(t, func() bool {
1575+
details.FlushStats()
1576+
ws, err := details.SDKClient.Workspace(context.Background(), details.Workspace.ID)
1577+
assert.NoError(t, err)
1578+
return ws.LastUsedAt.After(details.Workspace.LastUsedAt)
1579+
}, testutil.WaitShort, testutil.IntervalMedium, "workspace LastUsedAt not updated when it should have been")
1580+
}
1581+
1582+
// Except when it sometimes shouldn't (e.g. no access)
1583+
// NOTE: Despite our efforts with the flush channel, this is inherently racy.
1584+
func assertWorkspaceLastUsedAtNotUpdated(t testing.TB, details *Details) {
1585+
t.Helper()
1586+
1587+
details.FlushStats()
1588+
ws, err := details.SDKClient.Workspace(context.Background(), details.Workspace.ID)
1589+
require.NoError(t, err)
1590+
require.Equal(t, ws.LastUsedAt, details.Workspace.LastUsedAt, "workspace LastUsedAt updated when it should not have been")
1591+
}

‎coderd/workspaceapps/apptest/setup.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ type Deployment struct {
7171
SDKClient *codersdk.Client
7272
FirstUser codersdk.CreateFirstUserResponse
7373
PathAppBaseURL *url.URL
74+
FlushStats func()
7475
}
7576

7677
// DeploymentFactory generates a deployment with an API client, a path base URL,

‎coderd/workspaceapps/stats.go

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
"github.com/coder/coder/v2/coderd/database"
1414
"github.com/coder/coder/v2/coderd/database/dbauthz"
1515
"github.com/coder/coder/v2/coderd/database/dbtime"
16+
"github.com/coder/coder/v2/coderd/util/slice"
1617
)
1718

1819
const (
@@ -117,11 +118,25 @@ func (r *StatsDBReporter) Report(ctx context.Context, stats []StatsReport) error
117118
batch.Requests = batch.Requests[:0]
118119
}
119120
}
120-
if len(batch.UserID) > 0 {
121-
err := tx.InsertWorkspaceAppStats(ctx, batch)
122-
if err != nil {
123-
return err
124-
}
121+
if len(batch.UserID) == 0 {
122+
return nil
123+
}
124+
125+
if err := tx.InsertWorkspaceAppStats(ctx, batch); err != nil {
126+
return err
127+
}
128+
129+
// TODO: We currently measure workspace usage based on when we get stats from it.
130+
// There are currently two paths for this:
131+
// 1) From SSH -> workspace agent stats POSTed from agent
132+
// 2) From workspace apps / rpty -> workspace app stats (from coderd / wsproxy)
133+
// Ideally we would have a single code path for this.
134+
uniqueIDs := slice.Unique(batch.WorkspaceID)
135+
if err := tx.BatchUpdateWorkspaceLastUsedAt(ctx, database.BatchUpdateWorkspaceLastUsedAtParams{
136+
IDs: uniqueIDs,
137+
LastUsedAt: dbtime.Now(), // This isn't 100% accurate, but it's good enough.
138+
}); err != nil {
139+
return err
125140
}
126141

127142
return nil
@@ -234,6 +249,7 @@ func (sc *StatsCollector) Collect(report StatsReport) {
234249
}
235250
delete(sc.statsBySessionID, report.SessionID)
236251
}
252+
sc.opts.Logger.Debug(sc.ctx, "collected workspace app stats", slog.F("report", report))
237253
}
238254

239255
// rollup performs stats rollup for sessions that fall within the

0 commit comments

Comments
 (0)