Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 0dd942e

Browse files
f0sself0ssel
authored
fix: stop incrementing activity on empty agent stats (coder#15204)
1 parent cd890aa commit 0dd942e

File tree

8 files changed

+107
-86
lines changed

8 files changed

+107
-86
lines changed

‎coderd/agentapi/stats_test.go

+44-16
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,11 @@ func TestUpdateStates(t *testing.T) {
7070
}
7171
batcher = &workspacestatstest.StatsBatcher{}
7272
updateAgentMetricsFnCalled = false
73+
tickCh = make(chan time.Time)
74+
flushCh = make(chan int, 1)
75+
wut = workspacestats.NewTracker(dbM,
76+
workspacestats.TrackerWithTickFlush(tickCh, flushCh),
77+
)
7378

7479
req = &agentproto.UpdateStatsRequest{
7580
Stats: &agentproto.Stats{
@@ -109,6 +114,7 @@ func TestUpdateStates(t *testing.T) {
109114
Database: dbM,
110115
Pubsub: ps,
111116
StatsBatcher: batcher,
117+
UsageTracker: wut,
112118
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
113119
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
114120
updateAgentMetricsFnCalled = true
@@ -126,25 +132,26 @@ func TestUpdateStates(t *testing.T) {
126132
return now
127133
},
128134
}
135+
defer wut.Close()
129136

130137
// Workspace gets fetched.
131138
dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(workspace, nil)
132139

140+
// User gets fetched to hit the UpdateAgentMetricsFn.
141+
dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
142+
133143
// We expect an activity bump because ConnectionCount > 0.
134144
dbM.EXPECT().ActivityBumpWorkspace(gomock.Any(), database.ActivityBumpWorkspaceParams{
135145
WorkspaceID: workspace.ID,
136146
NextAutostart: time.Time{}.UTC(),
137147
}).Return(nil)
138148

139149
// Workspace last used at gets bumped.
140-
dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
141-
ID: workspace.ID,
150+
dbM.EXPECT().BatchUpdateWorkspaceLastUsedAt(gomock.Any(), database.BatchUpdateWorkspaceLastUsedAtParams{
151+
IDs: []uuid.UUID{workspace.ID},
142152
LastUsedAt: now,
143153
}).Return(nil)
144154

145-
// User gets fetched to hit the UpdateAgentMetricsFn.
146-
dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
147-
148155
// Ensure that pubsub notifications are sent.
149156
notifyDescription := make(chan []byte)
150157
ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, description []byte) {
@@ -159,6 +166,10 @@ func TestUpdateStates(t *testing.T) {
159166
ReportInterval: durationpb.New(10 * time.Second),
160167
}, resp)
161168

169+
tickCh <- now
170+
count := <-flushCh
171+
require.Equal(t, 1, count, "expected one flush with one id")
172+
162173
batcher.Mu.Lock()
163174
defer batcher.Mu.Unlock()
164175
require.Equal(t, int64(1), batcher.Called)
@@ -211,6 +222,7 @@ func TestUpdateStates(t *testing.T) {
211222
StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
212223
Database: dbM,
213224
Pubsub: ps,
225+
UsageTracker: workspacestats.NewTracker(dbM),
214226
StatsBatcher: batcher,
215227
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
216228
// Ignored when nil.
@@ -225,12 +237,6 @@ func TestUpdateStates(t *testing.T) {
225237
// Workspace gets fetched.
226238
dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(workspace, nil)
227239

228-
// Workspace last used at gets bumped.
229-
dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
230-
ID: workspace.ID,
231-
LastUsedAt: now,
232-
}).Return(nil)
233-
234240
_, err := api.UpdateStats(context.Background(), req)
235241
require.NoError(t, err)
236242
})
@@ -306,6 +312,11 @@ func TestUpdateStates(t *testing.T) {
306312
}
307313
batcher = &workspacestatstest.StatsBatcher{}
308314
updateAgentMetricsFnCalled = false
315+
tickCh = make(chan time.Time)
316+
flushCh = make(chan int, 1)
317+
wut = workspacestats.NewTracker(dbM,
318+
workspacestats.TrackerWithTickFlush(tickCh, flushCh),
319+
)
309320

