coder · mtojek · Apr 7, 2023 · Apr 3, 2023 · Apr 4, 2023 · Apr 4, 2023
diff --git a/cli/server.go b/cli/server.go
@@ -889,6 +889,15 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 				return xerrors.Errorf("create coder API: %w", err)
 			}
 
+			if cfg.Prometheus.Enable {
+				// Agent metrics require reference to the tailnet coordinator, so must be initiated after Coder API.
+				closeAgentsFunc, err := prometheusmetrics.Agents(ctx, logger, options.PrometheusRegistry, coderAPI.Database, &coderAPI.TailnetCoordinator, options.DERPMap, coderAPI.Options.AgentInactiveDisconnectTimeout, 0)
+				if err != nil {
+					return xerrors.Errorf("register agents prometheus metric: %w", err)
+				}
+				defer closeAgentsFunc()
+			}
+
 			client := codersdk.New(localURL)
 			if localURL.Scheme == "https" && isLocalhost(localURL.Hostname()) {
 				// The certificate will likely be self-signed or for a different

diff --git a/coderd/prometheusmetrics/collector.go b/coderd/prometheusmetrics/collector.go
@@ -0,0 +1,95 @@
+package prometheusmetrics
+
+import (
+	"sync"
+
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+// CachedGaugeVec is a wrapper for the prometheus.GaugeVec which allows
+// for staging changes in the metrics vector. Calling "WithLabelValues(...)"
+// will update the internal gauge value, but it will not be returned by
+// "Collect(...)" until the "Commit()" method is called. The "Commit()" method
+// resets the internal gauge and applies all staged changes to it.
+//
+// The Use of CachedGaugeVec is recommended for use cases when there is a risk
+// that the Prometheus collector receives incomplete metrics, collected
+// in the middle of metrics recalculation, between "Reset()" and the last
+// "WithLabelValues()" call.
+type CachedGaugeVec struct {
+	m sync.Mutex
+
+	gaugeVec *prometheus.GaugeVec
+	records  []vectorRecord
+}
+
+var _ prometheus.Collector = new(CachedGaugeVec)
+
+type VectorOperation int
+
+const (
+	VectorOperationAdd VectorOperation = iota
+	VectorOperationSet
+)
+
+type vectorRecord struct {
+	operation   VectorOperation
+	value       float64
+	labelValues []string
+}
+
+func NewCachedGaugeVec(gaugeVec *prometheus.GaugeVec) *CachedGaugeVec {
+	return &CachedGaugeVec{
+		gaugeVec: gaugeVec,
+	}
+}
+
+func (v *CachedGaugeVec) Describe(desc chan<- *prometheus.Desc) {
+	v.gaugeVec.Describe(desc)
+}
+
+func (v *CachedGaugeVec) Collect(ch chan<- prometheus.Metric) {
+	v.m.Lock()
+	defer v.m.Unlock()
+
+	v.gaugeVec.Collect(ch)
+}
+
+func (v *CachedGaugeVec) WithLabelValues(operation VectorOperation, value float64, labelValues ...string) {
+	switch operation {
+	case VectorOperationAdd:
+	case VectorOperationSet:
+	default:
+		panic("unsupported vector operation")
+	}
+
+	v.m.Lock()
+	defer v.m.Unlock()
+
+	v.records = append(v.records, vectorRecord{
+		operation:   operation,
+		value:       value,
+		labelValues: labelValues,
+	})
+}
+
+// Commit will set the internal value as the cached value to return from "Collect()".
+// The internal metric value is completely reset, so the caller should expect
+// the gauge to be empty for the next 'WithLabelValues' values.
+func (v *CachedGaugeVec) Commit() {
-func (v *CachedGaugeVec) Commit() {
+func (v *CachedGaugeVec) Commit() {
-func (v *CachedGaugeVec) Commit() {
+func (v *CachedGaugeVec) Commit() {
+	v.m.Lock()
+	defer v.m.Unlock()
+
+	v.gaugeVec.Reset()
+	for _, record := range v.records {
+		g := v.gaugeVec.WithLabelValues(record.labelValues...)
+		switch record.operation {
+		case VectorOperationAdd:
+			g.Add(record.value)
+		case VectorOperationSet:
+			g.Set(record.value)
+		}
+	}
+
+	v.records = nil
+}
diff --git a/coderd/prometheusmetrics/collector_test.go b/coderd/prometheusmetrics/collector_test.go
@@ -0,0 +1,140 @@
+package prometheusmetrics_test
+
+import (
+	"sort"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+	dto "github.com/prometheus/client_model/go"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/coderd/prometheusmetrics"
+)
+
+func TestCollector_Add(t *testing.T) {
+	t.Parallel()
+
+	// given
+	agentsGauge := prometheusmetrics.NewCachedGaugeVec(prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: "coderd",
+		Subsystem: "agents",
+		Name:      "up",
+		Help:      "The number of active agents per workspace.",
+	}, []string{"username", "workspace_name"}))
+
+	// when
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationAdd, 7, "first user", "my workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationAdd, 23, "second user", "your workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationAdd, 1, "first user", "my workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationAdd, 25, "second user", "your workspace")
+	agentsGauge.Commit()
+
+	// then
+	ch := make(chan prometheus.Metric, 2)
+	agentsGauge.Collect(ch)
+
+	metrics := collectAndSortMetrics(t, agentsGauge, 2)
+
+	assert.Equal(t, "first user", metrics[0].Label[0].GetValue())   // Username
+	assert.Equal(t, "my workspace", metrics[0].Label[1].GetValue()) // Workspace name
+	assert.Equal(t, 8, int(metrics[0].Gauge.GetValue()))            // Metric value
+
+	assert.Equal(t, "second user", metrics[1].Label[0].GetValue())    // Username
+	assert.Equal(t, "your workspace", metrics[1].Label[1].GetValue()) // Workspace name
+	assert.Equal(t, 48, int(metrics[1].Gauge.GetValue()))             // Metric value
+}
+
+func TestCollector_Set(t *testing.T) {
+	t.Parallel()
+
+	// given
+	agentsGauge := prometheusmetrics.NewCachedGaugeVec(prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: "coderd",
+		Subsystem: "agents",
+		Name:      "up",
+		Help:      "The number of active agents per workspace.",
+	}, []string{"username", "workspace_name"}))
+
+	// when
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationSet, 3, "first user", "my workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationSet, 4, "second user", "your workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationSet, 5, "first user", "my workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationSet, 6, "second user", "your workspace")
+	agentsGauge.Commit()
+
+	// then
+	ch := make(chan prometheus.Metric, 2)
+	agentsGauge.Collect(ch)
+
+	metrics := collectAndSortMetrics(t, agentsGauge, 2)
+
+	assert.Equal(t, "first user", metrics[0].Label[0].GetValue())   // Username
+	assert.Equal(t, "my workspace", metrics[0].Label[1].GetValue()) // Workspace name
+	assert.Equal(t, 5, int(metrics[0].Gauge.GetValue()))            // Metric value
+
+	assert.Equal(t, "second user", metrics[1].Label[0].GetValue())    // Username
+	assert.Equal(t, "your workspace", metrics[1].Label[1].GetValue()) // Workspace name
+	assert.Equal(t, 6, int(metrics[1].Gauge.GetValue()))              // Metric value
+}
+
+func TestCollector_Set_Add(t *testing.T) {
+	t.Parallel()
+
+	// given
+	agentsGauge := prometheusmetrics.NewCachedGaugeVec(prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: "coderd",
+		Subsystem: "agents",
+		Name:      "up",
+		Help:      "The number of active agents per workspace.",
+	}, []string{"username", "workspace_name"}))
+
+	// when
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationAdd, 9, "first user", "my workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationAdd, 8, "second user", "your workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationAdd, 7, "first user", "my workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationAdd, 6, "second user", "your workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationSet, 5, "first user", "my workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationSet, 4, "second user", "your workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationAdd, 3, "first user", "my workspace")
+	agentsGauge.WithLabelValues(prometheusmetrics.VectorOperationAdd, 2, "second user", "your workspace")
+	agentsGauge.Commit()
+
+	// then
+	ch := make(chan prometheus.Metric, 2)
+	agentsGauge.Collect(ch)
+
+	metrics := collectAndSortMetrics(t, agentsGauge, 2)
+
+	assert.Equal(t, "first user", metrics[0].Label[0].GetValue())   // Username
+	assert.Equal(t, "my workspace", metrics[0].Label[1].GetValue()) // Workspace name
+	assert.Equal(t, 8, int(metrics[0].Gauge.GetValue()))            // Metric value
+
+	assert.Equal(t, "second user", metrics[1].Label[0].GetValue())    // Username
+	assert.Equal(t, "your workspace", metrics[1].Label[1].GetValue()) // Workspace name
+	assert.Equal(t, 6, int(metrics[1].Gauge.GetValue()))              // Metric value
+}
+
+func collectAndSortMetrics(t *testing.T, collector prometheus.Collector, count int) []dto.Metric {
+	ch := make(chan prometheus.Metric, count)
+	defer close(ch)
+
+	var metrics []dto.Metric
+
+	collector.Collect(ch)
+	for i := 0; i < count; i++ {
+		m := <-ch
+
+		var metric dto.Metric
+		err := m.Write(&metric)
+		require.NoError(t, err)
+
+		metrics = append(metrics, metric)
+	}
+
+	// Ensure always the same order of metrics
+	sort.Slice(metrics, func(i, j int) bool {
+		return sort.StringsAreSorted([]string{metrics[i].Label[0].GetValue(), metrics[j].Label[1].GetValue()})
+	})
+	return metrics
+}