Add databasefake query for getting agent stats

kylecarbs · kylecarbs · commit ddf9841c2ccc · 2023-03-02T11:55:52.000-06:00
diff --git a/agent/agent.go b/agent/agent.go
@@ -60,7 +60,7 @@ const (
 
 	// MagicSSHSessionTypeEnvironmentVariable is used to track the purpose behind an SSH connection.
 	// This is stripped from any commands being executed, and is counted towards connection stats.
-	MagicSSHSessionTypeEnvironmentVariable = "__CODER_SSH_SESSION_TYPE"
+	MagicSSHSessionTypeEnvironmentVariable = "CODER_SSH_SESSION_TYPE"
 	// MagicSSHSessionTypeVSCode is set in the SSH config by the VS Code extension to identify itself.
 	MagicSSHSessionTypeVSCode = "vscode"
 	// MagicSSHSessionTypeJetBrains is set in the SSH config by the JetBrains extension to identify itself.
@@ -156,10 +156,6 @@ type agent struct {
 	network       *tailnet.Conn
 	connStatsChan chan *agentsdk.Stats
 
-	statRxPackets            atomic.Int64
-	statRxBytes              atomic.Int64
-	statTxPackets            atomic.Int64
-	statTxBytes              atomic.Int64
 	connCountVSCode          atomic.Int64
 	connCountJetBrains       atomic.Int64
 	connCountReconnectingPTY atomic.Int64
@@ -1188,14 +1184,12 @@ func (a *agent) startReportingConnectionStats(ctx context.Context) {
 			ConnectionCount:    int64(len(networkStats)),
 			ConnectionsByProto: map[string]int64{},
 		}
-		// Tailscale resets counts on every report!
-		// We'd rather have these compound, like Linux does!
 		for conn, counts := range networkStats {
 			stats.ConnectionsByProto[conn.Proto.String()]++
-			stats.RxBytes = a.statRxBytes.Add(int64(counts.RxBytes))
-			stats.RxPackets = a.statRxPackets.Add(int64(counts.RxPackets))
-			stats.TxBytes = a.statTxBytes.Add(int64(counts.TxBytes))
-			stats.TxPackets = a.statTxPackets.Add(int64(counts.TxPackets))
+			stats.RxBytes += int64(counts.RxBytes)
+			stats.RxPackets += int64(counts.RxPackets)
+			stats.TxBytes += int64(counts.TxBytes)
+			stats.TxPackets += int64(counts.TxPackets)
 		}
 
 		// The count of active sessions.
diff --git a/coderd/database/dbfake/databasefake.go b/coderd/database/dbfake/databasefake.go
@@ -263,6 +263,51 @@ func (*fakeQuerier) DeleteOldWorkspaceAgentStats(_ context.Context) error {
 	return nil
 }
 
+func (q *fakeQuerier) GetDeploymentWorkspaceAgentStats(_ context.Context, createdAfter time.Time) (database.GetDeploymentWorkspaceAgentStatsRow, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	agentStatsCreatedAfter := make([]database.WorkspaceAgentStat, 0)
+	for _, agentStat := range q.workspaceAgentStats {
+		if agentStat.CreatedAt.After(createdAfter) {
+			agentStatsCreatedAfter = append(agentStatsCreatedAfter, agentStat)
+		}
+	}
+
+	latestAgentStats := map[uuid.UUID]database.WorkspaceAgentStat{}
+	for _, agentStat := range q.workspaceAgentStats {
+		latestAgentStats[agentStat.AgentID] = agentStat
+	}
+
+	stat := database.GetDeploymentWorkspaceAgentStatsRow{}
+	for _, agentStat := range latestAgentStats {
+		stat.SessionCountVSCode += agentStat.SessionCountVSCode
+		stat.SessionCountJetBrains += agentStat.SessionCountJetBrains
+		stat.SessionCountReconnectingPTY += agentStat.SessionCountReconnectingPTY
+		stat.SessionCountSSH += agentStat.SessionCountSSH
+	}
+
+	latencies := make([]float64, 0)
+	for _, agentStat := range agentStatsCreatedAfter {
+		stat.WorkspaceRxBytes += agentStat.RxBytes
+		stat.WorkspaceTxBytes += agentStat.TxBytes
+		latencies = append(latencies, float64(agentStat.ConnectionMedianLatencyMS))
+	}
+
+	tryPercentile := func(fs []float64, p float64) float64 {
+		if len(fs) == 0 {
+			return -1
+		}
+		sort.Float64s(fs)
+		return fs[int(float64(len(fs))*p/100)]
+	}
+
+	stat.WorkspaceConnectionLatency50 = tryPercentile(latencies, 50)
+	stat.WorkspaceConnectionLatency95 = tryPercentile(latencies, 95)
+
+	return stat, nil
+}
+
 func (q *fakeQuerier) InsertWorkspaceAgentStat(_ context.Context, p database.InsertWorkspaceAgentStatParams) (database.WorkspaceAgentStat, error) {
 	if err := validateDatabaseType(p); err != nil {
 		return database.WorkspaceAgentStat{}, err
diff --git a/coderd/database/generate.sh b/coderd/database/generate.sh
@@ -18,7 +18,7 @@ SCRIPT_DIR=$(dirname "${BASH_SOURCE[0]}")
 	# Dump the updated schema (use make to utilize caching).
 	make -C ../.. --no-print-directory coderd/database/dump.sql
 	# The logic below depends on the exact version being correct :(
-	go run github.com/kyleconroy/sqlc/cmd/sqlc@v1.16.0 generate
+	sqlc generate
 
 	first=true
 	for fi in queries/*.sql.go; do
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
@@ -48,3 +48,21 @@ ORDER BY
 
 -- name: DeleteOldWorkspaceAgentStats :exec
 DELETE FROM workspace_agent_stats WHERE created_at < NOW() - INTERVAL '30 days';
+
+-- name: GetDeploymentWorkspaceAgentStats :one
+WITH agent_stats AS (
+	SELECT * FROM workspace_agent_stats
+		WHERE created_at > $1
+), latest_agent_stats AS (
+	SELECT * FROM agent_stats GROUP BY agent_id ORDER BY created_at
+)
+SELECT
+	SUM(latest_agent_stats.session_count_vscode) AS session_count_vscode,
+	SUM(latest_agent_stats.session_count_ssh) AS session_count_ssh,
+	SUM(latest_agent_stats.session_count_jetbrains) AS session_count_jetbrains,
+	SUM(latest_agent_stats.session_count_reconnecting_pty) AS session_count_reconnecting_pty,
+	SUM(agent_stats.rx_bytes) AS workspace_rx_bytes,
+	SUM(agent_stats.tx_bytes) AS workspace_tx_bytes,
+	coalesce((PERCENTILE_DISC(0.5) WITHIN GROUP(ORDER BY agent_stats.connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_50,
+	coalesce((PERCENTILE_DISC(0.95) WITHIN GROUP(ORDER BY agent_stats.connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_95
+ FROM agent_stats JOIN latest_agent_stats ON agent_stats.agent_id = latest_agent_stats.agent_id;
diff --git a/codersdk/deployment.go b/codersdk/deployment.go
@@ -528,3 +528,24 @@ func (c *Client) AppHost(ctx context.Context) (AppHostResponse, error) {
 	var host AppHostResponse
 	return host, json.NewDecoder(res.Body).Decode(&host)
 }
+
+type WorkspaceConnectionLatency struct {
+	P50 *int64
+	P95 *int64
+}
+
+type DeploymentStats struct {
+	WorkspacesByTransition       map[WorkspaceTransition]int
+	WorkspaceConnectionLatencyMS WorkspaceConnectionLatency
+
+	SessionCountVSCode          int
+	SessionCountSSH             int
+	SessionCountJetBrains       int
+	SessionCountReconnectingPTY int
+
+	WorkspaceRxBytes int
+	WorkspaceTxBytes int
+
+	// X downloaded in the past 10 minutes
+	// X uploaded in the past 10 minutes
+}
