Skip to content

Commit a52aca2

Browse files
sreyasreya
committedJul 14, 2023
add tests for updating workspace deleting_at
1 parent 056a454 commit a52aca2

File tree

3 files changed

+165
-84
lines changed

3 files changed

+165
-84
lines changed
 

‎enterprise/coderd/provisionerdaemons.go

Lines changed: 0 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import (
1010
"net"
1111
"net/http"
1212
"strings"
13-
"time"
1413

1514
"github.com/google/uuid"
1615
"github.com/hashicorp/yamux"
@@ -28,7 +27,6 @@ import (
2827
"github.com/coder/coder/coderd/httpmw"
2928
"github.com/coder/coder/coderd/provisionerdserver"
3029
"github.com/coder/coder/coderd/rbac"
31-
"github.com/coder/coder/coderd/schedule"
3230
"github.com/coder/coder/codersdk"
3331
"github.com/coder/coder/provisionerd/proto"
3432
)
@@ -310,85 +308,3 @@ func websocketNetConn(ctx context.Context, conn *websocket.Conn, msgType websock
310308
}
311309
}
312310

313-
type EnterpriseTemplateScheduleStore struct{}
314-
315-
var _ schedule.TemplateScheduleStore = &EnterpriseTemplateScheduleStore{}
316-
317-
func (*EnterpriseTemplateScheduleStore) GetTemplateScheduleOptions(ctx context.Context, db database.Store, templateID uuid.UUID) (schedule.TemplateScheduleOptions, error) {
318-
tpl, err := db.GetTemplateByID(ctx, templateID)
319-
if err != nil {
320-
return schedule.TemplateScheduleOptions{}, err
321-
}
322-
323-
return schedule.TemplateScheduleOptions{
324-
UserAutostartEnabled: tpl.AllowUserAutostart,
325-
UserAutostopEnabled: tpl.AllowUserAutostop,
326-
DefaultTTL: time.Duration(tpl.DefaultTTL),
327-
MaxTTL: time.Duration(tpl.MaxTTL),
328-
FailureTTL: time.Duration(tpl.FailureTTL),
329-
InactivityTTL: time.Duration(tpl.InactivityTTL),
330-
LockedTTL: time.Duration(tpl.LockedTTL),
331-
}, nil
332-
}
333-
334-
func (*EnterpriseTemplateScheduleStore) SetTemplateScheduleOptions(ctx context.Context, db database.Store, tpl database.Template, opts schedule.TemplateScheduleOptions) (database.Template, error) {
335-
if int64(opts.DefaultTTL) == tpl.DefaultTTL &&
336-
int64(opts.MaxTTL) == tpl.MaxTTL &&
337-
int64(opts.FailureTTL) == tpl.FailureTTL &&
338-
int64(opts.InactivityTTL) == tpl.InactivityTTL &&
339-
int64(opts.LockedTTL) == tpl.LockedTTL &&
340-
opts.UserAutostartEnabled == tpl.AllowUserAutostart &&
341-
opts.UserAutostopEnabled == tpl.AllowUserAutostop {
342-
// Avoid updating the UpdatedAt timestamp if nothing will be changed.
343-
return tpl, nil
344-
}
345-
346-
template, err := db.UpdateTemplateScheduleByID(ctx, database.UpdateTemplateScheduleByIDParams{
347-
ID: tpl.ID,
348-
UpdatedAt: database.Now(),
349-
AllowUserAutostart: opts.UserAutostartEnabled,
350-
AllowUserAutostop: opts.UserAutostopEnabled,
351-
DefaultTTL: int64(opts.DefaultTTL),
352-
MaxTTL: int64(opts.MaxTTL),
353-
FailureTTL: int64(opts.FailureTTL),
354-
InactivityTTL: int64(opts.InactivityTTL),
355-
LockedTTL: int64(opts.LockedTTL),
356-
})
357-
if err != nil {
358-
return database.Template{}, xerrors.Errorf("update template schedule: %w", err)
359-
}
360-
361-
// Update all workspaces using the template to set the user defined schedule
362-
// to be within the new bounds. This essentially does the following for each
363-
// workspace using the template.
364-
// if (template.ttl != NULL) {
365-
// workspace.ttl = min(workspace.ttl, template.ttl)
366-
// }
367-
//
368-
// NOTE: this does not apply to currently running workspaces as their
369-
// schedule information is committed to the workspace_build during start.
370-
// This limitation is displayed to the user while editing the template.
371-
if opts.MaxTTL > 0 {
372-
err = db.UpdateWorkspaceTTLToBeWithinTemplateMax(ctx, database.UpdateWorkspaceTTLToBeWithinTemplateMaxParams{
373-
TemplateID: template.ID,
374-
TemplateMaxTTL: int64(opts.MaxTTL),
375-
})
376-
if err != nil {
377-
return database.Template{}, xerrors.Errorf("update TTL of all workspaces on template to be within new template max TTL: %w", err)
378-
}
379-
}
380-
381-
// If we updated the locked_ttl we need to update all the workspaces deleting_at
382-
// to ensure workspaces are being cleaned up correctly.
383-
if opts.LockedTTL > 0 {
384-
err = db.UpdateWorkspacesDeletingAtByTemplateID(ctx, database.UpdateWorkspacesDeletingAtByTemplateIDParams{
385-
TemplateID: template.ID,
386-
LockedTtlMs: opts.LockedTTL.Milliseconds(),
387-
})
388-
if err != nil {
389-
return database.Template{}, xerrors.Errorf("update deleting_at of all workspaces for new locked_ttl %q: %w", opts.LockedTTL, err)
390-
}
391-
}
392-
393-
return template, nil
394-
}

