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

Skip to content

Commit 09a2dad

Browse files
kylecarbskylecarbs
committed
Fix aggregation query
1 parent 29719a4 commit 09a2dad

File tree

11 files changed

+207
-57
lines changed

11 files changed

+207
-57
lines changed

‎coderd/coderd.go

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,10 @@ func New(options *Options) *API {
222222
metricsCache := metricscache.New(
223223
options.Database,
224224
options.Logger.Named("metrics_cache"),
225-
options.MetricsCacheRefreshInterval,
225+
metricscache.Intervals{
226+
TemplateDAUs: options.MetricsCacheRefreshInterval,
227+
DeploymentStats: options.AgentStatsRefreshInterval,
228+
},
226229
)
227230

228231
staticHandler := site.Handler(site.FS(), binFS, binHashes)

‎coderd/database/dbgen/generator.go

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"crypto/sha256"
66
"database/sql"
77
"encoding/hex"
8+
"encoding/json"
89
"fmt"
910
"net"
1011
"testing"
@@ -437,3 +438,30 @@ func ParameterValue(t testing.TB, db database.Store, seed database.ParameterValu
437438
require.NoError(t, err, "insert parameter value")
438439
return scheme
439440
}
441+
442+
func WorkspaceAgentStat(t testing.TB, db database.Store, orig database.WorkspaceAgentStat) database.WorkspaceAgentStat {
443+
if orig.ConnectionsByProto == nil {
444+
orig.ConnectionsByProto = json.RawMessage([]byte("{}"))
445+
}
446+
scheme, err := db.InsertWorkspaceAgentStat(context.Background(), database.InsertWorkspaceAgentStatParams{
447+
ID: takeFirst(orig.ID, uuid.New()),
448+
CreatedAt: takeFirst(orig.CreatedAt, database.Now()),
449+
UserID: takeFirst(orig.UserID, uuid.New()),
450+
TemplateID: takeFirst(orig.TemplateID, uuid.New()),
451+
WorkspaceID: takeFirst(orig.WorkspaceID, uuid.New()),
452+
AgentID: takeFirst(orig.AgentID, uuid.New()),
453+
ConnectionsByProto: orig.ConnectionsByProto,
454+
ConnectionCount: takeFirst(orig.ConnectionCount, 0),
455+
RxPackets: takeFirst(orig.RxPackets, 0),
456+
RxBytes: takeFirst(orig.RxBytes, 0),
457+
TxPackets: takeFirst(orig.TxPackets, 0),
458+
TxBytes: takeFirst(orig.TxBytes, 0),
459+
SessionCountVSCode: takeFirst(orig.SessionCountVSCode, 0),
460+
SessionCountJetBrains: takeFirst(orig.SessionCountJetBrains, 0),
461+
SessionCountReconnectingPTY: takeFirst(orig.SessionCountReconnectingPTY, 0),
462+
SessionCountSSH: takeFirst(orig.SessionCountSSH, 0),
463+
ConnectionMedianLatencyMS: takeFirst(orig.ConnectionMedianLatencyMS, 0),
464+
})
465+
require.NoError(t, err, "insert workspace agent stat")
466+
return scheme
467+
}

‎coderd/database/dump.sql

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/migrations/000102_workspace_agent_stats_types.up.sql

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
ALTER TABLE workspace_agent_stats ADD COLUMN connection_median_latency_ms bigint DEFAULT -1 NOT NULL;
1+
ALTER TABLE workspace_agent_stats ADD COLUMN connection_median_latency_ms float DEFAULT -1 NOT NULL;
22
ALTER TABLE workspace_agent_stats ADD COLUMN session_count_vscode bigint DEFAULT 0 NOT NULL;
33
ALTER TABLE workspace_agent_stats ADD COLUMN session_count_jetbrains bigint DEFAULT 0 NOT NULL;
44
ALTER TABLE workspace_agent_stats ADD COLUMN session_count_reconnecting_pty bigint DEFAULT 0 NOT NULL;

‎coderd/database/models.go

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier_test.go

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
package database_test
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/google/uuid"
9+
"github.com/stretchr/testify/require"
10+
11+
"github.com/coder/coder/coderd/database"
12+
"github.com/coder/coder/coderd/database/dbgen"
13+
"github.com/coder/coder/coderd/database/migrations"
14+
)
15+
16+
func TestGetDeploymentWorkspaceAgentStats(t *testing.T) {
17+
t.Parallel()
18+
if testing.Short() {
19+
t.SkipNow()
20+
}
21+
t.Run("Aggregates", func(t *testing.T) {
22+
t.Parallel()
23+
sqlDB := testSQLDB(t)
24+
err := migrations.Up(sqlDB)
25+
require.NoError(t, err)
26+
db := database.New(sqlDB)
27+
ctx := context.Background()
28+
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
29+
TxBytes: 1,
30+
RxBytes: 1,
31+
ConnectionMedianLatencyMS: 1,
32+
SessionCountVSCode: 1,
33+
})
34+
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
35+
TxBytes: 1,
36+
RxBytes: 1,
37+
ConnectionMedianLatencyMS: 2,
38+
SessionCountVSCode: 1,
39+
})
40+
stats, err := db.GetDeploymentWorkspaceAgentStats(ctx, database.Now().Add(-time.Hour))
41+
require.NoError(t, err)
42+
43+
require.Equal(t, int64(2), stats.WorkspaceTxBytes)
44+
require.Equal(t, int64(2), stats.WorkspaceRxBytes)
45+
require.Equal(t, 1.5, stats.WorkspaceConnectionLatency50)
46+
require.Equal(t, 1.95, stats.WorkspaceConnectionLatency95)
47+
require.Equal(t, int64(2), stats.SessionCountVSCode)
48+
})
49+
50+
t.Run("GroupsByAgentID", func(t *testing.T) {
51+
t.Parallel()
52+
53+
sqlDB := testSQLDB(t)
54+
err := migrations.Up(sqlDB)
55+
require.NoError(t, err)
56+
db := database.New(sqlDB)
57+
ctx := context.Background()
58+
agentID := uuid.New()
59+
insertTime := database.Now()
60+
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
61+
CreatedAt: insertTime.Add(-time.Second),
62+
AgentID: agentID,
63+
TxBytes: 1,
64+
RxBytes: 1,
65+
ConnectionMedianLatencyMS: 1,
66+
SessionCountVSCode: 1,
67+
})
68+
dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
69+
// Ensure this stat is newer!
70+
CreatedAt: insertTime,
71+
AgentID: agentID,
72+
TxBytes: 1,
73+
RxBytes: 1,
74+
ConnectionMedianLatencyMS: 2,
75+
SessionCountVSCode: 1,
76+
})
77+
stats, err := db.GetDeploymentWorkspaceAgentStats(ctx, database.Now().Add(-time.Hour))
78+
require.NoError(t, err)
79+
80+
require.Equal(t, int64(2), stats.WorkspaceTxBytes)
81+
require.Equal(t, int64(2), stats.WorkspaceRxBytes)
82+
require.Equal(t, 1.5, stats.WorkspaceConnectionLatency50)
83+
require.Equal(t, 1.95, stats.WorkspaceConnectionLatency95)
84+
require.Equal(t, int64(1), stats.SessionCountVSCode)
85+
})
86+
}

