Merge pull request integrations#145 from terraform-providers/f-rate-l…

…imits

Implement & use RateLimitTransport

github/config.go

@@ -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)

github/transport.go

@@ -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
@@ -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)
+	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
+}

github/transport_test.go

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