Merge pull request #213 from JupiterOne/APP-15634-implement-backoff

APP-15634: Implement retry logic with backoff when we receive a 429

.gitignore

@@ -49,4 +49,4 @@ schema.json
 schema.graphql
 
 # Terrform lock files from testing
-.terraform.lock.hcl
+.terraform*

jupiterone/internal/client/client.go

@@ -3,9 +3,17 @@ package client
 //go:generate go run github.com/Khan/genqlient
 
 import (
+	"bytes"
 	"context"
+	"fmt"
+	"io"
+	"math"
+	"math/rand"
 	"net/http"
 	"os"
+	"strconv"
+	"sync"
+	"time"
 
 	"github.com/Khan/genqlient/graphql"
 	genql "github.com/Khan/genqlient/graphql"
@@ -15,6 +23,20 @@ import (
 )
 
 const DefaultRegion string = "us"
+const MaxBackoff = 60 * time.Second
+const MinBackoff = 15 * time.Second
+const PowerOfTwo = 2
+
+var (
+	lastNon429Response time.Time
+	timestampMutex     sync.Mutex
+)
+
+func updateLastNon429Response() {
+	timestampMutex.Lock()
+	defer timestampMutex.Unlock()
+	lastNon429Response = time.Now()
+}
 
 type JupiterOneClientConfig struct {
 	APIKey    string
@@ -37,6 +59,73 @@ func (t *jupiterOneTransport) RoundTrip(req *http.Request) (*http.Response, erro
 	return t.base.RoundTrip(req)
 }
 
+// RetryTransport is a custom RoundTripper that adds retry logic with backoff.
+type RetryTransport struct {
+	Transport  http.RoundTripper
+	MinBackoff time.Duration
+	MaxBackoff time.Duration
+}
+
+func (rt *RetryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	var resp *http.Response
+	var err error
+
+	ctx := req.Context()
+
+	// We need to keep a copy of the body because each request it gets consumed
+	var bodyBytes []byte
+	if req.Body != nil {
+		bodyBytes, err = io.ReadAll(req.Body)
+		if err != nil {
+			return nil, fmt.Errorf("failed to read request body: %v", err)
+		}
+	}
+
+	for i := 0; i >= 0; i++ {
+		// Setting the body for the request
+		req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
+
+		resp, _ = rt.Transport.RoundTrip(req)
+
+		if resp.StatusCode != http.StatusTooManyRequests {
+			updateLastNon429Response()
+			return resp, nil
+		}
+
+		// If this is not the first try, and the lastNon429Response
+		// was more than 90 seconds ago, we should break out of the loop
+		// and return the last response.
+		timestampMutex.Lock()
+		if i > 0 && time.Since(lastNon429Response) > 90*time.Second {
+			timestampMutex.Unlock()
+			tflog.Warn(ctx, "Not going to retry, we haven't got a non 429 in a while")
+			return resp, nil
+		}
+		timestampMutex.Unlock()
+
+		tflog.Debug(ctx, "Retrying after getting a 429 response")
+
+		// Calculate the backoff time using exponential backoff with jitter.
+		backoff := rt.MinBackoff * time.Duration(math.Pow(PowerOfTwo, float64(i)))
+		jitter := time.Duration(rand.Int63n(int64(rt.MinBackoff)))
+		sleepDuration := backoff + jitter
+
+		// Ensure we do not exceed the maximum backoff time.
+		if sleepDuration > rt.MaxBackoff {
+			sleepDuration = rt.MaxBackoff
+		}
+
+		// Convert duration to seconds
+		var backoffSeconds = int(sleepDuration.Seconds())
+
+		tflog.Debug(ctx, "Backoff info", map[string]interface{}{"sleepDurationSeconds": backoffSeconds, "retryCount": i})
+
+		time.Sleep(sleepDuration)
+	}
+
+	return resp, fmt.Errorf("after many attempts, last status: %s", strconv.Itoa(resp.StatusCode))
+}
+
 func (c *JupiterOneClientConfig) getRegion(ctx context.Context) string {
 	region := c.Region
 
@@ -76,6 +165,12 @@ func (c *JupiterOneClientConfig) Qlient(ctx context.Context) graphql.Client {
 	httpClient.Transport = &jupiterOneTransport{apiKey: c.APIKey, accountID: c.AccountID, base: httpClient.Transport}
 	httpClient.Transport = logging.NewLoggingHTTPTransport(httpClient.Transport)
 
+	httpClient.Transport = &RetryTransport{
+		Transport:  httpClient.Transport,
+		MinBackoff: MinBackoff,
+		MaxBackoff: MaxBackoff,
+	}
+
 	client := genql.NewClient(endpoint, httpClient)
 
 	return client