feat: use JWT ticket to avoid DB queries on apps (#6148)

Issue a JWT ticket on the first request with a short expiry that
contains details about which workspace/agent/app combo the ticket is
valid for.

coderd/database/dbauthz/querier.go

@@ -19,6 +19,14 @@ func (q *querier) Ping(ctx context.Context) (time.Duration, error) {
 	return q.db.Ping(ctx)
 }
 
+func (q *querier) AcquireLock(ctx context.Context, id int64) error {
+	return q.db.AcquireLock(ctx, id)
+}
+
+func (q *querier) TryAcquireLock(ctx context.Context, id int64) (bool, error) {
+	return q.db.TryAcquireLock(ctx, id)
+}
+
 // InTx runs the given function in a transaction.
 func (q *querier) InTx(function func(querier database.Store) error, txOpts *sql.TxOptions) error {
 	return q.db.InTx(func(tx database.Store) error {
@@ -317,6 +325,16 @@ func (q *querier) GetLogoURL(ctx context.Context) (string, error) {
 	return q.db.GetLogoURL(ctx)
 }
 
+func (q *querier) GetAppSigningKey(ctx context.Context) (string, error) {
+	// No authz checks
+	return q.db.GetAppSigningKey(ctx)
+}
+
+func (q *querier) InsertAppSigningKey(ctx context.Context, data string) error {
+	// No authz checks as this is done during startup
+	return q.db.InsertAppSigningKey(ctx, data)
+}
+
 func (q *querier) GetServiceBanner(ctx context.Context) (string, error) {
 	// No authz checks
 	return q.db.GetServiceBanner(ctx)

coderd/database/dbfake/databasefake.go

@@ -64,6 +64,7 @@ func New() database.Store {
 			workspaceApps:             make([]database.WorkspaceApp, 0),
 			workspaces:                make([]database.Workspace, 0),
 			licenses:                  make([]database.License, 0),
+			locks:                     map[int64]struct{}{},
 		},
 	}
 }
@@ -89,6 +90,11 @@ type fakeQuerier struct {
 	*data
 }
 
+type fakeTx struct {
+	*fakeQuerier
+	locks map[int64]struct{}
+}
+
 type data struct {
 	// Legacy tables
 	apiKeys             []database.APIKey
@@ -124,11 +130,15 @@ type data struct {
 	workspaceResources        []database.WorkspaceResource
 	workspaces                []database.Workspace
 
+	// Locks is a map of lock names. Any keys within the map are currently
+	// locked.
+	locks           map[int64]struct{}
 	deploymentID    string
 	derpMeshKey     string
 	lastUpdateCheck []byte
 	serviceBanner   []byte
 	logoURL         string
+	appSigningKey   string
 	lastLicenseID   int32
 }
 
@@ -196,11 +206,50 @@ func (*fakeQuerier) Ping(_ context.Context) (time.Duration, error) {
 	return 0, nil
 }
 
+func (*fakeQuerier) AcquireLock(_ context.Context, _ int64) error {
+	return xerrors.New("AcquireLock must only be called within a transaction")
+}
+
+func (*fakeQuerier) TryAcquireLock(_ context.Context, _ int64) (bool, error) {
+	return false, xerrors.New("TryAcquireLock must only be called within a transaction")
+}
+
+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)
+	}
+	tx.fakeQuerier.locks[id] = struct{}{}
+	tx.locks[id] = struct{}{}
+	return nil
+}
+
+func (tx *fakeTx) TryAcquireLock(_ context.Context, id int64) (bool, error) {
+	if _, ok := tx.fakeQuerier.locks[id]; ok {
+		return false, nil
+	}
+	tx.fakeQuerier.locks[id] = struct{}{}
+	tx.locks[id] = struct{}{}
+	return true, nil
+}
+
+func (tx *fakeTx) releaseLocks() {
+	for id := range tx.locks {
+		delete(tx.fakeQuerier.locks, id)
+	}
+	tx.locks = map[int64]struct{}{}
+}
+
 // InTx doesn't rollback data properly for in-memory yet.
 func (q *fakeQuerier) InTx(fn func(database.Store) error, _ *sql.TxOptions) error {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
-	return fn(&fakeQuerier{mutex: inTxMutex{}, data: q.data})
+	tx := &fakeTx{
+		fakeQuerier: &fakeQuerier{mutex: inTxMutex{}, data: q.data},
+		locks:       map[int64]struct{}{},
+	}
+	defer tx.releaseLocks()
+
+	return fn(tx)
 }
 
 func (q *fakeQuerier) AcquireProvisionerJob(_ context.Context, arg database.AcquireProvisionerJobParams) (database.ProvisionerJob, error) {
@@ -4004,6 +4053,21 @@ func (q *fakeQuerier) GetLogoURL(_ context.Context) (string, error) {
 	return q.logoURL, nil
 }
 
+func (q *fakeQuerier) GetAppSigningKey(_ context.Context) (string, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	return q.appSigningKey, nil
+}
+
+func (q *fakeQuerier) InsertAppSigningKey(_ context.Context, data string) error {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	q.appSigningKey = data
+	return nil
+}
+
 func (q *fakeQuerier) InsertLicense(
 	_ context.Context, arg database.InsertLicenseParams,
 ) (database.License, error) {

coderd/database/lock.go

@@ -0,0 +1,8 @@
+package database
+
+// Well-known lock IDs for lock functions in the database. These should not
+// change. If locks are deprecated, they should be kept to avoid reusing the
+// same ID.
+const (
+	LockIDDeploymentSetup = iota + 1
+)

coderd/database/querier.go

@@ -12,6 +12,13 @@ import (
 )
 
 type sqlcQuerier interface {
+	// Blocks until the lock is acquired.
+	//
+	// This must be called from within a transaction. The lock will be automatically
+	// released when the transaction ends.
+	//
+	// Use database.LockID() to generate a unique lock ID from a string.
+	AcquireLock(ctx context.Context, pgAdvisoryXactLock int64) error
 	// Acquires the lock for a single job that isn't started, completed,
 	// canceled, and that matches an array of provisioner types.
 	//
@@ -36,6 +43,7 @@ type sqlcQuerier interface {
 	GetAPIKeysByUserID(ctx context.Context, arg GetAPIKeysByUserIDParams) ([]APIKey, error)
 	GetAPIKeysLastUsedAfter(ctx context.Context, lastUsed time.Time) ([]APIKey, error)
 	GetActiveUserCount(ctx context.Context) (int64, error)
+	GetAppSigningKey(ctx context.Context) (string, error)
 	// GetAuditLogsBefore retrieves `row_limit` number of audit logs before the provided
 	// ID.
 	GetAuditLogsOffset(ctx context.Context, arg GetAuditLogsOffsetParams) ([]GetAuditLogsOffsetRow, error)
@@ -139,6 +147,7 @@ type sqlcQuerier interface {
 	// for simplicity since all users is
 	// every member of the org.
 	InsertAllUsersGroup(ctx context.Context, organizationID uuid.UUID) (Group, error)
+	InsertAppSigningKey(ctx context.Context, value string) error
 	InsertAuditLog(ctx context.Context, arg InsertAuditLogParams) (AuditLog, error)
 	InsertDERPMeshKey(ctx context.Context, value string) error
 	InsertDeploymentID(ctx context.Context, value string) error
@@ -177,6 +186,13 @@ type sqlcQuerier interface {
 	InsertWorkspaceResourceMetadata(ctx context.Context, arg InsertWorkspaceResourceMetadataParams) ([]WorkspaceResourceMetadatum, error)
 	ParameterValue(ctx context.Context, id uuid.UUID) (ParameterValue, error)
 	ParameterValues(ctx context.Context, arg ParameterValuesParams) ([]ParameterValue, error)
+	// Non blocking lock. Returns true if the lock was acquired, false otherwise.
+	//
+	// This must be called from within a transaction. The lock will be automatically
+	// released when the transaction ends.
+	//
+	// Use database.LockID() to generate a unique lock ID from a string.
+	TryAcquireLock(ctx context.Context, pgTryAdvisoryXactLock int64) (bool, error)
 	UpdateAPIKeyByID(ctx context.Context, arg UpdateAPIKeyByIDParams) error
 	UpdateGitAuthLink(ctx context.Context, arg UpdateGitAuthLinkParams) (GitAuthLink, error)
 	UpdateGitSSHKey(ctx context.Context, arg UpdateGitSSHKeyParams) (GitSSHKey, error)

coderd/database/queries.sql.go

@@ -1437,6 +1437,38 @@ func (q *sqlQuerier) InsertLicense(ctx context.Context, arg InsertLicenseParams)
 	return i, err
 }
 
+const acquireLock = `-- name: AcquireLock :exec
+SELECT pg_advisory_xact_lock($1)
+`
+
+// Blocks until the lock is acquired.
+//
+// This must be called from within a transaction. The lock will be automatically
+// released when the transaction ends.
+//
+// Use database.LockID() to generate a unique lock ID from a string.
+func (q *sqlQuerier) AcquireLock(ctx context.Context, pgAdvisoryXactLock int64) error {
+	_, err := q.db.ExecContext(ctx, acquireLock, pgAdvisoryXactLock)
+	return err
+}
+
+const tryAcquireLock = `-- name: TryAcquireLock :one
+SELECT pg_try_advisory_xact_lock($1)
+`
+
+// Non blocking lock. Returns true if the lock was acquired, false otherwise.
+//
+// This must be called from within a transaction. The lock will be automatically
+// released when the transaction ends.
+//
+// Use database.LockID() to generate a unique lock ID from a string.
+func (q *sqlQuerier) TryAcquireLock(ctx context.Context, pgTryAdvisoryXactLock int64) (bool, error) {
+	row := q.db.QueryRowContext(ctx, tryAcquireLock, pgTryAdvisoryXactLock)
+	var pg_try_advisory_xact_lock bool
+	err := row.Scan(&pg_try_advisory_xact_lock)
+	return pg_try_advisory_xact_lock, err
+}
+
 const getOrganizationIDsByMemberIDs = `-- name: GetOrganizationIDsByMemberIDs :many
 SELECT
     user_id, array_agg(organization_id) :: uuid [ ] AS "organization_IDs"
@@ -2909,6 +2941,17 @@ func (q *sqlQuerier) UpdateReplica(ctx context.Context, arg UpdateReplicaParams)
 	return i, err
 }
 
+const getAppSigningKey = `-- name: GetAppSigningKey :one
+SELECT value FROM site_configs WHERE key = 'app_signing_key'
+`
+
+func (q *sqlQuerier) GetAppSigningKey(ctx context.Context) (string, error) {
+	row := q.db.QueryRowContext(ctx, getAppSigningKey)
+	var value string
+	err := row.Scan(&value)
+	return value, err
+}
+
 const getDERPMeshKey = `-- name: GetDERPMeshKey :one
 SELECT value FROM site_configs WHERE key = 'derp_mesh_key'
 `
@@ -2964,6 +3007,15 @@ func (q *sqlQuerier) GetServiceBanner(ctx context.Context) (string, error) {
 	return value, err
 }
 
+const insertAppSigningKey = `-- name: InsertAppSigningKey :exec
+INSERT INTO site_configs (key, value) VALUES ('app_signing_key', $1)
+`
+
+func (q *sqlQuerier) InsertAppSigningKey(ctx context.Context, value string) error {
+	_, err := q.db.ExecContext(ctx, insertAppSigningKey, value)
+	return err
+}
+
 const insertDERPMeshKey = `-- name: InsertDERPMeshKey :exec
 INSERT INTO site_configs (key, value) VALUES ('derp_mesh_key', $1)
 `

coderd/database/queries/lock.sql

@@ -0,0 +1,17 @@
+-- name: AcquireLock :exec
+-- Blocks until the lock is acquired.
+--
+-- This must be called from within a transaction. The lock will be automatically
+-- released when the transaction ends.
+--
+-- Use database.LockID() to generate a unique lock ID from a string.
+SELECT pg_advisory_xact_lock($1);
+
+-- name: TryAcquireLock :one
+-- Non blocking lock. Returns true if the lock was acquired, false otherwise.
+--
+-- This must be called from within a transaction. The lock will be automatically
+-- released when the transaction ends.
+--
+-- Use database.LockID() to generate a unique lock ID from a string.
+SELECT pg_try_advisory_xact_lock($1);

coderd/database/queries/siteconfig.sql

@@ -30,3 +30,9 @@ ON CONFLICT (key) DO UPDATE SET value = $1 WHERE site_configs.key = 'logo_url';
 
 -- name: GetLogoURL :one
 SELECT value FROM site_configs WHERE key = 'logo_url';
+
+-- name: GetAppSigningKey :one
+SELECT value FROM site_configs WHERE key = 'app_signing_key';
+
+-- name: InsertAppSigningKey :exec
+INSERT INTO site_configs (key, value) VALUES ('app_signing_key', $1);