mirror of
https://github.com/coder/coder.git
synced 2026-06-03 04:58:23 +00:00
Spike Curtis
bddb808b25
Fixes all our Go file imports to match the preferred spec that we've _mostly_ been using. For example: ``` import ( "context" "time" "github.com/prometheus/client_golang/prometheus" "golang.org/x/xerrors" "gopkg.in/natefinch/lumberjack.v2" "cdr.dev/slog/v3" "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/serpent" ) ``` 3 groups: standard library, 3rd partly libs, Coder libs. This PR makes the change across the codebase. The PR in the stack above modifies our formatting to maintain this state of affairs, and is a separate PR so it's possible to review that one in detail.
259 lines
9.2 KiB
Go
259 lines
9.2 KiB
Go
package dbrollup_test
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
"go.uber.org/goleak"
|
|
|
|
"cdr.dev/slog/v3"
|
|
"cdr.dev/slog/v3/sloggers/slogtest"
|
|
"github.com/coder/coder/v2/coderd/database"
|
|
"github.com/coder/coder/v2/coderd/database/dbgen"
|
|
"github.com/coder/coder/v2/coderd/database/dbrollup"
|
|
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
|
"github.com/coder/coder/v2/coderd/database/dbtime"
|
|
"github.com/coder/coder/v2/testutil"
|
|
)
|
|
|
|
func TestMain(m *testing.M) {
|
|
goleak.VerifyTestMain(m, testutil.GoleakOptions...)
|
|
}
|
|
|
|
func TestRollup_Close(t *testing.T) {
|
|
t.Parallel()
|
|
db, _ := dbtestutil.NewDB(t)
|
|
rolluper := dbrollup.New(testutil.Logger(t), db, dbrollup.WithInterval(250*time.Millisecond))
|
|
err := rolluper.Close()
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
type wrapUpsertDB struct {
|
|
database.Store
|
|
resume <-chan struct{}
|
|
}
|
|
|
|
func (w *wrapUpsertDB) InTx(fn func(database.Store) error, opts *database.TxOptions) error {
|
|
return w.Store.InTx(func(tx database.Store) error {
|
|
return fn(&wrapUpsertDB{Store: tx, resume: w.resume})
|
|
}, opts)
|
|
}
|
|
|
|
func (w *wrapUpsertDB) UpsertTemplateUsageStats(ctx context.Context) error {
|
|
<-w.resume
|
|
return w.Store.UpsertTemplateUsageStats(ctx)
|
|
}
|
|
|
|
func TestRollup_TwoInstancesUseLocking(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
|
|
logger := testutil.Logger(t)
|
|
|
|
var (
|
|
org = dbgen.Organization(t, db, database.Organization{})
|
|
user = dbgen.User(t, db, database.User{Name: "user1"})
|
|
tpl = dbgen.Template(t, db, database.Template{OrganizationID: org.ID, CreatedBy: user.ID})
|
|
ver = dbgen.TemplateVersion(t, db, database.TemplateVersion{OrganizationID: org.ID, TemplateID: uuid.NullUUID{UUID: tpl.ID, Valid: true}, CreatedBy: user.ID})
|
|
ws = dbgen.Workspace(t, db, database.WorkspaceTable{OrganizationID: org.ID, TemplateID: tpl.ID, OwnerID: user.ID})
|
|
job = dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{OrganizationID: org.ID})
|
|
build = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{WorkspaceID: ws.ID, JobID: job.ID, TemplateVersionID: ver.ID})
|
|
res = dbgen.WorkspaceResource(t, db, database.WorkspaceResource{JobID: build.JobID})
|
|
agent = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ResourceID: res.ID})
|
|
)
|
|
|
|
refTime := dbtime.Now().Truncate(time.Hour)
|
|
_ = dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
|
|
TemplateID: tpl.ID,
|
|
WorkspaceID: ws.ID,
|
|
AgentID: agent.ID,
|
|
UserID: user.ID,
|
|
CreatedAt: refTime.Add(-time.Minute),
|
|
ConnectionMedianLatencyMS: 1,
|
|
ConnectionCount: 1,
|
|
SessionCountSSH: 1,
|
|
})
|
|
|
|
closeRolluper := func(rolluper *dbrollup.Rolluper, resume chan struct{}) {
|
|
close(resume)
|
|
err := rolluper.Close()
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
interval := dbrollup.WithInterval(250 * time.Millisecond)
|
|
events1 := make(chan dbrollup.Event)
|
|
resume1 := make(chan struct{}, 1)
|
|
rolluper1 := dbrollup.New(
|
|
logger.Named("dbrollup1"),
|
|
&wrapUpsertDB{Store: db, resume: resume1},
|
|
interval,
|
|
dbrollup.WithEventChannel(events1),
|
|
)
|
|
defer closeRolluper(rolluper1, resume1)
|
|
|
|
events2 := make(chan dbrollup.Event)
|
|
resume2 := make(chan struct{}, 1)
|
|
rolluper2 := dbrollup.New(
|
|
logger.Named("dbrollup2"),
|
|
&wrapUpsertDB{Store: db, resume: resume2},
|
|
interval,
|
|
dbrollup.WithEventChannel(events2),
|
|
)
|
|
defer closeRolluper(rolluper2, resume2)
|
|
|
|
_, _ = <-events1, <-events2 // Deplete init event, resume operation.
|
|
|
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
|
|
|
// One of the rollup instances should roll up and the other should not.
|
|
var ev1, ev2 dbrollup.Event
|
|
select {
|
|
case <-ctx.Done():
|
|
t.Fatal("timed out waiting for rollup to occur")
|
|
case ev1 = <-events1:
|
|
resume2 <- struct{}{}
|
|
ev2 = <-events2
|
|
case ev2 = <-events2:
|
|
resume1 <- struct{}{}
|
|
ev1 = <-events1
|
|
}
|
|
|
|
require.NotEqual(t, ev1, ev2, "one of the rollup instances should have rolled up and the other not")
|
|
|
|
rows, err := db.GetTemplateUsageStats(ctx, database.GetTemplateUsageStatsParams{
|
|
StartTime: refTime.Add(-time.Hour).Truncate(time.Hour),
|
|
EndTime: refTime,
|
|
})
|
|
require.NoError(t, err)
|
|
require.Len(t, rows, 1)
|
|
}
|
|
|
|
func TestRollupTemplateUsageStats(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
|
|
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
|
|
|
|
anHourAgo := dbtime.Now().Add(-time.Hour).Truncate(time.Hour).UTC()
|
|
anHourAndSixMonthsAgo := anHourAgo.AddDate(0, -6, 0).UTC()
|
|
|
|
var (
|
|
org = dbgen.Organization(t, db, database.Organization{})
|
|
user = dbgen.User(t, db, database.User{Name: "user1"})
|
|
tpl = dbgen.Template(t, db, database.Template{OrganizationID: org.ID, CreatedBy: user.ID})
|
|
ver = dbgen.TemplateVersion(t, db, database.TemplateVersion{OrganizationID: org.ID, TemplateID: uuid.NullUUID{UUID: tpl.ID, Valid: true}, CreatedBy: user.ID})
|
|
ws = dbgen.Workspace(t, db, database.WorkspaceTable{OrganizationID: org.ID, TemplateID: tpl.ID, OwnerID: user.ID})
|
|
job = dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{OrganizationID: org.ID})
|
|
build = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{WorkspaceID: ws.ID, JobID: job.ID, TemplateVersionID: ver.ID})
|
|
res = dbgen.WorkspaceResource(t, db, database.WorkspaceResource{JobID: build.JobID})
|
|
agent = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ResourceID: res.ID})
|
|
app = dbgen.WorkspaceApp(t, db, database.WorkspaceApp{AgentID: agent.ID})
|
|
)
|
|
|
|
// Stats inserted 6 months + 1 day ago, should be excluded.
|
|
_ = dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
|
|
TemplateID: tpl.ID,
|
|
WorkspaceID: ws.ID,
|
|
AgentID: agent.ID,
|
|
UserID: user.ID,
|
|
CreatedAt: anHourAndSixMonthsAgo.AddDate(0, 0, -1),
|
|
ConnectionMedianLatencyMS: 1,
|
|
ConnectionCount: 1,
|
|
SessionCountSSH: 1,
|
|
})
|
|
_ = dbgen.WorkspaceAppStat(t, db, database.WorkspaceAppStat{
|
|
UserID: user.ID,
|
|
WorkspaceID: ws.ID,
|
|
AgentID: agent.ID,
|
|
SessionStartedAt: anHourAndSixMonthsAgo.AddDate(0, 0, -1),
|
|
SessionEndedAt: anHourAndSixMonthsAgo.AddDate(0, 0, -1).Add(time.Minute),
|
|
SlugOrPort: app.Slug,
|
|
})
|
|
|
|
// Stats inserted 6 months - 1 day ago, should be rolled up.
|
|
wags1 := dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
|
|
TemplateID: tpl.ID,
|
|
WorkspaceID: ws.ID,
|
|
AgentID: agent.ID,
|
|
UserID: user.ID,
|
|
CreatedAt: anHourAndSixMonthsAgo.AddDate(0, 0, 1),
|
|
ConnectionMedianLatencyMS: 1,
|
|
ConnectionCount: 1,
|
|
SessionCountReconnectingPTY: 1,
|
|
})
|
|
wags2 := dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
|
|
TemplateID: tpl.ID,
|
|
WorkspaceID: ws.ID,
|
|
AgentID: agent.ID,
|
|
UserID: user.ID,
|
|
CreatedAt: wags1.CreatedAt.Add(time.Minute),
|
|
ConnectionMedianLatencyMS: 1,
|
|
ConnectionCount: 1,
|
|
SessionCountReconnectingPTY: 1,
|
|
})
|
|
// wags2 and waps1 overlap, so total usage is 4 - 1.
|
|
waps1 := dbgen.WorkspaceAppStat(t, db, database.WorkspaceAppStat{
|
|
UserID: user.ID,
|
|
WorkspaceID: ws.ID,
|
|
AgentID: agent.ID,
|
|
SessionStartedAt: wags2.CreatedAt,
|
|
SessionEndedAt: wags2.CreatedAt.Add(time.Minute),
|
|
SlugOrPort: app.Slug,
|
|
})
|
|
waps2 := dbgen.WorkspaceAppStat(t, db, database.WorkspaceAppStat{
|
|
UserID: user.ID,
|
|
WorkspaceID: ws.ID,
|
|
AgentID: agent.ID,
|
|
SessionStartedAt: waps1.SessionEndedAt,
|
|
SessionEndedAt: waps1.SessionEndedAt.Add(time.Minute),
|
|
SlugOrPort: app.Slug,
|
|
})
|
|
_ = waps2 // Keep the name for documentation.
|
|
|
|
// The data is already present, so we can rely on initial rollup to occur.
|
|
events := make(chan dbrollup.Event, 1)
|
|
rolluper := dbrollup.New(logger, db, dbrollup.WithInterval(250*time.Millisecond), dbrollup.WithEventChannel(events))
|
|
defer rolluper.Close()
|
|
|
|
<-events // Deplete init event, resume operation.
|
|
|
|
ctx := testutil.Context(t, testutil.WaitMedium)
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
t.Fatal("timed out waiting for rollup to occur")
|
|
case ev := <-events:
|
|
require.True(t, ev.TemplateUsageStats, "expected template usage stats to be rolled up")
|
|
}
|
|
|
|
stats, err := db.GetTemplateUsageStats(ctx, database.GetTemplateUsageStatsParams{
|
|
StartTime: anHourAndSixMonthsAgo.Add(-time.Minute),
|
|
EndTime: anHourAgo,
|
|
})
|
|
require.NoError(t, err)
|
|
require.Len(t, stats, 1)
|
|
|
|
// I do not know a better way to do this. Our database runs in a *random*
|
|
// timezone. So the returned time is in a random timezone and fails on the
|
|
// equal even though they are the same time if converted back to the same timezone.
|
|
stats[0].EndTime = stats[0].EndTime.UTC()
|
|
stats[0].StartTime = stats[0].StartTime.UTC()
|
|
|
|
require.Equal(t, database.TemplateUsageStat{
|
|
TemplateID: tpl.ID,
|
|
UserID: user.ID,
|
|
StartTime: wags1.CreatedAt,
|
|
EndTime: wags1.CreatedAt.Add(30 * time.Minute),
|
|
MedianLatencyMs: sql.NullFloat64{Float64: 1, Valid: true},
|
|
UsageMins: 3,
|
|
ReconnectingPtyMins: 2,
|
|
AppUsageMins: database.StringMapOfInt{
|
|
app.Slug: 2,
|
|
},
|
|
}, stats[0])
|
|
}
|