Skip to content

Commit d90035e

Browse files
mtojekmtojek
committed
Merge branch 'main' into 7964-order-agents
2 parents e93463d + 5a0d240 commit d90035e

File tree

83 files changed

+3343
-307
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

83 files changed

+3343
-307
lines changed

‎.github/workflows/ci.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ jobs:
141141
142142
# Check for any typos
143143
- name: Check for typos
144-
uses: crate-ci/typos@v1.18.0
144+
uses: crate-ci/typos@v1.18.2
145145
with:
146146
config: .github/workflows/typos.toml
147147

‎.github/workflows/release.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -411,7 +411,7 @@ jobs:
411411
- name: Sync fork
412412
run: gh repo sync cdrci/winget-pkgs -b master
413413
env:
414-
GH_TOKEN: ${{ github.token }}
414+
GH_TOKEN: ${{ secrets.CDRCI_GITHUB_TOKEN }}
415415

416416
- name: Checkout
417417
uses: actions/checkout@v4

‎.github/workflows/security.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ jobs:
122122
image_name: ${{ steps.build.outputs.image }}
123123

124124
- name: Run Trivy vulnerability scanner
125-
uses: aquasecurity/trivy-action@d43c1f16c00cfd3978dde6c07f4bbcf9eb6993ca
125+
uses: aquasecurity/trivy-action@84384bd6e777ef152729993b8145ea352e9dd3ef
126126
with:
127127
image-ref: ${{ steps.build.outputs.image }}
128128
format: sarif

‎coderd/agentapi/api.go

