diff --git a/coderd/agentapi/stats.go b/coderd/agentapi/stats.go index 4f6a6da1c8c66..226f06732d4ee 100644 --- 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 index 20ce616eab5ba..4e59724e7b781 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -421,6 +421,8 @@ 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"), @@ -428,6 +430,7 @@ func New(options *Options) *API { TemplateBuildTimes: options.MetricsCacheRefreshInterval, DeploymentStats: options.AgentStatsRefreshInterval, }, + experiments.Enabled(codersdk.ExperimentWorkspaceUsage), ) oauthConfigs := &httpmw.OAuth2Configs{ diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index f6bd03cc50e8b..3b0b56aaffc57 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -1418,6 +1418,10 @@ func (q *querier) GetDeploymentWorkspaceAgentStats(ctx context.Context, createdA return q.db.GetDeploymentWorkspaceAgentStats(ctx, createdAfter) } +func (q *querier) GetDeploymentWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) (database.GetDeploymentWorkspaceAgentUsageStatsRow, error) { + return q.db.GetDeploymentWorkspaceAgentUsageStats(ctx, createdAt) +} + func (q *querier) GetDeploymentWorkspaceStats(ctx context.Context) (database.GetDeploymentWorkspaceStatsRow, error) { return q.db.GetDeploymentWorkspaceStats(ctx) } @@ -2360,6 +2364,14 @@ func (q *querier) GetWorkspaceAgentStatsAndLabels(ctx context.Context, createdAf return q.db.GetWorkspaceAgentStatsAndLabels(ctx, createdAfter) } +func (q *querier) GetWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentUsageStatsRow, error) { + return q.db.GetWorkspaceAgentUsageStats(ctx, createdAt) +} + +func (q *querier) GetWorkspaceAgentUsageStatsAndLabels(ctx context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentUsageStatsAndLabelsRow, error) { + return q.db.GetWorkspaceAgentUsageStatsAndLabels(ctx, createdAt) +} + // GetWorkspaceAgentsByResourceIDs // The workspace/job is already fetched. func (q *querier) GetWorkspaceAgentsByResourceIDs(ctx context.Context, ids []uuid.UUID) ([]database.WorkspaceAgent, error) { diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index ccacb0dc0a995..4e31c5a55d32f 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -800,6 +800,7 @@ func WorkspaceAgentStat(t testing.TB, db database.Store, orig database.Workspace SessionCountReconnectingPTY: []int64{takeFirst(orig.SessionCountReconnectingPTY, 0)}, SessionCountSSH: []int64{takeFirst(orig.SessionCountSSH, 0)}, ConnectionMedianLatencyMS: []float64{takeFirst(orig.ConnectionMedianLatencyMS, 0)}, + Usage: []bool{takeFirst(orig.Usage, false)}, } err := db.InsertWorkspaceAgentStats(genCtx, params) require.NoError(t, err, "insert workspace agent stat") @@ -822,6 +823,7 @@ func WorkspaceAgentStat(t testing.TB, db database.Store, orig database.Workspace SessionCountJetBrains: params.SessionCountJetBrains[0], SessionCountReconnectingPTY: params.SessionCountReconnectingPTY[0], SessionCountSSH: params.SessionCountSSH[0], + Usage: params.Usage[0], } } diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go index b1d2178e66a29..f06c772c57c76 100644 --- a/coderd/database/dbmem/dbmem.go +++ b/coderd/database/dbmem/dbmem.go @@ -7,6 +7,7 @@ import ( "encoding/json" "errors" "fmt" + "math" "reflect" "regexp" "sort" @@ -213,6 +214,19 @@ type data struct { defaultProxyIconURL string } +func tryPercentile(fs []float64, p float64) float64 { + if len(fs) == 0 { + return -1 + } + sort.Float64s(fs) + pos := p * (float64(len(fs)) - 1) / 100 + lower, upper := int(pos), int(math.Ceil(pos)) + if lower == upper { + return fs[lower] + } + return fs[lower] + (fs[upper]-fs[lower])*(pos-float64(lower)) +} + func validateDatabaseTypeWithValid(v reflect.Value) (handled bool, err error) { if v.Kind() == reflect.Struct { return false, nil @@ -2418,17 +2432,59 @@ func (q *FakeQuerier) GetDeploymentWorkspaceAgentStats(_ context.Context, create latencies = append(latencies, agentStat.ConnectionMedianLatencyMS) } - tryPercentile := func(fs []float64, p float64) float64 { - if len(fs) == 0 { - return -1 + stat.WorkspaceConnectionLatency50 = tryPercentile(latencies, 50) + stat.WorkspaceConnectionLatency95 = tryPercentile(latencies, 95) + + return stat, nil +} + +func (q *FakeQuerier) GetDeploymentWorkspaceAgentUsageStats(_ context.Context, createdAt time.Time) (database.GetDeploymentWorkspaceAgentUsageStatsRow, error) { + q.mutex.RLock() + defer q.mutex.RUnlock() + + stat := database.GetDeploymentWorkspaceAgentUsageStatsRow{} + sessions := make(map[uuid.UUID]database.WorkspaceAgentStat) + agentStatsCreatedAfter := make([]database.WorkspaceAgentStat, 0) + for _, agentStat := range q.workspaceAgentStats { + // WHERE workspace_agent_stats.created_at > $1 + if agentStat.CreatedAt.After(createdAt) { + agentStatsCreatedAfter = append(agentStatsCreatedAfter, agentStat) + } + // WHERE + // created_at > $1 + // AND created_at < date_trunc('minute', now()) -- Exclude current partial minute + // AND usage = true + if agentStat.Usage && + (agentStat.CreatedAt.After(createdAt) || agentStat.CreatedAt.Equal(createdAt)) && + agentStat.CreatedAt.Before(time.Now().Truncate(time.Minute)) { + val, ok := sessions[agentStat.AgentID] + if !ok { + sessions[agentStat.AgentID] = agentStat + } else if val.CreatedAt.Before(agentStat.CreatedAt) { + sessions[agentStat.AgentID] = agentStat + } } - sort.Float64s(fs) - return fs[int(float64(len(fs))*p/100)] } + latencies := make([]float64, 0) + for _, agentStat := range agentStatsCreatedAfter { + if agentStat.ConnectionMedianLatencyMS <= 0 { + continue + } + stat.WorkspaceRxBytes += agentStat.RxBytes + stat.WorkspaceTxBytes += agentStat.TxBytes + latencies = append(latencies, agentStat.ConnectionMedianLatencyMS) + } stat.WorkspaceConnectionLatency50 = tryPercentile(latencies, 50) stat.WorkspaceConnectionLatency95 = tryPercentile(latencies, 95) + for _, agentStat := range sessions { + stat.SessionCountVSCode += agentStat.SessionCountVSCode + stat.SessionCountJetBrains += agentStat.SessionCountJetBrains + stat.SessionCountReconnectingPTY += agentStat.SessionCountReconnectingPTY + stat.SessionCountSSH += agentStat.SessionCountSSH + } + return stat, nil } @@ -3979,14 +4035,6 @@ func (q *FakeQuerier) GetTemplateAverageBuildTime(ctx context.Context, arg datab } } - tryPercentile := func(fs []float64, p float64) float64 { - if len(fs) == 0 { - return -1 - } - sort.Float64s(fs) - return fs[int(float64(len(fs))*p/100)] - } - var row database.GetTemplateAverageBuildTimeRow row.Delete50, row.Delete95 = tryPercentile(deleteTimes, 50), tryPercentile(deleteTimes, 95) row.Stop50, row.Stop95 = tryPercentile(stopTimes, 50), tryPercentile(stopTimes, 95) @@ -5014,14 +5062,6 @@ func (q *FakeQuerier) GetUserLatencyInsights(_ context.Context, arg database.Get seenTemplatesByUserID[stat.UserID] = uniqueSortedUUIDs(append(seenTemplatesByUserID[stat.UserID], stat.TemplateID)) } - tryPercentile := func(fs []float64, p float64) float64 { - if len(fs) == 0 { - return -1 - } - sort.Float64s(fs) - return fs[int(float64(len(fs))*p/100)] - } - var rows []database.GetUserLatencyInsightsRow for userID, latencies := range latenciesByUserID { user, err := q.getUserByIDNoLock(userID) @@ -5535,14 +5575,6 @@ func (q *FakeQuerier) GetWorkspaceAgentStats(_ context.Context, createdAfter tim latenciesByAgent[agentStat.AgentID] = append(latenciesByAgent[agentStat.AgentID], 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)] - } - for _, stat := range statByAgent { stat.AggregatedFrom = minimumDateByAgent[stat.AgentID] statByAgent[stat.AgentID] = stat @@ -5634,6 +5666,224 @@ func (q *FakeQuerier) GetWorkspaceAgentStatsAndLabels(ctx context.Context, creat return stats, nil } +func (q *FakeQuerier) GetWorkspaceAgentUsageStats(_ context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentUsageStatsRow, error) { + q.mutex.RLock() + defer q.mutex.RUnlock() + + type agentStatsKey struct { + UserID uuid.UUID + AgentID uuid.UUID + WorkspaceID uuid.UUID + TemplateID uuid.UUID + } + + type minuteStatsKey struct { + agentStatsKey + MinuteBucket time.Time + } + + latestAgentStats := map[agentStatsKey]database.GetWorkspaceAgentUsageStatsRow{} + latestAgentLatencies := map[agentStatsKey][]float64{} + for _, agentStat := range q.workspaceAgentStats { + key := agentStatsKey{ + UserID: agentStat.UserID, + AgentID: agentStat.AgentID, + WorkspaceID: agentStat.WorkspaceID, + TemplateID: agentStat.TemplateID, + } + if agentStat.CreatedAt.After(createdAt) { + val, ok := latestAgentStats[key] + if ok { + val.WorkspaceRxBytes += agentStat.RxBytes + val.WorkspaceTxBytes += agentStat.TxBytes + latestAgentStats[key] = val + } else { + latestAgentStats[key] = database.GetWorkspaceAgentUsageStatsRow{ + UserID: agentStat.UserID, + AgentID: agentStat.AgentID, + WorkspaceID: agentStat.WorkspaceID, + TemplateID: agentStat.TemplateID, + AggregatedFrom: createdAt, + WorkspaceRxBytes: agentStat.RxBytes, + WorkspaceTxBytes: agentStat.TxBytes, + } + } + + latencies, ok := latestAgentLatencies[key] + if !ok { + latestAgentLatencies[key] = []float64{agentStat.ConnectionMedianLatencyMS} + } else { + latestAgentLatencies[key] = append(latencies, agentStat.ConnectionMedianLatencyMS) + } + } + } + + for key, latencies := range latestAgentLatencies { + val, ok := latestAgentStats[key] + if ok { + val.WorkspaceConnectionLatency50 = tryPercentile(latencies, 50) + val.WorkspaceConnectionLatency95 = tryPercentile(latencies, 95) + } + latestAgentStats[key] = val + } + + type bucketRow struct { + database.GetWorkspaceAgentUsageStatsRow + MinuteBucket time.Time + } + + minuteBuckets := make(map[minuteStatsKey]bucketRow) + for _, agentStat := range q.workspaceAgentStats { + if agentStat.Usage && + (agentStat.CreatedAt.After(createdAt) || agentStat.CreatedAt.Equal(createdAt)) && + agentStat.CreatedAt.Before(time.Now().Truncate(time.Minute)) { + key := minuteStatsKey{ + agentStatsKey: agentStatsKey{ + UserID: agentStat.UserID, + AgentID: agentStat.AgentID, + WorkspaceID: agentStat.WorkspaceID, + TemplateID: agentStat.TemplateID, + }, + MinuteBucket: agentStat.CreatedAt.Truncate(time.Minute), + } + val, ok := minuteBuckets[key] + if ok { + val.SessionCountVSCode += agentStat.SessionCountVSCode + val.SessionCountJetBrains += agentStat.SessionCountJetBrains + val.SessionCountReconnectingPTY += agentStat.SessionCountReconnectingPTY + val.SessionCountSSH += agentStat.SessionCountSSH + } else { + minuteBuckets[key] = bucketRow{ + GetWorkspaceAgentUsageStatsRow: database.GetWorkspaceAgentUsageStatsRow{ + UserID: agentStat.UserID, + AgentID: agentStat.AgentID, + WorkspaceID: agentStat.WorkspaceID, + TemplateID: agentStat.TemplateID, + SessionCountVSCode: agentStat.SessionCountVSCode, + SessionCountSSH: agentStat.SessionCountSSH, + SessionCountJetBrains: agentStat.SessionCountJetBrains, + SessionCountReconnectingPTY: agentStat.SessionCountReconnectingPTY, + }, + MinuteBucket: agentStat.CreatedAt.Truncate(time.Minute), + } + } + } + } + + // Get the latest minute bucket for each agent. + latestBuckets := make(map[uuid.UUID]bucketRow) + for key, bucket := range minuteBuckets { + latest, ok := latestBuckets[key.AgentID] + if !ok || key.MinuteBucket.After(latest.MinuteBucket) { + latestBuckets[key.AgentID] = bucket + } + } + + for key, stat := range latestAgentStats { + bucket, ok := latestBuckets[stat.AgentID] + if ok { + stat.SessionCountVSCode = bucket.SessionCountVSCode + stat.SessionCountJetBrains = bucket.SessionCountJetBrains + stat.SessionCountReconnectingPTY = bucket.SessionCountReconnectingPTY + stat.SessionCountSSH = bucket.SessionCountSSH + } + latestAgentStats[key] = stat + } + return maps.Values(latestAgentStats), nil +} + +func (q *FakeQuerier) GetWorkspaceAgentUsageStatsAndLabels(_ context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentUsageStatsAndLabelsRow, error) { + q.mutex.RLock() + defer q.mutex.RUnlock() + + type statsKey struct { + AgentID uuid.UUID + UserID uuid.UUID + WorkspaceID uuid.UUID + } + + latestAgentStats := map[statsKey]database.WorkspaceAgentStat{} + maxConnMedianLatency := 0.0 + for _, agentStat := range q.workspaceAgentStats { + key := statsKey{ + AgentID: agentStat.AgentID, + UserID: agentStat.UserID, + WorkspaceID: agentStat.WorkspaceID, + } + // WHERE workspace_agent_stats.created_at > $1 + // GROUP BY user_id, agent_id, workspace_id + if agentStat.CreatedAt.After(createdAt) { + val, ok := latestAgentStats[key] + if !ok { + latestAgentStats[key] = agentStat + } else { + val.RxBytes += agentStat.RxBytes + val.TxBytes += agentStat.TxBytes + latestAgentStats[key] = val + } + } + // WHERE usage = true + // GROUP BY user_id, agent_id, workspace_id + if agentStat.Usage { + val, ok := latestAgentStats[key] + if !ok { + latestAgentStats[key] = agentStat + } else { + val.SessionCountVSCode += agentStat.SessionCountVSCode + val.SessionCountJetBrains += agentStat.SessionCountJetBrains + val.SessionCountReconnectingPTY += agentStat.SessionCountReconnectingPTY + val.SessionCountSSH += agentStat.SessionCountSSH + val.ConnectionCount += agentStat.ConnectionCount + latestAgentStats[key] = val + } + } + // SELECT + // agent_id, + // coalesce(MAX(connection_median_latency_ms), 0)::float AS connection_median_latency_ms + // FROM workspace_agent_stats + // GROUP BY user_id, agent_id, workspace_id + if agentStat.ConnectionMedianLatencyMS > maxConnMedianLatency { + val, ok := latestAgentStats[key] + if !ok { + latestAgentStats[key] = agentStat + } else { + val.ConnectionMedianLatencyMS = agentStat.ConnectionMedianLatencyMS + latestAgentStats[key] = val + } + } + } + + stats := make([]database.GetWorkspaceAgentUsageStatsAndLabelsRow, 0, len(latestAgentStats)) + for key, agentStat := range latestAgentStats { + user, err := q.getUserByIDNoLock(key.UserID) + if err != nil { + return nil, err + } + workspace, err := q.getWorkspaceByIDNoLock(context.Background(), key.WorkspaceID) + if err != nil { + return nil, err + } + agent, err := q.getWorkspaceAgentByIDNoLock(context.Background(), key.AgentID) + if err != nil { + return nil, err + } + stats = append(stats, database.GetWorkspaceAgentUsageStatsAndLabelsRow{ + Username: user.Username, + AgentName: agent.Name, + WorkspaceName: workspace.Name, + RxBytes: agentStat.RxBytes, + TxBytes: agentStat.TxBytes, + SessionCountVSCode: agentStat.SessionCountVSCode, + SessionCountSSH: agentStat.SessionCountSSH, + SessionCountJetBrains: agentStat.SessionCountJetBrains, + SessionCountReconnectingPTY: agentStat.SessionCountReconnectingPTY, + ConnectionCount: agentStat.ConnectionCount, + ConnectionMedianLatencyMS: agentStat.ConnectionMedianLatencyMS, + }) + } + return stats, nil +} + func (q *FakeQuerier) GetWorkspaceAgentsByResourceIDs(ctx context.Context, resourceIDs []uuid.UUID) ([]database.WorkspaceAgent, error) { q.mutex.RLock() defer q.mutex.RUnlock() @@ -7255,6 +7505,7 @@ func (q *FakeQuerier) InsertWorkspaceAgentStats(_ context.Context, arg database. SessionCountReconnectingPTY: arg.SessionCountReconnectingPTY[i], SessionCountSSH: arg.SessionCountSSH[i], ConnectionMedianLatencyMS: arg.ConnectionMedianLatencyMS[i], + Usage: arg.Usage[i], } q.workspaceAgentStats = append(q.workspaceAgentStats, stat) } diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go index 38289c143bfd9..486326025dc92 100644 --- a/coderd/database/dbmetrics/dbmetrics.go +++ b/coderd/database/dbmetrics/dbmetrics.go @@ -585,6 +585,13 @@ func (m metricsStore) GetDeploymentWorkspaceAgentStats(ctx context.Context, crea return row, err } +func (m metricsStore) GetDeploymentWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) (database.GetDeploymentWorkspaceAgentUsageStatsRow, error) { + start := time.Now() + r0, r1 := m.s.GetDeploymentWorkspaceAgentUsageStats(ctx, createdAt) + m.queryLatencies.WithLabelValues("GetDeploymentWorkspaceAgentUsageStats").Observe(time.Since(start).Seconds()) + return r0, r1 +} + func (m metricsStore) GetDeploymentWorkspaceStats(ctx context.Context) (database.GetDeploymentWorkspaceStatsRow, error) { start := time.Now() row, err := m.s.GetDeploymentWorkspaceStats(ctx) @@ -1348,6 +1355,20 @@ func (m metricsStore) GetWorkspaceAgentStatsAndLabels(ctx context.Context, creat return stats, err } +func (m metricsStore) GetWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentUsageStatsRow, error) { + start := time.Now() + r0, r1 := m.s.GetWorkspaceAgentUsageStats(ctx, createdAt) + m.queryLatencies.WithLabelValues("GetWorkspaceAgentUsageStats").Observe(time.Since(start).Seconds()) + return r0, r1 +} + +func (m metricsStore) GetWorkspaceAgentUsageStatsAndLabels(ctx context.Context, createdAt time.Time) ([]database.GetWorkspaceAgentUsageStatsAndLabelsRow, error) { + start := time.Now() + r0, r1 := m.s.GetWorkspaceAgentUsageStatsAndLabels(ctx, createdAt) + m.queryLatencies.WithLabelValues("GetWorkspaceAgentUsageStatsAndLabels").Observe(time.Since(start).Seconds()) + return r0, r1 +} + func (m metricsStore) GetWorkspaceAgentsByResourceIDs(ctx context.Context, ids []uuid.UUID) ([]database.WorkspaceAgent, error) { start := time.Now() agents, err := m.s.GetWorkspaceAgentsByResourceIDs(ctx, ids) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 1771807f26b2f..f497a1101698d 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -1149,6 +1149,21 @@ func (mr *MockStoreMockRecorder) GetDeploymentWorkspaceAgentStats(arg0, arg1 any return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDeploymentWorkspaceAgentStats", reflect.TypeOf((*MockStore)(nil).GetDeploymentWorkspaceAgentStats), arg0, arg1) } +// GetDeploymentWorkspaceAgentUsageStats mocks base method. +func (m *MockStore) GetDeploymentWorkspaceAgentUsageStats(arg0 context.Context, arg1 time.Time) (database.GetDeploymentWorkspaceAgentUsageStatsRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetDeploymentWorkspaceAgentUsageStats", arg0, arg1) + ret0, _ := ret[0].(database.GetDeploymentWorkspaceAgentUsageStatsRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetDeploymentWorkspaceAgentUsageStats indicates an expected call of GetDeploymentWorkspaceAgentUsageStats. +func (mr *MockStoreMockRecorder) GetDeploymentWorkspaceAgentUsageStats(arg0, arg1 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDeploymentWorkspaceAgentUsageStats", reflect.TypeOf((*MockStore)(nil).GetDeploymentWorkspaceAgentUsageStats), arg0, arg1) +} + // GetDeploymentWorkspaceStats mocks base method. func (m *MockStore) GetDeploymentWorkspaceStats(arg0 context.Context) (database.GetDeploymentWorkspaceStatsRow, error) { m.ctrl.T.Helper() @@ -2814,6 +2829,36 @@ func (mr *MockStoreMockRecorder) GetWorkspaceAgentStatsAndLabels(arg0, arg1 any) return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentStatsAndLabels", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentStatsAndLabels), arg0, arg1) } +// GetWorkspaceAgentUsageStats mocks base method. +func (m *MockStore) GetWorkspaceAgentUsageStats(arg0 context.Context, arg1 time.Time) ([]database.GetWorkspaceAgentUsageStatsRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetWorkspaceAgentUsageStats", arg0, arg1) + ret0, _ := ret[0].([]database.GetWorkspaceAgentUsageStatsRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetWorkspaceAgentUsageStats indicates an expected call of GetWorkspaceAgentUsageStats. +func (mr *MockStoreMockRecorder) GetWorkspaceAgentUsageStats(arg0, arg1 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentUsageStats", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentUsageStats), arg0, arg1) +} + +// GetWorkspaceAgentUsageStatsAndLabels mocks base method. +func (m *MockStore) GetWorkspaceAgentUsageStatsAndLabels(arg0 context.Context, arg1 time.Time) ([]database.GetWorkspaceAgentUsageStatsAndLabelsRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetWorkspaceAgentUsageStatsAndLabels", arg0, arg1) + ret0, _ := ret[0].([]database.GetWorkspaceAgentUsageStatsAndLabelsRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetWorkspaceAgentUsageStatsAndLabels indicates an expected call of GetWorkspaceAgentUsageStatsAndLabels. +func (mr *MockStoreMockRecorder) GetWorkspaceAgentUsageStatsAndLabels(arg0, arg1 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentUsageStatsAndLabels", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentUsageStatsAndLabels), arg0, arg1) +} + // GetWorkspaceAgentsByResourceIDs mocks base method. func (m *MockStore) GetWorkspaceAgentsByResourceIDs(arg0 context.Context, arg1 []uuid.UUID) ([]database.WorkspaceAgent, error) { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 989586dddf4ef..22cfade582357 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1352,7 +1352,8 @@ CREATE TABLE workspace_agent_stats ( session_count_vscode bigint DEFAULT 0 NOT NULL, session_count_jetbrains bigint DEFAULT 0 NOT NULL, session_count_reconnecting_pty bigint DEFAULT 0 NOT NULL, - session_count_ssh bigint DEFAULT 0 NOT NULL + session_count_ssh bigint DEFAULT 0 NOT NULL, + usage boolean DEFAULT false NOT NULL ); CREATE TABLE workspace_agents ( diff --git a/coderd/database/migrations/000249_agent_stats_usage.down.sql b/coderd/database/migrations/000249_agent_stats_usage.down.sql new file mode 100644 index 0000000000000..8cfc278493886 --- /dev/null +++ b/coderd/database/migrations/000249_agent_stats_usage.down.sql @@ -0,0 +1 @@ +ALTER TABLE workspace_agent_stats DROP COLUMN usage; diff --git a/coderd/database/migrations/000249_agent_stats_usage.up.sql b/coderd/database/migrations/000249_agent_stats_usage.up.sql new file mode 100644 index 0000000000000..92c839eb943a3 --- /dev/null +++ b/coderd/database/migrations/000249_agent_stats_usage.up.sql @@ -0,0 +1 @@ +ALTER TABLE workspace_agent_stats ADD COLUMN usage boolean NOT NULL DEFAULT false; diff --git a/coderd/database/models.go b/coderd/database/models.go index 959609d82eb79..2f6eeb6c2f04c 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -2823,6 +2823,7 @@ type WorkspaceAgentStat struct { 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"` + Usage bool `db:"usage" json:"usage"` } type WorkspaceApp struct { diff --git a/coderd/database/querier.go b/coderd/database/querier.go index c614a03834a9b..c0341e612b5f5 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -137,6 +137,7 @@ type sqlcQuerier interface { GetDeploymentDAUs(ctx context.Context, tzOffset int32) ([]GetDeploymentDAUsRow, error) GetDeploymentID(ctx context.Context) (string, error) GetDeploymentWorkspaceAgentStats(ctx context.Context, createdAt time.Time) (GetDeploymentWorkspaceAgentStatsRow, error) + GetDeploymentWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) (GetDeploymentWorkspaceAgentUsageStatsRow, error) GetDeploymentWorkspaceStats(ctx context.Context) (GetDeploymentWorkspaceStatsRow, error) GetExternalAuthLink(ctx context.Context, arg GetExternalAuthLinkParams) (ExternalAuthLink, error) GetExternalAuthLinksByUserID(ctx context.Context, userID uuid.UUID) ([]ExternalAuthLink, error) @@ -289,6 +290,8 @@ type sqlcQuerier interface { GetWorkspaceAgentScriptsByAgentIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceAgentScript, error) GetWorkspaceAgentStats(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentStatsRow, error) GetWorkspaceAgentStatsAndLabels(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentStatsAndLabelsRow, error) + GetWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentUsageStatsRow, error) + GetWorkspaceAgentUsageStatsAndLabels(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentUsageStatsAndLabelsRow, error) GetWorkspaceAgentsByResourceIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceAgent, error) GetWorkspaceAgentsCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceAgent, error) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx context.Context, workspaceID uuid.UUID) ([]WorkspaceAgent, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 7b7fd8b0a2823..bfa573853143a 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -100,6 +100,293 @@ func TestGetDeploymentWorkspaceAgentStats(t *testing.T) { }) } +func TestGetDeploymentWorkspaceAgentUsageStats(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + authz := rbac.NewAuthorizer(prometheus.NewRegistry()) + db = dbauthz.New(db, authz, slogtest.Make(t, &slogtest.Options{}), coderdtest.AccessControlStorePointer()) + ctx := context.Background() + agentID := uuid.New() + // Since the queries exclude the current minute + insertTime := dbtime.Now().Add(-time.Minute) + + // Old stats + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime.Add(-time.Minute), + AgentID: agentID, + TxBytes: 1, + RxBytes: 1, + ConnectionMedianLatencyMS: 1, + }) + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime.Add(-time.Minute), + AgentID: agentID, + SessionCountVSCode: 1, + Usage: true, + }) + + // Latest stats + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime, + AgentID: agentID, + TxBytes: 1, + RxBytes: 1, + ConnectionMedianLatencyMS: 2, + }) + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime, + AgentID: agentID, + SessionCountVSCode: 2, + Usage: true, + }) + + stats, err := db.GetDeploymentWorkspaceAgentUsageStats(ctx, dbtime.Now().Add(-time.Hour)) + require.NoError(t, err) + + require.Equal(t, int64(2), stats.WorkspaceTxBytes) + require.Equal(t, int64(2), stats.WorkspaceRxBytes) + require.Equal(t, 1.5, stats.WorkspaceConnectionLatency50) + require.Equal(t, 1.95, stats.WorkspaceConnectionLatency95) + require.Equal(t, int64(2), stats.SessionCountVSCode) +} + +func TestGetWorkspaceAgentUsageStats(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + authz := rbac.NewAuthorizer(prometheus.NewRegistry()) + db = dbauthz.New(db, authz, slogtest.Make(t, &slogtest.Options{}), coderdtest.AccessControlStorePointer()) + ctx := context.Background() + // Since the queries exclude the current minute + insertTime := dbtime.Now().Add(-time.Minute) + + agentID1 := uuid.New() + agentID2 := uuid.New() + workspaceID1 := uuid.New() + workspaceID2 := uuid.New() + templateID1 := uuid.New() + templateID2 := uuid.New() + userID1 := uuid.New() + userID2 := uuid.New() + + // Old workspace 1 stats + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime.Add(-time.Minute), + AgentID: agentID1, + WorkspaceID: workspaceID1, + TemplateID: templateID1, + UserID: userID1, + TxBytes: 1, + RxBytes: 1, + ConnectionMedianLatencyMS: 1, + }) + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime.Add(-time.Minute), + AgentID: agentID1, + WorkspaceID: workspaceID1, + TemplateID: templateID1, + UserID: userID1, + SessionCountVSCode: 1, + Usage: true, + }) + + // Latest workspace 1 stats + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime, + AgentID: agentID1, + WorkspaceID: workspaceID1, + TemplateID: templateID1, + UserID: userID1, + TxBytes: 2, + RxBytes: 2, + ConnectionMedianLatencyMS: 1, + }) + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime, + AgentID: agentID1, + WorkspaceID: workspaceID1, + TemplateID: templateID1, + UserID: userID1, + SessionCountVSCode: 2, + Usage: true, + }) + + // Latest workspace 2 stats + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime, + AgentID: agentID2, + WorkspaceID: workspaceID2, + TemplateID: templateID2, + UserID: userID2, + TxBytes: 4, + RxBytes: 8, + ConnectionMedianLatencyMS: 1, + }) + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime, + AgentID: agentID2, + WorkspaceID: workspaceID2, + TemplateID: templateID2, + UserID: userID2, + SessionCountVSCode: 5, + Usage: true, + }) + + reqTime := dbtime.Now().Add(-time.Hour) + stats, err := db.GetWorkspaceAgentUsageStats(ctx, reqTime) + require.NoError(t, err) + + ws1Stats, ws2Stats := stats[0], stats[1] + if ws1Stats.WorkspaceID != workspaceID1 { + ws1Stats, ws2Stats = ws2Stats, ws1Stats + } + require.Equal(t, ws1Stats.WorkspaceTxBytes, int64(3)) + require.Equal(t, ws1Stats.WorkspaceRxBytes, int64(3)) + require.Equal(t, ws1Stats.SessionCountVSCode, int64(2)) + + require.Equal(t, ws2Stats.WorkspaceTxBytes, int64(4)) + require.Equal(t, ws2Stats.WorkspaceRxBytes, int64(8)) + require.Equal(t, ws2Stats.SessionCountVSCode, int64(5)) +} + +func TestGetWorkspaceAgentUsageStatsAndLabels(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := context.Background() + // Since the queries exclude the current minute + insertTime := dbtime.Now().Add(-time.Minute) + + // Insert user, agent, template, workspace + user1 := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + job1 := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + OrganizationID: org.ID, + }) + resource1 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: job1.ID, + }) + agent1 := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: resource1.ID, + }) + template1 := dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + CreatedBy: user1.ID, + }) + workspace1 := dbgen.Workspace(t, db, database.Workspace{ + OwnerID: user1.ID, + OrganizationID: org.ID, + TemplateID: template1.ID, + }) + user2 := dbgen.User(t, db, database.User{}) + job2 := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + OrganizationID: org.ID, + }) + resource2 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + JobID: job2.ID, + }) + agent2 := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: resource2.ID, + }) + template2 := dbgen.Template(t, db, database.Template{ + CreatedBy: user1.ID, + OrganizationID: org.ID, + }) + workspace2 := dbgen.Workspace(t, db, database.Workspace{ + OwnerID: user2.ID, + OrganizationID: org.ID, + TemplateID: template2.ID, + }) + + // Old workspace 1 stats + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime.Add(-time.Minute), + AgentID: agent1.ID, + WorkspaceID: workspace1.ID, + TemplateID: template1.ID, + UserID: user1.ID, + TxBytes: 1, + RxBytes: 1, + ConnectionMedianLatencyMS: 1, + }) + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime.Add(-time.Minute), + AgentID: agent1.ID, + WorkspaceID: workspace1.ID, + TemplateID: template1.ID, + UserID: user1.ID, + SessionCountVSCode: 1, + Usage: true, + }) + + // Latest workspace 1 stats + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime, + AgentID: agent1.ID, + WorkspaceID: workspace1.ID, + TemplateID: template1.ID, + UserID: user1.ID, + TxBytes: 2, + RxBytes: 2, + ConnectionMedianLatencyMS: 1, + }) + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime, + AgentID: agent1.ID, + WorkspaceID: workspace1.ID, + TemplateID: template1.ID, + UserID: user1.ID, + SessionCountVSCode: 2, + Usage: true, + }) + + // Latest workspace 2 stats + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime, + AgentID: agent2.ID, + WorkspaceID: workspace2.ID, + TemplateID: template2.ID, + UserID: user2.ID, + TxBytes: 4, + RxBytes: 8, + ConnectionMedianLatencyMS: 1, + }) + dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{ + CreatedAt: insertTime, + AgentID: agent2.ID, + WorkspaceID: workspace2.ID, + TemplateID: template2.ID, + UserID: user2.ID, + SessionCountVSCode: 5, + Usage: true, + }) + + stats, err := db.GetWorkspaceAgentUsageStatsAndLabels(ctx, insertTime.Add(-time.Hour)) + require.NoError(t, err) + + require.Len(t, stats, 2) + require.Contains(t, stats, database.GetWorkspaceAgentUsageStatsAndLabelsRow{ + Username: user1.Username, + AgentName: agent1.Name, + WorkspaceName: workspace1.Name, + RxBytes: 3, + TxBytes: 3, + SessionCountVSCode: 3, + ConnectionMedianLatencyMS: 1, + }) + + require.Contains(t, stats, database.GetWorkspaceAgentUsageStatsAndLabelsRow{ + Username: user2.Username, + AgentName: agent2.Name, + WorkspaceName: workspace2.Name, + RxBytes: 8, + TxBytes: 4, + SessionCountVSCode: 5, + ConnectionMedianLatencyMS: 1, + }) +} + func TestInsertWorkspaceAgentLogs(t *testing.T) { t.Parallel() if testing.Short() { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index fc388e55247d0..f13ced09c142c 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -11602,7 +11602,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 FROM ( - SELECT id, created_at, user_id, agent_id, workspace_id, template_id, connections_by_proto, connection_count, rx_packets, rx_bytes, tx_packets, tx_bytes, connection_median_latency_ms, session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn + SELECT id, created_at, user_id, agent_id, workspace_id, template_id, connections_by_proto, connection_count, rx_packets, rx_bytes, tx_packets, tx_bytes, connection_median_latency_ms, session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, usage, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn FROM workspace_agent_stats WHERE created_at > $1 ) AS a WHERE a.rn = 1 ) @@ -11636,6 +11636,88 @@ func (q *sqlQuerier) GetDeploymentWorkspaceAgentStats(ctx context.Context, creat return i, err } +const getDeploymentWorkspaceAgentUsageStats = `-- name: GetDeploymentWorkspaceAgentUsageStats :one +WITH agent_stats AS ( + SELECT + coalesce(SUM(rx_bytes), 0)::bigint AS workspace_rx_bytes, + coalesce(SUM(tx_bytes), 0)::bigint AS workspace_tx_bytes, + coalesce((PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_50, + coalesce((PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_95 + FROM workspace_agent_stats + -- The greater than 0 is to support legacy agents that don't report connection_median_latency_ms. + WHERE workspace_agent_stats.created_at > $1 AND connection_median_latency_ms > 0 +), +minute_buckets AS ( + SELECT + 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 + 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 +), +latest_buckets AS ( + SELECT DISTINCT ON (agent_id) + agent_id, + minute_bucket, + session_count_vscode, + session_count_jetbrains, + session_count_reconnecting_pty, + session_count_ssh + FROM + minute_buckets + ORDER BY + agent_id, + minute_bucket DESC +), +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 + FROM + latest_buckets +) +SELECT workspace_rx_bytes, workspace_tx_bytes, workspace_connection_latency_50, workspace_connection_latency_95, session_count_vscode, session_count_ssh, session_count_jetbrains, session_count_reconnecting_pty FROM agent_stats, latest_agent_stats +` + +type GetDeploymentWorkspaceAgentUsageStatsRow struct { + WorkspaceRxBytes int64 `db:"workspace_rx_bytes" json:"workspace_rx_bytes"` + WorkspaceTxBytes int64 `db:"workspace_tx_bytes" json:"workspace_tx_bytes"` + WorkspaceConnectionLatency50 float64 `db:"workspace_connection_latency_50" json:"workspace_connection_latency_50"` + WorkspaceConnectionLatency95 float64 `db:"workspace_connection_latency_95" json:"workspace_connection_latency_95"` + SessionCountVSCode int64 `db:"session_count_vscode" json:"session_count_vscode"` + SessionCountSSH int64 `db:"session_count_ssh" json:"session_count_ssh"` + SessionCountJetBrains int64 `db:"session_count_jetbrains" json:"session_count_jetbrains"` + SessionCountReconnectingPTY int64 `db:"session_count_reconnecting_pty" json:"session_count_reconnecting_pty"` +} + +func (q *sqlQuerier) GetDeploymentWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) (GetDeploymentWorkspaceAgentUsageStatsRow, error) { + row := q.db.QueryRowContext(ctx, getDeploymentWorkspaceAgentUsageStats, createdAt) + var i GetDeploymentWorkspaceAgentUsageStatsRow + err := row.Scan( + &i.WorkspaceRxBytes, + &i.WorkspaceTxBytes, + &i.WorkspaceConnectionLatency50, + &i.WorkspaceConnectionLatency95, + &i.SessionCountVSCode, + &i.SessionCountSSH, + &i.SessionCountJetBrains, + &i.SessionCountReconnectingPTY, + ) + return i, err +} + const getTemplateDAUs = `-- name: GetTemplateDAUs :many SELECT (created_at at TIME ZONE cast($2::integer as text))::date as date, @@ -11707,7 +11789,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 FROM ( - SELECT id, created_at, user_id, agent_id, workspace_id, template_id, connections_by_proto, connection_count, rx_packets, rx_bytes, tx_packets, tx_bytes, connection_median_latency_ms, session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn + SELECT id, created_at, user_id, agent_id, workspace_id, template_id, connections_by_proto, connection_count, rx_packets, rx_bytes, tx_packets, tx_bytes, connection_median_latency_ms, session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, usage, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn FROM workspace_agent_stats WHERE created_at > $1 ) AS a WHERE a.rn = 1 GROUP BY a.user_id, a.agent_id, a.workspace_id, a.template_id ) @@ -11790,7 +11872,7 @@ WITH agent_stats AS ( coalesce(SUM(connection_count), 0)::bigint AS connection_count, coalesce(MAX(connection_median_latency_ms), 0)::float AS connection_median_latency_ms FROM ( - SELECT id, created_at, user_id, agent_id, workspace_id, template_id, connections_by_proto, connection_count, rx_packets, rx_bytes, tx_packets, tx_bytes, connection_median_latency_ms, session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn + SELECT id, created_at, user_id, agent_id, workspace_id, template_id, connections_by_proto, connection_count, rx_packets, rx_bytes, tx_packets, tx_bytes, connection_median_latency_ms, session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, usage, 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 @@ -11871,6 +11953,222 @@ func (q *sqlQuerier) GetWorkspaceAgentStatsAndLabels(ctx context.Context, create return items, nil } +const getWorkspaceAgentUsageStats = `-- name: GetWorkspaceAgentUsageStats :many +WITH agent_stats AS ( + SELECT + user_id, + agent_id, + workspace_id, + template_id, + MIN(created_at)::timestamptz AS aggregated_from, + coalesce(SUM(rx_bytes), 0)::bigint AS workspace_rx_bytes, + coalesce(SUM(tx_bytes), 0)::bigint AS workspace_tx_bytes, + coalesce((PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_50, + coalesce((PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_95 + FROM workspace_agent_stats + -- The greater than 0 is to support legacy agents that don't report connection_median_latency_ms. + WHERE workspace_agent_stats.created_at > $1 AND connection_median_latency_ms > 0 + GROUP BY user_id, agent_id, workspace_id, template_id +), +minute_buckets AS ( + SELECT + 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 + 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, + user_id, + agent_id, + workspace_id, + template_id +), +latest_buckets AS ( + SELECT DISTINCT ON (agent_id) + agent_id, + session_count_vscode, + session_count_ssh, + session_count_jetbrains, + session_count_reconnecting_pty + FROM + minute_buckets + ORDER BY + agent_id, + minute_bucket DESC +) +SELECT user_id, agent_stats.agent_id, workspace_id, template_id, aggregated_from, workspace_rx_bytes, workspace_tx_bytes, workspace_connection_latency_50, workspace_connection_latency_95, latest_buckets.agent_id, session_count_vscode, session_count_ssh, session_count_jetbrains, session_count_reconnecting_pty FROM agent_stats JOIN latest_buckets ON agent_stats.agent_id = latest_buckets.agent_id +` + +type GetWorkspaceAgentUsageStatsRow struct { + UserID uuid.UUID `db:"user_id" json:"user_id"` + AgentID uuid.UUID `db:"agent_id" json:"agent_id"` + WorkspaceID uuid.UUID `db:"workspace_id" json:"workspace_id"` + TemplateID uuid.UUID `db:"template_id" json:"template_id"` + AggregatedFrom time.Time `db:"aggregated_from" json:"aggregated_from"` + WorkspaceRxBytes int64 `db:"workspace_rx_bytes" json:"workspace_rx_bytes"` + WorkspaceTxBytes int64 `db:"workspace_tx_bytes" json:"workspace_tx_bytes"` + WorkspaceConnectionLatency50 float64 `db:"workspace_connection_latency_50" json:"workspace_connection_latency_50"` + WorkspaceConnectionLatency95 float64 `db:"workspace_connection_latency_95" json:"workspace_connection_latency_95"` + AgentID_2 uuid.UUID `db:"agent_id_2" json:"agent_id_2"` + SessionCountVSCode int64 `db:"session_count_vscode" json:"session_count_vscode"` + SessionCountSSH int64 `db:"session_count_ssh" json:"session_count_ssh"` + SessionCountJetBrains int64 `db:"session_count_jetbrains" json:"session_count_jetbrains"` + SessionCountReconnectingPTY int64 `db:"session_count_reconnecting_pty" json:"session_count_reconnecting_pty"` +} + +func (q *sqlQuerier) GetWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentUsageStatsRow, error) { + rows, err := q.db.QueryContext(ctx, getWorkspaceAgentUsageStats, createdAt) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetWorkspaceAgentUsageStatsRow + for rows.Next() { + var i GetWorkspaceAgentUsageStatsRow + if err := rows.Scan( + &i.UserID, + &i.AgentID, + &i.WorkspaceID, + &i.TemplateID, + &i.AggregatedFrom, + &i.WorkspaceRxBytes, + &i.WorkspaceTxBytes, + &i.WorkspaceConnectionLatency50, + &i.WorkspaceConnectionLatency95, + &i.AgentID_2, + &i.SessionCountVSCode, + &i.SessionCountSSH, + &i.SessionCountJetBrains, + &i.SessionCountReconnectingPTY, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getWorkspaceAgentUsageStatsAndLabels = `-- name: GetWorkspaceAgentUsageStatsAndLabels :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 + 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 + FROM workspace_agent_stats + WHERE usage = true + GROUP BY user_id, agent_id, workspace_id +), latest_agent_latencies AS ( + SELECT + agent_id, + coalesce(MAX(connection_median_latency_ms), 0)::float AS connection_median_latency_ms + FROM workspace_agent_stats + GROUP BY user_id, agent_id, 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 + latest_agent_latencies +ON + agent_stats.agent_id = latest_agent_latencies.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 +` + +type GetWorkspaceAgentUsageStatsAndLabelsRow struct { + Username string `db:"username" json:"username"` + AgentName string `db:"agent_name" json:"agent_name"` + WorkspaceName string `db:"workspace_name" json:"workspace_name"` + RxBytes int64 `db:"rx_bytes" json:"rx_bytes"` + TxBytes int64 `db:"tx_bytes" json:"tx_bytes"` + SessionCountVSCode int64 `db:"session_count_vscode" json:"session_count_vscode"` + SessionCountSSH int64 `db:"session_count_ssh" json:"session_count_ssh"` + SessionCountJetBrains int64 `db:"session_count_jetbrains" json:"session_count_jetbrains"` + SessionCountReconnectingPTY int64 `db:"session_count_reconnecting_pty" json:"session_count_reconnecting_pty"` + ConnectionCount int64 `db:"connection_count" json:"connection_count"` + ConnectionMedianLatencyMS float64 `db:"connection_median_latency_ms" json:"connection_median_latency_ms"` +} + +func (q *sqlQuerier) GetWorkspaceAgentUsageStatsAndLabels(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentUsageStatsAndLabelsRow, error) { + rows, err := q.db.QueryContext(ctx, getWorkspaceAgentUsageStatsAndLabels, createdAt) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetWorkspaceAgentUsageStatsAndLabelsRow + for rows.Next() { + var i GetWorkspaceAgentUsageStatsAndLabelsRow + if err := rows.Scan( + &i.Username, + &i.AgentName, + &i.WorkspaceName, + &i.RxBytes, + &i.TxBytes, + &i.SessionCountVSCode, + &i.SessionCountSSH, + &i.SessionCountJetBrains, + &i.SessionCountReconnectingPTY, + &i.ConnectionCount, + &i.ConnectionMedianLatencyMS, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const insertWorkspaceAgentStats = `-- name: InsertWorkspaceAgentStats :exec INSERT INTO workspace_agent_stats ( @@ -11890,7 +12188,8 @@ INSERT INTO session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, - connection_median_latency_ms + connection_median_latency_ms, + usage ) SELECT unnest($1 :: uuid[]) AS id, @@ -11909,7 +12208,8 @@ SELECT 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 + unnest($17 :: double precision[]) AS connection_median_latency_ms, + unnest($18 :: boolean[]) AS usage ` type InsertWorkspaceAgentStatsParams struct { @@ -11930,6 +12230,7 @@ type InsertWorkspaceAgentStatsParams struct { 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"` + Usage []bool `db:"usage" json:"usage"` } func (q *sqlQuerier) InsertWorkspaceAgentStats(ctx context.Context, arg InsertWorkspaceAgentStatsParams) error { @@ -11951,6 +12252,7 @@ func (q *sqlQuerier) InsertWorkspaceAgentStats(ctx context.Context, arg InsertWo pq.Array(arg.SessionCountReconnectingPTY), pq.Array(arg.SessionCountSSH), pq.Array(arg.ConnectionMedianLatencyMS), + pq.Array(arg.Usage), ) return err } diff --git a/coderd/database/queries/workspaceagentstats.sql b/coderd/database/queries/workspaceagentstats.sql index ccc6d7ae21380..ab24960769f02 100644 --- 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 @@ -119,6 +121,60 @@ WITH agent_stats AS ( ) SELECT * FROM agent_stats, latest_agent_stats; +-- name: GetDeploymentWorkspaceAgentUsageStats :one +WITH agent_stats AS ( + SELECT + coalesce(SUM(rx_bytes), 0)::bigint AS workspace_rx_bytes, + coalesce(SUM(tx_bytes), 0)::bigint AS workspace_tx_bytes, + coalesce((PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_50, + coalesce((PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_95 + FROM workspace_agent_stats + -- The greater than 0 is to support legacy agents that don't report connection_median_latency_ms. + WHERE workspace_agent_stats.created_at > $1 AND connection_median_latency_ms > 0 +), +minute_buckets AS ( + SELECT + 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 + 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 +), +latest_buckets AS ( + SELECT DISTINCT ON (agent_id) + agent_id, + minute_bucket, + session_count_vscode, + session_count_jetbrains, + session_count_reconnecting_pty, + session_count_ssh + FROM + minute_buckets + ORDER BY + agent_id, + minute_bucket DESC +), +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 + FROM + latest_buckets +) +SELECT * FROM agent_stats, latest_agent_stats; + -- name: GetWorkspaceAgentStats :many WITH agent_stats AS ( SELECT @@ -148,6 +204,60 @@ WITH agent_stats AS ( ) SELECT * FROM agent_stats JOIN latest_agent_stats ON agent_stats.agent_id = latest_agent_stats.agent_id; +-- name: GetWorkspaceAgentUsageStats :many +WITH agent_stats AS ( + SELECT + user_id, + agent_id, + workspace_id, + template_id, + MIN(created_at)::timestamptz AS aggregated_from, + coalesce(SUM(rx_bytes), 0)::bigint AS workspace_rx_bytes, + coalesce(SUM(tx_bytes), 0)::bigint AS workspace_tx_bytes, + coalesce((PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_50, + coalesce((PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_95 + FROM workspace_agent_stats + -- The greater than 0 is to support legacy agents that don't report connection_median_latency_ms. + WHERE workspace_agent_stats.created_at > $1 AND connection_median_latency_ms > 0 + GROUP BY user_id, agent_id, workspace_id, template_id +), +minute_buckets AS ( + SELECT + 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 + 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, + user_id, + agent_id, + workspace_id, + template_id +), +latest_buckets AS ( + SELECT DISTINCT ON (agent_id) + agent_id, + session_count_vscode, + session_count_ssh, + session_count_jetbrains, + session_count_reconnecting_pty + FROM + minute_buckets + ORDER BY + agent_id, + minute_bucket DESC +) +SELECT * FROM agent_stats JOIN latest_buckets ON agent_stats.agent_id = latest_buckets.agent_id; + -- name: GetWorkspaceAgentStatsAndLabels :many WITH agent_stats AS ( SELECT @@ -199,3 +309,59 @@ JOIN workspaces ON workspaces.id = agent_stats.workspace_id; + +-- name: GetWorkspaceAgentUsageStatsAndLabels :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 + 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 + FROM workspace_agent_stats + WHERE usage = true + GROUP BY user_id, agent_id, workspace_id +), latest_agent_latencies AS ( + SELECT + agent_id, + coalesce(MAX(connection_median_latency_ms), 0)::float AS connection_median_latency_ms + FROM workspace_agent_stats + GROUP BY user_id, agent_id, 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 + latest_agent_latencies +ON + agent_stats.agent_id = latest_agent_latencies.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; diff --git a/coderd/insights_test.go b/coderd/insights_test.go index 40daed5d0ce02..9bd3f66ddf9e6 100644 --- 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 index 13b2e1bb6856a..63817330fdebc 100644 --- a/coderd/metricscache/metricscache.go +++ b/coderd/metricscache/metricscache.go @@ -34,6 +34,10 @@ 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 { @@ -41,7 +45,7 @@ type Intervals struct { 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 index bcc9396d3cbc0..54f22322e32d8 100644 --- 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 index 62193b6d673f0..2c4f0b52c0116 100644 --- 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 diff --git a/coderd/workspacestats/batcher.go b/coderd/workspacestats/batcher.go index 2872c368dc61c..1f14c5cec5a17 100644 --- a/coderd/workspacestats/batcher.go +++ b/coderd/workspacestats/batcher.go @@ -25,7 +25,7 @@ const ( ) type Batcher interface { - Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error + Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats, usage bool) error } // DBBatcher holds a buffer of agent stats and periodically flushes them to @@ -138,6 +138,7 @@ func (b *DBBatcher) Add( userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats, + usage bool, ) error { b.mu.Lock() defer b.mu.Unlock() @@ -165,6 +166,7 @@ func (b *DBBatcher) Add( b.buf.SessionCountReconnectingPTY = append(b.buf.SessionCountReconnectingPTY, st.SessionCountReconnectingPty) b.buf.SessionCountSSH = append(b.buf.SessionCountSSH, st.SessionCountSsh) b.buf.ConnectionMedianLatencyMS = append(b.buf.ConnectionMedianLatencyMS, st.ConnectionMedianLatencyMs) + b.buf.Usage = append(b.buf.Usage, usage) // If the buffer is over 80% full, signal the flusher to flush immediately. // We want to trigger flushes early to reduce the likelihood of @@ -279,6 +281,7 @@ func (b *DBBatcher) initBuf(size int) { SessionCountReconnectingPTY: make([]int64, 0, b.batchSize), SessionCountSSH: make([]int64, 0, b.batchSize), ConnectionMedianLatencyMS: make([]float64, 0, b.batchSize), + Usage: make([]bool, 0, b.batchSize), } b.connectionsByProto = make([]map[string]int64, 0, size) @@ -302,5 +305,6 @@ func (b *DBBatcher) resetBuf() { b.buf.SessionCountReconnectingPTY = b.buf.SessionCountReconnectingPTY[:0] b.buf.SessionCountSSH = b.buf.SessionCountSSH[:0] b.buf.ConnectionMedianLatencyMS = b.buf.ConnectionMedianLatencyMS[:0] + b.buf.Usage = b.buf.Usage[:0] b.connectionsByProto = b.connectionsByProto[:0] } diff --git a/coderd/workspacestats/batcher_internal_test.go b/coderd/workspacestats/batcher_internal_test.go index 97fdaf9f2aec5..2f7a25b152127 100644 --- a/coderd/workspacestats/batcher_internal_test.go +++ b/coderd/workspacestats/batcher_internal_test.go @@ -63,7 +63,7 @@ func TestBatchStats(t *testing.T) { // Given: a single data point is added for workspace t2 := t1.Add(time.Second) t.Logf("inserting 1 stat") - require.NoError(t, b.Add(t2.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t))) + require.NoError(t, b.Add(t2.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t), false)) // When: it becomes time to report stats // Signal a tick and wait for a flush to complete. @@ -87,9 +87,9 @@ func TestBatchStats(t *testing.T) { t.Logf("inserting %d stats", defaultBufferSize) for i := 0; i < defaultBufferSize; i++ { if i%2 == 0 { - require.NoError(t, b.Add(t3.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t))) + require.NoError(t, b.Add(t3.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t), false)) } else { - require.NoError(t, b.Add(t3.Add(time.Millisecond), deps2.Agent.ID, deps2.User.ID, deps2.Template.ID, deps2.Workspace.ID, randStats(t))) + require.NoError(t, b.Add(t3.Add(time.Millisecond), deps2.Agent.ID, deps2.User.ID, deps2.Template.ID, deps2.Workspace.ID, randStats(t), false)) } } }() diff --git a/coderd/workspacestats/reporter.go b/coderd/workspacestats/reporter.go index c6b7afb3c68ad..fecfd1b1eda92 100644 --- a/coderd/workspacestats/reporter.go +++ b/coderd/workspacestats/reporter.go @@ -118,7 +118,7 @@ func (r *Reporter) ReportAppStats(ctx context.Context, stats []workspaceapps.Sta return nil } -func (r *Reporter) ReportAgentStats(ctx context.Context, now time.Time, workspace database.Workspace, workspaceAgent database.WorkspaceAgent, templateName string, stats *agentproto.Stats) error { +func (r *Reporter) ReportAgentStats(ctx context.Context, now time.Time, workspace database.Workspace, workspaceAgent database.WorkspaceAgent, templateName string, stats *agentproto.Stats, usage bool) error { if stats.ConnectionCount > 0 { var nextAutostart time.Time if workspace.AutostartSchedule.String != "" { @@ -143,7 +143,7 @@ func (r *Reporter) ReportAgentStats(ctx context.Context, now time.Time, workspac var errGroup errgroup.Group errGroup.Go(func() error { - err := r.opts.StatsBatcher.Add(now, workspaceAgent.ID, workspace.TemplateID, workspace.OwnerID, workspace.ID, stats) + err := r.opts.StatsBatcher.Add(now, workspaceAgent.ID, workspace.TemplateID, workspace.OwnerID, workspace.ID, stats, usage) if err != nil { r.opts.Logger.Error(ctx, "add agent stats to batcher", slog.Error(err)) return xerrors.Errorf("insert workspace agent stats batch: %w", err) diff --git a/coderd/workspacestats/workspacestatstest/batcher.go b/coderd/workspacestats/workspacestatstest/batcher.go index ad5ba60ad16d0..2f5dd7d13aa0a 100644 --- a/coderd/workspacestats/workspacestatstest/batcher.go +++ b/coderd/workspacestats/workspacestatstest/batcher.go @@ -20,11 +20,12 @@ type StatsBatcher struct { LastUserID uuid.UUID LastWorkspaceID uuid.UUID LastStats *agentproto.Stats + LastUsage bool } var _ workspacestats.Batcher = &StatsBatcher{} -func (b *StatsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error { +func (b *StatsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats, usage bool) error { b.Mu.Lock() defer b.Mu.Unlock() b.Called++ @@ -34,5 +35,6 @@ func (b *StatsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUI b.LastUserID = userID b.LastWorkspaceID = workspaceID b.LastStats = st + b.LastUsage = usage return nil }