‎enterprise/coderd/templates_test.go

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,76 @@ func TestTemplates(t *testing.T) {
255255
require.Equal(t, inactivityTTL, updated.InactivityTTLMillis)
256256
require.Equal(t, lockedTTL, updated.LockedTTLMillis)
257257
})
258+
259+
t.Run("UpdateLockedTTL", func(t *testing.T) {
260+
t.Parallel()
261+
262+
ctx := testutil.Context(t, testutil.WaitMedium)
263+
client := coderdenttest.New(t, &coderdenttest.Options{
264+
Options: &coderdtest.Options{
265+
IncludeProvisionerDaemon: true,
266+
},
267+
})
268+
user := coderdtest.CreateFirstUser(t, client)
269+
_ = coderdenttest.AddFullLicense(t, client)
270+
271+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
272+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
273+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
274+
275+
unlockedWorkspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
276+
lockedWorkspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
277+
require.Nil(t, unlockedWorkspace.DeletingAt)
278+
require.Nil(t, lockedWorkspace.DeletingAt)
279+
280+
_ = coderdtest.AwaitWorkspaceBuildJob(t, client, unlockedWorkspace.LatestBuild.ID)
281+
_ = coderdtest.AwaitWorkspaceBuildJob(t, client, lockedWorkspace.LatestBuild.ID)
282+
283+
err := client.UpdateWorkspaceLock(ctx, lockedWorkspace.ID, codersdk.UpdateWorkspaceLock{
284+
Lock: true,
285+
})
286+
require.NoError(t, err)
287+
288+
lockedWorkspace = coderdtest.MustWorkspace(t, client, lockedWorkspace.ID)
289+
require.NotNil(t, lockedWorkspace.LockedAt)
290+
// The deleting_at field should be nil since there is no template locked_ttl set.
291+
require.Nil(t, lockedWorkspace.DeletingAt)
292+
293+
lockedTTL := time.Minute
294+
updated, err := client.UpdateTemplateMeta(ctx, template.ID, codersdk.UpdateTemplateMeta{
295+
LockedTTLMillis: lockedTTL.Milliseconds(),
296+
})
297+
require.NoError(t, err)
298+
require.Equal(t, lockedTTL.Milliseconds(), updated.LockedTTLMillis)
299+
300+
unlockedWorkspace = coderdtest.MustWorkspace(t, client, unlockedWorkspace.ID)
301+
require.Nil(t, unlockedWorkspace.LockedAt)
302+
require.Nil(t, unlockedWorkspace.DeletingAt)
303+
304+
lockedWorkspace = coderdtest.MustWorkspace(t, client, lockedWorkspace.ID)
305+
require.NotNil(t, lockedWorkspace.LockedAt)
306+
require.NotNil(t, lockedWorkspace.DeletingAt)
307+
require.Equal(t, lockedWorkspace.LockedAt.Add(lockedTTL), *lockedWorkspace.DeletingAt)
308+
309+
// Disable the locked_ttl on the template, then we can assert that the workspaces
310+
// no longer have a deleting_at field.
311+
updated, err = client.UpdateTemplateMeta(ctx, template.ID, codersdk.UpdateTemplateMeta{
312+
LockedTTLMillis: 0,
313+
})
314+
require.NoError(t, err)
315+
require.EqualValues(t, 0, updated.LockedTTLMillis)
316+
317+
// The unlocked workspace should remain unchanged.
318+
unlockedWorkspace = coderdtest.MustWorkspace(t, client, unlockedWorkspace.ID)
319+
require.Nil(t, unlockedWorkspace.LockedAt)
320+
require.Nil(t, unlockedWorkspace.DeletingAt)
321+
322+
// Fetch the locked workspace. It should still be locked, but it should no
323+
// longer be scheduled for deletion.
324+
lockedWorkspace = coderdtest.MustWorkspace(t, client, lockedWorkspace.ID)
325+
require.NotNil(t, lockedWorkspace.LockedAt)
326+
require.Nil(t, lockedWorkspace.DeletingAt)
327+
})
258328
}
259329

