chore: Refactor autostart check, do it in Go

coder · mafredri · Sep 2, 2022 · Aug 29, 2022 · Aug 29, 2022 · Aug 29, 2022
commit 45e4141052ce78272916eebac09a2952e8ab7f73
diff --git a/coderd/autobuild/executor/lifecycle_executor.go b/coderd/autobuild/executor/lifecycle_executor.go
@@ -2,7 +2,6 @@ package executor
 
 import (
 	"context"
-	"database/sql"
 	"encoding/json"
 	"time"
 
@@ -101,36 +100,41 @@ func (e *Executor) runOnce(t time.Time) Stats {
 	// NOTE: If a workspace build is created with a given TTL and then the user either
 	//       changes or unsets the TTL, the deadline for the workspace build will not
 	//       have changed. This behavior is as expected per #2229.
-	eligibleWorkspaces, err := e.db.GetWorkspacesAutostart(e.ctx)
+	workspaces, err := e.db.GetWorkspaces(e.ctx, database.GetWorkspacesParams{
+		Deleted: false,
+	})
 	if err != nil {
-		e.log.Error(e.ctx, "get eligible workspaces for autostart or autostop", slog.Error(err))
+		e.log.Error(e.ctx, "get workspaces for autostart or autostop", slog.Error(err))
 		return stats
 	}
 
+	var eligibleWorkspaceIDs []uuid.UUID
+	for _, ws := range workspaces {
+		if isEligibleForAutoStartStop(ws) {
+			eligibleWorkspaceIDs = append(eligibleWorkspaceIDs, ws.ID)
+		}
+	}
+
 	// We only use errgroup here for convenience of API, not for early
 	// cancellation. This means we only return nil errors in th eg.Go.
 	eg := errgroup.Group{}
 	// Limit the concurrency to avoid overloading the database.
 	eg.SetLimit(10)
 
-	for _, ws := range eligibleWorkspaces {
-		ws := ws
-		log := e.log.With(slog.F("workspace_id", ws.ID))
+	for _, wsID := range eligibleWorkspaceIDs {
+		wsID := wsID
+		log := e.log.With(slog.F("workspace_id", wsID))
 
 		eg.Go(func() error {
 			err := e.db.InTx(func(db database.Store) error {
-				var err error
-
 				// Re-check eligibility since the first check was outside the
 				// transaction and the workspace settings may have changed.
-				ws, err = db.GetWorkspaceAutostart(e.ctx, ws.ID)
+				ws, err := db.GetWorkspaceByID(e.ctx, wsID)
 				if err != nil {
-					// Receiving ErrNoRows means the workspace settings changed
-					// and it is no longer eligible for autostart. Other errors
-					// means something went wrong.
-					if !xerrors.Is(err, sql.ErrNoRows) {
-						log.Error(e.ctx, "get workspace autostart failed", slog.Error(err))
-					}
+					log.Error(e.ctx, "get workspace autostart failed", slog.Error(err))
+					return nil
+				}
+				if !isEligibleForAutoStartStop(ws) {
 					return nil
 				}
 
@@ -191,6 +195,10 @@ func (e *Executor) runOnce(t time.Time) Stats {
 	return stats
 }
 
+func isEligibleForAutoStartStop(ws database.Workspace) bool {
+	return ws.AutostartSchedule.String != "" || ws.Ttl.Int64 > 0
+}
+
 func getNextTransition(
 	ws database.Workspace,
 	priorHistory database.WorkspaceBuild,

diff --git a/coderd/database/databasefake/databasefake.go b/coderd/database/databasefake/databasefake.go
@@ -588,33 +588,6 @@ func (q *fakeQuerier) GetWorkspaceAppsByAgentIDs(_ context.Context, ids []uuid.U
 	return apps, nil
 }
 
-func (q *fakeQuerier) GetWorkspaceAutostart(_ context.Context, workspaceID uuid.UUID) (database.Workspace, error) {
-	q.mutex.RLock()
-	defer q.mutex.RUnlock()
-	for _, ws := range q.workspaces {
-		if ws.ID == workspaceID {
-			if ws.AutostartSchedule.String != "" || ws.Ttl.Valid {
-				return ws, nil
-			}
-		}
-	}
-	return database.Workspace{}, sql.ErrNoRows
-}
-
-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)
-		} else if ws.Ttl.Valid {
-			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/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
@@ -50,37 +50,6 @@ WHERE
 	END
 ;
 
--- name: GetWorkspaceAutostart :one
-SELECT
-	*
-FROM
-	workspaces
-WHERE
-	deleted = false
-AND
-(
-	id = @workspace_id
-	AND (
-		(autostart_schedule IS NOT NULL AND autostart_schedule <> '')
-		OR
-		(ttl IS NOT NULL AND ttl > 0)
-	)
-);
-
--- name: GetWorkspacesAutostart :many
-SELECT
-	*
-FROM
-	workspaces
-WHERE
-	deleted = false
-AND
-(
-	(autostart_schedule IS NOT NULL AND autostart_schedule <> '')
-	OR
-	(ttl IS NOT NULL AND ttl > 0)
-);
-
 -- name: GetWorkspaceByOwnerIDAndName :one
 SELECT
 	*