Skip to content

perf: store workspace agent session counts as JSONB #28126

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
EhabY wants to merge 1 commit into from
+910 −710
Draft

perf: store workspace agent session counts as JSONB #28126

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 42 additions & 39 deletions coderd/database/dbgen/dbgen.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1673,54 +1673,57 @@ func TemplateVersionTerraformValues(t testing.TB, db database.Store, orig databa
return v
}

func WorkspaceAgentStat(t testing.TB, db database.Store, orig database.WorkspaceAgentStat) database.WorkspaceAgentStat {
// WorkspaceAgentStat inserts a workspace agent stat row. The optional map seeds
// its session_counts column.
func WorkspaceAgentStat(t testing.TB, db database.Store, orig database.WorkspaceAgentStat, sessionCounts ...map[string]int64) database.WorkspaceAgentStat {
if orig.ConnectionsByProto == nil {
orig.ConnectionsByProto = json.RawMessage([]byte("{}"))
}
jsonProto := []byte(fmt.Sprintf("[%s]", orig.ConnectionsByProto))

// The insert rejects null session count elements.
counts := map[string]int64{}
if len(sessionCounts) > 0 && sessionCounts[0] != nil {
counts = sessionCounts[0]
}
jsonSessionCounts, err := json.Marshal([]map[string]int64{counts})
require.NoError(t, err, "marshal session counts")

params := database.InsertWorkspaceAgentStatsParams{
ID: []uuid.UUID{takeFirst(orig.ID, uuid.New())},
CreatedAt: []time.Time{takeFirst(orig.CreatedAt, dbtime.Now())},
UserID: []uuid.UUID{takeFirst(orig.UserID, uuid.New())},
TemplateID: []uuid.UUID{takeFirst(orig.TemplateID, uuid.New())},
WorkspaceID: []uuid.UUID{takeFirst(orig.WorkspaceID, uuid.New())},
AgentID: []uuid.UUID{takeFirst(orig.AgentID, uuid.New())},
ConnectionsByProto: jsonProto,
ConnectionCount: []int64{takeFirst(orig.ConnectionCount, 0)},
RxPackets: []int64{takeFirst(orig.RxPackets, 0)},
RxBytes: []int64{takeFirst(orig.RxBytes, 0)},
TxPackets: []int64{takeFirst(orig.TxPackets, 0)},
TxBytes: []int64{takeFirst(orig.TxBytes, 0)},
SessionCountVSCode: []int64{takeFirst(orig.SessionCountVSCode, 0)},
SessionCountJetBrains: []int64{takeFirst(orig.SessionCountJetBrains, 0)},
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)
ID: []uuid.UUID{takeFirst(orig.ID, uuid.New())},
CreatedAt: []time.Time{takeFirst(orig.CreatedAt, dbtime.Now())},
UserID: []uuid.UUID{takeFirst(orig.UserID, uuid.New())},
TemplateID: []uuid.UUID{takeFirst(orig.TemplateID, uuid.New())},
WorkspaceID: []uuid.UUID{takeFirst(orig.WorkspaceID, uuid.New())},
AgentID: []uuid.UUID{takeFirst(orig.AgentID, uuid.New())},
ConnectionsByProto: jsonProto,
ConnectionCount: []int64{takeFirst(orig.ConnectionCount, 0)},
RxPackets: []int64{takeFirst(orig.RxPackets, 0)},
RxBytes: []int64{takeFirst(orig.RxBytes, 0)},
TxPackets: []int64{takeFirst(orig.TxPackets, 0)},
TxBytes: []int64{takeFirst(orig.TxBytes, 0)},
SessionCounts: jsonSessionCounts,
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")

return database.WorkspaceAgentStat{
ID: params.ID[0],
CreatedAt: params.CreatedAt[0],
UserID: params.UserID[0],
AgentID: params.AgentID[0],
WorkspaceID: params.WorkspaceID[0],
TemplateID: params.TemplateID[0],
ConnectionsByProto: orig.ConnectionsByProto,
ConnectionCount: params.ConnectionCount[0],
RxPackets: params.RxPackets[0],
RxBytes: params.RxBytes[0],
TxPackets: params.TxPackets[0],
TxBytes: params.TxBytes[0],
ConnectionMedianLatencyMS: params.ConnectionMedianLatencyMS[0],
SessionCountVSCode: params.SessionCountVSCode[0],
SessionCountJetBrains: params.SessionCountJetBrains[0],
SessionCountReconnectingPTY: params.SessionCountReconnectingPTY[0],
SessionCountSSH: params.SessionCountSSH[0],
Usage: params.Usage[0],
ID: params.ID[0],
CreatedAt: params.CreatedAt[0],
UserID: params.UserID[0],
AgentID: params.AgentID[0],
WorkspaceID: params.WorkspaceID[0],
TemplateID: params.TemplateID[0],
ConnectionsByProto: orig.ConnectionsByProto,
ConnectionCount: params.ConnectionCount[0],
RxPackets: params.RxPackets[0],
RxBytes: params.RxBytes[0],
TxPackets: params.TxPackets[0],
TxBytes: params.TxBytes[0],
ConnectionMedianLatencyMS: params.ConnectionMedianLatencyMS[0],
Usage: params.Usage[0],
}
}

Expand Down
9 changes: 3 additions & 6 deletions coderd/database/dbpurge/dbpurge_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -387,26 +387,23 @@ func TestDeleteOldWorkspaceAgentStats(t *testing.T) {
ConnectionCount: 1,
ConnectionMedianLatencyMS: 1,
RxBytes: 1111,
SessionCountSSH: 1,
})
}, map[string]int64{"ssh": 1})

