Skip to content

Commit 701c1b6

Browse files
ethanndicksonethanndickson
committed
chore: send workspace pubsub events by owner id
1 parent b22bd81 commit 701c1b6

18 files changed

+355
-199
lines changed

‎coderd/agentapi/api.go

Lines changed: 12 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,9 @@ type API struct {
5252
var _ agentproto.DRPCAgentServer = &API{}
5353

5454
type Options struct {
55-
AgentID uuid.UUID
55+
AgentID uuid.UUID
56+
OwnerID uuid.UUID
57+
WorkspaceID uuid.UUID
5658

5759
Ctx context.Context
5860
Log slog.Logger
@@ -62,7 +64,7 @@ type Options struct {
6264
TailnetCoordinator *atomic.Pointer[tailnet.Coordinator]
6365
StatsReporter *workspacestats.Reporter
6466
AppearanceFetcher *atomic.Pointer[appearance.Fetcher]
65-
PublishWorkspaceUpdateFn func(ctx context.Context, workspaceID uuid.UUID)
67+
PublishWorkspaceUpdateFn func(ctx context.Context, userID uuid.UUID, event codersdk.WorkspaceEvent)
6668
PublishWorkspaceAgentLogsUpdateFn func(ctx context.Context, workspaceAgentID uuid.UUID, msg agentsdk.LogsNotifyMessage)
6769
NetworkTelemetryHandler func(batch []*tailnetproto.TelemetryEvent)
6870

@@ -75,10 +77,6 @@ type Options struct {
7577
ExternalAuthConfigs []*externalauth.Config
7678
Experiments codersdk.Experiments
7779

78-
// Optional:
79-
// WorkspaceID avoids a future lookup to find the workspace ID by setting
80-
// the cache in advance.
81-
WorkspaceID uuid.UUID
8280
UpdateAgentMetricsFn func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
8381
}
8482

@@ -98,16 +96,7 @@ func New(opts Options) *API {
9896
AgentFn: api.agent,
9997
Database: opts.Database,
10098
DerpMapFn: opts.DerpMapFn,
101-
WorkspaceIDFn: func(ctx context.Context, wa *database.WorkspaceAgent) (uuid.UUID, error) {
102-
if opts.WorkspaceID != uuid.Nil {
103-
return opts.WorkspaceID, nil
104-
}
105-
ws, err := opts.Database.GetWorkspaceByAgentID(ctx, wa.ID)
106-
if err != nil {
107-
return uuid.Nil, err
108-
}
109-
return ws.Workspace.ID, nil
110-
},
99+
WorkspaceID: opts.WorkspaceID,
111100
}
112101

113102
api.AnnouncementBannerAPI = &AnnouncementBannerAPI{
@@ -125,7 +114,7 @@ func New(opts Options) *API {
125114

126115
api.LifecycleAPI = &LifecycleAPI{
127116
AgentFn: api.agent,
128-
WorkspaceIDFn: api.workspaceID,
117+
WorkspaceID: opts.WorkspaceID,
129118
Database: opts.Database,
130119
Log: opts.Log,
131120
PublishWorkspaceUpdateFn: api.publishWorkspaceUpdate,
@@ -242,6 +231,11 @@ func (a *API) publishWorkspaceUpdate(ctx context.Context, agent *database.Worksp
242231
return err
243232
}
244233

245-
a.opts.PublishWorkspaceUpdateFn(ctx, workspaceID)
234+
a.opts.PublishWorkspaceUpdateFn(ctx, a.opts.OwnerID, codersdk.WorkspaceEvent{
235+
Kind: codersdk.WorkspaceEventKindAgentUpdate,
236+
WorkspaceID: workspaceID,
237+
AgentID: &agent.ID,
238+
AgentName: &agent.Name,
239+
})
246240
return nil
247241
}

‎coderd/agentapi/lifecycle.go

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ func WithAPIVersion(ctx context.Context, version string) context.Context {
2525

2626
type LifecycleAPI struct {
2727
AgentFn func(context.Context) (database.WorkspaceAgent, error)
28-
WorkspaceIDFn func(context.Context, *database.WorkspaceAgent) (uuid.UUID, error)
28+
WorkspaceID uuid.UUID
2929
Database database.Store
3030
Log slog.Logger
3131
PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent) error
@@ -45,13 +45,9 @@ func (a *LifecycleAPI) UpdateLifecycle(ctx context.Context, req *agentproto.Upda
4545
if err != nil {
4646
return nil, err
4747
}
48-
workspaceID, err := a.WorkspaceIDFn(ctx, &workspaceAgent)
49-
if err != nil {
50-
return nil, err
51-
}
5248

5349
logger := a.Log.With(
54-
slog.F("workspace_id", workspaceID),
50+
slog.F("workspace_id", a.WorkspaceID),
5551
slog.F("payload", req),
5652
)
5753
logger.Debug(ctx, "workspace agent state report")
@@ -140,15 +136,11 @@ func (a *LifecycleAPI) UpdateStartup(ctx context.Context, req *agentproto.Update
140136
if err != nil {
141137
return nil, err
142138
}
143-
workspaceID, err := a.WorkspaceIDFn(ctx, &workspaceAgent)
144-
if err != nil {
145-
return nil, err
146-
}
147139

148140
a.Log.Debug(
149141
ctx,
150142
"post workspace agent version",
151-
slog.F("workspace_id", workspaceID),
143+
slog.F("workspace_id", a.WorkspaceID),
152144
slog.F("agent_version", req.Startup.Version),
153145
)
154146

‎coderd/agentapi/lifecycle_test.go

Lines changed: 25 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -69,11 +69,9 @@ func TestUpdateLifecycle(t *testing.T) {
6969
AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
7070
return agentCreated, nil
7171
},
72-
WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
73-
return workspaceID, nil
74-
},
75-
Database: dbM,
76-
Log: slogtest.Make(t, nil),
72+
WorkspaceID: workspaceID,
73+
Database: dbM,
74+
Log: slogtest.Make(t, nil),
7775
PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent) error {
7876
publishCalled = true
7977
return nil
@@ -111,11 +109,9 @@ func TestUpdateLifecycle(t *testing.T) {
111109
AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
112110
return agentStarting, nil
113111
},
114-
WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
115-
return workspaceID, nil
116-
},
117-
Database: dbM,
118-
Log: slogtest.Make(t, nil),
112+
WorkspaceID: workspaceID,
113+
Database: dbM,
114+
Log: slogtest.Make(t, nil),
119115
// Test that nil publish fn works.
120116
PublishWorkspaceUpdateFn: nil,
121117
}
@@ -156,11 +152,9 @@ func TestUpdateLifecycle(t *testing.T) {
156152
AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
157153
return agentCreated, nil
158154
},
159-
WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
160-
return workspaceID, nil
161-
},
162-
Database: dbM,
163-
Log: slogtest.Make(t, nil),
155+
WorkspaceID: workspaceID,
156+
Database: dbM,
157+
Log: slogtest.Make(t, nil),
164158
PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent) error {
165159
publishCalled = true
166160
return nil
@@ -204,9 +198,7 @@ func TestUpdateLifecycle(t *testing.T) {
204198
AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
205199
return agentCreated, nil
206200
},
207-
WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
208-
return workspaceID, nil
209-
},
201+
WorkspaceID: workspaceID,
210202
Database: dbM,
211203
Log: slogtest.Make(t, nil),
212204
PublishWorkspaceUpdateFn: nil,
@@ -239,11 +231,9 @@ func TestUpdateLifecycle(t *testing.T) {
239231
AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
240232
return agent, nil
241233
},
242-
WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
243-
return workspaceID, nil
244-
},
245-
Database: dbM,
246-
Log: slogtest.Make(t, nil),
234+
WorkspaceID: workspaceID,
235+
Database: dbM,
236+
Log: slogtest.Make(t, nil),
247237
PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent) error {
248238
atomic.AddInt64(&publishCalled, 1)
249239
return nil
@@ -314,11 +304,9 @@ func TestUpdateLifecycle(t *testing.T) {
314304
AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
315305
return agentCreated, nil
316306
},
317-
WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
318-
return workspaceID, nil
319-
},
320-
Database: dbM,
321-
Log: slogtest.Make(t, nil),
307+
WorkspaceID: workspaceID,
308+
Database: dbM,
309+
Log: slogtest.Make(t, nil),
322310
PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent) error {
323311
publishCalled = true
324312
return nil
@@ -354,11 +342,9 @@ func TestUpdateStartup(t *testing.T) {
354342
AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
355343
return agent, nil
356344
},
357-
WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
358-
return workspaceID, nil
359-
},
360-
Database: dbM,
361-
Log: slogtest.Make(t, nil),
345+
WorkspaceID: workspaceID,
346+
Database: dbM,
347+
Log: slogtest.Make(t, nil),
362348
// Not used by UpdateStartup.
363349
PublishWorkspaceUpdateFn: nil,
364350
}
@@ -402,11 +388,9 @@ func TestUpdateStartup(t *testing.T) {
402388
AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
403389
return agent, nil
404390
},
405-
WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
406-
return workspaceID, nil
407-
},
408-
Database: dbM,
409-
Log: slogtest.Make(t, nil),
391+
WorkspaceID: workspaceID,
392+
Database: dbM,
393+
Log: slogtest.Make(t, nil),
410394
// Not used by UpdateStartup.
411395
PublishWorkspaceUpdateFn: nil,
412396
}
@@ -435,11 +419,9 @@ func TestUpdateStartup(t *testing.T) {
435419
AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
436420
return agent, nil
437421
},
438-
WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
439-
return workspaceID, nil
440-
},
441-
Database: dbM,
442-
Log: slogtest.Make(t, nil),
422+
WorkspaceID: workspaceID,
423+
Database: dbM,
424+
Log: slogtest.Make(t, nil),
443425
// Not used by UpdateStartup.
444426
PublishWorkspaceUpdateFn: nil,
445427
}

