Skip to content

Commit 1955ae3

Browse files
kylecarbskylecarbs
committed
Add tests for resync and pubsub, and add back previous exp backoff retry
1 parent 5215fe9 commit 1955ae3

File tree

6 files changed

+166
-42
lines changed

6 files changed

+166
-42
lines changed

‎enterprise/coderd/coderd.go

Lines changed: 61 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010

1111
"golang.org/x/xerrors"
1212

13+
"github.com/cenkalti/backoff/v4"
1314
"github.com/go-chi/chi/v5"
1415

1516
"cdr.dev/slog"
@@ -64,7 +65,7 @@ func New(ctx context.Context, options *Options) (*API, error) {
6465
if err != nil {
6566
return nil, xerrors.Errorf("update entitlements: %w", err)
6667
}
67-
api.closeLicenseSubscribe, err = api.Pubsub.Subscribe(pubSubEventLicenses, func(ctx context.Context, message []byte) {
68+
api.closeLicenseSubscribe, err = api.Pubsub.Subscribe(PubsubEventLicenses, func(ctx context.Context, message []byte) {
6869
_ = api.updateEntitlements(ctx)
6970
})
7071
if err != nil {
@@ -101,23 +102,6 @@ func (api *API) Close() error {
101102
return api.AGPL.Close()
102103
}
103104

104-
func (api *API) runEntitlementsLoop(ctx context.Context) {
105-
ticker := time.NewTicker(api.EntitlementsUpdateInterval)
106-
defer ticker.Stop()
107-
for {
108-
select {
109-
case <-ctx.Done():
110-
return
111-
case <-ticker.C:
112-
}
113-
err := api.updateEntitlements(ctx)
114-
if err != nil {
115-
api.Logger.Warn(ctx, "failed to get feature entitlements", slog.Error(err))
116-
continue
117-
}
118-
}
119-
}
120-
121105
func (api *API) updateEntitlements(ctx context.Context) error {
122106
licenses, err := api.Database.GetUnexpiredLicenses(ctx)
123107
if err != nil {
@@ -169,7 +153,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
169153
auditor := agplaudit.NewNop()
170154
// A flag could be added to the options that would allow disabling
171155
// enhanced audit logging here!
172-
if api.auditLogs == codersdk.EntitlementEntitled && api.AuditLogging {
156+
if auditLogs == codersdk.EntitlementEntitled && api.AuditLogging {
173157
auditor = audit.NewAuditor(
174158
audit.DefaultFilter,
175159
backends.NewPostgres(api.Database, true),
@@ -231,6 +215,64 @@ func (api *API) entitlements(rw http.ResponseWriter, r *http.Request) {
231215
httpapi.Write(rw, http.StatusOK, resp)
232216
}
233217

218+
func (api *API) runEntitlementsLoop(ctx context.Context) {
219+
eb := backoff.NewExponentialBackOff()
220+
eb.MaxElapsedTime = 0 // retry indefinitely
221+
b := backoff.WithContext(eb, ctx)
222+
updates := make(chan struct{}, 1)
223+
subscribed := false
224+
225+
for {
226+
select {
227+
case <-ctx.Done():
228+
return
229+
default:
230+
// pass
231+
}
232+
if !subscribed {
233+
cancel, err := api.Pubsub.Subscribe(PubsubEventLicenses, func(_ context.Context, _ []byte) {
234+
// don't block. If the channel is full, drop the event, as there is a resync
235+
// scheduled already.
236+
select {
237+
case updates <- struct{}{}:
238+
// pass
239+
default:
240+
// pass
241+
}
242+
})
243+
if err != nil {
244+
api.Logger.Warn(ctx, "failed to subscribe to license updates", slog.Error(err))
245+
time.Sleep(b.NextBackOff())
246+
continue
247+
}
248+
// nolint: revive
249+
defer cancel()
250+
subscribed = true
251+
api.Logger.Debug(ctx, "successfully subscribed to pubsub")
252+
}
253+
254+
api.Logger.Info(ctx, "syncing licensed entitlements")
255+
err := api.updateEntitlements(ctx)
256+
if err != nil {
257+
api.Logger.Warn(ctx, "failed to get feature entitlements", slog.Error(err))
258+
time.Sleep(b.NextBackOff())
259+
continue
260+
}
261+
b.Reset()
262+
api.Logger.Debug(ctx, "synced licensed entitlements")
263+
264+
select {
265+
case <-ctx.Done():
266+
return
267+
case <-time.After(api.EntitlementsUpdateInterval):
268+
continue
269+
case <-updates:
270+
api.Logger.Debug(ctx, "got pubsub update")
271+
continue
272+
}
273+
}
274+
}
275+
234276
func max(a, b int64) int64 {
235277
if a > b {
236278
return a

‎enterprise/coderd/coderd_test.go

Lines changed: 79 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,22 @@ package coderd_test
22

33
import (
44
"context"
5+
"reflect"
56
"testing"
67
"time"
78

89
"github.com/stretchr/testify/assert"
910
"github.com/stretchr/testify/require"
1011
"go.uber.org/goleak"
1112

13+
agplaudit "github.com/coder/coder/coderd/audit"
1214
"github.com/coder/coder/coderd/coderdtest"
15+
"github.com/coder/coder/coderd/database"
1316
"github.com/coder/coder/codersdk"
17+
"github.com/coder/coder/enterprise/audit"
18+
"github.com/coder/coder/enterprise/coderd"
1419
"github.com/coder/coder/enterprise/coderd/coderdenttest"
20+
"github.com/coder/coder/testutil"
1521
)
1622

1723
func TestMain(m *testing.M) {
@@ -40,7 +46,7 @@ func TestEntitlements(t *testing.T) {
4046
t.Parallel()
4147
client := coderdenttest.New(t, nil)
4248
_ = coderdtest.CreateFirstUser(t, client)
43-
coderdenttest.AddLicense(t, client, coderdenttest.AddLicenseOptions{
49+
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
4450
UserLimit: 100,
4551
AuditLog: true,
4652
})
@@ -63,7 +69,7 @@ func TestEntitlements(t *testing.T) {
6369
t.Parallel()
6470
client := coderdenttest.New(t, nil)
6571
_ = coderdtest.CreateFirstUser(t, client)
66-
license := coderdenttest.AddLicense(t, client, coderdenttest.AddLicenseOptions{
72+
license := coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
6773
UserLimit: 100,
6874
AuditLog: true,
6975
})
@@ -79,7 +85,7 @@ func TestEntitlements(t *testing.T) {
7985

8086
res, err = client.Entitlements(context.Background())
8187
require.NoError(t, err)
82-
assert.True(t, res.HasLicense)
88+
assert.False(t, res.HasLicense)
8389
al = res.Features[codersdk.FeatureAuditLog]
8490
assert.Equal(t, codersdk.EntitlementNotEntitled, al.Entitlement)
8591
assert.True(t, al.Enabled)
@@ -91,7 +97,7 @@ func TestEntitlements(t *testing.T) {
9197
for i := 0; i < 4; i++ {
9298
coderdtest.CreateAnotherUser(t, client, first.OrganizationID)
9399
}
94-
coderdenttest.AddLicense(t, client, coderdenttest.AddLicenseOptions{
100+
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
95101
UserLimit: 4,
96102
AuditLog: true,
97103
GraceAt: time.Now().Add(-time.Second),
@@ -115,4 +121,73 @@ func TestEntitlements(t *testing.T) {
115121
assert.Contains(t, res.Warnings,
116122
"Audit logging is enabled but your license for this feature is expired.")
117123
})
124+
t.Run("Pubsub", func(t *testing.T) {
125+
t.Parallel()
126+
client, _, api := coderdenttest.NewWithAPI(t, nil)
127+
entitlements, err := client.Entitlements(context.Background())
128+
require.NoError(t, err)
129+
require.False(t, entitlements.HasLicense)
130+
coderdtest.CreateFirstUser(t, client)
131+
_, err = api.Database.InsertLicense(context.Background(), database.InsertLicenseParams{
132+
UploadedAt: database.Now(),
133+
Exp: database.Now().AddDate(1, 0, 0),
134+
JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
135+
AuditLog: true,
136+
}),
137+
})
138+
require.NoError(t, err)
139+
err = api.Pubsub.Publish(coderd.PubsubEventLicenses, []byte{})
140+
require.NoError(t, err)
141+
require.Eventually(t, func() bool {
142+
entitlements, err := client.Entitlements(context.Background())
143+
assert.NoError(t, err)
144+
return entitlements.HasLicense
145+
}, testutil.WaitShort, testutil.IntervalFast)
146+
})
147+
t.Run("Resync", func(t *testing.T) {
148+
t.Parallel()
149+
client, _, api := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
150+
EntitlementsUpdateInterval: 25 * time.Millisecond,
151+
})
152+
entitlements, err := client.Entitlements(context.Background())
153+
require.NoError(t, err)
154+
require.False(t, entitlements.HasLicense)
155+
coderdtest.CreateFirstUser(t, client)
156+
_, err = api.Database.InsertLicense(context.Background(), database.InsertLicenseParams{
157+
UploadedAt: database.Now(),
158+
Exp: database.Now().AddDate(1, 0, 0),
159+
JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
160+
AuditLog: true,
161+
}),
162+
})
163+
require.NoError(t, err)
164+
require.Eventually(t, func() bool {
165+
entitlements, err := client.Entitlements(context.Background())
166+
assert.NoError(t, err)
167+
return entitlements.HasLicense
168+
}, testutil.WaitShort, testutil.IntervalFast)
169+
})
170+
}
171+
172+
func TestAuditLogging(t *testing.T) {
173+
t.Parallel()
174+
t.Run("Enabled", func(t *testing.T) {
175+
t.Parallel()
176+
client, _, api := coderdenttest.NewWithAPI(t, nil)
177+
coderdtest.CreateFirstUser(t, client)
178+
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
179+
AuditLog: true,
180+
})
181+
_, auditor := api.AGPL.Auditor.Load(context.Background())
182+
ea := audit.NewAuditor(audit.DefaultFilter)
183+
assert.Equal(t, reflect.ValueOf(ea).Type(), reflect.ValueOf(auditor).Type())
184+
})
185+
t.Run("Disabled", func(t *testing.T) {
186+
t.Parallel()
187+
client, _, api := coderdenttest.NewWithAPI(t, nil)
188+
coderdtest.CreateFirstUser(t, client)
189+
_, auditor := api.AGPL.Auditor.Load(context.Background())
190+
ea := agplaudit.NewNop()
191+
assert.Equal(t, reflect.ValueOf(ea).Type(), reflect.ValueOf(auditor).Type())
192+
})
118193
}

‎enterprise/coderd/coderdenttest/coderdenttest.go

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ func init() {
3636

3737
type Options struct {
3838
*coderdtest.Options
39+
EntitlementsUpdateInterval time.Duration
3940
}
4041

4142
// New constructs a codersdk client connected to an in-memory Enterprise API instance.
@@ -53,8 +54,9 @@ func NewWithAPI(t *testing.T, options *Options) (*codersdk.Client, io.Closer, *c
5354
}
5455
srv, oop := coderdtest.NewOptions(t, options.Options)
5556
coderAPI, err := coderd.New(context.Background(), &coderd.Options{
56-
AuditLogging: true,
57-
Options: oop,
57+
AuditLogging: true,
58+
Options: oop,
59+
EntitlementsUpdateInterval: options.EntitlementsUpdateInterval,
5860
Keys: map[string]ed25519.PublicKey{
5961
testKeyID: testPublicKey,
6062
},
@@ -72,7 +74,7 @@ func NewWithAPI(t *testing.T, options *Options) (*codersdk.Client, io.Closer, *c
7274
return codersdk.New(coderAPI.AccessURL), provisionerCloser, coderAPI
7375
}
7476

75-
type AddLicenseOptions struct {
77+
type LicenseOptions struct {
7678
AccountType string
7779
AccountID string
7880
GraceAt time.Time
@@ -82,7 +84,16 @@ type AddLicenseOptions struct {
8284
}
8385

8486
// AddLicense generates a new license with the options provided and inserts it.
85-
func AddLicense(t *testing.T, client *codersdk.Client, options AddLicenseOptions) codersdk.License {
87+
func AddLicense(t *testing.T, client *codersdk.Client, options LicenseOptions) codersdk.License {
88+
license, err := client.AddLicense(context.Background(), codersdk.AddLicenseRequest{
89+
License: GenerateLicense(t, options),
90+
})
91+
require.NoError(t, err)
92+
return license
93+
}
94+
95+
// GenerateLicense returns a signed JWT using the test key.
96+
func GenerateLicense(t *testing.T, options LicenseOptions) string {
8697
if options.ExpiresAt.IsZero() {
8798
options.ExpiresAt = time.Now().Add(time.Hour)
8899
}
@@ -113,11 +124,7 @@ func AddLicense(t *testing.T, client *codersdk.Client, options AddLicenseOptions
113124
tok.Header[coderd.HeaderKeyID] = testKeyID
114125
signedTok, err := tok.SignedString(testPrivateKey)
115126
require.NoError(t, err)
116-
license, err := client.AddLicense(context.Background(), codersdk.AddLicenseRequest{
117-
License: signedTok,
118-
})
119-
require.NoError(t, err)
120-
return license
127+
return signedTok
121128
}
122129

123130
type nopcloser struct{}

‎enterprise/coderd/coderdenttest/coderdenttest_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -402,7 +402,7 @@ func newAuthTester(ctx context.Context, t *testing.T) *authTester {
402402
DestinationScheme: codersdk.ParameterDestinationSchemeProvisionerVariable,
403403
})
404404
require.NoError(t, err, "create template param")
405-
license := coderdenttest.AddLicense(t, client, coderdenttest.AddLicenseOptions{})
405+
license := coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{})
406406
urlParameters := map[string]string{
407407
"{organization}": admin.OrganizationID.String(),
408408
"{user}": admin.UserID.String(),

‎enterprise/coderd/licenses.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ const (
3131
AccountTypeSalesforce = "salesforce"
3232
VersionClaim = "version"
3333

34-
pubSubEventLicenses = "licenses"
34+
PubsubEventLicenses = "licenses"
3535
)
3636

3737
var ValidMethods = []string{"EdDSA"}
@@ -127,7 +127,7 @@ func (api *API) postLicense(rw http.ResponseWriter, r *http.Request) {
127127
})
128128
return
129129
}
130-
err = api.Pubsub.Publish(pubSubEventLicenses, []byte("add"))
130+
err = api.Pubsub.Publish(PubsubEventLicenses, []byte("add"))
131131
if err != nil {
132132
api.Logger.Error(context.Background(), "failed to publish license add", slog.Error(err))
133133
// don't fail the HTTP request, since we did write it successfully to the database
@@ -206,7 +206,7 @@ func (api *API) deleteLicense(rw http.ResponseWriter, r *http.Request) {
206206
})
207207
return
208208
}
209-
err = api.Pubsub.Publish(pubSubEventLicenses, []byte("delete"))
209+
err = api.Pubsub.Publish(PubsubEventLicenses, []byte("delete"))
210210
if err != nil {
211211
api.Logger.Error(context.Background(), "failed to publish license delete", slog.Error(err))
212212
// don't fail the HTTP request, since we did write it successfully to the database

‎enterprise/coderd/licenses_test.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ func TestPostLicense(t *testing.T) {
2424
t.Parallel()
2525
client := coderdenttest.New(t, nil)
2626
_ = coderdtest.CreateFirstUser(t, client)
27-
respLic := coderdenttest.AddLicense(t, client, coderdenttest.AddLicenseOptions{
27+
respLic := coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
2828
AccountType: coderd.AccountTypeSalesforce,
2929
AccountID: "testing",
3030
AuditLog: true,
@@ -55,7 +55,7 @@ func TestPostLicense(t *testing.T) {
5555
t.Parallel()
5656
client := coderdenttest.New(t, nil)
5757
_ = coderdtest.CreateFirstUser(t, client)
58-
coderdenttest.AddLicense(t, client, coderdenttest.AddLicenseOptions{})
58+
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{})
5959
_, err := client.AddLicense(context.Background(), codersdk.AddLicenseRequest{
6060
License: "invalid",
6161
})
@@ -77,12 +77,12 @@ func TestGetLicense(t *testing.T) {
7777
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
7878
defer cancel()
7979

80-
coderdenttest.AddLicense(t, client, coderdenttest.AddLicenseOptions{
80+
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
8181
AccountID: "testing",
8282
AuditLog: true,
8383
})
8484

85-
coderdenttest.AddLicense(t, client, coderdenttest.AddLicenseOptions{
85+
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
8686
AccountID: "testing2",
8787
AuditLog: true,
8888
UserLimit: 200,
@@ -144,11 +144,11 @@ func TestDeleteLicense(t *testing.T) {
144144
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
145145
defer cancel()
146146

147-
coderdenttest.AddLicense(t, client, coderdenttest.AddLicenseOptions{
147+
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
148148
AccountID: "testing",
149149
AuditLog: true,
150150
})
151-
coderdenttest.AddLicense(t, client, coderdenttest.AddLicenseOptions{
151+
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
152152
AccountID: "testing2",
153153
AuditLog: true,
154154
UserLimit: 200,

0 commit comments

Comments
 (0)