Skip to content

Commit 365f7f4

Browse files
ethanndicksonethanndickson
committed
feedback
1 parent 5386d9f commit 365f7f4

File tree

17 files changed

+177
-241
lines changed

17 files changed

+177
-241
lines changed

‎coderd/agentapi/api.go

Lines changed: 6 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,7 @@ type API struct {
4646
*ScriptsAPI
4747
*tailnet.DRPCService
4848

49-
mu sync.Mutex
50-
cachedWorkspaceID uuid.UUID
49+
mu sync.Mutex
5150
}
5251

5352
var _ agentproto.DRPCAgentServer = &API{}
@@ -83,9 +82,8 @@ type Options struct {
8382

8483
func New(opts Options) *API {
8584
api := &API{
86-
opts: opts,
87-
mu: sync.Mutex{},
88-
cachedWorkspaceID: opts.WorkspaceID,
85+
opts: opts,
86+
mu: sync.Mutex{},
8987
}
9088

9189
api.ManifestAPI = &ManifestAPI{
@@ -199,44 +197,11 @@ func (a *API) agent(ctx context.Context) (database.WorkspaceAgent, error) {
199197
return agent, nil
200198
}
201199

202-
func (a *API) workspaceID(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
203-
a.mu.Lock()
204-
if a.cachedWorkspaceID != uuid.Nil {
205-
id := a.cachedWorkspaceID
206-
a.mu.Unlock()
207-
return id, nil
208-
}
209-
210-
if agent == nil {
211-
agnt, err := a.agent(ctx)
212-
if err != nil {
213-
return uuid.Nil, err
214-
}
215-
agent = &agnt
216-
}
217-
218-
getWorkspaceAgentByIDRow, err := a.opts.Database.GetWorkspaceByAgentID(ctx, agent.ID)
219-
if err != nil {
220-
return uuid.Nil, xerrors.Errorf("get workspace by agent id %q: %w", agent.ID, err)
221-
}
222-
223-
a.mu.Lock()
224-
a.cachedWorkspaceID = getWorkspaceAgentByIDRow.ID
225-
a.mu.Unlock()
226-
return getWorkspaceAgentByIDRow.ID, nil
227-
}
228-
229-
func (a *API) publishWorkspaceUpdate(ctx context.Context, agent *database.WorkspaceAgent) error {
230-
workspaceID, err := a.workspaceID(ctx, agent)
231-
if err != nil {
232-
return err
233-
}
234-
200+
func (a *API) publishWorkspaceUpdate(ctx context.Context, agent *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
235201
a.opts.PublishWorkspaceUpdateFn(ctx, a.opts.OwnerID, wspubsub.WorkspaceEvent{
236-
Kind: wspubsub.WorkspaceEventKindAgentUpdate,
237-
WorkspaceID: workspaceID,
202+
Kind: kind,
203+
WorkspaceID: a.opts.WorkspaceID,
238204
AgentID: &agent.ID,
239-
AgentName: &agent.Name,
240205
})
241206
return nil
242207
}

‎coderd/agentapi/apps.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,14 @@ import (
99
"cdr.dev/slog"
1010
agentproto "github.com/coder/coder/v2/agent/proto"
1111
"github.com/coder/coder/v2/coderd/database"
12+
"github.com/coder/coder/v2/coderd/wspubsub"
1213
)
1314

1415
type AppsAPI struct {
1516
AgentFn func(context.Context) (database.WorkspaceAgent, error)
1617
Database database.Store
1718
Log slog.Logger
18-
PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent) error
19+
PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent, wspubsub.WorkspaceEventKind) error
1920
}
2021

2122
func (a *AppsAPI) BatchUpdateAppHealths(ctx context.Context, req *agentproto.BatchUpdateAppHealthRequest) (*agentproto.BatchUpdateAppHealthResponse, error) {
@@ -96,7 +97,7 @@ func (a *AppsAPI) BatchUpdateAppHealths(ctx context.Context, req *agentproto.Bat
9697
}
9798

9899
if a.PublishWorkspaceUpdateFn != nil && len(newApps) > 0 {
99-
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent)
100+
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent, wspubsub.WorkspaceEventKindAppHealthUpdate)
100101
if err != nil {
101102
return nil, xerrors.Errorf("publish workspace update: %w", err)
102103
}

‎coderd/agentapi/apps_test.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"github.com/coder/coder/v2/coderd/agentapi"
1515
"github.com/coder/coder/v2/coderd/database"
1616
"github.com/coder/coder/v2/coderd/database/dbmock"
17+
"github.com/coder/coder/v2/coderd/wspubsub"
1718
)
1819

1920
func TestBatchUpdateAppHealths(t *testing.T) {
@@ -62,7 +63,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
6263
},
6364
Database: dbM,
6465
Log: slogtest.Make(t, nil),
65-
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
66+
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
6667
publishCalled = true
6768
return nil
6869
},
@@ -100,7 +101,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
100101
},
101102
Database: dbM,
102103
Log: slogtest.Make(t, nil),
103-
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
104+
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
104105
publishCalled = true
105106
return nil
106107
},
@@ -139,7 +140,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
139140
},
140141
Database: dbM,
141142
Log: slogtest.Make(t, nil),
142-
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
143+
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
143144
publishCalled = true
144145
return nil
145146
},

