Skip to content

Commit bf9e534

Browse files
coadlercoadler
committed
fixup! feat: add prometheus metric for tracking user statuses
1 parent 9b0d26d commit bf9e534

File tree

3 files changed

+38
-16
lines changed

3 files changed

+38
-16
lines changed

‎cli/server.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -218,7 +218,7 @@ func enablePrometheus(
218218
}
219219
afterCtx(ctx, closeActiveUsersFunc)
220220

221-
closeUsersFunc, err := prometheusmetrics.Users(ctx, options.Logger.Named("user_metrics"), options.PrometheusRegistry, options.Database, 0)
221+
closeUsersFunc, err := prometheusmetrics.Users(ctx, options.Logger.Named("user_metrics"), options.Clock, options.PrometheusRegistry, options.Database, 0)
222222
if err != nil {
223223
return nil, xerrors.Errorf("register active users prometheus metric: %w", err)
224224
}

‎coderd/prometheusmetrics/prometheusmetrics.go

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212

1313
"github.com/google/uuid"
1414
"github.com/prometheus/client_golang/prometheus"
15+
"golang.org/x/xerrors"
1516
"tailscale.com/tailcfg"
1617

1718
"cdr.dev/slog"
@@ -22,6 +23,7 @@ import (
2223
"github.com/coder/coder/v2/coderd/database/dbtime"
2324
"github.com/coder/coder/v2/codersdk"
2425
"github.com/coder/coder/v2/tailnet"
26+
"github.com/coder/quartz"
2527
)
2628

2729
const defaultRefreshRate = time.Minute
@@ -58,7 +60,7 @@ func ActiveUsers(ctx context.Context, logger slog.Logger, registerer prometheus.
5860

5961
apiKeys, err := db.GetAPIKeysLastUsedAfter(ctx, dbtime.Now().Add(-1*time.Hour))
6062
if err != nil {
61-
logger.Error(ctx, "get api keys", slog.Error(err))
63+
logger.Error(ctx, "get api keys for active users prometheus metric", slog.Error(err))
6264
continue
6365
}
6466
distinctUsers := map[uuid.UUID]struct{}{}
@@ -74,8 +76,8 @@ func ActiveUsers(ctx context.Context, logger slog.Logger, registerer prometheus.
7476
}, nil
7577
}
7678

