diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index c64701fde482a..67b6d41970fa5 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -412,6 +412,21 @@ var ( policy.ActionCreate, policy.ActionDelete, policy.ActionRead, policy.ActionUpdate, policy.ActionWorkspaceStart, policy.ActionWorkspaceStop, }, + // Should be able to add the prebuilds system user as a member to any organization that needs prebuilds. + rbac.ResourceOrganizationMember.Type: { + policy.ActionCreate, + }, + // Needs to be able to assign roles to the system user in order to make it a member of an organization. + rbac.ResourceAssignOrgRole.Type: { + policy.ActionAssign, + }, + // Needs to be able to read users to determine which organizations the prebuild system user is a member of. + rbac.ResourceUser.Type: { + policy.ActionRead, + }, + rbac.ResourceOrganization.Type: { + policy.ActionRead, + }, }), }, }), diff --git a/coderd/rbac/roles.go b/coderd/rbac/roles.go index 8b98f5f2f2bc7..28ddc38462ce9 100644 --- a/coderd/rbac/roles.go +++ b/coderd/rbac/roles.go @@ -33,6 +33,8 @@ const ( orgUserAdmin string = "organization-user-admin" orgTemplateAdmin string = "organization-template-admin" orgWorkspaceCreationBan string = "organization-workspace-creation-ban" + + prebuildsOrchestrator string = "prebuilds-orchestrator" ) func init() { @@ -599,6 +601,9 @@ var assignRoles = map[string]map[string]bool{ orgUserAdmin: { orgMember: true, }, + prebuildsOrchestrator: { + orgMember: true, + }, } // ExpandableRoles is any type that can be expanded into a []Role. This is implemented diff --git a/enterprise/coderd/prebuilds/claim_test.go b/enterprise/coderd/prebuilds/claim_test.go index 5a18600a84602..83933f3a98cd3 100644 --- a/enterprise/coderd/prebuilds/claim_test.go +++ b/enterprise/coderd/prebuilds/claim_test.go @@ -19,13 +19,16 @@ import ( "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" agplprebuilds "github.com/coder/coder/v2/coderd/prebuilds" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" + "github.com/coder/coder/v2/enterprise/coderd/license" "github.com/coder/coder/v2/enterprise/coderd/prebuilds" "github.com/coder/coder/v2/provisioner/echo" + "github.com/coder/coder/v2/provisionersdk" "github.com/coder/coder/v2/provisionersdk/proto" "github.com/coder/coder/v2/testutil" ) @@ -105,7 +108,6 @@ func TestClaimPrebuild(t *testing.T) { expectPrebuildClaimed: true, markPrebuildsClaimable: true, }, - "no claimable prebuilt workspaces error is returned": { expectPrebuildClaimed: false, markPrebuildsClaimable: true, @@ -124,227 +126,248 @@ func TestClaimPrebuild(t *testing.T) { } for name, tc := range cases { - tc := tc - - t.Run(name, func(t *testing.T) { - t.Parallel() - - // Setup. - ctx := testutil.Context(t, testutil.WaitSuperLong) - db, pubsub := dbtestutil.NewDB(t) - - spy := newStoreSpy(db, tc.claimingErr) - expectedPrebuildsCount := desiredInstances * presetCount + // Ensure that prebuilt workspaces can be claimed in non-default organizations: + for _, useDefaultOrg := range []bool{true, false} { + tc := tc + t.Run(name, func(t *testing.T) { + t.Parallel() + + // Setup. + ctx := testutil.Context(t, testutil.WaitSuperLong) + db, pubsub := dbtestutil.NewDB(t) + + spy := newStoreSpy(db, tc.claimingErr) + expectedPrebuildsCount := desiredInstances * presetCount + + logger := testutil.Logger(t) + client, _, api, owner := coderdenttest.NewWithAPI(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + Database: spy, + Pubsub: pubsub, + }, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureExternalProvisionerDaemons: 1, + }, + }, - logger := testutil.Logger(t) - client, _, api, owner := coderdenttest.NewWithAPI(t, &coderdenttest.Options{ - Options: &coderdtest.Options{ - IncludeProvisionerDaemon: true, - Database: spy, - Pubsub: pubsub, - }, + EntitlementsUpdateInterval: time.Second, + }) - EntitlementsUpdateInterval: time.Second, - }) + orgID := owner.OrganizationID + if !useDefaultOrg { + secondOrg := dbgen.Organization(t, db, database.Organization{}) + orgID = secondOrg.ID + } - reconciler := prebuilds.NewStoreReconciler(spy, pubsub, codersdk.PrebuildsConfig{}, logger, quartz.NewMock(t), prometheus.NewRegistry(), newNoopEnqueuer()) - var claimer agplprebuilds.Claimer = prebuilds.NewEnterpriseClaimer(spy) - api.AGPL.PrebuildsClaimer.Store(&claimer) + provisionerCloser := coderdenttest.NewExternalProvisionerDaemon(t, client, orgID, map[string]string{ + provisionersdk.TagScope: provisionersdk.ScopeOrganization, + }) + defer provisionerCloser.Close() - version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, templateWithAgentAndPresetsWithPrebuilds(desiredInstances)) - _ = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) - coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID) - presets, err := client.TemplateVersionPresets(ctx, version.ID) - require.NoError(t, err) - require.Len(t, presets, presetCount) + reconciler := prebuilds.NewStoreReconciler(spy, pubsub, codersdk.PrebuildsConfig{}, logger, quartz.NewMock(t), prometheus.NewRegistry(), newNoopEnqueuer()) + var claimer agplprebuilds.Claimer = prebuilds.NewEnterpriseClaimer(spy) + api.AGPL.PrebuildsClaimer.Store(&claimer) - userClient, user := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID, rbac.RoleMember()) + version := coderdtest.CreateTemplateVersion(t, client, orgID, templateWithAgentAndPresetsWithPrebuilds(desiredInstances)) + _ = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) + coderdtest.CreateTemplate(t, client, orgID, version.ID) + presets, err := client.TemplateVersionPresets(ctx, version.ID) + require.NoError(t, err) + require.Len(t, presets, presetCount) - // Given: the reconciliation state is snapshot. - state, err := reconciler.SnapshotState(ctx, spy) - require.NoError(t, err) - require.Len(t, state.Presets, presetCount) + userClient, user := coderdtest.CreateAnotherUser(t, client, orgID, rbac.RoleMember()) - // When: a reconciliation is setup for each preset. - for _, preset := range presets { - ps, err := state.FilterByPreset(preset.ID) - require.NoError(t, err) - require.NotNil(t, ps) - actions, err := reconciler.CalculateActions(ctx, *ps) + // Given: the reconciliation state is snapshot. + state, err := reconciler.SnapshotState(ctx, spy) require.NoError(t, err) - require.NotNil(t, actions) + require.Len(t, state.Presets, presetCount) - require.NoError(t, reconciler.ReconcilePreset(ctx, *ps)) - } + // When: a reconciliation is setup for each preset. + for _, preset := range presets { + ps, err := state.FilterByPreset(preset.ID) + require.NoError(t, err) + require.NotNil(t, ps) + actions, err := reconciler.CalculateActions(ctx, *ps) + require.NoError(t, err) + require.NotNil(t, actions) - // Given: a set of running, eligible prebuilds eventually starts up. - runningPrebuilds := make(map[uuid.UUID]database.GetRunningPrebuiltWorkspacesRow, desiredInstances*presetCount) - require.Eventually(t, func() bool { - rows, err := spy.GetRunningPrebuiltWorkspaces(ctx) - if err != nil { - return false + require.NoError(t, reconciler.ReconcilePreset(ctx, *ps)) } - for _, row := range rows { - runningPrebuilds[row.CurrentPresetID.UUID] = row - - if !tc.markPrebuildsClaimable { - continue - } - - agents, err := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, row.ID) + // Given: a set of running, eligible prebuilds eventually starts up. + runningPrebuilds := make(map[uuid.UUID]database.GetRunningPrebuiltWorkspacesRow, desiredInstances*presetCount) + require.Eventually(t, func() bool { + rows, err := spy.GetRunningPrebuiltWorkspaces(ctx) if err != nil { return false } - // Workspaces are eligible once its agent is marked "ready". - for _, agent := range agents { - err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{ - ID: agent.ID, - LifecycleState: database.WorkspaceAgentLifecycleStateReady, - StartedAt: sql.NullTime{Time: time.Now().Add(time.Hour), Valid: true}, - ReadyAt: sql.NullTime{Time: time.Now().Add(-1 * time.Hour), Valid: true}, - }) + for _, row := range rows { + runningPrebuilds[row.CurrentPresetID.UUID] = row + + if !tc.markPrebuildsClaimable { + continue + } + + agents, err := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, row.ID) if err != nil { return false } - } - } - t.Logf("found %d running prebuilds so far, want %d", len(runningPrebuilds), expectedPrebuildsCount) - - return len(runningPrebuilds) == expectedPrebuildsCount - }, testutil.WaitSuperLong, testutil.IntervalSlow) - - // When: a user creates a new workspace with a preset for which prebuilds are configured. - workspaceName := strings.ReplaceAll(testutil.GetRandomName(t), "_", "-") - params := database.ClaimPrebuiltWorkspaceParams{ - NewUserID: user.ID, - NewName: workspaceName, - PresetID: presets[0].ID, - } - userWorkspace, err := userClient.CreateUserWorkspace(ctx, user.Username, codersdk.CreateWorkspaceRequest{ - TemplateVersionID: version.ID, - Name: workspaceName, - TemplateVersionPresetID: presets[0].ID, - }) + // Workspaces are eligible once its agent is marked "ready". + for _, agent := range agents { + err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{ + ID: agent.ID, + LifecycleState: database.WorkspaceAgentLifecycleStateReady, + StartedAt: sql.NullTime{Time: time.Now().Add(time.Hour), Valid: true}, + ReadyAt: sql.NullTime{Time: time.Now().Add(-1 * time.Hour), Valid: true}, + }) + if err != nil { + return false + } + } + } - isNoPrebuiltWorkspaces := errors.Is(tc.claimingErr, agplprebuilds.ErrNoClaimablePrebuiltWorkspaces) - isUnsupported := errors.Is(tc.claimingErr, agplprebuilds.ErrAGPLDoesNotSupportPrebuiltWorkspaces) + t.Logf("found %d running prebuilds so far, want %d", len(runningPrebuilds), expectedPrebuildsCount) - switch { - case tc.claimingErr != nil && (isNoPrebuiltWorkspaces || isUnsupported): - require.NoError(t, err) - coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, userWorkspace.LatestBuild.ID) + return len(runningPrebuilds) == expectedPrebuildsCount + }, testutil.WaitSuperLong, testutil.IntervalSlow) - // Then: the number of running prebuilds hasn't changed because claiming prebuild is failed and we fallback to creating new workspace. - currentPrebuilds, err := spy.GetRunningPrebuiltWorkspaces(ctx) - require.NoError(t, err) - require.Equal(t, expectedPrebuildsCount, len(currentPrebuilds)) - return + // When: a user creates a new workspace with a preset for which prebuilds are configured. + workspaceName := strings.ReplaceAll(testutil.GetRandomName(t), "_", "-") + params := database.ClaimPrebuiltWorkspaceParams{ + NewUserID: user.ID, + NewName: workspaceName, + PresetID: presets[0].ID, + } + userWorkspace, err := userClient.CreateUserWorkspace(ctx, user.Username, codersdk.CreateWorkspaceRequest{ + TemplateVersionID: version.ID, + Name: workspaceName, + TemplateVersionPresetID: presets[0].ID, + }) - case tc.claimingErr != nil && errors.Is(tc.claimingErr, unexpectedClaimingError): - // Then: unexpected error happened and was propagated all the way to the caller - require.Error(t, err) - require.ErrorContains(t, err, unexpectedClaimingError.Error()) + isNoPrebuiltWorkspaces := errors.Is(tc.claimingErr, agplprebuilds.ErrNoClaimablePrebuiltWorkspaces) + isUnsupported := errors.Is(tc.claimingErr, agplprebuilds.ErrAGPLDoesNotSupportPrebuiltWorkspaces) - // Then: the number of running prebuilds hasn't changed because claiming prebuild is failed. - currentPrebuilds, err := spy.GetRunningPrebuiltWorkspaces(ctx) - require.NoError(t, err) - require.Equal(t, expectedPrebuildsCount, len(currentPrebuilds)) - return + switch { + case tc.claimingErr != nil && (isNoPrebuiltWorkspaces || isUnsupported): + require.NoError(t, err) + build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, userWorkspace.LatestBuild.ID) + _ = build - default: - // tc.claimingErr is nil scenario - require.NoError(t, err) - coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, userWorkspace.LatestBuild.ID) - } + // Then: the number of running prebuilds hasn't changed because claiming prebuild is failed and we fallback to creating new workspace. + currentPrebuilds, err := spy.GetRunningPrebuiltWorkspaces(ctx) + require.NoError(t, err) + require.Equal(t, expectedPrebuildsCount, len(currentPrebuilds)) + return - // at this point we know that tc.claimingErr is nil + case tc.claimingErr != nil && errors.Is(tc.claimingErr, unexpectedClaimingError): + // Then: unexpected error happened and was propagated all the way to the caller + require.Error(t, err) + require.ErrorContains(t, err, unexpectedClaimingError.Error()) - // Then: a prebuild should have been claimed. - require.EqualValues(t, spy.claims.Load(), 1) - require.EqualValues(t, *spy.claimParams.Load(), params) + // Then: the number of running prebuilds hasn't changed because claiming prebuild is failed. + currentPrebuilds, err := spy.GetRunningPrebuiltWorkspaces(ctx) + require.NoError(t, err) + require.Equal(t, expectedPrebuildsCount, len(currentPrebuilds)) + return - if !tc.expectPrebuildClaimed { - require.Nil(t, spy.claimedWorkspace.Load()) - return - } + default: + // tc.claimingErr is nil scenario + require.NoError(t, err) + build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, userWorkspace.LatestBuild.ID) + require.Equal(t, build.Job.Status, codersdk.ProvisionerJobSucceeded) + } - require.NotNil(t, spy.claimedWorkspace.Load()) - claimed := *spy.claimedWorkspace.Load() - require.NotEqual(t, claimed.ID, uuid.Nil) + // at this point we know that tc.claimingErr is nil - // Then: the claimed prebuild must now be owned by the requester. - workspace, err := spy.GetWorkspaceByID(ctx, claimed.ID) - require.NoError(t, err) - require.Equal(t, user.ID, workspace.OwnerID) + // Then: a prebuild should have been claimed. + require.EqualValues(t, spy.claims.Load(), 1) + require.EqualValues(t, *spy.claimParams.Load(), params) - // Then: the number of running prebuilds has changed since one was claimed. - currentPrebuilds, err := spy.GetRunningPrebuiltWorkspaces(ctx) - require.NoError(t, err) - require.Equal(t, expectedPrebuildsCount-1, len(currentPrebuilds)) + if !tc.expectPrebuildClaimed { + require.Nil(t, spy.claimedWorkspace.Load()) + return + } - // Then: the claimed prebuild is now missing from the running prebuilds set. - found := slices.ContainsFunc(currentPrebuilds, func(prebuild database.GetRunningPrebuiltWorkspacesRow) bool { - return prebuild.ID == claimed.ID - }) - require.False(t, found, "claimed prebuild should not still be considered a running prebuild") + require.NotNil(t, spy.claimedWorkspace.Load()) + claimed := *spy.claimedWorkspace.Load() + require.NotEqual(t, claimed.ID, uuid.Nil) - // Then: reconciling at this point will provision a new prebuild to replace the claimed one. - { - // Given: the reconciliation state is snapshot. - state, err = reconciler.SnapshotState(ctx, spy) + // Then: the claimed prebuild must now be owned by the requester. + workspace, err := spy.GetWorkspaceByID(ctx, claimed.ID) require.NoError(t, err) + require.Equal(t, user.ID, workspace.OwnerID) - // When: a reconciliation is setup for each preset. - for _, preset := range presets { - ps, err := state.FilterByPreset(preset.ID) + // Then: the number of running prebuilds has changed since one was claimed. + currentPrebuilds, err := spy.GetRunningPrebuiltWorkspaces(ctx) + require.NoError(t, err) + require.Equal(t, expectedPrebuildsCount-1, len(currentPrebuilds)) + + // Then: the claimed prebuild is now missing from the running prebuilds set. + found := slices.ContainsFunc(currentPrebuilds, func(prebuild database.GetRunningPrebuiltWorkspacesRow) bool { + return prebuild.ID == claimed.ID + }) + require.False(t, found, "claimed prebuild should not still be considered a running prebuild") + + // Then: reconciling at this point will provision a new prebuild to replace the claimed one. + { + // Given: the reconciliation state is snapshot. + state, err = reconciler.SnapshotState(ctx, spy) require.NoError(t, err) - // Then: the reconciliation takes place without error. - require.NoError(t, reconciler.ReconcilePreset(ctx, *ps)) - } - } + // When: a reconciliation is setup for each preset. + for _, preset := range presets { + ps, err := state.FilterByPreset(preset.ID) + require.NoError(t, err) - require.Eventually(t, func() bool { - rows, err := spy.GetRunningPrebuiltWorkspaces(ctx) - if err != nil { - return false + // Then: the reconciliation takes place without error. + require.NoError(t, reconciler.ReconcilePreset(ctx, *ps)) + } } - t.Logf("found %d running prebuilds so far, want %d", len(rows), expectedPrebuildsCount) + require.Eventually(t, func() bool { + rows, err := spy.GetRunningPrebuiltWorkspaces(ctx) + if err != nil { + return false + } - return len(runningPrebuilds) == expectedPrebuildsCount - }, testutil.WaitSuperLong, testutil.IntervalSlow) + t.Logf("found %d running prebuilds so far, want %d", len(rows), expectedPrebuildsCount) - // Then: when restarting the created workspace (which claimed a prebuild), it should not try and claim a new prebuild. - // Prebuilds should ONLY be used for net-new workspaces. - // This is expected by default anyway currently since new workspaces and operations on existing workspaces - // take different code paths, but it's worth validating. + return len(runningPrebuilds) == expectedPrebuildsCount + }, testutil.WaitSuperLong, testutil.IntervalSlow) - spy.claims.Store(0) // Reset counter because we need to check if any new claim requests happen. + // Then: when restarting the created workspace (which claimed a prebuild), it should not try and claim a new prebuild. + // Prebuilds should ONLY be used for net-new workspaces. + // This is expected by default anyway currently since new workspaces and operations on existing workspaces + // take different code paths, but it's worth validating. - wp, err := userClient.WorkspaceBuildParameters(ctx, userWorkspace.LatestBuild.ID) - require.NoError(t, err) + spy.claims.Store(0) // Reset counter because we need to check if any new claim requests happen. - stopBuild, err := userClient.CreateWorkspaceBuild(ctx, workspace.ID, codersdk.CreateWorkspaceBuildRequest{ - TemplateVersionID: version.ID, - Transition: codersdk.WorkspaceTransitionStop, - }) - require.NoError(t, err) - coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, stopBuild.ID) + wp, err := userClient.WorkspaceBuildParameters(ctx, userWorkspace.LatestBuild.ID) + require.NoError(t, err) - startBuild, err := userClient.CreateWorkspaceBuild(ctx, workspace.ID, codersdk.CreateWorkspaceBuildRequest{ - TemplateVersionID: version.ID, - Transition: codersdk.WorkspaceTransitionStart, - RichParameterValues: wp, - }) - require.NoError(t, err) - coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, startBuild.ID) + stopBuild, err := userClient.CreateWorkspaceBuild(ctx, workspace.ID, codersdk.CreateWorkspaceBuildRequest{ + TemplateVersionID: version.ID, + Transition: codersdk.WorkspaceTransitionStop, + }) + require.NoError(t, err) + build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, stopBuild.ID) + require.Equal(t, build.Job.Status, codersdk.ProvisionerJobSucceeded) + + startBuild, err := userClient.CreateWorkspaceBuild(ctx, workspace.ID, codersdk.CreateWorkspaceBuildRequest{ + TemplateVersionID: version.ID, + Transition: codersdk.WorkspaceTransitionStart, + RichParameterValues: wp, + }) + require.NoError(t, err) + build = coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, startBuild.ID) + require.Equal(t, build.Job.Status, codersdk.ProvisionerJobSucceeded) - require.Zero(t, spy.claims.Load()) - }) + require.Zero(t, spy.claims.Load()) + }) + } } } diff --git a/enterprise/coderd/prebuilds/membership.go b/enterprise/coderd/prebuilds/membership.go new file mode 100644 index 0000000000000..079711bcbcc49 --- /dev/null +++ b/enterprise/coderd/prebuilds/membership.go @@ -0,0 +1,81 @@ +package prebuilds + +import ( + "context" + "database/sql" + "errors" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/quartz" +) + +// StoreMembershipReconciler encapsulates the responsibility of ensuring that the prebuilds system user is a member of all +// organizations for which prebuilt workspaces are requested. This is necessary because our data model requires that such +// prebuilt workspaces belong to a member of the organization of their eventual claimant. +type StoreMembershipReconciler struct { + store database.Store + clock quartz.Clock +} + +func NewStoreMembershipReconciler(store database.Store, clock quartz.Clock) StoreMembershipReconciler { + return StoreMembershipReconciler{ + store: store, + clock: clock, + } +} + +// ReconcileAll compares the current membership of a user to the membership required in order to create prebuilt workspaces. +// If the user in question is not yet a member of an organization that needs prebuilt workspaces, ReconcileAll will create +// the membership required. +// +// This method does not have an opinion on transaction or lock management. These responsibilities are left to the caller. +func (s StoreMembershipReconciler) ReconcileAll(ctx context.Context, userID uuid.UUID, presets []database.GetTemplatePresetsWithPrebuildsRow) error { + organizationMemberships, err := s.store.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{ + UserID: userID, + Deleted: sql.NullBool{ + Bool: false, + Valid: true, + }, + }) + if err != nil { + return xerrors.Errorf("determine prebuild organization membership: %w", err) + } + + systemUserMemberships := make(map[uuid.UUID]struct{}, 0) + defaultOrg, err := s.store.GetDefaultOrganization(ctx) + if err != nil { + return xerrors.Errorf("get default organization: %w", err) + } + systemUserMemberships[defaultOrg.ID] = struct{}{} + for _, o := range organizationMemberships { + systemUserMemberships[o.ID] = struct{}{} + } + + var membershipInsertionErrors error + for _, preset := range presets { + _, alreadyMember := systemUserMemberships[preset.OrganizationID] + if alreadyMember { + continue + } + // Add the organization to our list of memberships regardless of potential failure below + // to avoid a retry that will probably be doomed anyway. + systemUserMemberships[preset.OrganizationID] = struct{}{} + + // Insert the missing membership + _, err = s.store.InsertOrganizationMember(ctx, database.InsertOrganizationMemberParams{ + OrganizationID: preset.OrganizationID, + UserID: userID, + CreatedAt: s.clock.Now(), + UpdatedAt: s.clock.Now(), + Roles: []string{}, + }) + if err != nil { + membershipInsertionErrors = errors.Join(membershipInsertionErrors, xerrors.Errorf("insert membership for prebuilt workspaces: %w", err)) + continue + } + } + return membershipInsertionErrors +} diff --git a/enterprise/coderd/prebuilds/membership_test.go b/enterprise/coderd/prebuilds/membership_test.go new file mode 100644 index 0000000000000..6caa7178d9d60 --- /dev/null +++ b/enterprise/coderd/prebuilds/membership_test.go @@ -0,0 +1,127 @@ +package prebuilds_test + +import ( + "context" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/quartz" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + agplprebuilds "github.com/coder/coder/v2/coderd/prebuilds" + "github.com/coder/coder/v2/enterprise/coderd/prebuilds" +) + +// TestReconcileAll verifies that StoreMembershipReconciler correctly updates membership +// for the prebuilds system user. +func TestReconcileAll(t *testing.T) { + t.Parallel() + + ctx := context.Background() + clock := quartz.NewMock(t) + + // Helper to build a minimal Preset row belonging to a given org. + newPresetRow := func(orgID uuid.UUID) database.GetTemplatePresetsWithPrebuildsRow { + return database.GetTemplatePresetsWithPrebuildsRow{ + ID: uuid.New(), + OrganizationID: orgID, + } + } + + tests := []struct { + name string + includePreset bool + preExistingMembership bool + }{ + // The StoreMembershipReconciler acts based on the provided agplprebuilds.GlobalSnapshot. + // These test cases must therefore trust any valid snapshot, so the only relevant functional test cases are: + + // No presets to act on and the prebuilds user does not belong to any organizations. + // Reconciliation should be a no-op + {name: "no presets, no memberships", includePreset: false, preExistingMembership: false}, + // If we have a preset that requires prebuilds, but the prebuilds user is not a member of + // that organization, then we should add the membership. + {name: "preset, but no membership", includePreset: true, preExistingMembership: false}, + // If the prebuilds system user is already a member of the organization to which a preset belongs, + // then reconciliation should be a no-op: + {name: "preset, but already a member", includePreset: true, preExistingMembership: true}, + // If the prebuilds system user is a member of an organization that doesn't have need any prebuilds, + // then it must have required prebuilds in the past. The membership is not currently necessary, but + // the reconciler won't remove it, because there's little cost to keeping it and prebuilds might be + // enabled again. + {name: "member, but no presets", includePreset: false, preExistingMembership: true}, + } + + for _, tc := range tests { + tc := tc // capture + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + + defaultOrg, err := db.GetDefaultOrganization(ctx) + require.NoError(t, err) + + // introduce an unrelated organization to ensure that the membership reconciler don't interfere with it. + unrelatedOrg := dbgen.Organization(t, db, database.Organization{}) + targetOrg := dbgen.Organization(t, db, database.Organization{}) + + if !dbtestutil.WillUsePostgres() { + // dbmem doesn't ensure membership to the default organization + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + OrganizationID: defaultOrg.ID, + UserID: agplprebuilds.SystemUserID, + }) + } + + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: unrelatedOrg.ID, UserID: agplprebuilds.SystemUserID}) + if tc.preExistingMembership { + // System user already a member of both orgs. + dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: targetOrg.ID, UserID: agplprebuilds.SystemUserID}) + } + + presets := []database.GetTemplatePresetsWithPrebuildsRow{newPresetRow(unrelatedOrg.ID)} + if tc.includePreset { + presets = append(presets, newPresetRow(targetOrg.ID)) + } + + // Verify memberships before reconciliation. + preReconcileMemberships, err := db.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{ + UserID: agplprebuilds.SystemUserID, + }) + require.NoError(t, err) + expectedMembershipsBefore := []uuid.UUID{defaultOrg.ID, unrelatedOrg.ID} + if tc.preExistingMembership { + expectedMembershipsBefore = append(expectedMembershipsBefore, targetOrg.ID) + } + require.ElementsMatch(t, expectedMembershipsBefore, extractOrgIDs(preReconcileMemberships)) + + // Reconcile + reconciler := prebuilds.NewStoreMembershipReconciler(db, clock) + require.NoError(t, reconciler.ReconcileAll(ctx, agplprebuilds.SystemUserID, presets)) + + // Verify memberships after reconciliation. + postReconcileMemberships, err := db.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{ + UserID: agplprebuilds.SystemUserID, + }) + require.NoError(t, err) + expectedMembershipsAfter := expectedMembershipsBefore + if !tc.preExistingMembership && tc.includePreset { + expectedMembershipsAfter = append(expectedMembershipsAfter, targetOrg.ID) + } + require.ElementsMatch(t, expectedMembershipsAfter, extractOrgIDs(postReconcileMemberships)) + }) + } +} + +func extractOrgIDs(orgs []database.Organization) []uuid.UUID { + ids := make([]uuid.UUID, len(orgs)) + for i, o := range orgs { + ids[i] = o.ID + } + return ids +} diff --git a/enterprise/coderd/prebuilds/reconcile.go b/enterprise/coderd/prebuilds/reconcile.go index ebfcfaf2b3182..3a1ab66d009a7 100644 --- a/enterprise/coderd/prebuilds/reconcile.go +++ b/enterprise/coderd/prebuilds/reconcile.go @@ -251,8 +251,8 @@ func (c *StoreReconciler) ReconcileAll(ctx context.Context) error { logger.Debug(ctx, "starting reconciliation") - err := c.WithReconciliationLock(ctx, logger, func(ctx context.Context, db database.Store) error { - snapshot, err := c.SnapshotState(ctx, db) + err := c.WithReconciliationLock(ctx, logger, func(ctx context.Context, _ database.Store) error { + snapshot, err := c.SnapshotState(ctx, c.store) if err != nil { return xerrors.Errorf("determine current snapshot: %w", err) } @@ -264,6 +264,12 @@ func (c *StoreReconciler) ReconcileAll(ctx context.Context) error { return nil } + membershipReconciler := NewStoreMembershipReconciler(c.store, c.clock) + err = membershipReconciler.ReconcileAll(ctx, prebuilds.SystemUserID, snapshot.Presets) + if err != nil { + return xerrors.Errorf("reconcile prebuild membership: %w", err) + } + var eg errgroup.Group // Reconcile presets in parallel. Each preset in its own goroutine. for _, preset := range snapshot.Presets { diff --git a/enterprise/coderd/prebuilds/reconcile_test.go b/enterprise/coderd/prebuilds/reconcile_test.go index a0e1f9726d7d5..d2827999ba843 100644 --- a/enterprise/coderd/prebuilds/reconcile_test.go +++ b/enterprise/coderd/prebuilds/reconcile_test.go @@ -43,7 +43,7 @@ func TestNoReconciliationActionsIfNoPresets(t *testing.T) { t.Parallel() if !dbtestutil.WillUsePostgres() { - t.Skip("This test requires postgres") + t.Skip("dbmem times out on nesting transactions, postgres ignores the inner ones") } clock := quartz.NewMock(t) @@ -88,7 +88,7 @@ func TestNoReconciliationActionsIfNoPrebuilds(t *testing.T) { t.Parallel() if !dbtestutil.WillUsePostgres() { - t.Skip("This test requires postgres") + t.Skip("dbmem times out on nesting transactions, postgres ignores the inner ones") } clock := quartz.NewMock(t) diff --git a/enterprise/coderd/workspaceagents_test.go b/enterprise/coderd/workspaceagents_test.go index f0c9b37f3b2a3..1eea9ecda9ca8 100644 --- a/enterprise/coderd/workspaceagents_test.go +++ b/enterprise/coderd/workspaceagents_test.go @@ -11,7 +11,10 @@ import ( "testing" "time" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/provisionersdk" "github.com/coder/serpent" "github.com/google/uuid" @@ -84,8 +87,6 @@ func TestBlockNonBrowser(t *testing.T) { func TestReinitializeAgent(t *testing.T) { t.Parallel() - tempAgentLog := testutil.CreateTemp(t, "", "testReinitializeAgent") - if !dbtestutil.WillUsePostgres() { t.Skip("dbmem cannot currently claim a workspace") } @@ -94,159 +95,175 @@ func TestReinitializeAgent(t *testing.T) { t.Skip("test startup script is not supported on windows") } - startupScript := fmt.Sprintf("printenv >> %s; echo '---\n' >> %s", tempAgentLog.Name(), tempAgentLog.Name()) + // Ensure that workspace agents can reinitialize against claimed prebuilds in non-default organizations: + for _, useDefaultOrg := range []bool{true, false} { + t.Run("", func(t *testing.T) { + t.Parallel() - db, ps := dbtestutil.NewDB(t) - // GIVEN a live enterprise API with the prebuilds feature enabled - client, user := coderdenttest.New(t, &coderdenttest.Options{ - Options: &coderdtest.Options{ - Database: db, - Pubsub: ps, - DeploymentValues: coderdtest.DeploymentValues(t, func(dv *codersdk.DeploymentValues) { - dv.Prebuilds.ReconciliationInterval = serpent.Duration(time.Second) - dv.Experiments.Append(string(codersdk.ExperimentWorkspacePrebuilds)) - }), - IncludeProvisionerDaemon: true, - }, - LicenseOptions: &coderdenttest.LicenseOptions{ - Features: license.Features{ - codersdk.FeatureWorkspacePrebuilds: 1, - }, - }, - }) + tempAgentLog := testutil.CreateTemp(t, "", "testReinitializeAgent") - // GIVEN a template, template version, preset and a prebuilt workspace that uses them all - agentToken := uuid.UUID{3} - version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{ - Parse: echo.ParseComplete, - ProvisionPlan: []*proto.Response{ - { - Type: &proto.Response_Plan{ - Plan: &proto.PlanComplete{ - Presets: []*proto.Preset{ - { - Name: "test-preset", - Prebuild: &proto.Prebuild{ - Instances: 1, + startupScript := fmt.Sprintf("printenv >> %s; echo '---\n' >> %s", tempAgentLog.Name(), tempAgentLog.Name()) + + db, ps := dbtestutil.NewDB(t) + // GIVEN a live enterprise API with the prebuilds feature enabled + client, user := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + Database: db, + Pubsub: ps, + DeploymentValues: coderdtest.DeploymentValues(t, func(dv *codersdk.DeploymentValues) { + dv.Prebuilds.ReconciliationInterval = serpent.Duration(time.Second) + dv.Experiments.Append(string(codersdk.ExperimentWorkspacePrebuilds)) + }), + }, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureWorkspacePrebuilds: 1, + codersdk.FeatureExternalProvisionerDaemons: 1, + }, + }, + }) + + orgID := user.OrganizationID + if !useDefaultOrg { + secondOrg := dbgen.Organization(t, db, database.Organization{}) + orgID = secondOrg.ID + } + provisionerCloser := coderdenttest.NewExternalProvisionerDaemon(t, client, orgID, map[string]string{ + provisionersdk.TagScope: provisionersdk.ScopeOrganization, + }) + defer provisionerCloser.Close() + + // GIVEN a template, template version, preset and a prebuilt workspace that uses them all + agentToken := uuid.UUID{3} + version := coderdtest.CreateTemplateVersion(t, client, orgID, &echo.Responses{ + Parse: echo.ParseComplete, + ProvisionPlan: []*proto.Response{ + { + Type: &proto.Response_Plan{ + Plan: &proto.PlanComplete{ + Presets: []*proto.Preset{ + { + Name: "test-preset", + Prebuild: &proto.Prebuild{ + Instances: 1, + }, + }, }, - }, - }, - Resources: []*proto.Resource{ - { - Agents: []*proto.Agent{ + Resources: []*proto.Resource{ { - Name: "smith", - OperatingSystem: "linux", - Architecture: "i386", + Agents: []*proto.Agent{ + { + Name: "smith", + OperatingSystem: "linux", + Architecture: "i386", + }, + }, }, }, }, }, }, }, - }, - }, - ProvisionApply: []*proto.Response{ - { - Type: &proto.Response_Apply{ - Apply: &proto.ApplyComplete{ - Resources: []*proto.Resource{ - { - Type: "compute", - Name: "main", - Agents: []*proto.Agent{ + ProvisionApply: []*proto.Response{ + { + Type: &proto.Response_Apply{ + Apply: &proto.ApplyComplete{ + Resources: []*proto.Resource{ { - Name: "smith", - OperatingSystem: "linux", - Architecture: "i386", - Scripts: []*proto.Script{ + Type: "compute", + Name: "main", + Agents: []*proto.Agent{ { - RunOnStart: true, - Script: startupScript, + Name: "smith", + OperatingSystem: "linux", + Architecture: "i386", + Scripts: []*proto.Script{ + { + RunOnStart: true, + Script: startupScript, + }, + }, + Auth: &proto.Agent_Token{ + Token: agentToken.String(), + }, }, }, - Auth: &proto.Agent_Token{ - Token: agentToken.String(), - }, }, }, }, }, }, }, - }, - }, - }) - coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) + }) + coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) - coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID) + coderdtest.CreateTemplate(t, client, orgID, version.ID) - // Wait for prebuilds to create a prebuilt workspace - ctx := context.Background() - // ctx := testutil.Context(t, testutil.WaitLong) - var ( - prebuildID uuid.UUID - ) - require.Eventually(t, func() bool { - agentAndBuild, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, agentToken) - if err != nil { - return false - } - prebuildID = agentAndBuild.WorkspaceBuild.ID - return true - }, testutil.WaitLong, testutil.IntervalFast) + // Wait for prebuilds to create a prebuilt workspace + ctx := testutil.Context(t, testutil.WaitLong) + var prebuildID uuid.UUID + require.Eventually(t, func() bool { + agentAndBuild, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, agentToken) + if err != nil { + return false + } + prebuildID = agentAndBuild.WorkspaceBuild.ID + return true + }, testutil.WaitLong, testutil.IntervalFast) - prebuild := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, prebuildID) + prebuild := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, prebuildID) - preset, err := db.GetPresetByWorkspaceBuildID(ctx, prebuildID) - require.NoError(t, err) + preset, err := db.GetPresetByWorkspaceBuildID(ctx, prebuildID) + require.NoError(t, err) - // GIVEN a running agent - logDir := t.TempDir() - inv, _ := clitest.New(t, - "agent", - "--auth", "token", - "--agent-token", agentToken.String(), - "--agent-url", client.URL.String(), - "--log-dir", logDir, - ) - clitest.Start(t, inv) + // GIVEN a running agent + logDir := t.TempDir() + inv, _ := clitest.New(t, + "agent", + "--auth", "token", + "--agent-token", agentToken.String(), + "--agent-url", client.URL.String(), + "--log-dir", logDir, + ) + clitest.Start(t, inv) - // GIVEN the agent is in a happy steady state - waiter := coderdtest.NewWorkspaceAgentWaiter(t, client, prebuild.WorkspaceID) - waiter.WaitFor(coderdtest.AgentsReady) + // GIVEN the agent is in a happy steady state + waiter := coderdtest.NewWorkspaceAgentWaiter(t, client, prebuild.WorkspaceID) + waiter.WaitFor(coderdtest.AgentsReady) - // WHEN a workspace is created that can benefit from prebuilds - anotherClient, anotherUser := coderdtest.CreateAnotherUser(t, client, user.OrganizationID) - workspace, err := anotherClient.CreateUserWorkspace(ctx, anotherUser.ID.String(), codersdk.CreateWorkspaceRequest{ - TemplateVersionID: version.ID, - TemplateVersionPresetID: preset.ID, - Name: "claimed-workspace", - }) - require.NoError(t, err) + // WHEN a workspace is created that can benefit from prebuilds + anotherClient, anotherUser := coderdtest.CreateAnotherUser(t, client, orgID) + workspace, err := anotherClient.CreateUserWorkspace(ctx, anotherUser.ID.String(), codersdk.CreateWorkspaceRequest{ + TemplateVersionID: version.ID, + TemplateVersionPresetID: preset.ID, + Name: "claimed-workspace", + }) + require.NoError(t, err) - coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID) + coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID) - // THEN reinitialization completes - waiter.WaitFor(coderdtest.AgentsReady) + // THEN reinitialization completes + waiter.WaitFor(coderdtest.AgentsReady) - var matches [][]byte - require.Eventually(t, func() bool { - // THEN the agent script ran again and reused the same agent token - contents, err := os.ReadFile(tempAgentLog.Name()) - if err != nil { - return false - } - // UUID regex pattern (matches UUID v4-like strings) - uuidRegex := regexp.MustCompile(`\bCODER_AGENT_TOKEN=(.+)\b`) + var matches [][]byte + require.Eventually(t, func() bool { + // THEN the agent script ran again and reused the same agent token + contents, err := os.ReadFile(tempAgentLog.Name()) + if err != nil { + return false + } + // UUID regex pattern (matches UUID v4-like strings) + uuidRegex := regexp.MustCompile(`\bCODER_AGENT_TOKEN=(.+)\b`) - matches = uuidRegex.FindAll(contents, -1) - // When an agent reinitializes, we expect it to run startup scripts again. - // As such, we expect to have written the agent environment to the temp file twice. - // Once on initial startup and then once on reinitialization. - return len(matches) == 2 - }, testutil.WaitLong, testutil.IntervalMedium) - require.Equal(t, matches[0], matches[1]) + matches = uuidRegex.FindAll(contents, -1) + // When an agent reinitializes, we expect it to run startup scripts again. + // As such, we expect to have written the agent environment to the temp file twice. + // Once on initial startup and then once on reinitialization. + return len(matches) == 2 + }, testutil.WaitLong, testutil.IntervalMedium) + require.Equal(t, matches[0], matches[1]) + }) + } } type setupResp struct {