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 { pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy