fix: add grace period before showing replicas license error (#12989)

kylecarbs · web-flow · commit 227e63205398 · 2024-04-17T13:30:32.000-04:00
Fixes #8665.
diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"github.com/coder/coder/v2/coderd/appearance"
+	"github.com/coder/coder/v2/coderd/database"
 	agplportsharing "github.com/coder/coder/v2/coderd/portsharing"
 	"github.com/coder/coder/v2/enterprise/coderd/portsharing"
 
@@ -27,6 +28,7 @@ import (
 	"github.com/coder/coder/v2/coderd"
 	agplaudit "github.com/coder/coder/v2/coderd/audit"
 	agpldbauthz "github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/healthcheck"
 	"github.com/coder/coder/v2/coderd/httpapi"
 	"github.com/coder/coder/v2/coderd/httpmw"
@@ -64,6 +66,11 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
 	if options.Options.Authorizer == nil {
 		options.Options.Authorizer = rbac.NewCachingAuthorizer(options.PrometheusRegistry)
 	}
+	if options.ReplicaErrorGracePeriod == 0 {
+		// This will prevent the error from being shown for a minute
+		// from when an additional replica was started.
+		options.ReplicaErrorGracePeriod = time.Minute
+	}
 
 	ctx, cancelFunc := context.WithCancel(ctx)
 
@@ -429,6 +436,7 @@ type Options struct {
 
 	// Used for high availability.
 	ReplicaSyncUpdateInterval time.Duration
+	ReplicaErrorGracePeriod   time.Duration
 	DERPServerRelayAddress    string
 	DERPServerRegionID        int
 
@@ -525,9 +533,24 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 	api.entitlementsUpdateMu.Lock()
 	defer api.entitlementsUpdateMu.Unlock()
 
+	replicas := api.replicaManager.AllPrimary()
+	agedReplicas := make([]database.Replica, 0, len(replicas))
+	for _, replica := range replicas {
+		// If a replica is less than the update interval old, we don't
+		// want to display a warning. In the open-source version of Coder,
+		// Kubernetes Pods will start up before shutting down the other,
+		// and we don't want to display a warning in that case.
+		//
+		// Only display warnings for long-lived replicas!
+		if dbtime.Now().Sub(replica.StartedAt) < api.ReplicaErrorGracePeriod {
+			continue
+		}
+		agedReplicas = append(agedReplicas, replica)
+	}
+
 	entitlements, err := license.Entitlements(
 		ctx, api.Database,
-		api.Logger, len(api.replicaManager.AllPrimary()), len(api.ExternalAuthConfigs), api.LicenseKeys, map[codersdk.FeatureName]bool{
+		api.Logger, len(agedReplicas), len(api.ExternalAuthConfigs), api.LicenseKeys, map[codersdk.FeatureName]bool{
 			codersdk.FeatureAuditLog:                   api.AuditLogging,
 			codersdk.FeatureBrowserOnly:                api.BrowserOnly,
 			codersdk.FeatureSCIM:                       len(api.SCIMAPIKey) != 0,
diff --git a/enterprise/coderd/coderdenttest/coderdenttest.go b/enterprise/coderd/coderdenttest/coderdenttest.go
@@ -57,6 +57,7 @@ type Options struct {
 	DontAddLicense             bool
 	DontAddFirstUser           bool
 	ReplicaSyncUpdateInterval  time.Duration
+	ReplicaErrorGracePeriod    time.Duration
 	ExternalTokenEncryption    []dbcrypt.Cipher
 	ProvisionerDaemonPSK       string
 }
@@ -93,6 +94,7 @@ func NewWithAPI(t *testing.T, options *Options) (
 		DERPServerRelayAddress:     oop.AccessURL.String(),
 		DERPServerRegionID:         oop.BaseDERPMap.RegionIDs()[0],
 		ReplicaSyncUpdateInterval:  options.ReplicaSyncUpdateInterval,
+		ReplicaErrorGracePeriod:    options.ReplicaErrorGracePeriod,
 		Options:                    oop,
 		EntitlementsUpdateInterval: options.EntitlementsUpdateInterval,
 		LicenseKeys:                Keys,
diff --git a/enterprise/coderd/replicas_test.go b/enterprise/coderd/replicas_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/tls"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -22,9 +23,42 @@ import (
 func TestReplicas(t *testing.T) {
 	t.Parallel()
 	if !dbtestutil.WillUsePostgres() {
-		t.Skip("only test with real postgresF")
+		t.Skip("only test with real postgres")
 	}
 	t.Run("ErrorWithoutLicense", func(t *testing.T) {
+		t.Parallel()
+		// This will error because replicas are expected to instantly report
+		// errors when the license is not present.
+		db, pubsub := dbtestutil.NewDB(t)
+		firstClient, _ := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				IncludeProvisionerDaemon: true,
+				Database:                 db,
+				Pubsub:                   pubsub,
+			},
+			DontAddLicense:          true,
+			ReplicaErrorGracePeriod: time.Nanosecond,
+		})
+		secondClient, _, secondAPI, _ := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				Database: db,
+				Pubsub:   pubsub,
+			},
+			DontAddFirstUser:        true,
+			DontAddLicense:          true,
+			ReplicaErrorGracePeriod: time.Nanosecond,
+		})
+		secondClient.SetSessionToken(firstClient.SessionToken())
+		ents, err := secondClient.Entitlements(context.Background())
+		require.NoError(t, err)
+		require.Len(t, ents.Errors, 1)
+		_ = secondAPI.Close()
+
+		ents, err = firstClient.Entitlements(context.Background())
+		require.NoError(t, err)
+		require.Len(t, ents.Warnings, 0)
+	})
+	t.Run("DoesNotErrorBeforeGrace", func(t *testing.T) {
 		t.Parallel()
 		db, pubsub := dbtestutil.NewDB(t)
 		firstClient, _ := coderdenttest.New(t, &coderdenttest.Options{
@@ -46,12 +80,12 @@ func TestReplicas(t *testing.T) {
 		secondClient.SetSessionToken(firstClient.SessionToken())
 		ents, err := secondClient.Entitlements(context.Background())
 		require.NoError(t, err)
-		require.Len(t, ents.Errors, 1)
+		require.Len(t, ents.Errors, 0)
 		_ = secondAPI.Close()
 
 		ents, err = firstClient.Entitlements(context.Background())
 		require.NoError(t, err)
-		require.Len(t, ents.Warnings, 0)
+		require.Len(t, ents.Errors, 0)
 	})
 	t.Run("ConnectAcrossMultiple", func(t *testing.T) {
 		t.Parallel()
