feat: rate-limit konnect requests when receiving 429s

cmd/common_konnect.go

@@ -102,6 +102,7 @@ func getKongClientForKonnectMode(ctx context.Context) (*kong.Client, error) {
 		HTTPClient: httpClient,
 		Debug:      konnectConfig.Debug,
 		Headers:    konnectConfig.Headers,
+		Retryable:  true,
 	})
 }
 

go.mod

@@ -16,6 +16,7 @@ require (
 	github.com/google/go-querystring v1.1.0
 	github.com/google/uuid v1.3.0
 	github.com/hashicorp/go-memdb v1.3.3
+	github.com/hashicorp/go-retryablehttp v0.7.1
 	github.com/hexops/gotextdiff v1.0.3
 	github.com/imdario/mergo v0.3.12
 	github.com/kong/go-kong v0.29.0
@@ -47,6 +48,7 @@ require (
 	github.com/go-openapi/jsonpointer v0.19.5 // indirect
 	github.com/go-openapi/jsonreference v0.19.5 // indirect
 	github.com/go-openapi/swag v0.19.14 // indirect
+	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/hashicorp/go-immutable-radix v1.3.1 // indirect
 	github.com/hashicorp/golang-lru v0.5.4 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect

go.sum

@@ -180,11 +180,18 @@ github.com/googleapis/gnostic v0.5.1/go.mod h1:6U4PtQXGIEt/Z3h5MAT7FNofLnw9vXk2c
 github.com/googleapis/gnostic v0.5.5/go.mod h1:7+EbHbldMins07ALC74bsA81Ovc97DwqyJO1AENw9kA=
 github.com/googleapis/google-cloud-go-testing v0.0.0-20200911160855-bcd43fbb19e8/go.mod h1:dvDLG8qkwmyD9a/MJJN3XJcT3xFxOKAvTZGvuZmac9g=
 github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So=
+github.com/hashicorp/go-cleanhttp v0.5.1/go.mod h1:JpRdi6/HCYpAwUzNwuwqhbovhLtngrth3wmdIIUrZ80=
+github.com/hashicorp/go-cleanhttp v0.5.2 h1:035FKYIWjmULyFRBKPs8TBQoi0x6d9G4xc9neXJWAZQ=
+github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48=
+github.com/hashicorp/go-hclog v0.9.2/go.mod h1:5CU+agLiy3J7N7QjHK5d05KxGsuXiQLrjA0H7acj2lQ=
+github.com/hashicorp/go-hclog v1.2.0 h1:La19f8d7WIlm4ogzNHB0JGqs5AUDAZ2UfCY4sJXcJdM=
 github.com/hashicorp/go-immutable-radix v1.3.0/go.mod h1:0y9vanUI8NX6FsYoO3zeMjhV/C5i9g4Q3DwcSNZ4P60=
 github.com/hashicorp/go-immutable-radix v1.3.1 h1:DKHmCUm2hRBK510BaiZlwvpD40f8bJFeZnpfm2KLowc=
 github.com/hashicorp/go-immutable-radix v1.3.1/go.mod h1:0y9vanUI8NX6FsYoO3zeMjhV/C5i9g4Q3DwcSNZ4P60=
 github.com/hashicorp/go-memdb v1.3.3 h1:oGfEWrFuxtIUF3W2q/Jzt6G85TrMk9ey6XfYLvVe1Wo=
 github.com/hashicorp/go-memdb v1.3.3/go.mod h1:uBTr1oQbtuMgd1SSGoR8YV27eT3sBHbYiNm53bMpgSg=
+github.com/hashicorp/go-retryablehttp v0.7.1 h1:sUiuQAnLlbvmExtFQs72iFW/HXeUn8Z1aJLQ4LJJbTQ=
+github.com/hashicorp/go-retryablehttp v0.7.1/go.mod h1:vAew36LZh98gCBJNLH42IQ1ER/9wtLZZ8meHqQvEYWY=
 github.com/hashicorp/go-uuid v1.0.0 h1:RS8zrF7PhGwyNPOtxSClXXj9HA8feRnJzgnI1RJCSnM=
 github.com/hashicorp/go-uuid v1.0.0/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro=
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=

konnect/login_service.go

@@ -96,6 +96,8 @@ func (s *AuthService) UserInfo(ctx context.Context) (*UserInfo, error) {
 	}
 
 	info := &UserInfo{}
+	// info := map[string]string{}
+	//fmt.Println("EMEM ", s.client.client.Jar)
 	_, err = s.client.Do(ctx, req, info)
 	if err != nil {
 		return nil, err

utils/types.go

@@ -1,9 +1,11 @@
 package utils
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
+	"math"
 	"net"
 	"net/http"
 	"net/url"
@@ -13,6 +15,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/hashicorp/go-retryablehttp"
 	"github.com/kong/deck/konnect"
 	"github.com/kong/go-kong/kong"
 	"github.com/kong/go-kong/kong/custom"
@@ -100,6 +103,9 @@ type KongClientConfig struct {
 	TLSClientCert string
 
 	TLSClientKey string
+
+	// whether or not the client should retry on 429s
+	Retryable bool
 }
 
 type KonnectConfig struct {
@@ -119,6 +125,75 @@ func (kc *KongClientConfig) ForWorkspace(name string) KongClientConfig {
 	return result
 }
 
+// backoffStrategy provides a callback for Client.Backoff which
+// will perform exponential backoff based on the attempt number and limited
+// by the provided minimum and maximum durations.
+//
+// It also tries to parse Retry-After response header when a http.StatusTooManyRequests
+// (HTTP Code 429) is found in the resp parameter. Hence it will return the number of
+// seconds the server states it may be ready to process more requests from this client.
+//
+// This is the same as DefaultBackoff (https://github.com/hashicorp/go-retryablehttp/blob/master/client.go#L510)
+// except that here we are only retrying on 429s.
+func backoffStrategy(min, max time.Duration, attemptNum int, resp *http.Response) time.Duration {
+	const (
+		base            = 10
+		bitSize         = 64
+		baseExponential = 2
+	)
+	if resp != nil {
+		if resp.StatusCode == http.StatusTooManyRequests {
+			if s, ok := resp.Header["Retry-After"]; ok {
+				if sleep, err := strconv.ParseInt(s[0], base, bitSize); err == nil {
+					return time.Second * time.Duration(sleep)
+				}
+			}
+		}
+	}
+
+	mult := math.Pow(baseExponential, float64(attemptNum)) * float64(min)
+	sleep := time.Duration(mult)
+	if float64(sleep) != mult || sleep > max {
+		sleep = max
+	}
+	return sleep
+}
+
+// retryPolicy provides a callback for Client.CheckRetry, which
+// will retry on 429s errors.
+func retryPolicy(ctx context.Context, resp *http.Response, err error) (bool, error) {
+	// do not retry on context.Canceled or context.DeadlineExceeded
+	if ctx.Err() != nil {
+		return false, ctx.Err()
+	}
+
+	// 429 Too Many Requests is recoverable. Sometimes the server puts
+	// a Retry-After response header to indicate when the server is
+	// available to start processing request from client.
+	if resp.StatusCode == http.StatusTooManyRequests {
+		return true, nil
+	}
+	return false, nil
+}
+
+func getRetryableClient(client *http.Client) *http.Client {
+	const (
+		minRetryWait = 10 * time.Second
+		maxRetryWait = 60 * time.Second
+		retryMax     = 10
+	)
+	retryClient := retryablehttp.NewClient()
+	retryClient.HTTPClient = client
+	retryClient.Backoff = backoffStrategy
+	retryClient.CheckRetry = retryPolicy
+	retryClient.RetryMax = retryMax
+	retryClient.RetryWaitMax = maxRetryWait
+	retryClient.RetryWaitMin = minRetryWait
+	// logging is handled by deck.
+	retryClient.Logger = nil
+	return retryClient.StandardClient()
+}
+
 // GetKongClient returns a Kong client
 func GetKongClient(opt KongClientConfig) (*kong.Client, error) {
 	var tlsConfig tls.Config
@@ -163,6 +238,10 @@ func GetKongClient(opt KongClientConfig) (*kong.Client, error) {
 	}
 	c = kong.HTTPClientWithHeaders(c, headers)
 
+	if opt.Retryable {
+		c = getRetryableClient(c)
+	}
+
 	url, err := url.ParseRequestURI(address)
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse kong address: %w", err)