Skip to content

Commit 5331916

Browse files
authored
Merge pull request #71 from rolandg/f-rate-limit-ii
Add `rate_limit` option to provider configuration
2 parents af2733c + 57b11ac commit 5331916

File tree

5 files changed

+52
-1
lines changed

5 files changed

+52
-1
lines changed

README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ Have a look at the [examples directory](examples) for some use cases
3737
- `write_returns_object` (boolean, optional): Set this when the API returns the object created on all write operations (`POST`, `PUT`). This is used by the provider to refresh internal data structures. This can also be set with the environment variable `REST_API_WRO`.
3838
- `create_returns_object` (boolean, optional): Set this when the API returns the object created only on creation operations (`POST`). This is used by the provider to refresh internal data structures. This can also be set with the environment variable `REST_API_CRO`.
3939
- `xssi_prefix` (boolean, optional): Trim the xssi prefix from response string, if present, before parsing. This can also be set with the environment variable `REST_API_XSSI_PREFIX`.
40+
- `rate_limit` (float, optional): Set this to limit the number of requests per second made to the API.
4041
- `debug` (boolean, optional): Enabling this will cause lots of debug information to be printed to STDOUT by the API client. This can be gathered by setting `TF_LOG=1` environment variable. This can also be set with the environment variable `REST_API_DEBUG`.
4142

4243
 

restapi/api_client.go

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,19 @@ package restapi
22

33
import (
44
"bytes"
5+
"context"
56
"crypto/tls"
67
"errors"
78
"fmt"
89
"io/ioutil"
910
"log"
11+
"math"
1012
"net/http"
1113
"net/http/cookiejar"
1214
"strings"
1315
"time"
16+
17+
"golang.org/x/time/rate"
1418
)
1519

1620
type apiClientOpt struct {
@@ -31,6 +35,7 @@ type apiClientOpt struct {
3135
create_returns_object bool
3236
xssi_prefix string
3337
use_cookies bool
38+
rate_limit float64
3439
debug bool
3540
}
3641

@@ -52,6 +57,7 @@ type api_client struct {
5257
write_returns_object bool
5358
create_returns_object bool
5459
xssi_prefix string
60+
rate_limiter *rate.Limiter
5561
debug bool
5662
}
5763

@@ -101,12 +107,18 @@ func NewAPIClient(opt *apiClientOpt) (*api_client, error) {
101107
cookieJar, _ = cookiejar.New(nil)
102108
}
103109

110+
rateLimit := rate.Limit(opt.rate_limit)
111+
bucketSize := int(math.Max(math.Round(opt.rate_limit), 1))
112+
log.Printf("limit: %f bucket: %d", opt.rate_limit, bucketSize)
113+
rateLimiter := rate.NewLimiter(rateLimit, bucketSize)
114+
104115
client := api_client{
105116
http_client: &http.Client{
106117
Timeout: time.Second * time.Duration(opt.timeout),
107118
Transport: tr,
108119
Jar: cookieJar,
109120
},
121+
rate_limiter: rateLimiter,
110122
uri: opt.uri,
111123
insecure: opt.insecure,
112124
username: opt.username,
@@ -212,6 +224,14 @@ func (client *api_client) send_request(method string, path string, data string)
212224
log.Printf("%s\n", body)
213225
}
214226

227+
if client.rate_limiter != nil {
228+
// Rate limiting
229+
if client.debug {
230+
log.Printf("Waiting for rate limit availability\n")
231+
}
232+
_ = client.rate_limiter.Wait(context.Background())
233+
}
234+
215235
resp, err := client.http_client.Do(req)
216236

217237
if err != nil {

restapi/api_client_test.go

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,8 @@ func TestAPIClient(t *testing.T) {
2929
copy_keys: make([]string, 0),
3030
write_returns_object: false,
3131
create_returns_object: false,
32-
debug: debug,
32+
rate_limit: 1,
33+
debug: debug,
3334
}
3435
client, _ := NewAPIClient(opt)
3536

@@ -67,6 +68,20 @@ func TestAPIClient(t *testing.T) {
6768
t.Fatalf("client_test.go: Timeout did not trigger on slow request")
6869
}
6970

71+
if debug {
72+
log.Printf("api_client_test.go: Testing rate limited OK request\n")
73+
}
74+
startTime := time.Now().Unix()
75+
76+
for i := 0; i < 4; i++ {
77+
client.send_request("GET", "/ok", "")
78+
}
79+
80+
duration := time.Now().Unix() - startTime
81+
if duration < 3 {
82+
t.Fatalf("client_test.go: requests not delayed\n")
83+
}
84+
7085
if debug {
7186
log.Println("client_test.go: Stopping HTTP server")
7287
}

restapi/provider.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package restapi
22

33
import (
4+
"math"
5+
46
"github.com/hashicorp/terraform/helper/schema"
57
"github.com/hashicorp/terraform/terraform"
68
)
@@ -104,6 +106,12 @@ func Provider() terraform.ResourceProvider {
104106
DefaultFunc: schema.EnvDefaultFunc("REST_API_XSSI_PREFIX", nil),
105107
Description: "Trim the xssi prefix from response string, if present, before parsing.",
106108
},
109+
"rate_limit": &schema.Schema{
110+
Type: schema.TypeFloat,
111+
Optional: true,
112+
DefaultFunc: schema.EnvDefaultFunc("REST_API_RATE_LIMIT", math.MaxFloat64),
113+
Description: "Set this to limit the number of requests per second made to the API.",
114+
},
107115
"debug": &schema.Schema{
108116
Type: schema.TypeBool,
109117
Optional: true,
@@ -155,6 +163,7 @@ func configureProvider(d *schema.ResourceData) (interface{}, error) {
155163
write_returns_object: d.Get("write_returns_object").(bool),
156164
create_returns_object: d.Get("create_returns_object").(bool),
157165
xssi_prefix: d.Get("xssi_prefix").(string),
166+
rate_limit: d.Get("rate_limit").(float64),
158167
debug: d.Get("debug").(bool),
159168
}
160169

vendor/vendor.json

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1278,6 +1278,12 @@
12781278
"revision": "1c95b21c6cfffc50cb6b3c3c23b2257d28b61a9c",
12791279
"revisionTime": "2019-04-12T11:37:27Z"
12801280
},
1281+
{
1282+
"checksumSHA1": "7Ev/X4Xe8P3961myez/hBKO05ig=",
1283+
"path": "golang.org/x/time/rate",
1284+
"revision": "9d24e82272b4f38b78bc8cff74fa936d31ccd8ef",
1285+
"revisionTime": "2019-02-15T22:48:40Z"
1286+
},
12811287
{
12821288
"checksumSHA1": "IO1eKBGX7mKwOE+t1fzdYAHCZmQ=",
12831289
"origin": "github.com/hashicorp/terraform/vendor/google.golang.org/appengine",

0 commit comments

Comments
 (0)