Skip to content

Commit 5290571

Browse files
mtojekmtojek
committed
fix
1 parent 9c7aef8 commit 5290571

File tree

2 files changed

+29
-12
lines changed

2 files changed

+29
-12
lines changed

‎coderd/prometheusmetrics/prometheusmetrics.go

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import (
2323
)
2424

2525
// ActiveUsers tracks the number of users that have authenticated within the past hour.
26-
func ActiveUsers(ctx context.Context, registerer prometheus.Registerer, db database.Store, duration time.Duration) (context.CancelFunc, error) {
26+
func ActiveUsers(ctx context.Context, registerer prometheus.Registerer, db database.Store, duration time.Duration) (func(), error) {
2727
if duration == 0 {
2828
duration = 5 * time.Minute
2929
}
@@ -40,8 +40,10 @@ func ActiveUsers(ctx context.Context, registerer prometheus.Registerer, db datab
4040
}
4141

4242
ctx, cancelFunc := context.WithCancel(ctx)
43+
done := make(chan struct{})
4344
ticker := time.NewTicker(duration)
4445
go func() {
46+
defer close(done)
4547
defer ticker.Stop()
4648
for {
4749
select {
@@ -61,11 +63,14 @@ func ActiveUsers(ctx context.Context, registerer prometheus.Registerer, db datab
6163
gauge.Set(float64(len(distinctUsers)))
6264
}
6365
}()
64-
return cancelFunc, nil
66+
return func() {
67+
cancelFunc()
68+
<-done
69+
}, nil
6570
}
6671

6772
// Workspaces tracks the total number of workspaces with labels on status.
68-
func Workspaces(ctx context.Context, registerer prometheus.Registerer, db database.Store, duration time.Duration) (context.CancelFunc, error) {
73+
func Workspaces(ctx context.Context, registerer prometheus.Registerer, db database.Store, duration time.Duration) (func(), error) {
6974
if duration == 0 {
7075
duration = 5 * time.Minute
7176
}
@@ -85,8 +90,11 @@ func Workspaces(ctx context.Context, registerer prometheus.Registerer, db databa
8590
gauge.WithLabelValues("pending").Set(0)
8691

8792
ctx, cancelFunc := context.WithCancel(ctx)
93+
done := make(chan struct{})
94+
8895
ticker := time.NewTicker(duration)
8996
go func() {
97+
defer close(done)
9098
defer ticker.Stop()
9199
for {
92100
select {
@@ -115,11 +123,14 @@ func Workspaces(ctx context.Context, registerer prometheus.Registerer, db databa
115123
}
116124
}
117125
}()
118-
return cancelFunc, nil
126+
return func() {
127+
cancelFunc()
128+
<-done
129+
}, nil
119130
}
120131

121132
// Agents tracks the total number of workspaces with labels on status.
122-
func Agents(ctx context.Context, logger slog.Logger, registerer prometheus.Registerer, db database.Store, coordinator *atomic.Pointer[tailnet.Coordinator], derpMap *tailcfg.DERPMap, agentInactiveDisconnectTimeout, duration time.Duration) (context.CancelFunc, error) {
133+
func Agents(ctx context.Context, logger slog.Logger, registerer prometheus.Registerer, db database.Store, coordinator *atomic.Pointer[tailnet.Coordinator], derpMap *tailcfg.DERPMap, agentInactiveDisconnectTimeout, duration time.Duration) (func(), error) {
123134
if duration == 0 {
124135
duration = 1 * time.Minute
125136
}
@@ -182,8 +193,11 @@ func Agents(ctx context.Context, logger slog.Logger, registerer prometheus.Regis
182193

183194
// nolint:gocritic // Prometheus must collect metrics for all Coder users.
184195
ctx, cancelFunc := context.WithCancel(dbauthz.AsSystemRestricted(ctx))
196+
done := make(chan struct{})
197+
185198
ticker := time.NewTicker(duration)
186199
go func() {
200+
defer close(done)
187201
defer ticker.Stop()
188202
for {
189203
select {
@@ -290,7 +304,10 @@ func Agents(ctx context.Context, logger slog.Logger, registerer prometheus.Regis
290304
ticker.Reset(duration)
291305
}
292306
}()
293-
return cancelFunc, nil
307+
return func() {
308+
cancelFunc()
309+
<-done
310+
}, nil
294311
}
295312

296313
func AgentStats(ctx context.Context, logger slog.Logger, registerer prometheus.Registerer, db database.Store, initialCreateAfter time.Time, duration time.Duration) (func(), error) {

‎coderd/prometheusmetrics/prometheusmetrics_test.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -90,9 +90,9 @@ func TestActiveUsers(t *testing.T) {
9090
t.Run(tc.Name, func(t *testing.T) {
9191
t.Parallel()
9292
registry := prometheus.NewRegistry()
93-
cancel, err := prometheusmetrics.ActiveUsers(context.Background(), registry, tc.Database(t), time.Millisecond)
93+
closeFunc, err := prometheusmetrics.ActiveUsers(context.Background(), registry, tc.Database(t), time.Millisecond)
9494
require.NoError(t, err)
95-
t.Cleanup(cancel)
95+
t.Cleanup(closeFunc)
9696

9797
require.Eventually(t, func() bool {
9898
metrics, err := registry.Gather()
@@ -222,9 +222,9 @@ func TestWorkspaces(t *testing.T) {
222222
t.Run(tc.Name, func(t *testing.T) {
223223
t.Parallel()
224224
registry := prometheus.NewRegistry()
225-
cancel, err := prometheusmetrics.Workspaces(context.Background(), registry, tc.Database(), time.Millisecond)
225+
closeFunc, err := prometheusmetrics.Workspaces(context.Background(), registry, tc.Database(), time.Millisecond)
226226
require.NoError(t, err)
227-
t.Cleanup(cancel)
227+
t.Cleanup(closeFunc)
228228

229229
require.Eventually(t, func() bool {
230230
metrics, err := registry.Gather()
@@ -306,8 +306,8 @@ func TestAgents(t *testing.T) {
306306
registry := prometheus.NewRegistry()
307307

308308
// when
309-
cancel, err := prometheusmetrics.Agents(context.Background(), slogtest.Make(t, nil), registry, db, &coordinatorPtr, derpMap, agentInactiveDisconnectTimeout, time.Millisecond)
310-
t.Cleanup(cancel)
309+
closeFunc, err := prometheusmetrics.Agents(context.Background(), slogtest.Make(t, nil), registry, db, &coordinatorPtr, derpMap, agentInactiveDisconnectTimeout, time.Millisecond)
310+
t.Cleanup(closeFunc)
311311

312312
// then
313313
require.NoError(t, err)

0 commit comments

Comments
 (0)