Merge pull request #71 from rolandg/f-rate-limit-ii

Add `rate_limit` option to provider configuration

Author: DRuggeri

README.md

@@ -37,6 +37,7 @@ Have a look at the [examples directory](examples) for some use cases
 - `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`.
 - `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`.
 - `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`.
+- `rate_limit` (float, optional): Set this to limit the number of requests per second made to the API.
 - `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`.
 
  

restapi/api_client.go

@@ -2,15 +2,19 @@ package restapi
 
 import (
 	"bytes"
+	"context"
 	"crypto/tls"
 	"errors"
 	"fmt"
 	"io/ioutil"
 	"log"
+	"math"
 	"net/http"
 	"net/http/cookiejar"
 	"strings"
 	"time"
+
+	"golang.org/x/time/rate"
 )
 
 type apiClientOpt struct {
@@ -31,6 +35,7 @@ type apiClientOpt struct {
 	create_returns_object bool
 	xssi_prefix           string
 	use_cookies           bool
+	rate_limit            float64
 	debug                 bool
 }
 
@@ -52,6 +57,7 @@ type api_client struct {
 	write_returns_object  bool
 	create_returns_object bool
 	xssi_prefix           string
+	rate_limiter          *rate.Limiter
 	debug                 bool
 }
 
@@ -101,12 +107,18 @@ func NewAPIClient(opt *apiClientOpt) (*api_client, error) {
 		cookieJar, _ = cookiejar.New(nil)
 	}
 
+	rateLimit := rate.Limit(opt.rate_limit)
+	bucketSize := int(math.Max(math.Round(opt.rate_limit), 1))
+	log.Printf("limit: %f bucket: %d", opt.rate_limit, bucketSize)
+	rateLimiter := rate.NewLimiter(rateLimit, bucketSize)
+
 	client := api_client{
 		http_client: &http.Client{
 			Timeout:   time.Second * time.Duration(opt.timeout),
 			Transport: tr,
 			Jar:       cookieJar,
 		},
+		rate_limiter:          rateLimiter,
 		uri:                   opt.uri,
 		insecure:              opt.insecure,
 		username:              opt.username,
@@ -212,6 +224,14 @@ func (client *api_client) send_request(method string, path string, data string)
 		log.Printf("%s\n", body)
 	}
 
+	if client.rate_limiter != nil {
+		// Rate limiting
+		if client.debug {
+			log.Printf("Waiting for rate limit availability\n")
+		}
+		_ = client.rate_limiter.Wait(context.Background())
+	}
+
 	resp, err := client.http_client.Do(req)
 
 	if err != nil {

restapi/api_client_test.go

@@ -29,7 +29,8 @@ func TestAPIClient(t *testing.T) {
 		copy_keys:             make([]string, 0),
 		write_returns_object:  false,
 		create_returns_object: false,
-		debug: debug,
+		rate_limit:            1,
+		debug:                 debug,
 	}
 	client, _ := NewAPIClient(opt)
 
@@ -67,6 +68,20 @@ func TestAPIClient(t *testing.T) {
 		t.Fatalf("client_test.go: Timeout did not trigger on slow request")
 	}
 
+	if debug {
+		log.Printf("api_client_test.go: Testing rate limited OK request\n")
+	}
+	startTime := time.Now().Unix()
+
+	for i := 0; i < 4; i++ {
+		client.send_request("GET", "/ok", "")
+	}
+
+	duration := time.Now().Unix() - startTime
+	if duration < 3 {
+		t.Fatalf("client_test.go: requests not delayed\n")
+	}
+
 	if debug {
 		log.Println("client_test.go: Stopping HTTP server")
 	}

restapi/provider.go

@@ -1,6 +1,8 @@
 package restapi
 
 import (
+	"math"
+
 	"github.com/hashicorp/terraform/helper/schema"
 	"github.com/hashicorp/terraform/terraform"
 )
@@ -104,6 +106,12 @@ func Provider() terraform.ResourceProvider {
 				DefaultFunc: schema.EnvDefaultFunc("REST_API_XSSI_PREFIX", nil),
 				Description: "Trim the xssi prefix from response string, if present, before parsing.",
 			},
+			"rate_limit": &schema.Schema{
+				Type:        schema.TypeFloat,
+				Optional:    true,
+				DefaultFunc: schema.EnvDefaultFunc("REST_API_RATE_LIMIT", math.MaxFloat64),
+				Description: "Set this to limit the number of requests per second made to the API.",
+			},
 			"debug": &schema.Schema{
 				Type:        schema.TypeBool,
 				Optional:    true,
@@ -155,6 +163,7 @@ func configureProvider(d *schema.ResourceData) (interface{}, error) {
 		write_returns_object:  d.Get("write_returns_object").(bool),
 		create_returns_object: d.Get("create_returns_object").(bool),
 		xssi_prefix:           d.Get("xssi_prefix").(string),
+		rate_limit:            d.Get("rate_limit").(float64),
 		debug:                 d.Get("debug").(bool),
 	}
 

vendor/vendor.json

@@ -1278,6 +1278,12 @@
 			"revision": "1c95b21c6cfffc50cb6b3c3c23b2257d28b61a9c",
 			"revisionTime": "2019-04-12T11:37:27Z"
 		},
+		{
+			"checksumSHA1": "7Ev/X4Xe8P3961myez/hBKO05ig=",
+			"path": "golang.org/x/time/rate",
+			"revision": "9d24e82272b4f38b78bc8cff74fa936d31ccd8ef",
+			"revisionTime": "2019-02-15T22:48:40Z"
+		},
 		{
 			"checksumSHA1": "IO1eKBGX7mKwOE+t1fzdYAHCZmQ=",
 			"origin": "github.com/hashicorp/terraform/vendor/google.golang.org/appengine",