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

Skip to content

Commit b9c85a5

Browse files
mafredrimafredri
committed
fix: implement GetTemplateInsights in dbmem
1 parent e2e268c commit b9c85a5

File tree

1 file changed

+148
-48
lines changed

1 file changed

+148
-48
lines changed

‎coderd/database/dbmem/dbmem.go

Lines changed: 148 additions & 48 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,163 @@ 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)
3244+
/*
3245+
WITH
3246+
*/
32383247

3239-
q.mutex.RLock()
3240-
defer q.mutex.RUnlock()
3248+
/*
3249+
insights AS (
3250+
SELECT
3251+
user_id,
3252+
-- See motivation in GetTemplateInsights for LEAST(SUM(n), 30).
3253+
LEAST(SUM(usage_mins), 30) AS usage_mins,
3254+
LEAST(SUM(ssh_mins), 30) AS ssh_mins,
3255+
LEAST(SUM(sftp_mins), 30) AS sftp_mins,
3256+
LEAST(SUM(reconnecting_pty_mins), 30) AS reconnecting_pty_mins,
3257+
LEAST(SUM(vscode_mins), 30) AS vscode_mins,
3258+
LEAST(SUM(jetbrains_mins), 30) AS jetbrains_mins
3259+
FROM
3260+
template_usage_stats
3261+
WHERE
3262+
start_time >= @start_time::timestamptz
3263+
AND end_time <= @end_time::timestamptz
3264+
AND CASE WHEN COALESCE(array_length(@template_ids::uuid[], 1), 0) > 0 THEN template_id = ANY(@template_ids::uuid[]) ELSE TRUE END
3265+
GROUP BY
3266+
start_time, user_id
3267+
),
3268+
*/
32413269

3242-
for _, s := range q.workspaceAgentStats {
3243-
if s.CreatedAt.Before(arg.StartTime) || s.CreatedAt.Equal(arg.EndTime) || s.CreatedAt.After(arg.EndTime) {
3270+
type insightsGroupBy struct {
3271+
StartTime time.Time
3272+
UserID uuid.UUID
3273+
}
3274+
type insightsRow struct {
3275+
insightsGroupBy
3276+
UsageMins int16
3277+
SSHMins int16
3278+
SFTPMins int16
3279+
ReconnectingPTYMins int16
3280+
VSCodeMins int16
3281+
JetBrainsMins int16
3282+
}
3283+
insights := make(map[insightsGroupBy]insightsRow)
3284+
for _, stat := range q.templateUsageStats {
3285+
if stat.StartTime.Before(arg.StartTime) || stat.EndTime.After(arg.EndTime) {
32443286
continue
32453287
}
3246-
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, s.TemplateID) {
3288+
if len(arg.TemplateIDs) > 0 && !slices.Contains(arg.TemplateIDs, stat.TemplateID) {
32473289
continue
32483290
}
3249-
if s.ConnectionCount == 0 {
3250-
continue
3291+
key := insightsGroupBy{
3292+
StartTime: stat.StartTime,
3293+
UserID: stat.UserID,
32513294
}
3295+
row, ok := insights[key]
3296+
if !ok {
3297+
row = insightsRow{
3298+
insightsGroupBy: key,
3299+
}
3300+
}
3301+
row.UsageMins = least(row.UsageMins+stat.UsageMins, 30)
3302+
row.SSHMins = least(row.SSHMins+stat.SshMins, 30)
3303+
row.SFTPMins = least(row.SFTPMins+stat.SftpMins, 30)
3304+
row.ReconnectingPTYMins = least(row.ReconnectingPTYMins+stat.ReconnectingPtyMins, 30)
3305+
row.VSCodeMins = least(row.VSCodeMins+stat.VscodeMins, 30)
3306+
row.JetBrainsMins = least(row.JetBrainsMins+stat.JetbrainsMins, 30)
3307+
insights[key] = row
3308+
}
32523309

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

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-
}
3362+
/*
3363+
SELECT
3364+
COALESCE((SELECT template_ids FROM templates), '{}')::uuid[] AS template_ids, -- Includes app usage.
3365+
COALESCE((SELECT ssh_template_ids FROM templates), '{}')::uuid[] AS ssh_template_ids,
3366+
COALESCE((SELECT sftp_template_ids FROM templates), '{}')::uuid[] AS sftp_template_ids,
3367+
COALESCE((SELECT reconnecting_pty_template_ids FROM templates), '{}')::uuid[] AS reconnecting_pty_template_ids,
3368+
COALESCE((SELECT vscode_template_ids FROM templates), '{}')::uuid[] AS vscode_template_ids,
3369+
COALESCE((SELECT jetbrains_template_ids FROM templates), '{}')::uuid[] AS jetbrains_template_ids,
3370+
COALESCE(COUNT(DISTINCT user_id), 0)::bigint AS active_users, -- Includes app usage.
3371+
COALESCE(SUM(usage_mins) * 60, 0)::bigint AS usage_total_seconds, -- Includes app usage.
3372+
COALESCE(SUM(ssh_mins) * 60, 0)::bigint AS usage_ssh_seconds,
3373+
COALESCE(SUM(sftp_mins) * 60, 0)::bigint AS usage_sftp_seconds,
3374+
COALESCE(SUM(reconnecting_pty_mins) * 60, 0)::bigint AS usage_reconnecting_pty_seconds,
3375+
COALESCE(SUM(vscode_mins) * 60, 0)::bigint AS usage_vscode_seconds,
3376+
COALESCE(SUM(jetbrains_mins) * 60, 0)::bigint AS usage_jetbrains_seconds
3377+
FROM
3378+
insights;
3379+
*/
32873380

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
3381+
var row database.GetTemplateInsightsRow
3382+
row.TemplateIDs = uniqueSortedUUIDs(templates.TemplateIDs)
3383+
row.SshTemplateIds = uniqueSortedUUIDs(templates.SSHTemplateIDs)
3384+
row.SftpTemplateIds = uniqueSortedUUIDs(templates.SFTPTemplateIDs)
3385+
row.ReconnectingPtyTemplateIds = uniqueSortedUUIDs(templates.ReconnectingPTYIDs)
3386+
row.VscodeTemplateIds = uniqueSortedUUIDs(templates.VSCodeTemplateIDs)
3387+
row.JetbrainsTemplateIds = uniqueSortedUUIDs(templates.JetBrainsTemplateIDs)
3388+
activeUserIDs := make(map[uuid.UUID]struct{})
3389+
for _, insight := range insights {
3390+
activeUserIDs[insight.UserID] = struct{}{}
3391+
row.UsageTotalSeconds += int64(insight.UsageMins) * 60
3392+
row.UsageSshSeconds += int64(insight.SSHMins) * 60
3393+
row.UsageSftpSeconds += int64(insight.SFTPMins) * 60
3394+
row.UsageReconnectingPtySeconds += int64(insight.ReconnectingPTYMins) * 60
3395+
row.UsageVscodeSeconds += int64(insight.VSCodeMins) * 60
3396+
row.UsageJetbrainsSeconds += int64(insight.JetBrainsMins) * 60
3397+
}
3398+
row.ActiveUsers = int64(len(activeUserIDs))
3399+
3400+
return row, nil
33013401
}
33023402

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

0 commit comments

Comments
 (0)