feat: Add workspace metrics export to Prometheus (#3421)

This adds workspace totals indexed by status. It could be any
codersdk.ProvisionerJobStatus.

coderd/prometheusmetrics/prometheusmetrics.go

@@ -7,6 +7,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
 
+	"github.com/coder/coder/coderd"
 	"github.com/coder/coder/coderd/database"
 )
 
@@ -50,3 +51,56 @@ func ActiveUsers(ctx context.Context, registerer prometheus.Registerer, db datab
 	}()
 	return cancelFunc, nil
 }
+
+// Workspaces tracks the total number of workspaces with labels on status.
+func Workspaces(ctx context.Context, registerer prometheus.Registerer, db database.Store, duration time.Duration) (context.CancelFunc, error) {
+	if duration == 0 {
+		duration = 5 * time.Minute
+	}
+
+	gauge := prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: "coderd",
+		Subsystem: "api",
+		Name:      "workspace_latest_build_total",
+		Help:      "The latest workspace builds with a status.",
+	}, []string{"status"})
+	err := registerer.Register(gauge)
+	if err != nil {
+		return nil, err
+	}
+	// This exists so the prometheus metric exports immediately when set.
+	// It helps with tests so they don't have to wait for a tick.
+	gauge.WithLabelValues("pending").Set(0)
+
+	ctx, cancelFunc := context.WithCancel(ctx)
+	ticker := time.NewTicker(duration)
+	go func() {
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+			}
+			builds, err := db.GetLatestWorkspaceBuilds(ctx)
+			if err != nil {
+				continue
+			}
+			jobIDs := make([]uuid.UUID, 0, len(builds))
+			for _, build := range builds {
+				jobIDs = append(jobIDs, build.JobID)
+			}
+			jobs, err := db.GetProvisionerJobsByIDs(ctx, jobIDs)
+			if err != nil {
+				continue
+			}
+
+			gauge.Reset()
+			for _, job := range jobs {
+				status := coderd.ConvertProvisionerJobStatus(job)
+				gauge.WithLabelValues(string(status)).Add(1)
+			}
+		}
+	}()
+	return cancelFunc, nil
+}

coderd/prometheusmetrics/prometheusmetrics_test.go

@@ -2,6 +2,7 @@ package prometheusmetrics_test
 
 import (
 	"context"
+	"database/sql"
 	"testing"
 	"time"
 
@@ -13,6 +14,7 @@ import (
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/database/databasefake"
 	"github.com/coder/coder/coderd/prometheusmetrics"
+	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/testutil"
 )
 
@@ -81,13 +83,148 @@ func TestActiveUsers(t *testing.T) {
 			require.NoError(t, err)
 			t.Cleanup(cancel)
 
-			var result int
 			require.Eventually(t, func() bool {
 				metrics, err := registry.Gather()
 				assert.NoError(t, err)
-				result = int(*metrics[0].Metric[0].Gauge.Value)
+				result := int(*metrics[0].Metric[0].Gauge.Value)
 				return result == tc.Count
 			}, testutil.WaitShort, testutil.IntervalFast)
 		})
 	}
 }