260330
func TestTemplateACL(t *testing.T) {

‎enterprise/coderd/templateschedule.go

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
package coderd
2+
3+
import (
4+
"context"
5+
"time"
6+
7+
"github.com/google/uuid"
8+
"golang.org/x/xerrors"
9+
10+
"github.com/coder/coder/coderd/database"
11+
"github.com/coder/coder/coderd/schedule"
12+
)
13+
14+
type EnterpriseTemplateScheduleStore struct{}
15+
16+
var _ schedule.TemplateScheduleStore = &EnterpriseTemplateScheduleStore{}
17+
18+
func (*EnterpriseTemplateScheduleStore) GetTemplateScheduleOptions(ctx context.Context, db database.Store, templateID uuid.UUID) (schedule.TemplateScheduleOptions, error) {
19+
tpl, err := db.GetTemplateByID(ctx, templateID)
20+
if err != nil {
21+
return schedule.TemplateScheduleOptions{}, err
22+
}
23+
24+
return schedule.TemplateScheduleOptions{
25+
UserAutostartEnabled: tpl.AllowUserAutostart,
26+
UserAutostopEnabled: tpl.AllowUserAutostop,
27+
DefaultTTL: time.Duration(tpl.DefaultTTL),
28+
MaxTTL: time.Duration(tpl.MaxTTL),
29+
FailureTTL: time.Duration(tpl.FailureTTL),
30+
InactivityTTL: time.Duration(tpl.InactivityTTL),
31+
LockedTTL: time.Duration(tpl.LockedTTL),
32+
}, nil
33+
}
34+
35+
func (*EnterpriseTemplateScheduleStore) SetTemplateScheduleOptions(ctx context.Context, db database.Store, tpl database.Template, opts schedule.TemplateScheduleOptions) (database.Template, error) {
36+
if int64(opts.DefaultTTL) == tpl.DefaultTTL &&
37+
int64(opts.MaxTTL) == tpl.MaxTTL &&
38+
int64(opts.FailureTTL) == tpl.FailureTTL &&
39+
int64(opts.InactivityTTL) == tpl.InactivityTTL &&
40+
int64(opts.LockedTTL) == tpl.LockedTTL &&
41+
opts.UserAutostartEnabled == tpl.AllowUserAutostart &&
42+
opts.UserAutostopEnabled == tpl.AllowUserAutostop {
43+
// Avoid updating the UpdatedAt timestamp if nothing will be changed.
44+
return tpl, nil
45+
}
46+
47+
template, err := db.UpdateTemplateScheduleByID(ctx, database.UpdateTemplateScheduleByIDParams{
48+
ID: tpl.ID,
49+
UpdatedAt: database.Now(),
50+
AllowUserAutostart: opts.UserAutostartEnabled,
51+
AllowUserAutostop: opts.UserAutostopEnabled,
52+
DefaultTTL: int64(opts.DefaultTTL),
53+
MaxTTL: int64(opts.MaxTTL),
54+
FailureTTL: int64(opts.FailureTTL),
55+
InactivityTTL: int64(opts.InactivityTTL),
56+
LockedTTL: int64(opts.LockedTTL),
57+
})
58+
if err != nil {
59+
return database.Template{}, xerrors.Errorf("update template schedule: %w", err)
60+
}
61+
62+
// Update all workspaces using the template to set the user defined schedule
63+
// to be within the new bounds. This essentially does the following for each
64+
// workspace using the template.
65+
// if (template.ttl != NULL) {
66+
// workspace.ttl = min(workspace.ttl, template.ttl)
67+
// }
68+
//
69+
// NOTE: this does not apply to currently running workspaces as their
70+
// schedule information is committed to the workspace_build during start.
71+
// This limitation is displayed to the user while editing the template.
72+
if opts.MaxTTL > 0 {
73+
err = db.UpdateWorkspaceTTLToBeWithinTemplateMax(ctx, database.UpdateWorkspaceTTLToBeWithinTemplateMaxParams{
74+
TemplateID: template.ID,
75+
TemplateMaxTTL: int64(opts.MaxTTL),
76+
})
77+
if err != nil {
78+
return database.Template{}, xerrors.Errorf("update TTL of all workspaces on template to be within new template max TTL: %w", err)
79+
}
80+
}
81+
82+
// If we updated the locked_ttl we need to update all the workspaces deleting_at
83+
// to ensure workspaces are being cleaned up correctly. Similarly if we are
84+
// disabling it (by passing 0), then we want to delete nullify the deleting_at
85+
// fields of all the template workspaces.
86+
err = db.UpdateWorkspacesDeletingAtByTemplateID(ctx, database.UpdateWorkspacesDeletingAtByTemplateIDParams{
87+
TemplateID: template.ID,
88+
LockedTtlMs: opts.LockedTTL.Milliseconds(),
89+
})
90+
if err != nil {
91+
return database.Template{}, xerrors.Errorf("update deleting_at of all workspaces for new locked_ttl %q: %w", opts.LockedTTL, err)
92+
}
93+
94+
return template, nil
95+
}

0 commit comments

Comments
 (0)