add prebuild controller tests

coderd/database/querier_test.go

@@ -2916,210 +2916,6 @@ func TestGetUserStatusCounts(t *testing.T) {
 	}
 }
 
-func TestGetRunningPrebuilds(t *testing.T) {
-	// For each preset that requests a number of prebuilds, we need to determine how many are already running.
-	// so that we can create any additional or destroy any extraneous prebuilds.
-	
-	// Note that a running prebuild is not necessarily eligible to be claimed by a user creating a new workspace.
-	// Rather, the concept of a running prebuild is the same as a running workspace. It means that certain potentially
-	// costly ephemeral resources have been provisioned and have not yet been destroyed. This concept is used to 
-	// reconcile the desired number of prebuilds against the number of prebuilds that are consuming certain potentially 
-	// costly ephemeral resources. Even if a prebuild is not ready to be claimed, it is still running if it consumes 
-	// these resources.
-	
-	// This query finds all running prebuilds in a single transaction.
-	// It is the caller's responsibility to filter these by preset if that is required.
-	// Prebuilds are defined as workspaces that belong to the well known prebuilds user.
-
-	t.Parallel()
-
-	if !dbtestutil.WillUsePostgres() {
-		t.SkipNow()
-	}
-
-	db, _ := dbtestutil.NewDB(t)
-	ctx := testutil.Context(t, testutil.WaitShort)
-
-	// given there are no prebuilds
-	// when GetRunningPrebuilds is called
-	prebuilds, err := db.GetRunningPrebuilds(ctx)
-	// then the query runs successfully
-	require.NoError(t, err)
-	// and the query finds no prebuilds
-	require.Empty(t, prebuilds)
-
-	// given there are prebuilds
-	// but none are running:
-	// * one is not running because its latest build was a stop transition
-	// * another is not running because its latest build was a delete transition
-	// * a third is not running because its latest build was a start transition but the build failed
-	// * a fourth is not running because its latest build was a start transition but the build was canceled
-	// when GetRunningPrebuilds is called
-	prebuilds, err = db.GetRunningPrebuilds(ctx)
-	// then the query runs successfully
-	require.NoError(t, err)
-	// and the query finds no prebuilds
-	// because stopped, deleted and failed builds are not considered running in terms of the definition of "running" above.
-	require.Empty(t, prebuilds)
-
-	// given there are running prebuilds
-	// * one is running because its latest build was a start transition
-	// * another is running because its latest build attempted to stop it, but it failed
-	// * a third is running because its latest build attempted to stop it, but was canceled
-	// when GetRunningPrebuilds is called
-	prebuilds, err = db.GetRunningPrebuilds(ctx)
-	// then the query runs successfully
-	require.NoError(t, err)
-	// and the query finds the prebuilds that are running
-	require.NotEmpty(t, prebuilds)
-
-}
-
-func TestGetTemplatePresetsWithPrebuilds(t *testing.T) {
-	// GetTemplatePresetsWithPrebuilds returns a list of template version presets that define prebuilds.
-	// It is used in the prebuild reconciliation logic to establish the outer loop over prebuilds,
-	// so that each prebuild can be reconciled in turn. It has an optional template ID filter.
-	t.Parallel()
-	if !dbtestutil.WillUsePostgres() {
-		t.SkipNow()
-	}
-
-	t.Run("searching for all presets with prebuilds", func(t *testing.T) {
-		// ie. given a null template ID parameter
-
-		t.Parallel()
-
-		db, _ := dbtestutil.NewDB(t)
-		ctx := testutil.Context(t, testutil.WaitShort)
-
-		// given there are no presets
-		// when GetTemplatePresetsWithPrebuilds is called
-		presets, err := db.GetTemplatePresetsWithPrebuilds(ctx, uuid.NullUUID{})
-		// then the query runs successfully
-		require.NoError(t, err)
-		// but the query finds no presets
-		// because without presets, there are no prebuilds
-		// and without prebuilds, there is nothing to reconcile
-		require.Empty(t, presets)
-
-		// given there are presets, but no prebuilds
-		// TODO (sasswart): db setup
-
-		// when GetTemplatePresetsWithPrebuilds is called
-		presets, err = db.GetTemplatePresetsWithPrebuilds(ctx, uuid.NullUUID{})
-		// then the query runs successfully
-		require.NoError(t, err)
-		// but the query finds no presets
-		// because without prebuilds, there is nothing to reconcile
-		// even if there are presets
-		require.Empty(t, presets)
-
-		// given there are presets with prebuilds
-		// TODO (sasswart): db setup
-
-		// and there are also presets without prebuilds
-		// TODO (sasswart): db setup
-
-		// when GetTemplatePresetsWithPrebuilds is called
-		presets, err = db.GetTemplatePresetsWithPrebuilds(ctx, uuid.NullUUID{})
-		// then the query runs successfully
-		require.NoError(t, err)
-		// and the query finds the presets that have prebuilds
-		// but not the presets without prebuilds
-		// because we only want to reconcile prebuilds
-		// so we have no interest in presets without prebuilds
-		require.NotEmpty(t, presets)
-	})
-
-	t.Run("searching for presets with prebuilds for a specific template", func(t *testing.T) {
-		// ie. given a specific template ID as input parameter
-		t.Parallel()
-
-		db, _ := dbtestutil.NewDB(t)
-		ctx := testutil.Context(t, testutil.WaitShort)
-
-		// given there are no presets for the template in question
-		// TODO (sasswart): db setup
-		var templateID uuid.NullUUID
-
-		// for all tests below:
-		// given there are presets with prebuilds for unrelated templates
-		// (these should never be returned if we're filtering by a specific template)
-		// TODO (sasswart): db setup
-
-		// when GetTemplatePresetsWithPrebuilds is called with the template ID
-		presets, err := db.GetTemplatePresetsWithPrebuilds(ctx, templateID)
-		// then the query runs successfully
-		require.NoError(t, err)
-		// and the query doesn't find any presets, because the template we want doesn't have any even if other templates do
-		require.Empty(t, presets)
-
-		// given there are presets for the template in question, but they don't define any prebuilds
-		// TODO (sasswart): db setup
-
-		// when GetTemplatePresetsWithPrebuilds is called with the template ID
-		presets, err = db.GetTemplatePresetsWithPrebuilds(ctx, templateID)
-		// then the query runs successfully
-		require.NoError(t, err)
-		// and the query doesn't find any presets, because the template we want doesn't have any prebuilds even if other templates do
-		require.Empty(t, presets)
-
-		// given there are presets for the template in question, and they do define prebuilds
-		// TODO (sasswart): db setup
-
-		// when GetTemplatePresetsWithPrebuilds is called with the template ID
-		presets, err = db.GetTemplatePresetsWithPrebuilds(ctx, templateID)
-		// then the query runs successfully
-		require.NoError(t, err)
-		// and the query finds the presets that have prebuilds
-		// but not the presets without prebuilds
-		// nor the presets for other templates
-		require.NotEmpty(t, presets)
-	})
-
-	t.Run("presets from an inactive template version are still returned", func(t *testing.T) {
-		// eg. a new template version was pushed and promoted to active, replacing the previous active version
-		// We still want to find these presets because the old active version may still contain running prebuilds that need to be destroyed
-		t.Parallel()
-
-		db, _ := dbtestutil.NewDB(t)
-		ctx := testutil.Context(t, testutil.WaitShort)
-
-		// given an inactive template version
-		// TODO (sasswart): db setup
-		var templateID uuid.NullUUID
-
-		// when GetTemplatePresetsWithPrebuilds is called with the template ID
-		presets, err := db.GetTemplatePresetsWithPrebuilds(ctx, templateID)
-		// then the query runs successfully
-		require.NoError(t, err)
-		// and the query finds the presets that have prebuilds
-		// because the template version is not the active version
-		require.NotEmpty(t, presets)
-	})
-}
-
-func TestGetPrebuildsInProgress(t *testing.T) {
-	// When we want to ensure that the number of prebuilds matches the desired number,
-	// we consider all running prebuilds from GetRunningPrebuilds as tested in TestGetRunningPrebuilds.
-	// However, we also need to consider prebuilds that have already been queued but are not yet running,
-	// otherwise we might queue, and therefore eventually provision, too many prebuilds.
-	// This would cost users money, and we want to avoid that.
-	//
-	// GetPrebuildsInProgress returns a list of prebuilds that have been queued but are not yet running.
-	// It fills the potential gap between the number of running prebuilds and the desired number of prebuilds
-	// that may exist because of previously enqueued prebuilds that have not yet been provisioned.
-	t.Parallel()
-
-	if !dbtestutil.WillUsePostgres() {
-		t.SkipNow()
-	}
-
-	db, _ := dbtestutil.NewDB(t)
-	ctx := testutil.Context(t, testutil.WaitShort)
-
-}
-
 func requireUsersMatch(t testing.TB, expected []database.User, found []database.GetUsersRow, msg string) {
 	t.Helper()
 	require.ElementsMatch(t, expected, database.ConvertUserRows(found), msg)

enterprise/coderd/prebuilds/controller.go

@@ -52,7 +52,11 @@ func NewController(store database.Store, pubsub pubsub.Pubsub, cfg codersdk.Preb
 }
 
 func (c *Controller) Loop(ctx context.Context) error {
-	ticker := time.NewTicker(c.cfg.ReconciliationInterval.Value())
+	reconciliationInterval := c.cfg.ReconciliationInterval.Value()
+	if reconciliationInterval <= 0 { // avoids a panic
+		reconciliationInterval = 5 * time.Minute
+	}
+	ticker := time.NewTicker(reconciliationInterval)
 	defer ticker.Stop()
 
 	// TODO: create new authz role
@@ -88,9 +92,9 @@ func (c *Controller) isClosed() bool {
 	return c.closed.Load()
 }
 
-func (c *Controller) ReconcileTemplate(templateID uuid.UUID) {
+func (c *Controller) ReconcileTemplate(templateID *uuid.UUID) {
 	// TODO: replace this with pubsub listening
-	c.nudgeCh <- &templateID
+	c.nudgeCh <- templateID
 }
 
 // reconcile will attempt to resolve the desired vs actual state of all templates which have presets with prebuilds configured.
@@ -146,13 +150,7 @@ func (c *Controller) reconcile(ctx context.Context, templateID *uuid.UUID) {
 
 		logger.Debug(ctx, "acquired top-level prebuilds reconciliation lock", slog.F("acquire_wait_secs", fmt.Sprintf("%.4f", time.Since(start).Seconds())))
 
-		var id uuid.NullUUID
+		state, err := c.determineState(ctx, db, templateID)
-		if templateID != nil {
-			id.UUID = *templateID
-			id.Valid = true
-		}
-
-		state, err := c.determineState(ctx, db, id)
 		if err != nil {
 			return xerrors.Errorf("determine current state: %w", err)
 		}
@@ -200,7 +198,7 @@ func (c *Controller) reconcile(ctx context.Context, templateID *uuid.UUID) {
 
 // determineState determines the current state of prebuilds & the presets which define them.
 // An application-level lock is used
-func (c *Controller) determineState(ctx context.Context, store database.Store, id uuid.NullUUID) (*reconciliationState, error) {
+func (c *Controller) determineState(ctx context.Context, store database.Store, templateId *uuid.UUID) (*reconciliationState, error) {
 	if err := ctx.Err(); err != nil {
 		return nil, err
 	}
@@ -218,7 +216,13 @@ func (c *Controller) determineState(ctx context.Context, store database.Store, i
 
 		c.logger.Debug(ctx, "acquired state determination lock", slog.F("acquire_wait_secs", fmt.Sprintf("%.4f", time.Since(start).Seconds())))
 
-		presetsWithPrebuilds, err := db.GetTemplatePresetsWithPrebuilds(ctx, id)
+		var dbTemplateID uuid.NullUUID
+		if templateId != nil {
+			dbTemplateID.UUID = *templateId
+			dbTemplateID.Valid = true
+		}
+
+		presetsWithPrebuilds, err := db.GetTemplatePresetsWithPrebuilds(ctx, dbTemplateID)
 		if len(presetsWithPrebuilds) == 0 {
 			return nil
 		}

enterprise/coderd/prebuilds/controller_test.go

@@ -0,0 +1,297 @@
+package prebuilds
+
+import (
+	"database/sql"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbgen"
+	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/testutil"
+	"github.com/coder/serpent"
+)
+
+func TestNoReconciliationActionsIfNoPresets(t *testing.T) {
+	// Scenario: No reconciliation actions are taken if there are no presets
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	db, pubsub := dbtestutil.NewDB(t)
+	cfg := codersdk.PrebuildsConfig{
+		ReconciliationInterval: serpent.Duration(testutil.WaitLong),
+	}
+	logger := testutil.Logger(t)
+	controller := NewController(db, pubsub, cfg, logger)
+
+	// given a template version with no presets
+	org := dbgen.Organization(t, db, database.Organization{})
+	user := dbgen.User(t, db, database.User{})
+	template := dbgen.Template(t, db, database.Template{
+		CreatedBy:      user.ID,
+		OrganizationID: org.ID,
+	})
+	templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
+		TemplateID:     uuid.NullUUID{UUID: template.ID, Valid: true},
+		OrganizationID: org.ID,
+		CreatedBy:      user.ID,
+	})
+	// verify that the db state is correct
+	gotTemplateVersion, err := db.GetTemplateVersionByID(ctx, templateVersion.ID)
+	require.NoError(t, err)
+	require.Equal(t, templateVersion, gotTemplateVersion)
+
+	// when we trigger the reconciliation loop for all templates
+	controller.reconcile(ctx, nil)
+
+	// then no reconciliation actions are taken
+	// because without presets, there are no prebuilds
+	// and without prebuilds, there is nothing to reconcile
+	jobs, err := db.GetProvisionerJobsCreatedAfter(ctx, time.Now().Add(-time.Hour))
+	require.NoError(t, err)
+	require.Empty(t, jobs)
+}
+
+func TestNoReconciliationActionsIfNoPrebuilds(t *testing.T) {
+	// Scenario: No reconciliation actions are taken if there are no prebuilds
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	db, pubsub := dbtestutil.NewDB(t)
+	cfg := codersdk.PrebuildsConfig{
+		ReconciliationInterval: serpent.Duration(testutil.WaitLong),
+	}
+	logger := testutil.Logger(t)
+	controller := NewController(db, pubsub, cfg, logger)
+
+	// given there are presets, but no prebuilds
+	org := dbgen.Organization(t, db, database.Organization{})
+	user := dbgen.User(t, db, database.User{})
+	template := dbgen.Template(t, db, database.Template{
+		CreatedBy:      user.ID,
+		OrganizationID: org.ID,
+	})
+	templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
+		TemplateID:     uuid.NullUUID{UUID: template.ID, Valid: true},
+		OrganizationID: org.ID,
+		CreatedBy:      user.ID,
+	})
+	preset, err := db.InsertPreset(ctx, database.InsertPresetParams{
+		TemplateVersionID: templateVersion.ID,
+		Name:              "test",
+	})
+	require.NoError(t, err)
+	_, err = db.InsertPresetParameters(ctx, database.InsertPresetParametersParams{
+		TemplateVersionPresetID: preset.ID,
+		Names:                   []string{"test"},
+		Values:                  []string{"test"},
+	})
+	require.NoError(t, err)
+
+	// verify that the db state is correct
+	presetParameters, err := db.GetPresetParametersByTemplateVersionID(ctx, templateVersion.ID)
+	require.NoError(t, err)
+	require.NotEmpty(t, presetParameters)
+
+	// when we trigger the reconciliation loop for all templates
+	controller.reconcile(ctx, nil)
+
+	// then no reconciliation actions are taken
+	// because without prebuilds, there is nothing to reconcile
+	// even if there are presets
+	jobs, err := db.GetProvisionerJobsCreatedAfter(ctx, time.Now().Add(-time.Hour))
+	require.NoError(t, err)
+	require.Empty(t, jobs)
+}
+
+func TestPrebuildCreation(t *testing.T) {
+	t.Parallel()
+
+	// Scenario: Prebuilds are created if and only if they are needed
+	type testCase struct {
+		name                    string
+		prebuildStatus          database.WorkspaceStatus
+		shouldCreateNewPrebuild bool
+	}
+
+	testCases := []testCase{
+		{
+			name:                    "running prebuild",
+			prebuildStatus:          database.WorkspaceStatusRunning,
+			shouldCreateNewPrebuild: false,
+		},
+		{
+			name:                    "stopped prebuild",
+			prebuildStatus:          database.WorkspaceStatusStopped,
+			shouldCreateNewPrebuild: true,
+		},
+		{
+			name:                    "failed prebuild",
+			prebuildStatus:          database.WorkspaceStatusFailed,
+			shouldCreateNewPrebuild: true,
+		},
+		{
+			name:                    "canceled prebuild",
+			prebuildStatus:          database.WorkspaceStatusCanceled,
+			shouldCreateNewPrebuild: true,
+		},
+		{
+			name:                    "deleted prebuild",
+			prebuildStatus:          database.WorkspaceStatusDeleted,
+			shouldCreateNewPrebuild: true,
+		},
+		{
+			name:                    "pending prebuild",
+			prebuildStatus:          database.WorkspaceStatusPending,
+			shouldCreateNewPrebuild: false,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			ctx := testutil.Context(t, testutil.WaitShort)
+			db, pubsub := dbtestutil.NewDB(t)
+			cfg := codersdk.PrebuildsConfig{}
+			logger := testutil.Logger(t)
+			controller := NewController(db, pubsub, cfg, logger)
+
+			// given a user
+			org := dbgen.Organization(t, db, database.Organization{})
+			user := dbgen.User(t, db, database.User{})
+
+			template := dbgen.Template(t, db, database.Template{
+				CreatedBy:      user.ID,
+				OrganizationID: org.ID,
+			})
+			templateVersionJob := dbgen.ProvisionerJob(t, db, pubsub, database.ProvisionerJob{
+				ID:             uuid.New(),
+				CreatedAt:      time.Now().Add(-2 * time.Hour),
+				CompletedAt:    sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true},
+				OrganizationID: org.ID,
+				InitiatorID:    user.ID,
+			})
+			templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
+				TemplateID:     uuid.NullUUID{UUID: template.ID, Valid: true},
+				OrganizationID: org.ID,
+				CreatedBy:      user.ID,
+				JobID:          templateVersionJob.ID,
+			})
+			db.UpdateTemplateActiveVersionByID(ctx, database.UpdateTemplateActiveVersionByIDParams{
+				ID:              template.ID,
+				ActiveVersionID: templateVersion.ID,
+			})
+			preset, err := db.InsertPreset(ctx, database.InsertPresetParams{
+				TemplateVersionID: templateVersion.ID,
+				Name:              "test",
+			})
+			require.NoError(t, err)
+			_, err = db.InsertPresetParameters(ctx, database.InsertPresetParametersParams{
+				TemplateVersionPresetID: preset.ID,
+				Names:                   []string{"test"},
+				Values:                  []string{"test"},
+			})
+			require.NoError(t, err)
+			_, err = db.InsertPresetPrebuild(ctx, database.InsertPresetPrebuildParams{
+				ID:               uuid.New(),
+				PresetID:         preset.ID,
+				DesiredInstances: 1,
+			})
+			require.NoError(t, err)
+			completedAt := sql.NullTime{}
+			cancelledAt := sql.NullTime{}
+			transition := database.WorkspaceTransitionStart
+			deleted := false
+			buildError := sql.NullString{}
+			switch tc.prebuildStatus {
+			case database.WorkspaceStatusRunning:
+				completedAt = sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}
+			case database.WorkspaceStatusStopped:
+				completedAt = sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}
+				transition = database.WorkspaceTransitionStop
+			case database.WorkspaceStatusFailed:
+				completedAt = sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}
+				buildError = sql.NullString{String: "build failed", Valid: true}
+			case database.WorkspaceStatusCanceled:
+				completedAt = sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}
+				cancelledAt = sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}
+			case database.WorkspaceStatusDeleted:
+				completedAt = sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}
+				transition = database.WorkspaceTransitionDelete
+				deleted = true
+			case database.WorkspaceStatusPending:
+				completedAt = sql.NullTime{}
+				transition = database.WorkspaceTransitionStart
+			default:
+			}
+
+			workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
+				TemplateID:     template.ID,
+				OrganizationID: org.ID,
+				OwnerID:        OwnerID,
+				Deleted:        deleted,
+			})
+			job := dbgen.ProvisionerJob(t, db, pubsub, database.ProvisionerJob{
+				InitiatorID:    OwnerID,
+				CreatedAt:      time.Now().Add(-2 * time.Hour),
+				CompletedAt:    completedAt,
+				CanceledAt:     cancelledAt,
+				OrganizationID: org.ID,
+				Error:          buildError,
+			})
+			dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+				WorkspaceID:             workspace.ID,
+				InitiatorID:             OwnerID,
+				TemplateVersionID:       templateVersion.ID,
+				JobID:                   job.ID,
+				TemplateVersionPresetID: uuid.NullUUID{UUID: preset.ID, Valid: true},
+				Transition:              transition,
+			})
+
+			controller.reconcile(ctx, nil)
+			jobs, err := db.GetProvisionerJobsCreatedAfter(ctx, time.Now().Add(-time.Hour))
+			require.NoError(t, err)
+			require.Equal(t, tc.shouldCreateNewPrebuild, len(jobs) == 1)
+		})
+	}
+}
+
+func TestDeleteUnwantedPrebuilds(t *testing.T) {
+	// Scenario: Prebuilds are deleted if and only if they are extraneous
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitShort)
+	db, pubsub := dbtestutil.NewDB(t)
+	cfg := codersdk.PrebuildsConfig{}
+	logger := testutil.Logger(t)
+	controller := NewController(db, pubsub, cfg, logger)
+
+	// when does a prebuild get deleted?
+	// * when it is in some way permanently ineligible to be claimed
+	//   * this could be because the build failed or was canceled
+	//   * or it belongs to a template version that is no longer active
+	//   * or it belongs to a template version that is deprecated
+	// * when there are more prebuilds than the preset desires
+	//   * someone could have manually created a workspace for the prebuild user
+	// * any workspaces that were created for the prebuilds user and don't match a preset should be deleted - deferred
+
+	// given a preset that desires 2 prebuilds
+	// and there are 3 running prebuilds for the preset
+	// and there are 4 non-running prebuilds for the preset
+	// * one is not running because its latest build was a stop transition
+	// * another is not running because its latest build was a delete transition
+	// * a third is not running because its latest build was a start transition but the build failed
+	// * a fourth is not running because its latest build was a start transition but the build was canceled
+	// when we trigger the reconciliation loop for all templates
+	controller.reconcile(ctx, nil)
+	// then the four non running prebuilds are deleted
+	// and 1 of the running prebuilds is deleted
+	// because stopped, deleted and failed builds are not considered running in terms of the definition of "running" above.
+}
+
+// TODO (sasswart): test claim (success, fail) x (eligible)
+// TODO (sasswart): test idempotency of reconciliation
+// TODO (sasswart): test mutual exclusion