Merge pull request #84 from hortonworks/CB-17052

CB-17052 Rate limit aws api operations

aws/aws.go

@@ -12,6 +12,8 @@ import (
 	"sync"
 	"time"
 
+	"golang.org/x/time/rate"
+
 	"github.com/aws/aws-sdk-go/service/cloudformation"
 	elb "github.com/aws/aws-sdk-go/service/elbv2"
 	"github.com/hortonworks/cloud-haunter/utils"
@@ -46,6 +48,8 @@ const (
 
 var provider = awsProvider{}
 
+var rateLimiter = rate.NewLimiter(rate.Every(time.Duration(ctx.AwsApiOperationRateLimitIntervalInSeconds)*time.Second), ctx.AwsApiOperationRateLimit)
+
 type awsProvider struct {
 	ec2Clients           map[string]*ec2.EC2
 	autoScalingClients   map[string]*autoscaling.AutoScaling
@@ -58,6 +62,11 @@ type awsProvider struct {
 	govCloud             bool
 }
 
+type ThrottledTransport struct {
+	roundTripperWrap http.RoundTripper
+	ratelimiter      *rate.Limiter
+}
+
 func init() {
 	accessKeyId := os.Getenv("AWS_ACCESS_KEY_ID")
 	if len(accessKeyId) == 0 {
@@ -1834,11 +1843,26 @@ func newElbClient(region string) (*elb.ELBV2, error) {
 	return elb.New(awsSession), nil
 }
 
+func NewThrottledTransport(rateLimiter *rate.Limiter, transportWrap http.RoundTripper) http.RoundTripper {
+	return &ThrottledTransport{
+		roundTripperWrap: transportWrap,
+		ratelimiter:      rateLimiter,
+	}
+}
+
+func (c *ThrottledTransport) RoundTrip(r *http.Request) (*http.Response, error) {
+	err := c.ratelimiter.Wait(r.Context())
+	if err != nil {
+		return nil, err
+	}
+	return c.roundTripperWrap.RoundTrip(r)
+}
+
 func newSession(configure func(*aws.Config)) (*session.Session, error) {
 	httpClient := &http.Client{
-		Transport: &http.Transport{
+		Transport: NewThrottledTransport(rateLimiter, &http.Transport{
 			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
-		}}
+		})}
 	config := aws.Config{HTTPClient: httpClient}
 	if configure != nil {
 		configure(&config)

context/context.go

@@ -24,6 +24,12 @@ var (
 	AzureCreationTimeLabel = "creation-timestamp,cb-creation-timestamp,cdp-creation-timestamp"
 
 	AwsBulkOperationSize = 50
+
+	// AWS rate limit interval in seconds for each API calls regardless the region
+	AwsApiOperationRateLimitIntervalInSeconds = 30
+
+	// AWS rate limit for each API calls regardless the region
+	AwsApiOperationRateLimit = 1000
 )
 
 // DryRun is a global flag to skip concrete action

go.mod

@@ -43,6 +43,7 @@ require (
 	golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1 // indirect
 	golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1 // indirect
 	golang.org/x/text v0.3.3 // indirect
+	golang.org/x/time v0.0.0-20220411224347-583f2d630306
 	google.golang.org/appengine v1.1.0 // indirect
 	gopkg.in/airbrake/gobrake.v2 v2.0.9 // indirect
 	gopkg.in/gemnasium/logrus-airbrake-hook.v2 v2.1.2 // indirect

go.sum

@@ -96,6 +96,8 @@ golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
 golang.org/x/text v0.3.3 h1:cokOdA+Jmi5PJGXLlLllQSgYigAEfHXJAERHVMaCc2k=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/time v0.0.0-20220411224347-583f2d630306 h1:+gHMid33q6pen7kv9xvT+JRinntgeXO2AeZVd0AWD3w=
+golang.org/x/time v0.0.0-20220411224347-583f2d630306/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 google.golang.org/api v0.0.0-20180717000714-0025a57598c0 h1:H82u+VsvfY63/uydrxjj2DpwjFvroE43qWaIzpa2Xik=
 google.golang.org/api v0.0.0-20180717000714-0025a57598c0/go.mod h1:4mhQ8q/RsB7i+udVvVy5NUi08OU8ZlA0gRVgrF7VFY0=