feat(coderd): batch agent stats inserts (#8875)

This PR adds support for batching inserts to the workspace_agents_stats table.
Up to 1024 stats are batched, and flushed every second in a batch.

coderd/database/dbauthz/dbauthz.go

@@ -2016,6 +2016,14 @@ func (q *querier) InsertWorkspaceAgentStat(ctx context.Context, arg database.Ins
 	return q.db.InsertWorkspaceAgentStat(ctx, arg)
 }
 
+func (q *querier) InsertWorkspaceAgentStats(ctx context.Context, arg database.InsertWorkspaceAgentStatsParams) error {
+	if err := q.authorizeContext(ctx, rbac.ActionCreate, rbac.ResourceSystem); err != nil {
+		return err
+	}
+
+	return q.db.InsertWorkspaceAgentStats(ctx, arg)
+}
+
 func (q *querier) InsertWorkspaceApp(ctx context.Context, arg database.InsertWorkspaceAppParams) (database.WorkspaceApp, error) {
 	if err := q.authorizeContext(ctx, rbac.ActionCreate, rbac.ResourceSystem); err != nil {
 		return database.WorkspaceApp{}, err

coderd/database/dbfake/dbfake.go

@@ -2810,8 +2810,12 @@ func (q *FakeQuerier) GetWorkspaceAgentStats(_ context.Context, createdAfter tim
 	}
 
 	statByAgent := map[uuid.UUID]database.GetWorkspaceAgentStatsRow{}
-	for _, agentStat := range latestAgentStats {
-		stat := statByAgent[agentStat.AgentID]
+	for agentID, agentStat := range latestAgentStats {
+		stat := statByAgent[agentID]
+		stat.AgentID = agentStat.AgentID
+		stat.TemplateID = agentStat.TemplateID
+		stat.UserID = agentStat.UserID
+		stat.WorkspaceID = agentStat.WorkspaceID
 		stat.SessionCountVSCode += agentStat.SessionCountVSCode
 		stat.SessionCountJetBrains += agentStat.SessionCountJetBrains
 		stat.SessionCountReconnectingPTY += agentStat.SessionCountReconnectingPTY
@@ -4177,6 +4181,49 @@ func (q *FakeQuerier) InsertWorkspaceAgentStat(_ context.Context, p database.Ins
 	return stat, nil
 }
 
+func (q *FakeQuerier) InsertWorkspaceAgentStats(_ context.Context, arg database.InsertWorkspaceAgentStatsParams) error {
+	err := validateDatabaseType(arg)
+	if err != nil {
+		return err
+	}
+
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	var connectionsByProto []map[string]int64
+	if err := json.Unmarshal(arg.ConnectionsByProto, &connectionsByProto); err != nil {
+		return err
+	}
+	for i := 0; i < len(arg.ID); i++ {
+		cbp, err := json.Marshal(connectionsByProto[i])
+		if err != nil {
+			return xerrors.Errorf("failed to marshal connections_by_proto: %w", err)
+		}
+		stat := database.WorkspaceAgentStat{
+			ID:                          arg.ID[i],
+			CreatedAt:                   arg.CreatedAt[i],
+			WorkspaceID:                 arg.WorkspaceID[i],
+			AgentID:                     arg.AgentID[i],
+			UserID:                      arg.UserID[i],
+			ConnectionsByProto:          cbp,
+			ConnectionCount:             arg.ConnectionCount[i],
+			RxPackets:                   arg.RxPackets[i],
+			RxBytes:                     arg.RxBytes[i],
+			TxPackets:                   arg.TxPackets[i],
+			TxBytes:                     arg.TxBytes[i],
+			TemplateID:                  arg.TemplateID[i],
+			SessionCountVSCode:          arg.SessionCountVSCode[i],
+			SessionCountJetBrains:       arg.SessionCountJetBrains[i],
+			SessionCountReconnectingPTY: arg.SessionCountReconnectingPTY[i],
+			SessionCountSSH:             arg.SessionCountSSH[i],
+			ConnectionMedianLatencyMS:   arg.ConnectionMedianLatencyMS[i],
+		}
+		q.workspaceAgentStats = append(q.workspaceAgentStats, stat)
+	}
+
+	return nil
+}
+
 func (q *FakeQuerier) InsertWorkspaceApp(_ context.Context, arg database.InsertWorkspaceAppParams) (database.WorkspaceApp, error) {
 	if err := validateDatabaseType(arg); err != nil {
 		return database.WorkspaceApp{}, err

coderd/database/dbmetrics/dbmetrics.go

@@ -1236,6 +1236,13 @@ func (m metricsStore) InsertWorkspaceAgentStat(ctx context.Context, arg database
 	return stat, err
 }
 
+func (m metricsStore) InsertWorkspaceAgentStats(ctx context.Context, arg database.InsertWorkspaceAgentStatsParams) error {
+	start := time.Now()
+	r0 := m.s.InsertWorkspaceAgentStats(ctx, arg)
+	m.queryLatencies.WithLabelValues("InsertWorkspaceAgentStats").Observe(time.Since(start).Seconds())
+	return r0
+}
+
 func (m metricsStore) InsertWorkspaceApp(ctx context.Context, arg database.InsertWorkspaceAppParams) (database.WorkspaceApp, error) {
 	start := time.Now()
 	app, err := m.s.InsertWorkspaceApp(ctx, arg)

coderd/database/dbmock/dbmock.go

@@ -2598,6 +2598,20 @@ func (mr *MockStoreMockRecorder) InsertWorkspaceAgentStat(arg0, arg1 interface{}
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertWorkspaceAgentStat", reflect.TypeOf((*MockStore)(nil).InsertWorkspaceAgentStat), arg0, arg1)
 }
 
+// InsertWorkspaceAgentStats mocks base method.
+func (m *MockStore) InsertWorkspaceAgentStats(arg0 context.Context, arg1 database.InsertWorkspaceAgentStatsParams) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "InsertWorkspaceAgentStats", arg0, arg1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// InsertWorkspaceAgentStats indicates an expected call of InsertWorkspaceAgentStats.
+func (mr *MockStoreMockRecorder) InsertWorkspaceAgentStats(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertWorkspaceAgentStats", reflect.TypeOf((*MockStore)(nil).InsertWorkspaceAgentStats), arg0, arg1)
+}
+
 // InsertWorkspaceApp mocks base method.
 func (m *MockStore) InsertWorkspaceApp(arg0 context.Context, arg1 database.InsertWorkspaceAppParams) (database.WorkspaceApp, error) {
 	m.ctrl.T.Helper()

coderd/database/querier.go

@@ -225,6 +225,7 @@ type sqlcQuerier interface {
 	InsertWorkspaceAgentLogs(ctx context.Context, arg InsertWorkspaceAgentLogsParams) ([]WorkspaceAgentLog, error)
 	InsertWorkspaceAgentMetadata(ctx context.Context, arg InsertWorkspaceAgentMetadataParams) error
 	InsertWorkspaceAgentStat(ctx context.Context, arg InsertWorkspaceAgentStatParams) (WorkspaceAgentStat, error)
+	InsertWorkspaceAgentStats(ctx context.Context, arg InsertWorkspaceAgentStatsParams) error
 	InsertWorkspaceApp(ctx context.Context, arg InsertWorkspaceAppParams) (WorkspaceApp, error)
 	InsertWorkspaceBuild(ctx context.Context, arg InsertWorkspaceBuildParams) error
 	InsertWorkspaceBuildParameters(ctx context.Context, arg InsertWorkspaceBuildParametersParams) error

coderd/database/queries.sql.go

@@ -7418,6 +7418,90 @@ func (q *sqlQuerier) InsertWorkspaceAgentStat(ctx context.Context, arg InsertWor
 	return i, err
 }
 
+const insertWorkspaceAgentStats = `-- name: InsertWorkspaceAgentStats :exec
+INSERT INTO
+	workspace_agent_stats (
+		id,
+		created_at,
+		user_id,
+		workspace_id,
+		template_id,
+		agent_id,
+		connections_by_proto,
+		connection_count,
+		rx_packets,
+		rx_bytes,
+		tx_packets,
+		tx_bytes,
+		session_count_vscode,
+		session_count_jetbrains,
+		session_count_reconnecting_pty,
+		session_count_ssh,
+		connection_median_latency_ms
+	)
+SELECT
+	unnest($1 :: uuid[]) AS id,
+	unnest($2 :: timestamptz[]) AS created_at,
+	unnest($3 :: uuid[]) AS user_id,
+	unnest($4 :: uuid[]) AS workspace_id,
+	unnest($5 :: uuid[]) AS template_id,
+	unnest($6 :: uuid[]) AS agent_id,
+	jsonb_array_elements($7 :: jsonb) AS connections_by_proto,
+	unnest($8 :: bigint[]) AS connection_count,
+	unnest($9 :: bigint[]) AS rx_packets,
+	unnest($10 :: bigint[]) AS rx_bytes,
+	unnest($11 :: bigint[]) AS tx_packets,
+	unnest($12 :: bigint[]) AS tx_bytes,
+	unnest($13 :: bigint[]) AS session_count_vscode,
+	unnest($14 :: bigint[]) AS session_count_jetbrains,
+	unnest($15 :: bigint[]) AS session_count_reconnecting_pty,
+	unnest($16 :: bigint[]) AS session_count_ssh,
+	unnest($17 :: double precision[]) AS connection_median_latency_ms
+`
+
+type InsertWorkspaceAgentStatsParams struct {
+	ID                          []uuid.UUID     `db:"id" json:"id"`
+	CreatedAt                   []time.Time     `db:"created_at" json:"created_at"`
+	UserID                      []uuid.UUID     `db:"user_id" json:"user_id"`
+	WorkspaceID                 []uuid.UUID     `db:"workspace_id" json:"workspace_id"`
+	TemplateID                  []uuid.UUID     `db:"template_id" json:"template_id"`
+	AgentID                     []uuid.UUID     `db:"agent_id" json:"agent_id"`
+	ConnectionsByProto          json.RawMessage `db:"connections_by_proto" json:"connections_by_proto"`
+	ConnectionCount             []int64         `db:"connection_count" json:"connection_count"`
+	RxPackets                   []int64         `db:"rx_packets" json:"rx_packets"`
+	RxBytes                     []int64         `db:"rx_bytes" json:"rx_bytes"`
+	TxPackets                   []int64         `db:"tx_packets" json:"tx_packets"`
+	TxBytes                     []int64         `db:"tx_bytes" json:"tx_bytes"`
+	SessionCountVSCode          []int64         `db:"session_count_vscode" json:"session_count_vscode"`
+	SessionCountJetBrains       []int64         `db:"session_count_jetbrains" json:"session_count_jetbrains"`
+	SessionCountReconnectingPTY []int64         `db:"session_count_reconnecting_pty" json:"session_count_reconnecting_pty"`
+	SessionCountSSH             []int64         `db:"session_count_ssh" json:"session_count_ssh"`
+	ConnectionMedianLatencyMS   []float64       `db:"connection_median_latency_ms" json:"connection_median_latency_ms"`
+}
+
+func (q *sqlQuerier) InsertWorkspaceAgentStats(ctx context.Context, arg InsertWorkspaceAgentStatsParams) error {
+	_, err := q.db.ExecContext(ctx, insertWorkspaceAgentStats,
+		pq.Array(arg.ID),
+		pq.Array(arg.CreatedAt),
+		pq.Array(arg.UserID),
+		pq.Array(arg.WorkspaceID),
+		pq.Array(arg.TemplateID),
+		pq.Array(arg.AgentID),
+		arg.ConnectionsByProto,
+		pq.Array(arg.ConnectionCount),
+		pq.Array(arg.RxPackets),
+		pq.Array(arg.RxBytes),
+		pq.Array(arg.TxPackets),
+		pq.Array(arg.TxBytes),
+		pq.Array(arg.SessionCountVSCode),
+		pq.Array(arg.SessionCountJetBrains),
+		pq.Array(arg.SessionCountReconnectingPTY),
+		pq.Array(arg.SessionCountSSH),
+		pq.Array(arg.ConnectionMedianLatencyMS),
+	)
+	return err
+}
+
 const getWorkspaceAppByAgentIDAndSlug = `-- name: GetWorkspaceAppByAgentIDAndSlug :one
 SELECT id, created_at, agent_id, display_name, icon, command, url, healthcheck_url, healthcheck_interval, healthcheck_threshold, health, subdomain, sharing_level, slug, external FROM workspace_apps WHERE agent_id = $1 AND slug = $2
 `

coderd/database/queries/workspaceagentstats.sql

@@ -22,6 +22,46 @@ INSERT INTO
 VALUES
 	($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17) RETURNING *;
 
+-- name: InsertWorkspaceAgentStats :exec
+INSERT INTO
+	workspace_agent_stats (
+		id,
+		created_at,
+		user_id,
+		workspace_id,
+		template_id,
+		agent_id,
+		connections_by_proto,
+		connection_count,
+		rx_packets,
+		rx_bytes,
+		tx_packets,
+		tx_bytes,
+		session_count_vscode,
+		session_count_jetbrains,
+		session_count_reconnecting_pty,
+		session_count_ssh,
+		connection_median_latency_ms
+	)
+SELECT
+	unnest(@id :: uuid[]) AS id,
+	unnest(@created_at :: timestamptz[]) AS created_at,
+	unnest(@user_id :: uuid[]) AS user_id,
+	unnest(@workspace_id :: uuid[]) AS workspace_id,
+	unnest(@template_id :: uuid[]) AS template_id,
+	unnest(@agent_id :: uuid[]) AS agent_id,
+	jsonb_array_elements(@connections_by_proto :: jsonb) AS connections_by_proto,
+	unnest(@connection_count :: bigint[]) AS connection_count,
+	unnest(@rx_packets :: bigint[]) AS rx_packets,
+	unnest(@rx_bytes :: bigint[]) AS rx_bytes,
+	unnest(@tx_packets :: bigint[]) AS tx_packets,
+	unnest(@tx_bytes :: bigint[]) AS tx_bytes,
+	unnest(@session_count_vscode :: bigint[]) AS session_count_vscode,
+	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;
+
 -- name: GetTemplateDAUs :many
 SELECT
 	(created_at at TIME ZONE cast(@tz_offset::integer as text))::date as date,