Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement & use RateLimitTransport #145

Merged
merged 2 commits into from
Sep 7, 2018
Merged

Implement & use RateLimitTransport #145

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3a19bbf
Implement & use RateLimitTransport
radeksimko Sep 4, 2018
9e89aa8
Address code review comments
radeksimko Sep 7, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions github/config.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@ func (c *Config) Client() (interface{}, error) {

tc.Transport = NewEtagTransport(tc.Transport)

tc.Transport = NewRateLimitTransport(tc.Transport)

tc.Transport = logging.NewTransport("Github", tc.Transport)

org.client = github.NewClient(tc)
Expand Down
123 changes: 121 additions & 2 deletions github/transport.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,21 @@
package github

import (
"bytes"
"io"
"io/ioutil"
"log"
"net/http"
"sync"
"time"

"github.com/google/go-github/github"
)

const (
ctxEtag = "etag"
ctxId = "id"
ctxEtag = "etag"
ctxId = "id"
writeDelay = 1 * time.Second
)

// etagTransport allows saving API quota by passing previously stored Etag
Expand All @@ -29,3 +38,113 @@ func (ett *etagTransport) RoundTrip(req *http.Request) (*http.Response, error) {
func NewEtagTransport(rt http.RoundTripper) *etagTransport {
return &etagTransport{transport: rt}
}

// 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 {
transport http.RoundTripper
delayNextRequest bool
responseBody []byte

m sync.Mutex
}

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
rlt.lock(req)

// 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)
}

rlt.delayNextRequest = isWriteMethod(req.Method)

resp, err := rlt.transport.RoundTrip(req)
if err != nil {
rlt.unlock(req)
return resp, err
}

// Make response body accessible for retries & debugging
// (work around bug in GitHub SDK)
// See https://github.com/google/go-github/pull/986
r1, r2, err := drainBody(resp.Body)
if err != nil {
return nil, err
}
resp.Body = r1
ghErr := github.CheckResponse(resp)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm seeing now that github.CheckResponse accepts the response object, making the use of TeeReader impossible. The first time I read it I thought it accepted just a Reader, in that case TeeReader would have worked nicely

var r2 bytes.Buffer
r1 := io.TeeReader(resp.Body, &r2)
ghErr := github.CheckResponse(r1)
resp.Body = &r2

Also the interfaces might not have aligned, if resp.Body was ReadCloser you would have to wrap it just like you did in drainBody anyways nevermind me... always looking for an excuse to use io builtins 😛

resp.Body = r2

// When you have been limited, use the Retry-After response header to slow down.
if arlErr, ok := ghErr.(*github.AbuseRateLimitError); ok {
rlt.delayNextRequest = false
retryAfter := arlErr.GetRetryAfter()
log.Printf("[DEBUG] Abuse detection mechanism triggered, sleeping for %s before retrying",
retryAfter)
time.Sleep(retryAfter)
rlt.unlock(req)
return rlt.RoundTrip(req)
}

if rlErr, ok := ghErr.(*github.RateLimitError); ok {
rlt.delayNextRequest = false
retryAfter := rlErr.Rate.Reset.Sub(time.Now())
log.Printf("[DEBUG] Rate limit %d reached, sleeping for %s (until %s) before retrying",
rlErr.Rate.Limit, retryAfter, time.Now().Add(retryAfter))
time.Sleep(retryAfter)
rlt.unlock(req)
return rlt.RoundTrip(req)
}

rlt.unlock(req)

return resp, nil
}

func (rlt *rateLimitTransport) lock(req *http.Request) {
ctx := req.Context()
log.Printf("[TRACE] Aquiring lock for GitHub API request (%q)", ctx.Value(ctxId))
rlt.m.Lock()
}

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}
}

// drainBody reads all of b to memory and then returns two equivalent
// ReadClosers yielding the same bytes.
func drainBody(b io.ReadCloser) (r1, r2 io.ReadCloser, err error) {
if b == http.NoBody {
// No copying needed. Preserve the magic sentinel meaning of NoBody.
return http.NoBody, http.NoBody, nil
}
var buf bytes.Buffer
if _, err = buf.ReadFrom(b); err != nil {
return nil, b, err
}
if err = b.Close(); err != nil {
return nil, b, err
}
return ioutil.NopCloser(&buf), ioutil.NopCloser(bytes.NewReader(buf.Bytes())), nil
}

func isWriteMethod(method string) bool {
switch method {
case "POST", "PATCH", "PUT", "DELETE":
return true
}
return false
}
164 changes: 164 additions & 0 deletions github/transport_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,170 @@ func githubApiMock(responseSequence []*mockResponse) *httptest.Server {
}))
}

