Skip to content

Commit a7e3bb0

Browse files
deansheatherdeansheather
committed
feat: allow disabling autostart and custom autostop for template
1 parent 2612e32 commit a7e3bb0

33 files changed

+660
-214
lines changed

‎cli/server.go

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import (
3030
"strconv"
3131
"strings"
3232
"sync"
33+
"sync/atomic"
3334
"time"
3435

3536
"github.com/coreos/go-oidc/v3/oidc"
@@ -61,7 +62,6 @@ import (
6162
"github.com/coder/coder/cli/cliui"
6263
"github.com/coder/coder/cli/config"
6364
"github.com/coder/coder/coderd"
64-
"github.com/coder/coder/coderd/autobuild/executor"
6565
"github.com/coder/coder/coderd/database"
6666
"github.com/coder/coder/coderd/database/dbfake"
6767
"github.com/coder/coder/coderd/database/dbpurge"
@@ -72,6 +72,7 @@ import (
7272
"github.com/coder/coder/coderd/httpapi"
7373
"github.com/coder/coder/coderd/httpmw"
7474
"github.com/coder/coder/coderd/prometheusmetrics"
75+
"github.com/coder/coder/coderd/schedule"
7576
"github.com/coder/coder/coderd/telemetry"
7677
"github.com/coder/coder/coderd/tracing"
7778
"github.com/coder/coder/coderd/updatecheck"
@@ -632,6 +633,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
632633
LoginRateLimit: loginRateLimit,
633634
FilesRateLimit: filesRateLimit,
634635
HTTPClient: httpClient,
636+
TemplateScheduleStore: &atomic.Pointer[schedule.TemplateScheduleStore]{},
635637
SSHConfig: codersdk.SSHConfigResponse{
636638
HostnamePrefix: cfg.SSHConfig.DeploymentName.String(),
637639
SSHConfigOptions: configSSHOptions,
@@ -1017,11 +1019,6 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
10171019
return xerrors.Errorf("notify systemd: %w", err)
10181020
}
10191021

1020-
autobuildPoller := time.NewTicker(cfg.AutobuildPollInterval.Value())
1021-
defer autobuildPoller.Stop()
1022-
autobuildExecutor := executor.New(ctx, options.Database, logger, autobuildPoller.C)
1023-
autobuildExecutor.Run()
1024-
10251022
// Currently there is no way to ask the server to shut
10261023
// itself down, so any exit signal will result in a non-zero
10271024
// exit of the server.

‎coderd/activitybump_test.go

Lines changed: 5 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -16,29 +16,6 @@ import (
1616
"github.com/coder/coder/testutil"
1717
)
1818

19-
type mockTemplateScheduleStore struct {
20-
getFn func(ctx context.Context, db database.Store, templateID uuid.UUID) (schedule.TemplateScheduleOptions, error)
21-
setFn func(ctx context.Context, db database.Store, template database.Template, options schedule.TemplateScheduleOptions) (database.Template, error)
22-
}
23-
24-
var _ schedule.TemplateScheduleStore = mockTemplateScheduleStore{}
25-
26-
func (m mockTemplateScheduleStore) GetTemplateScheduleOptions(ctx context.Context, db database.Store, templateID uuid.UUID) (schedule.TemplateScheduleOptions, error) {
27-
if m.getFn != nil {
28-
return m.getFn(ctx, db, templateID)
29-
}
30-
31-
return schedule.NewAGPLTemplateScheduleStore().GetTemplateScheduleOptions(ctx, db, templateID)
32-
}
33-
34-
func (m mockTemplateScheduleStore) SetTemplateScheduleOptions(ctx context.Context, db database.Store, template database.Template, options schedule.TemplateScheduleOptions) (database.Template, error) {
35-
if m.setFn != nil {
36-
return m.setFn(ctx, db, template, options)
37-
}
38-
39-
return schedule.NewAGPLTemplateScheduleStore().SetTemplateScheduleOptions(ctx, db, template, options)
40-
}
41-
4219
func TestWorkspaceActivityBump(t *testing.T) {
4320
t.Parallel()
4421

@@ -57,12 +34,12 @@ func TestWorkspaceActivityBump(t *testing.T) {
5734
// Agent stats trigger the activity bump, so we want to report
5835
// very frequently in tests.
5936
AgentStatsRefreshInterval: time.Millisecond * 100,
60-
TemplateScheduleStore: mockTemplateScheduleStore{
61-
getFn: func(ctx context.Context, db database.Store, templateID uuid.UUID) (schedule.TemplateScheduleOptions, error) {
37+
TemplateScheduleStore: schedule.MockTemplateScheduleStore{
38+
GetFn: func(ctx context.Context, db database.Store, templateID uuid.UUID) (schedule.TemplateScheduleOptions, error) {
6239
return schedule.TemplateScheduleOptions{
63-
UserSchedulingEnabled: true,
64-
DefaultTTL: ttl,
65-
MaxTTL: maxTTL,
40+
UserAutoStopEnabled: true,
41+
DefaultTTL: ttl,
42+
MaxTTL: maxTTL,
6643
}, nil
6744
},
6845
},

‎coderd/apidoc/docs.go

Lines changed: 6 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/apidoc/swagger.json

Lines changed: 6 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/autobuild/executor/lifecycle_executor.go

Lines changed: 40 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package executor
33
import (
44
"context"
55
"encoding/json"
6+
"sync/atomic"
67
"time"
78

89
"github.com/google/uuid"
@@ -18,11 +19,12 @@ import (
1819

1920
// Executor automatically starts or stops workspaces.
2021
type Executor struct {
21-
ctx context.Context
22-
db database.Store
23-
log slog.Logger
24-
tick <-chan time.Time
25-
statsCh chan<- Stats
22+
ctx context.Context
23+
db database.Store
24+
templateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
25+
log slog.Logger
26+
tick <-chan time.Time
27+
statsCh chan<- Stats
2628
}
2729

2830
// Stats contains information about one run of Executor.
@@ -33,13 +35,14 @@ type Stats struct {
3335
}
3436

3537
// New returns a new autobuild executor.
36-
func New(ctx context.Context, db database.Store, log slog.Logger, tick <-chan time.Time) *Executor {
38+
func New(ctx context.Context, db database.Store, tss *atomic.Pointer[schedule.TemplateScheduleStore], log slog.Logger, tick <-chan time.Time) *Executor {
3739
le := &Executor{
3840
//nolint:gocritic // Autostart has a limited set of permissions.
39-
ctx: dbauthz.AsAutostart(ctx),
40-
db: db,
41-
tick: tick,
42-
log: log,
41+
ctx: dbauthz.AsAutostart(ctx),
42+
db: db,
43+
templateScheduleStore: tss,
44+
tick: tick,
45+
log: log,
4346
}
4447
return le
4548
}
@@ -102,30 +105,20 @@ func (e *Executor) runOnce(t time.Time) Stats {
102105
// NOTE: If a workspace build is created with a given TTL and then the user either
103106
// changes or unsets the TTL, the deadline for the workspace build will not
104107
// have changed. This behavior is as expected per #2229.
105-
workspaceRows, err := e.db.GetWorkspaces(e.ctx, database.GetWorkspacesParams{
106-
Deleted: false,
107-
})
108+
workspaces, err := e.db.GetWorkspacesEligibleForAutoStartStop(e.ctx, t)
108109
if err != nil {
109110
e.log.Error(e.ctx, "get workspaces for autostart or autostop", slog.Error(err))
110111
return stats
111112
}
112-
workspaces := database.ConvertWorkspaceRows(workspaceRows)
113-
114-
var eligibleWorkspaceIDs []uuid.UUID
115-
for _, ws := range workspaces {
116-
if isEligibleForAutoStartStop(ws) {
117-
eligibleWorkspaceIDs = append(eligibleWorkspaceIDs, ws.ID)
118-
}
119-
}
120113

121114
// We only use errgroup here for convenience of API, not for early
122115
// cancellation. This means we only return nil errors in th eg.Go.
123116
eg := errgroup.Group{}
124117
// Limit the concurrency to avoid overloading the database.
125118
eg.SetLimit(10)
126119

127-
for _, wsID := range eligibleWorkspaceIDs {
128-
wsID := wsID
120+
for _, ws := range workspaces {
121+
wsID := ws.ID
129122
log := e.log.With(slog.F("workspace_id", wsID))
130123

131124
eg.Go(func() error {
@@ -137,9 +130,6 @@ func (e *Executor) runOnce(t time.Time) Stats {
137130
log.Error(e.ctx, "get workspace autostart failed", slog.Error(err))
138131
return nil
139132
}
140-
if !isEligibleForAutoStartStop(ws) {
141-
return nil
142-
}
143133

144134
// Determine the workspace state based on its latest build.
145135
priorHistory, err := db.GetLatestWorkspaceBuildByWorkspaceID(e.ctx, ws.ID)
@@ -148,6 +138,16 @@ func (e *Executor) runOnce(t time.Time) Stats {
148138
return nil
149139
}
150140

141+
templateSchedule, err := (*(e.templateScheduleStore.Load())).GetTemplateScheduleOptions(e.ctx, db, ws.TemplateID)
142+
if err != nil {
143+
log.Warn(e.ctx, "get template schedule options", slog.Error(err))
144+
return nil
145+
}
146+
147+
if !isEligibleForAutoStartStop(ws, priorHistory, templateSchedule) {
148+
return nil
149+
}
150+
151151
priorJob, err := db.GetProvisionerJobByID(e.ctx, priorHistory.JobID)
152152
if err != nil {
153153
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 {
198198
return stats
199199
}
200200

201-
func isEligibleForAutoStartStop(ws database.Workspace) bool {
202-
return !ws.Deleted && (ws.AutostartSchedule.String != "" || ws.Ttl.Int64 > 0)
201+
func isEligibleForAutoStartStop(ws database.Workspace, priorHistory database.WorkspaceBuild, templateSchedule schedule.TemplateScheduleOptions) bool {
202+
if ws.Deleted {
203+
return false
204+
}
205+
if templateSchedule.UserAutoStartEnabled && ws.AutostartSchedule.Valid && ws.AutostartSchedule.String != "" {
206+
return true
207+
}
208+
// Don't check the template schedule to see whether it allows autostop, this
209+
// is done during the build when determining the deadline.
210+
if priorHistory.Transition == database.WorkspaceTransitionStart && !priorHistory.Deadline.IsZero() {
211+
return true
212+
}
213+
214+
return false
203215
}
204216

205217
func getNextTransition(

‎coderd/autobuild/executor/lifecycle_executor_test.go

Lines changed: 48 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,10 @@ import (
66
"testing"
77
"time"
88

9-
"go.uber.org/goleak"
10-
119
"github.com/google/uuid"
10+
"github.com/stretchr/testify/assert"
11+
"github.com/stretchr/testify/require"
12+
"go.uber.org/goleak"
1213

1314
"github.com/coder/coder/coderd/autobuild/executor"
1415
"github.com/coder/coder/coderd/coderdtest"
@@ -18,9 +19,6 @@ import (
1819
"github.com/coder/coder/codersdk"
1920
"github.com/coder/coder/provisioner/echo"
2021
"github.com/coder/coder/provisionersdk/proto"
21-
22-
"github.com/stretchr/testify/assert"
23-
"github.com/stretchr/testify/require"
2422
)
2523

2624
func TestExecutorAutostartOK(t *testing.T) {
@@ -605,6 +603,51 @@ func TestExecutorAutostartWithParameters(t *testing.T) {
605603
mustWorkspaceParameters(t, client, workspace.LatestBuild.ID)
606604
}
607605

606+
func TestExecutorAutostartTemplateDisabled(t *testing.T) {
607+
t.Parallel()
608+
609+
var (
610+
sched = mustSchedule(t, "CRON_TZ=UTC 0 * * * *")
611+
tickCh = make(chan time.Time)
612+
statsCh = make(chan executor.Stats)
613+
614+
client = coderdtest.New(t, &coderdtest.Options{
615+
AutobuildTicker: tickCh,
616+
IncludeProvisionerDaemon: true,
617+
AutobuildStats: statsCh,
618+
TemplateScheduleStore: schedule.MockTemplateScheduleStore{
619+
GetFn: func(_ context.Context, _ database.Store, _ uuid.UUID) (schedule.TemplateScheduleOptions, error) {
620+
return schedule.TemplateScheduleOptions{
621+
UserAutoStartEnabled: false,
622+
UserAutoStopEnabled: true,
623+
DefaultTTL: 0,
624+
MaxTTL: 0,
625+
}, nil
626+
},
627+
},
628+
})
629+
// futureTime = time.Now().Add(time.Hour)
630+
// futureTimeCron = fmt.Sprintf("%d %d * * *", futureTime.Minute(), futureTime.Hour())
631+
// Given: we have a user with a workspace configured to autostart some time in the future
632+
workspace = mustProvisionWorkspace(t, client, func(cwr *codersdk.CreateWorkspaceRequest) {
633+
cwr.AutostartSchedule = ptr.Ref(sched.String())
634+
})
635+
)
636+
// Given: workspace is stopped
637+
workspace = coderdtest.MustTransitionWorkspace(t, client, workspace.ID, database.WorkspaceTransitionStart, database.WorkspaceTransitionStop)
638+
639+
// When: the autobuild executor ticks before the next scheduled time
640+
go func() {
641+
tickCh <- sched.Next(workspace.LatestBuild.CreatedAt).Add(time.Minute)
642+
close(tickCh)
643+
}()
644+
645+
// Then: nothing should happen
646+
stats := <-statsCh
647+
assert.NoError(t, stats.Error)
648+
assert.Len(t, stats.Transitions, 0)
649+
}
650+
608651
func mustProvisionWorkspace(t *testing.T, client *codersdk.Client, mut ...func(*codersdk.CreateWorkspaceRequest)) codersdk.Workspace {
609652
t.Helper()
610653
user := coderdtest.CreateFirstUser(t, client)

‎coderd/coderd.go

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,7 @@ type Options struct {
120120
DERPMap *tailcfg.DERPMap
121121
SwaggerEndpoint bool
122122
SetUserGroups func(ctx context.Context, tx database.Store, userID uuid.UUID, groupNames []string) error
123-
TemplateScheduleStore schedule.TemplateScheduleStore
123+
TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
124124
// AppSigningKey denotes the symmetric key to use for signing app tickets.
125125
// The key must be 64 bytes long.
126126
AppSigningKey []byte
@@ -231,7 +231,11 @@ func New(options *Options) *API {
231231
}
232232
}
233233
if options.TemplateScheduleStore == nil {
234-
options.TemplateScheduleStore = schedule.NewAGPLTemplateScheduleStore()
234+
options.TemplateScheduleStore = &atomic.Pointer[schedule.TemplateScheduleStore]{}
235+
}
236+
if options.TemplateScheduleStore.Load() == nil {
237+
v := schedule.NewAGPLTemplateScheduleStore()
238+
options.TemplateScheduleStore.Store(&v)
235239
}
236240
if len(options.AppSigningKey) != 64 {
237241
panic("coderd: AppSigningKey must be 64 bytes long")
@@ -292,7 +296,7 @@ func New(options *Options) *API {
292296
),
293297
metricsCache: metricsCache,
294298
Auditor: atomic.Pointer[audit.Auditor]{},
295-
TemplateScheduleStore: atomic.Pointer[schedule.TemplateScheduleStore]{},
299+
TemplateScheduleStore: options.TemplateScheduleStore,
296300
Experiments: experiments,
297301
}
298302
if options.UpdateCheckOptions != nil {
@@ -309,7 +313,6 @@ func New(options *Options) *API {
309313
}
310314

311315
api.Auditor.Store(&options.Auditor)
312-
api.TemplateScheduleStore.Store(&options.TemplateScheduleStore)
313316
api.workspaceAgentCache = wsconncache.New(api.dialWorkspaceAgentTailnet, 0)
314317
api.TailnetCoordinator.Store(&options.TailnetCoordinator)
315318

@@ -745,7 +748,7 @@ type API struct {
745748
WorkspaceClientCoordinateOverride atomic.Pointer[func(rw http.ResponseWriter) bool]
746749
TailnetCoordinator atomic.Pointer[tailnet.Coordinator]
747750
QuotaCommitter atomic.Pointer[proto.QuotaCommitter]
748-
TemplateScheduleStore atomic.Pointer[schedule.TemplateScheduleStore]
751+
TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
749752

750753
HTTPAuth *HTTPAuthorizer
751754

@@ -855,7 +858,7 @@ func (api *API) CreateInMemoryProvisionerDaemon(ctx context.Context, debounce ti
855858
Tags: tags,
856859
QuotaCommitter: &api.QuotaCommitter,
857860
Auditor: &api.Auditor,
858-
TemplateScheduleStore: &api.TemplateScheduleStore,
861+
TemplateScheduleStore: api.TemplateScheduleStore,
859862
AcquireJobDebounce: debounce,
860863
Logger: api.Logger.Named(fmt.Sprintf("provisionerd-%s", daemon.Name)),
861864
})

0 commit comments

Comments
 (0)