chore: fix concurrent CommitQuota transactions for unrelated users/orgs (#15261)

The failure condition being fixed is `w1` and `w2` could belong
to different users, organizations, and templates and still cause a
serializable failure if run concurrently. This is because the old query 
did a `seq scan` on the `workspace_builds` table. Since that is the 
table being updated, we really want to prevent that.

So before this would fail for any 2 workspaces. Now it only fails if
`w1` and `w2` are owned by the same user and organization.

coderd/database/db.go

@@ -28,6 +28,7 @@ type Store interface {
 	wrapper
 
 	Ping(ctx context.Context) (time.Duration, error)
+	PGLocks(ctx context.Context) (PGLocks, error)
 	InTx(func(Store) error, *TxOptions) error
 }
 
@@ -48,13 +49,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 +118,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 +161,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 {
@@ -203,3 +219,10 @@ func (q *sqlQuerier) runTx(function func(Store) error, txOpts *sql.TxOptions) er
 	}
 	return nil
 }
+
+func safeString(s *string) string {
+	if s == nil {
+		return "<nil>"
+	}
+	return *s
+}

coderd/database/dbauthz/dbauthz.go

@@ -627,6 +627,10 @@ func (q *querier) Ping(ctx context.Context) (time.Duration, error) {
 	return q.db.Ping(ctx)
 }
 