310321
req = &agentproto.UpdateStatsRequest{
311322
Stats: &agentproto.Stats{
@@ -325,6 +336,7 @@ func TestUpdateStates(t *testing.T) {
325336
StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
326337
Database: dbM,
327338
Pubsub: ps,
339+
UsageTracker: wut,
328340
StatsBatcher: batcher,
329341
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
330342
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
@@ -343,6 +355,7 @@ func TestUpdateStates(t *testing.T) {
343355
return now
344356
},
345357
}
358+
defer wut.Close()
346359

347360
// Workspace gets fetched.
348361
dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(workspace, nil)
@@ -355,9 +368,9 @@ func TestUpdateStates(t *testing.T) {
355368
}).Return(nil)
356369

357370
// Workspace last used at gets bumped.
358-
dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
359-
ID: workspace.ID,
360-
LastUsedAt: now,
371+
dbM.EXPECT().BatchUpdateWorkspaceLastUsedAt(gomock.Any(), database.BatchUpdateWorkspaceLastUsedAtParams{
372+
IDs: []uuid.UUID{workspace.ID},
373+
LastUsedAt: now.UTC(),
361374
}).Return(nil)
362375

363376
// User gets fetched to hit the UpdateAgentMetricsFn.
@@ -369,6 +382,10 @@ func TestUpdateStates(t *testing.T) {
369382
ReportInterval: durationpb.New(15 * time.Second),
370383
}, resp)
371384

385+
tickCh <- now
386+
count := <-flushCh
387+
require.Equal(t, 1, count, "expected one flush with one id")
388+
372389
require.True(t, updateAgentMetricsFnCalled)
373390
})
374391

@@ -392,6 +409,11 @@ func TestUpdateStates(t *testing.T) {
392409
}
393410
batcher = &workspacestatstest.StatsBatcher{}
394411
updateAgentMetricsFnCalled = false
412+
tickCh = make(chan time.Time)
413+
flushCh = make(chan int, 1)
414+
wut = workspacestats.NewTracker(dbM,
415+
workspacestats.TrackerWithTickFlush(tickCh, flushCh),
416+
)
395417

396418
req = &agentproto.UpdateStatsRequest{
397419
Stats: &agentproto.Stats{
@@ -422,6 +444,7 @@ func TestUpdateStates(t *testing.T) {
422444
},
423445
}
424446
)
447+
defer wut.Close()
425448
api := agentapi.StatsAPI{
426449
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
427450
return agent, nil
@@ -431,6 +454,7 @@ func TestUpdateStates(t *testing.T) {
431454
Database: dbM,
432455
Pubsub: ps,
433456
StatsBatcher: batcher,
457+
UsageTracker: wut,
434458
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
435459
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
436460
updateAgentMetricsFnCalled = true
@@ -462,8 +486,8 @@ func TestUpdateStates(t *testing.T) {
462486
}).Return(nil)
463487

464488
// Workspace last used at gets bumped.
465-
dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
466-
ID: workspace.ID,
489+
dbM.EXPECT().BatchUpdateWorkspaceLastUsedAt(gomock.Any(), database.BatchUpdateWorkspaceLastUsedAtParams{
490+
IDs: []uuid.UUID{workspace.ID},
467491
LastUsedAt: now,
468492
}).Return(nil)
469493

@@ -484,6 +508,10 @@ func TestUpdateStates(t *testing.T) {
484508
ReportInterval: durationpb.New(10 * time.Second),
485509
}, resp)
486510

511+
tickCh <- now
512+
count := <-flushCh
513+
require.Equal(t, 1, count, "expected one flush with one id")
514+
487515
batcher.Mu.Lock()
488516
defer batcher.Mu.Unlock()
489517
require.EqualValues(t, 1, batcher.Called)

‎coderd/insights_test.go

