Merge pull request #1 from PaulPowershell/feat/budgetv1

Fix retry with header

Author: kevindelmont

metrics_azurerm_costs.go

@@ -258,7 +258,7 @@ func (m *MetricsCollectorAzureRmCosts) collectCostManagementMetrics(logger *zap.
 		params.TimePeriod = &timePeriod
 	}
 
-	result, err := m.sendCostQuery(m.Context(), logger, scope, params, nil)
+	result, err := m.sendCostQuery(m.Context(), logger, scope, params)
 	if err != nil {
 		logger.Panic(err)
 	}
@@ -399,10 +399,10 @@ func (m *MetricsCollectorAzureRmCosts) collectCostManagementMetrics(logger *zap.
 	time.Sleep(Config.Collectors.Costs.RequestDelay)
 }
 
-func (m *MetricsCollectorAzureRmCosts) sendCostQuery(ctx context.Context, logger *zap.SugaredLogger, scope string, parameters armcostmanagement.QueryDefinition, options *armcostmanagement.QueryClientUsageOptions) (armcostmanagement.QueryClientUsageResponse, error) {
+func (m *MetricsCollectorAzureRmCosts) sendCostQuery(ctx context.Context, logger *zap.SugaredLogger, scope string, parameters armcostmanagement.QueryDefinition) (armcostmanagement.QueryClientUsageResponse, error) {
 	clientOpts := AzureClient.NewArmClientOptions()
 
-	// cost queries should not retry soo fast, we have a strict rate limit on azure side
+	// Initialize the client with appropriate retry options.
 	clientOpts.Retry = policy.RetryOptions{
 		MaxRetries:    3,
 		RetryDelay:    30 * time.Second,
@@ -420,7 +420,7 @@ func (m *MetricsCollectorAzureRmCosts) sendCostQuery(ctx context.Context, logger
 		logger.Panic(err.Error())
 	}
 
-	// paging
+	// Set up the pipeline for paging.
 	pl, err := armruntime.NewPipeline("azurerm-costs", gitTag, AzureClient.GetCred(), runtime.PipelineOptions{}, AzureClient.NewArmClientOptions())
 	if err != nil {
 		logger.Panic(err.Error())
@@ -445,6 +445,18 @@ func (m *MetricsCollectorAzureRmCosts) sendCostQuery(ctx context.Context, logger
 				}
 				defer resp.Body.Close()
 
+				if resp.StatusCode == http.StatusTooManyRequests {
+					retryAfterHeader := resp.Header.Get("X-Ms-Ratelimit-Microsoft.costmanagement-Entity-Retry-After")
+					retryAfter, err := strconv.Atoi(retryAfterHeader)
+					if err != nil {
+						logger.Errorf("Unable to parse retry-after header: %v", retryAfterHeader)
+						return fmt.Errorf("unable to parse retry-after header: %v", retryAfterHeader)
+					}
+					logger.Errorf("Received 429 Too Many Requests. Retrying after %d seconds. Headers: %v", retryAfter, resp.Header)
+					time.Sleep(time.Duration(retryAfter) * time.Second)
+					return fmt.Errorf("received 429 Too Many Requests, retrying after %d seconds", retryAfter)
+				}
+
 				if runtime.HasStatusCode(resp, http.StatusOK) {
 					pagerResult := armcostmanagement.QueryClientUsageResponse{}
 					if err := runtime.UnmarshalAsJSON(resp, &pagerResult); err == nil {
@@ -454,12 +466,16 @@ func (m *MetricsCollectorAzureRmCosts) sendCostQuery(ctx context.Context, logger
 						logger.Panic(err.Error())
 					}
 				} else {
-					return fmt.Errorf(`unexpected status code: %v`, resp.StatusCode)
+					return fmt.Errorf("unexpected status code: %v", resp.StatusCode)
 				}
 
 				return nil
 			}()
 			if err != nil {
+				// If we encounter a rate limit error, retry after the specified delay.
+				if strings.Contains(err.Error(), "received 429 Too Many Requests") {
+					continue
+				}
 				return result, err
 			}