coder/coderd/database/dbrollup/dbrollup_test.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"
	"cdr.dev/slog/sloggers/slogtest"

	"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/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)
}

func TestRollup_Close(t *testing.T) {
	t.Parallel()
	rolluper := dbrollup.New(slogtest.Make(t, nil), dbmem.New(), 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 *sql.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()

	if !dbtestutil.WillUsePostgres() {
		t.Skip("Skipping test; only works with PostgreSQL.")
	}

	db, ps := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)

	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])
}