Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 3b66bab

Browse files
mafredrimafredri
committed
fix: implement GetTemplateInsights in dbmem
1 parent 81058e2 commit 3b66bab

File tree

1 file changed

+150
-47
lines changed

1 file changed

+150
-47
lines changed

‎coderd/database/dbmem/dbmem.go

Lines changed: 150 additions & 47 deletions
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
}
@@ -3233,71 +3241,166 @@ func (q *FakeQuerier) GetTemplateInsights(_ context.Context, arg database.GetTem
32333241
return database.GetTemplateInsightsRow{}, err
32343242
}
32353243

3236-
templateIDSet := make(map[uuid.UUID]struct{})
3237-
appUsageIntervalsByUser := make(map[uuid.UUID]map[time.Time]*database.GetTemplateInsightsRow)
3238-
32393244
q.mutex.RLock()
32403245
defer q.mutex.RUnlock()
32413246

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

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

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

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

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

0 commit comments

Comments
 (0)