connect usage flag to metrics cache

coder · ethanndickson · Sep 19, 2024 · Sep 5, 2024 · Sep 11, 2024 · Sep 17, 2024
commit f62203c9171f628ce8bcc10b24785e4b12492ff8
diff --git a/coderd/agentapi/stats.go b/coderd/agentapi/stats.go
@@ -74,6 +74,7 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
 		workspaceAgent,
 		getWorkspaceAgentByIDRow.TemplateName,
 		req.Stats,
+		false,
 	)
 	if err != nil {
 		return nil, xerrors.Errorf("report agent stats: %w", err)

diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -421,13 +421,16 @@ func New(options *Options) *API {
 		panic(xerrors.Errorf("read site bin failed: %w", err))
 	}
 
+	experiments.Enabled(codersdk.ExperimentWorkspaceUsage)
+
 	metricsCache := metricscache.New(
 		options.Database,
 		options.Logger.Named("metrics_cache"),
 		metricscache.Intervals{
 			TemplateBuildTimes: options.MetricsCacheRefreshInterval,
 			DeploymentStats:    options.AgentStatsRefreshInterval,
 		},
+		experiments.Enabled(codersdk.ExperimentWorkspaceUsage),
 	)
 
 	oauthConfigs := &httpmw.OAuth2Configs{

diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaceagentstats.sql b/coderd/database/queries/workspaceagentstats.sql
@@ -17,7 +17,8 @@ INSERT INTO
 		session_count_jetbrains,
 		session_count_reconnecting_pty,
 		session_count_ssh,
-		connection_median_latency_ms
+		connection_median_latency_ms,
+		usage
 	)
 SELECT
 	unnest(@id :: uuid[]) AS id,
@@ -36,7 +37,8 @@ SELECT
 	unnest(@session_count_jetbrains :: bigint[]) AS session_count_jetbrains,
 	unnest(@session_count_reconnecting_pty :: bigint[]) AS session_count_reconnecting_pty,
 	unnest(@session_count_ssh :: bigint[]) AS session_count_ssh,
-	unnest(@connection_median_latency_ms :: double precision[]) AS connection_median_latency_ms;
+	unnest(@connection_median_latency_ms :: double precision[]) AS connection_median_latency_ms,
+	unnest(@usage :: boolean[]) AS usage;
 
 -- name: GetTemplateDAUs :many
 SELECT
@@ -135,14 +137,15 @@ minute_buckets AS (
 		agent_id,
 		date_trunc('minute', created_at) AS minute_bucket,
 		SUM(session_count_vscode) AS session_count_vscode,
+		SUM(session_count_ssh) AS session_count_ssh,
 		SUM(session_count_jetbrains) AS session_count_jetbrains,
-		SUM(session_count_reconnecting_pty) AS session_count_reconnecting_pty,
-		SUM(session_count_ssh) AS session_count_ssh
+		SUM(session_count_reconnecting_pty) AS session_count_reconnecting_pty
 	FROM
 		workspace_agent_stats
 	WHERE
 		created_at >= $1
 		AND created_at < date_trunc('minute', now())  -- Exclude current partial minute
+		AND usage = true
 	GROUP BY
 		agent_id,
 		minute_bucket
@@ -164,9 +167,9 @@ latest_buckets AS (
 latest_agent_stats AS (
 	SELECT
 		SUM(session_count_vscode) AS session_count_vscode,
+		SUM(session_count_ssh) AS session_count_ssh,
 		SUM(session_count_jetbrains) AS session_count_jetbrains,
-		SUM(session_count_reconnecting_pty) AS session_count_reconnecting_pty,
-		SUM(session_count_ssh) AS session_count_ssh
+		SUM(session_count_reconnecting_pty) AS session_count_reconnecting_pty
 	FROM
 		latest_buckets
 )
@@ -223,14 +226,15 @@ minute_buckets AS (
 		agent_id,
 		date_trunc('minute', created_at) AS minute_bucket,
 		SUM(session_count_vscode) AS session_count_vscode,
+		SUM(session_count_ssh) AS session_count_ssh,		
 		SUM(session_count_jetbrains) AS session_count_jetbrains,
-		SUM(session_count_reconnecting_pty) AS session_count_reconnecting_pty,
-		SUM(session_count_ssh) AS session_count_ssh
+		SUM(session_count_reconnecting_pty) AS session_count_reconnecting_pty
 	FROM
 		workspace_agent_stats
 	WHERE
 		created_at >= $1
 		AND created_at < date_trunc('minute', now())  -- Exclude current partial minute
+		AND usage = true
 	GROUP BY
 		agent_id,
 		minute_bucket,
@@ -256,9 +260,9 @@ latest_buckets AS (
 	SELECT
 		agent_id,
 		SUM(session_count_vscode) AS session_count_vscode,
+		SUM(session_count_ssh) AS session_count_ssh,
 		SUM(session_count_jetbrains) AS session_count_jetbrains,
-		SUM(session_count_reconnecting_pty) AS session_count_reconnecting_pty,
-		SUM(session_count_ssh) AS session_count_ssh
+		SUM(session_count_reconnecting_pty) AS session_count_reconnecting_pty
 	FROM
 		latest_buckets
 )
@@ -336,6 +340,7 @@ WITH agent_stats AS (
 		coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty,
 		coalesce(SUM(connection_count), 0)::bigint AS connection_count
 	FROM workspace_agent_stats
+	WHERE usage = true
 	GROUP BY user_id, agent_id, workspace_id
 ), latest_agent_latencies AS (
 	SELECT

diff --git a/coderd/insights_test.go b/coderd/insights_test.go
@@ -706,7 +706,7 @@ func TestTemplateInsights_Golden(t *testing.T) {
 						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)
 				}
@@ -1605,7 +1605,7 @@ func TestUserActivityInsights_Golden(t *testing.T) {
 						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/metricscache/metricscache.go b/coderd/metricscache/metricscache.go
@@ -34,14 +34,18 @@ type Cache struct {
 
 	done   chan struct{}
 	cancel func()
+
+	// usage is a experiment flag to enable new workspace usage tracking behavior and will be
+	// removed when the experiment is complete.
+	usage bool
 }
 
 type Intervals struct {
 	TemplateBuildTimes time.Duration
 	DeploymentStats    time.Duration
 }
 
-func New(db database.Store, log slog.Logger, intervals Intervals) *Cache {
+func New(db database.Store, log slog.Logger, intervals Intervals, usage bool) *Cache {
 	if intervals.TemplateBuildTimes <= 0 {
 		intervals.TemplateBuildTimes = time.Hour
 	}
@@ -56,6 +60,7 @@ func New(db database.Store, log slog.Logger, intervals Intervals) *Cache {
 		log:       log,
 		done:      make(chan struct{}),
 		cancel:    cancel,
+		usage:     usage,
 	}
 	go func() {
 		var wg sync.WaitGroup
@@ -125,11 +130,25 @@ func (c *Cache) refreshTemplateBuildTimes(ctx context.Context) error {
 }
 
 func (c *Cache) refreshDeploymentStats(ctx context.Context) error {
-	from := dbtime.Now().Add(-15 * time.Minute)
-	agentStats, err := c.database.GetDeploymentWorkspaceAgentStats(ctx, from)
-	if err != nil {
-		return err
+	var (
+		from       = dbtime.Now().Add(-15 * time.Minute)
+		agentStats database.GetDeploymentWorkspaceAgentStatsRow
+		err        error
+	)
+
+	if c.usage {
+		agentUsageStats, err := c.database.GetDeploymentWorkspaceAgentUsageStats(ctx, from)
+		if err != nil {
+			return err
+		}
+		agentStats = database.GetDeploymentWorkspaceAgentStatsRow(agentUsageStats)
+	} else {
+		agentStats, err = c.database.GetDeploymentWorkspaceAgentStats(ctx, from)
+		if err != nil {
+			return err
+		}
 	}
+
 	workspaceStats, err := c.database.GetDeploymentWorkspaceStats(ctx)
 	if err != nil {
 		return err

diff --git a/coderd/metricscache/metricscache_test.go b/coderd/metricscache/metricscache_test.go
@@ -32,7 +32,7 @@ func TestCache_TemplateWorkspaceOwners(t *testing.T) {
 		db    = dbmem.New()
 		cache = metricscache.New(db, slogtest.Make(t, nil), metricscache.Intervals{
 			TemplateBuildTimes: testutil.IntervalFast,
-		})
+		}, false)
 	)
 
 	defer cache.Close()
@@ -183,7 +183,7 @@ func TestCache_BuildTime(t *testing.T) {
 				db    = dbmem.New()
 				cache = metricscache.New(db, slogtest.Make(t, nil), metricscache.Intervals{
 					TemplateBuildTimes: testutil.IntervalFast,
-				})
+				}, false)
 			)
 
 			defer cache.Close()
@@ -278,7 +278,7 @@ func TestCache_DeploymentStats(t *testing.T) {
 	db := dbmem.New()
 	cache := metricscache.New(db, slogtest.Make(t, nil), metricscache.Intervals{
 		DeploymentStats: testutil.IntervalFast,
-	})
+	}, false)
 	defer cache.Close()
 
 	err := db.InsertWorkspaceAgentStats(context.Background(), database.InsertWorkspaceAgentStatsParams{

diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -1340,7 +1340,7 @@ func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	err = api.statsReporter.ReportAgentStats(ctx, dbtime.Now(), workspace, agent, template.Name, stat)
+	err = api.statsReporter.ReportAgentStats(ctx, dbtime.Now(), workspace, agent, template.Name, stat, true)
 	if err != nil {
 		httpapi.InternalServerError(rw, err)
 		return