From d9188d055bbdc6435257eb499de142d72a5cc8e7 Mon Sep 17 00:00:00 2001 From: Kemal Hadimli Date: Tue, 29 Oct 2024 14:09:56 +0000 Subject: [PATCH 1/6] feat: Support for quota query interval header --- premium/monitor.go | 12 ++++++++++-- premium/monitor_test.go | 16 ++++++++-------- premium/usage.go | 36 +++++++++++++++++++++++++++--------- premium/usage_test.go | 4 ++-- 4 files changed, 47 insertions(+), 21 deletions(-) diff --git a/premium/monitor.go b/premium/monitor.go index 02573b4672..e15bcc3c9e 100644 --- a/premium/monitor.go +++ b/premium/monitor.go @@ -65,8 +65,11 @@ func (qc quotaChecker) checkInitialQuota(ctx context.Context) error { if err != nil { return err } + if hasQuota.SuggestedQueryInterval > 0 { + qc.duration = hasQuota.SuggestedQueryInterval + } - if !hasQuota { + if !hasQuota.HasQuota { return ErrNoQuota{team: qc.qm.TeamName()} } @@ -94,9 +97,14 @@ func (qc quotaChecker) startQuotaMonitor(ctx context.Context) context.Context { } continue } + if hasQuota.SuggestedQueryInterval > 0 && qc.duration != hasQuota.SuggestedQueryInterval { + qc.duration = hasQuota.SuggestedQueryInterval + ticker.Stop() + ticker = time.NewTicker(qc.duration) + } consecutiveFailures = 0 hasQuotaErrors = nil - if !hasQuota { + if !hasQuota.HasQuota { cancelWithCause(ErrNoQuota{team: qc.qm.TeamName()}) return } diff --git a/premium/monitor_test.go b/premium/monitor_test.go index 0ce3876bda..0578d16ac4 100644 --- a/premium/monitor_test.go +++ b/premium/monitor_test.go @@ -10,7 +10,7 @@ import ( ) type quotaResponse struct { - hasQuota bool + hasQuota QuotaCheckResult err error } @@ -23,7 +23,7 @@ type fakeQuotaMonitor struct { calls int } -func (f *fakeQuotaMonitor) HasQuota(_ context.Context) (bool, error) { +func (f *fakeQuotaMonitor) HasQuota(_ context.Context) (QuotaCheckResult, error) { resp := f.responses[f.calls] if f.calls < len(f.responses)-1 { f.calls++ @@ -39,7 +39,7 @@ func TestWithCancelOnQuotaExceeded_NoInitialQuota(t *testing.T) { ctx := context.Background() responses := []quotaResponse{ - {false, nil}, + {QuotaCheckResult{HasQuota: false}, nil}, } _, err := WithCancelOnQuotaExceeded(ctx, newFakeQuotaMonitor(responses...)) @@ -50,8 +50,8 @@ func TestWithCancelOnQuotaExceeded_NoQuota(t *testing.T) { ctx := context.Background() responses := []quotaResponse{ - {true, nil}, - {false, nil}, + {QuotaCheckResult{HasQuota: true}, nil}, + {QuotaCheckResult{HasQuota: false}, nil}, } ctx, err := WithCancelOnQuotaExceeded(ctx, newFakeQuotaMonitor(responses...), WithQuotaCheckPeriod(1*time.Millisecond)) require.NoError(t, err) @@ -65,9 +65,9 @@ func TestWithCancelOnQuotaCheckConsecutiveFailures(t *testing.T) { ctx := context.Background() responses := []quotaResponse{ - {true, nil}, - {false, errors.New("test2")}, - {false, errors.New("test3")}, + {QuotaCheckResult{HasQuota: true}, nil}, + {QuotaCheckResult{HasQuota: false}, errors.New("test2")}, + {QuotaCheckResult{HasQuota: false}, errors.New("test3")}, } ctx, err := WithCancelOnQuotaExceeded(ctx, newFakeQuotaMonitor(responses...), diff --git a/premium/usage.go b/premium/usage.go index 50f6265a28..bb965292ea 100644 --- a/premium/usage.go +++ b/premium/usage.go @@ -43,6 +43,7 @@ const ( BatchLimitHeader = "x-cq-batch-limit" MinimumUpdateIntervalHeader = "x-cq-minimum-update-interval" MaximumUpdateIntervalHeader = "x-cq-maximum-update-interval" + QueryIntervalHeader = "x-cq-query-interval" ) //go:generate mockgen -package=mocks -destination=../premium/mocks/marketplacemetering.go -source=usage.go AWSMarketplaceClientInterface @@ -55,11 +56,16 @@ type TokenClient interface { GetTokenType() auth.TokenType } +type QuotaCheckResult struct { + HasQuota bool + SuggestedQueryInterval time.Duration +} + type QuotaMonitor interface { // TeamName returns the team name TeamName() string // HasQuota returns true if the quota has not been exceeded - HasQuota(context.Context) (bool, error) + HasQuota(context.Context) (QuotaCheckResult, error) } type UsageClient interface { @@ -359,21 +365,33 @@ func (u *BatchUpdater) TeamName() string { return u.teamName } -func (u *BatchUpdater) HasQuota(ctx context.Context) (bool, error) { +func (u *BatchUpdater) HasQuota(ctx context.Context) (QuotaCheckResult, error) { if u.awsMarketplaceClient != nil { - return true, nil + return QuotaCheckResult{HasQuota: true}, nil } 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") usage, err := u.apiClient.GetTeamPluginUsageWithResponse(ctx, u.teamName, u.pluginMeta.Team, u.pluginMeta.Kind, u.pluginMeta.Name) if err != nil { - return false, fmt.Errorf("failed to get usage: %w", err) + return QuotaCheckResult{HasQuota: false}, fmt.Errorf("failed to get usage: %w", err) } if usage.StatusCode() != http.StatusOK { - return false, fmt.Errorf("failed to get usage: %s", usage.Status()) + return QuotaCheckResult{HasQuota: false}, fmt.Errorf("failed to get usage: %s", usage.Status()) } - hasQuota := usage.JSON200.RemainingRows == nil || *usage.JSON200.RemainingRows > 0 - return hasQuota, nil + res := QuotaCheckResult{ + HasQuota: usage.JSON200.RemainingRows == nil || *usage.JSON200.RemainingRows > 0, + } + if usage.HTTPResponse == nil { + return res, nil + } + if headerValue := usage.HTTPResponse.Header.Get(QueryIntervalHeader); headerValue != "" { + if interval, err := strconv.ParseUint(headerValue, 10, 32); err != nil { + u.logger.Warn().Err(err).Str(QueryIntervalHeader, headerValue).Msg("failed to parse query interval") + } else if interval > 0 { + res.SuggestedQueryInterval = time.Duration(interval) * time.Second + } + } + return res, nil } func (u *BatchUpdater) Close() error { @@ -697,8 +715,8 @@ func (n *NoOpUsageClient) TeamName() string { return n.TeamNameValue } -func (NoOpUsageClient) HasQuota(_ context.Context) (bool, error) { - return true, nil +func (NoOpUsageClient) HasQuota(_ context.Context) (QuotaCheckResult, error) { + return QuotaCheckResult{HasQuota: true}, nil } func (NoOpUsageClient) Increase(_ uint32) error { diff --git a/premium/usage_test.go b/premium/usage_test.go index 9a3a17b3a6..b2935a4a47 100644 --- a/premium/usage_test.go +++ b/premium/usage_test.go @@ -116,7 +116,7 @@ func TestUsageService_HasQuota_NoRowsRemaining(t *testing.T) { hasQuota, err := usageClient.HasQuota(ctx) require.NoError(t, err) - assert.False(t, hasQuota, "should not have quota") + assert.False(t, hasQuota.HasQuota, "should not have quota") } func TestUsageService_HasQuota_WithRowsRemaining(t *testing.T) { @@ -133,7 +133,7 @@ func TestUsageService_HasQuota_WithRowsRemaining(t *testing.T) { hasQuota, err := usageClient.HasQuota(ctx) require.NoError(t, err) - assert.True(t, hasQuota, "should have quota") + assert.True(t, hasQuota.HasQuota, "should have quota") } func TestUsageService_Increase_ZeroBatchSize(t *testing.T) { From 8c1686ef5931025966f3daef5f98675152fabb09 Mon Sep 17 00:00:00 2001 From: Kemal Hadimli Date: Tue, 29 Oct 2024 14:19:38 +0000 Subject: [PATCH 2/6] lint actually caught something --- premium/monitor.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/premium/monitor.go b/premium/monitor.go index e15bcc3c9e..f178d0eb9c 100644 --- a/premium/monitor.go +++ b/premium/monitor.go @@ -60,7 +60,7 @@ func WithCancelOnQuotaExceeded(ctx context.Context, qm QuotaMonitor, ops ...Quot return newCtx, nil } -func (qc quotaChecker) checkInitialQuota(ctx context.Context) error { +func (qc *quotaChecker) checkInitialQuota(ctx context.Context) error { hasQuota, err := qc.qm.HasQuota(ctx) if err != nil { return err @@ -76,7 +76,7 @@ func (qc quotaChecker) checkInitialQuota(ctx context.Context) error { return nil } -func (qc quotaChecker) startQuotaMonitor(ctx context.Context) context.Context { +func (qc *quotaChecker) startQuotaMonitor(ctx context.Context) context.Context { newCtx, cancelWithCause := context.WithCancelCause(ctx) go func() { ticker := time.NewTicker(qc.duration) From 6677d34750c66b23dfc0cb86ab935dc14940dd3e Mon Sep 17 00:00:00 2001 From: Kemal Hadimli Date: Tue, 29 Oct 2024 15:35:27 +0000 Subject: [PATCH 3/6] CR: Swap names around: `HasQuota()` => `CheckQuota()`, `QuotaCheckResult` => `CheckQuotaResult`, `hasQuota` => `result` --- premium/monitor.go | 16 ++++++++-------- premium/monitor_test.go | 16 ++++++++-------- premium/usage.go | 25 ++++++++++++++----------- premium/usage_test.go | 8 ++++---- 4 files changed, 34 insertions(+), 31 deletions(-) diff --git a/premium/monitor.go b/premium/monitor.go index f178d0eb9c..97a70824c9 100644 --- a/premium/monitor.go +++ b/premium/monitor.go @@ -61,15 +61,15 @@ func WithCancelOnQuotaExceeded(ctx context.Context, qm QuotaMonitor, ops ...Quot } func (qc *quotaChecker) checkInitialQuota(ctx context.Context) error { - hasQuota, err := qc.qm.HasQuota(ctx) + result, err := qc.qm.CheckQuota(ctx) if err != nil { return err } - if hasQuota.SuggestedQueryInterval > 0 { - qc.duration = hasQuota.SuggestedQueryInterval + if result.SuggestedQueryInterval > 0 { + qc.duration = result.SuggestedQueryInterval } - if !hasQuota.HasQuota { + if !result.HasQuota { return ErrNoQuota{team: qc.qm.TeamName()} } @@ -87,7 +87,7 @@ func (qc *quotaChecker) startQuotaMonitor(ctx context.Context) context.Context { case <-newCtx.Done(): return case <-ticker.C: - hasQuota, err := qc.qm.HasQuota(newCtx) + result, err := qc.qm.CheckQuota(newCtx) if err != nil { consecutiveFailures++ hasQuotaErrors = errors.Join(hasQuotaErrors, err) @@ -97,14 +97,14 @@ func (qc *quotaChecker) startQuotaMonitor(ctx context.Context) context.Context { } continue } - if hasQuota.SuggestedQueryInterval > 0 && qc.duration != hasQuota.SuggestedQueryInterval { - qc.duration = hasQuota.SuggestedQueryInterval + if result.SuggestedQueryInterval > 0 && qc.duration != result.SuggestedQueryInterval { + qc.duration = result.SuggestedQueryInterval ticker.Stop() ticker = time.NewTicker(qc.duration) } consecutiveFailures = 0 hasQuotaErrors = nil - if !hasQuota.HasQuota { + if !result.HasQuota { cancelWithCause(ErrNoQuota{team: qc.qm.TeamName()}) return } diff --git a/premium/monitor_test.go b/premium/monitor_test.go index 0578d16ac4..8cff1cdb23 100644 --- a/premium/monitor_test.go +++ b/premium/monitor_test.go @@ -10,7 +10,7 @@ import ( ) type quotaResponse struct { - hasQuota QuotaCheckResult + hasQuota CheckQuotaResult err error } @@ -23,7 +23,7 @@ type fakeQuotaMonitor struct { calls int } -func (f *fakeQuotaMonitor) HasQuota(_ context.Context) (QuotaCheckResult, error) { +func (f *fakeQuotaMonitor) CheckQuota(_ context.Context) (CheckQuotaResult, error) { resp := f.responses[f.calls] if f.calls < len(f.responses)-1 { f.calls++ @@ -39,7 +39,7 @@ func TestWithCancelOnQuotaExceeded_NoInitialQuota(t *testing.T) { ctx := context.Background() responses := []quotaResponse{ - {QuotaCheckResult{HasQuota: false}, nil}, + {CheckQuotaResult{HasQuota: false}, nil}, } _, err := WithCancelOnQuotaExceeded(ctx, newFakeQuotaMonitor(responses...)) @@ -50,8 +50,8 @@ func TestWithCancelOnQuotaExceeded_NoQuota(t *testing.T) { ctx := context.Background() responses := []quotaResponse{ - {QuotaCheckResult{HasQuota: true}, nil}, - {QuotaCheckResult{HasQuota: false}, nil}, + {CheckQuotaResult{HasQuota: true}, nil}, + {CheckQuotaResult{HasQuota: false}, nil}, } ctx, err := WithCancelOnQuotaExceeded(ctx, newFakeQuotaMonitor(responses...), WithQuotaCheckPeriod(1*time.Millisecond)) require.NoError(t, err) @@ -65,9 +65,9 @@ func TestWithCancelOnQuotaCheckConsecutiveFailures(t *testing.T) { ctx := context.Background() responses := []quotaResponse{ - {QuotaCheckResult{HasQuota: true}, nil}, - {QuotaCheckResult{HasQuota: false}, errors.New("test2")}, - {QuotaCheckResult{HasQuota: false}, errors.New("test3")}, + {CheckQuotaResult{HasQuota: true}, nil}, + {CheckQuotaResult{HasQuota: false}, errors.New("test2")}, + {CheckQuotaResult{HasQuota: false}, errors.New("test3")}, } ctx, err := WithCancelOnQuotaExceeded(ctx, newFakeQuotaMonitor(responses...), diff --git a/premium/usage.go b/premium/usage.go index bb965292ea..d6eaea8007 100644 --- a/premium/usage.go +++ b/premium/usage.go @@ -56,16 +56,19 @@ type TokenClient interface { GetTokenType() auth.TokenType } -type QuotaCheckResult struct { - HasQuota bool +type CheckQuotaResult struct { + // HasQuota is true if the quota has not been exceeded + HasQuota bool + + // SuggestedQueryInterval is the suggested interval to wait before querying the API again SuggestedQueryInterval time.Duration } type QuotaMonitor interface { // TeamName returns the team name TeamName() string - // HasQuota returns true if the quota has not been exceeded - HasQuota(context.Context) (QuotaCheckResult, error) + // CheckQuota checks if the quota has been exceeded + CheckQuota(context.Context) (CheckQuotaResult, error) } type UsageClient interface { @@ -365,20 +368,20 @@ func (u *BatchUpdater) TeamName() string { return u.teamName } -func (u *BatchUpdater) HasQuota(ctx context.Context) (QuotaCheckResult, error) { +func (u *BatchUpdater) CheckQuota(ctx context.Context) (CheckQuotaResult, error) { if u.awsMarketplaceClient != nil { - return QuotaCheckResult{HasQuota: true}, nil + return CheckQuotaResult{HasQuota: true}, nil } 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") usage, err := u.apiClient.GetTeamPluginUsageWithResponse(ctx, u.teamName, u.pluginMeta.Team, u.pluginMeta.Kind, u.pluginMeta.Name) if err != nil { - return QuotaCheckResult{HasQuota: false}, fmt.Errorf("failed to get usage: %w", err) + return CheckQuotaResult{HasQuota: false}, fmt.Errorf("failed to get usage: %w", err) } if usage.StatusCode() != http.StatusOK { - return QuotaCheckResult{HasQuota: false}, fmt.Errorf("failed to get usage: %s", usage.Status()) + return CheckQuotaResult{HasQuota: false}, fmt.Errorf("failed to get usage: %s", usage.Status()) } - res := QuotaCheckResult{ + res := CheckQuotaResult{ HasQuota: usage.JSON200.RemainingRows == nil || *usage.JSON200.RemainingRows > 0, } if usage.HTTPResponse == nil { @@ -715,8 +718,8 @@ func (n *NoOpUsageClient) TeamName() string { return n.TeamNameValue } -func (NoOpUsageClient) HasQuota(_ context.Context) (QuotaCheckResult, error) { - return QuotaCheckResult{HasQuota: true}, nil +func (NoOpUsageClient) CheckQuota(_ context.Context) (CheckQuotaResult, error) { + return CheckQuotaResult{HasQuota: true}, nil } func (NoOpUsageClient) Increase(_ uint32) error { diff --git a/premium/usage_test.go b/premium/usage_test.go index b2935a4a47..aa17e2f58b 100644 --- a/premium/usage_test.go +++ b/premium/usage_test.go @@ -113,10 +113,10 @@ func TestUsageService_HasQuota_NoRowsRemaining(t *testing.T) { usageClient := newClient(t, apiClient, WithBatchLimit(0)) - hasQuota, err := usageClient.HasQuota(ctx) + result, err := usageClient.CheckQuota(ctx) require.NoError(t, err) - assert.False(t, hasQuota.HasQuota, "should not have quota") + assert.False(t, result.HasQuota, "should not have quota") } func TestUsageService_HasQuota_WithRowsRemaining(t *testing.T) { @@ -130,10 +130,10 @@ func TestUsageService_HasQuota_WithRowsRemaining(t *testing.T) { usageClient := newClient(t, apiClient, WithBatchLimit(0)) - hasQuota, err := usageClient.HasQuota(ctx) + result, err := usageClient.CheckQuota(ctx) require.NoError(t, err) - assert.True(t, hasQuota.HasQuota, "should have quota") + assert.True(t, result.HasQuota, "should have quota") } func TestUsageService_Increase_ZeroBatchSize(t *testing.T) { From 75486639cd254debd1eddda2582f63c898b6b09b Mon Sep 17 00:00:00 2001 From: Kemal Hadimli Date: Tue, 29 Oct 2024 15:36:54 +0000 Subject: [PATCH 4/6] another CR rename --- premium/monitor_test.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/premium/monitor_test.go b/premium/monitor_test.go index 8cff1cdb23..a319e84550 100644 --- a/premium/monitor_test.go +++ b/premium/monitor_test.go @@ -10,8 +10,8 @@ import ( ) type quotaResponse struct { - hasQuota CheckQuotaResult - err error + result CheckQuotaResult + err error } func newFakeQuotaMonitor(hasQuota ...quotaResponse) *fakeQuotaMonitor { @@ -28,7 +28,7 @@ func (f *fakeQuotaMonitor) CheckQuota(_ context.Context) (CheckQuotaResult, erro if f.calls < len(f.responses)-1 { f.calls++ } - return resp.hasQuota, resp.err + return resp.result, resp.err } func (*fakeQuotaMonitor) TeamName() string { From 8be40547e9b3c4a52e55ea459dfbbeddfbbb1802 Mon Sep 17 00:00:00 2001 From: Kemal Hadimli Date: Tue, 29 Oct 2024 15:44:37 +0000 Subject: [PATCH 5/6] use `ticker.Reset` --- premium/monitor.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/premium/monitor.go b/premium/monitor.go index 97a70824c9..cf3deb99f6 100644 --- a/premium/monitor.go +++ b/premium/monitor.go @@ -99,8 +99,7 @@ func (qc *quotaChecker) startQuotaMonitor(ctx context.Context) context.Context { } if result.SuggestedQueryInterval > 0 && qc.duration != result.SuggestedQueryInterval { qc.duration = result.SuggestedQueryInterval - ticker.Stop() - ticker = time.NewTicker(qc.duration) + ticker.Reset(qc.duration) } consecutiveFailures = 0 hasQuotaErrors = nil From 818d9f4a983266bee5b051cf3d65c180ecb57e98 Mon Sep 17 00:00:00 2001 From: Kemal Hadimli Date: Tue, 29 Oct 2024 15:45:44 +0000 Subject: [PATCH 6/6] CheckQuota response handler: Use the same logic to protect/log against invalid interval values --- premium/usage.go | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/premium/usage.go b/premium/usage.go index d6eaea8007..d141c3ff93 100644 --- a/premium/usage.go +++ b/premium/usage.go @@ -388,10 +388,11 @@ func (u *BatchUpdater) CheckQuota(ctx context.Context) (CheckQuotaResult, error) return res, nil } if headerValue := usage.HTTPResponse.Header.Get(QueryIntervalHeader); headerValue != "" { - if interval, err := strconv.ParseUint(headerValue, 10, 32); err != nil { - u.logger.Warn().Err(err).Str(QueryIntervalHeader, headerValue).Msg("failed to parse query interval") - } else if interval > 0 { + interval, err := strconv.ParseUint(headerValue, 10, 32) + if interval > 0 { res.SuggestedQueryInterval = time.Duration(interval) * time.Second + } else { + u.logger.Warn().Err(err).Str(QueryIntervalHeader, headerValue).Msg("failed to parse query interval") } } return res, nil