‎coderd/agentapi/manifest.go

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,23 +29,18 @@ type ManifestAPI struct {
2929
ExternalAuthConfigs []*externalauth.Config
3030
DisableDirectConnections bool
3131
DerpForceWebSockets bool
32+
WorkspaceID uuid.UUID
3233

33-
AgentFn func(context.Context) (database.WorkspaceAgent, error)
34-
WorkspaceIDFn func(context.Context, *database.WorkspaceAgent) (uuid.UUID, error)
35-
Database database.Store
36-
DerpMapFn func() *tailcfg.DERPMap
34+
AgentFn func(context.Context) (database.WorkspaceAgent, error)
35+
Database database.Store
36+
DerpMapFn func() *tailcfg.DERPMap
3737
}
3838

3939
func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifestRequest) (*agentproto.Manifest, error) {
4040
workspaceAgent, err := a.AgentFn(ctx)
4141
if err != nil {
4242
return nil, err
4343
}
44-
workspaceID, err := a.WorkspaceIDFn(ctx, &workspaceAgent)
45-
if err != nil {
46-
return nil, err
47-
}
48-
4944
var (
5045
dbApps []database.WorkspaceApp
5146
scripts []database.WorkspaceAgentScript
@@ -75,7 +70,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
7570
return err
7671
})
7772
eg.Go(func() (err error) {
78-
workspace, err = a.Database.GetWorkspaceByID(ctx, workspaceID)
73+
workspace, err = a.Database.GetWorkspaceByID(ctx, a.WorkspaceID)
7974
if err != nil {
8075
return xerrors.Errorf("getting workspace by id: %w", err)
8176
}

‎coderd/agentapi/manifest_test.go

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -288,11 +288,9 @@ func TestGetManifest(t *testing.T) {
288288
AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
289289
return agent, nil
290290
},
291-
WorkspaceIDFn: func(ctx context.Context, _ *database.WorkspaceAgent) (uuid.UUID, error) {
292-
return workspace.ID, nil
293-
},
294-
Database: mDB,
295-
DerpMapFn: derpMapFn,
291+
WorkspaceID: workspace.ID,
292+
Database: mDB,
293+
DerpMapFn: derpMapFn,
296294
}
297295

298296
mDB.EXPECT().GetWorkspaceAppsByAgentID(gomock.Any(), agent.ID).Return(apps, nil)
@@ -355,11 +353,9 @@ func TestGetManifest(t *testing.T) {
355353
AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
356354
return agent, nil
357355
},
358-
WorkspaceIDFn: func(ctx context.Context, _ *database.WorkspaceAgent) (uuid.UUID, error) {
359-
return workspace.ID, nil
360-
},
361-
Database: mDB,
362-
DerpMapFn: derpMapFn,
356+
WorkspaceID: workspace.ID,
357+
Database: mDB,
358+
DerpMapFn: derpMapFn,
363359
}
364360