+1
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@ func New(opts Options) *API {
114114
api.StatsAPI = &StatsAPI{
115115
AgentFn: api.agent,
116116
Database: opts.Database,
117+
Pubsub: opts.Pubsub,
117118
Log: opts.Log,
118119
StatsBatcher: opts.StatsBatcher,
119120
TemplateScheduleStore: opts.TemplateScheduleStore,

‎coderd/agentapi/stats.go

+14
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,10 @@ import (
1616
"github.com/coder/coder/v2/coderd/autobuild"
1717
"github.com/coder/coder/v2/coderd/database"
1818
"github.com/coder/coder/v2/coderd/database/dbtime"
19+
"github.com/coder/coder/v2/coderd/database/pubsub"
1920
"github.com/coder/coder/v2/coderd/prometheusmetrics"
2021
"github.com/coder/coder/v2/coderd/schedule"
22+
"github.com/coder/coder/v2/codersdk"
2123
)
2224

2325
type StatsBatcher interface {
@@ -27,6 +29,7 @@ type StatsBatcher interface {
2729
type StatsAPI struct {
2830
AgentFn func(context.Context) (database.WorkspaceAgent, error)
2931
Database database.Store
32+
Pubsub pubsub.Pubsub
3033
Log slog.Logger
3134
StatsBatcher StatsBatcher
3235
TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
@@ -130,5 +133,16 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
130133
return nil, xerrors.Errorf("update stats in database: %w", err)
131134
}
132135

136+
// Tell the frontend about the new agent report, now that everything is updated
137+
a.publishWorkspaceAgentStats(ctx, workspace.ID)
138+
133139
return res, nil
134140
}
141+
142+
func (a *StatsAPI) publishWorkspaceAgentStats(ctx context.Context, workspaceID uuid.UUID) {
143+
err := a.Pubsub.Publish(codersdk.WorkspaceNotifyChannel(workspaceID), codersdk.WorkspaceNotifyDescriptionAgentStatsOnly)
144+
if err != nil {
145+
a.Log.Warn(ctx, "failed to publish workspace agent stats",
146+
slog.F("workspace_id", workspaceID), slog.Error(err))
147+
}
148+
}

‎coderd/agentapi/stats_test.go

+40-11
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package agentapi_test
22

33
import (
4+
"bytes"
45
"context"
56
"database/sql"
67
"sync"
@@ -19,8 +20,11 @@ import (
1920
"github.com/coder/coder/v2/coderd/database"
2021
"github.com/coder/coder/v2/coderd/database/dbmock"
2122
"github.com/coder/coder/v2/coderd/database/dbtime"
23+
"github.com/coder/coder/v2/coderd/database/pubsub"
2224
"github.com/coder/coder/v2/coderd/prometheusmetrics"
2325
"github.com/coder/coder/v2/coderd/schedule"
26+
"github.com/coder/coder/v2/codersdk"
27+
"github.com/coder/coder/v2/testutil"
2428
)
2529

2630
type statsBatcher struct {
@@ -78,8 +82,10 @@ func TestUpdateStates(t *testing.T) {
7882
t.Parallel()
7983

8084
var (
81-
now = dbtime.Now()
82-
dbM = dbmock.NewMockStore(gomock.NewController(t))
85+
now = dbtime.Now()
86+
dbM = dbmock.NewMockStore(gomock.NewController(t))
87+
ps = pubsub.NewInMemory()
88+
8389
templateScheduleStore = schedule.MockTemplateScheduleStore{
8490
GetFn: func(context.Context, database.Store, uuid.UUID) (schedule.TemplateScheduleOptions, error) {
8591
panic("should not be called")
@@ -125,6 +131,7 @@ func TestUpdateStates(t *testing.T) {
125131
return agent, nil
126132
},
127133
Database: dbM,
134+
Pubsub: ps,
128135
StatsBatcher: batcher,
129136
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
130137
AgentStatsRefreshInterval: 10 * time.Second,
@@ -164,6 +171,15 @@ func TestUpdateStates(t *testing.T) {
164171
// User gets fetched to hit the UpdateAgentMetricsFn.
165172
dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
166173

174+
// Ensure that pubsub notifications are sent.
175+
publishAgentStats := make(chan bool)
176+
ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, description []byte) {
177+
go func() {
178+
publishAgentStats <- bytes.Equal(description, codersdk.WorkspaceNotifyDescriptionAgentStatsOnly)
179+
close(publishAgentStats)
180+
}()
181+
})
182+
167183
resp, err := api.UpdateStats(context.Background(), req)
168184
require.NoError(t, err)
169185
require.Equal(t, &agentproto.UpdateStatsResponse{
@@ -179,7 +195,13 @@ func TestUpdateStates(t *testing.T) {
179195
require.Equal(t, user.ID, batcher.lastUserID)
180196
require.Equal(t, workspace.ID, batcher.lastWorkspaceID)
181197
require.Equal(t, req.Stats, batcher.lastStats)
182-
198+
ctx := testutil.Context(t, testutil.WaitShort)
199+
select {
200+
case <-ctx.Done():
201+
t.Error("timed out while waiting for pubsub notification")
202+
case wasAgentStatsOnly := <-publishAgentStats:
203+
require.Equal(t, wasAgentStatsOnly, true)
204+
}
183205
require.True(t, updateAgentMetricsFnCalled)
184206
})
185207

@@ -189,6 +211,7 @@ func TestUpdateStates(t *testing.T) {
189211
var (
190212
now = dbtime.Now()
191213
dbM = dbmock.NewMockStore(gomock.NewController(t))
214+
ps = pubsub.NewInMemory()
192215
templateScheduleStore = schedule.MockTemplateScheduleStore{
193216
GetFn: func(context.Context, database.Store, uuid.UUID) (schedule.TemplateScheduleOptions, error) {
194217
panic("should not be called")
@@ -214,6 +237,7 @@ func TestUpdateStates(t *testing.T) {
214237
return agent, nil
215238
},
216239
Database: dbM,
240+
Pubsub: ps,
217241
StatsBatcher: batcher,
218242
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
219243
AgentStatsRefreshInterval: 10 * time.Second,
@@ -244,7 +268,8 @@ func TestUpdateStates(t *testing.T) {
244268
t.Parallel()
245269

246270
var (
247-
dbM = dbmock.NewMockStore(gomock.NewController(t))
271+
db = dbmock.NewMockStore(gomock.NewController(t))
272+
ps = pubsub.NewInMemory()
248273
req = &agentproto.UpdateStatsRequest{
249274
Stats: &agentproto.Stats{
250275
ConnectionsByProto: map[string]int64{}, // len() == 0
@@ -255,7 +280,8 @@ func TestUpdateStates(t *testing.T) {
255280
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
256281
return agent, nil
257282
},
258-
Database: dbM,
283+
Database: db,
284+
Pubsub: ps,
259285
StatsBatcher: nil, // should not be called
260286
TemplateScheduleStore: nil, // should not be called
261287
AgentStatsRefreshInterval: 10 * time.Second,
@@ -290,7 +316,9 @@ func TestUpdateStates(t *testing.T) {
290316
nextAutostart := now.Add(30 * time.Minute).UTC() // always sent to DB as UTC
291317

292318
var (
293-
dbM = dbmock.NewMockStore(gomock.NewController(t))
319+
db = dbmock.NewMockStore(gomock.NewController(t))
320+
ps = pubsub.NewInMemory()
321+
294322
templateScheduleStore = schedule.MockTemplateScheduleStore{
295323
GetFn: func(context.Context, database.Store, uuid.UUID) (schedule.TemplateScheduleOptions, error) {
296324
return schedule.TemplateScheduleOptions{
@@ -321,7 +349,8 @@ func TestUpdateStates(t *testing.T) {
321349
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
322350
return agent, nil
323351
},
324-
Database: dbM,
352+
Database: db,
353+
Pubsub: ps,
325354
StatsBatcher: batcher,
326355
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
327356
AgentStatsRefreshInterval: 15 * time.Second,
@@ -341,26 +370,26 @@ func TestUpdateStates(t *testing.T) {
341370
}
342371

343372
// Workspace gets fetched.
344-
dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(database.GetWorkspaceByAgentIDRow{
373+
db.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(database.GetWorkspaceByAgentIDRow{
345374
Workspace: workspace,
346375
TemplateName: template.Name,
347376
}, nil)
348377

349378
// We expect an activity bump because ConnectionCount > 0. However, the
350379
// next autostart time will be set on the bump.
351-
dbM.EXPECT().ActivityBumpWorkspace(gomock.Any(), database.ActivityBumpWorkspaceParams{
380+
db.EXPECT().ActivityBumpWorkspace(gomock.Any(), database.ActivityBumpWorkspaceParams{
352381
WorkspaceID: workspace.ID,
353382
NextAutostart: nextAutostart,
354383
}).Return(nil)
355384

356385
// Workspace last used at gets bumped.
357-
dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
386+
db.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
358387
ID: workspace.ID,
359388
LastUsedAt: now,
360389
}).Return(nil)
361390

362391
// User gets fetched to hit the UpdateAgentMetricsFn.
363-
dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
392+
db.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
364393

365394
resp, err := api.UpdateStats(context.Background(), req)
366395
require.NoError(t, err)

0 commit comments

Comments
 (0)