Merge pull request #800 from xanzy/svh/f-rate-limit

Implement a rate limiting stategy

Author: svanharmelen

.travis.yml

@@ -1,8 +1,8 @@
 language: go
 
 go:
-  - 1.12.x
   - 1.13.x
+  - 1.14.x
   - master
 
 stages:

gitlab.go

@@ -25,6 +25,7 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"math/rand"
 	"net/http"
 	"net/url"
 	"sort"
@@ -33,13 +34,19 @@ import (
 	"time"
 
 	"github.com/google/go-querystring/query"
+	"github.com/hashicorp/go-cleanhttp"
+	retryablehttp "github.com/hashicorp/go-retryablehttp"
 	"golang.org/x/oauth2"
+	"golang.org/x/time/rate"
 )
 
 const (
 	defaultBaseURL = "https://gitlab.com/"
 	apiVersionPath = "api/v4/"
 	userAgent      = "go-gitlab"
+
+	headerRateLimit = "RateLimit-Limit"
+	headerRateReset = "RateLimit-Reset"
 )
 
 // authType represents an authentication type within GitLab.
@@ -321,13 +328,16 @@ const (
 // A Client manages communication with the GitLab API.
 type Client struct {
 	// HTTP client used to communicate with the API.
-	client *http.Client
+	client *retryablehttp.Client
 
 	// Base URL for API requests. Defaults to the public GitLab API, but can be
 	// set to a domain endpoint to use with a self hosted GitLab server. baseURL
 	// should always be specified with a trailing slash.
 	baseURL *url.URL
 
+	// Limiter is used to limit API calls and prevent 429 responses.
+	limiter *rate.Limiter
+
 	// Token type used to make authenticated API calls.
 	authType authType
 
@@ -479,15 +489,25 @@ func NewOAuthClient(httpClient *http.Client, token string) *Client {
 
 func newClient(httpClient *http.Client) *Client {
 	if httpClient == nil {
-		httpClient = http.DefaultClient
+		httpClient = cleanhttp.DefaultPooledClient()
 	}
 
-	c := &Client{client: httpClient, UserAgent: userAgent}
-	if err := c.SetBaseURL(defaultBaseURL); err != nil {
-		// Should never happen since defaultBaseURL is our constant.
-		panic(err)
+	c := &Client{UserAgent: userAgent}
+
+	// Configure the HTTP client.
+	c.client = &retryablehttp.Client{
+		Backoff:      c.retryHTTPBackoff,
+		CheckRetry:   c.retryHTTPCheck,
+		ErrorHandler: retryablehttp.PassthroughErrorHandler,
+		HTTPClient:   httpClient,
+		RetryWaitMin: 100 * time.Millisecond,
+		RetryWaitMax: 400 * time.Millisecond,
+		RetryMax:     30,
 	}
 
+	// Set the default base URL.
+	_ = c.SetBaseURL(defaultBaseURL)
+
 	// Create the internal timeStats service.
 	timeStats := &timeStatsService{client: c}
 
@@ -566,6 +586,64 @@ func newClient(httpClient *http.Client) *Client {
 	return c
 }
 
+// retryHTTPCheck provides a callback for Client.CheckRetry which
+// will retry both rate limit (429) and server (>= 500) errors.
+func (c *Client) retryHTTPCheck(ctx context.Context, resp *http.Response, err error) (bool, error) {
+	if ctx.Err() != nil {
+		return false, ctx.Err()
+	}
+	if err != nil {
+		return false, err
+	}
+	if resp.StatusCode == 429 || resp.StatusCode >= 500 {
+		return true, nil
+	}
+	return false, nil
+}
+
+// retryHTTPBackoff provides a generic callback for Client.Backoff which
+// will pass through all calls based on the status code of the response.
+func (c *Client) retryHTTPBackoff(min, max time.Duration, attemptNum int, resp *http.Response) time.Duration {
+	// Use the rate limit backoff function when we are rate limited.
+	if resp != nil && resp.StatusCode == 429 {
+		return rateLimitBackoff(min, max, attemptNum, resp)
+	}
+
+	// Set custom duration's when we experience a service interruption.
+	min = 700 * time.Millisecond
+	max = 900 * time.Millisecond
+
+	return retryablehttp.LinearJitterBackoff(min, max, attemptNum, resp)
+}
+
+// rateLimitBackoff provides a callback for Client.Backoff which will use the
+// RateLimit-Reset header to determine the time to wait. We add some jitter
+// to prevent a thundering herd.
+//
+// min and max are mainly used for bounding the jitter that will be added to
+// the reset time retrieved from the headers. But if the final wait time is
+// less then min, min will be used instead.
+func rateLimitBackoff(min, max time.Duration, attemptNum int, resp *http.Response) time.Duration {
+	// rnd is used to generate pseudo-random numbers.
+	rnd := rand.New(rand.NewSource(time.Now().UnixNano()))
+
+	// First create some jitter bounded by the min and max durations.
+	jitter := time.Duration(rnd.Float64() * float64(max-min))
+
+	if resp != nil {
+		if v := resp.Header.Get(headerRateReset); v != "" {
+			if reset, _ := strconv.ParseInt(v, 10, 64); reset > 0 {
+				// Only update min if the given time to wait is longer.
+				if wait := time.Until(time.Unix(reset, 0)); wait > min {
+					min = wait
+				}
+			}
+		}
+	}
+
+	return min + jitter
+}
+
 // BaseURL return a copy of the baseURL.
 func (c *Client) BaseURL() *url.URL {
 	u := *c.baseURL
@@ -592,6 +670,49 @@ func (c *Client) SetBaseURL(urlStr string) error {
 	// Update the base URL of the client.
 	c.baseURL = baseURL
 
+	// Reconfigure the rate limiter.
+	return c.configureLimiter()
+}
+
+// configureLimiter configures the rate limiter.
+func (c *Client) configureLimiter() error {
+	// Set default values for when rate limiting is disabled.
+	limit := rate.Inf
+	burst := 0
+
+	defer func() {
+		// Create a new limiter using the calculated values.
+		c.limiter = rate.NewLimiter(limit, burst)
+	}()
+
+	// Create a new request.
+	req, err := http.NewRequest("GET", c.baseURL.String(), nil)
+	if err != nil {
+		return err
+	}
+
+	// Make a single request to retrieve the rate limit headers.
+	resp, err := c.client.HTTPClient.Do(req)
+	if err != nil {
+		return err
+	}
+	resp.Body.Close()
+
+	if v := resp.Header.Get(headerRateLimit); v != "" {
+		if rateLimit, _ := strconv.ParseFloat(v, 64); rateLimit > 0 {
+			// The rate limit is based on requests per minute, so for our limiter to
+			// work correctly we devide the limit by 60 to get the limit per second.
+			rateLimit /= 60
+			// Configure the limit and burst using a split of 2/3 for the limit and
+			// 1/3 for the burst. This enables clients to burst 1/3 of the allowed
+			// calls before the limiter kicks in. The remaining calls will then be
+			// spread out evenly using intervals of time.Second / limit which should
+			// prevent hitting the rate limit.
+			limit = rate.Limit(rateLimit * 0.66)
+			burst = int(rateLimit * 0.33)
+		}
+	}
+
 	return nil
 }
 
@@ -600,7 +721,7 @@ func (c *Client) SetBaseURL(urlStr string) error {
 // Relative URL paths should always be specified without a preceding slash. If
 // specified, the value pointed to by body is JSON encoded and included as the
 // request body.
-func (c *Client) NewRequest(method, path string, opt interface{}, options []OptionFunc) (*http.Request, error) {
+func (c *Client) NewRequest(method, path string, opt interface{}, options []OptionFunc) (*retryablehttp.Request, error) {
 	u := *c.baseURL
 	unescaped, err := url.PathUnescape(path)
 	if err != nil {
@@ -618,54 +739,53 @@ func (c *Client) NewRequest(method, path string, opt interface{}, options []Opti
 		}
 		u.RawQuery = q.Encode()
 	}
+	// Create a request specific headers map.
+	reqHeaders := make(http.Header)
+	reqHeaders.Set("Accept", "application/json")
 
-	req := &http.Request{
-		Method:     method,
-		URL:        &u,
-		Proto:      "HTTP/1.1",
-		ProtoMajor: 1,
-		ProtoMinor: 1,
-		Header:     make(http.Header),
-		Host:       u.Host,
+	switch c.authType {
+	case basicAuth, oAuthToken:
+		reqHeaders.Set("Authorization", "Bearer "+c.token)
+	case privateToken:
+		reqHeaders.Set("PRIVATE-TOKEN", c.token)
 	}
 
-	for _, fn := range options {
-		if fn == nil {
-			continue
-		}
-
-		if err := fn(req); err != nil {
-			return nil, err
-		}
+	if c.UserAgent != "" {
+		reqHeaders.Set("User-Agent", c.UserAgent)
 	}
 
+	var body interface{}
 	if method == "POST" || method == "PUT" {
-		bodyBytes, err := json.Marshal(opt)
-		if err != nil {
-			return nil, err
-		}
-		bodyReader := bytes.NewReader(bodyBytes)
-
 		u.RawQuery = ""
-		req.Body = ioutil.NopCloser(bodyReader)
-		req.GetBody = func() (io.ReadCloser, error) {
-			return ioutil.NopCloser(bodyReader), nil
+		reqHeaders.Set("Content-Type", "application/json")
+
+		if opt != nil {
+			bodyBytes, err := json.Marshal(opt)
+			if err != nil {
+				return nil, err
+			}
+			body = bytes.NewReader(bodyBytes)
 		}
-		req.ContentLength = int64(bodyReader.Len())
-		req.Header.Set("Content-Type", "application/json")
 	}
 
-	req.Header.Set("Accept", "application/json")
+	req, err := retryablehttp.NewRequest(method, u.String(), body)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, fn := range options {
+		if fn == nil {
+			continue
+		}
 
-	switch c.authType {
-	case basicAuth, oAuthToken:
-		req.Header.Set("Authorization", "Bearer "+c.token)
-	case privateToken:
-		req.Header.Set("PRIVATE-TOKEN", c.token)
+		if err := fn(req); err != nil {
+			return nil, err
+		}
 	}
 
-	if c.UserAgent != "" {
-		req.Header.Set("User-Agent", c.UserAgent)
+	// Set the request specific headers.
+	for k, v := range reqHeaders {
+		req.Header[k] = v
 	}
 
 	return req, nil
@@ -733,7 +853,12 @@ func (r *Response) populatePageValues() {
 // error if an API error has occurred. If v implements the io.Writer
 // interface, the raw response body will be written to v, without attempting to
 // first decode it.
-func (c *Client) Do(req *http.Request, v interface{}) (*Response, error) {
+func (c *Client) Do(req *retryablehttp.Request, v interface{}) (*Response, error) {
+	// Wait will block until the limiter can obtain a new token.
+	if err := c.limiter.Wait(req.Context()); err != nil {
+		return nil, err
+	}
+
 	resp, err := c.client.Do(req)
 	if err != nil {
 		return nil, err
@@ -752,8 +877,8 @@ func (c *Client) Do(req *http.Request, v interface{}) (*Response, error) {
 
 	err = CheckResponse(resp)
 	if err != nil {
-		// even though there was an error, we still return the response
-		// in case the caller wants to inspect it further
+		// Even though there was an error, we still return the response
+		// in case the caller wants to inspect it further.
 		return response, err
 	}
 
@@ -872,11 +997,11 @@ func parseError(raw interface{}) string {
 // another user, provided your private token is from an administrator account.
 //
 // GitLab docs: https://docs.gitlab.com/ce/api/README.html#sudo
-type OptionFunc func(*http.Request) error
+type OptionFunc func(*retryablehttp.Request) error
 
 // WithSudo takes either a username or user ID and sets the SUDO request header
 func WithSudo(uid interface{}) OptionFunc {
-	return func(req *http.Request) error {
+	return func(req *retryablehttp.Request) error {
 		user, err := parseID(uid)
 		if err != nil {
 			return err
@@ -888,7 +1013,7 @@ func WithSudo(uid interface{}) OptionFunc {
 
 // WithContext runs the request with the provided context
 func WithContext(ctx context.Context) OptionFunc {
-	return func(req *http.Request) error {
+	return func(req *retryablehttp.Request) error {
 		*req = *req.WithContext(ctx)
 		return nil
 	}

gitlab_test.go

@@ -13,6 +13,8 @@ import (
 	"os"
 	"strings"
 	"testing"
+
+	retryablehttp "github.com/hashicorp/go-retryablehttp"
 )
 
 // setup sets up a test HTTP server along with a gitlab.Client that is
@@ -73,7 +75,7 @@ func mustWriteHTTPResponse(t *testing.T, w io.Writer, fixturePath string) {
 	}
 }
 
-func errorOption(*http.Request) error {
+func errorOption(*retryablehttp.Request) error {
 	return errors.New("OptionFunc returns an error")
 }
 
@@ -93,8 +95,8 @@ func TestSetBaseURL(t *testing.T) {
 	expectedBaseURL := "http://gitlab.local/foo/" + apiVersionPath
 	c := NewClient(nil, "")
 	err := c.SetBaseURL("http://gitlab.local/foo")
-	if err != nil {
-		t.Fatalf("Failed to SetBaseURL: %v", err)
+	if err == nil {
+		t.Errorf("Expected a 'no such host' error, got: %v", err)
 	}
 	if c.BaseURL().String() != expectedBaseURL {
 		t.Errorf("BaseURL is %s, want %s", c.BaseURL().String(), expectedBaseURL)
@@ -108,7 +110,7 @@ func TestCheckResponse(t *testing.T) {
 	}
 
 	resp := &http.Response{
-		Request:    req,
+		Request:    req.Request,
 		StatusCode: http.StatusBadRequest,
 		Body: ioutil.NopCloser(strings.NewReader(`
 		{

go.mod

@@ -2,10 +2,13 @@ module github.com/xanzy/go-gitlab
 
 require (
 	github.com/google/go-querystring v1.0.0
+	github.com/hashicorp/go-cleanhttp v0.5.1
+	github.com/hashicorp/go-retryablehttp v0.6.4
 	github.com/stretchr/testify v1.4.0
 	golang.org/x/net v0.0.0-20181108082009-03003ca0c849 // indirect
 	golang.org/x/oauth2 v0.0.0-20181106182150-f42d05182288
 	golang.org/x/sync v0.0.0-20181108010431-42b317875d0f // indirect
+	golang.org/x/time v0.0.0-20191024005414-555d28b269f0
 	google.golang.org/appengine v1.3.0 // indirect
 )