stop accessing db directly, only query workspaces with autostart enabled

coder · johnstcn · May 11, 2022 · Apr 19, 2022 · Apr 26, 2022 · Apr 26, 2022
commit cfd0d1e521eb00c4ef0985fb6fb1bfdb74a26b78
diff --git a/coderd/autostart/lifecycle/lifecycle_executor.go b/coderd/autostart/lifecycle/lifecycle_executor.go
@@ -53,16 +53,12 @@ func (e *Executor) Run() error {
 func (e *Executor) runOnce(t time.Time) error {
 	currentTick := t.Round(time.Minute)
 	return e.db.InTx(func(db database.Store) error {
-		allWorkspaces, err := db.GetWorkspaces(e.ctx)
+		autostartWorkspaces, err := db.GetWorkspacesAutostart(e.ctx)
 		if err != nil {
 			return xerrors.Errorf("get all workspaces: %w", err)
 		}
 
-		for _, ws := range allWorkspaces {
-			// We only care about workspaces with autostart enabled.
-			if ws.AutostartSchedule.String == "" {
-				continue
-			}
+		for _, ws := range autostartWorkspaces {
 			sched, err := schedule.Weekly(ws.AutostartSchedule.String)
 			if err != nil {
 				e.log.Warn(e.ctx, "workspace has invalid autostart schedule",
@@ -73,7 +69,6 @@ func (e *Executor) runOnce(t time.Time) error {
 			}
 
 			// Determine the workspace state based on its latest build. We expect it to be stopped.
-			// TODO(cian): is this **guaranteed** to be the latest build???
 			latestBuild, err := db.GetWorkspaceBuildByWorkspaceIDWithoutAfter(e.ctx, ws.ID)
 			if err != nil {
 				return xerrors.Errorf("get latest build for workspace %q: %w", ws.ID, err)

diff --git a/coderd/autostart/lifecycle/lifecycle_executor_test.go b/coderd/autostart/lifecycle/lifecycle_executor_test.go
@@ -5,14 +5,9 @@ import (
 	"testing"
 	"time"
 
-	"cdr.dev/slog"
-	"cdr.dev/slog/sloggers/slogtest"
-
-	"github.com/coder/coder/coderd/autostart/lifecycle"
 	"github.com/coder/coder/coderd/autostart/schedule"
 	"github.com/coder/coder/coderd/coderdtest"
 	"github.com/coder/coder/coderd/database"
-	"github.com/coder/coder/coderd/database/databasefake"
 	"github.com/coder/coder/codersdk"
 
 	"github.com/stretchr/testify/require"
@@ -25,15 +20,11 @@ func Test_Executor_Run(t *testing.T) {
 		t.Parallel()
 
 		var (
-			ctx               = context.Background()
-			cancelCtx, cancel = context.WithCancel(context.Background())
-			log               = slogtest.Make(t, nil).Named("lifecycle.executor").Leveled(slog.LevelDebug)
-			err               error
-			tickCh            = make(chan time.Time)
-			db                = databasefake.New()
-			le                = lifecycle.NewExecutor(cancelCtx, db, log, tickCh)
-			client            = coderdtest.New(t, &coderdtest.Options{
-				Ticker: tickCh,
+			ctx    = context.Background()
+			err    error
+			tickCh = make(chan time.Time)
+			client = coderdtest.New(t, &coderdtest.Options{
+				LifecycleTicker: tickCh,
 			})
 			// Given: we have a user with a workspace
 			_         = coderdtest.NewProvisionerDaemon(t, client)
@@ -71,31 +62,25 @@ func Test_Executor_Run(t *testing.T) {
 		// When: the lifecycle executor ticks
 		go func() {
 			tickCh <- time.Now().UTC().Add(time.Minute)
-			cancel()
 		}()
-		require.NoError(t, le.Run())
 
 		// Then: the workspace should be started
 		require.Eventually(t, func() bool {
 			ws := coderdtest.MustWorkspace(t, client, workspace.ID)
 			return ws.LatestBuild.Job.Status == codersdk.ProvisionerJobSucceeded &&
 				ws.LatestBuild.Transition == database.WorkspaceTransitionStart
-		}, 10*time.Second, 1000*time.Millisecond)
+		}, 5*time.Second, 250*time.Millisecond)
 	})
 
 	t.Run("AlreadyRunning", func(t *testing.T) {
 		t.Parallel()
 
 		var (
-			ctx               = context.Background()
-			cancelCtx, cancel = context.WithCancel(context.Background())
-			log               = slogtest.Make(t, nil).Named("lifecycle.executor").Leveled(slog.LevelDebug)
-			err               error
-			tickCh            = make(chan time.Time)
-			db                = databasefake.New()
-			le                = lifecycle.NewExecutor(cancelCtx, db, log, tickCh)
-			client            = coderdtest.New(t, &coderdtest.Options{
-				Ticker: tickCh,
+			ctx    = context.Background()
+			err    error
+			tickCh = make(chan time.Time)
+			client = coderdtest.New(t, &coderdtest.Options{
+				LifecycleTicker: tickCh,
 			})
 			// Given: we have a user with a workspace
 			_         = coderdtest.NewProvisionerDaemon(t, client)
@@ -123,14 +108,48 @@ func Test_Executor_Run(t *testing.T) {
 		// When: the lifecycle executor ticks
 		go func() {
 			tickCh <- time.Now().UTC().Add(time.Minute)
-			cancel()
 		}()
-		require.NoError(t, le.Run())
 
 		// Then: the workspace should not be started.
 		require.Never(t, func() bool {
 			ws := coderdtest.MustWorkspace(t, client, workspace.ID)
 			return ws.LatestBuild.ID != workspace.LatestBuild.ID
-		}, 10*time.Second, 1000*time.Millisecond)
+		}, 5*time.Second, 250*time.Millisecond)
+	})
+
+	t.Run("NotEnabled", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			tickCh = make(chan time.Time)
+			client = coderdtest.New(t, &coderdtest.Options{
+				LifecycleTicker: tickCh,
+			})
+			// Given: we have a user with a workspace
+			_         = coderdtest.NewProvisionerDaemon(t, client)
+			user      = coderdtest.CreateFirstUser(t, client)
+			version   = coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
+			template  = coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+			_         = coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+			workspace = coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+			_         = coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+		)
+
+		// Given: we ensure the workspace is now in a stopped state
+		require.Equal(t, database.WorkspaceTransitionStart, workspace.LatestBuild.Transition)
+
+		// Given: the workspace has autostart disabled
+		require.Empty(t, workspace.AutostartSchedule)
+
+		// When: the lifecycle executor ticks
+		go func() {
+			tickCh <- time.Now().UTC().Add(time.Minute)
+		}()
+
+		// Then: the workspace should not be started.
+		require.Never(t, func() bool {
+			ws := coderdtest.MustWorkspace(t, client, workspace.ID)
+			return ws.LatestBuild.ID != workspace.LatestBuild.ID
+		}, 5*time.Second, 250*time.Millisecond)
 	})
 }
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -59,7 +59,7 @@ type Options struct {
 	GoogleTokenValidator *idtoken.Validator
 	SSHKeygenAlgorithm   gitsshkey.Algorithm
 	APIRateLimit         int
-	Ticker               <-chan time.Time
+	LifecycleTicker      <-chan time.Time
 }
 
 // New constructs an in-memory coderd instance and returns
@@ -75,10 +75,10 @@ func New(t *testing.T, options *Options) *codersdk.Client {
 		options.GoogleTokenValidator, err = idtoken.NewValidator(ctx, option.WithoutAuthentication())
 		require.NoError(t, err)
 	}
-	if options.Ticker == nil {
-		ticker := time.NewTicker(time.Second)
-		options.Ticker = ticker.C
-		t.Cleanup(ticker.Stop)
+	if options.LifecycleTicker == nil {
+		ticker := make(chan time.Time)
+		options.LifecycleTicker = ticker
+		t.Cleanup(func() { close(ticker) })
 	}
 
 	// This can be hotswapped for a live database instance.
@@ -109,7 +109,7 @@ func New(t *testing.T, options *Options) *codersdk.Client {
 		ctx,
 		db,
 		slogtest.Make(t, nil).Named("lifecycle.executor").Leveled(slog.LevelDebug),
-		options.Ticker,
+		options.LifecycleTicker,
 	)
 	go lifecycleExecutor.Run()
 

diff --git a/coderd/database/databasefake/databasefake.go b/coderd/database/databasefake/databasefake.go
@@ -333,6 +333,18 @@ func (q *fakeQuerier) GetWorkspaces(_ context.Context) ([]database.Workspace, er
 	return workspaces, nil
 }
 
+func (q *fakeQuerier) GetWorkspacesAutostart(_ context.Context) ([]database.Workspace, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+	workspaces := make([]database.Workspace, 0)
+	for _, ws := range q.workspaces {
+		if ws.AutostartSchedule.String != "" {
+			workspaces = append(workspaces, ws)
+		}
+	}
+	return workspaces, nil
+}
+
 func (q *fakeQuerier) GetWorkspaceOwnerCountsByTemplateIDs(_ context.Context, templateIDs []uuid.UUID) ([]database.GetWorkspaceOwnerCountsByTemplateIDsRow, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()

diff --git a/coderd/database/models.go b/coderd/database/models.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaces.sql b/coderd/database/queries/workspaces.sql
@@ -22,6 +22,17 @@ FROM
 WHERE
 	deleted = false;
 
+-- name: GetWorkspacesAutostart :many
+SELECT
+	*
+FROM
+	workspaces
+WHERE
+	deleted = false
+AND
+	autostart_schedule <> ''
+;
+
 -- name: GetWorkspacesByTemplateID :many
 SELECT
 	*