coder · Emyrk · Nov 1, 2024 · Oct 23, 2024 · Oct 23, 2024 · Oct 24, 2024
diff --git a/coderd/database/db.go b/coderd/database/db.go
@@ -48,13 +48,26 @@ type DBTX interface {
 	GetContext(ctx context.Context, dest interface{}, query string, args ...interface{}) error
 }
 
+func WithSerialRetryCount(count int) func(*sqlQuerier) {
+	return func(q *sqlQuerier) {
+		q.serialRetryCount = count
+	}
+}
+
 // New creates a new database store using a SQL database connection.
-func New(sdb *sql.DB) Store {
+func New(sdb *sql.DB, opts ...func(*sqlQuerier)) Store {
 	dbx := sqlx.NewDb(sdb, "postgres")
-	return &sqlQuerier{
+	q := &sqlQuerier{
 		db:  dbx,
 		sdb: dbx,
+		// This is an arbitrary number.
+		serialRetryCount: 3,
 	}
+
+	for _, opt := range opts {
+		opt(q)
+	}
+	return q
 }
 
 // TxOptions is used to pass some execution metadata to the callers.
@@ -104,6 +117,10 @@ type querier interface {
 type sqlQuerier struct {
 	sdb *sqlx.DB
 	db  DBTX
+
+	// serialRetryCount is the number of times to retry a transaction
+	// if it fails with a serialization error.
+	serialRetryCount int
 }
 
 func (*sqlQuerier) Wrappers() []string {
@@ -143,11 +160,9 @@ func (q *sqlQuerier) InTx(function func(Store) error, txOpts *TxOptions) error {
 	// If we are in a transaction already, the parent InTx call will handle the retry.
 	// We do not want to duplicate those retries.
 	if !inTx && sqlOpts.Isolation == sql.LevelSerializable {
-		// This is an arbitrarily chosen number.
-		const retryAmount = 3
 		var err error
 		attempts := 0
-		for attempts = 0; attempts < retryAmount; attempts++ {
+		for attempts = 0; attempts < q.serialRetryCount; attempts++ {
 			txOpts.executionCount++
 			err = q.runTx(function, sqlOpts)
 			if err == nil {

diff --git a/coderd/database/dbfake/builder.go b/coderd/database/dbfake/builder.go
@@ -0,0 +1,127 @@
+package dbfake
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/coder/v2/coderd/database/dbgen"
+	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/testutil"
+)
+
+type OrganizationBuilder struct {
+	t                 *testing.T
+	db                database.Store
+	seed              database.Organization
+	allUsersAllowance int32
+	members           []uuid.UUID
+	groups            map[database.Group][]uuid.UUID
+}
+
+func Organization(t *testing.T, db database.Store) OrganizationBuilder {
+	return OrganizationBuilder{
+		t:       t,
+		db:      db,
+		members: []uuid.UUID{},
+		groups:  make(map[database.Group][]uuid.UUID),
+	}
+}
+
+type OrganizationResponse struct {
+	Org           database.Organization
+	AllUsersGroup database.Group
+	Members       []database.OrganizationMember
+	Groups        []database.Group
+}
+
+func (b OrganizationBuilder) EveryoneAllowance(allowance int) OrganizationBuilder {
+	//nolint: revive // returns modified struct
+	b.allUsersAllowance = int32(allowance)
+	return b
+}
+
+func (b OrganizationBuilder) Seed(seed database.Organization) OrganizationBuilder {
+	//nolint: revive // returns modified struct
+	b.seed = seed
+	return b
+}
+
+func (b OrganizationBuilder) Members(users ...database.User) OrganizationBuilder {
+	for _, u := range users {
+		//nolint: revive // returns modified struct
+		b.members = append(b.members, u.ID)
+	}
+	return b
+}
+
+func (b OrganizationBuilder) Group(seed database.Group, members ...database.User) OrganizationBuilder {
+	//nolint: revive // returns modified struct
+	b.groups[seed] = []uuid.UUID{}
+	for _, u := range members {
+		//nolint: revive // returns modified struct
+		b.groups[seed] = append(b.groups[seed], u.ID)
+	}
+	return b
+}
+
+func (b OrganizationBuilder) Do() OrganizationResponse {
+	org := dbgen.Organization(b.t, b.db, b.seed)
+
+	ctx := testutil.Context(b.t, testutil.WaitShort)
+	//nolint:gocritic // builder code needs perms
+	ctx = dbauthz.AsSystemRestricted(ctx)
+	everyone, err := b.db.InsertAllUsersGroup(ctx, org.ID)
+	require.NoError(b.t, err)
+
+	if b.allUsersAllowance > 0 {
+		everyone, err = b.db.UpdateGroupByID(ctx, database.UpdateGroupByIDParams{
+			Name:           everyone.Name,
+			DisplayName:    everyone.DisplayName,
+			AvatarURL:      everyone.AvatarURL,
+			QuotaAllowance: b.allUsersAllowance,
+			ID:             everyone.ID,
+		})
+		require.NoError(b.t, err)
+	}
+
+	members := make([]database.OrganizationMember, 0)
+	if len(b.members) > 0 {
+		for _, u := range b.members {
+			newMem := dbgen.OrganizationMember(b.t, b.db, database.OrganizationMember{
+				UserID:         u,
+				OrganizationID: org.ID,
+				CreatedAt:      dbtime.Now(),
+				UpdatedAt:      dbtime.Now(),
+				Roles:          nil,
+			})
+			members = append(members, newMem)
+		}
+	}
+
+	groups := make([]database.Group, 0)
+	if len(b.groups) > 0 {
+		for g, users := range b.groups {
+			g.OrganizationID = org.ID
+			group := dbgen.Group(b.t, b.db, g)
+			groups = append(groups, group)
+
+			for _, u := range users {
+				dbgen.GroupMember(b.t, b.db, database.GroupMemberTable{
+					UserID:  u,
+					GroupID: group.ID,
+				})
+			}
+		}
+	}
+
+	return OrganizationResponse{
+		Org:           org,
+		AllUsersGroup: everyone,
+		Members:       members,
+		Groups:        groups,
+	}
+}
diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -407,6 +407,8 @@ func OrganizationMember(t testing.TB, db database.Store, orig database.Organizat
 }
 
 func Group(t testing.TB, db database.Store, orig database.Group) database.Group {
+	t.Helper()
+
 	name := takeFirst(orig.Name, testutil.GetRandomName(t))
 	group, err := db.InsertGroup(genCtx, database.InsertGroupParams{
 		ID:             takeFirst(orig.ID, uuid.New()),

diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go
@@ -135,7 +135,8 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) {
 		if o.dumpOnFailure {
 			t.Cleanup(func() { DumpOnFailure(t, connectionURL) })
 		}
-		db = database.New(sqlDB)
+		// Unit tests should not retry serial transaction failures.
+		db = database.New(sqlDB, database.WithSerialRetryCount(1))
 
 		ps, err = pubsub.New(context.Background(), o.logger, sqlDB, connectionURL)
 		require.NoError(t, err)

diff --git a/coderd/database/dbtestutil/tx.go b/coderd/database/dbtestutil/tx.go
@@ -0,0 +1,73 @@
+package dbtestutil
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/coderd/database"
+)
+
+type DBTx struct {
+	database.Store
+	mu       sync.Mutex
+	done     chan error
+	finalErr chan error
+}
+
+// StartTx starts a transaction and returns a DBTx object. This allows running
+// 2 transactions concurrently in a test more easily.
+// Example:
+//
+//	a := StartTx(t, db, opts)
+//	b := StartTx(t, db, opts)
+//
+//	a.GetUsers(...)
+//	b.GetUsers(...)
+//
+//	require.NoError(t, a.Done()
+func StartTx(t *testing.T, db database.Store, opts *database.TxOptions) *DBTx {
+	done := make(chan error)
+	finalErr := make(chan error)
+	txC := make(chan database.Store)
+
+	go func() {
+		t.Helper()
+		once := sync.Once{}
+		count := 0
+
+		err := db.InTx(func(store database.Store) error {
+			// InTx can be retried
+			once.Do(func() {
+				txC <- store
+			})
+			count++
+			if count > 1 {
+				// If you recursively call InTx, then don't use this.
+				t.Logf("InTx called more than once: %d", count)
+				assert.NoError(t, xerrors.New("InTx called more than once, this is not allowed with the StartTx helper"))
+			}
+
+			<-done
+			// Just return nil. The caller should be checking their own errors.
+			return nil
+		}, opts)
+		finalErr <- err
+	}()
+
+	txStore := <-txC
+	close(txC)
+
+	return &DBTx{Store: txStore, done: done, finalErr: finalErr}
+}
+
+// Done can only be called once. If you call it twice, it will panic.
+func (tx *DBTx) Done() error {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+
+	close(tx.done)
+	return <-tx.finalErr
+}
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/quotas.sql b/coderd/database/queries/quotas.sql
@@ -18,23 +18,34 @@ INNER JOIN groups ON
 WITH latest_builds AS (
 SELECT
 	DISTINCT ON
-	(workspace_id) id,
-	workspace_id,
-	daily_cost
+	(wb.workspace_id) wb.id,
 -- name: GetQuotaConsumedForUser :one 
 WITH latest_builds AS ( 
 SELECT 
 	DISTINCT ON 
 	(workspace_id) id, 
 	workspace_id, 
 	daily_cost 
 FROM 
 	workspace_builds wb 
 ORDER BY 
 	workspace_id, 
 	created_at DESC 
 ) 
 SELECT 
 	coalesce(SUM(daily_cost), 0)::BIGINT 
 FROM 
 	workspaces 
 JOIN latest_builds ON 
 	latest_builds.workspace_id = workspaces.id 
 WHERE NOT deleted AND workspaces.owner_id = $1; 
 -- name: GetQuotaConsumedForUser :one 
 WITH latest_builds AS ( 
 SELECT 
 	DISTINCT ON 
 	(workspace_id) id, 
 	workspace_id, 
 	daily_cost 
 FROM 
 	workspace_builds wb 
 ORDER BY 
 	workspace_id, 
 	created_at DESC 
 ) 
 SELECT 
 	coalesce(SUM(daily_cost), 0)::BIGINT 
 FROM 
 	workspaces 
 JOIN latest_builds ON 
 	latest_builds.workspace_id = workspaces.id 
 WHERE NOT deleted AND workspaces.owner_id = $1; 
+	wb.workspace_id,
+	wb.daily_cost
 FROM
 	workspace_builds wb
+-- This INNER JOIN prevents a seq scan of the workspace_builds table.
+-- Limit the rows to the absolute minimum required, which is all workspaces
+-- in a given organization for a given user.
+INNER JOIN
+	workspaces on wb.workspace_id = workspaces.id
+WHERE
+	workspaces.owner_id = @owner_id AND
+	workspaces.organization_id = @organization_id
 ORDER BY
-	workspace_id,
-	created_at DESC
+	wb.workspace_id,
+	wb.created_at DESC
 )
 SELECT
 	coalesce(SUM(daily_cost), 0)::BIGINT
 FROM
 	workspaces
 JOIN latest_builds ON
 	latest_builds.workspace_id = workspaces.id
-WHERE NOT
-	deleted AND
+WHERE
+	NOT deleted AND
+	-- We can likely remove these conditions since we check above.
+	-- But it does not hurt to be defensive and make sure future query changes
+	-- do not break anything.
 	workspaces.owner_id = @owner_id AND
 	workspaces.organization_id = @organization_id
 ;