WIP

coder · mtojek · Oct 13, 2023 · Oct 10, 2023 · Oct 10, 2023 · Oct 10, 2023
commit 404896564e91e6cc8425d2af04668ee1eb4bae49
diff --git a/cli/server.go b/cli/server.go
@@ -814,6 +814,12 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 					return xerrors.Errorf("register agents prometheus metric: %w", err)
 				}
 				defer closeAgentsFunc()
+
+				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
@@ -343,6 +343,16 @@ func New(options *Options) *API {
 	if err != nil {
 		panic(xerrors.Errorf("get deployment ID: %w", err))
 	}
+
+	licenseMetrics, err := prometheusmetrics.NewLicenseMetrics(&prometheusmetrics.LicenseMetricsOptions{
+		Database: options.Database,
+		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,
@@ -377,6 +387,8 @@ func New(options *Options) *API {
 			options.Logger.Named("acquirer"),
 			options.Database,
 			options.Pubsub),
+
+		LicenseMetrics: licenseMetrics,
 	}
 	if options.UpdateCheckOptions != nil {
 		api.updateChecker = updatecheck.New(

diff --git a/coderd/prometheusmetrics/license.go b/coderd/prometheusmetrics/license.go
@@ -5,13 +5,12 @@ import (
 	"sync/atomic"
 	"time"
 
-	"cdr.dev/slog"
-	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/xerrors"
 
+	"cdr.dev/slog"
+
 	"github.com/coder/coder/v2/coderd/database"
-	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/codersdk"
 )
 
@@ -51,14 +50,13 @@ func NewLicenseMetrics(opts *LicenseMetricsOptions) (*LicenseMetrics, error) {
 }
 
 func (lm *LicenseMetrics) Collect(ctx context.Context) (func(), error) {
-
 	licenseLimitGauge := prometheus.NewGauge(prometheus.GaugeOpts{
 		Namespace: "coderd",
 		Subsystem: "license",
 		Name:      "user_limit",
 		Help:      `The user seats limit based on the current license. "Zero" means unlimited or a disabled feature.`,
 	})
-	err := registerer.Register(licenseLimitGauge)
+	err := lm.registry.Register(licenseLimitGauge)
 	if err != nil {
 		return nil, err
 	}
@@ -69,14 +67,25 @@ func (lm *LicenseMetrics) Collect(ctx context.Context) (func(), error) {
 		Name:      "active_users",
 		Help:      "The number of active users.",
 	})
-	err = registerer.Register(activeUsersGauge)
+	err = lm.registry.Register(activeUsersGauge)
+	if err != nil {
+		return nil, err
+	}
+
+	userLimitGauge := prometheus.NewGauge(prometheus.GaugeOpts{
+		Namespace: "coderd",
+		Subsystem: "license",
+		Name:      "user_limit",
+		Help:      "The user seats limit based on the active Coder license.",
+	})
+	err = lm.registry.Register(activeUsersGauge)
 	if err != nil {
 		return nil, err
 	}
 
 	ctx, cancelFunc := context.WithCancel(ctx)
 	done := make(chan struct{})
-	ticker := time.NewTicker(duration)
+	ticker := time.NewTicker(lm.interval)
 	go func() {
 		defer close(done)
 		defer ticker.Stop()
@@ -87,15 +96,6 @@ func (lm *LicenseMetrics) Collect(ctx context.Context) (func(), error) {
 			case <-ticker.C:
 			}
 
-			apiKeys, err := db.GetAPIKeysLastUsedAfter(ctx, dbtime.Now().Add(-1*time.Hour))
-			if err != nil {
-				continue
-			}
-			distinctUsers := map[uuid.UUID]struct{}{}
-			for _, apiKey := range apiKeys {
-				distinctUsers[apiKey.UserID] = struct{}{}
-			}
-			gauge.Set(float64(len(distinctUsers)))
 		}
 	}()
 	return func() {

diff --git a/coderd/prometheusmetrics/prometheusmetrics.go b/coderd/prometheusmetrics/prometheusmetrics.go
@@ -502,61 +502,3 @@ func AgentStats(ctx context.Context, logger slog.Logger, registerer prometheus.R
 		<-done
 	}, nil
 }
-
-// License function tracks the number of occupied license seats, based on the entitlements.
-func License(ctx context.Context, registerer prometheus.Registerer, db database.Store, entitlementsFetcher func() codersdk.Entitlements, duration time.Duration) (func(), error) {
-	if duration == 0 {
-		duration = 1 * time.Minute
-	}
-
-	licenseLimitGauge := prometheus.NewGauge(prometheus.GaugeOpts{
-		Namespace: "coderd",
-		Subsystem: "license",
-		Name:      "user_limit",
-		Help:      `The user seats limit based on the current license. "Zero" means unlimited or a disabled feature.`,
-	})
-	err := registerer.Register(licenseLimitGauge)
-	if err != nil {
-		return nil, err
-	}
-
-	activeUsersGauge := prometheus.NewGauge(prometheus.GaugeOpts{
-		Namespace: "coderd",
-		Subsystem: "license",
-		Name:      "active_users",
-		Help:      "The number of active users.",
-	})
-	err = registerer.Register(activeUsersGauge)
-	if err != nil {
-		return nil, err
-	}
-
-	ctx, cancelFunc := context.WithCancel(ctx)
-	done := make(chan struct{})
-	ticker := time.NewTicker(duration)
-	go func() {
-		defer close(done)
-		defer ticker.Stop()
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			case <-ticker.C:
-			}
-
-			apiKeys, err := db.GetAPIKeysLastUsedAfter(ctx, dbtime.Now().Add(-1*time.Hour))
-			if err != nil {
-				continue
-			}
-			distinctUsers := map[uuid.UUID]struct{}{}
-			for _, apiKey := range apiKeys {
-				distinctUsers[apiKey.UserID] = struct{}{}
-			}
-			gauge.Set(float64(len(distinctUsers)))
-		}
-	}()
-	return func() {
-		cancelFunc()
-		<-done
-	}, nil
-}
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
 }