fix: serialize updateEntitlements() (#14974)

fixes #14961

Adding the license and updating entitlements is flaky, especially at the start of our `coderdent` testing because, while the actual modifications to the `entitlements.Set` were threadsafe, we could have multiple goroutines reading from the database and writing to the set, so we could end up writing stale data.

This enforces serialization on updates, so that if you modify the database and kick off an update, you know the state of the `Set` is at least as fresh as your database update.

coderd/entitlements/entitlements.go

@@ -1,21 +1,31 @@
 package entitlements
 
 import (
+	"context"
 	"encoding/json"
 	"net/http"
 	"sync"
 	"time"
 
+	"golang.org/x/exp/slices"
+	"golang.org/x/xerrors"
+
 	"github.com/coder/coder/v2/codersdk"
 )
 
 type Set struct {
 	entitlementsMu sync.RWMutex
 	entitlements   codersdk.Entitlements
+	// right2Update works like a semaphore. Reading from the chan gives the right to update the set,
+	// and you send on the chan when you are done. We only allow one simultaneous update, so this
+	// serve to serialize them.  You MUST NOT attempt to read from this channel while holding the
+	// entitlementsMu lock. It is permissible to acquire the entitlementsMu lock while holding the
+	// right2Update token.
+	right2Update chan struct{}
 }
 
 func New() *Set {
-	return &Set{
+	s := &Set{
 		// Some defaults for an unlicensed instance.
 		// These will be updated when coderd is initialized.
 		entitlements: codersdk.Entitlements{
@@ -27,7 +37,44 @@ func New() *Set {
 			RequireTelemetry: false,
 			RefreshedAt:      time.Time{},
 		},
+		right2Update: make(chan struct{}, 1),
 	}
+	s.right2Update <- struct{}{} // one token, serialized updates
+	return s
+}
+
+// ErrLicenseRequiresTelemetry is an error returned by a fetch passed to Update to indicate that the
+// fetched license cannot be used because it requires telemetry.
+var ErrLicenseRequiresTelemetry = xerrors.New("License requires telemetry but telemetry is disabled")
+
+func (l *Set) Update(ctx context.Context, fetch func(context.Context) (codersdk.Entitlements, error)) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-l.right2Update:
+		defer func() {
+			l.right2Update <- struct{}{}
+		}()
+	}
+	ents, err := fetch(ctx)
+	if xerrors.Is(err, ErrLicenseRequiresTelemetry) {
+		// We can't fail because then the user couldn't remove the offending
+		// license w/o a restart.
+		//
+		// We don't simply append to entitlement.Errors since we don't want any
+		// enterprise features enabled.
+		l.Modify(func(entitlements *codersdk.Entitlements) {
+			entitlements.Errors = []string{err.Error()}
+		})
+		return nil
+	}
+	if err != nil {
+		return err
+	}
+	l.entitlementsMu.Lock()
+	defer l.entitlementsMu.Unlock()
+	l.entitlements = ents
+	return nil
 }
 
 // AllowRefresh returns whether the entitlements are allowed to be refreshed.
@@ -74,14 +121,7 @@ func (l *Set) AsJSON() json.RawMessage {
 	return b
 }
 
-func (l *Set) Replace(entitlements codersdk.Entitlements) {
+func (l *Set) Modify(do func(entitlements *codersdk.Entitlements)) {
-	l.entitlementsMu.Lock()
-	defer l.entitlementsMu.Unlock()
-
-	l.entitlements = entitlements
-}
-
-func (l *Set) Update(do func(entitlements *codersdk.Entitlements)) {
 	l.entitlementsMu.Lock()
 	defer l.entitlementsMu.Unlock()
 
@@ -107,3 +147,9 @@ func (l *Set) WriteEntitlementWarningHeaders(header http.Header) {
 		header.Add(codersdk.EntitlementsWarningHeader, warning)
 	}
 }
+
+func (l *Set) Errors() []string {
+	l.entitlementsMu.RLock()
+	defer l.entitlementsMu.RUnlock()
+	return slices.Clone(l.entitlements.Errors)
+}

coderd/entitlements/entitlements_test.go

@@ -1,6 +1,7 @@
 package entitlements_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -8,15 +9,16 @@ import (
 
 	"github.com/coder/coder/v2/coderd/entitlements"
 	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/testutil"
 )
 
-func TestUpdate(t *testing.T) {
+func TestModify(t *testing.T) {
 	t.Parallel()
 
 	set := entitlements.New()
 	require.False(t, set.Enabled(codersdk.FeatureMultipleOrganizations))
 
-	set.Update(func(entitlements *codersdk.Entitlements) {
+	set.Modify(func(entitlements *codersdk.Entitlements) {
 		entitlements.Features[codersdk.FeatureMultipleOrganizations] = codersdk.Feature{
 			Enabled:     true,
 			Entitlement: codersdk.EntitlementEntitled,
@@ -30,7 +32,7 @@ func TestAllowRefresh(t *testing.T) {
 
 	now := time.Now()
 	set := entitlements.New()
-	set.Update(func(entitlements *codersdk.Entitlements) {
+	set.Modify(func(entitlements *codersdk.Entitlements) {
 		entitlements.RefreshedAt = now
 	})
 
@@ -38,7 +40,7 @@ func TestAllowRefresh(t *testing.T) {
 	require.False(t, ok)
 	require.InDelta(t, time.Minute.Seconds(), wait.Seconds(), 5)
 
-	set.Update(func(entitlements *codersdk.Entitlements) {
+	set.Modify(func(entitlements *codersdk.Entitlements) {
 		entitlements.RefreshedAt = now.Add(time.Minute * -2)
 	})
 
@@ -47,17 +49,76 @@ func TestAllowRefresh(t *testing.T) {
 	require.Equal(t, time.Duration(0), wait)
 }
 
-func TestReplace(t *testing.T) {
+func TestUpdate(t *testing.T) {
 	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
 
 	set := entitlements.New()
 	require.False(t, set.Enabled(codersdk.FeatureMultipleOrganizations))
-	set.Replace(codersdk.Entitlements{
+	fetchStarted := make(chan struct{})
+	firstDone := make(chan struct{})
+	errCh := make(chan error, 2)
+	go func() {
+		err := set.Update(ctx, func(_ context.Context) (codersdk.Entitlements, error) {
+			close(fetchStarted)
+			select {
+			case <-firstDone:
+				// OK!
+			case <-ctx.Done():
+				t.Error("timeout")
+				return codersdk.Entitlements{}, ctx.Err()
+			}
+			return codersdk.Entitlements{
 				Features: map[codersdk.FeatureName]codersdk.Feature{
 					codersdk.FeatureMultipleOrganizations: {
 						Enabled: true,
 					},
 				},
+			}, nil
 		})
+		errCh <- err
+	}()
+	testutil.RequireRecvCtx(ctx, t, fetchStarted)
+	require.False(t, set.Enabled(codersdk.FeatureMultipleOrganizations))
+	// start a second update while the first one is in progress
+	go func() {
+		err := set.Update(ctx, func(_ context.Context) (codersdk.Entitlements, error) {
+			return codersdk.Entitlements{
+				Features: map[codersdk.FeatureName]codersdk.Feature{
+					codersdk.FeatureMultipleOrganizations: {
+						Enabled: true,
+					},
+					codersdk.FeatureAppearance: {
+						Enabled: true,
+					},
+				},
+			}, nil
+		})
+		errCh <- err
+	}()
+	close(firstDone)
+	err := testutil.RequireRecvCtx(ctx, t, errCh)
+	require.NoError(t, err)
+	err = testutil.RequireRecvCtx(ctx, t, errCh)
+	require.NoError(t, err)
 	require.True(t, set.Enabled(codersdk.FeatureMultipleOrganizations))
+	require.True(t, set.Enabled(codersdk.FeatureAppearance))
+}
+
+func TestUpdate_LicenseRequiresTelemetry(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+	set := entitlements.New()
+	set.Modify(func(entitlements *codersdk.Entitlements) {
+		entitlements.Errors = []string{"some error"}
+		entitlements.Features[codersdk.FeatureAppearance] = codersdk.Feature{
+			Enabled: true,
+		}
+	})
+	err := set.Update(ctx, func(_ context.Context) (codersdk.Entitlements, error) {
+		return codersdk.Entitlements{}, entitlements.ErrLicenseRequiresTelemetry
+	})
+	require.NoError(t, err)
+	require.True(t, set.Enabled(codersdk.FeatureAppearance))
+	require.Equal(t, []string{entitlements.ErrLicenseRequiresTelemetry.Error()}, set.Errors())
 }

enterprise/coderd/coderd.go

@@ -613,6 +613,7 @@ func (api *API) Close() error {
 }
 
 func (api *API) updateEntitlements(ctx context.Context) error {
+	return api.Entitlements.Update(ctx, func(ctx context.Context) (codersdk.Entitlements, error) {
 		replicas := api.replicaManager.AllPrimary()
 		agedReplicas := make([]database.Replica, 0, len(replicas))
 		for _, replica := range replicas {
@@ -645,23 +646,12 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 				codersdk.FeatureControlSharedPorts:         true,
 			})
 		if err != nil {
-		return err
+			return codersdk.Entitlements{}, err
 		}
 
 		if reloadedEntitlements.RequireTelemetry && !api.DeploymentValues.Telemetry.Enable.Value() {
-		// We can't fail because then the user couldn't remove the offending
-		// license w/o a restart.
-		//
-		// We don't simply append to entitlement.Errors since we don't want any
-		// enterprise features enabled.
-		api.Entitlements.Update(func(entitlements *codersdk.Entitlements) {
-			entitlements.Errors = []string{
-				"License requires telemetry but telemetry is disabled",
-			}
-		})
-
 			api.Logger.Error(ctx, "license requires telemetry enabled")
-		return nil
+			return codersdk.Entitlements{}, entitlements.ErrLicenseRequiresTelemetry
 		}
 
 		featureChanged := func(featureName codersdk.FeatureName) (initial, changed, enabled bool) {
@@ -830,9 +820,8 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 			reloadedEntitlements.Warnings = append(reloadedEntitlements.Warnings, msg)
 		}
 		reloadedEntitlements.Features[codersdk.FeatureExternalTokenEncryption] = featureExternalTokenEncryption
-
+		return reloadedEntitlements, nil
-	api.Entitlements.Replace(reloadedEntitlements)
+	})
-	return nil
 }
 
 // getProxyDERPStartingRegionID returns the starting region ID that should be

enterprise/coderd/enidpsync/groups_test.go

@@ -19,7 +19,7 @@ func TestEnterpriseParseGroupClaims(t *testing.T) {
 	t.Parallel()
 
 	entitled := entitlements.New()
-	entitled.Update(func(entitlements *codersdk.Entitlements) {
+	entitled.Modify(func(entitlements *codersdk.Entitlements) {
 		entitlements.Features[codersdk.FeatureTemplateRBAC] = codersdk.Feature{
 			Entitlement: codersdk.EntitlementEntitled,
 			Enabled:     true,

enterprise/coderd/enidpsync/organizations_test.go

@@ -70,7 +70,7 @@ func TestOrganizationSync(t *testing.T) {
 	}
 
 	entitled := entitlements.New()
-	entitled.Update(func(entitlements *codersdk.Entitlements) {
+	entitled.Modify(func(entitlements *codersdk.Entitlements) {
 		entitlements.Features[codersdk.FeatureMultipleOrganizations] = codersdk.Feature{
 			Entitlement: codersdk.EntitlementEntitled,
 			Enabled:     true,

enterprise/coderd/enidpsync/role_test.go

@@ -20,7 +20,7 @@ func TestEnterpriseParseRoleClaims(t *testing.T) {
 	t.Parallel()
 
 	entitled := entitlements.New()
-	entitled.Update(func(en *codersdk.Entitlements) {
+	entitled.Modify(func(en *codersdk.Entitlements) {
 		en.Features[codersdk.FeatureUserRoleManagement] = codersdk.Feature{
 			Entitlement: codersdk.EntitlementEntitled,
 			Enabled:     true,

enterprise/coderd/license/metricscollector_test.go

@@ -27,7 +27,7 @@ func TestCollectLicenseMetrics(t *testing.T) {
 		userLimit   = 7
 	)
 	sut.Entitlements = entitlements.New()
-	sut.Entitlements.Update(func(entitlements *codersdk.Entitlements) {
+	sut.Entitlements.Modify(func(entitlements *codersdk.Entitlements) {
 		entitlements.Features[codersdk.FeatureUserLimit] = codersdk.Feature{
 			Enabled: true,
 			Actual:  ptr.Int64(actualUsers),