feat(coderd): insert provisioner daemons (#11207)

* Adds UpdateProvisionerDaemonLastSeenAt
* Adds heartbeat to provisioner daemons
* Inserts provisioner daemons to database upon start
* Ensures TagOwner is an empty string and not nil
* Adds COALESCE() in idx_provisioner_daemons_name_owner_key

coderd/database/dbauthz/dbauthz.go

@@ -232,6 +232,7 @@ var (
 					rbac.ResourceOrganization.Type:       {rbac.ActionCreate},
 					rbac.ResourceOrganizationMember.Type: {rbac.ActionCreate},
 					rbac.ResourceOrgRoleAssignment.Type:  {rbac.ActionCreate},
+					rbac.ResourceProvisionerDaemon.Type:  {rbac.ActionCreate, rbac.ActionUpdate},
 					rbac.ResourceUser.Type:               {rbac.ActionCreate, rbac.ActionUpdate, rbac.ActionDelete},
 					rbac.ResourceUserData.Type:           {rbac.ActionCreate, rbac.ActionUpdate},
 					rbac.ResourceWorkspace.Type:          {rbac.ActionUpdate},
@@ -2499,6 +2500,13 @@ func (q *querier) UpdateMemberRoles(ctx context.Context, arg database.UpdateMemb
 	return q.db.UpdateMemberRoles(ctx, arg)
 }
 
+func (q *querier) UpdateProvisionerDaemonLastSeenAt(ctx context.Context, arg database.UpdateProvisionerDaemonLastSeenAtParams) error {
+	if err := q.authorizeContext(ctx, rbac.ActionUpdate, rbac.ResourceProvisionerDaemon); err != nil {
+		return err
+	}
+	return q.db.UpdateProvisionerDaemonLastSeenAt(ctx, arg)
+}
+
 // TODO: We need to create a ProvisionerJob resource type
 func (q *querier) UpdateProvisionerJobByID(ctx context.Context, arg database.UpdateProvisionerJobByIDParams) error {
 	// if err := q.authorizeContext(ctx, rbac.ActionUpdate, rbac.ResourceSystem); err != nil {

coderd/database/dbauthz/dbauthz_test.go

@@ -1592,6 +1592,18 @@ func (s *MethodTestSuite) TestExtraMethods() {
 		s.NoError(err, "insert provisioner daemon")
 		check.Args().Asserts(rbac.ResourceSystem, rbac.ActionDelete)
 	}))
+	s.Run("UpdateProvisionerDaemonLastSeenAt", s.Subtest(func(db database.Store, check *expects) {
+		d, err := db.UpsertProvisionerDaemon(context.Background(), database.UpsertProvisionerDaemonParams{
+			Tags: database.StringMap(map[string]string{
+				provisionersdk.TagScope: provisionersdk.ScopeOrganization,
+			}),
+		})
+		s.NoError(err, "insert provisioner daemon")
+		check.Args(database.UpdateProvisionerDaemonLastSeenAtParams{
+			ID:         d.ID,
+			LastSeenAt: sql.NullTime{Time: dbtime.Now(), Valid: true},
+		}).Asserts(rbac.ResourceProvisionerDaemon, rbac.ActionUpdate)
+	}))
 }
 
 // All functions in this method test suite are not implemented in dbmem, but

coderd/database/dbauthz/setup_test.go

@@ -28,7 +28,7 @@ import (
 	"github.com/coder/coder/v2/coderd/util/slice"
 )
 
