feat: respect X-Ratelimit-Reset when retrying API requests (#434)

Author: jonathan-ostrander

internal/ld/ld.go

@@ -74,6 +74,32 @@ func IsTransient(err error) bool {
 	return !errors.As(err, &e)
 }
 
+// LaunchDarkly API uses the X-Ratelimit-Reset header to communicate when to retry after a 429
+// Fallback to default backoff if header can't be parsed
+// https://apidocs.launchdarkly.com/#section/Overview/Rate-limiting
+// Method is curried in order to avoid stubbing the time package and fallback Backoff in unit tests
+func RateLimitBackoff(now func() time.Time, fallbackBackoff h.Backoff) func(time.Duration, time.Duration, int, *http.Response) time.Duration {
+	return func(min, max time.Duration, attemptNum int, resp *http.Response) time.Duration {
+		if resp != nil {
+			if resp.StatusCode == http.StatusTooManyRequests {
+				if s, ok := resp.Header["X-Ratelimit-Reset"]; ok {
+					if sleepUntil, err := strconv.ParseInt(s[0], 10, 64); err == nil {
+						sleep := sleepUntil - now().UnixMilli()
+
+						if sleep > 0 {
+							return time.Millisecond * time.Duration(sleep)
+						} else {
+							return time.Duration(0)
+						}
+					}
+				}
+			}
+		}
+
+		return fallbackBackoff(min, max, attemptNum, resp)
+	}
+}
+
 func InitApiClient(options ApiOptions) ApiClient {
 	if options.BaseUri == "" {
 		options.BaseUri = "https://app.launchdarkly.com"
@@ -83,9 +109,7 @@ func InitApiClient(options ApiOptions) ApiClient {
 	if options.RetryMax != nil && *options.RetryMax >= 0 {
 		client.RetryMax = *options.RetryMax
 	}
-	client.RetryWaitMin = 1 * time.Second
-	client.RetryWaitMax = 30 * time.Second
-	client.Backoff = h.LinearJitterBackoff
+	client.Backoff = RateLimitBackoff(time.Now, h.LinearJitterBackoff)
 
 	return ApiClient{
 		ldClient: ldapi.NewAPIClient(&ldapi.Configuration{

internal/ld/ld_test.go

@@ -5,7 +5,9 @@ import (
 	"net/http/httptest"
 	"os"
 	"testing"
+	"time"
 
+	h "github.com/hashicorp/go-retryablehttp"
 	"github.com/stretchr/testify/require"
 
 	"github.com/launchdarkly/ld-find-code-refs/v2/internal/log"
@@ -234,3 +236,41 @@ func TestCountByProjectAndFlag(t *testing.T) {
 	require.Equal(t, count, want)
 
 }
+
+func TestRateLimitBackoff(t *testing.T) {
+	// Backoff instance where the time is always 0
+	backoff := RateLimitBackoff(func() time.Time { return time.Unix(0, 0) }, h.DefaultBackoff)
+
+	defaultBackoff := time.Second * time.Duration(1)
+
+	invalidRateLimitReset := "abc"
+	validRateLimitReset := "2000"
+	pastRateLimitReset := "-1000"
+	specs := []struct {
+		name           string
+		status         int
+		rateLimitReset *string
+		expected       time.Duration
+	}{
+		{"falls back to default backoff due to status", http.StatusBadGateway, nil, defaultBackoff},
+		{"falls back to default backoff due to missing header", http.StatusTooManyRequests, nil, defaultBackoff},
+		{"falls back to default backoff due to invalid header", http.StatusTooManyRequests, &invalidRateLimitReset, defaultBackoff},
+		{"returns difference between reset and current time", http.StatusTooManyRequests, &validRateLimitReset, time.Second * time.Duration(2)},
+		{"returns 0 because reset is in past", http.StatusTooManyRequests, &pastRateLimitReset, time.Duration(0)},
+	}
+	for _, tt := range specs {
+		t.Run(tt.name, func(t *testing.T) {
+			resp := &http.Response{
+				StatusCode: tt.status,
+				Header:     make(http.Header),
+			}
+
+			if tt.rateLimitReset != nil {
+				resp.Header.Set("X-Ratelimit-Reset", *tt.rateLimitReset)
+			}
+
+			actual := backoff(defaultBackoff, time.Second*time.Duration(10), 0, resp)
+			require.Equal(t, tt.expected, actual)
+		})
+	}
+}