Merge pull request #115 from TheJokersThief/thejokersthief/impose-api…

…-ratelimit

Enforce and configure an API limit for GCP

cmd/drone-autoscaler/main.go

@@ -249,6 +249,7 @@ func setupProvider(c config.Config) (autoscaler.Provider, error) {
 			google.WithUserDataFile(c.Google.UserDataFile),
 			google.WithZones(c.Google.Zone...),
 			google.WithUserDataKey(c.Google.UserDataKey),
+			google.WithRateLimit(c.Google.RateLimit),
 		)
 	case c.DigitalOcean.Token != "":
 		return digitalocean.New(

config/config.go

@@ -173,6 +173,7 @@ type (
 			UserDataFile        string            `envconfig:"DRONE_GOOGLE_USERDATA_FILE"`
 			Zone                []string          `envconfig:"DRONE_GOOGLE_ZONE"`
 			UserDataKey         string            `envconfig:"DRONE_GOOGLE_USERDATA_KEY" default:"user-data"`
+			RateLimit           int               `envconfig:"DRONE_GOOGLE_READ_RATELIMIT" default:"25"`
 		}
 
 		HetznerCloud struct {

config/load_test.go

@@ -117,6 +117,7 @@ func TestLoad(t *testing.T) {
 		"DRONE_GOOGLE_TAGS":                "drone,agent,prod",
 		"DRONE_GOOGLE_USERDATA":            "#cloud-init",
 		"DRONE_GOOGLE_USERDATA_FILE":       "/path/to/cloud/init.yml",
+		"DRONE_GOOGLE_READ_RATELIMIT":      "20",
 		"DRONE_AMAZON_IMAGE":               "ami-80ca47e6",
 		"DRONE_AMAZON_INSTANCE":            "t2.medium",
 		"DRONE_AMAZON_PRIVATE_IP":          "true",
@@ -295,7 +296,8 @@ var jsonConfig = []byte(`{
     ],
     "UserData": "#cloud-init",
     "UserDataFile": "/path/to/cloud/init.yml",
-    "UserDataKey": "user-data"
+    "UserDataKey": "user-data",
+	"RateLimit": 20
   },
   "HetznerCloud": {
     "Token": "12345678",

drivers/google/option.go

@@ -7,8 +7,10 @@ package google
 import (
 	"io/ioutil"
 	"net/http"
+	"time"
 
 	"github.com/drone/autoscaler/drivers/internal/userdata"
+	"golang.org/x/time/rate"
 
 	"google.golang.org/api/compute/v1"
 )
@@ -153,3 +155,10 @@ func WithServiceAccountEmail(email string) Option {
 		p.serviceAccountEmail = email
 	}
 }
+
+func WithRateLimit(limitAmount int) Option {
+	return func(p *provider) {
+		limit := rate.Every(1 * time.Second / time.Duration(limitAmount))
+		p.rateLimiter = rate.NewLimiter(limit, 1)
+	}
+}

drivers/google/provider.go

@@ -16,6 +16,7 @@ import (
 	"github.com/drone/autoscaler/drivers/internal/userdata"
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/google"
+	"golang.org/x/time/rate"
 	compute "google.golang.org/api/compute/v1"
 	"google.golang.org/api/googleapi"
 )
@@ -53,6 +54,8 @@ type provider struct {
 	userdata            *template.Template
 	userdataKey         string
 
+	rateLimiter *rate.Limiter
+
 	service *compute.Service
 }
 
@@ -95,6 +98,13 @@ func New(opts ...Option) (autoscaler.Provider, error) {
 	if p.serviceAccountEmail == "" {
 		p.serviceAccountEmail = "default"
 	}
+
+	if p.rateLimiter == nil {
+		// If unspecified, set to the max read rate limit for the API 25/s
+		// Source: https://cloud.google.com/compute/docs/api-rate-limits
+		p.rateLimiter = rate.NewLimiter(rate.Every(time.Second/25), 1)
+	}
+
 	if p.service == nil {
 		client, err := google.DefaultClient(oauth2.NoContext, compute.ComputeScope)
 		if err != nil {
@@ -110,35 +120,39 @@ func New(opts ...Option) (autoscaler.Provider, error) {
 
 func (p *provider) waitZoneOperation(ctx context.Context, name string, zone string) error {
 	for {
-		op, err := p.service.ZoneOperations.Get(p.project, zone, name).Context(ctx).Do()
-		if err != nil {
-			if gerr, ok := err.(*googleapi.Error); ok &&
-				gerr.Code == http.StatusNotFound {
-				return autoscaler.ErrInstanceNotFound
+		if p.rateLimiter.Allow() {
+			op, err := p.service.ZoneOperations.Get(p.project, zone, name).Context(ctx).Do()
+			if err != nil {
+				if gerr, ok := err.(*googleapi.Error); ok &&
+					gerr.Code == http.StatusNotFound {
+					return autoscaler.ErrInstanceNotFound
+				}
+				return err
+			}
+			if op.Error != nil {
+				return errors.New(op.Error.Errors[0].Message)
+			}
+			if op.Status == "DONE" {
+				return nil
 			}
-			return err
-		}
-		if op.Error != nil {
-			return errors.New(op.Error.Errors[0].Message)
-		}
-		if op.Status == "DONE" {
-			return nil
 		}
 		time.Sleep(time.Second)
 	}
 }
 
 func (p *provider) waitGlobalOperation(ctx context.Context, name string) error {
 	for {
-		op, err := p.service.GlobalOperations.Get(p.project, name).Context(ctx).Do()
-		if err != nil {
-			return err
-		}
-		if op.Error != nil {
-			return errors.New(op.Error.Errors[0].Message)
-		}
-		if op.Status == "DONE" {
-			return nil
+		if p.rateLimiter.Allow() {
+			op, err := p.service.GlobalOperations.Get(p.project, name).Context(ctx).Do()
+			if err != nil {
+				return err
+			}
+			if op.Error != nil {
+				return errors.New(op.Error.Errors[0].Message)
+			}
+			if op.Status == "DONE" {
+				return nil
+			}
 		}
 		time.Sleep(time.Second)
 	}

go.mod

@@ -68,7 +68,7 @@ require (
 	golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3 // indirect
 	golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be
 	golang.org/x/sync v0.0.0-20190423024810-112230192c58
-	golang.org/x/time v0.0.0-20200416051211-89c76fbcd5d1 // indirect
+	golang.org/x/time v0.0.0-20220224211638-0e9765cccd65 // indirect
 	golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135 // indirect
 	google.golang.org/api v0.0.0-20180921000521-920bb1beccf7
 	google.golang.org/appengine v1.4.0 // indirect

go.sum

@@ -186,6 +186,8 @@ golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/time v0.0.0-20200416051211-89c76fbcd5d1 h1:NusfzzA6yGQ+ua51ck7E3omNUX/JuqbFSaRGqU8CcLI=
 golang.org/x/time v0.0.0-20200416051211-89c76fbcd5d1/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
+golang.org/x/time v0.0.0-20220224211638-0e9765cccd65 h1:M73Iuj3xbbb9Uk1DYhzydthsj6oOd6l9bpuFcNoUvTs=
+golang.org/x/time v0.0.0-20220224211638-0e9765cccd65/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=