‎coderd/agentapi/lifecycle.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
agentproto "github.com/coder/coder/v2/agent/proto"
1616
"github.com/coder/coder/v2/coderd/database"
1717
"github.com/coder/coder/v2/coderd/database/dbtime"
18+
"github.com/coder/coder/v2/coderd/wspubsub"
1819
)
1920

2021
type contextKeyAPIVersion struct{}
@@ -28,7 +29,7 @@ type LifecycleAPI struct {
2829
WorkspaceID uuid.UUID
2930
Database database.Store
3031
Log slog.Logger
31-
PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent) error
32+
PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent, wspubsub.WorkspaceEventKind) error
3233

3334
TimeNowFn func() time.Time // defaults to dbtime.Now()
3435
}
@@ -118,7 +119,7 @@ func (a *LifecycleAPI) UpdateLifecycle(ctx context.Context, req *agentproto.Upda
118119
}
119120

120121
if a.PublishWorkspaceUpdateFn != nil {
121-
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent)
122+
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent, wspubsub.WorkspaceEventKindAgentLifecycleUpdate)
122123
if err != nil {
123124
return nil, xerrors.Errorf("publish workspace update: %w", err)
124125
}

‎coderd/agentapi/lifecycle_test.go

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"github.com/coder/coder/v2/coderd/database"
2020
"github.com/coder/coder/v2/coderd/database/dbmock"
2121
"github.com/coder/coder/v2/coderd/database/dbtime"
22+
"github.com/coder/coder/v2/coderd/wspubsub"
2223
)
2324

2425
func TestUpdateLifecycle(t *testing.T) {
@@ -72,7 +73,7 @@ func TestUpdateLifecycle(t *testing.T) {
7273
WorkspaceID: workspaceID,
7374
Database: dbM,
7475
Log: slogtest.Make(t, nil),
75-
PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent) error {
76+
PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
7677
publishCalled = true
7778
return nil
7879
},
@@ -155,7 +156,7 @@ func TestUpdateLifecycle(t *testing.T) {
155156
WorkspaceID: workspaceID,
156157
Database: dbM,
157158
Log: slogtest.Make(t, nil),
158-
PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent) error {
159+
PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
159160
publishCalled = true
160161
return nil
161162
},
@@ -234,7 +235,7 @@ func TestUpdateLifecycle(t *testing.T) {
234235
WorkspaceID: workspaceID,
235236
Database: dbM,
236237
Log: slogtest.Make(t, nil),
237-
PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent) error {
238+
PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
238239
atomic.AddInt64(&publishCalled, 1)
239240
return nil
240241
},
@@ -307,7 +308,7 @@ func TestUpdateLifecycle(t *testing.T) {
307308
WorkspaceID: workspaceID,
308309
Database: dbM,
309310
Log: slogtest.Make(t, nil),
310-
PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent) error {
311+
PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
311312
publishCalled = true
312313
return nil
313314
},

‎coderd/agentapi/logs.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,15 @@ import (
1111
agentproto "github.com/coder/coder/v2/agent/proto"
1212
"github.com/coder/coder/v2/coderd/database"
1313
"github.com/coder/coder/v2/coderd/database/dbtime"
14+
"github.com/coder/coder/v2/coderd/wspubsub"
1415
"github.com/coder/coder/v2/codersdk/agentsdk"
1516
)
1617

1718
type LogsAPI struct {
1819
AgentFn func(context.Context) (database.WorkspaceAgent, error)
1920
Database database.Store
2021
Log slog.Logger
21-
PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent) error
22+
PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent, wspubsub.WorkspaceEventKind) error
2223
PublishWorkspaceAgentLogsUpdateFn func(ctx context.Context, workspaceAgentID uuid.UUID, msg agentsdk.LogsNotifyMessage)
2324

