Skip to content

Commit 37fc664

Browse files
mafredrimafredri
authored
perf(coderd/database): limit GetLatestWorkspaceAppStatusByAppID to 1 row (#20917)
## Description This PR fixes an issue where `GetLatestWorkspaceAppStatusesByAppID` returned an unbounded number of rows for a given app ID, which could cause performance issues for noisy or long-running AI tasks. ## Impact This change reduces database query overhead for workspace app status updates, particularly for busy AI tasks that update their status frequently. Previously, fetching the latest status would return all historical statuses, now it returns only the most recent one. Fixes #20862 --- 🤖 This change was written by Claude Sonnet 4.5 Thinking using [mux](https://github.com/coder/mux) and reviewed by a human 🏄🏻‍♂️
1 parent 5213023 commit 37fc664

File tree

8 files changed

+45
-56
lines changed

8 files changed

+45
-56
lines changed

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2426,11 +2426,11 @@ func (q *querier) GetLatestCryptoKeyByFeature(ctx context.Context, feature datab
24262426
return q.db.GetLatestCryptoKeyByFeature(ctx, feature)
24272427
}
24282428

2429-
func (q *querier) GetLatestWorkspaceAppStatusesByAppID(ctx context.Context, appID uuid.UUID) ([]database.WorkspaceAppStatus, error) {
2429+
func (q *querier) GetLatestWorkspaceAppStatusByAppID(ctx context.Context, appID uuid.UUID) (database.WorkspaceAppStatus, error) {
24302430
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
2431-
return nil, err
2431+
return database.WorkspaceAppStatus{}, err
24322432
}
2433-
return q.db.GetLatestWorkspaceAppStatusesByAppID(ctx, appID)
2433+
return q.db.GetLatestWorkspaceAppStatusByAppID(ctx, appID)
24342434
}
24352435

24362436
func (q *querier) GetLatestWorkspaceAppStatusesByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]database.WorkspaceAppStatus, error) {

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2864,9 +2864,9 @@ func (s *MethodTestSuite) TestSystemFunctions() {
28642864
dbm.EXPECT().UpdateUserLinkedID(gomock.Any(), arg).Return(l, nil).AnyTimes()
28652865
check.Args(arg).Asserts(rbac.ResourceSystem, policy.ActionUpdate).Returns(l)
28662866
}))
2867-
s.Run("GetLatestWorkspaceAppStatusesByAppID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
2867+
s.Run("GetLatestWorkspaceAppStatusByAppID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
28682868
appID := uuid.New()
2869-
dbm.EXPECT().GetLatestWorkspaceAppStatusesByAppID(gomock.Any(), appID).Return([]database.WorkspaceAppStatus{}, nil).AnyTimes()
2869+
dbm.EXPECT().GetLatestWorkspaceAppStatusByAppID(gomock.Any(), appID).Return(database.WorkspaceAppStatus{}, nil).AnyTimes()
28702870
check.Args(appID).Asserts(rbac.ResourceSystem, policy.ActionRead)
28712871
}))
28722872
s.Run("GetLatestWorkspaceAppStatusesByWorkspaceIDs", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {

‎coderd/database/dbmetrics/querymetrics.go‎

Lines changed: 3 additions & 3 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: 7 additions & 7 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 & 1 deletion
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: 16 additions & 31 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/workspaceapps.sql‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,11 +73,12 @@ RETURNING *;
7373
-- name: GetWorkspaceAppStatusesByAppIDs :many
7474
SELECT * FROM workspace_app_statuses WHERE app_id = ANY(@ids :: uuid [ ]);
7575

76-
-- name: GetLatestWorkspaceAppStatusesByAppID :many
76+
-- name: GetLatestWorkspaceAppStatusByAppID :one
7777
SELECT *
7878
FROM workspace_app_statuses
7979
WHERE app_id = @app_id::uuid
80-
ORDER BY created_at DESC, id DESC;
80+
ORDER BY created_at DESC, id DESC
81+
LIMIT 1;
8182

8283
-- name: GetLatestWorkspaceAppStatusesByWorkspaceIDs :many
8384
SELECT DISTINCT ON (workspace_id)

‎coderd/workspaceagents.go‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -388,16 +388,17 @@ func (api *API) patchWorkspaceAgentAppStatus(rw http.ResponseWriter, r *http.Req
388388
// Treat the message as untrusted input.
389389
cleaned := strutil.UISanitize(req.Message)
390390

391-
// Get the latest statuses for the workspace app to detect no-op updates
391+
// Get the latest status for the workspace app to detect no-op updates
392392
// nolint:gocritic // This is a system restricted operation.
393-
latestAppStatus, err := api.Database.GetLatestWorkspaceAppStatusesByAppID(dbauthz.AsSystemRestricted(ctx), app.ID)
394-
if err != nil {
393+
latestAppStatus, err := api.Database.GetLatestWorkspaceAppStatusByAppID(dbauthz.AsSystemRestricted(ctx), app.ID)
394+
if err != nil && !errors.Is(err, sql.ErrNoRows) {
395395
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
396-
Message: "Failed to get latest workspace app statuses.",
396+
Message: "Failed to get latest workspace app status.",
397397
Detail: err.Error(),
398398
})
399399
return
400400
}
401+
// If no rows found, latestAppStatus will be a zero-value struct (ID == uuid.Nil)
401402

402403
// nolint:gocritic // This is a system restricted operation.
403404
_, err = api.Database.InsertWorkspaceAppStatus(dbauthz.AsSystemRestricted(ctx), database.InsertWorkspaceAppStatusParams{
@@ -442,7 +443,7 @@ func (api *API) patchWorkspaceAgentAppStatus(rw http.ResponseWriter, r *http.Req
442443
func (api *API) enqueueAITaskStateNotification(
443444
ctx context.Context,
444445
appID uuid.UUID,
445-
latestAppStatus []database.WorkspaceAppStatus,
446+
latestAppStatus database.WorkspaceAppStatus,
446447
newAppStatus codersdk.WorkspaceAppStatusState,
447448
workspace database.Workspace,
448449
agent database.WorkspaceAgent,
@@ -492,14 +493,16 @@ func (api *API) enqueueAITaskStateNotification(
492493
}
493494

494495
// Skip if the latest persisted state equals the new state (no new transition)
495-
if len(latestAppStatus) > 0 && latestAppStatus[0].State == database.WorkspaceAppStatusState(newAppStatus) {
496+
// Note: uuid.Nil check is valid here. If no previous status exists,
497+
// GetLatestWorkspaceAppStatusByAppID returns sql.ErrNoRows and we get a zero-value struct.
498+
if latestAppStatus.ID != uuid.Nil && latestAppStatus.State == database.WorkspaceAppStatusState(newAppStatus) {
496499
return
497500
}
498501

499502
// Skip the initial "Working" notification when task first starts.
500503
// This is obvious to the user since they just created the task.
501504
// We still notify on first "Idle" status and all subsequent transitions.
502-
if len(latestAppStatus) == 0 && newAppStatus == codersdk.WorkspaceAppStatusStateWorking {
505+
if latestAppStatus.ID == uuid.Nil && newAppStatus == codersdk.WorkspaceAppStatusStateWorking {
503506
return
504507
}
505508

0 commit comments

Comments
 (0)