Make rate-limiting configurable (#907)

* Make struct type public

* Scaffold functional pattern for providing writeDelay

* Add WriteDelay to config and plumb it down

* Add configuration for write delay

* Validate writeDelay is positive

* Add website description of write_delay argument

* Use convenience func to pass writeDelay into rate limiter

* Move convenience func closer to implementation; add convenience type

* Rename write_delay --> write_delay_ms

* Update github/provider.go

Co-authored-by: Jeremy Udit <[email protected]>

* Describe default for write_delay_ms in help text

* Update github/transport.go

Co-authored-by: Jeremy Udit <[email protected]>

* Update website/docs/index.html.markdown

Co-authored-by: Jeremy Udit <[email protected]>

Co-authored-by: Jeremy Udit <[email protected]>

Author: kfcampbell

github/config.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/url"
 	"path"
+	"time"
 
 	"github.com/google/go-github/v39/github"
 	"github.com/hashicorp/terraform-plugin-sdk/helper/logging"
@@ -14,10 +15,11 @@ import (
 )
 
 type Config struct {
-	Token    string
-	Owner    string
-	BaseURL  string
-	Insecure bool
+	Token      string
+	Owner      string
+	BaseURL    string
+	Insecure   bool
+	WriteDelay time.Duration
 }
 
 type Owner struct {
@@ -29,14 +31,13 @@ type Owner struct {
 	IsOrganization bool
 }
 
-func RateLimitedHTTPClient(client *http.Client) *http.Client {
+func RateLimitedHTTPClient(client *http.Client, writeDelay time.Duration) *http.Client {
 
 	client.Transport = NewEtagTransport(client.Transport)
-	client.Transport = NewRateLimitTransport(client.Transport)
+	client.Transport = NewRateLimitTransport(client.Transport, WithWriteDelay(writeDelay))
 	client.Transport = logging.NewTransport("Github", client.Transport)
 
 	return client
-
 }
 
 func (c *Config) AuthenticatedHTTPClient() *http.Client {
@@ -47,8 +48,7 @@ func (c *Config) AuthenticatedHTTPClient() *http.Client {
 	)
 	client := oauth2.NewClient(ctx, ts)
 
-	return RateLimitedHTTPClient(client)
-
+	return RateLimitedHTTPClient(client, c.WriteDelay)
 }
 
 func (c *Config) Anonymous() bool {
@@ -57,7 +57,7 @@ func (c *Config) Anonymous() bool {
 
 func (c *Config) AnonymousHTTPClient() *http.Client {
 	client := &http.Client{Transport: &http.Transport{}}
-	return RateLimitedHTTPClient(client)
+	return RateLimitedHTTPClient(client, c.WriteDelay)
 }
 
 func (c *Config) NewGraphQLClient(client *http.Client) (*githubv4.Client, error) {
@@ -74,7 +74,6 @@ func (c *Config) NewGraphQLClient(client *http.Client) (*githubv4.Client, error)
 	}
 
 	return githubv4.NewEnterpriseClient(uv4.String(), client), nil
-
 }
 
 func (c *Config) NewRESTClient(client *http.Client) (*github.Client, error) {
@@ -96,7 +95,6 @@ func (c *Config) NewRESTClient(client *http.Client) (*github.Client, error) {
 	v3client.BaseURL = uv3
 
 	return v3client, nil
-
 }
 
 func (c *Config) ConfigureOwner(owner *Owner) (*Owner, error) {

github/provider.go

@@ -3,6 +3,7 @@ package github
 import (
 	"fmt"
 	"log"
+	"time"
 
 	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
 	"github.com/hashicorp/terraform-plugin-sdk/terraform"
@@ -42,6 +43,12 @@ func Provider() terraform.ResourceProvider {
 				Default:     false,
 				Description: descriptions["insecure"],
 			},
+			"write_delay_ms": {
+				Type:        schema.TypeInt,
+				Optional:    true,
+				Default:     1000,
+				Description: descriptions["write_delay_ms"],
+			},
 			"app_auth": {
 				Type:        schema.TypeList,
 				Optional:    true,
@@ -158,6 +165,8 @@ func init() {
 		"app_auth.id":              "The GitHub App ID.",
 		"app_auth.installation_id": "The GitHub App installation instance ID.",
 		"app_auth.pem_file":        "The GitHub App PEM file contents.",
+		"write_delay_ms": "Amount of time in milliseconds to sleep in between writes to GitHub API. " +
+			"Defaults to 1000ms or 1s if not set.",
 	}
 }
 
@@ -221,11 +230,18 @@ func providerConfigure(p *schema.Provider) schema.ConfigureFunc {
 			token = appToken
 		}
 
+		writeDelay := d.Get("write_delay_ms").(int)
+		if writeDelay <= 0 {
+			return nil, fmt.Errorf("write_delay_ms must be greater than 0ms")
+		}
+		log.Printf("[DEBUG] Setting write_delay_ms to %d", writeDelay)
+
 		config := Config{
-			Token:    token,
-			BaseURL:  baseURL,
-			Insecure: insecure,
-			Owner:    owner,
+			Token:      token,
+			BaseURL:    baseURL,
+			Insecure:   insecure,
+			Owner:      owner,
+			WriteDelay: time.Duration(writeDelay) * time.Millisecond,
 		}
 
 		meta, err := config.Meta()

github/transport.go

@@ -13,9 +13,8 @@ import (
 )
 
 const (
-	ctxEtag    = ctxEtagType("etag")
-	ctxId      = ctxIdType("id")
-	writeDelay = 1 * time.Second
+	ctxEtag = ctxEtagType("etag")
+	ctxId   = ctxIdType("id")
 )
 
 // ctxIdType is used to avoid collisions between packages using context
@@ -45,17 +44,18 @@ func NewEtagTransport(rt http.RoundTripper) *etagTransport {
 	return &etagTransport{transport: rt}
 }
 
-// rateLimitTransport implements GitHub's best practices
+// RateLimitTransport implements GitHub's best practices
 // for avoiding rate limits
 // https://developer.github.com/v3/guides/best-practices-for-integrators/#dealing-with-abuse-rate-limits
-type rateLimitTransport struct {
+type RateLimitTransport struct {
 	transport        http.RoundTripper
 	delayNextRequest bool
+	writeDelay       time.Duration
 
 	m sync.Mutex
 }
 
-func (rlt *rateLimitTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+func (rlt *RateLimitTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 	// Make requests for a single user or client ID serially
 	// This is also necessary for safely saving
 	// and restoring bodies between retries below
@@ -64,8 +64,8 @@ func (rlt *rateLimitTransport) RoundTrip(req *http.Request) (*http.Response, err
 	// If you're making a large number of POST, PATCH, PUT, or DELETE requests
 	// for a single user or client ID, wait at least one second between each request.
 	if rlt.delayNextRequest {
-		log.Printf("[DEBUG] Sleeping %s between write operations", writeDelay)
-		time.Sleep(writeDelay)
+		log.Printf("[DEBUG] Sleeping %s between write operations", rlt.writeDelay)
+		time.Sleep(rlt.writeDelay)
 	}
 
 	rlt.delayNextRequest = isWriteMethod(req.Method)
@@ -113,20 +113,39 @@ func (rlt *rateLimitTransport) RoundTrip(req *http.Request) (*http.Response, err
 	return resp, nil
 }
 
-func (rlt *rateLimitTransport) lock(req *http.Request) {
+func (rlt *RateLimitTransport) lock(req *http.Request) {
 	ctx := req.Context()
 	log.Printf("[TRACE] Acquiring lock for GitHub API request (%q)", ctx.Value(ctxId))
 	rlt.m.Lock()
 }
 
-func (rlt *rateLimitTransport) unlock(req *http.Request) {
+func (rlt *RateLimitTransport) unlock(req *http.Request) {
 	ctx := req.Context()
 	log.Printf("[TRACE] Releasing lock for GitHub API request (%q)", ctx.Value(ctxId))
 	rlt.m.Unlock()
 }
 
-func NewRateLimitTransport(rt http.RoundTripper) *rateLimitTransport {
-	return &rateLimitTransport{transport: rt}
+type RateLimitTransportOption func(*RateLimitTransport)
+
+// NewRateLimitTransport takes in an http.RoundTripper and a variadic list of
+// optional functions that modify the RateLimitTransport struct itself. This
+// may be used to alter the write delay in between requests, for example.
+func NewRateLimitTransport(rt http.RoundTripper, options ...RateLimitTransportOption) *RateLimitTransport {
+	// Default to 1 second of write delay if none is provided
+	rlt := &RateLimitTransport{transport: rt, writeDelay: 1 * time.Second}
+
+	for _, opt := range options {
+		opt(rlt)
+	}
+
+	return rlt
+}
+
+// WithWriteDelay is used to set the write delay between requests
+func WithWriteDelay(d time.Duration) RateLimitTransportOption {
+	return func(rlt *RateLimitTransport) {
+		rlt.writeDelay = d
+	}
 }
 
 // drainBody reads all of b to memory and then returns two equivalent

website/docs/index.html.markdown

@@ -107,6 +107,8 @@ The following arguments are supported in the `provider` block:
   * `installation_id` - (Required) This is the ID of the GitHub App installation. It can sourced from the `GITHUB_APP_INSTALLATION_ID` environment variable.
   * `pem_file` - (Required) This is the contents of the GitHub App private key PEM file. It can also be sourced from the `GITHUB_APP_PEM_FILE` environment variable.
 
+* `write_delay_ms` - (Optional) The number of milliseconds to sleep in between write operations in order to satisfy the GitHub API rate limits. Defaults to 1000ms or 1 second if not provided.
+
 Note: If you have a PEM file on disk, you can pass it in via `pem_file = file("path/to/file.pem")`.
 
 For backwards compatibility, if more than one of `owner`, `organization`,