Add retry logic when rate limit (#1896)

Adding enable retry in the API client and tests

---------

Co-authored-by: ci.datadog-api-spec <packages@datadoghq.com>
Co-authored-by: Thomas Hervé <thomas.herve@datadoghq.com>

.generator/src/generator/templates/client.j2

@@ -12,6 +12,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"math"
 	"mime/multipart"
 	"net/http"
 	"net/http/httputil"
@@ -29,8 +30,9 @@ import (
 )
 
 var (
-	jsonCheck = regexp.MustCompile(`(?i:(?:application|text)/(?:vnd\.[^;]+\+)?json)`)
-	xmlCheck  = regexp.MustCompile(`(?i:(?:application|text)/xml)`)
+	jsonCheck            = regexp.MustCompile(`(?i:(?:application|text)/(?:vnd\.[^;]+\+)?json)`)
+	xmlCheck             = regexp.MustCompile(`(?i:(?:application|text)/xml)`)
+	rateLimitResetHeader = "X-Ratelimit-Reset"
 )
 
 // APIClient manages communication with the {{ openapi.info.title }} API v{{ openapi.info.version}}.
@@ -115,35 +117,91 @@ func ParameterToString(obj interface{}, collectionFormat string) string {
 {# The method is used in Terraform client and needs to be public. -#}
 // CallAPI do the request.
 func (c *APIClient) CallAPI(request *http.Request) (*http.Response, error) {
-	if c.Cfg.Debug {
-		dump, err := httputil.DumpRequestOut(request, true)
-		if err != nil {
-			return nil, err
+	var rawBody []byte
+	if request.Body != nil && request.Body != http.NoBody {
+		rawBody, _ = io.ReadAll(request.Body)
+		request.Body.Close()
+	}
+	ctx, ccancel := context.WithTimeout(request.Context(), c.Cfg.RetryConfiguration.HTTPRetryTimeout)
+	defer ccancel()
+	retryCount := 0
+	for {
+		newRequest := copyRequest(request, &rawBody)
+		if c.Cfg.Debug {
+			dump, err := httputil.DumpRequestOut(newRequest, true)
+			if err != nil {
+				return nil, err
+			}
+			// Strip any api keys from the response being logged
+			keys, ok := newRequest.Context().Value(ContextAPIKeys).(map[string]APIKey)
+			if keys != nil && ok {
+				for _, apiKey := range keys {
+					valueRegex := regexp.MustCompile(fmt.Sprintf("(?m)%s", apiKey.Key))
+					dump = valueRegex.ReplaceAll(dump, []byte("REDACTED"))
+				}
+			}
+			log.Printf("\n%s\n", string(dump))
+		}
+		resp, requestErr := c.Cfg.HTTPClient.Do(newRequest)
+
+		if requestErr != nil {
+			return resp, requestErr
 		}
-		// Strip any api keys from the response being logged
-		keys, ok := request.Context().Value(ContextAPIKeys).(map[string]APIKey)
-		if keys != nil && ok {
-			for _, apiKey := range keys {
-				valueRegex := regexp.MustCompile(fmt.Sprintf("(?m)%s", apiKey.Key))
-				dump = valueRegex.ReplaceAll(dump, []byte("REDACTED"))
+
+		if c.Cfg.Debug {
+			dump, _ := httputil.DumpResponse(resp, true)
+			if c.Cfg.RetryConfiguration.EnableRetry {
+				log.Println("Max retries:", c.Cfg.RetryConfiguration.MaxRetries, " Current retry:", retryCount)
+				if retryCount == c.Cfg.RetryConfiguration.MaxRetries {
+					log.Println("Max retries reached")
+				}
 			}
+			log.Printf("\n%s\n", string(dump))
+		}
+
+		retryDuration, shouldRetry := c.shouldRetryRequest(resp, retryCount)
+		if !shouldRetry {
+			return resp, requestErr
+		}
+
+		select {
+		case <-ctx.Done():
+			return resp, requestErr
+		case <-time.After(*retryDuration):
+			retryCount++
+			continue
 		}
-		log.Printf("\n%s\n", string(dump))
-	}
 
-	resp, err := c.Cfg.HTTPClient.Do(request)
-	if err != nil {
-		return resp, err
 	}
+}
 
-	if c.Cfg.Debug {
-		dump, err := httputil.DumpResponse(resp, true)
-		if err != nil {
-			return resp, err
+// Determine if a request should be retried 
+func (c *APIClient) shouldRetryRequest(response *http.Response, retryCount int) (*time.Duration, bool) {
+	enableRetry := c.Cfg.RetryConfiguration.EnableRetry
+	maxRetries := c.Cfg.RetryConfiguration.MaxRetries
+	if !enableRetry || retryCount == maxRetries {
+		return nil, false
+	}
+	var err error
+	if v := response.Header.Get(rateLimitResetHeader); response.StatusCode == 429 && v != "" {
+		vInt, err := strconv.ParseInt(v, 10, 64)
+		if err == nil {
+			retryDuration := time.Duration(vInt) * time.Second
+			return &retryDuration, true
 		}
-		log.Printf("\n%s\n", string(dump))
 	}
-	return resp, err
+
+	// Calculate retry for 5xx errors or if unable to parse value of rateLimitResetHeader
+	// or if the `rateLimitResetHeader` header is missing or if status code >= 500.
+	if err != nil || response.StatusCode == 429 || response.StatusCode >= 500 {
+		// Calculate the retry val (base * multiplier^retryCount)
+		retryVal := c.Cfg.RetryConfiguration.BackOffBase * math.Pow(c.Cfg.RetryConfiguration.BackOffMultiplier, float64(retryCount))
+		// retry duration shouldn't exceed default timeout period
+		retryVal = math.Min(float64(c.Cfg.HTTPClient.Timeout/time.Second), retryVal)
+		retryDuration := time.Duration(retryVal) * time.Second
+		return &retryDuration, true
+	}
+	return nil, false
 }
 
 // GetConfig allows modification of underlying config for alternate implementations and testing.

.generator/src/generator/templates/configuration.j2

@@ -9,6 +9,7 @@ import (
 	"os"
 	"runtime"
 	"strings"
+	"time"
 
 	client "{{ module }}"
 )
@@ -100,6 +101,16 @@ type Configuration struct {
 {#withCustomMiddlewareFunction
 	Middleware         MiddlewareFunction
 #}	unstableOperations map[string]bool
+	RetryConfiguration RetryConfiguration
+}
+
+// RetryConfiguration stores the configuration of the retry behavior of the api client
+type RetryConfiguration struct {
+	EnableRetry       bool
+	BackOffMultiplier float64
+	BackOffBase       float64
+	HTTPRetryTimeout  time.Duration
+	MaxRetries        int
 }
 
 {%- macro server_configuration(server) -%}
@@ -165,6 +176,13 @@ func NewConfiguration() *Configuration {
 		{%- endfor %}
         {%- endfor %}
 		},
+		RetryConfiguration: RetryConfiguration{
+			EnableRetry:       false,
+			BackOffMultiplier: 2,
+			BackOffBase:       2,
+			HTTPRetryTimeout:  60 * time.Second,
+			MaxRetries:        3,
+		},
 	}
 	return cfg
 }

.generator/src/generator/templates/utils.j2

@@ -2,7 +2,10 @@
 package {{ common_package_name }}
 
 import (
+	"bytes"
 	"encoding/json"
+	"io"
+	"net/http"
 	"reflect"
 	"strings"
 	"time"
@@ -443,3 +446,14 @@ func ContainsUnparsedObject(i interface{}) (bool, interface{}) {
 func Strlen(s string) int {
 	return utf8.RuneCountInString(s)
 }
+
+// Copy the original request so it doesn't get lost when retrying
+func copyRequest(r *http.Request, rawBody *[]byte) *http.Request {
+	newRequest := *r
+
+	if r.Body == nil || r.Body == http.NoBody {
+		return &newRequest
+	}
+	newRequest.Body = io.NopCloser(bytes.NewBuffer(*rawBody))
+	return &newRequest
+}

README.md

@@ -124,6 +124,20 @@ If you want to enable requests logging, set the `debug` flag on your configurati
     configuration.Debug = true
 ```
 
+### Enable retry
+
+If you want to enable retry when getting status code `429` rate-limited, set `EnableRetry` to `true`
+
+```go
+    configuration.RetryConfiguration.EnableRetry = true
+```
+
+The default max retry is `3`, you can change it with `MaxRetries`
+
+```go
+    configuration.RetryConfiguration.MaxRetries = 3
+```
+
 ### Configure proxy
 
 If you want to configure proxy, set env var `HTTP_PROXY`, and `HTTPS_PROXY` or set custom

api/datadog/client.go

@@ -15,6 +15,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"math"
 	"mime/multipart"
 	"net/http"
 	"net/http/httputil"
@@ -23,15 +24,17 @@ import (
 	"path/filepath"
 	"reflect"
 	"regexp"
+	"strconv"
 	"strings"
 	"time"
 
 	"golang.org/x/oauth2"
 )
 
 var (
-	jsonCheck = regexp.MustCompile(`(?i:(?:application|text)/(?:vnd\.[^;]+\+)?json)`)
-	xmlCheck  = regexp.MustCompile(`(?i:(?:application|text)/xml)`)
+	jsonCheck            = regexp.MustCompile(`(?i:(?:application|text)/(?:vnd\.[^;]+\+)?json)`)
+	xmlCheck             = regexp.MustCompile(`(?i:(?:application|text)/xml)`)
+	rateLimitResetHeader = "X-Ratelimit-Reset"
 )
 
 // APIClient manages communication with the Datadog API V2 Collection API v1.0.
@@ -115,35 +118,91 @@ func ParameterToString(obj interface{}, collectionFormat string) string {
 
 // CallAPI do the request.
 func (c *APIClient) CallAPI(request *http.Request) (*http.Response, error) {
-	if c.Cfg.Debug {
-		dump, err := httputil.DumpRequestOut(request, true)
-		if err != nil {
-			return nil, err
+	var rawBody []byte
+	if request.Body != nil && request.Body != http.NoBody {
+		rawBody, _ = io.ReadAll(request.Body)
+		request.Body.Close()
+	}
+	ctx, ccancel := context.WithTimeout(request.Context(), c.Cfg.RetryConfiguration.HTTPRetryTimeout)
+	defer ccancel()
+	retryCount := 0
+	for {
+		newRequest := copyRequest(request, &rawBody)
+		if c.Cfg.Debug {
+			dump, err := httputil.DumpRequestOut(newRequest, true)
+			if err != nil {
+				return nil, err
+			}
+			// Strip any api keys from the response being logged
+			keys, ok := newRequest.Context().Value(ContextAPIKeys).(map[string]APIKey)
+			if keys != nil && ok {
+				for _, apiKey := range keys {
+					valueRegex := regexp.MustCompile(fmt.Sprintf("(?m)%s", apiKey.Key))
+					dump = valueRegex.ReplaceAll(dump, []byte("REDACTED"))
+				}
+			}
+			log.Printf("\n%s\n", string(dump))
+		}
+		resp, requestErr := c.Cfg.HTTPClient.Do(newRequest)
+
+		if requestErr != nil {
+			return resp, requestErr
 		}
-		// Strip any api keys from the response being logged
-		keys, ok := request.Context().Value(ContextAPIKeys).(map[string]APIKey)
-		if keys != nil && ok {
-			for _, apiKey := range keys {
-				valueRegex := regexp.MustCompile(fmt.Sprintf("(?m)%s", apiKey.Key))
-				dump = valueRegex.ReplaceAll(dump, []byte("REDACTED"))
+
+		if c.Cfg.Debug {
+			dump, _ := httputil.DumpResponse(resp, true)
+			if c.Cfg.RetryConfiguration.EnableRetry {
+				log.Println("Max retries:", c.Cfg.RetryConfiguration.MaxRetries, " Current retry:", retryCount)
+				if retryCount == c.Cfg.RetryConfiguration.MaxRetries {
+					log.Println("Max retries reached")
+				}
 			}
+			log.Printf("\n%s\n", string(dump))
+		}
+
+		retryDuration, shouldRetry := c.shouldRetryRequest(resp, retryCount)
+		if !shouldRetry {
+			return resp, requestErr
+		}
+
+		select {
+		case <-ctx.Done():
+			return resp, requestErr
+		case <-time.After(*retryDuration):
+			retryCount++
+			continue
 		}
-		log.Printf("\n%s\n", string(dump))
-	}
 
-	resp, err := c.Cfg.HTTPClient.Do(request)
-	if err != nil {
-		return resp, err
 	}
+}
 
-	if c.Cfg.Debug {
-		dump, err := httputil.DumpResponse(resp, true)
-		if err != nil {
-			return resp, err
+// Determine if a request should be retried
+func (c *APIClient) shouldRetryRequest(response *http.Response, retryCount int) (*time.Duration, bool) {
+	enableRetry := c.Cfg.RetryConfiguration.EnableRetry
+	maxRetries := c.Cfg.RetryConfiguration.MaxRetries
+	if !enableRetry || retryCount == maxRetries {
+		return nil, false
+	}
+	var err error
+	if v := response.Header.Get(rateLimitResetHeader); response.StatusCode == 429 && v != "" {
+		vInt, err := strconv.ParseInt(v, 10, 64)
+		if err == nil {
+			retryDuration := time.Duration(vInt) * time.Second
+			return &retryDuration, true
 		}
-		log.Printf("\n%s\n", string(dump))
 	}
-	return resp, err
+
+	// Calculate retry for 5xx errors or if unable to parse value of rateLimitResetHeader
+	// or if the `rateLimitResetHeader` header is missing or if status code >= 500.
+	if err != nil || response.StatusCode == 429 || response.StatusCode >= 500 {
+		// Calculate the retry val (base * multiplier^retryCount)
+		retryVal := c.Cfg.RetryConfiguration.BackOffBase * math.Pow(c.Cfg.RetryConfiguration.BackOffMultiplier, float64(retryCount))
+		// retry duration shouldn't exceed default timeout period
+		retryVal = math.Min(float64(c.Cfg.HTTPClient.Timeout/time.Second), retryVal)
+		retryDuration := time.Duration(retryVal) * time.Second
+		return &retryDuration, true
+	}
+	return nil, false
 }
 
 // GetConfig allows modification of underlying config for alternate implementations and testing.

api/datadog/configuration.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"runtime"
 	"strings"
+	"time"
 
 	client "github.com/DataDog/datadog-api-client-go/v2"
 )
@@ -96,6 +97,16 @@ type Configuration struct {
 	OperationServers   map[string]ServerConfigurations
 	HTTPClient         *http.Client
 	unstableOperations map[string]bool
+	RetryConfiguration RetryConfiguration
+}
+
+// RetryConfiguration stores the configuration of the retry behavior of the api client
+type RetryConfiguration struct {
+	EnableRetry       bool
+	BackOffMultiplier float64
+	BackOffBase       float64
+	HTTPRetryTimeout  time.Duration
+	MaxRetries        int
 }
 
 // NewConfiguration returns a new Configuration object.
@@ -387,6 +398,13 @@ func NewConfiguration() *Configuration {
 			"v2.ListIncidentTeams":         false,
 			"v2.UpdateIncidentTeam":        false,
 		},
+		RetryConfiguration: RetryConfiguration{
+			EnableRetry:       false,
+			BackOffMultiplier: 2,
+			BackOffBase:       2,
+			HTTPRetryTimeout:  60 * time.Second,
+			MaxRetries:        3,
+		},
 	}
 	return cfg
 }