// Stat inserted 180 days - 2 hour ago, should not be deleted before rollup.
second := dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: now.AddDate(0, 0, -180).Add(2 * time.Hour),
ConnectionCount: 1,
ConnectionMedianLatencyMS: 1,
RxBytes: 2222,
SessionCountSSH: 1,
})
}, map[string]int64{"ssh": 1})

// Stat inserted 179 days - 4 hour ago, should not be deleted at all.
third := dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
CreatedAt: now.AddDate(0, 0, -179).Add(4 * time.Hour),
ConnectionCount: 1,
ConnectionMedianLatencyMS: 1,
RxBytes: 3333,
SessionCountSSH: 1,
})
}, map[string]int64{"ssh": 1})

// when
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
Expand Down
40 changes: 18 additions & 22 deletions coderd/database/dbrollup/dbrollup_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,7 @@ func TestRollup_TwoInstancesUseLocking(t *testing.T) {
CreatedAt: refTime.Add(-time.Minute),
ConnectionMedianLatencyMS: 1,
ConnectionCount: 1,
SessionCountSSH: 1,
})
}, map[string]int64{"ssh": 1})

closeRolluper := func(rolluper *dbrollup.Rolluper, resume chan struct{}) {
close(resume)
Expand Down Expand Up @@ -163,8 +162,7 @@ func TestRollupTemplateUsageStats(t *testing.T) {
CreatedAt: anHourAndSixMonthsAgo.AddDate(0, 0, -1),
ConnectionMedianLatencyMS: 1,
ConnectionCount: 1,
SessionCountSSH: 1,
})
}, map[string]int64{"ssh": 1})
_ = dbgen.WorkspaceAppStat(t, db, database.WorkspaceAppStat{
UserID: user.ID,
WorkspaceID: ws.ID,
Expand All @@ -176,25 +174,23 @@ func TestRollupTemplateUsageStats(t *testing.T) {

// Stats inserted 6 months - 1 day ago, should be rolled up.
wags1 := dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
TemplateID: tpl.ID,
WorkspaceID: ws.ID,
AgentID: agent.ID,
UserID: user.ID,
CreatedAt: anHourAndSixMonthsAgo.AddDate(0, 0, 1),
ConnectionMedianLatencyMS: 1,
ConnectionCount: 1,
SessionCountReconnectingPTY: 1,
})
TemplateID: tpl.ID,
WorkspaceID: ws.ID,
AgentID: agent.ID,
UserID: user.ID,
CreatedAt: anHourAndSixMonthsAgo.AddDate(0, 0, 1),
ConnectionMedianLatencyMS: 1,
ConnectionCount: 1,
}, map[string]int64{"reconnecting_pty": 1})
wags2 := dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
TemplateID: tpl.ID,
WorkspaceID: ws.ID,
AgentID: agent.ID,
UserID: user.ID,
CreatedAt: wags1.CreatedAt.Add(time.Minute),
ConnectionMedianLatencyMS: 1,
ConnectionCount: 1,
SessionCountReconnectingPTY: 1,
})
TemplateID: tpl.ID,
WorkspaceID: ws.ID,
AgentID: agent.ID,
UserID: user.ID,
CreatedAt: wags1.CreatedAt.Add(time.Minute),
ConnectionMedianLatencyMS: 1,
ConnectionCount: 1,
}, map[string]int64{"reconnecting_pty": 1})
// wags2 and waps1 overlap, so total usage is 4 - 1.
waps1 := dbgen.WorkspaceAppStat(t, db, database.WorkspaceAppStat{
UserID: user.ID,
Expand Down
11 changes: 5 additions & 6 deletions coderd/database/dump.sql
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

22 changes: 22 additions & 0 deletions coderd/database/migrations/000569_workspace_agent_session_counts.down.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
ALTER TABLE workspace_agent_stats
ADD COLUMN session_count_vscode bigint DEFAULT 0 NOT NULL,
ADD COLUMN session_count_jetbrains bigint DEFAULT 0 NOT NULL,
ADD COLUMN session_count_reconnecting_pty bigint DEFAULT 0 NOT NULL,
ADD COLUMN session_count_ssh bigint DEFAULT 0 NOT NULL;

-- Restore the four known session counts. Other keys are discarded.
UPDATE workspace_agent_stats
SET
session_count_vscode = COALESCE((session_counts ->> 'vscode')::bigint, 0),
session_count_jetbrains = COALESCE((session_counts ->> 'jetbrains')::bigint, 0),
session_count_reconnecting_pty = COALESCE((session_counts ->> 'reconnecting_pty')::bigint, 0),
session_count_ssh = COALESCE((session_counts ->> 'ssh')::bigint, 0);

DROP INDEX workspace_agent_stats_template_id_created_at_user_id_idx;

ALTER TABLE workspace_agent_stats
DROP COLUMN session_counts;

CREATE INDEX workspace_agent_stats_template_id_created_at_user_id_idx ON workspace_agent_stats USING btree (template_id, created_at, user_id) INCLUDE (session_count_vscode, session_count_jetbrains, session_count_reconnecting_pty, session_count_ssh, connection_median_latency_ms) WHERE (connection_count > 0);

COMMENT ON INDEX workspace_agent_stats_template_id_created_at_user_id_idx IS 'Support index for template insights endpoint to build interval reports faster.';
70 changes: 70 additions & 0 deletions coderd/database/migrations/000569_workspace_agent_session_counts.up.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
LOCK TABLE workspace_agent_stats IN ACCESS EXCLUSIVE MODE;

DO $$
DECLARE
latest_rollup_start timestamptz;
migration_cutoff timestamptz;
BEGIN
SELECT MAX(start_time) INTO latest_rollup_start FROM template_usage_stats;
migration_cutoff := COALESCE(latest_rollup_start - interval '1 day', statement_timestamp() - interval '180 days');

IF (latest_rollup_start IS NULL OR latest_rollup_start < statement_timestamp() - interval '24 hours')
AND EXISTS (
SELECT 1
FROM workspace_agent_stats
WHERE created_at >= migration_cutoff
AND (
session_count_vscode > 0
OR session_count_jetbrains > 0
OR session_count_reconnecting_pty > 0
OR session_count_ssh > 0
)
)
THEN
RAISE EXCEPTION 'migration 000569 requires template usage stats rolled up within the last 24 hours; run the previous Coder version until template usage stats roll up, then retry the upgrade'
USING DETAIL = format(
'Latest template_usage_stats.start_time: %s.',
COALESCE(latest_rollup_start::text, 'none')
),
HINT = 'Check coderd logs for "failed to rollup data" if the timestamp does not advance.';
END IF;
END
$$;

ALTER TABLE workspace_agent_stats
ADD COLUMN session_counts jsonb DEFAULT '{}'::jsonb NOT NULL;

COMMENT ON COLUMN workspace_agent_stats.session_counts IS 'Positive session counts keyed by the canonical app name reported by the agent.';

-- Convert the raw window still used by rollups and operational statistics.
-- Older usage already lives in template usage rollups.
UPDATE workspace_agent_stats
SET session_counts = jsonb_strip_nulls(jsonb_build_object(
'vscode', CASE WHEN session_count_vscode > 0 THEN session_count_vscode END,
'jetbrains', CASE WHEN session_count_jetbrains > 0 THEN session_count_jetbrains END,
'reconnecting_pty', CASE WHEN session_count_reconnecting_pty > 0 THEN session_count_reconnecting_pty END,
'ssh', CASE WHEN session_count_ssh > 0 THEN session_count_ssh END
))
WHERE created_at >= (
SELECT COALESCE(MAX(start_time) - interval '1 day', statement_timestamp() - interval '180 days')
FROM template_usage_stats
)
AND (
session_count_vscode > 0
OR session_count_jetbrains > 0
OR session_count_reconnecting_pty > 0
OR session_count_ssh > 0
);

-- Recreate the index without the dropped columns in its INCLUDE list.
DROP INDEX workspace_agent_stats_template_id_created_at_user_id_idx;

ALTER TABLE workspace_agent_stats
DROP COLUMN session_count_vscode,
DROP COLUMN session_count_jetbrains,
DROP COLUMN session_count_reconnecting_pty,
DROP COLUMN session_count_ssh;

CREATE INDEX workspace_agent_stats_template_id_created_at_user_id_idx ON workspace_agent_stats USING btree (template_id, created_at, user_id) INCLUDE (connection_median_latency_ms) WHERE (connection_count > 0);

COMMENT ON INDEX workspace_agent_stats_template_id_created_at_user_id_idx IS 'Support index for template insights endpoint to build interval reports faster.';
Loading
Loading