Skip to content

Commit 33c681f

Browse files
mafredrimafredri
committed
fix: implement GetTemplateInsights in dbmem
1 parent 3febbe4 commit 33c681f

File tree

1 file changed

+150
-47
lines changed

1 file changed

+150
-47
lines changed

‎coderd/database/dbmem/dbmem.go

+150-47
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616

1717
"github.com/google/uuid"
1818
"github.com/lib/pq"
19+
"golang.org/x/exp/constraints"
1920
"golang.org/x/exp/maps"
2021
"golang.org/x/exp/slices"
2122
"golang.org/x/xerrors"
@@ -791,6 +792,13 @@ func tagsSubset(m1, m2 map[string]string) bool {
791792
// default tags when no tag is specified for a provisioner or job
792793
var tagsUntagged = provisionersdk.MutateTags(uuid.Nil, nil)
793794

795+
func least[T constraints.Ordered](a, b T) T {
796+
if a < b {
797+
return a
798+
}
799+
return b
800+
}
801+
794802
func (*FakeQuerier) AcquireLock(_ context.Context, _ int64) error {
795803
return xerrors.New("AcquireLock must only be called within a transaction")
796804
}
@@ -3237,71 +3245,166 @@ func (q *FakeQuerier) GetTemplateInsights(_ context.Context, arg database.GetTem
32373245
return database.GetTemplateInsightsRow{}, err
32383246
}
32393247

3240-
templateIDSet := make(map[uuid.UUID]struct{})
3241-
appUsageIntervalsByUser := make(map[uuid.UUID]map[time.Time]*database.GetTemplateInsightsRow)
3242-
32433248
q.mutex.RLock()
32443249
defer q.mutex.RUnlock()
32453250

3246-
for _, s := range q.workspaceAgentStats {
3247-
if s.CreatedAt.Before(arg.StartTime) || s.CreatedAt.Equal(arg.EndTime) || s.CreatedAt.After(arg.EndTime) {
3251+
/*
3252+
WITH
3253+
*/
3254+
3255+
/*
3256+
insights AS (
3257+
SELECT
3258+
user_id,
3259+
-- See motivation in GetTemplateInsights for LEAST(SUM(n), 30).
3260+
LEAST(SUM(usage_mins), 30) AS usage_mins,
3261+
LEAST(SUM(ssh_mins), 30) AS ssh_mins,
3262+
LEAST(SUM(sftp_mins), 30) AS sftp_mins,
3263+
LEAST(SUM(reconnecting_pty_mins), 30) AS reconnecting_pty_mins,
3264+
LEAST(SUM(vscode_mins), 30) AS vscode_mins,
3265+
LEAST(SUM(jetbrains_mins), 30) AS jetbrains_mins
3266+
FROM
3267+
template_usage_stats
3268+
WHERE
3269+
start_time >= @start_time::timestamptz
3270+
AND end_time <= @end_time::timestamptz
3271+
AND CASE WHEN COALESCE(array_length(@template_ids::uuid[], 1), 0) > 0 THEN template_id = ANY(@template_ids::uuid[]) ELSE TRUE END
3272+
GROUP BY
3273+
start_time, user_id
3274+
),
3275+
*/
3276+
3277+
type insightsGroupBy struct {
3278+
StartTime time.Time
3279+
UserID uuid.UUID
3280+
}
3281+
type insightsRow struct {
3282+
insightsGroupBy
3283+
UsageMins int16
3284+
SSHMins int16
3285+
SFTPMins int16
3286+
ReconnectingPTYMins int16
3287+
VSCodeMins int16
3288+
JetBrainsMins int16
3289+
}
3290+
insights := make(map[insightsGroupBy]insightsRow)
3291+
for _, stat := range q.templateUsageStats {
3292+
if stat.StartTime.Before(arg.StartTime) || stat.EndTime.After(arg.EndTime) {
32483293
continue
32493294
}
3250-
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, s.TemplateID) {
3295+
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, stat.TemplateID) {
32513296
continue
32523297
}
3253-
if s.ConnectionCount == 0 {
3254-
continue
3298+
key := insightsGroupBy{
3299+
StartTime: stat.StartTime,
3300+
UserID: stat.UserID,
3301+
}
3302+
row, ok := insights[key]
3303+
if !ok {
3304+
row = insightsRow{
3305+
insightsGroupBy: key,
3306+
}
32553307
}
3308+
row.UsageMins = least(row.UsageMins+stat.UsageMins, 30)
3309+
row.SSHMins = least(row.SSHMins+stat.SshMins, 30)
3310+
row.SFTPMins = least(row.SFTPMins+stat.SftpMins, 30)
3311+
row.ReconnectingPTYMins = least(row.ReconnectingPTYMins+stat.ReconnectingPtyMins, 30)
3312+
row.VSCodeMins = least(row.VSCodeMins+stat.VscodeMins, 30)
3313+
row.JetBrainsMins = least(row.JetBrainsMins+stat.JetbrainsMins, 30)
3314+
insights[key] = row
3315+
}
32563316

3257-
templateIDSet[s.TemplateID] = struct{}{}
3258-
if appUsageIntervalsByUser[s.UserID] == nil {
3259-
appUsageIntervalsByUser[s.UserID] = make(map[time.Time]*database.GetTemplateInsightsRow)
3317+
/*
3318+
templates AS (
3319+
SELECT
3320+
array_agg(DISTINCT template_id) AS template_ids,
3321+
array_agg(DISTINCT template_id) FILTER (WHERE ssh_mins > 0) AS ssh_template_ids,
3322+
array_agg(DISTINCT template_id) FILTER (WHERE sftp_mins > 0) AS sftp_template_ids,
3323+
array_agg(DISTINCT template_id) FILTER (WHERE reconnecting_pty_mins > 0) AS reconnecting_pty_template_ids,
3324+
array_agg(DISTINCT template_id) FILTER (WHERE vscode_mins > 0) AS vscode_template_ids,
3325+
array_agg(DISTINCT template_id) FILTER (WHERE jetbrains_mins > 0) AS jetbrains_template_ids
3326+
FROM
3327+
template_usage_stats
3328+
WHERE
3329+
start_time >= @start_time::timestamptz
3330+
AND end_time <= @end_time::timestamptz
3331+
AND CASE WHEN COALESCE(array_length(@template_ids::uuid[], 1), 0) > 0 THEN template_id = ANY(@template_ids::uuid[]) ELSE TRUE END
3332+
)
3333+
*/
3334+
3335+
type templateRow struct {
3336+
TemplateIDs []uuid.UUID
3337+
SSHTemplateIDs []uuid.UUID
3338+
SFTPTemplateIDs []uuid.UUID
3339+
ReconnectingPTYIDs []uuid.UUID
3340+
VSCodeTemplateIDs []uuid.UUID
3341+
JetBrainsTemplateIDs []uuid.UUID
3342+
}
3343+
templates := templateRow{}
3344+
for _, stat := range q.templateUsageStats {
3345+
if stat.StartTime.Before(arg.StartTime) || stat.EndTime.After(arg.EndTime) {
3346+
continue
32603347
}
3261-
t := s.CreatedAt.Truncate(time.Minute)
3262-
if _, ok := appUsageIntervalsByUser[s.UserID][t]; !ok {
3263-
appUsageIntervalsByUser[s.UserID][t] = &database.GetTemplateInsightsRow{}
3348+
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, stat.TemplateID) {
3349+
continue
32643350
}
3265-
3266-
if s.SessionCountJetBrains > 0 {
3267-
appUsageIntervalsByUser[s.UserID][t].UsageJetbrainsSeconds = 60
3351+
templates.TemplateIDs = append(templates.TemplateIDs, stat.TemplateID)
3352+
if stat.SshMins > 0 {
3353+
templates.SSHTemplateIDs = append(templates.SSHTemplateIDs, stat.TemplateID)
32683354
}
3269-
if s.SessionCountVSCode > 0 {
3270-
appUsageIntervalsByUser[s.UserID][t].UsageVscodeSeconds = 60
3355+
if stat.SftpMins > 0 {
3356+
templates.SFTPTemplateIDs = append(templates.SFTPTemplateIDs, stat.TemplateID)
32713357
}
3272-
if s.SessionCountReconnectingPTY > 0 {
3273-
appUsageIntervalsByUser[s.UserID][t].UsageReconnectingPtySeconds = 60
3358+
if stat.ReconnectingPtyMins > 0 {
3359+
templates.ReconnectingPTYIDs = append(templates.ReconnectingPTYIDs, stat.TemplateID)
32743360
}
3275-
if s.SessionCountSSH > 0 {
3276-
appUsageIntervalsByUser[s.UserID][t].UsageSshSeconds = 60
3361+
if stat.VscodeMins > 0 {
3362+
templates.VSCodeTemplateIDs = append(templates.VSCodeTemplateIDs, stat.TemplateID)
3363+
}
3364+
if stat.JetbrainsMins > 0 {
3365+
templates.JetBrainsTemplateIDs = append(templates.JetBrainsTemplateIDs, stat.TemplateID)
32773366
}
32783367
}
32793368

3280-
templateIDs := make([]uuid.UUID, 0, len(templateIDSet))
3281-
for templateID := range templateIDSet {
3282-
templateIDs = append(templateIDs, templateID)
3283-
}
3284-
slices.SortFunc(templateIDs, func(a, b uuid.UUID) int {
3285-
return slice.Ascending(a.String(), b.String())
3286-
})
3287-
activeUserIDs := make([]uuid.UUID, 0, len(appUsageIntervalsByUser))
3288-
for userID := range appUsageIntervalsByUser {
3289-
activeUserIDs = append(activeUserIDs, userID)
3290-
}
3369+
/*
3370+
SELECT
3371+
COALESCE((SELECT template_ids FROM templates), '{}')::uuid[] AS template_ids, -- Includes app usage.
3372+
COALESCE((SELECT ssh_template_ids FROM templates), '{}')::uuid[] AS ssh_template_ids,
3373+
COALESCE((SELECT sftp_template_ids FROM templates), '{}')::uuid[] AS sftp_template_ids,
3374+
COALESCE((SELECT reconnecting_pty_template_ids FROM templates), '{}')::uuid[] AS reconnecting_pty_template_ids,
3375+
COALESCE((SELECT vscode_template_ids FROM templates), '{}')::uuid[] AS vscode_template_ids,
3376+
COALESCE((SELECT jetbrains_template_ids FROM templates), '{}')::uuid[] AS jetbrains_template_ids,
3377+
COALESCE(COUNT(DISTINCT user_id), 0)::bigint AS active_users, -- Includes app usage.
3378+
COALESCE(SUM(usage_mins) * 60, 0)::bigint AS usage_total_seconds, -- Includes app usage.
3379+
COALESCE(SUM(ssh_mins) * 60, 0)::bigint AS usage_ssh_seconds,
3380+
COALESCE(SUM(sftp_mins) * 60, 0)::bigint AS usage_sftp_seconds,
3381+
COALESCE(SUM(reconnecting_pty_mins) * 60, 0)::bigint AS usage_reconnecting_pty_seconds,
3382+
COALESCE(SUM(vscode_mins) * 60, 0)::bigint AS usage_vscode_seconds,
3383+
COALESCE(SUM(jetbrains_mins) * 60, 0)::bigint AS usage_jetbrains_seconds
3384+
FROM
3385+
insights;
3386+
*/
32913387

3292-
result := database.GetTemplateInsightsRow{
3293-
TemplateIDs: templateIDs,
3294-
ActiveUsers: int64(len(activeUserIDs)),
3295-
}
3296-
for _, intervals := range appUsageIntervalsByUser {
3297-
for _, interval := range intervals {
3298-
result.UsageJetbrainsSeconds += interval.UsageJetbrainsSeconds
3299-
result.UsageVscodeSeconds += interval.UsageVscodeSeconds
3300-
result.UsageReconnectingPtySeconds += interval.UsageReconnectingPtySeconds
3301-
result.UsageSshSeconds += interval.UsageSshSeconds
3302-
}
3303-
}
3304-
return result, nil
3388+
var row database.GetTemplateInsightsRow
3389+
row.TemplateIDs = uniqueSortedUUIDs(templates.TemplateIDs)
3390+
row.SshTemplateIds = uniqueSortedUUIDs(templates.SSHTemplateIDs)
3391+
row.SftpTemplateIds = uniqueSortedUUIDs(templates.SFTPTemplateIDs)
3392+
row.ReconnectingPtyTemplateIds = uniqueSortedUUIDs(templates.ReconnectingPTYIDs)
3393+
row.VscodeTemplateIds = uniqueSortedUUIDs(templates.VSCodeTemplateIDs)
3394+
row.JetbrainsTemplateIds = uniqueSortedUUIDs(templates.JetBrainsTemplateIDs)
3395+
activeUserIDs := make(map[uuid.UUID]struct{})
3396+
for _, insight := range insights {
3397+
activeUserIDs[insight.UserID] = struct{}{}
3398+
row.UsageTotalSeconds += int64(insight.UsageMins) * 60
3399+
row.UsageSshSeconds += int64(insight.SSHMins) * 60
3400+
row.UsageSftpSeconds += int64(insight.SFTPMins) * 60
3401+
row.UsageReconnectingPtySeconds += int64(insight.ReconnectingPTYMins) * 60
3402+
row.UsageVscodeSeconds += int64(insight.VSCodeMins) * 60
3403+
row.UsageJetbrainsSeconds += int64(insight.JetBrainsMins) * 60
3404+
}
3405+
row.ActiveUsers = int64(len(activeUserIDs))
3406+
3407+
return row, nil
33053408
}
33063409

33073410
func (q *FakeQuerier) GetTemplateInsightsByInterval(ctx context.Context, arg database.GetTemplateInsightsByIntervalParams) ([]database.GetTemplateInsightsByIntervalRow, error) {

0 commit comments

Comments
 (0)