func TestRateLimitTransport_abuseLimit_get(t *testing.T) {
ts := githubApiMock([]*mockResponse{
{
ExpectedUri: "/repos/test/blah",
ResponseBody: `{
"message": "You have triggered an abuse detection mechanism and have been temporarily blocked from content creation. Please retry your request again later.",
"documentation_url": "https://developer.github.com/v3/#abuse-rate-limits"
}`,
StatusCode: 403,
ResponseHeaders: map[string]string{
"Retry-After": "0.1",
},
},
{
ExpectedUri: "/repos/test/blah",
ResponseBody: `{
"message": "You have triggered an abuse detection mechanism and have been temporarily blocked from content creation. Please retry your request again later.",
"documentation_url": "https://developer.github.com/v3/#abuse-rate-limits"
}`,
StatusCode: 403,
ResponseHeaders: map[string]string{
"Retry-After": "0.1",
},
},
{
ExpectedUri: "/repos/test/blah",
ResponseBody: `{"id": 1234}`,
StatusCode: 200,
},
})
defer ts.Close()

httpClient := http.DefaultClient
httpClient.Transport = NewRateLimitTransport(http.DefaultTransport)

client := github.NewClient(httpClient)
u, _ := url.Parse(ts.URL + "/")
client.BaseURL = u

ctx := context.WithValue(context.Background(), ctxId, t.Name())
r, _, err := client.Repositories.Get(ctx, "test", "blah")
if err != nil {
t.Fatal(err)
}

if r.GetID() != 1234 {
t.Fatalf("Expected ID to be 1234, got: %d", r.GetID())
}
}

func TestRateLimitTransport_abuseLimit_post(t *testing.T) {
ts := githubApiMock([]*mockResponse{
{
ExpectedUri: "/orgs/tada/repos",
ExpectedMethod: "POST",
ExpectedBody: []byte(`{"name":"radek-example-48","description":""}
`),
ResponseBody: `{
"message": "You have triggered an abuse detection mechanism and have been temporarily blocked from content creation. Please retry your request again later.",
"documentation_url": "https://developer.github.com/v3/#abuse-rate-limits"
}`,
StatusCode: 403,
ResponseHeaders: map[string]string{
"Retry-After": "0.1",
},
},
{
ExpectedUri: "/orgs/tada/repos",
ExpectedMethod: "POST",
ExpectedBody: []byte(`{"name":"radek-example-48","description":""}
`),
ResponseBody: `{"id": 1234}`,
StatusCode: 200,
},
})
defer ts.Close()

httpClient := http.DefaultClient
httpClient.Transport = NewRateLimitTransport(http.DefaultTransport)

client := github.NewClient(httpClient)
u, _ := url.Parse(ts.URL + "/")
client.BaseURL = u

ctx := context.WithValue(context.Background(), ctxId, t.Name())
r, _, err := client.Repositories.Create(ctx, "tada", &github.Repository{
Name: github.String("radek-example-48"),
Description: github.String(""),
})
if err != nil {
t.Fatal(err)
}

if r.GetID() != 1234 {
t.Fatalf("Expected ID to be 1234, got: %d", r.GetID())
}
}

func TestRateLimitTransport_abuseLimit_post_error(t *testing.T) {
ts := githubApiMock([]*mockResponse{
{
ExpectedUri: "/orgs/tada/repos",
ExpectedMethod: "POST",
ExpectedBody: []byte(`{"name":"radek-example-48","description":""}
`),
ResponseBody: `{
"message": "You have triggered an abuse detection mechanism and have been temporarily blocked from content creation. Please retry your request again later.",
"documentation_url": "https://developer.github.com/v3/#abuse-rate-limits"
}`,
StatusCode: 403,
ResponseHeaders: map[string]string{
"Retry-After": "0.1",
},
},
{
ExpectedUri: "/orgs/tada/repos",
ExpectedMethod: "POST",
ExpectedBody: []byte(`{"name":"radek-example-48","description":""}
`),
ResponseBody: `{
"message": "Repository creation failed.",
"errors": [
{
"resource": "Repository",
"code": "custom",
"field": "name",
"message": "name already exists on this account"
}
],
"documentation_url": "https://developer.github.com/v3/repos/#create"
}
`,
StatusCode: 422,
},
})
defer ts.Close()

httpClient := http.DefaultClient
httpClient.Transport = NewRateLimitTransport(http.DefaultTransport)

client := github.NewClient(httpClient)
u, _ := url.Parse(ts.URL + "/")
client.BaseURL = u

ctx := context.WithValue(context.Background(), ctxId, t.Name())
_, _, err := client.Repositories.Create(ctx, "tada", &github.Repository{
Name: github.String("radek-example-48"),
Description: github.String(""),
})
if err == nil {
t.Fatal("Expected 422 error, got nil")
}

ghErr, ok := err.(*github.ErrorResponse)
if !ok {
t.Fatalf("Expected github.ErrorResponse, got: %#v", err)
}

expectedMessage := "Repository creation failed."
if ghErr.Message != expectedMessage {
t.Fatalf("Expected message %q, got: %q", expectedMessage, ghErr.Message)
}
}

type mockResponse struct {
ExpectedUri string
ExpectedMethod string
Expand Down