77-
// Users tracks the total user count, broken out by status..
78-
func Users(ctx context.Context, logger slog.Logger, registerer prometheus.Registerer, db database.Store, duration time.Duration) (func(), error) {
79+
// Users tracks the total number of registered users, partitioned by status.
80+
func Users(ctx context.Context, logger slog.Logger, clk quartz.Clock, registerer prometheus.Registerer, db database.Store, duration time.Duration) (func(), error) {
7981
if duration == 0 {
8082
// It's not super important this tracks real-time.
8183
duration = defaultRefreshRate * 5
@@ -85,16 +87,16 @@ func Users(ctx context.Context, logger slog.Logger, registerer prometheus.Regist
8587
Namespace: "coderd",
8688
Subsystem: "api",
8789
Name: "total_user_count",
88-
Help: "The total number of users, broken out by status.",
90+
Help: "The total number of registered users, partitioned by status.",
8991
}, []string{"status"})
9092
err := registerer.Register(gauge)
9193
if err != nil {
92-
return nil, err
94+
return nil, xerrors.Errorf("register total_user_count gauge: %w", err)
9395
}
9496

9597
ctx, cancelFunc := context.WithCancel(ctx)
9698
done := make(chan struct{})
97-
ticker := time.NewTicker(duration)
99+
ticker := clk.NewTicker(duration)
98100
go func() {
99101
defer close(done)
100102
defer ticker.Stop()

‎coderd/prometheusmetrics/prometheusmetrics_test.go

Lines changed: 29 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ import (
3838
"github.com/coder/coder/v2/tailnet"
3939
"github.com/coder/coder/v2/tailnet/tailnettest"
4040
"github.com/coder/coder/v2/testutil"
41+
"github.com/coder/quartz"
4142
)
4243

4344
func TestActiveUsers(t *testing.T) {
@@ -118,21 +119,21 @@ func TestUsers(t *testing.T) {
118119
for _, tc := range []struct {
119120
Name string
120121
Database func(t *testing.T) database.Store
121-
Count map[string]int
122+
Count map[database.UserStatus]int
122123
}{{
123124
Name: "None",
124125
Database: func(t *testing.T) database.Store {
125126
return dbmem.New()
126127
},
127-
Count: map[string]int{},
128+
Count: map[database.UserStatus]int{},
128129
}, {
129130
Name: "One",
130131
Database: func(t *testing.T) database.Store {
131132
db := dbmem.New()
132133
dbgen.User(t, db, database.User{Status: database.UserStatusActive})
133134
return db
134135
},
135-
Count: map[string]int{"active": 1},
136+
Count: map[database.UserStatus]int{database.UserStatusActive: 1},
136137
}, {
137138
Name: "MultipleStatuses",
138139
Database: func(t *testing.T) database.Store {
@@ -143,7 +144,7 @@ func TestUsers(t *testing.T) {
143144

144145
return db
145146
},
146-
Count: map[string]int{"active": 1, "dormant": 1},
147+
Count: map[database.UserStatus]int{database.UserStatusActive: 1, database.UserStatusDormant: 1},
147148
}, {
148149
Name: "MultipleActive",
149150
Database: func(t *testing.T) database.Store {
@@ -153,17 +154,25 @@ func TestUsers(t *testing.T) {
153154
dbgen.User(t, db, database.User{Status: database.UserStatusActive})
154155
return db
155156
},
156-
Count: map[string]int{"active": 3},
157+
Count: map[database.UserStatus]int{database.UserStatusActive: 3},
157158
}} {
158159
tc := tc
159160
t.Run(tc.Name, func(t *testing.T) {
160161
t.Parallel()
162+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
163+
defer cancel()
164+
161165
registry := prometheus.NewRegistry()
162-
closeFunc, err := prometheusmetrics.Users(context.Background(), slogtest.Make(t, nil), registry, tc.Database(t), time.Millisecond)
166+
mClock := quartz.NewMock(t)
167+
db := tc.Database(t)
168+
closeFunc, err := prometheusmetrics.Users(context.Background(), slogtest.Make(t, nil), mClock, registry, db, time.Millisecond)
163169
require.NoError(t, err)
164170
t.Cleanup(closeFunc)
165171

166-
require.Eventually(t, func() bool {
172+
_, w := mClock.AdvanceNext()
173+
w.MustWait(ctx)
174+
175+
checkFn := func() bool {
167176
metrics, err := registry.Gather()
168177
if err != nil {
169178
return false
@@ -176,13 +185,24 @@ func TestUsers(t *testing.T) {
176185
}
177186

178187
for _, metric := range metrics[0].Metric {
179-
if tc.Count[*metric.Label[0].Value] != int(metric.Gauge.GetValue()) {
188+
if tc.Count[database.UserStatus(*metric.Label[0].Value)] != int(metric.Gauge.GetValue()) {
180189
return false
181190
}
182191
}
183192

184193
return true
185-
}, testutil.WaitShort, testutil.IntervalSlow)
194+
}
195+
196+
require.Eventually(t, checkFn, testutil.WaitShort, testutil.IntervalFast)
197+
198+
// Add another dormant user and ensure it updates
199+
dbgen.User(t, db, database.User{Status: database.UserStatusDormant})
200+
tc.Count[database.UserStatusDormant]++
201+
202+
_, w = mClock.AdvanceNext()
203+
w.MustWait(ctx)
204+
205+
require.Eventually(t, checkFn, testutil.WaitShort, testutil.IntervalFast)
186206
})
187207
}
188208
}

0 commit comments

Comments
 (0)