chore: move stat reporting into workspacestats package (#13386)

coderd/agentapi/activitybump.go

@@ -1,62 +0,0 @@
-package agentapi
-
-import (
-	"context"
-	"time"
-
-	"github.com/google/uuid"
-	"golang.org/x/xerrors"
-
-	"cdr.dev/slog"
-	"github.com/coder/coder/v2/coderd/database"
-)
-
-// ActivityBumpWorkspace automatically bumps the workspace's auto-off timer
-// if it is set to expire soon. The deadline will be bumped by 1 hour*.
-// If the bump crosses over an autostart time, the workspace will be
-// bumped by the workspace ttl instead.
-//
-// If nextAutostart is the zero value or in the past, the workspace
-// will be bumped by 1 hour.
-// It handles the edge case in the example:
-//  1. Autostart is set to 9am.
-//  2. User works all day, and leaves a terminal open to the workspace overnight.
-//  3. The open terminal continually bumps the workspace deadline.
-//  4. 9am the next day, the activity bump pushes to 10am.
-//  5. If the user goes inactive for 1 hour during the day, the workspace will
-//     now stop, because it has been extended by 1 hour durations. Despite the TTL
-//     being set to 8hrs from the autostart time.
-//
-// So the issue is that when the workspace is bumped across an autostart
-// deadline, we should treat the workspace as being "started" again and
-// extend the deadline by the autostart time + workspace ttl instead.
-//
-// The issue still remains with build_max_deadline. We need to respect the original
-// maximum deadline, so that will need to be handled separately.
-// A way to avoid this is to configure the max deadline to something that will not
-// span more than 1 day. This will force the workspace to restart and reset the deadline
-// each morning when it autostarts.
-func ActivityBumpWorkspace(ctx context.Context, log slog.Logger, db database.Store, workspaceID uuid.UUID, nextAutostart time.Time) {
-	// We set a short timeout so if the app is under load, these
-	// low priority operations fail first.
-	ctx, cancel := context.WithTimeout(ctx, time.Second*15)
-	defer cancel()
-	// nolint:gocritic // (#13146) Will be moved soon as part of refactor.
-	err := db.ActivityBumpWorkspace(ctx, database.ActivityBumpWorkspaceParams{
-		NextAutostart: nextAutostart.UTC(),
-		WorkspaceID:   workspaceID,
-	})
-	if err != nil {
-		if !xerrors.Is(err, context.Canceled) && !database.IsQueryCanceledError(err) {
-			// Bump will fail if the context is canceled, but this is ok.
-			log.Error(ctx, "activity bump failed", slog.Error(err),
-				slog.F("workspace_id", workspaceID),
-			)
-		}
-		return
-	}
-
-	log.Debug(ctx, "bumped deadline from activity",
-		slog.F("workspace_id", workspaceID),
-	)
-}

coderd/agentapi/activitybump_test.go

@@ -1,319 +0,0 @@
-package agentapi_test
-
-import (
-	"database/sql"
-	"testing"
-	"time"
-
-	"github.com/google/uuid"
-
-	"cdr.dev/slog/sloggers/slogtest"
-	"github.com/coder/coder/v2/coderd/agentapi"
-	"github.com/coder/coder/v2/coderd/database"
-	"github.com/coder/coder/v2/coderd/database/dbgen"
-	"github.com/coder/coder/v2/coderd/database/dbtestutil"
-	"github.com/coder/coder/v2/coderd/database/dbtime"
-	"github.com/coder/coder/v2/coderd/util/ptr"
-	"github.com/coder/coder/v2/testutil"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-)
-
-func Test_ActivityBumpWorkspace(t *testing.T) {
-	t.Parallel()
-
-	// We test the below in multiple timezones specifically
-	// chosen to trigger timezone-related bugs.
-	timezones := []string{
-		"Asia/Kolkata",        // No DST, positive fractional offset
-		"Canada/Newfoundland", // DST, negative fractional offset
-		"Europe/Paris",        // DST, positive offset
-		"US/Arizona",          // No DST, negative offset
-		"UTC",                 // Baseline
-	}
-
-	for _, tt := range []struct {
-		name                          string
-		transition                    database.WorkspaceTransition
-		jobCompletedAt                sql.NullTime
-		buildDeadlineOffset           *time.Duration
-		maxDeadlineOffset             *time.Duration
-		workspaceTTL                  time.Duration
-		templateTTL                   time.Duration
-		templateActivityBump          time.Duration
-		templateDisallowsUserAutostop bool
-		expectedBump                  time.Duration
-		// If the tests get queued, we need to be able to set the next autostart
-		// based on the actual time the unit test is running.
-		nextAutostart func(now time.Time) time.Time
-	}{
-		{
-			name:                "NotFinishedYet",
-			transition:          database.WorkspaceTransitionStart,
-			jobCompletedAt:      sql.NullTime{},
-			buildDeadlineOffset: ptr.Ref(8 * time.Hour),
-			workspaceTTL:        8 * time.Hour,
-			expectedBump:        0,
-		},
-		{
-			name:                "ManualShutdown",
-			transition:          database.WorkspaceTransitionStart,
-			jobCompletedAt:      sql.NullTime{Valid: true, Time: dbtime.Now()},
-			buildDeadlineOffset: nil,
-			expectedBump:        0,
-		},
-		{
-			name:                "NotTimeToBumpYet",
-			transition:          database.WorkspaceTransitionStart,
-			jobCompletedAt:      sql.NullTime{Valid: true, Time: dbtime.Now()},
-			buildDeadlineOffset: ptr.Ref(8 * time.Hour),
-			workspaceTTL:        8 * time.Hour,
-			expectedBump:        0,
-		},
-		{
-			// Expected bump is 0 because the original deadline is more than 1 hour
-			// out, so a bump would decrease the deadline.
-			name:                "BumpLessThanDeadline",
-			transition:          database.WorkspaceTransitionStart,
-			jobCompletedAt:      sql.NullTime{Valid: true, Time: dbtime.Now().Add(-30 * time.Minute)},
-			buildDeadlineOffset: ptr.Ref(8*time.Hour - 30*time.Minute),
-			workspaceTTL:        8 * time.Hour,
-			expectedBump:        0,
-		},
-		{
-			name:                "TimeToBump",
-			transition:          database.WorkspaceTransitionStart,
-			jobCompletedAt:      sql.NullTime{Valid: true, Time: dbtime.Now().Add(-30 * time.Minute)},
-			buildDeadlineOffset: ptr.Ref(-30 * time.Minute),
-			workspaceTTL:        8 * time.Hour,
-			expectedBump:        time.Hour,
-		},
-		{
-			name:                "TimeToBumpNextAutostart",
-			transition:          database.WorkspaceTransitionStart,
-			jobCompletedAt:      sql.NullTime{Valid: true, Time: dbtime.Now().Add(-30 * time.Minute)},
-			buildDeadlineOffset: ptr.Ref(-30 * time.Minute),
-			workspaceTTL:        8 * time.Hour,
-			expectedBump:        8*time.Hour + 30*time.Minute,
-			nextAutostart:       func(now time.Time) time.Time { return now.Add(time.Minute * 30) },
-		},
-		{
-			name:                "MaxDeadline",
-			transition:          database.WorkspaceTransitionStart,
-			jobCompletedAt:      sql.NullTime{Valid: true, Time: dbtime.Now().Add(-24 * time.Minute)},
-			buildDeadlineOffset: ptr.Ref(time.Minute), // last chance to bump!
-			maxDeadlineOffset:   ptr.Ref(time.Minute * 30),
-			workspaceTTL:        8 * time.Hour,
-			expectedBump:        time.Minute * 30,
-		},
-		{
-			// A workspace that is still running, has passed its deadline, but has not
-			// yet been auto-stopped should still bump the deadline.
-			name:                "PastDeadlineStillBumps",
-			transition:          database.WorkspaceTransitionStart,
-			jobCompletedAt:      sql.NullTime{Valid: true, Time: dbtime.Now().Add(-24 * time.Minute)},
-			buildDeadlineOffset: ptr.Ref(-time.Minute),
-			workspaceTTL:        8 * time.Hour,
-			expectedBump:        time.Hour,
-		},
-		{
-			// A stopped workspace should never bump.
-			name:                "StoppedWorkspace",
-			transition:          database.WorkspaceTransitionStop,
-			jobCompletedAt:      sql.NullTime{Valid: true, Time: dbtime.Now().Add(-time.Minute)},
-			buildDeadlineOffset: ptr.Ref(-time.Minute),
-			workspaceTTL:        8 * time.Hour,
-		},
-		{
-			// A workspace built from a template that disallows user autostop should bump
-			// by the template TTL instead.
-			name:                          "TemplateDisallowsUserAutostop",
-			transition:                    database.WorkspaceTransitionStart,
-			jobCompletedAt:                sql.NullTime{Valid: true, Time: dbtime.Now().Add(-3 * time.Hour)},
-			buildDeadlineOffset:           ptr.Ref(-30 * time.Minute),
-			workspaceTTL:                  2 * time.Hour,
-			templateTTL:                   10 * time.Hour,
-			templateDisallowsUserAutostop: true,
-			expectedBump:                  10*time.Hour + (time.Minute * 30),
-			nextAutostart:                 func(now time.Time) time.Time { return now.Add(time.Minute * 30) },
-		},
-		{
-			// Custom activity bump duration specified on the template.
-			name:                 "TemplateCustomActivityBump",
-			transition:           database.WorkspaceTransitionStart,
-			jobCompletedAt:       sql.NullTime{Valid: true, Time: dbtime.Now().Add(-30 * time.Minute)},
-			buildDeadlineOffset:  ptr.Ref(-30 * time.Minute),
-			workspaceTTL:         8 * time.Hour,
-			templateActivityBump: 5 * time.Hour, // instead of default 1h
-			expectedBump:         5 * time.Hour,
-		},
-		{
-			// Activity bump duration is 0.
-			name:                 "TemplateCustomActivityBumpZero",
-			transition:           database.WorkspaceTransitionStart,
-			jobCompletedAt:       sql.NullTime{Valid: true, Time: dbtime.Now().Add(-30 * time.Minute)},
-			buildDeadlineOffset:  ptr.Ref(-30 * time.Minute),
-			workspaceTTL:         8 * time.Hour,
-			templateActivityBump: -1, // negative values get changed to 0 in the test
-			expectedBump:         0,
-		},
-	} {
-		tt := tt
-		for _, tz := range timezones {
-			tz := tz
-			t.Run(tt.name+"/"+tz, func(t *testing.T) {
-				t.Parallel()
-				nextAutostart := tt.nextAutostart
-				if tt.nextAutostart == nil {
-					nextAutostart = func(now time.Time) time.Time { return time.Time{} }
-				}
-
-				var (
-					now   = dbtime.Now()
-					ctx   = testutil.Context(t, testutil.WaitShort)
-					log   = slogtest.Make(t, nil)
-					db, _ = dbtestutil.NewDB(t, dbtestutil.WithTimezone(tz))
-					org   = dbgen.Organization(t, db, database.Organization{})
-					user  = dbgen.User(t, db, database.User{
-						Status: database.UserStatusActive,
-					})
-					_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
-						UserID:         user.ID,
-						OrganizationID: org.ID,
-					})
-					templateVersion = dbgen.TemplateVersion(t, db, database.TemplateVersion{
-						OrganizationID: org.ID,
-						CreatedBy:      user.ID,
-					})
-					template = dbgen.Template(t, db, database.Template{
-						OrganizationID:  org.ID,
-						ActiveVersionID: templateVersion.ID,
-						CreatedBy:       user.ID,
-					})
-					ws = dbgen.Workspace(t, db, database.Workspace{
-						OwnerID:        user.ID,
-						OrganizationID: org.ID,
-						TemplateID:     template.ID,
-						Ttl:            sql.NullInt64{Valid: true, Int64: int64(tt.workspaceTTL)},
-					})
-					job = dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
-						OrganizationID: org.ID,
-						CompletedAt:    tt.jobCompletedAt,
-					})
-					_ = dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
-						JobID: job.ID,
-					})
-					buildID = uuid.New()
-				)
-
-				activityBump := 1 * time.Hour
-				if tt.templateActivityBump < 0 {
-					// less than 0 => 0
-					activityBump = 0
-				} else if tt.templateActivityBump != 0 {
-					activityBump = tt.templateActivityBump
-				}
-				require.NoError(t, db.UpdateTemplateScheduleByID(ctx, database.UpdateTemplateScheduleByIDParams{
-					ID:                template.ID,
-					UpdatedAt:         dbtime.Now(),
-					AllowUserAutostop: !tt.templateDisallowsUserAutostop,
-					DefaultTTL:        int64(tt.templateTTL),
-					ActivityBump:      int64(activityBump),
-				}), "unexpected error updating template schedule")
-
-				var buildNumber int32 = 1
-				// Insert a number of previous workspace builds.
-				for i := 0; i < 5; i++ {
-					insertPrevWorkspaceBuild(t, db, org.ID, templateVersion.ID, ws.ID, database.WorkspaceTransitionStart, buildNumber)
-					buildNumber++
-					insertPrevWorkspaceBuild(t, db, org.ID, templateVersion.ID, ws.ID, database.WorkspaceTransitionStop, buildNumber)
-					buildNumber++
-				}
-
-				// dbgen.WorkspaceBuild automatically sets deadline to now+1 hour if not set
-				var buildDeadline time.Time
-				if tt.buildDeadlineOffset != nil {
-					buildDeadline = now.Add(*tt.buildDeadlineOffset)
-				}
-				var maxDeadline time.Time
-				if tt.maxDeadlineOffset != nil {
-					maxDeadline = now.Add(*tt.maxDeadlineOffset)
-				}
-				err := db.InsertWorkspaceBuild(ctx, database.InsertWorkspaceBuildParams{
-					ID:                buildID,
-					CreatedAt:         dbtime.Now(),
-					UpdatedAt:         dbtime.Now(),
-					BuildNumber:       buildNumber,
-					InitiatorID:       user.ID,
-					Reason:            database.BuildReasonInitiator,
-					WorkspaceID:       ws.ID,
-					JobID:             job.ID,
-					TemplateVersionID: templateVersion.ID,
-					Transition:        tt.transition,
-					Deadline:          buildDeadline,
-					MaxDeadline:       maxDeadline,
-				})
-				require.NoError(t, err, "unexpected error inserting workspace build")
-				bld, err := db.GetWorkspaceBuildByID(ctx, buildID)
-				require.NoError(t, err, "unexpected error fetching inserted workspace build")
-
-				// Validate our initial state before bump
-				require.Equal(t, tt.transition, bld.Transition, "unexpected transition before bump")
-				require.Equal(t, tt.jobCompletedAt.Time.UTC(), job.CompletedAt.Time.UTC(), "unexpected job completed at before bump")
-				require.Equal(t, buildDeadline.UTC(), bld.Deadline.UTC(), "unexpected build deadline before bump")
-				require.Equal(t, maxDeadline.UTC(), bld.MaxDeadline.UTC(), "unexpected max deadline before bump")
-				require.Equal(t, tt.workspaceTTL, time.Duration(ws.Ttl.Int64), "unexpected workspace TTL before bump")
-
-				// Wait a bit before bumping as dbtime is rounded to the nearest millisecond.
-				// This should also hopefully be enough for Windows time resolution to register
-				// a tick (win32 max timer resolution is apparently between 0.5 and 15.6ms)
-				<-time.After(testutil.IntervalFast)
-
-				// Bump duration is measured from the time of the bump, so we measure from here.
-				start := dbtime.Now()
-				agentapi.ActivityBumpWorkspace(ctx, log, db, bld.WorkspaceID, nextAutostart(start))
-				end := dbtime.Now()
-
-				// Validate our state after bump
-				updatedBuild, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, bld.WorkspaceID)
-				require.NoError(t, err, "unexpected error getting latest workspace build")
-				require.Equal(t, bld.MaxDeadline.UTC(), updatedBuild.MaxDeadline.UTC(), "max_deadline should not have changed")
-				if tt.expectedBump == 0 {
-					assert.Equal(t, bld.UpdatedAt.UTC(), updatedBuild.UpdatedAt.UTC(), "should not have bumped updated_at")
-					assert.Equal(t, bld.Deadline.UTC(), updatedBuild.Deadline.UTC(), "should not have bumped deadline")
-					return
-				}
-				assert.NotEqual(t, bld.UpdatedAt.UTC(), updatedBuild.UpdatedAt.UTC(), "should have bumped updated_at")
-				if tt.maxDeadlineOffset != nil {
-					assert.Equal(t, bld.MaxDeadline.UTC(), updatedBuild.MaxDeadline.UTC(), "new deadline must equal original max deadline")
-					return
-				}
-
-				// Assert that the bump occurred between start and end. 1min buffer on either side.
-				expectedDeadlineStart := start.Add(tt.expectedBump).Add(time.Minute * -1)
-				expectedDeadlineEnd := end.Add(tt.expectedBump).Add(time.Minute)
-				require.GreaterOrEqual(t, updatedBuild.Deadline, expectedDeadlineStart, "new deadline should be greater than or equal to start")
-				require.LessOrEqual(t, updatedBuild.Deadline, expectedDeadlineEnd, "new deadline should be less than or equal to end")
-			})
-		}
-	}
-}
-
-func insertPrevWorkspaceBuild(t *testing.T, db database.Store, orgID, tvID, workspaceID uuid.UUID, transition database.WorkspaceTransition, buildNumber int32) {
-	t.Helper()
-
-	job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
-		OrganizationID: orgID,
-	})
-	_ = dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
-		JobID: job.ID,
-	})
-	_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
-		BuildNumber:       buildNumber,
-		WorkspaceID:       workspaceID,
-		JobID:             job.ID,
-		TemplateVersionID: tvID,
-		Transition:        transition,
-	})
-}

