Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit9d47f64

Browse files
committed
perf: optimize prebuilds membership reconciliation to check orgs not presets
1 parentcf93c34 commit9d47f64

File tree

10 files changed

+363
-210
lines changed

10 files changed

+363
-210
lines changed

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2648,6 +2648,13 @@ func (q *querier) GetOrganizationsByUserID(ctx context.Context, userID database.
26482648
returnfetchWithPostFilter(q.auth,policy.ActionRead,q.db.GetOrganizationsByUserID)(ctx,userID)
26492649
}
26502650

2651+
func (q*querier)GetOrganizationsWithPrebuildStatus(ctx context.Context,arg database.GetOrganizationsWithPrebuildStatusParams) ([]database.GetOrganizationsWithPrebuildStatusRow,error) {
2652+
iferr:=q.authorizeContext(ctx,policy.ActionRead,rbac.ResourceOrganization.All());err!=nil {
2653+
returnnil,err
2654+
}
2655+
returnq.db.GetOrganizationsWithPrebuildStatus(ctx,arg)
2656+
}
2657+
26512658
func (q*querier)GetParameterSchemasByJobID(ctx context.Context,jobID uuid.UUID) ([]database.ParameterSchema,error) {
26522659
version,err:=q.db.GetTemplateVersionByJobID(ctx,jobID)
26532660
iferr!=nil {

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3756,6 +3756,14 @@ func (s *MethodTestSuite) TestPrebuilds() {
37563756
dbm.EXPECT().GetPrebuildMetrics(gomock.Any()).Return([]database.GetPrebuildMetricsRow{},nil).AnyTimes()
37573757
check.Args().Asserts(rbac.ResourceWorkspace.All(),policy.ActionRead)
37583758
}))
3759+
s.Run("GetOrganizationsWithPrebuildStatus",s.Mocked(func(dbm*dbmock.MockStore,faker*gofakeit.Faker,check*expects) {
3760+
arg:= database.GetOrganizationsWithPrebuildStatusParams{
3761+
UserID:uuid.New(),
3762+
GroupName:"test",
3763+
}
3764+
dbm.EXPECT().GetOrganizationsWithPrebuildStatus(gomock.Any(),arg).Return([]database.GetOrganizationsWithPrebuildStatusRow{},nil).AnyTimes()
3765+
check.Args(arg).Asserts(rbac.ResourceOrganization.All(),policy.ActionRead)
3766+
}))
37593767
s.Run("GetPrebuildsSettings",s.Mocked(func(dbm*dbmock.MockStore,_*gofakeit.Faker,check*expects) {
37603768
dbm.EXPECT().GetPrebuildsSettings(gomock.Any()).Return("{}",nil).AnyTimes()
37613769
check.Args().Asserts()

‎coderd/database/dbmetrics/querymetrics.go‎

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go‎

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/database/querier.go‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go‎

Lines changed: 87 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/database/queries/prebuilds.sql‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -325,3 +325,47 @@ WHERE id IN (
325325
ANDpj.completed_at ISNULL
326326
)
327327
RETURNING id;
328+
329+
-- name: GetOrganizationsWithPrebuildStatus :many
330+
-- GetOrganizationsWithPrebuildStatus returns organizations with prebuilds configured and their
331+
-- membership status for the prebuilds system user (org membership, group existence, group membership).
332+
WITH orgs_with_prebuildsAS (
333+
-- Get unique organizations that have presets with prebuilds configured
334+
SELECT DISTINCTo.id,o.name
335+
FROM organizations o
336+
INNER JOIN templates tONt.organization_id=o.id
337+
INNER JOIN template_versions tvONtv.template_id=t.id
338+
INNER JOIN template_version_presets tvpONtvp.template_version_id=tv.id
339+
WHEREtvp.desired_instancesIS NOT NULL
340+
),
341+
prebuild_user_membershipAS (
342+
-- Check if the user is a member of the organizations
343+
SELECTom.organization_id
344+
FROM organization_members om
345+
INNER JOIN orgs_with_prebuilds owpONowp.id=om.organization_id
346+
WHEREom.user_id= @user_id::uuid
347+
),
348+
prebuild_groupsAS (
349+
-- Check if the organizations have the prebuilds group
350+
SELECTg.organization_id,g.idas group_id
351+
FROM groups g
352+
INNER JOIN orgs_with_prebuilds owpONowp.id=g.organization_id
353+
WHEREg.name= @group_name::text
354+
),
355+
prebuild_group_membershipAS (
356+
-- Check if the user is in the prebuilds group
357+
SELECTpg.organization_id
358+
FROM prebuild_groups pg
359+
INNER JOIN group_members gmONgm.group_id=pg.group_id
360+
WHEREgm.user_id= @user_id::uuid
361+
)
362+
SELECT
363+
owp.idAS organization_id,
364+
owp.nameAS organization_name,
365+
(pum.organization_idIS NOT NULL)::booleanAS has_prebuild_user,
366+
pg.group_idAS prebuilds_group_id,
367+
(pgm.organization_idIS NOT NULL)::booleanAS has_prebuild_user_in_group
368+
FROM orgs_with_prebuilds owp
369+
LEFT JOIN prebuild_groups pgONpg.organization_id=owp.id
370+
LEFT JOIN prebuild_user_membership pumONpum.organization_id=owp.id
371+
LEFT JOIN prebuild_group_membership pgmONpgm.organization_id=owp.id;

‎enterprise/coderd/prebuilds/membership.go‎

Lines changed: 43 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package prebuilds
22

33
import (
44
"context"
5-
"database/sql"
65
"errors"
76

87
"github.com/google/uuid"
@@ -32,103 +31,78 @@ func NewStoreMembershipReconciler(store database.Store, clock quartz.Clock) Stor
3231
}
3332
}
3433

35-
// ReconcileAll compares the current organization and group memberships of a user to the memberships required
36-
// in order to create prebuilt workspaces. If the user in question is not yet a member of an organization that
37-
// needs prebuilt workspaces, ReconcileAll will create the membership required.
34+
// ReconcileAll ensures the prebuilds system user has the necessary memberships to create prebuilt workspaces.
35+
// For each organization with prebuilds configured, it ensures:
36+
// * The user is a member of the organization
37+
// * A group exists with quota 0
38+
// * The user is a member of that group
3839
//
39-
// To facilitate quota management, ReconcileAll will ensure:
40-
// * the existence of a group (defined by PrebuiltWorkspacesGroupName) in each organization that needs prebuilt workspaces
41-
// * that the prebuilds system user belongs to the group in each organization that needs prebuilt workspaces
42-
// * that the group has a quota of 0 by default, which users can adjust based on their needs.
40+
// Unique constraint violations are safely ignored (concurrent creation).
4341
//
4442
// ReconcileAll does not have an opinion on transaction or lock management. These responsibilities are left to the caller.
45-
func (sStoreMembershipReconciler)ReconcileAll(ctx context.Context,userID uuid.UUID,presets []database.GetTemplatePresetsWithPrebuildsRow)error {
46-
organizationMemberships,err:=s.store.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{
47-
UserID:userID,
48-
Deleted: sql.NullBool{
49-
Bool:false,
50-
Valid:true,
51-
},
43+
func (sStoreMembershipReconciler)ReconcileAll(ctx context.Context,userID uuid.UUID,groupNamestring)error {
44+
orgStatuses,err:=s.store.GetOrganizationsWithPrebuildStatus(ctx, database.GetOrganizationsWithPrebuildStatusParams{
45+
UserID:userID,
46+
GroupName:groupName,
5247
})
5348
iferr!=nil {
54-
returnxerrors.Errorf("determine prebuild organization membership: %w",err)
55-
}
56-
57-
orgMemberships:=make(map[uuid.UUID]struct{},0)
58-
defaultOrg,err:=s.store.GetDefaultOrganization(ctx)
59-
iferr!=nil {
60-
returnxerrors.Errorf("get default organization: %w",err)
61-
}
62-
orgMemberships[defaultOrg.ID]=struct{}{}
63-
for_,o:=rangeorganizationMemberships {
64-
orgMemberships[o.ID]=struct{}{}
49+
returnxerrors.Errorf("get organizations with prebuild status: %w",err)
6550
}
6651

6752
varmembershipInsertionErrorserror
68-
for_,preset:=rangepresets {
69-
_,alreadyOrgMember:=orgMemberships[preset.OrganizationID]
70-
if!alreadyOrgMember {
71-
// Add the organization to our list of memberships regardless of potential failure below
72-
// to avoid a retry that will probably be doomed anyway.
73-
orgMemberships[preset.OrganizationID]=struct{}{}
74-
75-
// Insert the missing membership
53+
for_,orgStatus:=rangeorgStatuses {
54+
// Add user to org if needed
55+
if!orgStatus.HasPrebuildUser {
7656
_,err=s.store.InsertOrganizationMember(ctx, database.InsertOrganizationMemberParams{
77-
OrganizationID:preset.OrganizationID,
57+
OrganizationID:orgStatus.OrganizationID,
7858
UserID:userID,
7959
CreatedAt:s.clock.Now(),
8060
UpdatedAt:s.clock.Now(),
8161
Roles: []string{},
8262
})
83-
iferr!=nil {
84-
membershipInsertionErrors=errors.Join(membershipInsertionErrors,xerrors.Errorf("insert membership for prebuilt workspaces: %w",err))
85-
continue
86-
}
87-
}
88-
89-
// determine whether the org already has a prebuilds group
90-
prebuildsGroupExists:=true
91-
prebuildsGroup,err:=s.store.GetGroupByOrgAndName(ctx, database.GetGroupByOrgAndNameParams{
92-
OrganizationID:preset.OrganizationID,
93-
Name:PrebuiltWorkspacesGroupName,
94-
})
95-
iferr!=nil {
96-
if!xerrors.Is(err,sql.ErrNoRows) {
97-
membershipInsertionErrors=errors.Join(membershipInsertionErrors,xerrors.Errorf("get prebuilds group: %w",err))
63+
// Unique violation means membership was created after status check, safe to ignore.
64+
iferr!=nil&&!database.IsUniqueViolation(err) {
65+
membershipInsertionErrors=errors.Join(membershipInsertionErrors,err)
9866
continue
9967
}
100-
prebuildsGroupExists=false
10168
}
10269

103-
//if the prebuildsgroupdoes not exist, create it
104-
if!prebuildsGroupExists {
105-
// create a "prebuilds" group in the organization and add the system user to it
106-
//this group will have a quota of 0 by default, which users can adjust based on their needs
107-
prebuildsGroup,err=s.store.InsertGroup(ctx, database.InsertGroupParams{
70+
//Creategroupif it doesn't exist
71+
vargroupID uuid.UUID
72+
if!orgStatus.PrebuildsGroupID.Valid {
73+
//Group doesn't exist, create it
74+
group,err:=s.store.InsertGroup(ctx, database.InsertGroupParams{
10875
ID:uuid.New(),
10976
Name:PrebuiltWorkspacesGroupName,
11077
DisplayName:PrebuiltWorkspacesGroupDisplayName,
111-
OrganizationID:preset.OrganizationID,
78+
OrganizationID:orgStatus.OrganizationID,
11279
AvatarURL:"",
113-
QuotaAllowance:0,// Default quota of 0, users should set this based on their needs
80+
QuotaAllowance:0,
11481
})
115-
iferr!=nil {
116-
membershipInsertionErrors=errors.Join(membershipInsertionErrors,xerrors.Errorf("create prebuilds group: %w",err))
82+
// Unique violation means membership was created after status check, safe to ignore.
83+
iferr!=nil&&!database.IsUniqueViolation(err) {
84+
membershipInsertionErrors=errors.Join(membershipInsertionErrors,err)
11785
continue
11886
}
87+
groupID=group.ID
88+
}else {
89+
// Group exists
90+
groupID=orgStatus.PrebuildsGroupID.UUID
11991
}
12092

121-
// add the system user to the prebuilds group
122-
err=s.store.InsertGroupMember(ctx, database.InsertGroupMemberParams{
123-
GroupID:prebuildsGroup.ID,
124-
UserID:userID,
125-
})
126-
iferr!=nil {
127-
// ignore unique violation errors as the user might already be in the group
128-
if!database.IsUniqueViolation(err) {
129-
membershipInsertionErrors=errors.Join(membershipInsertionErrors,xerrors.Errorf("add system user to prebuilds group: %w",err))
93+
// Add user to group if needed
94+
if!orgStatus.HasPrebuildUserInGroup {
95+
err=s.store.InsertGroupMember(ctx, database.InsertGroupMemberParams{
96+
GroupID:groupID,
97+
UserID:userID,
98+
})
99+
// Unique violation means membership was created after status check, safe to ignore.
100+
iferr!=nil&&!database.IsUniqueViolation(err) {
101+
membershipInsertionErrors=errors.Join(membershipInsertionErrors,err)
102+
continue
130103
}
131104
}
132105
}
106+
133107
returnmembershipInsertionErrors
134108
}

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp