coder · mtojek · Oct 13, 2023 · Oct 10, 2023 · Oct 10, 2023 · Oct 10, 2023
diff --git a/cli/server.go b/cli/server.go
@@ -228,6 +228,7 @@ func enablePrometheus(
 	), nil
 }
 
+// nolint:gocognit // This function makes more sense being large for now, until refactored.
 func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.API, io.Closer, error)) *clibase.Cmd {
 	if newAPI == nil {
 		newAPI = func(_ context.Context, o *coderd.Options) (*coderd.API, io.Closer, error) {
@@ -814,6 +815,15 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 					return xerrors.Errorf("register agents prometheus metric: %w", err)
 				}
 				defer closeAgentsFunc()
+
+				if coderAPI.LicenseMetrics == nil {
+					return xerrors.Errorf("developer error: LicenseMetrics should be not nil")
+				}
+				closeLicenseMetricsFunc, err := coderAPI.LicenseMetrics.Collect(ctx)
+				if err != nil {
+					return xerrors.Errorf("register license metric: %w", err)
+				}
+				defer closeLicenseMetricsFunc()
 			}
 
 			client := codersdk.New(localURL)

diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -38,6 +38,7 @@ import (
 	// Used for swagger docs.
 	_ "github.com/coder/coder/v2/coderd/apidoc"
 	"github.com/coder/coder/v2/coderd/externalauth"
+	"github.com/coder/coder/v2/coderd/prometheusmetrics"
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/v2/buildinfo"
@@ -342,6 +343,15 @@ func New(options *Options) *API {
 	if err != nil {
 		panic(xerrors.Errorf("get deployment ID: %w", err))
 	}
+
+	licenseMetrics, err := prometheusmetrics.NewLicenseMetrics(&prometheusmetrics.LicenseMetricsOptions{
+		Logger:   options.Logger,
+		Registry: options.PrometheusRegistry,
+	})
+	if err != nil {
+		panic(xerrors.Errorf("unable to initialize license metrics: %w", err))
+	}
+
 	api := &API{
 		ctx:          ctx,
 		cancel:       cancel,
@@ -376,6 +386,8 @@ func New(options *Options) *API {
 			options.Logger.Named("acquirer"),
 			options.Database,
 			options.Pubsub),
+
+		LicenseMetrics: licenseMetrics,
 	}
 	if options.UpdateCheckOptions != nil {
 		api.updateChecker = updatecheck.New(
@@ -1035,6 +1047,8 @@ type API struct {
 	statsBatcher *batchstats.Batcher
 
 	Acquirer *provisionerdserver.Acquirer
+
+	LicenseMetrics *prometheusmetrics.LicenseMetrics
 }
 
 // Close waits for all WebSocket connections to drain before returning.

diff --git a/coderd/prometheusmetrics/license.go b/coderd/prometheusmetrics/license.go
@@ -0,0 +1,125 @@
+package prometheusmetrics
+
+import (
+	"context"
+	"fmt"
+	"sync/atomic"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+
+	"cdr.dev/slog"
+
+	"github.com/coder/coder/v2/codersdk"
+)
+
+type LicenseMetrics struct {
+	interval time.Duration
+	logger   slog.Logger
+	registry *prometheus.Registry
+
+	Entitlements atomic.Pointer[codersdk.Entitlements]
+}
+
+type LicenseMetricsOptions struct {
+	Interval time.Duration
+	Logger   slog.Logger
+	Registry *prometheus.Registry
+}
+
+func NewLicenseMetrics(opts *LicenseMetricsOptions) (*LicenseMetrics, error) {
+	if opts.Interval == 0 {
+		opts.Interval = 1 * time.Minute
+	}
+	if opts.Registry == nil {
+		opts.Registry = prometheus.NewRegistry()
+	}
+
+	return &LicenseMetrics{
+		interval: opts.Interval,
+		logger:   opts.Logger,
+		registry: opts.Registry,
+	}, nil
+}
+
+func (lm *LicenseMetrics) Collect(ctx context.Context) (func(), error) {
+	activeUsersGauge := NewCachedGaugeVec(prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: "coderd",
+		Subsystem: "license",
+		Name:      "active_users",
+		Help:      `The number of active users.`,
+	}, []string{"entitled"}))
+	err := lm.registry.Register(activeUsersGauge)
+	if err != nil {
+		return nil, err
+	}
+
+	userLimitGauge := NewCachedGaugeVec(prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: "coderd",
+		Subsystem: "license",
+		Name:      "user_limit",
+		Help:      "The user seats limit based on the active Coder license.",
+	}, []string{"entitled"}))
+	err = lm.registry.Register(userLimitGauge)
+	if err != nil {
+		return nil, err
+	}
+
+	ctx, cancelFunc := context.WithCancel(ctx)
+	done := make(chan struct{})
+	ticker := time.NewTicker(time.Nanosecond)
+
+	doTick := func() {
+		defer ticker.Reset(lm.interval)
+
+		entitlements := lm.Entitlements.Load()
+		if entitlements == nil {
+			lm.logger.Warn(ctx, `entitlements have not been loaded yet`)
+			return
+		}
+
+		if entitlements.Features == nil {
+			lm.logger.Warn(ctx, `entitlements features are undefined`)
+			return
+		}
+
+		userLimitEntitlement, ok := entitlements.Features[codersdk.FeatureUserLimit]
+		if !ok {
+			lm.logger.Warn(ctx, `"user_limit" entitlement is not present`)
+			return
+		}
+
+		enabled := fmt.Sprintf("%v", userLimitEntitlement.Enabled)
+		if userLimitEntitlement.Actual != nil {
+			activeUsersGauge.WithLabelValues(VectorOperationSet, float64(*userLimitEntitlement.Actual), enabled)
+		} else {
+			activeUsersGauge.WithLabelValues(VectorOperationSet, 0, enabled)
+		}
+
+		if userLimitEntitlement.Limit != nil {
+			userLimitGauge.WithLabelValues(VectorOperationSet, float64(*userLimitEntitlement.Limit), enabled)
+		} else {
+			userLimitGauge.WithLabelValues(VectorOperationSet, 0, enabled)
+		}
+
+		activeUsersGauge.Commit()
+		userLimitGauge.Commit()
+	}
+
+	go func() {
+		defer close(done)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				doTick()
+			}
+		}
+	}()
+	return func() {
+		cancelFunc()
+		<-done
+	}, nil
+}
diff --git a/coderd/prometheusmetrics/license_test.go b/coderd/prometheusmetrics/license_test.go
@@ -0,0 +1,89 @@
+package prometheusmetrics_test
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+	"reflect"
+	"testing"
+	"time"
+
+	"github.com/aws/smithy-go/ptr"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/slogtest"
+
+	"github.com/coder/coder/v2/coderd/prometheusmetrics"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/testutil"
+)
+
+func TestCollectLicenseMetrics(t *testing.T) {
+	t.Parallel()
+
+	// Given
+	registry := prometheus.NewRegistry()
+	sut, err := prometheusmetrics.NewLicenseMetrics(&prometheusmetrics.LicenseMetricsOptions{
+		Interval: time.Millisecond,
+		Logger:   slogtest.Make(t, nil).Leveled(slog.LevelDebug),
+		Registry: registry,
+	})
+	require.NoError(t, err)
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	t.Cleanup(cancelFunc)
+
+	const (
+		actualUsers = 4
+		userLimit   = 7
+	)
+	sut.Entitlements.Store(&codersdk.Entitlements{
+		Features: map[codersdk.FeatureName]codersdk.Feature{
+			codersdk.FeatureUserLimit: {
+				Enabled: true,
+				Actual:  ptr.Int64(actualUsers),
+				Limit:   ptr.Int64(userLimit),
+			},
+		},
+	})
+
+	// When
+	closeFunc, err := sut.Collect(ctx)
+	require.NoError(t, err)
+	t.Cleanup(closeFunc)
+
+	// Then
+	goldenFile, err := os.ReadFile("testdata/license-metrics.json")
+	require.NoError(t, err)
+	golden := map[string]int{}
+	err = json.Unmarshal(goldenFile, &golden)
+	require.NoError(t, err)
+
+	collected := map[string]int{}
+
+	assert.Eventually(t, func() bool {
+		metrics, err := registry.Gather()
+		assert.NoError(t, err)
+
+		if len(metrics) < 1 {
+			return false
+		}
+
+		for _, metric := range metrics {
+			switch metric.GetName() {
+			case "coderd_license_active_users", "coderd_license_user_limit":
+				for _, m := range metric.Metric {
+					collected[m.Label[0].GetName()+"="+m.Label[0].GetValue()+":"+metric.GetName()] = int(m.Gauge.GetValue())
+				}
+			default:
+				require.FailNowf(t, "unexpected metric collected", "metric: %s", metric.GetName())
+			}
+		}
+		return reflect.DeepEqual(golden, collected)
+	}, testutil.WaitShort, testutil.IntervalFast)
+
+	assert.EqualValues(t, golden, collected)
+}
diff --git a/coderd/prometheusmetrics/testdata/license-metrics.json b/coderd/prometheusmetrics/testdata/license-metrics.json
@@ -0,0 +1,4 @@
+{
+  "entitled=true:coderd_license_active_users": 4,
+  "entitled=true:coderd_license_user_limit": 7
+}
diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
@@ -661,6 +661,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 	defer api.entitlementsMu.Unlock()
 	api.entitlements = entitlements
 	api.AGPL.SiteHandler.Entitlements.Store(&entitlements)
+	api.AGPL.LicenseMetrics.Entitlements.Store(&entitlements)
 
 	return nil
 }