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

## 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

coderd/database/dbauthz/dbauthz.go

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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;

enterprise/coderd/prebuilds/membership.go

@@ -2,12 +2,13 @@ package prebuilds
 
 import (
 	"context"
-	"database/sql"
 	"errors"
 
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
 
+	"cdr.dev/slog"
+
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/quartz"
 )
@@ -21,114 +22,117 @@ const (
 // 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
+	store  database.Store
+	clock  quartz.Clock
+	logger slog.Logger
 }
 
-func NewStoreMembershipReconciler(store database.Store, clock quartz.Clock) StoreMembershipReconciler {
+func NewStoreMembershipReconciler(store database.Store, clock quartz.Clock, logger slog.Logger) StoreMembershipReconciler {
 	return StoreMembershipReconciler{
-		store: store,
-		clock: clock,
+		store:  store,
+		clock:  clock,
+		logger: logger,
 	}
 }
 
-// ReconcileAll compares the current organization and group memberships of a user to the memberships 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.
+// ReconcileAll ensures the prebuilds system user has the necessary memberships to create prebuilt workspaces.
+// For each organization with prebuilds configured, it ensures:
+// * The user is a member of the organization
+// * A group exists with quota 0
+// * The user is a member of that group
 //
-// To facilitate quota management, ReconcileAll will ensure:
-// * the existence of a group (defined by PrebuiltWorkspacesGroupName) in each organization that needs prebuilt workspaces
-// * that the prebuilds system user belongs to the group in each organization that needs prebuilt workspaces
-// * that the group has a quota of 0 by default, which users can adjust based on their needs.
+// Unique constraint violations are safely ignored (concurrent creation).
 //
 // ReconcileAll 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,
-		},
+func (s StoreMembershipReconciler) ReconcileAll(ctx context.Context, userID uuid.UUID, groupName string) error {
+	orgStatuses, err := s.store.GetOrganizationsWithPrebuildStatus(ctx, database.GetOrganizationsWithPrebuildStatusParams{
+		UserID:    userID,
+		GroupName: groupName,
 	})
 	if err != nil {
-		return xerrors.Errorf("determine prebuild organization membership: %w", err)
-	}
-
-	orgMemberships := make(map[uuid.UUID]struct{}, 0)
-	defaultOrg, err := s.store.GetDefaultOrganization(ctx)
-	if err != nil {
-		return xerrors.Errorf("get default organization: %w", err)
-	}
-	orgMemberships[defaultOrg.ID] = struct{}{}
-	for _, o := range organizationMemberships {
-		orgMemberships[o.ID] = struct{}{}
+		return xerrors.Errorf("get organizations with prebuild status: %w", err)
 	}
 
 	var membershipInsertionErrors error
-	for _, preset := range presets {
-		_, alreadyOrgMember := orgMemberships[preset.OrganizationID]
-		if !alreadyOrgMember {
-			// Add the organization to our list of memberships regardless of potential failure below
-			// to avoid a retry that will probably be doomed anyway.
-			orgMemberships[preset.OrganizationID] = struct{}{}
+	for _, orgStatus := range orgStatuses {
+		s.logger.Debug(ctx, "organization prebuild status",
+			slog.F("organization_id", orgStatus.OrganizationID),
+			slog.F("organization_name", orgStatus.OrganizationName),
+			slog.F("has_prebuild_user", orgStatus.HasPrebuildUser),
+			slog.F("has_prebuild_group", orgStatus.PrebuildsGroupID.Valid),
+			slog.F("has_prebuild_user_in_group", orgStatus.HasPrebuildUserInGroup))
 
-			// Insert the missing membership
+		// Add user to org if needed
+		if !orgStatus.HasPrebuildUser {
 			_, err = s.store.InsertOrganizationMember(ctx, database.InsertOrganizationMemberParams{
-				OrganizationID: preset.OrganizationID,
+				OrganizationID: orgStatus.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))
+			// Unique violation means organization membership was created after status check, safe to ignore.
+			if err != nil && !database.IsUniqueViolation(err) {
+				membershipInsertionErrors = errors.Join(membershipInsertionErrors, err)
 				continue
 			}
+			if err == nil {
+				s.logger.Info(ctx, "added prebuilds user to organization",
+					slog.F("organization_id", orgStatus.OrganizationID),
+					slog.F("organization_name", orgStatus.OrganizationName),
+					slog.F("prebuilds_user", userID.String()))
+			}
 		}
 
-		// determine whether the org already has a prebuilds group
-		prebuildsGroupExists := true
-		prebuildsGroup, err := s.store.GetGroupByOrgAndName(ctx, database.GetGroupByOrgAndNameParams{
-			OrganizationID: preset.OrganizationID,
-			Name:           PrebuiltWorkspacesGroupName,
-		})
-		if err != nil {
-			if !xerrors.Is(err, sql.ErrNoRows) {
-				membershipInsertionErrors = errors.Join(membershipInsertionErrors, xerrors.Errorf("get prebuilds group: %w", err))
-				continue
-			}
-			prebuildsGroupExists = false
-		}
-
-		// if the prebuilds group does not exist, create it
-		if !prebuildsGroupExists {
-			// create a "prebuilds" group in the organization and add the system user to it
-			// this group will have a quota of 0 by default, which users can adjust based on their needs
-			prebuildsGroup, err = s.store.InsertGroup(ctx, database.InsertGroupParams{
+		// Create group if it doesn't exist
+		var groupID uuid.UUID
+		if !orgStatus.PrebuildsGroupID.Valid {
+			// Group doesn't exist, create it
+			group, err := s.store.InsertGroup(ctx, database.InsertGroupParams{
 				ID:             uuid.New(),
 				Name:           PrebuiltWorkspacesGroupName,
 				DisplayName:    PrebuiltWorkspacesGroupDisplayName,
-				OrganizationID: preset.OrganizationID,
+				OrganizationID: orgStatus.OrganizationID,
 				AvatarURL:      "",
-				QuotaAllowance: 0, // Default quota of 0, users should set this based on their needs
+				QuotaAllowance: 0,
 			})
-			if err != nil {
-				membershipInsertionErrors = errors.Join(membershipInsertionErrors, xerrors.Errorf("create prebuilds group: %w", err))
+			// Unique violation means group was created after status check, safe to ignore.
+			if err != nil && !database.IsUniqueViolation(err) {
+				membershipInsertionErrors = errors.Join(membershipInsertionErrors, err)
 				continue
 			}
+			if err == nil {
+				s.logger.Info(ctx, "created prebuilds group in organization",
+					slog.F("organization_id", orgStatus.OrganizationID),
+					slog.F("organization_name", orgStatus.OrganizationName),
+					slog.F("prebuilds_group", group.ID.String()))
+			}
+			groupID = group.ID
+		} else {
+			// Group exists
+			groupID = orgStatus.PrebuildsGroupID.UUID
 		}
 
-		// add the system user to the prebuilds group
-		err = s.store.InsertGroupMember(ctx, database.InsertGroupMemberParams{
-			GroupID: prebuildsGroup.ID,
-			UserID:  userID,
-		})
-		if err != nil {
-			// ignore unique violation errors as the user might already be in the group
-			if !database.IsUniqueViolation(err) {
-				membershipInsertionErrors = errors.Join(membershipInsertionErrors, xerrors.Errorf("add system user to prebuilds group: %w", err))
+		// Add user to group if needed
+		if !orgStatus.HasPrebuildUserInGroup {
+			err = s.store.InsertGroupMember(ctx, database.InsertGroupMemberParams{
+				GroupID: groupID,
+				UserID:  userID,
+			})
+			// Unique violation means group membership was created after status check, safe to ignore.
+			if err != nil && !database.IsUniqueViolation(err) {
+				membershipInsertionErrors = errors.Join(membershipInsertionErrors, err)
+				continue
+			}
+			if err == nil {
+				s.logger.Info(ctx, "added prebuilds user to prebuilds group",
+					slog.F("organization_id", orgStatus.OrganizationID),
+					slog.F("organization_name", orgStatus.OrganizationName),
+					slog.F("prebuilds_user", userID.String()),
+					slog.F("prebuilds_group", groupID.String()))
 			}
 		}
 	}
+
 	return membershipInsertionErrors
 }

enterprise/coderd/prebuilds/membership_test.go

@@ -7,16 +7,17 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
-	"tailscale.com/types/ptr"
 
-	"github.com/coder/quartz"
+	"cdr.dev/slog/sloggers/slogtest"
 
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/coder/v2/coderd/database/dbfake"
 	"github.com/coder/coder/v2/coderd/database/dbgen"
 	"github.com/coder/coder/v2/enterprise/coderd/prebuilds"
 	"github.com/coder/coder/v2/testutil"
+	"github.com/coder/quartz"
 )
 
 // TestReconcileAll verifies that StoreMembershipReconciler correctly updates membership
@@ -26,169 +27,178 @@ func TestReconcileAll(t *testing.T) {
 
 	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
+		includePreset              bool
 		preExistingOrgMembership   []bool
 		preExistingGroup           []bool
 		preExistingGroupMembership []bool
 		// Expected outcomes
-		expectOrgMembershipExists *bool
-		expectGroupExists         *bool
-		expectUserInGroup         *bool
+		expectOrgMembershipExists bool
+		expectGroupExists         bool
+		expectUserInGroup         bool
 	}{
 		{
 			name:                       "if there are no presets, membership reconciliation is a no-op",
-			includePreset:              []bool{false},
+			includePreset:              false,
 			preExistingOrgMembership:   []bool{true, false},
 			preExistingGroup:           []bool{true, false},
 			preExistingGroupMembership: []bool{true, false},
-			expectOrgMembershipExists:  ptr.To(false),
-			expectGroupExists:          ptr.To(false),
+			expectOrgMembershipExists:  false,
+			expectGroupExists:          false,
+			expectUserInGroup:          false,
 		},
 		{
 			name:                       "if there is a preset, then we should enforce org and group membership in all cases",
-			includePreset:              []bool{true},
+			includePreset:              true,
 			preExistingOrgMembership:   []bool{true, false},
 			preExistingGroup:           []bool{true, false},
 			preExistingGroupMembership: []bool{true, false},
-			expectOrgMembershipExists:  ptr.To(true),
-			expectGroupExists:          ptr.To(true),
-			expectUserInGroup:          ptr.To(true),
+			expectOrgMembershipExists:  true,
+			expectGroupExists:          true,
+			expectUserInGroup:          true,
 		},
 	}
 
 	for _, tc := range tests {
 		tc := tc
-		for _, includePreset := range tc.includePreset {
-			includePreset := includePreset
-			for _, preExistingOrgMembership := range tc.preExistingOrgMembership {
-				preExistingOrgMembership := preExistingOrgMembership
-				for _, preExistingGroup := range tc.preExistingGroup {
-					preExistingGroup := preExistingGroup
-					for _, preExistingGroupMembership := range tc.preExistingGroupMembership {
-						preExistingGroupMembership := preExistingGroupMembership
-						t.Run(tc.name, func(t *testing.T) {
-							t.Parallel()
+		includePreset := tc.includePreset
+		for _, preExistingOrgMembership := range tc.preExistingOrgMembership {
+			preExistingOrgMembership := preExistingOrgMembership
+			for _, preExistingGroup := range tc.preExistingGroup {
+				preExistingGroup := preExistingGroup
+				for _, preExistingGroupMembership := range tc.preExistingGroupMembership {
+					preExistingGroupMembership := preExistingGroupMembership
+					t.Run(tc.name, func(t *testing.T) {
+						t.Parallel()
 
-							// nolint:gocritic // Reconciliation happens as prebuilds system user, not a human user.
-							ctx := dbauthz.AsPrebuildsOrchestrator(testutil.Context(t, testutil.WaitLong))
-							_, db := coderdtest.NewWithDatabase(t, nil)
+						// nolint:gocritic // Reconciliation happens as prebuilds system user, not a human user.
+						ctx := dbauthz.AsPrebuildsOrchestrator(testutil.Context(t, testutil.WaitLong))
+						client, db := coderdtest.NewWithDatabase(t, nil)
+						owner := coderdtest.CreateFirstUser(t, client)
 
-							defaultOrg, err := db.GetDefaultOrganization(ctx)
-							require.NoError(t, err)
+						defaultOrg, err := db.GetDefaultOrganization(ctx)
+						require.NoError(t, err)
 
-							// introduce an unrelated organization to ensure that the membership reconciler doesn't interfere with it.
-							unrelatedOrg := dbgen.Organization(t, db, database.Organization{})
-							targetOrg := dbgen.Organization(t, db, database.Organization{})
+						// Introduce an unrelated organization to ensure that the membership reconciler doesn't interfere with it.
+						unrelatedOrg := dbgen.Organization(t, db, database.Organization{})
+						dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: unrelatedOrg.ID, UserID: database.PrebuildsSystemUserID})
 
-							// Ensure membership to unrelated org.
-							dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: unrelatedOrg.ID, UserID: database.PrebuildsSystemUserID})
+						// Organization to test
+						targetOrg := dbgen.Organization(t, db, database.Organization{})
 
-							if preExistingOrgMembership {
-								// System user already a member of both orgs.
-								dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: targetOrg.ID, UserID: database.PrebuildsSystemUserID})
-							}
+						// Prebuilds system user is a member of the organization
+						if preExistingOrgMembership {
+							dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: targetOrg.ID, UserID: database.PrebuildsSystemUserID})
+						}
 
-							// Create pre-existing prebuilds group if required by test case
-							var prebuildsGroup database.Group
-							if preExistingGroup {
-								prebuildsGroup = dbgen.Group(t, db, database.Group{
-									Name:           prebuilds.PrebuiltWorkspacesGroupName,
-									DisplayName:    prebuilds.PrebuiltWorkspacesGroupDisplayName,
-									OrganizationID: targetOrg.ID,
-									QuotaAllowance: 0,
-								})
-
-								// Add the system user to the group if preExistingGroupMembership is true
-								if preExistingGroupMembership {
-									dbgen.GroupMember(t, db, database.GroupMemberTable{
-										GroupID: prebuildsGroup.ID,
-										UserID:  database.PrebuildsSystemUserID,
-									})
-								}
-							}
-
-							presets := []database.GetTemplatePresetsWithPrebuildsRow{newPresetRow(unrelatedOrg.ID)}
-							if includePreset {
-								presets = append(presets, newPresetRow(targetOrg.ID))
-							}
-
-							// Verify memberships before reconciliation.
-							preReconcileMemberships, err := db.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{
-								UserID: database.PrebuildsSystemUserID,
-							})
-							require.NoError(t, err)
-							expectedMembershipsBefore := []uuid.UUID{defaultOrg.ID, unrelatedOrg.ID}
-							if preExistingOrgMembership {
-								expectedMembershipsBefore = append(expectedMembershipsBefore, targetOrg.ID)
-							}
-							require.ElementsMatch(t, expectedMembershipsBefore, extractOrgIDs(preReconcileMemberships))
-
-							// Reconcile
-							reconciler := prebuilds.NewStoreMembershipReconciler(db, clock)
-							require.NoError(t, reconciler.ReconcileAll(ctx, database.PrebuildsSystemUserID, presets))
-
-							// Verify memberships after reconciliation.
-							postReconcileMemberships, err := db.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{
-								UserID: database.PrebuildsSystemUserID,
-							})
-							require.NoError(t, err)
-							expectedMembershipsAfter := expectedMembershipsBefore
-							if !preExistingOrgMembership && tc.expectOrgMembershipExists != nil && *tc.expectOrgMembershipExists {
-								expectedMembershipsAfter = append(expectedMembershipsAfter, targetOrg.ID)
-							}
-							require.ElementsMatch(t, expectedMembershipsAfter, extractOrgIDs(postReconcileMemberships))
-
-							// Verify prebuilds group behavior based on expected outcomes
-							prebuildsGroup, err = db.GetGroupByOrgAndName(ctx, database.GetGroupByOrgAndNameParams{
-								OrganizationID: targetOrg.ID,
+						// Organization has the prebuilds group
+						var prebuildsGroup database.Group
+						if preExistingGroup {
+							prebuildsGroup = dbgen.Group(t, db, database.Group{
 								Name:           prebuilds.PrebuiltWorkspacesGroupName,
+								DisplayName:    prebuilds.PrebuiltWorkspacesGroupDisplayName,
+								OrganizationID: targetOrg.ID,
+								QuotaAllowance: 0,
 							})
-							if tc.expectGroupExists != nil && *tc.expectGroupExists {
-								require.NoError(t, err)
-								require.Equal(t, prebuilds.PrebuiltWorkspacesGroupName, prebuildsGroup.Name)
-								require.Equal(t, prebuilds.PrebuiltWorkspacesGroupDisplayName, prebuildsGroup.DisplayName)
-								require.Equal(t, int32(0), prebuildsGroup.QuotaAllowance) // Default quota should be 0
 
-								if tc.expectUserInGroup != nil && *tc.expectUserInGroup {
-									// Check that the system user is a member of the prebuilds group
-									groupMembers, err := db.GetGroupMembersByGroupID(ctx, database.GetGroupMembersByGroupIDParams{
-										GroupID:       prebuildsGroup.ID,
-										IncludeSystem: true,
-									})
-									require.NoError(t, err)
-									require.Len(t, groupMembers, 1)
-									require.Equal(t, database.PrebuildsSystemUserID, groupMembers[0].UserID)
-								}
-
-								// If no preset exists, then we do not enforce group membership:
-								if tc.expectUserInGroup != nil && !*tc.expectUserInGroup {
-									// Check that the system user is NOT a member of the prebuilds group
-									groupMembers, err := db.GetGroupMembersByGroupID(ctx, database.GetGroupMembersByGroupIDParams{
-										GroupID:       prebuildsGroup.ID,
-										IncludeSystem: true,
-									})
-									require.NoError(t, err)
-									require.Len(t, groupMembers, 0)
-								}
+							// Add the system user to the group if required by test case
+							if preExistingGroupMembership {
+								dbgen.GroupMember(t, db, database.GroupMemberTable{
+									GroupID: prebuildsGroup.ID,
+									UserID:  database.PrebuildsSystemUserID,
+								})
 							}
+						}
 
-							if !preExistingGroup && tc.expectGroupExists != nil && !*tc.expectGroupExists {
-								// Verify that no prebuilds group exists
-								require.Error(t, err)
-								require.True(t, errors.Is(err, sql.ErrNoRows))
-							}
+						// Setup unrelated org preset
+						dbfake.TemplateVersion(t, db).Seed(database.TemplateVersion{
+							OrganizationID: unrelatedOrg.ID,
+							CreatedBy:      owner.UserID,
+						}).Preset(database.TemplateVersionPreset{
+							DesiredInstances: sql.NullInt32{
+								Int32: 1,
+								Valid: true,
+							},
+						}).Do()
+
+						// Setup target org preset
+						dbfake.TemplateVersion(t, db).Seed(database.TemplateVersion{
+							OrganizationID: targetOrg.ID,
+							CreatedBy:      owner.UserID,
+						}).Preset(database.TemplateVersionPreset{
+							DesiredInstances: sql.NullInt32{
+								Int32: 0,
+								Valid: includePreset,
+							},
+						}).Do()
+
+						// Verify memberships before reconciliation.
+						preReconcileMemberships, err := db.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{
+							UserID: database.PrebuildsSystemUserID,
 						})
-					}
+						require.NoError(t, err)
+						expectedMembershipsBefore := []uuid.UUID{defaultOrg.ID, unrelatedOrg.ID}
+						if preExistingOrgMembership {
+							expectedMembershipsBefore = append(expectedMembershipsBefore, targetOrg.ID)
+						}
+						require.ElementsMatch(t, expectedMembershipsBefore, extractOrgIDs(preReconcileMemberships))
+
+						// Reconcile
+						reconciler := prebuilds.NewStoreMembershipReconciler(db, clock, slogtest.Make(t, nil))
+						require.NoError(t, reconciler.ReconcileAll(ctx, database.PrebuildsSystemUserID, prebuilds.PrebuiltWorkspacesGroupName))
+
+						// Verify memberships after reconciliation.
+						postReconcileMemberships, err := db.GetOrganizationsByUserID(ctx, database.GetOrganizationsByUserIDParams{
+							UserID: database.PrebuildsSystemUserID,
+						})
+						require.NoError(t, err)
+						expectedMembershipsAfter := expectedMembershipsBefore
+						if !preExistingOrgMembership && tc.expectOrgMembershipExists {
+							expectedMembershipsAfter = append(expectedMembershipsAfter, targetOrg.ID)
+						}
+						require.ElementsMatch(t, expectedMembershipsAfter, extractOrgIDs(postReconcileMemberships))
+
+						// Verify prebuilds group behavior based on expected outcomes
+						prebuildsGroup, err = db.GetGroupByOrgAndName(ctx, database.GetGroupByOrgAndNameParams{
+							OrganizationID: targetOrg.ID,
+							Name:           prebuilds.PrebuiltWorkspacesGroupName,
+						})
+						if tc.expectGroupExists {
+							require.NoError(t, err)
+							require.Equal(t, prebuilds.PrebuiltWorkspacesGroupName, prebuildsGroup.Name)
+							require.Equal(t, prebuilds.PrebuiltWorkspacesGroupDisplayName, prebuildsGroup.DisplayName)
+							require.Equal(t, int32(0), prebuildsGroup.QuotaAllowance) // Default quota should be 0
+
+							if tc.expectUserInGroup {
+								// Check that the system user is a member of the prebuilds group
+								groupMembers, err := db.GetGroupMembersByGroupID(ctx, database.GetGroupMembersByGroupIDParams{
+									GroupID:       prebuildsGroup.ID,
+									IncludeSystem: true,
+								})
+								require.NoError(t, err)
+								require.Len(t, groupMembers, 1)
+								require.Equal(t, database.PrebuildsSystemUserID, groupMembers[0].UserID)
+							}
+
+							// If no preset exists, then we do not enforce group membership:
+							if !tc.expectUserInGroup {
+								// Check that the system user is NOT a member of the prebuilds group
+								groupMembers, err := db.GetGroupMembersByGroupID(ctx, database.GetGroupMembersByGroupIDParams{
+									GroupID:       prebuildsGroup.ID,
+									IncludeSystem: true,
+								})
+								require.NoError(t, err)
+								require.Len(t, groupMembers, 0)
+							}
+						}
+
+						if !preExistingGroup && !tc.expectGroupExists {
+							// Verify that no prebuilds group exists
+							require.Error(t, err)
+							require.True(t, errors.Is(err, sql.ErrNoRows))
+						}
+					})
 				}
 			}
 		}

enterprise/coderd/prebuilds/reconcile.go

@@ -322,6 +322,12 @@ func (c *StoreReconciler) ReconcileAll(ctx context.Context) (stats prebuilds.Rec
 			return nil
 		}
 
+		membershipReconciler := NewStoreMembershipReconciler(c.store, c.clock, logger)
+		err = membershipReconciler.ReconcileAll(ctx, database.PrebuildsSystemUserID, PrebuiltWorkspacesGroupName)
+		if err != nil {
+			return xerrors.Errorf("reconcile prebuild membership: %w", err)
+		}
+
 		snapshot, err := c.SnapshotState(ctx, c.store)
 		if err != nil {
 			return xerrors.Errorf("determine current snapshot: %w", err)
@@ -334,12 +340,6 @@ func (c *StoreReconciler) ReconcileAll(ctx context.Context) (stats prebuilds.Rec
 			return nil
 		}
 
-		membershipReconciler := NewStoreMembershipReconciler(c.store, c.clock)
-		err = membershipReconciler.ReconcileAll(ctx, database.PrebuildsSystemUserID, 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 {