Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 4c6a501

Browse files
deansheatherdeansheather
committed
chore: split enterprise code for max_ttl into interface
1 parent 5509824 commit 4c6a501

14 files changed

+240
-109
lines changed

‎coderd/coderd.go

Lines changed: 28 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -110,12 +110,13 @@ type Options struct {
110110
RealIPConfig *httpmw.RealIPConfig
111111
TrialGenerator func(ctx context.Context, email string) error
112112
// TLSCertificates is used to mesh DERP servers securely.
113-
TLSCertificates []tls.Certificate
114-
TailnetCoordinator tailnet.Coordinator
115-
DERPServer *derp.Server
116-
DERPMap *tailcfg.DERPMap
117-
SwaggerEndpoint bool
118-
SetUserGroups func(ctx context.Context, tx database.Store, userID uuid.UUID, groupNames []string) error
113+
TLSCertificates []tls.Certificate
114+
TailnetCoordinator tailnet.Coordinator
115+
DERPServer *derp.Server
116+
DERPMap *tailcfg.DERPMap
117+
SwaggerEndpoint bool
118+
SetUserGroups func(ctx context.Context, tx database.Store, userID uuid.UUID, groupNames []string) error
119+
TemplateScheduleStore provisionerdserver.TemplateScheduleStore
119120

120121
// APIRateLimit is the minutely throughput rate limit per user or ip.
121122
// Setting a rate limit <0 will disable the rate limiter across the entire
@@ -206,6 +207,9 @@ func New(options *Options) *API {
206207
if options.SetUserGroups == nil {
207208
options.SetUserGroups = func(context.Context, database.Store, uuid.UUID, []string) error { return nil }
208209
}
210+
if options.TemplateScheduleStore == nil {
211+
options.TemplateScheduleStore = provisionerdserver.NewAGPLTemplateScheduleStore()
212+
}
209213

210214
siteCacheDir := options.CacheDir
211215
if siteCacheDir != "" {
@@ -232,9 +236,10 @@ func New(options *Options) *API {
232236
Authorizer: options.Authorizer,
233237
Logger: options.Logger,
234238
},
235-
metricsCache: metricsCache,
236-
Auditor: atomic.Pointer[audit.Auditor]{},
237-
Experiments: experiments,
239+
metricsCache: metricsCache,
240+
Auditor: atomic.Pointer[audit.Auditor]{},
241+
TemplateScheduleStore: atomic.Pointer[provisionerdserver.TemplateScheduleStore]{},
242+
Experiments: experiments,
238243
}
239244
if options.UpdateCheckOptions != nil {
240245
api.updateChecker = updatecheck.New(
@@ -244,6 +249,7 @@ func New(options *Options) *API {
244249
)
245250
}
246251
api.Auditor.Store(&options.Auditor)
252+
api.TemplateScheduleStore.Store(&options.TemplateScheduleStore)
247253
api.workspaceAgentCache = wsconncache.New(api.dialWorkspaceAgentTailnet, 0)
248254
api.TailnetCoordinator.Store(&options.TailnetCoordinator)
249255
oauthConfigs := &httpmw.OAuth2Configs{
@@ -679,6 +685,7 @@ type API struct {
679685
WorkspaceClientCoordinateOverride atomic.Pointer[func(rw http.ResponseWriter) bool]
680686
TailnetCoordinator atomic.Pointer[tailnet.Coordinator]
681687
QuotaCommitter atomic.Pointer[proto.QuotaCommitter]
688+
TemplateScheduleStore atomic.Pointer[provisionerdserver.TemplateScheduleStore]
682689

683690
HTTPAuth *HTTPAuthorizer
684691

@@ -770,17 +777,18 @@ func (api *API) CreateInMemoryProvisionerDaemon(ctx context.Context, debounce ti
770777

771778
mux := drpcmux.New()
772779
err = proto.DRPCRegisterProvisionerDaemon(mux, &provisionerdserver.Server{
773-
AccessURL: api.AccessURL,
774-
ID: daemon.ID,
775-
Database: api.Database,
776-
Pubsub: api.Pubsub,
777-
Provisioners: daemon.Provisioners,
778-
Telemetry: api.Telemetry,
779-
Tags: tags,
780-
QuotaCommitter: &api.QuotaCommitter,
781-
Auditor: &api.Auditor,
782-
AcquireJobDebounce: debounce,
783-
Logger: api.Logger.Named(fmt.Sprintf("provisionerd-%s", daemon.Name)),
780+
AccessURL: api.AccessURL,
781+
ID: daemon.ID,
782+
Database: api.Database,
783+
Pubsub: api.Pubsub,
784+
Provisioners: daemon.Provisioners,
785+
Telemetry: api.Telemetry,
786+
Tags: tags,
787+
QuotaCommitter: &api.QuotaCommitter,
788+
Auditor: &api.Auditor,
789+
TemplateScheduleStore: &api.TemplateScheduleStore,
790+
AcquireJobDebounce: debounce,
791+
Logger: api.Logger.Named(fmt.Sprintf("provisionerd-%s", daemon.Name)),
784792
})
785793
if err != nil {
786794
return nil, err

‎coderd/database/models.go

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go

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

‎coderd/database/sqlc.yaml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ overrides:
4242
group_acl: GroupACL
4343
troubleshooting_url: TroubleshootingURL
4444
default_ttl: DefaultTTL
45+
max_ttl: MaxTTL
4546
motd_file: MOTDFile
4647

4748
sql:

‎coderd/provisionerdserver/provisionerdserver.go

Lines changed: 51 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -40,16 +40,17 @@ var (
4040
)
4141

4242
type Server struct {
43-
AccessURL *url.URL
44-
ID uuid.UUID
45-
Logger slog.Logger
46-
Provisioners []database.ProvisionerType
47-
Tags json.RawMessage
48-
Database database.Store
49-
Pubsub database.Pubsub
50-
Telemetry telemetry.Reporter
51-
QuotaCommitter *atomic.Pointer[proto.QuotaCommitter]
52-
Auditor *atomic.Pointer[audit.Auditor]
43+
AccessURL *url.URL
44+
ID uuid.UUID
45+
Logger slog.Logger
46+
Provisioners []database.ProvisionerType
47+
Tags json.RawMessage
48+
Database database.Store
49+
Pubsub database.Pubsub
50+
Telemetry telemetry.Reporter
51+
QuotaCommitter *atomic.Pointer[proto.QuotaCommitter]
52+
Auditor *atomic.Pointer[audit.Auditor]
53+
TemplateScheduleStore *atomic.Pointer[TemplateScheduleStore]
5354

5455
AcquireJobDebounce time.Duration
5556
}
@@ -520,12 +521,12 @@ func (server *Server) FailJob(ctx context.Context, failJob *proto.FailedJob) (*p
520521

521522
var build database.WorkspaceBuild
522523
err := server.Database.InTx(func(db database.Store) error {
523-
workspaceBuild, err := server.Database.GetWorkspaceBuildByID(ctx, input.WorkspaceBuildID)
524+
workspaceBuild, err := db.GetWorkspaceBuildByID(ctx, input.WorkspaceBuildID)
524525
if err != nil {
525526
return xerrors.Errorf("get workspace build: %w", err)
526527
}
527528

528-
build, err = server.Database.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
529+
build, err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
529530
ID: input.WorkspaceBuildID,
530531
UpdatedAt: database.Now(),
531532
ProvisionerState: jobType.WorkspaceBuild.State,
@@ -611,6 +612,8 @@ func (server *Server) FailJob(ctx context.Context, failJob *proto.FailedJob) (*p
611612
}
612613

613614
// CompleteJob is triggered by a provision daemon to mark a provisioner job as completed.
615+
//
616+
//nolint:gocyclo
614617
func (server *Server) CompleteJob(ctx context.Context, completed *proto.CompletedJob) (*proto.Empty, error) {
615618
jobID, err := uuid.Parse(completed.JobId)
616619
if err != nil {
@@ -715,39 +718,49 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
715718
var getWorkspaceError error
716719

717720
err = server.Database.InTx(func(db database.Store) error {
718-
now := database.Now()
719-
var workspaceDeadline time.Time
721+
var (
722+
now = database.Now()
723+
// deadline is the time when the workspace will be stopped. The
724+
// value can be bumped by user activity or manually by the user
725+
// via the UI.
726+
deadline time.Time
727+
// maxDeadline is the maximum value for deadline.
728+
maxDeadline time.Time
729+
)
730+
720731
workspace, getWorkspaceError = db.GetWorkspaceByID(ctx, workspaceBuild.WorkspaceID)
721-
if getWorkspaceError == nil {
722-
if workspace.Ttl.Valid {
723-
workspaceDeadline = now.Add(time.Duration(workspace.Ttl.Int64))
724-
}
725-
} else {
726-
// Huh? Did the workspace get deleted?
727-
// In any case, since this is just for the TTL, try and continue anyway.
732+
if getWorkspaceError != nil {
728733
server.Logger.Error(ctx,
729734
"fetch workspace for build",
730735
slog.F("workspace_build_id", workspaceBuild.ID),
731736
slog.F("workspace_id", workspaceBuild.WorkspaceID),
732737
)
738+
return getWorkspaceError
739+
}
740+
if workspace.Ttl.Valid {
741+
deadline = now.Add(time.Duration(workspace.Ttl.Int64))
733742
}
734743

735-
var workspaceMaxDeadline time.Time
736-
template, err := db.GetTemplateByID(ctx, workspace.TemplateID)
737-
if err == nil {
738-
if template.MaxTtl > 0 {
739-
workspaceMaxDeadline = now.Add(time.Duration(template.MaxTtl))
740-
if !workspaceDeadline.IsZero() && workspaceMaxDeadline.Before(workspaceDeadline) {
741-
workspaceDeadline = workspaceMaxDeadline
742-
}
744+
templateSchedule, err := (*server.TemplateScheduleStore.Load()).GetTemplateScheduleOptions(ctx, db, workspace.TemplateID)
745+
if err != nil {
746+
return xerrors.Errorf("get template schedule options: %w", err)
747+
}
748+
if !templateSchedule.UserSchedulingEnabled {
749+
// The user is not permitted to set their own TTL.
750+
deadline = time.Time{}
751+
}
752+
if deadline.IsZero() && templateSchedule.DefaultTTL > 0 {
753+
deadline = now.Add(templateSchedule.DefaultTTL)
754+
}
755+
if templateSchedule.MaxTTL > 0 {
756+
maxDeadline = now.Add(templateSchedule.MaxTTL)
757+
758+
if deadline.IsZero() || maxDeadline.Before(deadline) {
759+
// If the workspace doesn't have a deadline or the max
760+
// deadline is sooner than the workspace deadline, use the
761+
// max deadline as the actual deadline.
762+
deadline = maxDeadline
743763
}
744-
} else {
745-
server.Logger.Error(ctx,
746-
"fetch template for build",
747-
slog.F("workspace_build_id", workspaceBuild.ID),
748-
slog.F("workspace_id", workspaceBuild.WorkspaceID),
749-
slog.F("template_id", workspace.TemplateID),
750-
)
751764
}
752765

753766
err = db.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
@@ -763,8 +776,8 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
763776
}
764777
_, err = db.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
765778
ID: workspaceBuild.ID,
766-
Deadline: workspaceDeadline,
767-
MaxDeadline: workspaceMaxDeadline,
779+
Deadline: deadline,
780+
MaxDeadline: maxDeadline,
768781
ProvisionerState: jobType.WorkspaceBuild.State,
769782
UpdatedAt: now,
770783
})

0 commit comments

Comments
 (0)