365361
mDB.EXPECT().GetWorkspaceAppsByAgentID(gomock.Any(), agent.ID).Return(apps, nil)

‎coderd/agentapi/stats_test.go

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -148,12 +148,15 @@ func TestUpdateStates(t *testing.T) {
148148
dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
149149

150150
// Ensure that pubsub notifications are sent.
151-
notifyDescription := make(chan []byte)
152-
ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, description []byte) {
153-
go func() {
154-
notifyDescription <- description
155-
}()
156-
})
151+
notifyDescription := make(chan struct{})
152+
ps.Subscribe(codersdk.WorkspaceEventChannel(workspace.OwnerID),
153+
codersdk.HandleWorkspaceEvent(func(_ context.Context, e codersdk.WorkspaceEvent) {
154+
if e.Kind == codersdk.WorkspaceEventKindUpdatedStats && e.WorkspaceID == workspace.ID {
155+
go func() {
156+
notifyDescription <- struct{}{}
157+
}()
158+
}
159+
}))
157160

158161
resp, err := api.UpdateStats(context.Background(), req)
159162
require.NoError(t, err)
@@ -175,7 +178,7 @@ func TestUpdateStates(t *testing.T) {
175178
case <-ctx.Done():
176179
t.Error("timed out while waiting for pubsub notification")
177180
case description := <-notifyDescription:
178-
require.Equal(t, description, []byte{})
181+
require.Equal(t, description, struct{}{})
179182
}
180183
require.True(t, updateAgentMetricsFnCalled)
181184
})
@@ -482,12 +485,15 @@ func TestUpdateStates(t *testing.T) {
482485
dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
483486

484487
// Ensure that pubsub notifications are sent.
485-
notifyDescription := make(chan []byte)
486-
ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, description []byte) {
487-
go func() {
488-
notifyDescription <- description
489-
}()
490-
})
488+
notifyDescription := make(chan struct{})
489+
ps.Subscribe(codersdk.WorkspaceEventChannel(workspace.OwnerID),
490+
codersdk.HandleWorkspaceEvent(func(_ context.Context, e codersdk.WorkspaceEvent) {
491+
if e.Kind == codersdk.WorkspaceEventKindUpdatedStats && e.WorkspaceID == workspace.ID {
492+
go func() {
493+
notifyDescription <- struct{}{}
494+
}()
495+
}
496+
}))
491497