2425
TimeNowFn func() time.Time // defaults to dbtime.Now()
@@ -123,7 +124,7 @@ func (a *LogsAPI) BatchCreateLogs(ctx context.Context, req *agentproto.BatchCrea
123124
}
124125

125126
if a.PublishWorkspaceUpdateFn != nil {
126-
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent)
127+
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent, wspubsub.WorkspaceEventKindAgentLogsOverflow)
127128
if err != nil {
128129
return nil, xerrors.Errorf("publish workspace update: %w", err)
129130
}
@@ -143,7 +144,7 @@ func (a *LogsAPI) BatchCreateLogs(ctx context.Context, req *agentproto.BatchCrea
143144
if workspaceAgent.LogsLength == 0 && a.PublishWorkspaceUpdateFn != nil {
144145
// If these are the first logs being appended, we publish a UI update
145146
// to notify the UI that logs are now available.
146-
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent)
147+
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent, wspubsub.WorkspaceEventKindAgentLogsUpdate)
147148
if err != nil {
148149
return nil, xerrors.Errorf("publish workspace update: %w", err)
149150
}

‎coderd/agentapi/logs_test.go

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"github.com/coder/coder/v2/coderd/database"
2020
"github.com/coder/coder/v2/coderd/database/dbmock"
2121
"github.com/coder/coder/v2/coderd/database/dbtime"
22+
"github.com/coder/coder/v2/coderd/wspubsub"
2223
"github.com/coder/coder/v2/codersdk/agentsdk"
2324
)
2425

@@ -50,7 +51,7 @@ func TestBatchCreateLogs(t *testing.T) {
5051
},
5152
Database: dbM,
5253
Log: slogtest.Make(t, nil),
53-
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
54+
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
5455
publishWorkspaceUpdateCalled = true
5556
return nil
5657
},
@@ -154,7 +155,7 @@ func TestBatchCreateLogs(t *testing.T) {
154155
},
155156
Database: dbM,
156157
Log: slogtest.Make(t, nil),
157-
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
158+
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
158159
publishWorkspaceUpdateCalled = true
159160
return nil
160161
},
@@ -202,7 +203,7 @@ func TestBatchCreateLogs(t *testing.T) {
202203
},
203204
Database: dbM,
204205
Log: slogtest.Make(t, nil),
205-
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
206+
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
206207
publishWorkspaceUpdateCalled = true
207208
return nil
208209
},
@@ -295,7 +296,7 @@ func TestBatchCreateLogs(t *testing.T) {
295296
},
296297
Database: dbM,
297298
Log: slogtest.Make(t, nil),
298-
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
299+
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
299300
publishWorkspaceUpdateCalled = true
300301
return nil
301302
},
@@ -339,7 +340,7 @@ func TestBatchCreateLogs(t *testing.T) {
339340
},
340341
Database: dbM,
341342
Log: slogtest.Make(t, nil),
342-
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
343+
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
343344
publishWorkspaceUpdateCalled = true
344345
return nil
345346
},
@@ -386,7 +387,7 @@ func TestBatchCreateLogs(t *testing.T) {
386387
},
387388
Database: dbM,
388389
Log: slogtest.Make(t, nil),
389-
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
390+
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
390391
publishWorkspaceUpdateCalled = true
391392
return nil
392393
},

‎coderd/agentapi/stats_test.go

