Implement stats as a Prometheus collector

coder · code-asher · May 13, 2022 · May 13, 2022 · May 13, 2022 · May 16, 2022
commit 63252b1c9d6d41d8736dca5bb918c0694754de45
diff --git a/coderd/monitoring/collector.go b/coderd/monitoring/collector.go
@@ -0,0 +1,122 @@
+package monitoring
+
+import (
+	"context"
+	"database/sql"
+	"sync"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/coderd/database"
+)
+
+// Collector implements prometheus.Collector and collects statistics from the
+// provided database.
+type Collector struct {
+	ctx                context.Context
+	db                 database.Store
+	users              *prometheus.Desc
+	workspaces         *prometheus.Desc
+	workspaceResources *prometheus.Desc
+}
+
+func NewCollector(ctx context.Context, db database.Store) *Collector {
+	return &Collector{
+		ctx: ctx,
+		db:  db,
+		users: prometheus.NewDesc(
+			"coder_users",
+			"The users in a Coder deployment.",
+			nil,
+			nil,
+		),
+		workspaces: prometheus.NewDesc(
+			"coder_workspaces",
+			"The workspaces in a Coder deployment.",
+			nil,
+			nil,
+		),
+		workspaceResources: prometheus.NewDesc(
+			"coder_workspace_resources",
+			"The workspace resources in a Coder deployment.",
+			[]string{
+				"workspace_resource_type",
+			},
+			nil,
+		),
+	}
+}
+
+// Describe implements prometheus.Collector.
+func (c *Collector) Describe(ch chan<- *prometheus.Desc) {
+	ch <- c.users
+	ch <- c.workspaces
+	ch <- c.workspaceResources
+}
+
+// Collect implements prometheus.Collector.
+func (c *Collector) Collect(ch chan<- prometheus.Metric) {
+	var wg sync.WaitGroup
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		dbUsers, err := c.db.GetUsers(c.ctx, database.GetUsersParams{})
+		if err != nil && !xerrors.Is(err, sql.ErrNoRows) {
+			ch <- prometheus.NewInvalidMetric(c.users, err)
+			return
+		}
+
+		ch <- prometheus.MustNewConstMetric(
+			c.users,
+			prometheus.GaugeValue,
+			float64(len(dbUsers)),
+		)
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		dbWorkspaces, err := c.db.GetWorkspaces(c.ctx, false)
+		if err != nil && !xerrors.Is(err, sql.ErrNoRows) {
+			ch <- prometheus.NewInvalidMetric(c.workspaces, err)
+			return
+		}
+
+		ch <- prometheus.MustNewConstMetric(
+			c.workspaces,
+			prometheus.GaugeValue,
+			float64(len(dbWorkspaces)),
+		)
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		dbWorkspaceResources, err := c.db.GetWorkspaceResources(c.ctx)
+		if err != nil && !xerrors.Is(err, sql.ErrNoRows) {
+			ch <- prometheus.NewInvalidMetric(c.workspaceResources, err)
+			return
+		}
+
+		resourcesByType := map[string][]database.WorkspaceResource{}
+		for _, dbwr := range dbWorkspaceResources {
+			resourcesByType[dbwr.Type] = append(resourcesByType[dbwr.Type], dbwr)
+		}
+
+		for resourceType, resources := range resourcesByType {
+			ch <- prometheus.MustNewConstMetric(
+				c.workspaceResources,
+				prometheus.GaugeValue,
+				float64(len(resources)),
+				resourceType,
+			)
+		}
+	}()
+
+	wg.Wait()
+}
diff --git a/coderd/monitoring/collector_test.go b/coderd/monitoring/collector_test.go
@@ -0,0 +1,100 @@
+package monitoring_test
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/coder/coder/coderd/database"
+	"github.com/coder/coder/coderd/database/databasefake"
+	"github.com/coder/coder/coderd/monitoring"
+	"github.com/google/uuid"
+	"github.com/prometheus/client_golang/prometheus/testutil"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCollector(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	db := databasefake.New()
+	populateDB(ctx, db)
+
+	collector := monitoring.NewCollector(ctx, db)
+	expected := `
+		# HELP coder_users The users in a Coder deployment.
+        # TYPE coder_users gauge
+        coder_users 1
+        # HELP coder_workspace_resources The workspace resources in a Coder deployment.
+        # TYPE coder_workspace_resources gauge
+        coder_workspace_resources{workspace_resource_type="google_compute_instance"} 2
+        # HELP coder_workspaces The workspaces in a Coder deployment.
+        # TYPE coder_workspaces gauge
+        coder_workspaces 2
+	`
+	require.NoError(t, testutil.CollectAndCompare(collector, strings.NewReader(expected)))
+}
+
+func populateDB(ctx context.Context, db database.Store) {
+	user, _ := db.InsertUser(ctx, database.InsertUserParams{
+		ID:       uuid.New(),
+		Username: "kyle",
+	})
+	org, _ := db.InsertOrganization(ctx, database.InsertOrganizationParams{
+		ID:   uuid.New(),
+		Name: "potato",
+	})
+	template, _ := db.InsertTemplate(ctx, database.InsertTemplateParams{
+		ID:             uuid.New(),
+		Name:           "something",
+		OrganizationID: org.ID,
+	})
+	workspace, _ := db.InsertWorkspace(ctx, database.InsertWorkspaceParams{
+		ID:             uuid.New(),
+		OwnerID:        user.ID,
+		OrganizationID: org.ID,
+		TemplateID:     template.ID,
+		Name:           "banana1",
+	})
+	job, _ := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{
+		ID:             uuid.New(),
+		OrganizationID: org.ID,
+	})
+	version, _ := db.InsertTemplateVersion(ctx, database.InsertTemplateVersionParams{
+		ID: uuid.New(),
+		TemplateID: uuid.NullUUID{
+			UUID:  template.ID,
+			Valid: true,
+		},
+		CreatedAt:      database.Now(),
+		OrganizationID: org.ID,
+		JobID:          job.ID,
+	})
+	db.InsertWorkspaceBuild(ctx, database.InsertWorkspaceBuildParams{
+		ID:                uuid.New(),
+		JobID:             job.ID,
+		WorkspaceID:       workspace.ID,
+		TemplateVersionID: version.ID,
+		Transition:        database.WorkspaceTransitionStart,
+	})
+	db.InsertWorkspaceResource(ctx, database.InsertWorkspaceResourceParams{
+		ID:    uuid.New(),
+		JobID: job.ID,
+		Type:  "google_compute_instance",
+		Name:  "banana2",
+	})
+	db.InsertWorkspaceResource(ctx, database.InsertWorkspaceResourceParams{
+		ID:    uuid.New(),
+		JobID: job.ID,
+		Type:  "google_compute_instance",
+		Name:  "banana3",
+	})
+	db.InsertWorkspace(ctx, database.InsertWorkspaceParams{
+		ID:             uuid.New(),
+		OwnerID:        user.ID,
+		OrganizationID: org.ID,
+		TemplateID:     template.ID,
+		Name:           "banana4",
+	})
+}