coder · mtojek · Apr 14, 2023 · Apr 3, 2023 · Apr 4, 2023 · Apr 4, 2023
diff --git a/cli/server.go b/cli/server.go
@@ -844,6 +844,14 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 				}
 				defer closeWorkspacesFunc()
 
+				if cfg.Prometheus.CollectAgentStats {
+					closeAgentStatsFunc, err := prometheusmetrics.AgentStats(ctx, logger, options.PrometheusRegistry, options.Database, time.Now(), 0)
+					if err != nil {
+						return xerrors.Errorf("register agent stats prometheus metric: %w", err)
+					}
+					defer closeAgentStatsFunc()
+				}
+
 				//nolint:revive
 				defer serveHandler(ctx, logger, promhttp.InstrumentMetricHandler(
 					options.PrometheusRegistry, promhttp.HandlerFor(options.PrometheusRegistry, promhttp.HandlerOpts{}),

diff --git a/cli/testdata/coder_server_--help.golden b/cli/testdata/coder_server_--help.golden
@@ -90,6 +90,9 @@ Use a YAML configuration file when your server launch become unwieldy.
       --prometheus-address host:port, $CODER_PROMETHEUS_ADDRESS (default: 127.0.0.1:2112)
           The bind address to serve prometheus metrics.
 
+      --prometheus-collect-agent-stats bool, $CODER_PROMETHEUS_COLLECT_AGENT_STATS
+          Collect agent stats (may increase charges for metrics storage).
+
       --prometheus-enable bool, $CODER_PROMETHEUS_ENABLE
           Serve prometheus metrics on the address defined by prometheus address.
 

diff --git a/cli/testdata/server-config.yaml.golden b/cli/testdata/server-config.yaml.golden
@@ -146,6 +146,9 @@ introspection:
     # The bind address to serve prometheus metrics.
     # (default: 127.0.0.1:2112, type: host:port)
     address: 127.0.0.1:2112
+    # Collect agent stats (may increase charges for metrics storage).
+    # (default: <unset>, type: bool)
+    collect_agent_stats: false
   pprof:
     # Serve pprof metrics on the address defined by pprof address.
     # (default: <unset>, type: bool)

diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/database/dbauthz/system.go b/coderd/database/dbauthz/system.go
@@ -302,6 +302,10 @@ func (q *querier) GetWorkspaceAgentStats(ctx context.Context, createdAfter time.
 	return q.db.GetWorkspaceAgentStats(ctx, createdAfter)
 }
 
+func (q *querier) GetWorkspaceAgentStatsAndLabels(ctx context.Context, createdAfter time.Time) ([]database.GetWorkspaceAgentStatsAndLabelsRow, error) {
+	return q.db.GetWorkspaceAgentStatsAndLabels(ctx, createdAfter)
+}
+
 func (q *querier) GetDeploymentWorkspaceStats(ctx context.Context) (database.GetDeploymentWorkspaceStatsRow, error) {
 	return q.db.GetDeploymentWorkspaceStats(ctx)
 }

diff --git a/coderd/database/dbfake/databasefake.go b/coderd/database/dbfake/databasefake.go
@@ -3998,6 +3998,77 @@ func (q *fakeQuerier) GetWorkspaceAgentStats(_ context.Context, createdAfter tim
 	return stats, nil
 }
 
+func (q *fakeQuerier) GetWorkspaceAgentStatsAndLabels(ctx context.Context, createdAfter time.Time) ([]database.GetWorkspaceAgentStatsAndLabelsRow, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	agentStatsCreatedAfter := make([]database.WorkspaceAgentStat, 0)
+	latestAgentStats := map[uuid.UUID]database.WorkspaceAgentStat{}
+
+	for _, agentStat := range q.workspaceAgentStats {
+		if agentStat.CreatedAt.After(createdAfter) {
+			agentStatsCreatedAfter = append(agentStatsCreatedAfter, agentStat)
+			latestAgentStats[agentStat.AgentID] = agentStat
+		}
+	}
+
+	statByAgent := map[uuid.UUID]database.GetWorkspaceAgentStatsAndLabelsRow{}
+
+	// Session and connection metrics
+	for _, agentStat := range latestAgentStats {
+		stat := statByAgent[agentStat.AgentID]
+		stat.SessionCountVSCode += agentStat.SessionCountVSCode
+		stat.SessionCountJetBrains += agentStat.SessionCountJetBrains
+		stat.SessionCountReconnectingPTY += agentStat.SessionCountReconnectingPTY
+		stat.SessionCountSSH += agentStat.SessionCountSSH
+		stat.ConnectionCount += agentStat.ConnectionCount
+		if agentStat.ConnectionMedianLatencyMS >= 0 && stat.ConnectionMedianLatencyMS < agentStat.ConnectionMedianLatencyMS {
+			stat.ConnectionMedianLatencyMS = agentStat.ConnectionMedianLatencyMS
+		}
+		statByAgent[agentStat.AgentID] = stat
+	}
+
+	// Tx, Rx metrics
+	for _, agentStat := range agentStatsCreatedAfter {
+		stat := statByAgent[agentStat.AgentID]
+		stat.RxBytes += agentStat.RxBytes
+		stat.TxBytes += agentStat.TxBytes
+		statByAgent[agentStat.AgentID] = stat
+	}
+
+	// Labels
+	for _, agentStat := range agentStatsCreatedAfter {
+		stat := statByAgent[agentStat.AgentID]
+
+		user, err := q.getUserByIDNoLock(agentStat.UserID)
+		if err != nil {
+			return nil, err
+		}
+
+		stat.Username = user.Username
+
+		workspace, err := q.GetWorkspaceByID(ctx, agentStat.WorkspaceID)
+		if err != nil {
+			return nil, err
+		}
+		stat.WorkspaceName = workspace.Name
+
+		agent, err := q.GetWorkspaceAgentByID(ctx, agentStat.AgentID)
+		if err != nil {
+			return nil, err
+		}
+		stat.AgentName = agent.Name
+
+		statByAgent[agentStat.AgentID] = stat
+	}
+
+	stats := make([]database.GetWorkspaceAgentStatsAndLabelsRow, 0, len(statByAgent))
+	for _, agent := range statByAgent {
+		stats = append(stats, agent)
+	}
+	return stats, nil
+}
+
 func (q *fakeQuerier) GetWorkspacesEligibleForAutoStartStop(ctx context.Context, now time.Time) ([]database.Workspace, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()

diff --git a/coderd/database/querier.go b/coderd/database/querier.go
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
@@ -103,3 +103,55 @@ WITH agent_stats AS (
 	) AS a WHERE a.rn = 1 GROUP BY a.user_id, a.agent_id, a.workspace_id, a.template_id
 )
 SELECT * FROM agent_stats JOIN latest_agent_stats ON agent_stats.agent_id = latest_agent_stats.agent_id;
+
+-- name: GetWorkspaceAgentStatsAndLabels :many
+WITH agent_stats AS (
+	SELECT
+		user_id,
+		agent_id,
+		workspace_id,
+		coalesce(SUM(rx_bytes), 0)::bigint AS rx_bytes,
+		coalesce(SUM(tx_bytes), 0)::bigint AS tx_bytes
+	 FROM workspace_agent_stats
+		WHERE workspace_agent_stats.created_at > $1
+		GROUP BY user_id, agent_id, workspace_id
+), latest_agent_stats AS (
+	SELECT
+		a.agent_id,
+		coalesce(SUM(session_count_vscode), 0)::bigint AS session_count_vscode,
+		coalesce(SUM(session_count_ssh), 0)::bigint AS session_count_ssh,
+		coalesce(SUM(session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
+		coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty,
+		coalesce(SUM(connection_count), 0)::bigint AS connection_count,
+		coalesce(MAX(connection_median_latency_ms), 0)::float AS connection_median_latency_ms
+	 FROM (
+		SELECT *, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn
+		FROM workspace_agent_stats
+		-- The greater than 0 is to support legacy agents that don't report connection_median_latency_ms.
+		WHERE created_at > $1 AND connection_median_latency_ms > 0
+	) AS a
+	WHERE a.rn = 1
+	GROUP BY a.user_id, a.agent_id, a.workspace_id
+)
+SELECT
+	users.username, workspace_agents.name AS agent_name, workspaces.name AS workspace_name, rx_bytes, tx_bytes,
+	session_count_vscode, session_count_ssh, session_count_jetbrains, session_count_reconnecting_pty,
+	connection_count, connection_median_latency_ms
+FROM
+	agent_stats
+JOIN
+	latest_agent_stats
+ON
+	agent_stats.agent_id = latest_agent_stats.agent_id
+JOIN
+	users
+ON
+	users.id = agent_stats.user_id
+JOIN
+	workspace_agents
+ON
+	workspace_agents.id = agent_stats.agent_id
+JOIN
+	workspaces
+ON
+	workspaces.id = agent_stats.workspace_id;