feat: allow disabling autostart and custom autostop for template

coder · deansheather · Apr 4, 2023 · Mar 31, 2023 · Mar 31, 2023 · Mar 31, 2023
commit a7e3bb04841226dee4e6ec32c9d7293b126204d3
diff --git a/cli/server.go b/cli/server.go
@@ -30,6 +30,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/coreos/go-oidc/v3/oidc"
@@ -61,7 +62,6 @@ import (
 	"github.com/coder/coder/cli/cliui"
 	"github.com/coder/coder/cli/config"
 	"github.com/coder/coder/coderd"
-	"github.com/coder/coder/coderd/autobuild/executor"
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/database/dbfake"
 	"github.com/coder/coder/coderd/database/dbpurge"
@@ -72,6 +72,7 @@ import (
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/coderd/httpmw"
 	"github.com/coder/coder/coderd/prometheusmetrics"
+	"github.com/coder/coder/coderd/schedule"
 	"github.com/coder/coder/coderd/telemetry"
 	"github.com/coder/coder/coderd/tracing"
 	"github.com/coder/coder/coderd/updatecheck"
@@ -632,6 +633,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 				LoginRateLimit:              loginRateLimit,
 				FilesRateLimit:              filesRateLimit,
 				HTTPClient:                  httpClient,
+				TemplateScheduleStore:       &atomic.Pointer[schedule.TemplateScheduleStore]{},
 				SSHConfig: codersdk.SSHConfigResponse{
 					HostnamePrefix:   cfg.SSHConfig.DeploymentName.String(),
 					SSHConfigOptions: configSSHOptions,
@@ -1017,11 +1019,6 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 				return xerrors.Errorf("notify systemd: %w", err)
 			}
 
-			autobuildPoller := time.NewTicker(cfg.AutobuildPollInterval.Value())
-			defer autobuildPoller.Stop()
-			autobuildExecutor := executor.New(ctx, options.Database, logger, autobuildPoller.C)
-			autobuildExecutor.Run()
-
 			// Currently there is no way to ask the server to shut
 			// itself down, so any exit signal will result in a non-zero
 			// exit of the server.

diff --git a/coderd/activitybump_test.go b/coderd/activitybump_test.go
@@ -16,29 +16,6 @@ import (
 	"github.com/coder/coder/testutil"
 )
 
-type mockTemplateScheduleStore struct {
-	getFn func(ctx context.Context, db database.Store, templateID uuid.UUID) (schedule.TemplateScheduleOptions, error)
-	setFn func(ctx context.Context, db database.Store, template database.Template, options schedule.TemplateScheduleOptions) (database.Template, error)
-}
-
-var _ schedule.TemplateScheduleStore = mockTemplateScheduleStore{}
-
-func (m mockTemplateScheduleStore) GetTemplateScheduleOptions(ctx context.Context, db database.Store, templateID uuid.UUID) (schedule.TemplateScheduleOptions, error) {
-	if m.getFn != nil {
-		return m.getFn(ctx, db, templateID)
-	}
-
-	return schedule.NewAGPLTemplateScheduleStore().GetTemplateScheduleOptions(ctx, db, templateID)
-}
-
-func (m mockTemplateScheduleStore) SetTemplateScheduleOptions(ctx context.Context, db database.Store, template database.Template, options schedule.TemplateScheduleOptions) (database.Template, error) {
-	if m.setFn != nil {
-		return m.setFn(ctx, db, template, options)
-	}
-
-	return schedule.NewAGPLTemplateScheduleStore().SetTemplateScheduleOptions(ctx, db, template, options)
-}
-
 func TestWorkspaceActivityBump(t *testing.T) {
 	t.Parallel()
 
@@ -57,12 +34,12 @@ func TestWorkspaceActivityBump(t *testing.T) {
 			// Agent stats trigger the activity bump, so we want to report
 			// very frequently in tests.
 			AgentStatsRefreshInterval: time.Millisecond * 100,
-			TemplateScheduleStore: mockTemplateScheduleStore{
-				getFn: func(ctx context.Context, db database.Store, templateID uuid.UUID) (schedule.TemplateScheduleOptions, error) {
+			TemplateScheduleStore: schedule.MockTemplateScheduleStore{
+				GetFn: func(ctx context.Context, db database.Store, templateID uuid.UUID) (schedule.TemplateScheduleOptions, error) {
 					return schedule.TemplateScheduleOptions{
-						UserSchedulingEnabled: true,
-						DefaultTTL:            ttl,
-						MaxTTL:                maxTTL,
+						UserAutoStopEnabled: true,
+						DefaultTTL:          ttl,
+						MaxTTL:              maxTTL,
 					}, nil
 				},
 			},

diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/autobuild/executor/lifecycle_executor.go b/coderd/autobuild/executor/lifecycle_executor.go
@@ -3,6 +3,7 @@ package executor
 import (
 	"context"
 	"encoding/json"
+	"sync/atomic"
 	"time"
 
 	"github.com/google/uuid"
@@ -18,11 +19,12 @@ import (
 
 // Executor automatically starts or stops workspaces.
 type Executor struct {
-	ctx     context.Context
-	db      database.Store
-	log     slog.Logger
-	tick    <-chan time.Time
-	statsCh chan<- Stats
+	ctx                   context.Context
+	db                    database.Store
+	templateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
+	log                   slog.Logger
+	tick                  <-chan time.Time
+	statsCh               chan<- Stats
 }
 
 // Stats contains information about one run of Executor.
@@ -33,13 +35,14 @@ type Stats struct {
 }
 
 // New returns a new autobuild executor.
-func New(ctx context.Context, db database.Store, log slog.Logger, tick <-chan time.Time) *Executor {
+func New(ctx context.Context, db database.Store, tss *atomic.Pointer[schedule.TemplateScheduleStore], log slog.Logger, tick <-chan time.Time) *Executor {
 	le := &Executor{
 		//nolint:gocritic // Autostart has a limited set of permissions.
-		ctx:  dbauthz.AsAutostart(ctx),
-		db:   db,
-		tick: tick,
-		log:  log,
+		ctx:                   dbauthz.AsAutostart(ctx),
+		db:                    db,
+		templateScheduleStore: tss,
+		tick:                  tick,
+		log:                   log,
 	}
 	return le
 }
@@ -102,30 +105,20 @@ 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.
-	workspaceRows, err := e.db.GetWorkspaces(e.ctx, database.GetWorkspacesParams{
-		Deleted: false,
-	})
+	workspaces, err := e.db.GetWorkspacesEligibleForAutoStartStop(e.ctx, t)
 	if err != nil {
 		e.log.Error(e.ctx, "get workspaces for autostart or autostop", slog.Error(err))
 		return stats
 	}
-	workspaces := database.ConvertWorkspaceRows(workspaceRows)
-
-	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 _, wsID := range eligibleWorkspaceIDs {
-		wsID := wsID
+	for _, ws := range workspaces {
+		wsID := ws.ID
 		log := e.log.With(slog.F("workspace_id", wsID))
 
 		eg.Go(func() error {
@@ -137,9 +130,6 @@ func (e *Executor) runOnce(t time.Time) Stats {
 					log.Error(e.ctx, "get workspace autostart failed", slog.Error(err))
 					return nil
 				}
-				if !isEligibleForAutoStartStop(ws) {
-					return nil
-				}
 
 				// Determine the workspace state based on its latest build.
 				priorHistory, err := db.GetLatestWorkspaceBuildByWorkspaceID(e.ctx, ws.ID)
@@ -148,6 +138,16 @@ func (e *Executor) runOnce(t time.Time) Stats {
 					return nil
 				}
 
+				templateSchedule, err := (*(e.templateScheduleStore.Load())).GetTemplateScheduleOptions(e.ctx, db, ws.TemplateID)
+				if err != nil {
+					log.Warn(e.ctx, "get template schedule options", slog.Error(err))
+					return nil
+				}
+
+				if !isEligibleForAutoStartStop(ws, priorHistory, templateSchedule) {
+					return nil
+				}
+
 				priorJob, err := db.GetProvisionerJobByID(e.ctx, priorHistory.JobID)
 				if err != nil {
 					log.Warn(e.ctx, "get last provisioner job for workspace %q: %w", slog.Error(err))
@@ -198,8 +198,20 @@ func (e *Executor) runOnce(t time.Time) Stats {
 	return stats
 }
 
-func isEligibleForAutoStartStop(ws database.Workspace) bool {
-	return !ws.Deleted && (ws.AutostartSchedule.String != "" || ws.Ttl.Int64 > 0)
+func isEligibleForAutoStartStop(ws database.Workspace, priorHistory database.WorkspaceBuild, templateSchedule schedule.TemplateScheduleOptions) bool {
+	if ws.Deleted {
+		return false
+	}
+	if templateSchedule.UserAutoStartEnabled && ws.AutostartSchedule.Valid && ws.AutostartSchedule.String != "" {
+		return true
+	}
+	// Don't check the template schedule to see whether it allows autostop, this
+	// is done during the build when determining the deadline.
+	if priorHistory.Transition == database.WorkspaceTransitionStart && !priorHistory.Deadline.IsZero() {
+		return true
+	}
+
+	return false
 }
 
 func getNextTransition(

diff --git a/coderd/autobuild/executor/lifecycle_executor_test.go b/coderd/autobuild/executor/lifecycle_executor_test.go
@@ -6,9 +6,10 @@ import (
 	"testing"
 	"time"
 
-	"go.uber.org/goleak"
-
 	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
 
 	"github.com/coder/coder/coderd/autobuild/executor"
 	"github.com/coder/coder/coderd/coderdtest"
@@ -18,9 +19,6 @@ import (
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
 func TestExecutorAutostartOK(t *testing.T) {
@@ -605,6 +603,51 @@ func TestExecutorAutostartWithParameters(t *testing.T) {
 	mustWorkspaceParameters(t, client, workspace.LatestBuild.ID)
 }
 
+func TestExecutorAutostartTemplateDisabled(t *testing.T) {
+	t.Parallel()
+
+	var (
+		sched   = mustSchedule(t, "CRON_TZ=UTC 0 * * * *")
+		tickCh  = make(chan time.Time)
+		statsCh = make(chan executor.Stats)
+
+		client = coderdtest.New(t, &coderdtest.Options{
+			AutobuildTicker:          tickCh,
+			IncludeProvisionerDaemon: true,
+			AutobuildStats:           statsCh,
+			TemplateScheduleStore: schedule.MockTemplateScheduleStore{
+				GetFn: func(_ context.Context, _ database.Store, _ uuid.UUID) (schedule.TemplateScheduleOptions, error) {
+					return schedule.TemplateScheduleOptions{
+						UserAutoStartEnabled: false,
+						UserAutoStopEnabled:  true,
+						DefaultTTL:           0,
+						MaxTTL:               0,
+					}, nil
+				},
+			},
+		})
+		// futureTime     = time.Now().Add(time.Hour)
+		// futureTimeCron = fmt.Sprintf("%d %d * * *", futureTime.Minute(), futureTime.Hour())
+		// Given: we have a user with a workspace configured to autostart some time in the future
+		workspace = mustProvisionWorkspace(t, client, func(cwr *codersdk.CreateWorkspaceRequest) {
+			cwr.AutostartSchedule = ptr.Ref(sched.String())
+		})
+	)
+	// Given: workspace is stopped
+	workspace = coderdtest.MustTransitionWorkspace(t, client, workspace.ID, database.WorkspaceTransitionStart, database.WorkspaceTransitionStop)
+
+	// When: the autobuild executor ticks before the next scheduled time
+	go func() {
+		tickCh <- sched.Next(workspace.LatestBuild.CreatedAt).Add(time.Minute)
+		close(tickCh)
+	}()
+
+	// Then: nothing should happen
+	stats := <-statsCh
+	assert.NoError(t, stats.Error)
+	assert.Len(t, stats.Transitions, 0)
+}
+
 func mustProvisionWorkspace(t *testing.T, client *codersdk.Client, mut ...func(*codersdk.CreateWorkspaceRequest)) codersdk.Workspace {
 	t.Helper()
 	user := coderdtest.CreateFirstUser(t, client)

diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -120,7 +120,7 @@ type Options struct {
 	DERPMap               *tailcfg.DERPMap
 	SwaggerEndpoint       bool
 	SetUserGroups         func(ctx context.Context, tx database.Store, userID uuid.UUID, groupNames []string) error
-	TemplateScheduleStore schedule.TemplateScheduleStore
+	TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
 	// AppSigningKey denotes the symmetric key to use for signing app tickets.
 	// The key must be 64 bytes long.
 	AppSigningKey []byte
@@ -231,7 +231,11 @@ func New(options *Options) *API {
 		}
 	}
 	if options.TemplateScheduleStore == nil {
-		options.TemplateScheduleStore = schedule.NewAGPLTemplateScheduleStore()
+		options.TemplateScheduleStore = &atomic.Pointer[schedule.TemplateScheduleStore]{}
+	}
+	if options.TemplateScheduleStore.Load() == nil {
+		v := schedule.NewAGPLTemplateScheduleStore()
+		options.TemplateScheduleStore.Store(&v)
 	}
 	if len(options.AppSigningKey) != 64 {
 		panic("coderd: AppSigningKey must be 64 bytes long")
@@ -292,7 +296,7 @@ func New(options *Options) *API {
 		),
 		metricsCache:          metricsCache,
 		Auditor:               atomic.Pointer[audit.Auditor]{},
-		TemplateScheduleStore: atomic.Pointer[schedule.TemplateScheduleStore]{},
+		TemplateScheduleStore: options.TemplateScheduleStore,
 		Experiments:           experiments,
 	}
 	if options.UpdateCheckOptions != nil {
@@ -309,7 +313,6 @@ func New(options *Options) *API {
 	}
 
 	api.Auditor.Store(&options.Auditor)
-	api.TemplateScheduleStore.Store(&options.TemplateScheduleStore)
 	api.workspaceAgentCache = wsconncache.New(api.dialWorkspaceAgentTailnet, 0)
 	api.TailnetCoordinator.Store(&options.TailnetCoordinator)
 
@@ -745,7 +748,7 @@ type API struct {
 	WorkspaceClientCoordinateOverride atomic.Pointer[func(rw http.ResponseWriter) bool]
 	TailnetCoordinator                atomic.Pointer[tailnet.Coordinator]
 	QuotaCommitter                    atomic.Pointer[proto.QuotaCommitter]
-	TemplateScheduleStore             atomic.Pointer[schedule.TemplateScheduleStore]
+	TemplateScheduleStore             *atomic.Pointer[schedule.TemplateScheduleStore]
 
 	HTTPAuth *HTTPAuthorizer
 
@@ -855,7 +858,7 @@ func (api *API) CreateInMemoryProvisionerDaemon(ctx context.Context, debounce ti
 		Tags:                  tags,
 		QuotaCommitter:        &api.QuotaCommitter,
 		Auditor:               &api.Auditor,
-		TemplateScheduleStore: &api.TemplateScheduleStore,
+		TemplateScheduleStore: api.TemplateScheduleStore,
 		AcquireJobDebounce:    debounce,
 		Logger:                api.Logger.Named(fmt.Sprintf("provisionerd-%s", daemon.Name)),
 	})