Revert "remove workspace tracker"

This reverts commit dc818d7.
coder · f0ssel · May 21, 2024 · May 23, 2024 · May 23, 2024 · May 23, 2024
commit a60de2b8274dbc7279f1ae0dc9c04c35636ad978
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -68,6 +68,7 @@ import (
 	"github.com/coder/coder/v2/coderd/updatecheck"
 	"github.com/coder/coder/v2/coderd/util/slice"
 	"github.com/coder/coder/v2/coderd/workspaceapps"
+	"github.com/coder/coder/v2/coderd/workspaceusage"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/drpc"
 	"github.com/coder/coder/v2/codersdk/healthsdk"
@@ -203,6 +204,8 @@ type Options struct {
 	// DatabaseRolluper rolls up template usage stats from raw agent and app
 	// stats. This is used to provide insights in the WebUI.
 	DatabaseRolluper *dbrollup.Rolluper
+	// WorkspaceUsageTracker tracks workspace usage by the CLI.
+	WorkspaceUsageTracker *workspaceusage.Tracker
 }
 
 // @title Coder API
@@ -379,6 +382,12 @@ func New(options *Options) *API {
 		options.DatabaseRolluper = dbrollup.New(options.Logger.Named("dbrollup"), options.Database)
 	}
 
+	if options.WorkspaceUsageTracker == nil {
+		options.WorkspaceUsageTracker = workspaceusage.New(options.Database,
+			workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
+		)
+	}
+
 	ctx, cancel := context.WithCancel(context.Background())
 	r := chi.NewRouter()
 
@@ -423,7 +432,8 @@ func New(options *Options) *API {
 			options.Database,
 			options.Pubsub,
 		),
-		dbRolluper: options.DatabaseRolluper,
+		dbRolluper:            options.DatabaseRolluper,
+		workspaceUsageTracker: options.WorkspaceUsageTracker,
 	}
 
 	api.AppearanceFetcher.Store(&appearance.DefaultFetcher)
@@ -1265,13 +1275,13 @@ type API struct {
 	healthCheckGroup *singleflight.Group[string, *healthsdk.HealthcheckReport]
 	healthCheckCache atomic.Pointer[healthsdk.HealthcheckReport]
 
-	statsBatcher   *batchstats.Batcher
-	statsCollector workspaceapps.StatsCollector
+	statsBatcher *batchstats.Batcher
 
 	Acquirer *provisionerdserver.Acquirer
 	// dbRolluper rolls up template usage stats from raw agent and app
 	// stats. This is used to provide insights in the WebUI.
-	dbRolluper *dbrollup.Rolluper
+	dbRolluper            *dbrollup.Rolluper
+	workspaceUsageTracker *workspaceusage.Tracker
 }
 
 // Close waits for all WebSocket connections to drain before returning.
@@ -1310,6 +1320,7 @@ func (api *API) Close() error {
 		_ = (*coordinator).Close()
 	}
 	_ = api.agentProvider.Close()
+	api.workspaceUsageTracker.Close()
 	return nil
 }
 

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -71,6 +71,7 @@ import (
 	"github.com/coder/coder/v2/coderd/util/ptr"
 	"github.com/coder/coder/v2/coderd/workspaceapps"
 	"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
+	"github.com/coder/coder/v2/coderd/workspaceusage"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/codersdk/drpc"
@@ -334,6 +335,12 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 	if options.WorkspaceUsageTrackerTick == nil {
 		options.WorkspaceUsageTrackerTick = make(chan time.Time, 1) // buffering just in case
 	}
+	// Close is called by API.Close()
+	wuTracker := workspaceusage.New(
+		options.Database,
+		workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
+		workspaceusage.WithTickFlush(options.WorkspaceUsageTrackerTick, options.WorkspaceUsageTrackerFlush),
+	)
 
 	var mutex sync.RWMutex
 	var handler http.Handler
@@ -488,6 +495,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 			AllowWorkspaceRenames:              options.AllowWorkspaceRenames,
 			NewTicker:                          options.NewTicker,
 			DatabaseRolluper:                   options.DatabaseRolluper,
+			WorkspaceUsageTracker:              wuTracker,
 		}
 }
 

diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -36,7 +36,7 @@ import (
 	"github.com/coder/coder/v2/coderd/httpmw"
 	"github.com/coder/coder/v2/coderd/prometheusmetrics"
 	"github.com/coder/coder/v2/coderd/rbac/policy"
-	"github.com/coder/coder/v2/coderd/workspaceapps"
+	"github.com/coder/coder/v2/coderd/schedule"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/codersdk/workspacesdk"
@@ -1167,6 +1167,35 @@ func (api *API) workspaceAgentReportStats(rw http.ResponseWriter, r *http.Reques
 		slog.F("payload", req),
 	)
 
+	if req.ConnectionCount > 0 {
+		var nextAutostart time.Time
+		if workspace.AutostartSchedule.String != "" {
+			templateSchedule, err := (*(api.TemplateScheduleStore.Load())).Get(ctx, api.Database, workspace.TemplateID)
+			// If the template schedule fails to load, just default to bumping without the next transition and log it.
+			if err != nil {
+				// There's nothing we can do if the query was canceled, the
+				// client most likely went away so we just return an internal
+				// server error.
+				if database.IsQueryCanceledError(err) {
+					httpapi.InternalServerError(rw, err)
+					return
+				}
+				api.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(time.Now(), workspace.AutostartSchedule.String, templateSchedule)
+				if allowed {
+					nextAutostart = next
+				}
+			}
+		}
+		agentapi.ActivityBumpWorkspace(ctx, api.Logger.Named("activity_bump"), api.Database, workspace.ID, nextAutostart)
+	}
+
+	now := dbtime.Now()
 	protoStats := &agentproto.Stats{
 		ConnectionsByProto:          req.ConnectionsByProto,
 		ConnectionCount:             req.ConnectionCount,
@@ -1213,6 +1242,19 @@ func (api *API) workspaceAgentReportStats(rw http.ResponseWriter, r *http.Reques
 		}
 		return nil
 	})
+	if req.SessionCount() > 0 {
+		errGroup.Go(func() error {
+			// nolint:gocritic // (#13146) Will be moved soon as part of refactor.
+			err := api.Database.UpdateWorkspaceLastUsedAt(ctx, database.UpdateWorkspaceLastUsedAtParams{
+				ID:         workspace.ID,
+				LastUsedAt: now,
+			})
+			if err != nil {
+				return xerrors.Errorf("can't update workspace LastUsedAt: %w", err)
+			}
+			return nil
+		})
+	}
 	if api.Options.UpdateAgentMetrics != nil {
 		errGroup.Go(func() error {
 			user, err := api.Database.GetUserByID(ctx, workspace.OwnerID)
@@ -1235,13 +1277,6 @@ func (api *API) workspaceAgentReportStats(rw http.ResponseWriter, r *http.Reques
 		return
 	}
 
-	// Flushing the stats collector will update last_used_at,
-	// dealine for the workspace, and will publish a workspace update event.
-	api.statsCollector.CollectAndFlush(ctx, workspaceapps.StatsReport{
-		WorkspaceID: workspace.ID,
-		// TODO: fill out
-	})
-
 	httpapi.Write(ctx, rw, http.StatusOK, agentsdk.StatsResponse{
 		ReportInterval: api.AgentStatsRefreshInterval,
 	})

diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -29,7 +29,6 @@ import (
 	"github.com/coder/coder/v2/coderd/searchquery"
 	"github.com/coder/coder/v2/coderd/telemetry"
 	"github.com/coder/coder/v2/coderd/util/ptr"
-	"github.com/coder/coder/v2/coderd/workspaceapps"
 	"github.com/coder/coder/v2/coderd/wsbuilder"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -1116,10 +1115,7 @@ func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	api.statsCollector.CollectAndFlush(r.Context(), workspaceapps.StatsReport{
-		WorkspaceID: workspace.ID,
-	})
-
+	api.workspaceUsageTracker.Add(workspace.ID)
 	rw.WriteHeader(http.StatusNoContent)
 }