+2-4
Original file line numberDiff line numberDiff line change
@@ -700,14 +700,13 @@ func TestTemplateInsights_Golden(t *testing.T) {
700700
connectionCount = 0
701701
}
702702
for createdAt.Before(stat.endedAt) {
703-
err = batcher.Add(createdAt, workspace.agentID, workspace.template.id, workspace.user.(*testUser).sdk.ID, workspace.id, &agentproto.Stats{
703+
batcher.Add(createdAt, workspace.agentID, workspace.template.id, workspace.user.(*testUser).sdk.ID, workspace.id, &agentproto.Stats{
704704
ConnectionCount: connectionCount,
705705
SessionCountVscode: stat.sessionCountVSCode,
706706
SessionCountJetbrains: stat.sessionCountJetBrains,
707707
SessionCountReconnectingPty: stat.sessionCountReconnectingPTY,
708708
SessionCountSsh: stat.sessionCountSSH,
709709
}, false)
710-
require.NoError(t, err, "want no error inserting agent stats")
711710
createdAt = createdAt.Add(30 * time.Second)
712711
}
713712
}
@@ -1599,14 +1598,13 @@ func TestUserActivityInsights_Golden(t *testing.T) {
15991598
connectionCount = 0
16001599
}
16011600
for createdAt.Before(stat.endedAt) {
1602-
err = batcher.Add(createdAt, workspace.agentID, workspace.template.id, workspace.user.(*testUser).sdk.ID, workspace.id, &agentproto.Stats{
1601+
batcher.Add(createdAt, workspace.agentID, workspace.template.id, workspace.user.(*testUser).sdk.ID, workspace.id, &agentproto.Stats{
16031602
ConnectionCount: connectionCount,
16041603
SessionCountVscode: stat.sessionCountVSCode,
16051604
SessionCountJetbrains: stat.sessionCountJetBrains,
16061605
SessionCountReconnectingPty: stat.sessionCountReconnectingPTY,
16071606
SessionCountSsh: stat.sessionCountSSH,
16081607
}, false)
1609-
require.NoError(t, err, "want no error inserting agent stats")
16101608
createdAt = createdAt.Add(30 * time.Second)
16111609
}
16121610
}

‎coderd/workspaceagentsrpc_test.go

+12-1
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package coderd_test
33
import (
44
"context"
55
"testing"
6+
"time"
67

78
"github.com/stretchr/testify/assert"
89
"github.com/stretchr/testify/require"
@@ -11,6 +12,7 @@ import (
1112
"github.com/coder/coder/v2/coderd/coderdtest"
1213
"github.com/coder/coder/v2/coderd/database"
1314
"github.com/coder/coder/v2/coderd/database/dbfake"
15+
"github.com/coder/coder/v2/coderd/database/dbtime"
1416
"github.com/coder/coder/v2/codersdk/agentsdk"
1517
"github.com/coder/coder/v2/provisionersdk/proto"
1618
"github.com/coder/coder/v2/testutil"
@@ -20,7 +22,12 @@ import (
2022
func TestWorkspaceAgentReportStats(t *testing.T) {
2123
t.Parallel()
2224

23-
client, db := coderdtest.NewWithDatabase(t, nil)
25+
tickCh := make(chan time.Time)
26+
flushCh := make(chan int, 1)
27+
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
28+
WorkspaceUsageTrackerFlush: flushCh,
29+
WorkspaceUsageTrackerTick: tickCh,
30+
})
2431
user := coderdtest.CreateFirstUser(t, client)
2532
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
2633
OrganizationID: user.OrganizationID,
@@ -53,6 +60,10 @@ func TestWorkspaceAgentReportStats(t *testing.T) {
5360
})
5461
require.NoError(t, err)
5562

63+
tickCh <- dbtime.Now()
64+
count := <-flushCh
65+
require.Equal(t, 1, count, "expected one flush with one id")
66+
5667
newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID)
5768
require.NoError(t, err)
5869

‎coderd/workspacestats/batcher.go

+2-3
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ const (
2525
)
2626

2727
type Batcher interface {
28-
Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats, usage bool) error
28+
Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats, usage bool)
2929
}
3030

3131
// DBBatcher holds a buffer of agent stats and periodically flushes them to
@@ -139,7 +139,7 @@ func (b *DBBatcher) Add(
139139
workspaceID uuid.UUID,
140140
st *agentproto.Stats,
141141
usage bool,
142-
) error {
142+
) {
143143
b.mu.Lock()
144144
defer b.mu.Unlock()
145145

@@ -176,7 +176,6 @@ func (b *DBBatcher) Add(
176176
b.flushLever <- struct{}{}
177177
b.flushForced.Store(true)
178178
}
179-
return nil
180179
}
181180

182181
// Run runs the batcher.

‎coderd/workspacestats/batcher_internal_test.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ func TestBatchStats(t *testing.T) {
6363
// Given: a single data point is added for workspace
6464
t2 := t1.Add(time.Second)
6565
t.Logf("inserting 1 stat")
66-
require.NoError(t, b.Add(t2.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t), false))
66+
b.Add(t2.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t), false)
6767

6868
// When: it becomes time to report stats
6969
// Signal a tick and wait for a flush to complete.
@@ -87,9 +87,9 @@ func TestBatchStats(t *testing.T) {
8787
t.Logf("inserting %d stats", defaultBufferSize)
8888
for i := 0; i < defaultBufferSize; i++ {
8989
if i%2 == 0 {
90-
require.NoError(t, b.Add(t3.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t), false))
90+
b.Add(t3.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t), false)
9191
} else {
92-
require.NoError(t, b.Add(t3.Add(time.Millisecond), deps2.Agent.ID, deps2.User.ID, deps2.Template.ID, deps2.Workspace.ID, randStats(t), false))
92+
b.Add(t3.Add(time.Millisecond), deps2.Agent.ID, deps2.User.ID, deps2.Template.ID, deps2.Workspace.ID, randStats(t), false)
9393
}
9494
}
9595
}()

0 commit comments

Comments
 (0)