opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

perf: optimize prebuilds membership reconciliation to check orgs not presets (#20493)

Browse Source 
## Description

The membership reconciliation ensures the prebuilds system user is a
member of all organizations with prebuilds configured. To support
prebuilds quota management, each organization must have a prebuilds
group that the system user belongs to.

## Problem

Previously, membership reconciliation iterated over all presets to check
and update membership status. This meant database queries
`GetGroupByOrgAndName` and `InsertGroupMember` were executed for each
preset. Since presets are unique combinations of `(organization,
template, template version, preset)`, this resulted in several redundant
checks for the same organization.

In dogfood, `InsertGroupMember` was called thousands of times per day,
even though memberships were already configured ([internal Grafana
dashboard link](https://grafana.dev.coder.com/goto/46MZ1UgDg?orgId=1))

<img width="5382" height="1788" alt="Screenshot 2025-10-28 at 16 01 36"
src="https://github.com/user-attachments/assets/757b7253-106f-4f72-8586-8e2ede9f18db"
/>

## Solution

This PR introduces `GetOrganizationsWithPrebuildStatus`, a single query
that returns:
* All unique organizations with prebuilds configured
* Whether the prebuilds user is a member of each organization
* Whether the prebuilds group exists in each organization
* Whether the prebuilds user is in the prebuilds group

The membership reconciliation logic now:
* Fetches status for all organizations in one query
* Only performs inserts for organizations missing required memberships
or groups
* Safely handles concurrent operations via unique constraint violations
* This reduces database load from `O(presets)` to `O(organizations)` per
reconciliation loop, with a single read query when everything is
configured.

## Changes

* Add `GetOrganizationsWithPrebuildStatus` SQL query
* Update `membership.ReconcileAll` to use organization-based
reconciliation instead of preset-based
* Update tests to reflect new behavior

Related to internal thread:
https://codercom.slack.com/archives/C07GRNNRW03/p1760535570381369
This commit is contained in:
Susana Ferreira
2025-10-29 14:24:29 +00:00
committed by GitHub
parent dd28eef5b4
commit 7e8fcb4b0f
10 changed files with 394 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/dbauthz/dbauthz.go
+7
View File
@@ -2648,6 +2648,13 @@ func (q *querier) GetOrganizationsByUserID(ctx context.Context, userID database.
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetOrganizationsByUserID)(ctx, userID)
}
func (q *querier) GetOrganizationsWithPrebuildStatus(ctx context.Context, arg database.GetOrganizationsWithPrebuildStatusParams) ([]database.GetOrganizationsWithPrebuildStatusRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceOrganization.All()); err != nil {
return nil, err
}
return q.db.GetOrganizationsWithPrebuildStatus(ctx, arg)
}
func (q *querier) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]database.ParameterSchema, error) {
version, err := q.db.GetTemplateVersionByJobID(ctx, jobID)
if err != nil {
coderd/database/dbauthz/dbauthz_test.go
+8
View File
@@ -3759,6 +3759,14 @@ func (s *MethodTestSuite) TestPrebuilds() {
dbm.EXPECT().GetPrebuildMetrics(gomock.Any()).Return([]database.GetPrebuildMetricsRow{}, nil).AnyTimes()
check.Args().Asserts(rbac.ResourceWorkspace.All(), policy.ActionRead)
}))
s.Run("GetOrganizationsWithPrebuildStatus", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
arg := database.GetOrganizationsWithPrebuildStatusParams{
UserID: uuid.New(),
GroupName: "test",
}
dbm.EXPECT().GetOrganizationsWithPrebuildStatus(gomock.Any(), arg).Return([]database.GetOrganizationsWithPrebuildStatusRow{}, nil).AnyTimes()
check.Args(arg).Asserts(rbac.ResourceOrganization.All(), policy.ActionRead)
}))
s.Run("GetPrebuildsSettings", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
dbm.EXPECT().GetPrebuildsSettings(gomock.Any()).Return("{}", nil).AnyTimes()
check.Args().Asserts()
coderd/database/dbmetrics/querymetrics.go
+7
View File
@@ -1243,6 +1243,13 @@ func (m queryMetricsStore) GetOrganizationsByUserID(ctx context.Context, userID
return organizations, err
}
func (m queryMetricsStore) GetOrganizationsWithPrebuildStatus(ctx context.Context, arg database.GetOrganizationsWithPrebuildStatusParams) ([]database.GetOrganizationsWithPrebuildStatusRow, error) {
start := time.Now()
r0, r1 := m.s.GetOrganizationsWithPrebuildStatus(ctx, arg)
m.queryLatencies.WithLabelValues("GetOrganizationsWithPrebuildStatus").Observe(time.Since(start).Seconds())
return r0, r1
}
func (m queryMetricsStore) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]database.ParameterSchema, error) {
start := time.Now()
schemas, err := m.s.GetParameterSchemasByJobID(ctx, jobID)
coderd/database/dbmock/dbmock.go
+15
View File
@@ -2622,6 +2622,21 @@ func (mr *MockStoreMockRecorder) GetOrganizationsByUserID(ctx, arg any) *gomock.
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetOrganizationsByUserID", reflect.TypeOf((*MockStore)(nil).GetOrganizationsByUserID), ctx, arg)
}
// GetOrganizationsWithPrebuildStatus mocks base method.
func (m *MockStore) GetOrganizationsWithPrebuildStatus(ctx context.Context, arg database.GetOrganizationsWithPrebuildStatusParams) ([]database.GetOrganizationsWithPrebuildStatusRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetOrganizationsWithPrebuildStatus", ctx, arg)
ret0, _ := ret[0].([]database.GetOrganizationsWithPrebuildStatusRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetOrganizationsWithPrebuildStatus indicates an expected call of GetOrganizationsWithPrebuildStatus.
func (mr *MockStoreMockRecorder) GetOrganizationsWithPrebuildStatus(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetOrganizationsWithPrebuildStatus", reflect.TypeOf((*MockStore)(nil).GetOrganizationsWithPrebuildStatus), ctx, arg)
}
// GetParameterSchemasByJobID mocks base method.
func (m *MockStore) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]database.ParameterSchema, error) {
m.ctrl.T.Helper()
coderd/database/querier.go
+3
View File
@@ -269,6 +269,9 @@ type sqlcQuerier interface {
GetOrganizationResourceCountByID(ctx context.Context, organizationID uuid.UUID) (GetOrganizationResourceCountByIDRow, error)
GetOrganizations(ctx context.Context, arg GetOrganizationsParams) ([]Organization, error)
GetOrganizationsByUserID(ctx context.Context, arg GetOrganizationsByUserIDParams) ([]Organization, error)
// GetOrganizationsWithPrebuildStatus returns organizations with prebuilds configured and their
// membership status for the prebuilds system user (org membership, group existence, group membership).
GetOrganizationsWithPrebuildStatus(ctx context.Context, arg GetOrganizationsWithPrebuildStatusParams) ([]GetOrganizationsWithPrebuildStatusRow, error)
GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]ParameterSchema, error)
GetPrebuildMetrics(ctx context.Context) ([]GetPrebuildMetricsRow, error)
GetPrebuildsSettings(ctx context.Context) (string, error)
coderd/database/queries.sql.go
+87
View File
@@ -8285,6 +8285,93 @@ func (q *sqlQuerier) FindMatchingPresetID(ctx context.Context, arg FindMatchingP
return template_version_preset_id, err
}
const getOrganizationsWithPrebuildStatus = `-- name: GetOrganizationsWithPrebuildStatus :many
WITH orgs_with_prebuilds AS (
-- Get unique organizations that have presets with prebuilds configured
SELECT DISTINCT o.id, o.name
FROM organizations o
INNER JOIN templates t ON t.organization_id = o.id
INNER JOIN template_versions tv ON tv.template_id = t.id
INNER JOIN template_version_presets tvp ON tvp.template_version_id = tv.id
WHERE tvp.desired_instances IS NOT NULL
),
prebuild_user_membership AS (
-- Check if the user is a member of the organizations
SELECT om.organization_id
FROM organization_members om
INNER JOIN orgs_with_prebuilds owp ON owp.id = om.organization_id
WHERE om.user_id = $1::uuid
),
prebuild_groups AS (
-- Check if the organizations have the prebuilds group
SELECT g.organization_id, g.id as group_id
FROM groups g
INNER JOIN orgs_with_prebuilds owp ON owp.id = g.organization_id
WHERE g.name = $2::text
),
prebuild_group_membership AS (
-- Check if the user is in the prebuilds group
SELECT pg.organization_id
FROM prebuild_groups pg
INNER JOIN group_members gm ON gm.group_id = pg.group_id
WHERE gm.user_id = $1::uuid
)
SELECT
owp.id AS organization_id,
owp.name AS organization_name,
(pum.organization_id IS NOT NULL)::boolean AS has_prebuild_user,
pg.group_id AS prebuilds_group_id,
(pgm.organization_id IS NOT NULL)::boolean AS has_prebuild_user_in_group
FROM orgs_with_prebuilds owp
LEFT JOIN prebuild_groups pg ON pg.organization_id = owp.id
LEFT JOIN prebuild_user_membership pum ON pum.organization_id = owp.id
LEFT JOIN prebuild_group_membership pgm ON pgm.organization_id = owp.id
`
type GetOrganizationsWithPrebuildStatusParams struct {
UserID uuid.UUID `db:"user_id" json:"user_id"`
GroupName string `db:"group_name" json:"group_name"`
}
type GetOrganizationsWithPrebuildStatusRow struct {
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
OrganizationName string `db:"organization_name" json:"organization_name"`
HasPrebuildUser bool `db:"has_prebuild_user" json:"has_prebuild_user"`
PrebuildsGroupID uuid.NullUUID `db:"prebuilds_group_id" json:"prebuilds_group_id"`
HasPrebuildUserInGroup bool `db:"has_prebuild_user_in_group" json:"has_prebuild_user_in_group"`
}
// GetOrganizationsWithPrebuildStatus returns organizations with prebuilds configured and their
// membership status for the prebuilds system user (org membership, group existence, group membership).
func (q *sqlQuerier) GetOrganizationsWithPrebuildStatus(ctx context.Context, arg GetOrganizationsWithPrebuildStatusParams) ([]GetOrganizationsWithPrebuildStatusRow, error) {
rows, err := q.db.QueryContext(ctx, getOrganizationsWithPrebuildStatus, arg.UserID, arg.GroupName)
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetOrganizationsWithPrebuildStatusRow
for rows.Next() {
var i GetOrganizationsWithPrebuildStatusRow
if err := rows.Scan(
&i.OrganizationID,
&i.OrganizationName,
&i.HasPrebuildUser,
&i.PrebuildsGroupID,
&i.HasPrebuildUserInGroup,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const getPrebuildMetrics = `-- name: GetPrebuildMetrics :many
SELECT
t.name as template_name,
coderd/database/queries/prebuilds.sql
+44
View File
@@ -327,3 +327,47 @@ SET
FROM jobs_to_cancel
WHERE provisioner_jobs.id = jobs_to_cancel.id
RETURNING jobs_to_cancel.id, jobs_to_cancel.workspace_id, jobs_to_cancel.template_id, jobs_to_cancel.template_version_preset_id;
-- name: GetOrganizationsWithPrebuildStatus :many
-- GetOrganizationsWithPrebuildStatus returns organizations with prebuilds configured and their
-- membership status for the prebuilds system user (org membership, group existence, group membership).
WITH orgs_with_prebuilds AS (
-- Get unique organizations that have presets with prebuilds configured
SELECT DISTINCT o.id, o.name
FROM organizations o
INNER JOIN templates t ON t.organization_id = o.id
INNER JOIN template_versions tv ON tv.template_id = t.id
INNER JOIN template_version_presets tvp ON tvp.template_version_id = tv.id
WHERE tvp.desired_instances IS NOT NULL
),
prebuild_user_membership AS (
-- Check if the user is a member of the organizations
SELECT om.organization_id
FROM organization_members om
INNER JOIN orgs_with_prebuilds owp ON owp.id = om.organization_id
WHERE om.user_id = @user_id::uuid
),
prebuild_groups AS (
-- Check if the organizations have the prebuilds group
SELECT g.organization_id, g.id as group_id
FROM groups g
INNER JOIN orgs_with_prebuilds owp ON owp.id = g.organization_id
WHERE g.name = @group_name::text
),
prebuild_group_membership AS (
-- Check if the user is in the prebuilds group
SELECT pg.organization_id
FROM prebuild_groups pg
INNER JOIN group_members gm ON gm.group_id = pg.group_id
WHERE gm.user_id = @user_id::uuid
)
SELECT
owp.id AS organization_id,
owp.name AS organization_name,
(pum.organization_id IS NOT NULL)::boolean AS has_prebuild_user,
pg.group_id AS prebuilds_group_id,
(pgm.organization_id IS NOT NULL)::boolean AS has_prebuild_user_in_group
FROM orgs_with_prebuilds owp
LEFT JOIN prebuild_groups pg ON pg.organization_id = owp.id
LEFT JOIN prebuild_user_membership pum ON pum.organization_id = owp.id
LEFT JOIN prebuild_group_membership pgm ON pgm.organization_id = owp.id;