+func (q *querier) PGLocks(ctx context.Context) (database.PGLocks, error) {
+	return q.db.PGLocks(ctx)
+}
+
 // InTx runs the given function in a transaction.
 func (q *querier) InTx(function func(querier database.Store) error, txOpts *database.TxOptions) error {
 	return q.db.InTx(func(tx database.Store) error {

coderd/database/dbauthz/dbauthz_test.go

@@ -152,7 +152,10 @@ func TestDBAuthzRecursive(t *testing.T) {
 		for i := 2; i < method.Type.NumIn(); i++ {
 			ins = append(ins, reflect.New(method.Type.In(i)).Elem())
 		}
-		if method.Name == "InTx" || method.Name == "Ping" || method.Name == "Wrappers" {
+		if method.Name == "InTx" ||
+			method.Name == "Ping" ||
+			method.Name == "Wrappers" ||
+			method.Name == "PGLocks" {
 			continue
 		}
 		// Log the name of the last method, so if there is a panic, it is

coderd/database/dbauthz/setup_test.go

@@ -34,6 +34,7 @@ var errMatchAny = xerrors.New("match any error")
 var skipMethods = map[string]string{
 	"InTx":           "Not relevant",
 	"Ping":           "Not relevant",
+	"PGLocks":        "Not relevant",
 	"Wrappers":       "Not relevant",
 	"AcquireLock":    "Not relevant",
 	"TryAcquireLock": "Not relevant",

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,
+	}
+}

coderd/database/dbgen/dbgen.go

@@ -408,6 +408,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()),

coderd/database/dbmem/dbmem.go

@@ -339,6 +339,10 @@ func (*FakeQuerier) Ping(_ context.Context) (time.Duration, error) {
 	return 0, nil
 }
 
+func (*FakeQuerier) PGLocks(_ context.Context) (database.PGLocks, error) {
+	return []database.PGLock{}, nil
+}
+
 func (tx *fakeTx) AcquireLock(_ context.Context, id int64) error {
 	if _, ok := tx.FakeQuerier.locks[id]; ok {
 		return xerrors.Errorf("cannot acquire lock %d: already held", id)

coderd/database/dbmetrics/querymetrics.go

@@ -66,6 +66,13 @@ func (m queryMetricsStore) Ping(ctx context.Context) (time.Duration, error) {
 	return duration, err
 }
 
+func (m queryMetricsStore) PGLocks(ctx context.Context) (database.PGLocks, error) {
+	start := time.Now()
+	locks, err := m.s.PGLocks(ctx)
+	m.queryLatencies.WithLabelValues("PGLocks").Observe(time.Since(start).Seconds())
+	return locks, err
+}
+
 func (m queryMetricsStore) InTx(f func(database.Store) error, options *database.TxOptions) error {
 	return m.dbMetrics.InTx(f, options)
 }

coderd/database/dbmock/dbmock.go

@@ -4329,6 +4329,21 @@ func (mr *MockStoreMockRecorder) OrganizationMembers(arg0, arg1 any) *gomock.Cal
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OrganizationMembers", reflect.TypeOf((*MockStore)(nil).OrganizationMembers), arg0, arg1)
 }
 
+// PGLocks mocks base method.
+func (m *MockStore) PGLocks(arg0 context.Context) (database.PGLocks, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "PGLocks", arg0)
+	ret0, _ := ret[0].(database.PGLocks)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// PGLocks indicates an expected call of PGLocks.
+func (mr *MockStoreMockRecorder) PGLocks(arg0 any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PGLocks", reflect.TypeOf((*MockStore)(nil).PGLocks), arg0)
+}
+
 // Ping mocks base method.
 func (m *MockStore) Ping(arg0 context.Context) (time.Duration, error) {
 	m.ctrl.T.Helper()

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)

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
+}

coderd/database/pglocks.go

@@ -0,0 +1,119 @@
+package database
+
+import (
+	"context"
+	"fmt"
+	"reflect"
+	"sort"
+	"strings"
+	"time"
+
+	"github.com/jmoiron/sqlx"
+
+	"github.com/coder/coder/v2/coderd/util/slice"
+)
+
+// PGLock docs see: https://www.postgresql.org/docs/current/view-pg-locks.html#VIEW-PG-LOCKS
+type PGLock struct {
+	// LockType see: https://www.postgresql.org/docs/current/monitoring-stats.html#WAIT-EVENT-LOCK-TABLE
+	LockType           *string    `db:"locktype"`
+	Database           *string    `db:"database"` // oid
+	Relation           *string    `db:"relation"` // oid
+	RelationName       *string    `db:"relation_name"`
+	Page               *int       `db:"page"`
+	Tuple              *int       `db:"tuple"`
+	VirtualXID         *string    `db:"virtualxid"`
+	TransactionID      *string    `db:"transactionid"` // xid
+	ClassID            *string    `db:"classid"`       // oid
+	ObjID              *string    `db:"objid"`         // oid
+	ObjSubID           *int       `db:"objsubid"`
+	VirtualTransaction *string    `db:"virtualtransaction"`
+	PID                int        `db:"pid"`
+	Mode               *string    `db:"mode"`
+	Granted            bool       `db:"granted"`
+	FastPath           *bool      `db:"fastpath"`
+	WaitStart          *time.Time `db:"waitstart"`
+}
+
+func (l PGLock) Equal(b PGLock) bool {
+	// Lazy, but hope this works
+	return reflect.DeepEqual(l, b)
+}
+
+func (l PGLock) String() string {
+	granted := "granted"
+	if !l.Granted {
+		granted = "waiting"
+	}
+	var details string
+	switch safeString(l.LockType) {
+	case "relation":
+		details = ""
+	case "page":
+		details = fmt.Sprintf("page=%d", *l.Page)
+	case "tuple":
+		details = fmt.Sprintf("page=%d tuple=%d", *l.Page, *l.Tuple)
+	case "virtualxid":
+		details = "waiting to acquire virtual tx id lock"
+	default:
+		details = "???"
+	}
+	return fmt.Sprintf("%d-%5s [%s] %s/%s/%s: %s",
+		l.PID,
+		safeString(l.TransactionID),
+		granted,
+		safeString(l.RelationName),
+		safeString(l.LockType),
+		safeString(l.Mode),
+		details,
+	)
+}
+
+// PGLocks returns a list of all locks in the database currently in use.
+func (q *sqlQuerier) PGLocks(ctx context.Context) (PGLocks, error) {
+	rows, err := q.sdb.QueryContext(ctx, `
+	SELECT
+		relation::regclass AS relation_name,
+	    *
+	FROM pg_locks;
+	`)
+	if err != nil {
+		return nil, err
+	}
+
+	defer rows.Close()
+
+	var locks []PGLock
+	err = sqlx.StructScan(rows, &locks)
+	if err != nil {
+		return nil, err
+	}
+
+	return locks, err
+}
+
+type PGLocks []PGLock
+
+func (l PGLocks) String() string {
+	// Try to group things together by relation name.
+	sort.Slice(l, func(i, j int) bool {
+		return safeString(l[i].RelationName) < safeString(l[j].RelationName)
+	})
+
+	var out strings.Builder
+	for i, lock := range l {
+		if i != 0 {
+			_, _ = out.WriteString("\n")
+		}
+		_, _ = out.WriteString(lock.String())
+	}
+	return out.String()
+}
+
+// Difference returns the difference between two sets of locks.
+// This is helpful to determine what changed between the two sets.
+func (l PGLocks) Difference(to PGLocks) (new PGLocks, removed PGLocks) {
+	return slice.SymmetricDifferenceFunc(l, to, func(a, b PGLock) bool {
+		return a.Equal(b)
+	})
+}

coderd/database/queries.sql.go

@@ -6736,23 +6736,33 @@ const getQuotaConsumedForUser = `-- name: GetQuotaConsumedForUser :one
 WITH latest_builds AS (
 SELECT
 	DISTINCT ON
-	(workspace_id) id,
-	workspace_id,
-	daily_cost
+	(wb.workspace_id) 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 = $1 AND
+	workspaces.organization_id = $2
 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
+INNER 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 = $1 AND
 	workspaces.organization_id = $2
 `

coderd/database/queries/quotas.sql

@@ -18,23 +18,33 @@ INNER JOIN groups ON
 WITH latest_builds AS (
 SELECT
 	DISTINCT ON
-	(workspace_id) id,
-	workspace_id,
-	daily_cost
+	(wb.workspace_id) 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
+INNER 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
 ;