Lines changed: 22 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ import (
1313
"go.uber.org/mock/gomock"
1414
"google.golang.org/protobuf/types/known/durationpb"
1515

16+
"cdr.dev/slog/sloggers/slogtest"
17+
1618
agentproto "github.com/coder/coder/v2/agent/proto"
1719
"github.com/coder/coder/v2/coderd/agentapi"
1820
"github.com/coder/coder/v2/coderd/database"
@@ -156,13 +158,15 @@ func TestUpdateStates(t *testing.T) {
156158
// Ensure that pubsub notifications are sent.
157159
notifyDescription := make(chan struct{})
158160
ps.Subscribe(wspubsub.WorkspaceEventChannel(workspace.OwnerID),
159-
wspubsub.HandleWorkspaceEvent(func(_ context.Context, e wspubsub.WorkspaceEvent) {
160-
if e.Kind == wspubsub.WorkspaceEventKindUpdatedStats && e.WorkspaceID == workspace.ID {
161-
go func() {
162-
notifyDescription <- struct{}{}
163-
}()
164-
}
165-
}))
161+
wspubsub.HandleWorkspaceEvent(
162+
slogtest.Make(t, nil),
163+
func(_ context.Context, e wspubsub.WorkspaceEvent) {
164+
if e.Kind == wspubsub.WorkspaceEventKindStatsUpdate && e.WorkspaceID == workspace.ID {
165+
go func() {
166+
notifyDescription <- struct{}{}
167+
}()
168+
}
169+
}))
166170

167171
resp, err := api.UpdateStats(context.Background(), req)
168172
require.NoError(t, err)
@@ -187,8 +191,7 @@ func TestUpdateStates(t *testing.T) {
187191
select {
188192
case <-ctx.Done():
189193
t.Error("timed out while waiting for pubsub notification")
190-
case description := <-notifyDescription:
191-
require.Equal(t, description, struct{}{})
194+
case <-notifyDescription:
192195
}
193196
require.True(t, updateAgentMetricsFnCalled)
194197
})
@@ -501,13 +504,15 @@ func TestUpdateStates(t *testing.T) {
501504
// Ensure that pubsub notifications are sent.
502505
notifyDescription := make(chan struct{})
503506
ps.Subscribe(wspubsub.WorkspaceEventChannel(workspace.OwnerID),
504-
wspubsub.HandleWorkspaceEvent(func(_ context.Context, e wspubsub.WorkspaceEvent) {
505-
if e.Kind == wspubsub.WorkspaceEventKindUpdatedStats && e.WorkspaceID == workspace.ID {
506-
go func() {
507-
notifyDescription <- struct{}{}
508-
}()
509-
}
510-
}))
507+
wspubsub.HandleWorkspaceEvent(
508+
slogtest.Make(t, nil),
509+
func(_ context.Context, e wspubsub.WorkspaceEvent) {
510+
if e.Kind == wspubsub.WorkspaceEventKindStatsUpdate && e.WorkspaceID == workspace.ID {
511+
go func() {
512+
notifyDescription <- struct{}{}
513+
}()
514+
}
515+
}))
511516

512517
resp, err := api.UpdateStats(context.Background(), req)
513518
require.NoError(t, err)
@@ -530,8 +535,7 @@ func TestUpdateStates(t *testing.T) {
530535
select {
531536
case <-ctx.Done():
532537
t.Error("timed out while waiting for pubsub notification")
533-
case description := <-notifyDescription:
534-
require.Equal(t, description, struct{}{})
538+
case <-notifyDescription:
535539
}
536540
require.True(t, updateAgentMetricsFnCalled)
537541
})

‎coderd/database/dbfake/dbfake.go

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ import (
1919
"github.com/coder/coder/v2/coderd/provisionerdserver"
2020
"github.com/coder/coder/v2/coderd/rbac"
2121
"github.com/coder/coder/v2/coderd/telemetry"
22-
"github.com/coder/coder/v2/coderd/util/ptr"
2322
"github.com/coder/coder/v2/coderd/wspubsub"
2423
"github.com/coder/coder/v2/provisionersdk"
2524
sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
@@ -227,14 +226,11 @@ func (b WorkspaceBuildBuilder) Do() WorkspaceResponse {
227226

228227
if b.ps != nil {
229228
msg, err := json.Marshal(wspubsub.WorkspaceEvent{
230-
Kind: wspubsub.WorkspaceEventKindStateChange,
231-
WorkspaceID: resp.Workspace.ID,
232-
WorkspaceName: ptr.Ref(resp.Workspace.Name),
233-
Transition: ptr.Ref(resp.Build.Transition),
234-
JobStatus: ptr.Ref(job.JobStatus),
229+
Kind: wspubsub.WorkspaceEventKindStateChange,
230+
WorkspaceID: resp.Workspace.ID,
235231
})
236232
require.NoError(b.t, err)
237-
err = b.ps.Publish(wspubsub.WorkspaceEventChannel(resp.Build.InitiatorID), msg)
233+
err = b.ps.Publish(wspubsub.WorkspaceEventChannel(resp.Workspace.OwnerID), msg)
238234
require.NoError(b.t, err)
239235
}
240236

0 commit comments

Comments
 (0)