fix: stop incrementing activity on empty agent stats (#15204)

This commit is contained in:
Garrett Delfosse
2024-10-25 16:49:44 +00:00
committed by GitHub
parent cd890aa3a0
commit 0dd942e197
8 changed files with 110 additions and 89 deletions
+44 -16
View File
@@ -70,6 +70,11 @@ func TestUpdateStates(t *testing.T) {
}
batcher = &workspacestatstest.StatsBatcher{}
updateAgentMetricsFnCalled = false
tickCh = make(chan time.Time)
flushCh = make(chan int, 1)
wut = workspacestats.NewTracker(dbM,
workspacestats.TrackerWithTickFlush(tickCh, flushCh),
)
req = &agentproto.UpdateStatsRequest{
Stats: &agentproto.Stats{
@@ -109,6 +114,7 @@ func TestUpdateStates(t *testing.T) {
Database: dbM,
Pubsub: ps,
StatsBatcher: batcher,
UsageTracker: wut,
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
updateAgentMetricsFnCalled = true
@@ -126,10 +132,14 @@ func TestUpdateStates(t *testing.T) {
return now
},
}
defer wut.Close()
// Workspace gets fetched.
dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(workspace, nil)
// User gets fetched to hit the UpdateAgentMetricsFn.
dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
// We expect an activity bump because ConnectionCount > 0.
dbM.EXPECT().ActivityBumpWorkspace(gomock.Any(), database.ActivityBumpWorkspaceParams{
WorkspaceID: workspace.ID,
@@ -137,14 +147,11 @@ func TestUpdateStates(t *testing.T) {
}).Return(nil)
// Workspace last used at gets bumped.
dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
ID: workspace.ID,
dbM.EXPECT().BatchUpdateWorkspaceLastUsedAt(gomock.Any(), database.BatchUpdateWorkspaceLastUsedAtParams{
IDs: []uuid.UUID{workspace.ID},
LastUsedAt: now,
}).Return(nil)
// User gets fetched to hit the UpdateAgentMetricsFn.
dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
// Ensure that pubsub notifications are sent.
notifyDescription := make(chan []byte)
ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, description []byte) {
@@ -159,6 +166,10 @@ func TestUpdateStates(t *testing.T) {
ReportInterval: durationpb.New(10 * time.Second),
}, resp)
tickCh <- now
count := <-flushCh
require.Equal(t, 1, count, "expected one flush with one id")
batcher.Mu.Lock()
defer batcher.Mu.Unlock()
require.Equal(t, int64(1), batcher.Called)
@@ -211,6 +222,7 @@ func TestUpdateStates(t *testing.T) {
StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
Database: dbM,
Pubsub: ps,
UsageTracker: workspacestats.NewTracker(dbM),
StatsBatcher: batcher,
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
// Ignored when nil.
@@ -225,12 +237,6 @@ func TestUpdateStates(t *testing.T) {
// Workspace gets fetched.
dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(workspace, nil)
// Workspace last used at gets bumped.
dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
ID: workspace.ID,
LastUsedAt: now,
}).Return(nil)
_, err := api.UpdateStats(context.Background(), req)
require.NoError(t, err)
})
@@ -306,6 +312,11 @@ func TestUpdateStates(t *testing.T) {
}
batcher = &workspacestatstest.StatsBatcher{}
updateAgentMetricsFnCalled = false
tickCh = make(chan time.Time)
flushCh = make(chan int, 1)
wut = workspacestats.NewTracker(dbM,
workspacestats.TrackerWithTickFlush(tickCh, flushCh),
)
req = &agentproto.UpdateStatsRequest{
Stats: &agentproto.Stats{
@@ -325,6 +336,7 @@ func TestUpdateStates(t *testing.T) {
StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
Database: dbM,
Pubsub: ps,
UsageTracker: wut,
StatsBatcher: batcher,
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
@@ -343,6 +355,7 @@ func TestUpdateStates(t *testing.T) {
return now
},
}
defer wut.Close()
// Workspace gets fetched.
dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(workspace, nil)
@@ -355,9 +368,9 @@ func TestUpdateStates(t *testing.T) {
}).Return(nil)
// Workspace last used at gets bumped.
dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
ID: workspace.ID,
LastUsedAt: now,
dbM.EXPECT().BatchUpdateWorkspaceLastUsedAt(gomock.Any(), database.BatchUpdateWorkspaceLastUsedAtParams{
IDs: []uuid.UUID{workspace.ID},
LastUsedAt: now.UTC(),
}).Return(nil)
// User gets fetched to hit the UpdateAgentMetricsFn.
@@ -369,6 +382,10 @@ func TestUpdateStates(t *testing.T) {
ReportInterval: durationpb.New(15 * time.Second),
}, resp)
tickCh <- now
count := <-flushCh
require.Equal(t, 1, count, "expected one flush with one id")
require.True(t, updateAgentMetricsFnCalled)
})
@@ -392,6 +409,11 @@ func TestUpdateStates(t *testing.T) {
}
batcher = &workspacestatstest.StatsBatcher{}
updateAgentMetricsFnCalled = false
tickCh = make(chan time.Time)
flushCh = make(chan int, 1)
wut = workspacestats.NewTracker(dbM,
workspacestats.TrackerWithTickFlush(tickCh, flushCh),
)
req = &agentproto.UpdateStatsRequest{
Stats: &agentproto.Stats{
@@ -422,6 +444,7 @@ func TestUpdateStates(t *testing.T) {
},
}
)
defer wut.Close()
api := agentapi.StatsAPI{
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
return agent, nil
@@ -431,6 +454,7 @@ func TestUpdateStates(t *testing.T) {
Database: dbM,
Pubsub: ps,
StatsBatcher: batcher,
UsageTracker: wut,
TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
updateAgentMetricsFnCalled = true
@@ -462,8 +486,8 @@ func TestUpdateStates(t *testing.T) {
}).Return(nil)
// Workspace last used at gets bumped.
dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{
ID: workspace.ID,
dbM.EXPECT().BatchUpdateWorkspaceLastUsedAt(gomock.Any(), database.BatchUpdateWorkspaceLastUsedAtParams{
IDs: []uuid.UUID{workspace.ID},
LastUsedAt: now,
}).Return(nil)
@@ -484,6 +508,10 @@ func TestUpdateStates(t *testing.T) {
ReportInterval: durationpb.New(10 * time.Second),
}, resp)
tickCh <- now
count := <-flushCh
require.Equal(t, 1, count, "expected one flush with one id")
batcher.Mu.Lock()
defer batcher.Mu.Unlock()
require.EqualValues(t, 1, batcher.Called)