+
+func TestWorkspaces(t *testing.T) {
+	t.Parallel()
+
+	insertRunning := func(db database.Store) database.ProvisionerJob {
+		job, _ := db.InsertProvisionerJob(context.Background(), database.InsertProvisionerJobParams{
+			ID:          uuid.New(),
+			CreatedAt:   database.Now(),
+			UpdatedAt:   database.Now(),
+			Provisioner: database.ProvisionerTypeEcho,
+		})
+		_, _ = db.InsertWorkspaceBuild(context.Background(), database.InsertWorkspaceBuildParams{
+			ID:          uuid.New(),
+			WorkspaceID: uuid.New(),
+			JobID:       job.ID,
+			BuildNumber: 1,
+		})
+		// This marks the job as started.
+		_, _ = db.AcquireProvisionerJob(context.Background(), database.AcquireProvisionerJobParams{
+			StartedAt: sql.NullTime{
+				Time:  database.Now(),
+				Valid: true,
+			},
+			Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
+		})
+		return job
+	}
+
+	insertCanceled := func(db database.Store) {
+		job := insertRunning(db)
+		_ = db.UpdateProvisionerJobWithCancelByID(context.Background(), database.UpdateProvisionerJobWithCancelByIDParams{
+			ID: job.ID,
+			CanceledAt: sql.NullTime{
+				Time:  database.Now(),
+				Valid: true,
+			},
+		})
+		_ = db.UpdateProvisionerJobWithCompleteByID(context.Background(), database.UpdateProvisionerJobWithCompleteByIDParams{
+			ID: job.ID,
+			CompletedAt: sql.NullTime{
+				Time:  database.Now(),
+				Valid: true,
+			},
+		})
+	}
+
+	insertFailed := func(db database.Store) {
+		job := insertRunning(db)
+		_ = db.UpdateProvisionerJobWithCompleteByID(context.Background(), database.UpdateProvisionerJobWithCompleteByIDParams{
+			ID: job.ID,
+			CompletedAt: sql.NullTime{
+				Time:  database.Now(),
+				Valid: true,
+			},
+			Error: sql.NullString{
+				String: "failed",
+				Valid:  true,
+			},
+		})
+	}
+
+	insertSuccess := func(db database.Store) {
+		job := insertRunning(db)
+		_ = db.UpdateProvisionerJobWithCompleteByID(context.Background(), database.UpdateProvisionerJobWithCompleteByIDParams{
+			ID: job.ID,
+			CompletedAt: sql.NullTime{
+				Time:  database.Now(),
+				Valid: true,
+			},
+		})
+	}
+
+	for _, tc := range []struct {
+		Name     string
+		Database func() database.Store
+		Total    int
+		Status   map[codersdk.ProvisionerJobStatus]int
+	}{{
+		Name: "None",
+		Database: func() database.Store {
+			return databasefake.New()
+		},
+		Total: 0,
+	}, {
+		Name: "Multiple",
+		Database: func() database.Store {
+			db := databasefake.New()
+			insertCanceled(db)
+			insertFailed(db)
+			insertFailed(db)
+			insertSuccess(db)
+			insertSuccess(db)
+			insertSuccess(db)
+			insertRunning(db)
+			return db
+		},
+		Total: 7,
+		Status: map[codersdk.ProvisionerJobStatus]int{
+			codersdk.ProvisionerJobCanceled:  1,
+			codersdk.ProvisionerJobFailed:    2,
+			codersdk.ProvisionerJobSucceeded: 3,
+			codersdk.ProvisionerJobRunning:   1,
+		},
+	}} {
+		tc := tc
+		t.Run(tc.Name, func(t *testing.T) {
+			t.Parallel()
+			registry := prometheus.NewRegistry()
+			cancel, err := prometheusmetrics.Workspaces(context.Background(), registry, tc.Database(), time.Millisecond)
+			require.NoError(t, err)
+			t.Cleanup(cancel)
+
+			require.Eventually(t, func() bool {
+				metrics, err := registry.Gather()
+				assert.NoError(t, err)
+				if len(metrics) < 1 {
+					return false
+				}
+				sum := 0
+				for _, metric := range metrics[0].Metric {
+					count, ok := tc.Status[codersdk.ProvisionerJobStatus(metric.Label[0].GetValue())]
+					if metric.Gauge.GetValue() == 0 {
+						continue
+					}
+					if !ok {
+						t.Fail()
+					}
+					require.Equal(t, count, int(metric.Gauge.GetValue()), "invalid count for %s", metric.Label[0].GetValue())
+					sum += int(metric.Gauge.GetValue())
+				}
+				t.Logf("sum %d == total %d", sum, tc.Total)
+				return sum == tc.Total
+			}, testutil.WaitShort, testutil.IntervalFast)
+		})
+	}
+}