Fix ratelimits (#28)

Author: vstpme

client.go

@@ -1,23 +1,17 @@
 package aapi
 
 import (
-	"context"
 	"encoding/json"
 	"fmt"
 	"io"
 	"io/ioutil"
-	"log"
-	"math/rand"
 	"net/http"
 	"net/url"
 	"sort"
-	"strconv"
 	"strings"
-	"time"
 
 	"github.com/google/go-querystring/query"
 	"github.com/hashicorp/go-retryablehttp"
-	"golang.org/x/time/rate"
 )
 
 const (
@@ -37,13 +31,11 @@ type PaginatedResponse struct {
 
 type Client struct {
 	// HTTP client used to communicate with the API.
-	client         *retryablehttp.Client
-	token          string
-	baseURL        *url.URL
-	grafanaURL     *url.URL
-	disableRetries bool
-	limiter        *rate.Limiter
-	UserAgent      string
+	client     *retryablehttp.Client
+	token      string
+	baseURL    *url.URL
+	grafanaURL *url.URL
+	UserAgent  string
 	// List of Services. Keep in sync with func newClient
 	Alerts                *AlertService
 	Integrations          *IntegrationService
@@ -87,18 +79,8 @@ func New(base_url, token string) (*Client, error) {
 func newClient(url, grafana_url string) (*Client, error) {
 	c := &Client{}
 
-	// Configure the HTTP client.
-	c.client = &retryablehttp.Client{
-		Backoff:      c.retryHTTPBackoff,
-		CheckRetry:   c.retryHTTPCheck,
-		RetryWaitMin: 100 * time.Millisecond,
-		RetryWaitMax: 400 * time.Millisecond,
-		RetryMax:     5,
-	}
-	// https://grafana.com/docs/grafana-cloud/oncall/oncall-api-reference/#rate-limits
-	baseLimit := 50.0 / 60
-	limit := rate.Limit(baseLimit)
-	c.limiter = rate.NewLimiter(limit, 50)
+	// retryablehttp.Client will retry up to 4 times on recoverable errors (429, 5xx, and low-level network errors)
+	c.client = retryablehttp.NewClient()
 
 	// Set the default base URL. _ suppress error handling
 	err := c.setBaseURL(url)
@@ -214,12 +196,6 @@ func (c *Client) NewRequest(method, path string, opt interface{}) (*retryablehtt
 // JSON decoded and stored in the value pointed to by v, or returned as an
 // error if an API error has occurred.
 func (c *Client) Do(req *retryablehttp.Request, v interface{}) (*http.Response, error) {
-	err := c.limiter.Wait(req.Context())
-	if err != nil {
-		log.Println("limiter")
-		return nil, err
-	}
-
 	resp, err := c.client.Do(req)
 	if err != nil {
 		return nil, err
@@ -307,43 +283,6 @@ func (e *ErrorResponse) Error() string {
 	return fmt.Sprintf("%s %s: %d %s", e.Response.Request.Method, u, e.Response.StatusCode, e.Message)
 }
 
-func (c *Client) retryHTTPCheck(ctx context.Context, resp *http.Response, err error) (bool, error) {
-	if ctx.Err() != nil {
-		return false, ctx.Err()
-	}
-	if err != nil {
-		return false, err
-	}
-	if !c.disableRetries && (resp.StatusCode == 429 || resp.StatusCode >= 500) {
-		return true, nil
-	}
-	return false, nil
-}
-
-func (c *Client) retryHTTPBackoff(min, max time.Duration, attemptNum int, resp *http.Response) time.Duration {
-	if resp != nil && resp.StatusCode == 429 {
-		return rateLimitBackoff(min, max, attemptNum, resp)
-	}
-
-	return retryablehttp.LinearJitterBackoff(min, max, attemptNum, resp)
-}
-
-func rateLimitBackoff(min, max time.Duration, attemptNum int, resp *http.Response) time.Duration {
-	rnd := rand.New(rand.NewSource(time.Now().UnixNano()))
-	jitter := time.Duration(rnd.Float64() * float64(max-min))
-	log.Printf("[DEBUG] ratelimited call")
-	if resp != nil {
-		if v := resp.Header.Get("RateLimit-Reset"); v != "" {
-			if reset, _ := strconv.ParseInt(v, 10, 64); reset > 0 {
-				log.Printf("[DEBUG] reset in '%d", reset)
-				min = time.Duration(reset) * time.Second
-			}
-		}
-	}
-
-	return min + jitter
-}
-
 func (c *Client) BaseURL() *url.URL {
 	u := *c.baseURL
 	return &u

client_test.go

@@ -157,3 +157,30 @@ func TestCheckResponse(t *testing.T) {
 		t.Errorf("Expected error: %s, got %s", want, errResp.Error())
 	}
 }
+
+func TestRatelimitRetry(t *testing.T) {
+	mux, server, client := setup(t)
+	defer teardown(server)
+
+	requestCount := 0
+	mux.HandleFunc("/api/v1/test", func(w http.ResponseWriter, r *http.Request) {
+		requestCount++
+		w.Header().Set("Retry-After", "0")
+		w.WriteHeader(http.StatusTooManyRequests)
+	})
+
+	req, err := client.NewRequest("GET", "test", nil)
+	if err != nil {
+		t.Fatalf("Failed to create request: %v", err)
+	}
+
+	_, err = client.Do(req, nil)
+	expectedErr := fmt.Sprintf("GET %s giving up after 5 attempt(s)", req.URL)
+	if err.Error() != expectedErr {
+		t.Fatalf("Expected error: %s, got %s", expectedErr, err.Error())
+	}
+
+	if requestCount != 5 {
+		t.Errorf("Expected 5 requests (1 original + 4 retries), got %d", requestCount)
+	}
+}