coderd/agentapi/api.go

@@ -24,6 +24,7 @@ import (
 	"github.com/coder/coder/v2/coderd/prometheusmetrics"
 	"github.com/coder/coder/v2/coderd/schedule"
 	"github.com/coder/coder/v2/coderd/tracing"
+	"github.com/coder/coder/v2/coderd/workspacestats"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/tailnet"
 	tailnetproto "github.com/coder/coder/v2/tailnet/proto"
@@ -59,7 +60,7 @@ type Options struct {
 	DerpMapFn                         func() *tailcfg.DERPMap
 	TailnetCoordinator                *atomic.Pointer[tailnet.Coordinator]
 	TemplateScheduleStore             *atomic.Pointer[schedule.TemplateScheduleStore]
-	StatsBatcher                      StatsBatcher
+	StatsReporter                     *workspacestats.Reporter
 	AppearanceFetcher                 *atomic.Pointer[appearance.Fetcher]
 	PublishWorkspaceUpdateFn          func(ctx context.Context, workspaceID uuid.UUID)
 	PublishWorkspaceAgentLogsUpdateFn func(ctx context.Context, workspaceAgentID uuid.UUID, msg agentsdk.LogsNotifyMessage)
@@ -114,12 +115,9 @@ func New(opts Options) *API {
 	api.StatsAPI = &StatsAPI{
 		AgentFn:                   api.agent,
 		Database:                  opts.Database,
-		Pubsub:                    opts.Pubsub,
 		Log:                       opts.Log,
-		StatsBatcher:              opts.StatsBatcher,
-		TemplateScheduleStore:     opts.TemplateScheduleStore,
+		StatsReporter:             opts.StatsReporter,
 		AgentStatsRefreshInterval: opts.AgentStatsRefreshInterval,
-		UpdateAgentMetricsFn:      opts.UpdateAgentMetricsFn,
 	}
 
 	api.LifecycleAPI = &LifecycleAPI{

coderd/agentapi/stats.go

@@ -2,10 +2,8 @@ package agentapi
 
 import (
 	"context"
-	"sync/atomic"
 	"time"
 
-	"golang.org/x/sync/errgroup"
 	"golang.org/x/xerrors"
 	"google.golang.org/protobuf/types/known/durationpb"
 
@@ -15,10 +13,7 @@ import (
 	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
-	"github.com/coder/coder/v2/coderd/database/pubsub"
-	"github.com/coder/coder/v2/coderd/prometheusmetrics"
-	"github.com/coder/coder/v2/coderd/schedule"
-	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/coderd/workspacestats"
 )
 
 type StatsBatcher interface {
@@ -28,12 +23,9 @@ type StatsBatcher interface {
 type StatsAPI struct {
 	AgentFn                   func(context.Context) (database.WorkspaceAgent, error)
 	Database                  database.Store
-	Pubsub                    pubsub.Pubsub
 	Log                       slog.Logger
-	StatsBatcher              StatsBatcher
-	TemplateScheduleStore     *atomic.Pointer[schedule.TemplateScheduleStore]
+	StatsReporter             *workspacestats.Reporter
 	AgentStatsRefreshInterval time.Duration
-	UpdateAgentMetricsFn      func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
 
 	TimeNowFn func() time.Time // defaults to dbtime.Now()
 }
@@ -69,80 +61,17 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
 		slog.F("payload", req),
 	)
 
-	now := a.now()
-	if req.Stats.ConnectionCount > 0 {
-		var nextAutostart time.Time
-		if workspace.AutostartSchedule.String != "" {
-			templateSchedule, err := (*(a.TemplateScheduleStore.Load())).Get(ctx, a.Database, workspace.TemplateID)
-			// If the template schedule fails to load, just default to bumping
-			// without the next transition and log it.
-			if err != nil {
-				a.Log.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, a.Log.Named("activity_bump"), a.Database, workspace.ID, nextAutostart)
-	}
-
-	var errGroup errgroup.Group
-	errGroup.Go(func() error {
-		err := a.StatsBatcher.Add(now, workspaceAgent.ID, workspace.TemplateID, workspace.OwnerID, workspace.ID, req.Stats)
-		if err != nil {
-			a.Log.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 {
-		// nolint:gocritic // (#13146) Will be moved soon as part of refactor.
-		err := a.Database.UpdateWorkspaceLastUsedAt(ctx, database.UpdateWorkspaceLastUsedAtParams{
-			ID:         workspace.ID,
-			LastUsedAt: now,
-		})
-		if err != nil {
-			return xerrors.Errorf("update workspace LastUsedAt: %w", err)
-		}
-		return nil
-	})
-	if a.UpdateAgentMetricsFn != nil {
-		errGroup.Go(func() error {
-			user, err := a.Database.GetUserByID(ctx, workspace.OwnerID)
-			if err != nil {
-				return xerrors.Errorf("get user: %w", err)
-			}
-
-			a.UpdateAgentMetricsFn(ctx, prometheusmetrics.AgentMetricLabels{
-				Username:      user.Username,
-				WorkspaceName: workspace.Name,
-				AgentName:     workspaceAgent.Name,
-				TemplateName:  getWorkspaceAgentByIDRow.TemplateName,
-			}, req.Stats.Metrics)
-			return nil
-		})
-	}
-	err = errGroup.Wait()
+	err = a.StatsReporter.ReportAgentStats(
+		ctx,
+		a.now(),
+		workspace,
+		workspaceAgent,
+		getWorkspaceAgentByIDRow.TemplateName,
+		req.Stats,
+	)
 	if err != nil {
-		return nil, xerrors.Errorf("update stats in database: %w", err)
+		return nil, xerrors.Errorf("report agent stats: %w", err)
 	}
 
-	// Tell the frontend about the new agent report, now that everything is updated
-	a.publishWorkspaceAgentStats(ctx, workspace.ID)
-
 	return res, nil
 }
-
-func (a *StatsAPI) publishWorkspaceAgentStats(ctx context.Context, workspaceID uuid.UUID) {
-	err := a.Pubsub.Publish(codersdk.WorkspaceNotifyChannel(workspaceID), []byte{})
-	if err != nil {
-		a.Log.Warn(ctx, "failed to publish workspace agent stats",
-			slog.F("workspace_id", workspaceID), slog.Error(err))
-	}
-}

coderd/agentapi/stats_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"github.com/coder/coder/v2/coderd/prometheusmetrics"
 	"github.com/coder/coder/v2/coderd/schedule"
+	"github.com/coder/coder/v2/coderd/workspacestats"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/testutil"
 )
@@ -129,21 +130,24 @@ func TestUpdateStates(t *testing.T) {
 			AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			Database:                  dbM,
-			Pubsub:                    ps,
-			StatsBatcher:              batcher,
-			TemplateScheduleStore:     templateScheduleStorePtr(templateScheduleStore),
+			Database: dbM,
+			StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
+				Database:              dbM,
+				Pubsub:                ps,
+				StatsBatcher:          batcher,
+				TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
+				UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
+					updateAgentMetricsFnCalled = true
+					assert.Equal(t, prometheusmetrics.AgentMetricLabels{
+						Username:      user.Username,
+						WorkspaceName: workspace.Name,
+						AgentName:     agent.Name,
+						TemplateName:  template.Name,
+					}, labels)
+					assert.Equal(t, req.Stats.Metrics, metrics)
+				},
+			}),
 			AgentStatsRefreshInterval: 10 * time.Second,
-			UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
-				updateAgentMetricsFnCalled = true
-				assert.Equal(t, prometheusmetrics.AgentMetricLabels{
-					Username:      user.Username,
-					WorkspaceName: workspace.Name,
-					AgentName:     agent.Name,
-					TemplateName:  template.Name,
-				}, labels)
-				assert.Equal(t, req.Stats.Metrics, metrics)
-			},
 			TimeNowFn: func() time.Time {
 				return now
 			},
@@ -232,13 +236,16 @@ func TestUpdateStates(t *testing.T) {
 			AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			Database:                  dbM,
-			Pubsub:                    ps,
-			StatsBatcher:              batcher,
-			TemplateScheduleStore:     templateScheduleStorePtr(templateScheduleStore),
+			Database: dbM,
+			StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
+				Database:              dbM,
+				Pubsub:                ps,
+				StatsBatcher:          batcher,
+				TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
+				// Ignored when nil.
+				UpdateAgentMetricsFn: nil,
+			}),
 			AgentStatsRefreshInterval: 10 * time.Second,
-			// Ignored when nil.
-			UpdateAgentMetricsFn: nil,
 			TimeNowFn: func() time.Time {
 				return now
 			},
@@ -274,12 +281,15 @@ func TestUpdateStates(t *testing.T) {
 			AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			Database:                  dbM,
-			Pubsub:                    ps,
-			StatsBatcher:              nil, // should not be called
-			TemplateScheduleStore:     nil, // should not be called
+			Database: dbM,
+			StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
+				Database:              dbM,
+				Pubsub:                ps,
+				StatsBatcher:          nil, // should not be called
+				TemplateScheduleStore: nil, // should not be called
+				UpdateAgentMetricsFn:  nil, // should not be called
+			}),
 			AgentStatsRefreshInterval: 10 * time.Second,
-			UpdateAgentMetricsFn:      nil, // should not be called
 			TimeNowFn: func() time.Time {
 				panic("should not be called")
 			},
@@ -343,21 +353,24 @@ func TestUpdateStates(t *testing.T) {
 			AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			Database:                  dbM,
-			Pubsub:                    ps,
-			StatsBatcher:              batcher,
-			TemplateScheduleStore:     templateScheduleStorePtr(templateScheduleStore),
+			Database: dbM,
+			StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
+				Database:              dbM,
+				Pubsub:                ps,
+				StatsBatcher:          batcher,
+				TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore),
+				UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
+					updateAgentMetricsFnCalled = true
+					assert.Equal(t, prometheusmetrics.AgentMetricLabels{
+						Username:      user.Username,
+						WorkspaceName: workspace.Name,
+						AgentName:     agent.Name,
+						TemplateName:  template.Name,
+					}, labels)
+					assert.Equal(t, req.Stats.Metrics, metrics)
+				},
+			}),
 			AgentStatsRefreshInterval: 15 * time.Second,
-			UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) {
-				updateAgentMetricsFnCalled = true
-				assert.Equal(t, prometheusmetrics.AgentMetricLabels{
-					Username:      user.Username,
-					WorkspaceName: workspace.Name,
-					AgentName:     agent.Name,
-					TemplateName:  template.Name,
-				}, labels)
-				assert.Equal(t, req.Stats.Metrics, metrics)
-			},
 			TimeNowFn: func() time.Time {
 				return now
 			},