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

Skip to content

Commit 2040b87

Browse files
mafredrimafredri
committed
refactor and add test for locking
1 parent c6de43a commit 2040b87

File tree

4 files changed

+207
-78
lines changed

4 files changed

+207
-78
lines changed

‎coderd/coderd.go

+11-12
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,6 @@ type Options struct {
181181

182182
UpdateAgentMetrics func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
183183
StatsBatcher *batchstats.Batcher
184-
DBRollupInterval time.Duration
185184

186185
WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
187186

@@ -194,6 +193,9 @@ type Options struct {
194193
// NewTicker is used for unit tests to replace "time.NewTicker".
195194
NewTicker func(duration time.Duration) (tick <-chan time.Time, done func())
196195

196+
// DatabaseRolluper rolls up template usage stats from raw agent and app
197+
// stats. This is used to provide insights in the WebUI.
198+
DatabaseRolluper *dbrollup.Rolluper
197199
// WorkspaceUsageTracker tracks workspace usage by the CLI.
198200
WorkspaceUsageTracker *workspaceusage.Tracker
199201
}
@@ -344,9 +346,6 @@ func New(options *Options) *API {
344346
if options.StatsBatcher == nil {
345347
panic("developer error: options.StatsBatcher is nil")
346348
}
347-
if options.DBRollupInterval == 0 {
348-
options.DBRollupInterval = dbrollup.DefaultInterval
349-
}
350349

351350
siteCacheDir := options.CacheDir
352351
if siteCacheDir != "" {
@@ -371,6 +370,10 @@ func New(options *Options) *API {
371370
OIDC: options.OIDCConfig,
372371
}
373372

373+
if options.DatabaseRolluper == nil {
374+
options.DatabaseRolluper = dbrollup.New(options.Logger.Named("dbrollup"), options.Database)
375+
}
376+
374377
if options.WorkspaceUsageTracker == nil {
375378
options.WorkspaceUsageTracker = workspaceusage.New(options.Database,
376379
workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
@@ -421,11 +424,7 @@ func New(options *Options) *API {
421424
options.Database,
422425
options.Pubsub,
423426
),
424-
rolluper: dbrollup.New(
425-
options.Logger,
426-
options.Database,
427-
options.DBRollupInterval,
428-
),
427+
dbRolluper: options.DatabaseRolluper,
429428
workspaceUsageTracker: options.WorkspaceUsageTracker,
430429
}
431430

@@ -1208,9 +1207,9 @@ type API struct {
12081207
statsBatcher *batchstats.Batcher
12091208

12101209
Acquirer *provisionerdserver.Acquirer
1211-
// rolluper rolls up template usage stats from raw agent and app
1210+
// dbRolluper rolls up template usage stats from raw agent and app
12121211
// stats. This is used to provide insights in the WebUI.
1213-
rolluper *dbrollup.Rolluper
1212+
dbRolluper *dbrollup.Rolluper
12141213
workspaceUsageTracker *workspaceusage.Tracker
12151214
}
12161215

@@ -1225,7 +1224,7 @@ func (api *API) Close() error {
12251224
api.WebsocketWaitGroup.Wait()
12261225
api.WebsocketWaitMutex.Unlock()
12271226

1228-
api.rolluper.Close()
1227+
api.dbRolluper.Close()
12291228
api.metricsCache.Close()
12301229
if api.updateChecker != nil {
12311230
api.updateChecker.Close()

‎coderd/coderdtest/coderdtest.go

+3-2
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ import (
5757
"github.com/coder/coder/v2/coderd/batchstats"
5858
"github.com/coder/coder/v2/coderd/database"
5959
"github.com/coder/coder/v2/coderd/database/dbauthz"
60+
"github.com/coder/coder/v2/coderd/database/dbrollup"
6061
"github.com/coder/coder/v2/coderd/database/dbtestutil"
6162
"github.com/coder/coder/v2/coderd/database/pubsub"
6263
"github.com/coder/coder/v2/coderd/externalauth"
@@ -147,7 +148,7 @@ type Options struct {
147148
WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
148149
AllowWorkspaceRenames bool
149150
NewTicker func(duration time.Duration) (<-chan time.Time, func())
150-
DBRollupInterval time.Duration
151+
DatabaseRolluper *dbrollup.Rolluper
151152
WorkspaceUsageTrackerFlush chan int
152153
WorkspaceUsageTrackerTick chan time.Time
153154
}
@@ -492,7 +493,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
492493
WorkspaceAppsStatsCollectorOptions: options.WorkspaceAppsStatsCollectorOptions,
493494
AllowWorkspaceRenames: options.AllowWorkspaceRenames,
494495
NewTicker: options.NewTicker,
495-
DBRollupInterval: options.DBRollupInterval,
496+
DatabaseRolluper: options.DatabaseRolluper,
496497
WorkspaceUsageTracker: wuTracker,
497498
}
498499
}

‎coderd/database/dbrollup/dbrollup.go

+64-21
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package dbrollup
22

33
import (
44
"context"
5+
"flag"
56
"time"
67

78
"golang.org/x/sync/errgroup"
@@ -19,36 +20,68 @@ const (
1920
DefaultInterval = 5 * time.Minute
2021
)
2122

23+
type Event struct {
24+
TemplateUsageStats bool
25+
}
26+
2227
type Rolluper struct {
23-
cancel context.CancelFunc
24-
closed chan struct{}
25-
db database.Store
26-
logger slog.Logger
28+
cancel context.CancelFunc
29+
closed chan struct{}
30+
db database.Store
31+
logger slog.Logger
32+
interval time.Duration
33+
event chan<- Event
34+
}
35+
36+
type Option func(*Rolluper)
37+
38+
// WithInterval sets the interval between rollups.
39+
func WithInterval(interval time.Duration) Option {
40+
return func(r *Rolluper) {
41+
r.interval = interval
42+
}
43+
}
44+
45+
// WithEventChannel sets the event channel to use for rollup events.
46+
//
47+
// This is only used for testing.
48+
func WithEventChannel(ch chan<- Event) Option {
49+
if flag.Lookup("test.v") == nil {
50+
panic("developer error: WithEventChannel is not to be used outside of tests")
51+
}
52+
return func(r *Rolluper) {
53+
r.event = ch
54+
}
2755
}
2856

2957
// New creates a new DB rollup service that periodically runs rollup queries.
3058
// It is the caller's responsibility to call Close on the returned instance.
3159
//
3260
// This is for e.g. generating insights data (template_usage_stats) from
3361
// raw data (workspace_agent_stats, workspace_app_stats).
34-
func New(logger slog.Logger, db database.Store, interval time.Duration) *Rolluper {
62+
func New(logger slog.Logger, db database.Store, opts ...Option) *Rolluper {
3563
ctx, cancel := context.WithCancel(context.Background())
3664

3765
r := &Rolluper{
38-
cancel: cancel,
39-
closed: make(chan struct{}),
40-
db: db,
41-
logger: logger.Named("dbrollup"),
66+
cancel: cancel,
67+
closed: make(chan struct{}),
68+
db: db,
69+
logger: logger,
70+
interval: DefaultInterval,
71+
}
72+
73+
for _, opt := range opts {
74+
opt(r)
4275
}
4376

4477
//nolint:gocritic // The system rolls up database tables without user input.
4578
ctx = dbauthz.AsSystemRestricted(ctx)
46-
go r.start(ctx, interval)
79+
go r.start(ctx)
4780

4881
return r
4982
}
5083

51-
func (r *Rolluper) start(ctx context.Context, interval time.Duration) {
84+
func (r *Rolluper) start(ctx context.Context) {
5285
defer close(r.closed)
5386

5487
do := func() {
@@ -58,7 +91,7 @@ func (r *Rolluper) start(ctx context.Context, interval time.Duration) {
5891
now := time.Now()
5992

6093
// Track whether or not we performed a rollup (we got the advisory lock).
61-
templateUsageStats := false
94+
var ev Event
6295

6396
eg.Go(func() error {
6497
return r.db.InTx(func(tx database.Store) error {
@@ -72,7 +105,7 @@ func (r *Rolluper) start(ctx context.Context, interval time.Duration) {
72105
return nil
73106
}
74107

75-
templateUsageStats = true
108+
ev.TemplateUsageStats = true
76109
return tx.UpsertTemplateUsageStats(ctx)
77110
}, nil)
78111
})
@@ -86,12 +119,22 @@ func (r *Rolluper) start(ctx context.Context, interval time.Duration) {
86119
if ctx.Err() == nil {
87120
r.logger.Error(ctx, "failed to rollup data", slog.Error(err))
88121
}
89-
} else {
90-
r.logger.Debug(ctx,
91-
"rolled up data",
92-
slog.F("took", time.Since(now)),
93-
slog.F("template_usage_stats", templateUsageStats),
94-
)
122+
return
123+
}
124+
125+
r.logger.Debug(ctx,
126+
"rolled up data",
127+
slog.F("took", time.Since(now)),
128+
slog.F("event", ev),
129+
)
130+
131+
// For testing.
132+
if r.event != nil {
133+
select {
134+
case <-ctx.Done():
135+
return
136+
case r.event <- ev:
137+
}
95138
}
96139
}
97140

@@ -108,11 +151,11 @@ func (r *Rolluper) start(ctx context.Context, interval time.Duration) {
108151
case <-t.C:
109152
// Ensure we're on the interval.
110153
now := time.Now()
111-
next := now.Add(interval).Truncate(interval) // Ensure we're on the interval and synced with the clock.
154+
next := now.Add(r.interval).Truncate(r.interval) // Ensure we're on the interval and synced with the clock.
112155
d := next.Sub(now)
113156
// Safety check (shouldn't be possible).
114157
if d <= 0 {
115-
d = interval
158+
d = r.interval
116159
}
117160
t.Reset(d)
118161

0 commit comments

Comments
 (0)