From 0dd942e19728af9551406f04fa76ad6a807246e1 Mon Sep 17 00:00:00 2001 From: Garrett Delfosse Date: Fri, 25 Oct 2024 12:49:44 -0400 Subject: [PATCH] fix: stop incrementing activity on empty agent stats (#15204) --- coderd/agentapi/stats_test.go | 60 +++++++--- coderd/insights_test.go | 6 +- coderd/workspaceagentsrpc_test.go | 13 ++- coderd/workspacestats/batcher.go | 5 +- .../workspacestats/batcher_internal_test.go | 6 +- coderd/workspacestats/reporter.go | 105 ++++++++---------- coderd/workspacestats/tracker.go | 1 - .../workspacestatstest/batcher.go | 3 +- 8 files changed, 110 insertions(+), 89 deletions(-) diff --git a/coderd/agentapi/stats_test.go b/coderd/agentapi/stats_test.go index d2c8e4f163..83edb8cccc 100644 --- a/coderd/agentapi/stats_test.go +++ b/coderd/agentapi/stats_test.go @@ -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) diff --git a/coderd/insights_test.go b/coderd/insights_test.go index 06fe8d46ca..bf8aa4bc44 100644 --- a/coderd/insights_test.go +++ b/coderd/insights_test.go @@ -700,14 +700,13 @@ func TestTemplateInsights_Golden(t *testing.T) { connectionCount = 0 } for createdAt.Before(stat.endedAt) { - err = batcher.Add(createdAt, workspace.agentID, workspace.template.id, workspace.user.(*testUser).sdk.ID, workspace.id, &agentproto.Stats{ + batcher.Add(createdAt, workspace.agentID, workspace.template.id, workspace.user.(*testUser).sdk.ID, workspace.id, &agentproto.Stats{ ConnectionCount: connectionCount, SessionCountVscode: stat.sessionCountVSCode, SessionCountJetbrains: stat.sessionCountJetBrains, SessionCountReconnectingPty: stat.sessionCountReconnectingPTY, SessionCountSsh: stat.sessionCountSSH, }, false) - require.NoError(t, err, "want no error inserting agent stats") createdAt = createdAt.Add(30 * time.Second) } } @@ -1599,14 +1598,13 @@ func TestUserActivityInsights_Golden(t *testing.T) { connectionCount = 0 } for createdAt.Before(stat.endedAt) { - err = batcher.Add(createdAt, workspace.agentID, workspace.template.id, workspace.user.(*testUser).sdk.ID, workspace.id, &agentproto.Stats{ + batcher.Add(createdAt, workspace.agentID, workspace.template.id, workspace.user.(*testUser).sdk.ID, workspace.id, &agentproto.Stats{ ConnectionCount: connectionCount, SessionCountVscode: stat.sessionCountVSCode, SessionCountJetbrains: stat.sessionCountJetBrains, SessionCountReconnectingPty: stat.sessionCountReconnectingPTY, SessionCountSsh: stat.sessionCountSSH, }, false) - require.NoError(t, err, "want no error inserting agent stats") createdAt = createdAt.Add(30 * time.Second) } } diff --git a/coderd/workspaceagentsrpc_test.go b/coderd/workspaceagentsrpc_test.go index 817aa11c4c..3f1f1a2b8a 100644 --- a/coderd/workspaceagentsrpc_test.go +++ b/coderd/workspaceagentsrpc_test.go @@ -3,6 +3,7 @@ package coderd_test import ( "context" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -11,6 +12,7 @@ import ( "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbfake" + "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/provisionersdk/proto" "github.com/coder/coder/v2/testutil" @@ -20,7 +22,12 @@ import ( func TestWorkspaceAgentReportStats(t *testing.T) { t.Parallel() - client, db := coderdtest.NewWithDatabase(t, nil) + tickCh := make(chan time.Time) + flushCh := make(chan int, 1) + client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{ + WorkspaceUsageTrackerFlush: flushCh, + WorkspaceUsageTrackerTick: tickCh, + }) user := coderdtest.CreateFirstUser(t, client) r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OrganizationID: user.OrganizationID, @@ -53,6 +60,10 @@ func TestWorkspaceAgentReportStats(t *testing.T) { }) require.NoError(t, err) + tickCh <- dbtime.Now() + count := <-flushCh + require.Equal(t, 1, count, "expected one flush with one id") + newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID) require.NoError(t, err) diff --git a/coderd/workspacestats/batcher.go b/coderd/workspacestats/batcher.go index 1f14c5cec5..46efc69170 100644 --- a/coderd/workspacestats/batcher.go +++ b/coderd/workspacestats/batcher.go @@ -25,7 +25,7 @@ const ( ) type Batcher interface { - Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats, usage bool) error + Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats, usage bool) } // DBBatcher holds a buffer of agent stats and periodically flushes them to @@ -139,7 +139,7 @@ func (b *DBBatcher) Add( workspaceID uuid.UUID, st *agentproto.Stats, usage bool, -) error { +) { b.mu.Lock() defer b.mu.Unlock() @@ -176,7 +176,6 @@ func (b *DBBatcher) Add( b.flushLever <- struct{}{} b.flushForced.Store(true) } - return nil } // Run runs the batcher. diff --git a/coderd/workspacestats/batcher_internal_test.go b/coderd/workspacestats/batcher_internal_test.go index 3e106f07e4..874acd7667 100644 --- a/coderd/workspacestats/batcher_internal_test.go +++ b/coderd/workspacestats/batcher_internal_test.go @@ -63,7 +63,7 @@ func TestBatchStats(t *testing.T) { // Given: a single data point is added for workspace t2 := t1.Add(time.Second) t.Logf("inserting 1 stat") - require.NoError(t, b.Add(t2.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t), false)) + b.Add(t2.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t), false) // When: it becomes time to report stats // Signal a tick and wait for a flush to complete. @@ -87,9 +87,9 @@ func TestBatchStats(t *testing.T) { t.Logf("inserting %d stats", defaultBufferSize) for i := 0; i < defaultBufferSize; i++ { if i%2 == 0 { - require.NoError(t, b.Add(t3.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t), false)) + b.Add(t3.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t), false) } else { - require.NoError(t, b.Add(t3.Add(time.Millisecond), deps2.Agent.ID, deps2.User.ID, deps2.Template.ID, deps2.Workspace.ID, randStats(t), false)) + b.Add(t3.Add(time.Millisecond), deps2.Agent.ID, deps2.User.ID, deps2.Template.ID, deps2.Workspace.ID, randStats(t), false) } } }() diff --git a/coderd/workspacestats/reporter.go b/coderd/workspacestats/reporter.go index fecfd1b1ed..6bb1b2dea4 100644 --- a/coderd/workspacestats/reporter.go +++ b/coderd/workspacestats/reporter.go @@ -6,7 +6,6 @@ import ( "time" "github.com/google/uuid" - "golang.org/x/sync/errgroup" "golang.org/x/xerrors" "cdr.dev/slog" @@ -119,69 +118,57 @@ func (r *Reporter) ReportAppStats(ctx context.Context, stats []workspaceapps.Sta } func (r *Reporter) ReportAgentStats(ctx context.Context, now time.Time, workspace database.Workspace, workspaceAgent database.WorkspaceAgent, templateName string, stats *agentproto.Stats, usage bool) error { - if stats.ConnectionCount > 0 { - var nextAutostart time.Time - if workspace.AutostartSchedule.String != "" { - templateSchedule, err := (*(r.opts.TemplateScheduleStore.Load())).Get(ctx, r.opts.Database, workspace.TemplateID) - // If the template schedule fails to load, just default to bumping - // without the next transition and log it. - if err != nil { - r.opts.Logger.Error(ctx, "failed to load template schedule bumping activity, defaulting to bumping by 60min", - slog.F("workspace_id", workspace.ID), - slog.F("template_id", workspace.TemplateID), - slog.Error(err), - ) - } else { - next, allowed := schedule.NextAutostart(now, workspace.AutostartSchedule.String, templateSchedule) - if allowed { - nextAutostart = next - } - } - } - ActivityBumpWorkspace(ctx, r.opts.Logger.Named("activity_bump"), r.opts.Database, workspace.ID, nextAutostart) - } + // update agent stats + r.opts.StatsBatcher.Add(now, workspaceAgent.ID, workspace.TemplateID, workspace.OwnerID, workspace.ID, stats, usage) - var errGroup errgroup.Group - errGroup.Go(func() error { - err := r.opts.StatsBatcher.Add(now, workspaceAgent.ID, workspace.TemplateID, workspace.OwnerID, workspace.ID, stats, usage) - if err != nil { - r.opts.Logger.Error(ctx, "add agent stats to batcher", slog.Error(err)) - return xerrors.Errorf("insert workspace agent stats batch: %w", err) - } - return nil - }) - errGroup.Go(func() error { - err := r.opts.Database.UpdateWorkspaceLastUsedAt(ctx, database.UpdateWorkspaceLastUsedAtParams{ - ID: workspace.ID, - LastUsedAt: now, - }) - if err != nil { - return xerrors.Errorf("update workspace LastUsedAt: %w", err) - } - return nil - }) + // update prometheus metrics if r.opts.UpdateAgentMetricsFn != nil { - errGroup.Go(func() error { - user, err := r.opts.Database.GetUserByID(ctx, workspace.OwnerID) - if err != nil { - return xerrors.Errorf("get user: %w", err) + user, err := r.opts.Database.GetUserByID(ctx, workspace.OwnerID) + if err != nil { + return xerrors.Errorf("get user: %w", err) + } + + r.opts.UpdateAgentMetricsFn(ctx, prometheusmetrics.AgentMetricLabels{ + Username: user.Username, + WorkspaceName: workspace.Name, + AgentName: workspaceAgent.Name, + TemplateName: templateName, + }, stats.Metrics) + } + + // if no active connections we do not bump activity + if stats.ConnectionCount == 0 { + return nil + } + + // check next autostart + var nextAutostart time.Time + if workspace.AutostartSchedule.String != "" { + templateSchedule, err := (*(r.opts.TemplateScheduleStore.Load())).Get(ctx, r.opts.Database, workspace.TemplateID) + // If the template schedule fails to load, just default to bumping + // without the next transition and log it. + if err != nil { + r.opts.Logger.Error(ctx, "failed to load template schedule bumping activity, defaulting to bumping by 60min", + slog.F("workspace_id", workspace.ID), + slog.F("template_id", workspace.TemplateID), + slog.Error(err), + ) + } else { + next, allowed := schedule.NextAutostart(now, workspace.AutostartSchedule.String, templateSchedule) + if allowed { + nextAutostart = next } - - r.opts.UpdateAgentMetricsFn(ctx, prometheusmetrics.AgentMetricLabels{ - Username: user.Username, - WorkspaceName: workspace.Name, - AgentName: workspaceAgent.Name, - TemplateName: templateName, - }, stats.Metrics) - return nil - }) - } - err := errGroup.Wait() - if err != nil { - return xerrors.Errorf("update stats in database: %w", err) + } } - err = r.opts.Pubsub.Publish(codersdk.WorkspaceNotifyChannel(workspace.ID), []byte{}) + // bump workspace activity + ActivityBumpWorkspace(ctx, r.opts.Logger.Named("activity_bump"), r.opts.Database, workspace.ID, nextAutostart) + + // bump workspace last_used_at + r.opts.UsageTracker.Add(workspace.ID) + + // notify workspace update + err := r.opts.Pubsub.Publish(codersdk.WorkspaceNotifyChannel(workspace.ID), []byte{}) if err != nil { r.opts.Logger.Warn(ctx, "failed to publish workspace agent stats", slog.F("workspace_id", workspace.ID), slog.Error(err)) diff --git a/coderd/workspacestats/tracker.go b/coderd/workspacestats/tracker.go index 33532247b3..f55edde3b5 100644 --- a/coderd/workspacestats/tracker.go +++ b/coderd/workspacestats/tracker.go @@ -130,7 +130,6 @@ func (tr *UsageTracker) flush(now time.Time) { authCtx := dbauthz.AsSystemRestricted(ctx) tr.flushLock.Lock() defer tr.flushLock.Unlock() - // nolint:gocritic // (#13146) Will be moved soon as part of refactor. if err := tr.s.BatchUpdateWorkspaceLastUsedAt(authCtx, database.BatchUpdateWorkspaceLastUsedAtParams{ LastUsedAt: now, IDs: ids, diff --git a/coderd/workspacestats/workspacestatstest/batcher.go b/coderd/workspacestats/workspacestatstest/batcher.go index 2f5dd7d13a..592e244518 100644 --- a/coderd/workspacestats/workspacestatstest/batcher.go +++ b/coderd/workspacestats/workspacestatstest/batcher.go @@ -25,7 +25,7 @@ type StatsBatcher struct { var _ workspacestats.Batcher = &StatsBatcher{} -func (b *StatsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats, usage bool) error { +func (b *StatsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats, usage bool) { b.Mu.Lock() defer b.Mu.Unlock() b.Called++ @@ -36,5 +36,4 @@ func (b *StatsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUI b.LastWorkspaceID = workspaceID b.LastStats = st b.LastUsage = usage - return nil }