Skip to content

Commit 837b21b

Browse files
sreyasreya
committed
add tests for workspacelock
1 parent daf60f5 commit 837b21b

File tree

8 files changed

+123
-125
lines changed

8 files changed

+123
-125
lines changed

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2583,8 +2583,6 @@ func (q *querier) UpdateWorkspaceTTLToBeWithinTemplateMax(ctx context.Context, a
25832583
}
25842584

25852585
func (q *querier) UpdateWorkspacesDeletingAtByTemplateID(ctx context.Context, arg database.UpdateWorkspacesDeletingAtByTemplateIDParams) error {
2586-
// TODO: This is kinda messed up. We want to update this whenever we update the locked_ttl on a template. But a template
2587-
// admin may not have permissions to edit workspaces so this won't work.
25882586
fetch := func(ctx context.Context, arg database.UpdateWorkspacesDeletingAtByTemplateIDParams) (database.Template, error) {
25892587
return q.db.GetTemplateByID(ctx, arg.TemplateID)
25902588
}

‎coderd/database/dbfake/dbfake.go

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5266,11 +5266,12 @@ func (q *fakeQuerier) UpdateWorkspaceLockedDeletingAt(ctx context.Context, arg d
52665266
workspace.LockedAt = arg.LockedAt
52675267
if workspace.LockedAt.Time.IsZero() {
52685268
workspace.LastUsedAt = database.Now()
5269+
workspace.DeletingAt = sql.NullTime{}
52695270
}
5270-
if arg.LockedTtlMs > 0 {
5271+
if arg.LockedTtlMs > 0 && !workspace.LockedAt.Time.IsZero() {
52715272
workspace.DeletingAt = sql.NullTime{
52725273
Valid: true,
5273-
Time: workspace.LastUsedAt.Add(time.Duration(arg.LockedTtlMs) * time.Millisecond),
5274+
Time: workspace.LockedAt.Time.Add(time.Duration(arg.LockedTtlMs) * time.Millisecond),
52745275
}
52755276
}
52765277
q.workspaces[index] = workspace
@@ -5360,13 +5361,30 @@ func (q *fakeQuerier) UpdateWorkspaceTTLToBeWithinTemplateMax(_ context.Context,
53605361
return nil
53615362
}
53625363

5363-
func (q *fakeQuerier) UpdateWorkspacesDeletingAtByTemplateID(ctx context.Context, arg database.UpdateWorkspacesDeletingAtByTemplateIDParams) error {
5364+
func (q *fakeQuerier) UpdateWorkspacesDeletingAtByTemplateID(_ context.Context, arg database.UpdateWorkspacesDeletingAtByTemplateIDParams) error {
5365+
q.mutex.Lock()
5366+
defer q.mutex.Unlock()
5367+
53645368
err := validateDatabaseType(arg)
53655369
if err != nil {
53665370
return err
53675371
}
53685372

5369-
panic("not implemented")
5373+
if arg.LockedTtlMs == 0 {
5374+
return nil
5375+
}
5376+
5377+
for _, ws := range q.workspaces {
5378+
if ws.LockedAt.Time.IsZero() {
5379+
continue
5380+
}
5381+
ws.DeletingAt = sql.NullTime{
5382+
Valid: true,
5383+
Time: ws.LockedAt.Time.Add(time.Duration(arg.LockedTtlMs) * time.Millisecond),
5384+
}
5385+
}
5386+
5387+
return nil
53705388
}
53715389

53725390
func (q *fakeQuerier) UpsertAppSecurityKey(_ context.Context, data string) error {

‎coderd/database/queries.sql.go

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

‎coderd/database/queries/workspaces.sql

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -473,21 +473,21 @@ WHERE
473473
-- name: UpdateWorkspaceLockedDeletingAt :exec
474474
UPDATE workspaces
475475
SET
476-
locked_at = @locked_at,
477-
-- If locked_at is null (meaning unlocked) or the template-defined locked_ttl is 0 we should set
478-
-- deleting_at to NULL else set it to the locked_at + locked_ttl duration.
479-
deleting_at = CASE WHEN @locked_at IS NULL OR @locked_ttl_ms::bigint = 0 THEN NULL ELSE '@locked_at' + interval '1 millisecond' * @locked_ttl_ms::bigint END,
476+
locked_at = $2,
480477
-- When a workspace is unlocked we want to update the last_used_at to avoid the workspace getting re-locked.
481478
-- if we're locking the workspace then we leave it alone.
482-
last_used_at = CASE WHEN @locked_at IS NULL THEN now() at time zone 'utc' ELSE last_used_at END
479+
last_used_at = CASE WHEN $2::timestamptz IS NULL THEN now() at time zone 'utc' ELSE last_used_at END,
480+
-- If locked_at is null (meaning unlocked) or the template-defined locked_ttl is 0 we should set
481+
-- deleting_at to NULL else set it to the locked_at + locked_ttl duration.
482+
deleting_at = CASE WHEN $2::timestamptz IS NULL OR @locked_ttl_ms::bigint = 0 THEN NULL ELSE $2::timestamptz + interval '1 millisecond' * @locked_ttl_ms::bigint END
483483
WHERE
484-
id = @id;
484+
id = $1;
485485

486486
-- name: UpdateWorkspacesDeletingAtByTemplateID :exec
487487
UPDATE
488488
workspaces
489489
SET
490-
deleting_at = locked_at + interval '1 milliseconds' * @locked_ttl_ms::bigint
490+
deleting_at = CASE WHEN @locked_ttl_ms::bigint = 0 THEN NULL ELSE locked_at + interval '1 milliseconds' * @locked_ttl_ms::bigint END
491491
WHERE
492492
template_id = @template_id
493493
AND

‎coderd/workspaces.go

Lines changed: 7 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ import (
1212

1313
"github.com/go-chi/chi/v5"
1414
"github.com/google/uuid"
15-
"golang.org/x/exp/slices"
1615
"golang.org/x/xerrors"
1716

1817
"cdr.dev/slog"
@@ -1120,6 +1119,11 @@ func convertWorkspace(
11201119
lockedAt = &workspace.LockedAt.Time
11211120
}
11221121

1122+
var deletedAt *time.Time
1123+
if workspace.DeletingAt.Valid {
1124+
deletedAt = &workspace.DeletingAt.Time
1125+
}
1126+
11231127
failingAgents := []uuid.UUID{}
11241128
for _, resource := range workspaceBuild.Resources {
11251129
for _, agent := range resource.Agents {
@@ -1130,8 +1134,7 @@ func convertWorkspace(
11301134
}
11311135

11321136
var (
1133-
ttlMillis = convertWorkspaceTTLMillis(workspace.Ttl)
1134-
deletingAt = calculateDeletingAt(workspace, template, workspaceBuild)
1137+
ttlMillis = convertWorkspaceTTLMillis(workspace.Ttl)
11351138
)
11361139

11371140
return codersdk.Workspace{
@@ -1152,7 +1155,7 @@ func convertWorkspace(
11521155
AutostartSchedule: autostartSchedule,
11531156
TTLMillis: ttlMillis,
11541157
LastUsedAt: workspace.LastUsedAt,
1155-
DeletingAt: deletingAt,
1158+
DeletingAt: deletedAt,
11561159
LockedAt: lockedAt,
11571160
Health: codersdk.WorkspaceHealth{
11581161
Healthy: len(failingAgents) == 0,
@@ -1170,19 +1173,6 @@ func convertWorkspaceTTLMillis(i sql.NullInt64) *int64 {
11701173
return &millis
11711174
}
11721175

1173-
// Calculate the time of the upcoming workspace deletion, if applicable; otherwise, return nil.
1174-
// Workspaces may have impending deletions if InactivityTTL feature is turned on and the workspace is inactive.
1175-
func calculateDeletingAt(workspace database.Workspace, template database.Template, build codersdk.WorkspaceBuild) *time.Time {
1176-
inactiveStatuses := []codersdk.WorkspaceStatus{codersdk.WorkspaceStatusStopped, codersdk.WorkspaceStatusCanceled, codersdk.WorkspaceStatusFailed, codersdk.WorkspaceStatusDeleted}
1177-
isInactive := slices.Contains(inactiveStatuses, build.Status)
1178-
// If InactivityTTL is turned off (set to 0) or if the workspace is active, there is no impending deletion
1179-
if template.InactivityTTL == 0 || !isInactive {
1180-
return nil
1181-
}
1182-
1183-
return ptr.Ref(workspace.LastUsedAt.Add(time.Duration(template.InactivityTTL) * time.Nanosecond))
1184-
}
1185-
11861176
func validWorkspaceTTLMillis(millis *int64, templateDefault, templateMax time.Duration) (sql.NullInt64, error) {
11871177
if templateDefault == 0 && templateMax != 0 || (templateMax > 0 && templateDefault > templateMax) {
11881178
templateDefault = templateMax

‎coderd/workspaces_internal_test.go

Lines changed: 0 additions & 82 deletions
This file was deleted.

‎coderd/workspaces_test.go

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2653,25 +2653,34 @@ func TestWorkspaceLock(t *testing.T) {
26532653
user = coderdtest.CreateFirstUser(t, client)
26542654
version = coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
26552655
_ = coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
2656-
template = coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
2657-
workspace = coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
2658-
_ = coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
2656+
lockedTTL = time.Minute
26592657
)
26602658

2659+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID, func(ctr *codersdk.CreateTemplateRequest) {
2660+
ctr.LockedTTLMillis = ptr.Ref[int64](lockedTTL.Milliseconds())
2661+
})
2662+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
2663+
_ = coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
2664+
26612665
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
26622666
defer cancel()
26632667

2668+
lastUsedAt := workspace.LastUsedAt
26642669
err := client.UpdateWorkspaceLock(ctx, workspace.ID, codersdk.UpdateWorkspaceLock{
26652670
Lock: true,
26662671
})
26672672
require.NoError(t, err)
26682673

26692674
workspace = coderdtest.MustWorkspace(t, client, workspace.ID)
26702675
require.NoError(t, err, "fetch provisioned workspace")
2676+
// The template doesn't have a locked_ttl set so this should be nil.
2677+
require.Nil(t, workspace.DeletingAt)
26712678
require.NotNil(t, workspace.LockedAt)
26722679
require.WithinRange(t, *workspace.LockedAt, time.Now().Add(-time.Second*10), time.Now())
2680+
require.Equal(t, lastUsedAt, workspace.LastUsedAt)
26732681

2674-
lastUsedAt := workspace.LastUsedAt
2682+
workspace = coderdtest.MustWorkspace(t, client, workspace.ID)
2683+
lastUsedAt = workspace.LastUsedAt
26752684
err = client.UpdateWorkspaceLock(ctx, workspace.ID, codersdk.UpdateWorkspaceLock{
26762685
Lock: false,
26772686
})
@@ -2680,6 +2689,9 @@ func TestWorkspaceLock(t *testing.T) {
26802689
workspace, err = client.Workspace(ctx, workspace.ID)
26812690
require.NoError(t, err, "fetch provisioned workspace")
26822691
require.Nil(t, workspace.LockedAt)
2692+
// The template doesn't have a locked_ttl set so this should be nil.
2693+
require.Nil(t, workspace.DeletingAt)
2694+
// The last_used_at should get updated when we unlock the workspace.
26832695
require.True(t, workspace.LastUsedAt.After(lastUsedAt))
26842696
})
26852697

‎enterprise/coderd/workspaces_test.go

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -667,3 +667,65 @@ func TestWorkspacesFiltering(t *testing.T) {
667667
assert.Equal(t, workspace.ID, res.Workspaces[0].ID)
668668
})
669669
}
670+
671+
func TestWorkspaceLock(t *testing.T) {
672+
t.Parallel()
673+
674+
t.Run("TemplateLockedTTL", func(t *testing.T) {
675+
t.Parallel()
676+
var (
677+
client = coderdenttest.New(t, &coderdenttest.Options{
678+
Options: &coderdtest.Options{
679+
IncludeProvisionerDaemon: true,
680+
TemplateScheduleStore: &coderd.EnterpriseTemplateScheduleStore{},
681+
},
682+
})
683+
684+
user = coderdtest.CreateFirstUser(t, client)
685+
_ = coderdenttest.AddFullLicense(t, client)
686+
version = coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
687+
_ = coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
688+
lockedTTL = time.Minute
689+
)
690+
691+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID, func(ctr *codersdk.CreateTemplateRequest) {
692+
ctr.LockedTTLMillis = ptr.Ref[int64](lockedTTL.Milliseconds())
693+
})
694+
695+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
696+
_ = coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
697+
698+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
699+
defer cancel()
700+
701+
lastUsedAt := workspace.LastUsedAt
702+
err := client.UpdateWorkspaceLock(ctx, workspace.ID, codersdk.UpdateWorkspaceLock{
703+
Lock: true,
704+
})
705+
require.NoError(t, err)
706+
707+
workspace = coderdtest.MustWorkspace(t, client, workspace.ID)
708+
require.NoError(t, err, "fetch provisioned workspace")
709+
require.NotNil(t, workspace.DeletingAt)
710+
require.NotNil(t, workspace.LockedAt)
711+
require.Equal(t, workspace.LockedAt.Add(lockedTTL), *workspace.DeletingAt)
712+
require.WithinRange(t, *workspace.LockedAt, time.Now().Add(-time.Second*10), time.Now())
713+
// Locking a workspace shouldn't update the last_used_at.
714+
require.Equal(t, lastUsedAt, workspace.LastUsedAt)
715+
716+
workspace = coderdtest.MustWorkspace(t, client, workspace.ID)
717+
lastUsedAt = workspace.LastUsedAt
718+
err = client.UpdateWorkspaceLock(ctx, workspace.ID, codersdk.UpdateWorkspaceLock{
719+
Lock: false,
720+
})
721+
require.NoError(t, err)
722+
723+
workspace, err = client.Workspace(ctx, workspace.ID)
724+
require.NoError(t, err, "fetch provisioned workspace")
725+
require.Nil(t, workspace.LockedAt)
726+
// Unlocking a workspace should cause the deleting_at to be unset.
727+
require.Nil(t, workspace.DeletingAt)
728+
// The last_used_at should get updated when we unlock the workspace.
729+
require.True(t, workspace.LastUsedAt.After(lastUsedAt))
730+
})
731+
}

0 commit comments

Comments
 (0)