Skip to content

Commit

Permalink
Make rate-limiting configurable (integrations#907)
Browse files Browse the repository at this point in the history
* 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 <jcudit@github.com>

* Describe default for write_delay_ms in help text

* Update github/transport.go

Co-authored-by: Jeremy Udit <jcudit@github.com>

* Update website/docs/index.html.markdown

Co-authored-by: Jeremy Udit <jcudit@github.com>

Co-authored-by: Jeremy Udit <jcudit@github.com>
  • Loading branch information
kfcampbell and Jeremy Udit authored Oct 5, 2021
1 parent 422041f commit 98140ef
Show file tree
Hide file tree
Showing 4 changed files with 63 additions and 28 deletions.
22 changes: 10 additions & 12 deletions github/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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 {
Expand All @@ -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 {
Expand All @@ -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 {
Expand All @@ -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) {
Expand All @@ -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) {
Expand All @@ -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) {
Expand Down
24 changes: 20 additions & 4 deletions github/provider.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.",
}
}

Expand Down Expand Up @@ -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()
Expand Down
43 changes: 31 additions & 12 deletions github/transport.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions website/docs/index.html.markdown
Original file line number Diff line number Diff line change
Expand Up @@ -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`,
Expand Down

0 comments on commit 98140ef

Please sign in to comment.