-var errMatchAny = errors.New("match any error")
+var errMatchAny = xerrors.New("match any error")
 
 var skipMethods = map[string]string{
 	"InTx":           "Not relevant",

coderd/database/dbmem/dbmem.go

@@ -5945,6 +5945,28 @@ func (q *FakeQuerier) UpdateMemberRoles(_ context.Context, arg database.UpdateMe
 	return database.OrganizationMember{}, sql.ErrNoRows
 }
 
+func (q *FakeQuerier) UpdateProvisionerDaemonLastSeenAt(_ context.Context, arg database.UpdateProvisionerDaemonLastSeenAtParams) error {
+	err := validateDatabaseType(arg)
+	if err != nil {
+		return err
+	}
+
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	for idx := range q.provisionerDaemons {
+		if q.provisionerDaemons[idx].ID != arg.ID {
+			continue
+		}
+		if q.provisionerDaemons[idx].LastSeenAt.Time.After(arg.LastSeenAt.Time) {
+			continue
+		}
+		q.provisionerDaemons[idx].LastSeenAt = arg.LastSeenAt
+		return nil
+	}
+	return sql.ErrNoRows
+}
+
 func (q *FakeQuerier) UpdateProvisionerJobByID(_ context.Context, arg database.UpdateProvisionerJobByIDParams) error {
 	if err := validateDatabaseType(arg); err != nil {
 		return err

coderd/database/dbmetrics/dbmetrics.go

@@ -1593,6 +1593,13 @@ func (m metricsStore) UpdateMemberRoles(ctx context.Context, arg database.Update
 	return member, err
 }
 
+func (m metricsStore) UpdateProvisionerDaemonLastSeenAt(ctx context.Context, arg database.UpdateProvisionerDaemonLastSeenAtParams) error {
+	start := time.Now()
+	r0 := m.s.UpdateProvisionerDaemonLastSeenAt(ctx, arg)
+	m.queryLatencies.WithLabelValues("UpdateProvisionerDaemonLastSeenAt").Observe(time.Since(start).Seconds())
+	return r0
+}
+
 func (m metricsStore) UpdateProvisionerJobByID(ctx context.Context, arg database.UpdateProvisionerJobByIDParams) error {
 	start := time.Now()
 	err := m.s.UpdateProvisionerJobByID(ctx, arg)

coderd/database/dbmock/dbmock.go

@@ -3362,6 +3362,20 @@ func (mr *MockStoreMockRecorder) UpdateMemberRoles(arg0, arg1 interface{}) *gomo
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateMemberRoles", reflect.TypeOf((*MockStore)(nil).UpdateMemberRoles), arg0, arg1)
 }
 
+// UpdateProvisionerDaemonLastSeenAt mocks base method.
+func (m *MockStore) UpdateProvisionerDaemonLastSeenAt(arg0 context.Context, arg1 database.UpdateProvisionerDaemonLastSeenAtParams) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateProvisionerDaemonLastSeenAt", arg0, arg1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// UpdateProvisionerDaemonLastSeenAt indicates an expected call of UpdateProvisionerDaemonLastSeenAt.
+func (mr *MockStoreMockRecorder) UpdateProvisionerDaemonLastSeenAt(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateProvisionerDaemonLastSeenAt", reflect.TypeOf((*MockStore)(nil).UpdateProvisionerDaemonLastSeenAt), arg0, arg1)
+}
+
 // UpdateProvisionerJobByID mocks base method.
 func (m *MockStore) UpdateProvisionerJobByID(arg0 context.Context, arg1 database.UpdateProvisionerJobByIDParams) error {
 	m.ctrl.T.Helper()

coderd/database/dump.sql

@@ -1417,9 +1417,9 @@ CREATE UNIQUE INDEX idx_organization_name ON organizations USING btree (name);
 
 CREATE UNIQUE INDEX idx_organization_name_lower ON organizations USING btree (lower(name));
 
-CREATE UNIQUE INDEX idx_provisioner_daemons_name_owner_key ON provisioner_daemons USING btree (name, lower((tags ->> 'owner'::text)));
+CREATE UNIQUE INDEX idx_provisioner_daemons_name_owner_key ON provisioner_daemons USING btree (name, lower(COALESCE((tags ->> 'owner'::text), ''::text)));
 
-COMMENT ON INDEX idx_provisioner_daemons_name_owner_key IS 'Relax uniqueness constraint for provisioner daemon names';
+COMMENT ON INDEX idx_provisioner_daemons_name_owner_key IS 'Allow unique provisioner daemon names by user';
 
 CREATE INDEX idx_tailnet_agents_coordinator ON tailnet_agents USING btree (coordinator_id);
 

coderd/database/migrations/000181_coalesce_provisioner_daemon_idx_owner.down.sql

@@ -0,0 +1,8 @@
+DROP INDEX IF EXISTS idx_provisioner_daemons_name_owner_key;
+
+CREATE UNIQUE INDEX IF NOT EXISTS idx_provisioner_daemons_name_owner_key
+	ON provisioner_daemons
+	USING btree (name, lower((tags->>'owner')::text));
+
+COMMENT ON INDEX idx_provisioner_daemons_name_owner_key
+	IS 'Relax uniqueness constraint for provisioner daemon names';

coderd/database/migrations/000181_coalesce_provisioner_daemon_idx_owner.up.sql

@@ -0,0 +1,8 @@
+DROP INDEX IF EXISTS idx_provisioner_daemons_name_owner_key;
+
+CREATE UNIQUE INDEX IF NOT EXISTS idx_provisioner_daemons_name_owner_key
+	ON provisioner_daemons
+	USING btree (name, LOWER(COALESCE(tags->>'owner', '')::text));
+
+COMMENT ON INDEX idx_provisioner_daemons_name_owner_key
+	IS 'Allow unique provisioner daemon names by user';

coderd/database/querier.go

@@ -318,6 +318,7 @@ type sqlcQuerier interface {
 	UpdateGroupByID(ctx context.Context, arg UpdateGroupByIDParams) (Group, error)
 	UpdateInactiveUsersToDormant(ctx context.Context, arg UpdateInactiveUsersToDormantParams) ([]UpdateInactiveUsersToDormantRow, error)
 	UpdateMemberRoles(ctx context.Context, arg UpdateMemberRolesParams) (OrganizationMember, error)
+	UpdateProvisionerDaemonLastSeenAt(ctx context.Context, arg UpdateProvisionerDaemonLastSeenAtParams) error
 	UpdateProvisionerJobByID(ctx context.Context, arg UpdateProvisionerJobByIDParams) error
 	UpdateProvisionerJobWithCancelByID(ctx context.Context, arg UpdateProvisionerJobWithCancelByIDParams) error
 	UpdateProvisionerJobWithCompleteByID(ctx context.Context, arg UpdateProvisionerJobWithCompleteByIDParams) error

coderd/database/queries.sql.go

@@ -3057,6 +3057,26 @@ func (q *sqlQuerier) GetProvisionerDaemons(ctx context.Context) ([]ProvisionerDa
 	return items, nil
 }
 
+const updateProvisionerDaemonLastSeenAt = `-- name: UpdateProvisionerDaemonLastSeenAt :exec
+UPDATE provisioner_daemons
+SET
+	last_seen_at = $1
+WHERE
+	id = $2
+AND
+	last_seen_at <= $1
+`
+
+type UpdateProvisionerDaemonLastSeenAtParams struct {
+	LastSeenAt sql.NullTime `db:"last_seen_at" json:"last_seen_at"`
+	ID         uuid.UUID    `db:"id" json:"id"`
+}
+
+func (q *sqlQuerier) UpdateProvisionerDaemonLastSeenAt(ctx context.Context, arg UpdateProvisionerDaemonLastSeenAtParams) error {
+	_, err := q.db.ExecContext(ctx, updateProvisionerDaemonLastSeenAt, arg.LastSeenAt, arg.ID)
+	return err
+}
+
 const upsertProvisionerDaemon = `-- name: UpsertProvisionerDaemon :one
 INSERT INTO
 	provisioner_daemons (
@@ -3078,7 +3098,7 @@ VALUES (
 	$5,
 	$6,
 	$7
-) ON CONFLICT("name", lower((tags ->> 'owner'::text))) DO UPDATE SET
+) ON CONFLICT("name", LOWER(COALESCE(tags ->> 'owner'::text, ''::text))) DO UPDATE SET
 	provisioners = $3,
 	tags = $4,
 	last_seen_at = $5,

coderd/database/queries/provisionerdaemons.sql

@@ -35,7 +35,7 @@ VALUES (
 	@last_seen_at,
 	@version,
 	@api_version
-) ON CONFLICT("name", lower((tags ->> 'owner'::text))) DO UPDATE SET
+) ON CONFLICT("name", LOWER(COALESCE(tags ->> 'owner'::text, ''::text))) DO UPDATE SET
 	provisioners = @provisioners,
 	tags = @tags,
 	last_seen_at = @last_seen_at,
@@ -45,3 +45,12 @@ WHERE
 	-- Only ones with the same tags are allowed clobber
 	provisioner_daemons.tags <@ @tags :: jsonb
 RETURNING *;
+
+-- name: UpdateProvisionerDaemonLastSeenAt :exec
+UPDATE provisioner_daemons
+SET
+	last_seen_at = @last_seen_at
+WHERE
+	id = @id
+AND
+	last_seen_at <= @last_seen_at;

coderd/database/unique_constraint.go

@@ -65,7 +65,7 @@ const (
 	UniqueIndexAPIKeyName                                   UniqueConstraint = "idx_api_key_name"                                         // CREATE UNIQUE INDEX idx_api_key_name ON api_keys USING btree (user_id, token_name) WHERE (login_type = 'token'::login_type);
 	UniqueIndexOrganizationName                             UniqueConstraint = "idx_organization_name"                                    // CREATE UNIQUE INDEX idx_organization_name ON organizations USING btree (name);
 	UniqueIndexOrganizationNameLower                        UniqueConstraint = "idx_organization_name_lower"                              // CREATE UNIQUE INDEX idx_organization_name_lower ON organizations USING btree (lower(name));
-	UniqueIndexProvisionerDaemonsNameOwnerKey               UniqueConstraint = "idx_provisioner_daemons_name_owner_key"                   // CREATE UNIQUE INDEX idx_provisioner_daemons_name_owner_key ON provisioner_daemons USING btree (name, lower((tags ->> 'owner'::text)));
+	UniqueIndexProvisionerDaemonsNameOwnerKey               UniqueConstraint = "idx_provisioner_daemons_name_owner_key"                   // CREATE UNIQUE INDEX idx_provisioner_daemons_name_owner_key ON provisioner_daemons USING btree (name, lower(COALESCE((tags ->> 'owner'::text), ''::text)));
 	UniqueIndexUsersEmail                                   UniqueConstraint = "idx_users_email"                                          // CREATE UNIQUE INDEX idx_users_email ON users USING btree (email) WHERE (deleted = false);
 	UniqueIndexUsersUsername                                UniqueConstraint = "idx_users_username"                                       // CREATE UNIQUE INDEX idx_users_username ON users USING btree (username) WHERE (deleted = false);
 	UniqueTemplatesOrganizationIDNameIndex                  UniqueConstraint = "templates_organization_id_name_idx"                       // CREATE UNIQUE INDEX templates_organization_id_name_idx ON templates USING btree (organization_id, lower((name)::text)) WHERE (deleted = false);