Skip to content

Commit 98c1f93

Browse files
bradrydzewskibradrydzewski
authored
Merge pull request #115 from TheJokersThief/thejokersthief/impose-api-ratelimit
Enforce and configure an API limit for GCP
2 parents 8cd545f + 333b47a commit 98c1f93

File tree

7 files changed

+52
-23
lines changed

7 files changed

+52
-23
lines changed

‎cmd/drone-autoscaler/main.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -249,6 +249,7 @@ func setupProvider(c config.Config) (autoscaler.Provider, error) {
249249
google.WithUserDataFile(c.Google.UserDataFile),
250250
google.WithZones(c.Google.Zone...),
251251
google.WithUserDataKey(c.Google.UserDataKey),
252+
google.WithRateLimit(c.Google.RateLimit),
252253
)
253254
case c.DigitalOcean.Token != "":
254255
return digitalocean.New(

‎config/config.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,7 @@ type (
173173
UserDataFile string `envconfig:"DRONE_GOOGLE_USERDATA_FILE"`
174174
Zone []string `envconfig:"DRONE_GOOGLE_ZONE"`
175175
UserDataKey string `envconfig:"DRONE_GOOGLE_USERDATA_KEY" default:"user-data"`
176+
RateLimit int `envconfig:"DRONE_GOOGLE_READ_RATELIMIT" default:"25"`
176177
}
177178

178179
HetznerCloud struct {

‎config/load_test.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ func TestLoad(t *testing.T) {
117117
"DRONE_GOOGLE_TAGS": "drone,agent,prod",
118118
"DRONE_GOOGLE_USERDATA": "#cloud-init",
119119
"DRONE_GOOGLE_USERDATA_FILE": "/path/to/cloud/init.yml",
120+
"DRONE_GOOGLE_READ_RATELIMIT": "20",
120121
"DRONE_AMAZON_IMAGE": "ami-80ca47e6",
121122
"DRONE_AMAZON_INSTANCE": "t2.medium",
122123
"DRONE_AMAZON_PRIVATE_IP": "true",
@@ -295,7 +296,8 @@ var jsonConfig = []byte(`{
295296
],
296297
"UserData": "#cloud-init",
297298
"UserDataFile": "/path/to/cloud/init.yml",
298-
"UserDataKey": "user-data"
299+
"UserDataKey": "user-data",
300+
"RateLimit": 20
299301
},
300302
"HetznerCloud": {
301303
"Token": "12345678",

‎drivers/google/option.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@ package google
77
import (
88
"io/ioutil"
99
"net/http"
10+
"time"
1011

1112
"github.com/drone/autoscaler/drivers/internal/userdata"
13+
"golang.org/x/time/rate"
1214

1315
"google.golang.org/api/compute/v1"
1416
)
@@ -153,3 +155,10 @@ func WithServiceAccountEmail(email string) Option {
153155
p.serviceAccountEmail = email
154156
}
155157
}
158+
159+
func WithRateLimit(limitAmount int) Option {
160+
return func(p *provider) {
161+
limit := rate.Every(1 * time.Second / time.Duration(limitAmount))
162+
p.rateLimiter = rate.NewLimiter(limit, 1)
163+
}
164+
}

‎drivers/google/provider.go

Lines changed: 35 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"github.com/drone/autoscaler/drivers/internal/userdata"
1717
"golang.org/x/oauth2"
1818
"golang.org/x/oauth2/google"
19+
"golang.org/x/time/rate"
1920
compute "google.golang.org/api/compute/v1"
2021
"google.golang.org/api/googleapi"
2122
)
@@ -53,6 +54,8 @@ type provider struct {
5354
userdata *template.Template
5455
userdataKey string
5556

57+
rateLimiter *rate.Limiter
58+
5659
service *compute.Service
5760
}
5861

@@ -95,6 +98,13 @@ func New(opts ...Option) (autoscaler.Provider, error) {
9598
if p.serviceAccountEmail == "" {
9699
p.serviceAccountEmail = "default"
97100
}
101+
102+
if p.rateLimiter == nil {
103+
// If unspecified, set to the max read rate limit for the API 25/s
104+
// Source: https://cloud.google.com/compute/docs/api-rate-limits
105+
p.rateLimiter = rate.NewLimiter(rate.Every(time.Second/25), 1)
106+
}
107+
98108
if p.service == nil {
99109
client, err := google.DefaultClient(oauth2.NoContext, compute.ComputeScope)
100110
if err != nil {
@@ -110,35 +120,39 @@ func New(opts ...Option) (autoscaler.Provider, error) {
110120

111121
func (p *provider) waitZoneOperation(ctx context.Context, name string, zone string) error {
112122
for {
113-
op, err := p.service.ZoneOperations.Get(p.project, zone, name).Context(ctx).Do()
114-
if err != nil {
115-
if gerr, ok := err.(*googleapi.Error); ok &&
116-
gerr.Code == http.StatusNotFound {
117-
return autoscaler.ErrInstanceNotFound
123+
if p.rateLimiter.Allow() {
124+
op, err := p.service.ZoneOperations.Get(p.project, zone, name).Context(ctx).Do()
125+
if err != nil {
126+
if gerr, ok := err.(*googleapi.Error); ok &&
127+
gerr.Code == http.StatusNotFound {
128+
return autoscaler.ErrInstanceNotFound
129+
}
130+
return err
131+
}
132+
if op.Error != nil {
133+
return errors.New(op.Error.Errors[0].Message)
134+
}
135+
if op.Status == "DONE" {
136+
return nil
118137
}
119-
return err
120-
}
121-
if op.Error != nil {
122-
return errors.New(op.Error.Errors[0].Message)
123-
}
124-
if op.Status == "DONE" {
125-
return nil
126138
}
127139
time.Sleep(time.Second)
128140
}
129141
}
130142

131143
func (p *provider) waitGlobalOperation(ctx context.Context, name string) error {
132144
for {
133-
op, err := p.service.GlobalOperations.Get(p.project, name).Context(ctx).Do()
134-
if err != nil {
135-
return err
136-
}
137-
if op.Error != nil {
138-
return errors.New(op.Error.Errors[0].Message)
139-
}
140-
if op.Status == "DONE" {
141-
return nil
145+
if p.rateLimiter.Allow() {
146+
op, err := p.service.GlobalOperations.Get(p.project, name).Context(ctx).Do()
147+
if err != nil {
148+
return err
149+
}
150+
if op.Error != nil {
151+
return errors.New(op.Error.Errors[0].Message)
152+
}
153+
if op.Status == "DONE" {
154+
return nil
155+
}
142156
}
143157
time.Sleep(time.Second)
144158
}

‎go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ require (
6868
golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3 // indirect
6969
golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be
7070
golang.org/x/sync v0.0.0-20190423024810-112230192c58
71-
golang.org/x/time v0.0.0-20200416051211-89c76fbcd5d1 // indirect
71+
golang.org/x/time v0.0.0-20220224211638-0e9765cccd65 // indirect
7272
golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135 // indirect
7373
google.golang.org/api v0.0.0-20180921000521-920bb1beccf7
7474
google.golang.org/appengine v1.4.0 // indirect

‎go.sum

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,8 @@ golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg=
186186
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
187187
golang.org/x/time v0.0.0-20200416051211-89c76fbcd5d1 h1:NusfzzA6yGQ+ua51ck7E3omNUX/JuqbFSaRGqU8CcLI=
188188
golang.org/x/time v0.0.0-20200416051211-89c76fbcd5d1/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
189+
golang.org/x/time v0.0.0-20220224211638-0e9765cccd65 h1:M73Iuj3xbbb9Uk1DYhzydthsj6oOd6l9bpuFcNoUvTs=
190+
golang.org/x/time v0.0.0-20220224211638-0e9765cccd65/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
189191
golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
190192
golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=
191193
golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=

0 commit comments

Comments
 (0)