492498
resp, err := api.UpdateStats(context.Background(), req)
493499
require.NoError(t, err)
@@ -507,7 +513,7 @@ func TestUpdateStates(t *testing.T) {
507513
case <-ctx.Done():
508514
t.Error("timed out while waiting for pubsub notification")
509515
case description := <-notifyDescription:
510-
require.Equal(t, description, []byte{})
516+
require.Equal(t, description, struct{}{})
511517
}
512518
require.True(t, updateAgentMetricsFnCalled)
513519
})

‎coderd/database/dbfake/dbfake.go

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ 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"
2223
"github.com/coder/coder/v2/codersdk"
2324
"github.com/coder/coder/v2/provisionersdk"
2425
sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
@@ -225,7 +226,14 @@ func (b WorkspaceBuildBuilder) Do() WorkspaceResponse {
225226
_ = dbgen.WorkspaceBuildParameters(b.t, b.db, b.params)
226227

227228
if b.ps != nil {
228-
err = b.ps.Publish(codersdk.WorkspaceNotifyChannel(resp.Build.WorkspaceID), []byte{})
229+
msg, err := json.Marshal(codersdk.WorkspaceEvent{
230+
Kind: codersdk.WorkspaceEventKindStateChange,
231+
WorkspaceID: resp.Workspace.ID,
232+
Transition: ptr.Ref(codersdk.WorkspaceTransition(resp.Build.Transition)),
233+
JobStatus: ptr.Ref(codersdk.ProvisionerJobStatus(job.JobStatus)),
234+
})
235+
require.NoError(b.t, err)
236+
err = b.ps.Publish(codersdk.WorkspaceEventChannel(resp.Build.InitiatorID), msg)
229237
require.NoError(b.t, err)
230238
}
231239

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy