From 333b47ad151494f1a54242318043183b389a25c1 Mon Sep 17 00:00:00 2001 From: Evan Smith Date: Mon, 4 Apr 2022 21:49:10 +0100 Subject: [PATCH] Enforce and configure an API limit for GCP --- cmd/drone-autoscaler/main.go | 1 + config/config.go | 1 + config/load_test.go | 4 ++- drivers/google/option.go | 9 ++++++ drivers/google/provider.go | 56 ++++++++++++++++++++++-------------- go.mod | 2 +- go.sum | 2 ++ 7 files changed, 52 insertions(+), 23 deletions(-) diff --git a/cmd/drone-autoscaler/main.go b/cmd/drone-autoscaler/main.go index 188084b9..fad0bff1 100644 --- a/cmd/drone-autoscaler/main.go +++ b/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( diff --git a/config/config.go b/config/config.go index 441960a6..bdd34bcb 100644 --- a/config/config.go +++ b/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 { diff --git a/config/load_test.go b/config/load_test.go index 31960ce2..996b73c0 100644 --- a/config/load_test.go +++ b/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", diff --git a/drivers/google/option.go b/drivers/google/option.go index 0d3dbcb4..8d739c4d 100644 --- a/drivers/google/option.go +++ b/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) + } +} diff --git a/drivers/google/provider.go b/drivers/google/provider.go index bd7071cb..365953e0 100644 --- a/drivers/google/provider.go +++ b/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,19 +120,21 @@ 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) } @@ -130,15 +142,17 @@ func (p *provider) waitZoneOperation(ctx context.Context, name string, zone stri 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) } diff --git a/go.mod b/go.mod index 697d0ca3..1fecd497 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 801300b2..974853e5 100644 --- a/go.sum +++ b/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=