Skip to content

chore: move Batcher and Tracker to workspacestats #13418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jun 10, 2024
Merged

chore: move Batcher and Tracker to workspacestats #13418

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
530e334
chore: move Batcher and Tracker to workspacestats
f0ssel May 30, 2024
efed3bc
remove duplicate interface
f0ssel May 30, 2024
9295149
gotta hit save sometimes
f0ssel May 30, 2024
b5e83cb
rename interface
f0ssel Jun 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 6 additions & 7 deletions cli/server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,6 @@ import (
"github.com/coder/coder/v2/cli/config"
"github.com/coder/coder/v2/coderd"
"github.com/coder/coder/v2/coderd/autobuild"
"github.com/coder/coder/v2/coderd/batchstats"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/awsiamrds"
"github.com/coder/coder/v2/coderd/database/dbmem"
Expand All @@ -87,7 +86,7 @@ import (
stringutil "github.com/coder/coder/v2/coderd/util/strings"
"github.com/coder/coder/v2/coderd/workspaceapps"
"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
"github.com/coder/coder/v2/coderd/workspaceusage"
"github.com/coder/coder/v2/coderd/workspacestats"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/drpc"
"github.com/coder/coder/v2/cryptorand"
Expand Down Expand Up @@ -870,9 +869,9 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
options.SwaggerEndpoint = vals.Swagger.Enable.Value()
}

batcher, closeBatcher, err := batchstats.New(ctx,
batchstats.WithLogger(options.Logger.Named("batchstats")),
batchstats.WithStore(options.Database),
batcher, closeBatcher, err := workspacestats.NewBatcher(ctx,
workspacestats.BatcherWithLogger(options.Logger.Named("batchstats")),
workspacestats.BatcherWithStore(options.Database),
Comment on lines +872 to +874
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not particularly married to the functional opts here if you want to use a more conventional opts struct instead or just straight-up function args.

)
if err != nil {
return xerrors.Errorf("failed to create agent stats batcher: %w", err)
Expand Down Expand Up @@ -977,8 +976,8 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
defer purger.Close()

// Updates workspace usage
tracker := workspaceusage.New(options.Database,
workspaceusage.WithLogger(logger.Named("workspace_usage_tracker")),
tracker := workspacestats.NewTracker(options.Database,
workspacestats.TrackerWithLogger(logger.Named("workspace_usage_tracker")),
)
options.WorkspaceUsageTracker = tracker
defer tracker.Close()
Expand Down
6 changes: 0 additions & 6 deletions coderd/agentapi/stats.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,19 +7,13 @@ import (
"golang.org/x/xerrors"
"google.golang.org/protobuf/types/known/durationpb"

"github.com/google/uuid"

"cdr.dev/slog"
agentproto "github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/workspacestats"
)

type StatsBatcher interface {
Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error
}

type StatsAPI struct {
AgentFn func(context.Context) (database.WorkspaceAgent, error)
Database database.Store
Expand Down
2 changes: 1 addition & 1 deletion coderd/agentapi/stats_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ type statsBatcher struct {
lastStats *agentproto.Stats
}

var _ agentapi.StatsBatcher = &statsBatcher{}
var _ workspacestats.Batcher = &statsBatcher{}

func (b *statsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error {
b.mu.Lock()
Expand Down
19 changes: 8 additions & 11 deletions coderd/coderd.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,6 @@ import (
"github.com/coder/coder/v2/coderd/appearance"
"github.com/coder/coder/v2/coderd/audit"
"github.com/coder/coder/v2/coderd/awsidentity"
"github.com/coder/coder/v2/coderd/batchstats"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbrollup"
Expand All @@ -69,7 +68,6 @@ import (
"github.com/coder/coder/v2/coderd/util/slice"
"github.com/coder/coder/v2/coderd/workspaceapps"
"github.com/coder/coder/v2/coderd/workspacestats"
"github.com/coder/coder/v2/coderd/workspaceusage"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/drpc"
"github.com/coder/coder/v2/codersdk/healthsdk"
Expand Down Expand Up @@ -189,7 +187,7 @@ type Options struct {
HTTPClient *http.Client

UpdateAgentMetrics func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
StatsBatcher *batchstats.Batcher
StatsBatcher *workspacestats.DBBatcher

WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions

Expand All @@ -206,7 +204,7 @@ type Options struct {
// stats. This is used to provide insights in the WebUI.
DatabaseRolluper *dbrollup.Rolluper
// WorkspaceUsageTracker tracks workspace usage by the CLI.
WorkspaceUsageTracker *workspaceusage.Tracker
WorkspaceUsageTracker *workspacestats.UsageTracker
}

// @title Coder API
Expand Down Expand Up @@ -384,8 +382,8 @@ func New(options *Options) *API {
}

if options.WorkspaceUsageTracker == nil {
options.WorkspaceUsageTracker = workspaceusage.New(options.Database,
workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
options.WorkspaceUsageTracker = workspacestats.NewTracker(options.Database,
workspacestats.TrackerWithLogger(options.Logger.Named("workspace_usage_tracker")),
)
}

Expand Down Expand Up @@ -434,8 +432,7 @@ func New(options *Options) *API {
options.Database,
options.Pubsub,
),
dbRolluper: options.DatabaseRolluper,
workspaceUsageTracker: options.WorkspaceUsageTracker,
dbRolluper: options.DatabaseRolluper,
}

var customRoleHandler CustomRoleHandler = &agplCustomRoleHandler{}
Expand Down Expand Up @@ -557,6 +554,7 @@ func New(options *Options) *API {
Pubsub: options.Pubsub,
TemplateScheduleStore: options.TemplateScheduleStore,
StatsBatcher: options.StatsBatcher,
UsageTracker: options.WorkspaceUsageTracker,
UpdateAgentMetricsFn: options.UpdateAgentMetrics,
AppStatBatchSize: workspaceapps.DefaultStatsDBReporterBatchSize,
})
Expand Down Expand Up @@ -1301,8 +1299,7 @@ type API struct {
Acquirer *provisionerdserver.Acquirer
// dbRolluper rolls up template usage stats from raw agent and app
// stats. This is used to provide insights in the WebUI.
dbRolluper *dbrollup.Rolluper
workspaceUsageTracker *workspaceusage.Tracker
dbRolluper *dbrollup.Rolluper
}

// Close waits for all WebSocket connections to drain before returning.
Expand Down Expand Up @@ -1341,7 +1338,7 @@ func (api *API) Close() error {
_ = (*coordinator).Close()
}
_ = api.agentProvider.Close()
api.workspaceUsageTracker.Close()
_ = api.statsReporter.Close()
return nil
}

Expand Down
17 changes: 8 additions & 9 deletions coderd/coderdtest/coderdtest.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,6 @@ import (
"github.com/coder/coder/v2/coderd/audit"
"github.com/coder/coder/v2/coderd/autobuild"
"github.com/coder/coder/v2/coderd/awsidentity"
"github.com/coder/coder/v2/coderd/batchstats"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbrollup"
Expand All @@ -71,7 +70,7 @@ import (
"github.com/coder/coder/v2/coderd/util/ptr"
"github.com/coder/coder/v2/coderd/workspaceapps"
"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
"github.com/coder/coder/v2/coderd/workspaceusage"
"github.com/coder/coder/v2/coderd/workspacestats"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
"github.com/coder/coder/v2/codersdk/drpc"
Expand Down Expand Up @@ -145,7 +144,7 @@ type Options struct {
// Logger should only be overridden if you expect errors
// as part of your test.
Logger *slog.Logger
StatsBatcher *batchstats.Batcher
StatsBatcher *workspacestats.DBBatcher

WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions
AllowWorkspaceRenames bool
Expand Down Expand Up @@ -272,10 +271,10 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
if options.StatsBatcher == nil {
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
batcher, closeBatcher, err := batchstats.New(ctx,
batchstats.WithStore(options.Database),
batcher, closeBatcher, err := workspacestats.NewBatcher(ctx,
workspacestats.BatcherWithStore(options.Database),
// Avoid cluttering up test output.
batchstats.WithLogger(slog.Make(sloghuman.Sink(io.Discard))),
workspacestats.BatcherWithLogger(slog.Make(sloghuman.Sink(io.Discard))),
)
require.NoError(t, err, "create stats batcher")
options.StatsBatcher = batcher
Expand Down Expand Up @@ -337,10 +336,10 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
options.WorkspaceUsageTrackerTick = make(chan time.Time, 1) // buffering just in case
}
// Close is called by API.Close()
wuTracker := workspaceusage.New(
wuTracker := workspacestats.NewTracker(
options.Database,
workspaceusage.WithLogger(options.Logger.Named("workspace_usage_tracker")),
workspaceusage.WithTickFlush(options.WorkspaceUsageTrackerTick, options.WorkspaceUsageTrackerFlush),
workspacestats.TrackerWithLogger(options.Logger.Named("workspace_usage_tracker")),
workspacestats.TrackerWithTickFlush(options.WorkspaceUsageTrackerTick, options.WorkspaceUsageTrackerFlush),
)

var mutex sync.RWMutex
Expand Down
17 changes: 8 additions & 9 deletions coderd/insights_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ import (
"cdr.dev/slog/sloggers/slogtest"
"github.com/coder/coder/v2/agent/agenttest"
agentproto "github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/coderd/batchstats"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
Expand Down Expand Up @@ -684,11 +683,11 @@ func TestTemplateInsights_Golden(t *testing.T) {
// NOTE(mafredri): Ideally we would pass batcher as a coderd option and
// insert using the agentClient, but we have a circular dependency on
// the database.
batcher, batcherCloser, err := batchstats.New(
batcher, batcherCloser, err := workspacestats.NewBatcher(
ctx,
batchstats.WithStore(db),
batchstats.WithLogger(logger.Named("batchstats")),
batchstats.WithInterval(time.Hour),
workspacestats.BatcherWithStore(db),
workspacestats.BatcherWithLogger(logger.Named("batchstats")),
workspacestats.BatcherWithInterval(time.Hour),
)
require.NoError(t, err)
defer batcherCloser() // Flushes the stats, this is to ensure they're written.
Expand Down Expand Up @@ -1583,11 +1582,11 @@ func TestUserActivityInsights_Golden(t *testing.T) {
// NOTE(mafredri): Ideally we would pass batcher as a coderd option and
// insert using the agentClient, but we have a circular dependency on
// the database.
batcher, batcherCloser, err := batchstats.New(
batcher, batcherCloser, err := workspacestats.NewBatcher(
ctx,
batchstats.WithStore(db),
batchstats.WithLogger(logger.Named("batchstats")),
batchstats.WithInterval(time.Hour),
workspacestats.BatcherWithStore(db),
workspacestats.BatcherWithLogger(logger.Named("batchstats")),
workspacestats.BatcherWithInterval(time.Hour),
)
require.NoError(t, err)
defer batcherCloser() // Flushes the stats, this is to ensure they're written.
Expand Down
8 changes: 4 additions & 4 deletions coderd/prometheusmetrics/prometheusmetrics_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,14 +21,14 @@ import (
"cdr.dev/slog/sloggers/slogtest"

"github.com/coder/coder/v2/coderd/agentmetrics"
"github.com/coder/coder/v2/coderd/batchstats"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbgen"
"github.com/coder/coder/v2/coderd/database/dbmem"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/prometheusmetrics"
"github.com/coder/coder/v2/coderd/workspacestats"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
"github.com/coder/coder/v2/cryptorand"
Expand Down Expand Up @@ -391,14 +391,14 @@ func TestAgentStats(t *testing.T) {
db, pubsub := dbtestutil.NewDB(t)
log := slogtest.Make(t, nil).Leveled(slog.LevelDebug)

batcher, closeBatcher, err := batchstats.New(ctx,
batcher, closeBatcher, err := workspacestats.NewBatcher(ctx,
// We had previously set the batch size to 1 here, but that caused
// intermittent test flakes due to a race between the batcher completing
// its flush and the test asserting that the metrics were collected.
// Instead, we close the batcher after all stats have been posted, which
// forces a flush.
batchstats.WithStore(db),
batchstats.WithLogger(log),
workspacestats.BatcherWithStore(db),
workspacestats.BatcherWithLogger(log),
)
require.NoError(t, err, "create stats batcher failed")
t.Cleanup(closeBatcher)
Expand Down
2 changes: 1 addition & 1 deletion coderd/workspaces.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1115,7 +1115,7 @@ func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
return
}

api.workspaceUsageTracker.Add(workspace.ID)
api.statsReporter.TrackUsage(workspace.ID)
rw.WriteHeader(http.StatusNoContent)
}

Expand Down
Loading
Loading