‎coderd/database/queries.sql.go

Lines changed: 25 additions & 22 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/workspaceagentstats.sql

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -51,19 +51,22 @@ DELETE FROM workspace_agent_stats WHERE created_at < NOW() - INTERVAL '30 days';
5151

5252
-- name: GetDeploymentWorkspaceAgentStats :one
5353
WITH agent_stats AS (
54-
SELECT * FROM workspace_agent_stats
54+
SELECT
55+
coalesce(SUM(rx_bytes), 0)::bigint AS workspace_rx_bytes,
56+
coalesce(SUM(tx_bytes), 0)::bigint AS workspace_tx_bytes,
57+
coalesce((PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_50,
58+
coalesce((PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_95
59+
FROM workspace_agent_stats
5560
WHERE workspace_agent_stats.created_at > $1
5661
), latest_agent_stats AS (
57-
SELECT *, ROW_NUMBER() OVER(PARTITION BY id ORDER BY created_at DESC) AS rn
58-
FROM agent_stats
62+
SELECT
63+
coalesce(SUM(session_count_vscode), 0)::bigint AS session_count_vscode,
64+
coalesce(SUM(session_count_ssh), 0)::bigint AS session_count_ssh,
65+
coalesce(SUM(session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
66+
coalesce(SUM(session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty
67+
FROM (
68+
SELECT *, ROW_NUMBER() OVER(PARTITION BY agent_id ORDER BY created_at DESC) AS rn
69+
FROM workspace_agent_stats
70+
) AS a WHERE a.rn = 1
5971
)
60-
SELECT
61-
coalesce(SUM(latest_agent_stats.session_count_vscode), 0)::bigint AS session_count_vscode,
62-
coalesce(SUM(latest_agent_stats.session_count_ssh), 0)::bigint AS session_count_ssh,
63-
coalesce(SUM(latest_agent_stats.session_count_jetbrains), 0)::bigint AS session_count_jetbrains,
64-
coalesce(SUM(latest_agent_stats.session_count_reconnecting_pty), 0)::bigint AS session_count_reconnecting_pty,
65-
coalesce(SUM(agent_stats.rx_bytes), 0)::bigint AS workspace_rx_bytes,
66-
coalesce(SUM(agent_stats.tx_bytes), 0)::bigint AS workspace_tx_bytes,
67-
coalesce((PERCENTILE_DISC(0.5) WITHIN GROUP(ORDER BY agent_stats.connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_50,
68-
coalesce((PERCENTILE_DISC(0.95) WITHIN GROUP(ORDER BY agent_stats.connection_median_latency_ms)), -1)::FLOAT AS workspace_connection_latency_95
69-
FROM agent_stats JOIN latest_agent_stats ON agent_stats.agent_id = latest_agent_stats.agent_id AND rn = 1;
72+
SELECT * FROM agent_stats CROSS JOIN latest_agent_stats;

‎coderd/metricscache/metricscache.go

Lines changed: 35 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package metricscache
33
import (
44
"context"
55
"database/sql"
6+
"sync"
67
"sync/atomic"
78
"time"
89

@@ -36,13 +37,19 @@ type Cache struct {
3637

3738
done chan struct{}
3839
cancel func()
40+
}
3941

40-
interval time.Duration
42+
type Intervals struct {
43+
TemplateDAUs time.Duration
44+
DeploymentStats time.Duration
4145
}
4246

43-
func New(db database.Store, log slog.Logger, interval time.Duration) *Cache {
44-
if interval <= 0 {
45-
interval = time.Hour
47+
func New(db database.Store, log slog.Logger, intervals Intervals) *Cache {
48+
if intervals.TemplateDAUs <= 0 {
49+
intervals.TemplateDAUs = time.Hour
50+
}
51+
if intervals.DeploymentStats <= 0 {
52+
intervals.DeploymentStats = time.Minute
4653
}
4754
ctx, cancel := context.WithCancel(context.Background())
4855

@@ -51,9 +58,22 @@ func New(db database.Store, log slog.Logger, interval time.Duration) *Cache {
5158
log: log,
5259
done: make(chan struct{}),
5360
cancel: cancel,
54-
interval: interval,
5561
}
56-
go c.run(ctx)
62+
go func() {
63+
var wg sync.WaitGroup
64+
defer close(c.done)
65+
wg.Add(1)
66+
go func() {
67+
wg.Done()
68+
c.run(ctx, intervals.TemplateDAUs, c.refreshTemplateDAUs)
69+
}()
70+
wg.Add(1)
71+
go func() {
72+
wg.Done()
73+
c.run(ctx, intervals.DeploymentStats, c.refreshDeploymentStats)
74+
}()
75+
wg.Wait()
76+
}()
5777
return c
5878
}
5979

@@ -143,7 +163,7 @@ func countUniqueUsers(rows []database.GetTemplateDAUsRow) int {
143163
return len(seen)
144164
}
145165

146-
func (c *Cache) refresh(ctx context.Context) error {
166+
func (c *Cache) refreshTemplateDAUs(ctx context.Context) error {
147167
//nolint:gocritic // This is a system service.
148168
ctx = dbauthz.AsSystemRestricted(ctx)
149169
err := c.database.DeleteOldWorkspaceAgentStats(ctx)
@@ -197,6 +217,10 @@ func (c *Cache) refresh(ctx context.Context) error {
197217
c.templateUniqueUsers.Store(&templateUniqueUsers)
198218
c.templateAverageBuildTime.Store(&templateAverageBuildTimes)
199219

220+
return nil
221+
}
222+
223+
func (c *Cache) refreshDeploymentStats(ctx context.Context) error {
200224
from := database.Now().Add(-15 * time.Minute)
201225
deploymentStats, err := c.database.GetDeploymentWorkspaceAgentStats(ctx, from)
202226
if err != nil {
@@ -216,20 +240,17 @@ func (c *Cache) refresh(ctx context.Context) error {
216240
WorkspaceRxBytes: deploymentStats.WorkspaceRxBytes,
217241
WorkspaceTxBytes: deploymentStats.WorkspaceTxBytes,
218242
})
219-
220243
return nil
221244
}
222245

223-
func (c *Cache) run(ctx context.Context) {
224-
defer close(c.done)
225-
226-
ticker := time.NewTicker(c.interval)
246+
func (c *Cache) run(ctx context.Context, interval time.Duration, refresh func(context.Context) error) {
247+
ticker := time.NewTicker(interval)
227248
defer ticker.Stop()
228249

229250
for {
230251
for r := retry.New(time.Millisecond*100, time.Minute); r.Wait(ctx); {
231252
start := time.Now()
232-
err := c.refresh(ctx)
253+
err := refresh(ctx)
233254
if err != nil {
234255
if ctx.Err() != nil {
235256
return
@@ -241,7 +262,7 @@ func (c *Cache) run(ctx context.Context) {
241262
ctx,
242263
"metrics refreshed",
243264
slog.F("took", time.Since(start)),
244-
slog.F("interval", c.interval),
265+
slog.F("interval", interval),
245266
)
246267
break
247268
}

0 commit comments

Comments
 (0)