coder · Emyrk · Nov 1, 2024 · Oct 31, 2024 · Oct 31, 2024 · Oct 31, 2024
diff --git a/coderd/database/db.go b/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
 }
 
@@ -218,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
+}
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -603,6 +603,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 {

diff --git a/coderd/database/dbauthz/dbauthz_test.go b/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

diff --git a/coderd/database/dbauthz/setup_test.go b/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",

diff --git a/coderd/database/dbmem/dbmem.go b/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)

diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/pglocks.go b/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)
+	})
+}