feat: collect

coder · mtojek · Apr 14, 2023 · Apr 3, 2023 · Apr 4, 2023 · Apr 4, 2023
commit 1ed37b45ba89e43d5c81683c5243f432aec19031
diff --git a/cli/server.go b/cli/server.go
@@ -844,7 +844,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 				}
 				defer closeWorkspacesFunc()
 
-				closeAgentStatsFunc, err := prometheusmetrics.AgentStats(ctx, logger, options.PrometheusRegistry, options.Database, 0)
+				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)
 				}

diff --git a/coderd/database/dbfake/databasefake.go b/coderd/database/dbfake/databasefake.go
@@ -3998,8 +3998,75 @@ func (q *fakeQuerier) GetWorkspaceAgentStats(_ context.Context, createdAfter tim
 	return stats, nil
 }
 
-func (q *fakeQuerier) GetWorkspaceAgentStatsAndLabels(_ context.Context, createdAfter time.Time) ([]database.GetWorkspaceAgentStatsAndLabelsRow, error) {
-	panic("not implemented yet")
+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.WorkspaceRxBytes += agentStat.RxBytes
+		stat.WorkspaceTxBytes += 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) {

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
@@ -123,7 +123,7 @@ WITH agent_stats AS (
 		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(SUM(connection_median_latency_ms), 0)::float AS connection_median_latency_ms
+		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

diff --git a/coderd/prometheusmetrics/prometheusmetrics.go b/coderd/prometheusmetrics/prometheusmetrics.go
@@ -291,7 +291,7 @@ func Agents(ctx context.Context, logger slog.Logger, registerer prometheus.Regis
 	return cancelFunc, nil
 }
 
-func AgentStats(ctx context.Context, logger slog.Logger, registerer prometheus.Registerer, db database.Store, duration time.Duration) (context.CancelFunc, error) {
+func AgentStats(ctx context.Context, logger slog.Logger, registerer prometheus.Registerer, db database.Store, initialCreateAfter time.Time, duration time.Duration) (context.CancelFunc, error) {
 	if duration == 0 {
 		duration = 1 * time.Minute
 	}
@@ -344,8 +344,8 @@ func AgentStats(ctx context.Context, logger slog.Logger, registerer prometheus.R
 	agentStatsConnectionMedianLatencyGauge := NewCachedGaugeVec(prometheus.NewGaugeVec(prometheus.GaugeOpts{
 		Namespace: "coderd",
 		Subsystem: "agentstats",
-		Name:      "connection_median_latency",
-		Help:      "The median agent connection latency",
+		Name:      "connection_median_latency_seconds",
+		Help:      "The median agent connection latency in seconds",
 	}, []string{"agent_name", "username", "workspace_name"}))
 	err = registerer.Register(agentStatsConnectionMedianLatencyGauge)
 	if err != nil {
@@ -396,7 +396,7 @@ func AgentStats(ctx context.Context, logger slog.Logger, registerer prometheus.R
 		return nil, err
 	}
 
-	createdAfter := time.Now()
+	createdAfter := initialCreateAfter
 	ctx, cancelFunc := context.WithCancel(ctx)
 	ticker := time.NewTicker(duration)
 	go func() {
@@ -411,12 +411,18 @@ func AgentStats(ctx context.Context, logger slog.Logger, registerer prometheus.R
 			logger.Debug(ctx, "Agent metrics collection is starting")
 			timer := prometheus.NewTimer(metricsCollectorAgentStats)
 
+			checkpoint := time.Now()
 			stats, err := db.GetWorkspaceAgentStatsAndLabels(ctx, createdAfter)
+
 			if err != nil {
 				logger.Error(ctx, "can't get agent stats", slog.Error(err))
 				goto done
 			}
 
+			if len(stats) == 0 {
+				goto done
+			}
+
 			for _, agentStat := range stats {
 				agentStatsRxBytesGauge.WithLabelValues(VectorOperationAdd, float64(agentStat.WorkspaceTxBytes), agentStat.AgentName, agentStat.Username, agentStat.WorkspaceName)
 				agentStatsTxBytesGauge.WithLabelValues(VectorOperationAdd, float64(agentStat.WorkspaceRxBytes), agentStat.AgentName, agentStat.Username, agentStat.WorkspaceName)
@@ -445,7 +451,7 @@ func AgentStats(ctx context.Context, logger slog.Logger, registerer prometheus.R
 			logger.Debug(ctx, "Agent metrics collection is done")
 			metricsCollectorAgentStats.Observe(timer.ObserveDuration().Seconds())
 
-			createdAfter = time.Now()
+			createdAfter = checkpoint
 		}
 	}()
 	return cancelFunc, nil

diff --git a/coderd/prometheusmetrics/prometheusmetrics_test.go b/coderd/prometheusmetrics/prometheusmetrics_test.go
@@ -20,6 +20,7 @@ import (
 	"github.com/coder/coder/coderd/database/dbgen"
 	"github.com/coder/coder/coderd/prometheusmetrics"
 	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/codersdk/agentsdk"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
 	"github.com/coder/coder/tailnet"
@@ -352,3 +353,86 @@ func TestAgents(t *testing.T) {
 		return agentsUp && agentsConnections && agentsApps && agentsExecutionInSeconds
 	}, testutil.WaitShort, testutil.IntervalFast)
 }
+
+func TestAgentStats(t *testing.T) {
+	t.Parallel()
+
+	// Build a sample workspace with test agent and fake agent client
+	client, _, api := coderdtest.NewWithAPI(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+	db := api.Database
+
+	user := coderdtest.CreateFirstUser(t, client)
+	authToken := uuid.NewString()
+	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+		Parse:          echo.ParseComplete,
+		ProvisionPlan:  echo.ProvisionComplete,
+		ProvisionApply: echo.ProvisionApplyWithAgent(authToken),
+	})
+	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+	coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+	workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+	coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+
+	agentClient := agentsdk.New(client.URL)
+	agentClient.SetSessionToken(authToken)
+
+	registry := prometheus.NewRegistry()
+
+	// given
+	cancel, err := prometheusmetrics.AgentStats(context.Background(), slogtest.Make(t, nil), registry, db, time.Now(), time.Second)
+	require.NoError(t, err)
+	t.Cleanup(cancel)
+
+	// when
+	_, err = agentClient.PostStats(context.Background(), &agentsdk.Stats{
+		ConnectionsByProto:          map[string]int64{"TCP": 1},
+		ConnectionCount:             2,
+		RxPackets:                   3,
+		RxBytes:                     4,
+		TxPackets:                   5,
+		TxBytes:                     6,
+		SessionCountVSCode:          7,
+		SessionCountJetBrains:       8,
+		SessionCountReconnectingPTY: 9,
+		SessionCountSSH:             10,
+		ConnectionMedianLatencyMS:   10000,
+	})
+	require.NoError(t, err)
+
+	// then
+	require.NoError(t, err)
+
+	collectedMetrics := map[string]struct{}{}
+	require.Eventually(t, func() bool {
+		metrics, err := registry.Gather()
+		assert.NoError(t, err)
+
+		if len(metrics) < 1 {
+			return false
+		}
+
+		for _, metric := range metrics {
+			switch metric.GetName() {
+			case "coderd_prometheusmetrics_agentstats_execution_seconds":
+				collectedMetrics[metric.GetName()] = struct{}{}
+			case "coderd_agentstats_connection_count",
+				"coderd_agentstats_connection_median_latency_seconds",
+				"coderd_agentstats_rx_bytes",
+				"coderd_agentstats_tx_bytes",
+				"coderd_agentstats_session_count_jetbrains",
+				"coderd_agentstats_session_count_reconnecting_pty",
+				"coderd_agentstats_session_count_ssh",
+				"coderd_agentstats_session_count_vscode":
+				collectedMetrics[metric.GetName()] = struct{}{}
+				assert.Equal(t, "example", metric.Metric[0].Label[0].GetValue())            // Agent name
+				assert.Equal(t, "testuser", metric.Metric[0].Label[1].GetValue())           // Username
+				assert.Equal(t, workspace.Name, metric.Metric[0].Label[2].GetValue())       // Workspace name
+				assert.NotZero(t, int(metric.Metric[0].Gauge.GetValue()), metric.GetName()) // Metric value
+			default:
+				require.FailNowf(t, "unexpected metric collected", "metric: %s", metric.GetName())
+			}
+		}
+
+		return len(collectedMetrics) == 9
+	}, testutil.WaitShort, testutil.IntervalFast, "collected metrics: %v", collectedMetrics)
+}