Skip to content

Commit bfce6fe

Browse files
disqdisq
authored
feat: Support for quota query interval header (#1948)
The query interval is something we're introducing in the backend.
1 parent 1c4e9e6 commit bfce6fe

4 files changed

Lines changed: 59 additions & 30 deletions

File tree

‎premium/monitor.go‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -60,20 +60,23 @@ func WithCancelOnQuotaExceeded(ctx context.Context, qm QuotaMonitor, ops ...Quot
6060
return newCtx, nil
6161
}
6262

63-
func (qc quotaChecker) checkInitialQuota(ctx context.Context) error {
64-
hasQuota, err := qc.qm.HasQuota(ctx)
63+
func (qc *quotaChecker) checkInitialQuota(ctx context.Context) error {
64+
result, err := qc.qm.CheckQuota(ctx)
6565
if err != nil {
6666
return err
6767
}
68+
if result.SuggestedQueryInterval > 0 {
69+
qc.duration = result.SuggestedQueryInterval
70+
}
6871

69-
if !hasQuota {
72+
if !result.HasQuota {
7073
return ErrNoQuota{team: qc.qm.TeamName()}
7174
}
7275

7376
return nil
7477
}
7578

76-
func (qc quotaChecker) startQuotaMonitor(ctx context.Context) context.Context {
79+
func (qc *quotaChecker) startQuotaMonitor(ctx context.Context) context.Context {
7780
newCtx, cancelWithCause := context.WithCancelCause(ctx)
7881
go func() {
7982
ticker := time.NewTicker(qc.duration)
@@ -84,7 +87,7 @@ func (qc quotaChecker) startQuotaMonitor(ctx context.Context) context.Context {
8487
case <-newCtx.Done():
8588
return
8689
case <-ticker.C:
87-
hasQuota, err := qc.qm.HasQuota(newCtx)
90+
result, err := qc.qm.CheckQuota(newCtx)
8891
if err != nil {
8992
consecutiveFailures++
9093
hasQuotaErrors = errors.Join(hasQuotaErrors, err)
@@ -94,9 +97,13 @@ func (qc quotaChecker) startQuotaMonitor(ctx context.Context) context.Context {
9497
}
9598
continue
9699
}
100+
if result.SuggestedQueryInterval > 0 && qc.duration != result.SuggestedQueryInterval {
101+
qc.duration = result.SuggestedQueryInterval
102+
ticker.Reset(qc.duration)
103+
}
97104
consecutiveFailures = 0
98105
hasQuotaErrors = nil
99-
if !hasQuota {
106+
if !result.HasQuota {
100107
cancelWithCause(ErrNoQuota{team: qc.qm.TeamName()})
101108
return
102109
}

‎premium/monitor_test.go‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@ import (
1010
)
1111

1212
type quotaResponse struct {
13-
hasQuota bool
14-
err error
13+
result CheckQuotaResult
14+
err error
1515
}
1616

1717
func newFakeQuotaMonitor(hasQuota ...quotaResponse) *fakeQuotaMonitor {
@@ -23,12 +23,12 @@ type fakeQuotaMonitor struct {
2323
calls int
2424
}
2525

26-
func (f *fakeQuotaMonitor) HasQuota(_ context.Context) (bool, error) {
26+
func (f *fakeQuotaMonitor) CheckQuota(_ context.Context) (CheckQuotaResult, error) {
2727
resp := f.responses[f.calls]
2828
if f.calls < len(f.responses)-1 {
2929
f.calls++
3030
}
31-
return resp.hasQuota, resp.err
31+
return resp.result, resp.err
3232
}
3333

3434
func (*fakeQuotaMonitor) TeamName() string {
@@ -39,7 +39,7 @@ func TestWithCancelOnQuotaExceeded_NoInitialQuota(t *testing.T) {
3939
ctx := context.Background()
4040

4141
responses := []quotaResponse{
42-
{false, nil},
42+
{CheckQuotaResult{HasQuota: false}, nil},
4343
}
4444
_, err := WithCancelOnQuotaExceeded(ctx, newFakeQuotaMonitor(responses...))
4545

@@ -50,8 +50,8 @@ func TestWithCancelOnQuotaExceeded_NoQuota(t *testing.T) {
5050
ctx := context.Background()
5151

5252
responses := []quotaResponse{
53-
{true, nil},
54-
{false, nil},
53+
{CheckQuotaResult{HasQuota: true}, nil},
54+
{CheckQuotaResult{HasQuota: false}, nil},
5555
}
5656
ctx, err := WithCancelOnQuotaExceeded(ctx, newFakeQuotaMonitor(responses...), WithQuotaCheckPeriod(1*time.Millisecond))
5757
require.NoError(t, err)
@@ -65,9 +65,9 @@ func TestWithCancelOnQuotaCheckConsecutiveFailures(t *testing.T) {
6565
ctx := context.Background()
6666

6767
responses := []quotaResponse{
68-
{true, nil},
69-
{false, errors.New("test2")},
70-
{false, errors.New("test3")},
68+
{CheckQuotaResult{HasQuota: true}, nil},
69+
{CheckQuotaResult{HasQuota: false}, errors.New("test2")},
70+
{CheckQuotaResult{HasQuota: false}, errors.New("test3")},
7171
}
7272
ctx, err := WithCancelOnQuotaExceeded(ctx,
7373
newFakeQuotaMonitor(responses...),

‎premium/usage.go‎

Lines changed: 32 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ const (
4343
BatchLimitHeader = "x-cq-batch-limit"
4444
MinimumUpdateIntervalHeader = "x-cq-minimum-update-interval"
4545
MaximumUpdateIntervalHeader = "x-cq-maximum-update-interval"
46+
QueryIntervalHeader = "x-cq-query-interval"
4647
)
4748

4849
//go:generate mockgen -package=mocks -destination=../premium/mocks/marketplacemetering.go -source=usage.go AWSMarketplaceClientInterface
@@ -55,11 +56,19 @@ type TokenClient interface {
5556
GetTokenType() auth.TokenType
5657
}
5758

59+
type CheckQuotaResult struct {
60+
// HasQuota is true if the quota has not been exceeded
61+
HasQuota bool
62+
63+
// SuggestedQueryInterval is the suggested interval to wait before querying the API again
64+
SuggestedQueryInterval time.Duration
65+
}
66+
5867
type QuotaMonitor interface {
5968
// TeamName returns the team name
6069
TeamName() string
61-
// HasQuota returns true if the quota has not been exceeded
62-
HasQuota(context.Context) (bool, error)
70+
// CheckQuota checks if the quota has been exceeded
71+
CheckQuota(context.Context) (CheckQuotaResult, error)
6372
}
6473

6574
type UsageClient interface {
@@ -359,21 +368,34 @@ func (u *BatchUpdater) TeamName() string {
359368
return u.teamName
360369
}
361370

362-
func (u *BatchUpdater) HasQuota(ctx context.Context) (bool, error) {
371+
func (u *BatchUpdater) CheckQuota(ctx context.Context) (CheckQuotaResult, error) {
363372
if u.awsMarketplaceClient != nil {
364-
return true, nil
373+
return CheckQuotaResult{HasQuota: true}, nil
365374
}
366375
u.logger.Debug().Str("url", u.url).Str("team", u.teamName).Str("pluginTeam", u.pluginMeta.Team).Str("pluginKind", string(u.pluginMeta.Kind)).Str("pluginName", u.pluginMeta.Name).Msg("checking quota")
367376
usage, err := u.apiClient.GetTeamPluginUsageWithResponse(ctx, u.teamName, u.pluginMeta.Team, u.pluginMeta.Kind, u.pluginMeta.Name)
368377
if err != nil {
369-
return false, fmt.Errorf("failed to get usage: %w", err)
378+
return CheckQuotaResult{HasQuota: false}, fmt.Errorf("failed to get usage: %w", err)
370379
}
371380
if usage.StatusCode() != http.StatusOK {
372-
return false, fmt.Errorf("failed to get usage: %s", usage.Status())
381+
return CheckQuotaResult{HasQuota: false}, fmt.Errorf("failed to get usage: %s", usage.Status())
373382
}
374383

375-
hasQuota := usage.JSON200.RemainingRows == nil || *usage.JSON200.RemainingRows > 0
376-
return hasQuota, nil
384+
res := CheckQuotaResult{
385+
HasQuota: usage.JSON200.RemainingRows == nil || *usage.JSON200.RemainingRows > 0,
386+
}
387+
if usage.HTTPResponse == nil {
388+
return res, nil
389+
}
390+
if headerValue := usage.HTTPResponse.Header.Get(QueryIntervalHeader); headerValue != "" {
391+
interval, err := strconv.ParseUint(headerValue, 10, 32)
392+
if interval > 0 {
393+
res.SuggestedQueryInterval = time.Duration(interval) * time.Second
394+
} else {
395+
u.logger.Warn().Err(err).Str(QueryIntervalHeader, headerValue).Msg("failed to parse query interval")
396+
}
397+
}
398+
return res, nil
377399
}
378400

379401
func (u *BatchUpdater) Close() error {
@@ -700,8 +722,8 @@ func (n *NoOpUsageClient) TeamName() string {
700722
return n.TeamNameValue
701723
}
702724

703-
func (NoOpUsageClient) HasQuota(_ context.Context) (bool, error) {
704-
return true, nil
725+
func (NoOpUsageClient) CheckQuota(_ context.Context) (CheckQuotaResult, error) {
726+
return CheckQuotaResult{HasQuota: true}, nil
705727
}
706728

707729
func (NoOpUsageClient) Increase(_ uint32) error {

‎premium/usage_test.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -113,10 +113,10 @@ func TestUsageService_HasQuota_NoRowsRemaining(t *testing.T) {
113113

114114
usageClient := newClient(t, apiClient, WithBatchLimit(0))
115115

116-
hasQuota, err := usageClient.HasQuota(ctx)
116+
result, err := usageClient.CheckQuota(ctx)
117117
require.NoError(t, err)
118118

119-
assert.False(t, hasQuota, "should not have quota")
119+
assert.False(t, result.HasQuota, "should not have quota")
120120
}
121121

122122
func TestUsageService_HasQuota_WithRowsRemaining(t *testing.T) {
@@ -130,10 +130,10 @@ func TestUsageService_HasQuota_WithRowsRemaining(t *testing.T) {
130130

131131
usageClient := newClient(t, apiClient, WithBatchLimit(0))
132132

133-
hasQuota, err := usageClient.HasQuota(ctx)
133+
result, err := usageClient.CheckQuota(ctx)
134134
require.NoError(t, err)
135135

136-
assert.True(t, hasQuota, "should have quota")
136+
assert.True(t, result.HasQuota, "should have quota")
137137
}
138138

139139
func TestUsageService_Increase_ZeroBatchSize(t *testing.T) {

0 commit comments

Comments
 (0)