From c422f3408a301200323433793107db78d5571ae2 Mon Sep 17 00:00:00 2001 From: Marcin Tojek Date: Thu, 4 Jul 2024 13:13:33 +0200 Subject: [PATCH 1/7] feat: autostop workspaces owned by suspended users --- coderd/autobuild/lifecycle_executor.go | 10 +++++++--- coderd/database/dbmem/dbmem.go | 9 +++++++++ coderd/database/queries/workspaces.sql | 8 ++++++++ 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/coderd/autobuild/lifecycle_executor.go b/coderd/autobuild/lifecycle_executor.go index e0d804328b2d3..4bbbaba667c7e 100644 --- a/coderd/autobuild/lifecycle_executor.go +++ b/coderd/autobuild/lifecycle_executor.go @@ -316,7 +316,7 @@ func getNextTransition( error, ) { switch { - case isEligibleForAutostop(ws, latestBuild, latestJob, currentTick): + case isEligibleForAutostop(user, ws, latestBuild, latestJob, currentTick): return database.WorkspaceTransitionStop, database.BuildReasonAutostop, nil case isEligibleForAutostart(user, ws, latestBuild, latestJob, templateSchedule, currentTick): return database.WorkspaceTransitionStart, database.BuildReasonAutostart, nil @@ -376,8 +376,8 @@ func isEligibleForAutostart(user database.User, ws database.Workspace, build dat return !currentTick.Before(nextTransition) } -// isEligibleForAutostart returns true if the workspace should be autostopped. -func isEligibleForAutostop(ws database.Workspace, build database.WorkspaceBuild, job database.ProvisionerJob, currentTick time.Time) bool { +// isEligibleForAutostop returns true if the workspace should be autostopped. +func isEligibleForAutostop(user database.User, ws database.Workspace, build database.WorkspaceBuild, job database.ProvisionerJob, currentTick time.Time) bool { if job.JobStatus == database.ProvisionerJobStatusFailed { return false } @@ -387,6 +387,10 @@ func isEligibleForAutostop(ws database.Workspace, build database.WorkspaceBuild, return false } + if build.Transition == database.WorkspaceTransitionStart && user.Status == database.UserStatusSuspended { + return true + } + // A workspace must be started in order for it to be auto-stopped. return build.Transition == database.WorkspaceTransitionStart && !build.Deadline.IsZero() && diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go index d19d218556b8d..ec7becdfd39c9 100644 --- a/coderd/database/dbmem/dbmem.go +++ b/coderd/database/dbmem/dbmem.go @@ -5844,6 +5844,15 @@ func (q *FakeQuerier) GetWorkspacesEligibleForTransition(ctx context.Context, no workspaces = append(workspaces, workspace) continue } + + user, err := q.getUserByIDNoLock(workspace.OwnerID) + if err != nil { + return nil, xerrors.Errorf("get user by ID: %w", err) + } + if user.Status == database.UserStatusSuspended && build.Transition == database.WorkspaceTransitionStart { + workspaces = append(workspaces, workspace) + continue + } } return workspaces, nil diff --git a/coderd/database/queries/workspaces.sql b/coderd/database/queries/workspaces.sql index 616e83a2bae16..ec8767e1f2be5 100644 --- a/coderd/database/queries/workspaces.sql +++ b/coderd/database/queries/workspaces.sql @@ -557,6 +557,8 @@ INNER JOIN provisioner_jobs ON workspace_builds.job_id = provisioner_jobs.id INNER JOIN templates ON workspaces.template_id = templates.id +INNER JOIN + users ON workspaces.owner_id = users.id WHERE workspace_builds.build_number = ( SELECT @@ -608,6 +610,12 @@ WHERE ( templates.time_til_dormant_autodelete > 0 AND workspaces.dormant_at IS NOT NULL + ) OR + + -- If the user account is suspended, and the workspace is running. + ( + users.status = 'suspended'::user_status AND + workspace_builds.transition = 'start'::workspace_transition ) ) AND workspaces.deleted = 'false'; From 63ddb62b0b5b79b84ab79d2de35a00673f08c27b Mon Sep 17 00:00:00 2001 From: Marcin Tojek Date: Thu, 4 Jul 2024 13:36:47 +0200 Subject: [PATCH 2/7] Version? --- coderd/database/queries.sql.go | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index ff7b7f6f955bd..e66a0b24e6cf4 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -3079,7 +3079,7 @@ func (q *sqlQuerier) GetJFrogXrayScanByWorkspaceAndAgentID(ctx context.Context, } const upsertJFrogXrayScanByWorkspaceAndAgentID = `-- name: UpsertJFrogXrayScanByWorkspaceAndAgentID :exec -INSERT INTO +INSERT INTO jfrog_xray_scans ( agent_id, workspace_id, @@ -3088,7 +3088,7 @@ INSERT INTO medium, results_url ) -VALUES +VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT (agent_id, workspace_id) DO UPDATE SET critical = $3, high = $4, medium = $5, results_url = $6 @@ -13426,6 +13426,8 @@ INNER JOIN provisioner_jobs ON workspace_builds.job_id = provisioner_jobs.id INNER JOIN templates ON workspaces.template_id = templates.id +INNER JOIN + users ON workspaces.owner_id = users.id WHERE workspace_builds.build_number = ( SELECT @@ -13477,6 +13479,12 @@ WHERE ( templates.time_til_dormant_autodelete > 0 AND workspaces.dormant_at IS NOT NULL + ) OR + + -- If the user account is suspended, and the workspace is running. + ( + users.status = 'suspended'::user_status AND + workspace_builds.transition = 'start'::workspace_transition ) ) AND workspaces.deleted = 'false' ` From 627f68a7899dfacb853fc16b745ef40aefe6cd99 Mon Sep 17 00:00:00 2001 From: Marcin Tojek Date: Thu, 4 Jul 2024 14:20:57 +0200 Subject: [PATCH 3/7] unit test --- coderd/autobuild/lifecycle_executor_test.go | 47 +++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/coderd/autobuild/lifecycle_executor_test.go b/coderd/autobuild/lifecycle_executor_test.go index 54ceb53254680..e37e1a3e9e5d8 100644 --- a/coderd/autobuild/lifecycle_executor_test.go +++ b/coderd/autobuild/lifecycle_executor_test.go @@ -563,6 +563,53 @@ func TestExecutorWorkspaceAutostopBeforeDeadline(t *testing.T) { assert.Len(t, stats.Transitions, 0) } +func TestExecuteAutostopSuspendedUser(t *testing.T) { + t.Parallel() + + var ( + ctx = testutil.Context(t, testutil.WaitShort) + sched = mustSchedule(t, "CRON_TZ=UTC 0 * * * *") + tickCh = make(chan time.Time) + statsCh = make(chan autobuild.Stats) + client = coderdtest.New(t, &coderdtest.Options{ + AutobuildTicker: tickCh, + IncludeProvisionerDaemon: true, + AutobuildStats: statsCh, + }) + ) + + admin := coderdtest.CreateFirstUser(t, client) + version := coderdtest.CreateTemplateVersion(t, client, admin.OrganizationID, nil) + coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) + template := coderdtest.CreateTemplate(t, client, admin.OrganizationID, version.ID) + userClient, user := coderdtest.CreateAnotherUser(t, client, admin.OrganizationID) + workspace := coderdtest.CreateWorkspace(t, userClient, admin.OrganizationID, template.ID, func(cwr *codersdk.CreateWorkspaceRequest) { + cwr.AutostartSchedule = ptr.Ref(sched.String()) + }) + coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, workspace.LatestBuild.ID) + + // Given: workspace is started, and the user is suspended. + workspace = coderdtest.MustWorkspace(t, userClient, workspace.ID) + _, err := client.UpdateUserStatus(ctx, user.ID.String(), codersdk.UserStatusSuspended) + require.NoError(t, err, "update user status") + + // When: the autobuild executor ticks after the scheduled time + go func() { + tickCh <- sched.Next(workspace.LatestBuild.CreatedAt) + close(tickCh) + }() + + // Then: nothing should happen + stats := <-statsCh + assert.Len(t, stats.Errors, 0) + assert.Len(t, stats.Transitions, 1) + assert.Equal(t, stats.Transitions[workspace.ID], database.WorkspaceTransitionStop) + + // Wait for stop to complete + workspace = coderdtest.MustWorkspace(t, client, workspace.ID) + _ = coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID) +} + func TestExecutorWorkspaceAutostopNoWaitChangedMyMind(t *testing.T) { t.Parallel() From 1bb5f920495d4e8c454e256dd4589ab78ed44e0d Mon Sep 17 00:00:00 2001 From: Marcin Tojek Date: Thu, 4 Jul 2024 12:31:56 +0000 Subject: [PATCH 4/7] make gen --- coderd/database/queries.sql.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index e66a0b24e6cf4..cd48412c2ff40 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -3079,7 +3079,7 @@ func (q *sqlQuerier) GetJFrogXrayScanByWorkspaceAndAgentID(ctx context.Context, } const upsertJFrogXrayScanByWorkspaceAndAgentID = `-- name: UpsertJFrogXrayScanByWorkspaceAndAgentID :exec -INSERT INTO +INSERT INTO jfrog_xray_scans ( agent_id, workspace_id, @@ -3088,7 +3088,7 @@ INSERT INTO medium, results_url ) -VALUES +VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT (agent_id, workspace_id) DO UPDATE SET critical = $3, high = $4, medium = $5, results_url = $6 From 7c35a55e80f606dbb39495a5d1016079ffcbb120 Mon Sep 17 00:00:00 2001 From: Marcin Tojek Date: Thu, 4 Jul 2024 14:47:37 +0200 Subject: [PATCH 5/7] fix: comment --- coderd/autobuild/lifecycle_executor_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/coderd/autobuild/lifecycle_executor_test.go b/coderd/autobuild/lifecycle_executor_test.go index e37e1a3e9e5d8..cf945bb0f31d9 100644 --- a/coderd/autobuild/lifecycle_executor_test.go +++ b/coderd/autobuild/lifecycle_executor_test.go @@ -599,7 +599,7 @@ func TestExecuteAutostopSuspendedUser(t *testing.T) { close(tickCh) }() - // Then: nothing should happen + // Then: the workspace should be stopped stats := <-statsCh assert.Len(t, stats.Errors, 0) assert.Len(t, stats.Transitions, 1) From 9e952072dc9849b23c28bbe4ccbbc3f4e7997f15 Mon Sep 17 00:00:00 2001 From: Marcin Tojek Date: Thu, 4 Jul 2024 15:18:59 +0200 Subject: [PATCH 6/7] Validate workspace status --- coderd/autobuild/lifecycle_executor_test.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/coderd/autobuild/lifecycle_executor_test.go b/coderd/autobuild/lifecycle_executor_test.go index cf945bb0f31d9..c229ca844c8c7 100644 --- a/coderd/autobuild/lifecycle_executor_test.go +++ b/coderd/autobuild/lifecycle_executor_test.go @@ -588,8 +588,9 @@ func TestExecuteAutostopSuspendedUser(t *testing.T) { }) coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, workspace.LatestBuild.ID) - // Given: workspace is started, and the user is suspended. + // Given: workspace is running, and the user is suspended. workspace = coderdtest.MustWorkspace(t, userClient, workspace.ID) + require.Equal(t, codersdk.WorkspaceStatusRunning, workspace.LatestBuild.Status) _, err := client.UpdateUserStatus(ctx, user.ID.String(), codersdk.UserStatusSuspended) require.NoError(t, err, "update user status") @@ -607,7 +608,8 @@ func TestExecuteAutostopSuspendedUser(t *testing.T) { // Wait for stop to complete workspace = coderdtest.MustWorkspace(t, client, workspace.ID) - _ = coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID) + workspaceBuild := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID) + assert.Equal(t, codersdk.WorkspaceStatusStopped, workspaceBuild.Status) } func TestExecutorWorkspaceAutostopNoWaitChangedMyMind(t *testing.T) { From b9908cfa833444616d26a9337435c1e61ba99e47 Mon Sep 17 00:00:00 2001 From: Marcin Tojek Date: Thu, 4 Jul 2024 15:25:45 +0200 Subject: [PATCH 7/7] remove sched --- coderd/autobuild/lifecycle_executor_test.go | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/coderd/autobuild/lifecycle_executor_test.go b/coderd/autobuild/lifecycle_executor_test.go index c229ca844c8c7..bc480b97e4aa2 100644 --- a/coderd/autobuild/lifecycle_executor_test.go +++ b/coderd/autobuild/lifecycle_executor_test.go @@ -568,7 +568,6 @@ func TestExecuteAutostopSuspendedUser(t *testing.T) { var ( ctx = testutil.Context(t, testutil.WaitShort) - sched = mustSchedule(t, "CRON_TZ=UTC 0 * * * *") tickCh = make(chan time.Time) statsCh = make(chan autobuild.Stats) client = coderdtest.New(t, &coderdtest.Options{ @@ -583,9 +582,7 @@ func TestExecuteAutostopSuspendedUser(t *testing.T) { coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) template := coderdtest.CreateTemplate(t, client, admin.OrganizationID, version.ID) userClient, user := coderdtest.CreateAnotherUser(t, client, admin.OrganizationID) - workspace := coderdtest.CreateWorkspace(t, userClient, admin.OrganizationID, template.ID, func(cwr *codersdk.CreateWorkspaceRequest) { - cwr.AutostartSchedule = ptr.Ref(sched.String()) - }) + workspace := coderdtest.CreateWorkspace(t, userClient, admin.OrganizationID, template.ID) coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, workspace.LatestBuild.ID) // Given: workspace is running, and the user is suspended. @@ -596,7 +593,7 @@ func TestExecuteAutostopSuspendedUser(t *testing.T) { // When: the autobuild executor ticks after the scheduled time go func() { - tickCh <- sched.Next(workspace.LatestBuild.CreatedAt) + tickCh <- time.Unix(0, 0) // the exact time is not important close(tickCh) }()