github.com/munnerz/test-infra@v0.0.0-20190108210205-ce3d181dc989/pkg/ghclient/core.go (about)

     1  /*
     2  Copyright 2017 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  // Package ghclient provides a github client that wraps go-github with retry logic, rate limiting,
    18  // and depagination where necessary.
    19  package ghclient
    20  
    21  import (
    22  	"fmt"
    23  	"math"
    24  	"net/http"
    25  	"time"
    26  
    27  	"github.com/golang/glog"
    28  
    29  	"github.com/google/go-github/github"
    30  	"golang.org/x/oauth2"
    31  )
    32  
    33  // Client is an augmentation of the go-github client that adds retry logic, rate limiting, and pagination
    34  // handling to applicable the client functions.
    35  type Client struct {
    36  	issueService issueService
    37  	prService    pullRequestService
    38  	repoService  repositoryService
    39  	userService  usersService
    40  
    41  	retries             int
    42  	retryInitialBackoff time.Duration
    43  
    44  	tokenReserve int
    45  	dryRun       bool
    46  }
    47  
    48  // NewClient makes a new Client with the specified token and dry-run status.
    49  func NewClient(token string, dryRun bool) *Client {
    50  	httpClient := &http.Client{
    51  		Transport: &oauth2.Transport{
    52  			Base:   http.DefaultTransport,
    53  			Source: oauth2.ReuseTokenSource(nil, oauth2.StaticTokenSource(&oauth2.Token{AccessToken: token})),
    54  		},
    55  	}
    56  	client := github.NewClient(httpClient)
    57  	return &Client{
    58  		issueService:        client.Issues,
    59  		prService:           client.PullRequests,
    60  		repoService:         client.Repositories,
    61  		userService:         client.Users,
    62  		retries:             5,
    63  		retryInitialBackoff: time.Second,
    64  		tokenReserve:        50,
    65  		dryRun:              dryRun,
    66  	}
    67  }
    68  
    69  func (c *Client) sleepForAttempt(retryCount int) {
    70  	maxDelay := 20 * time.Second
    71  	delay := c.retryInitialBackoff * time.Duration(math.Exp2(float64(retryCount)))
    72  	if delay > maxDelay {
    73  		delay = maxDelay
    74  	}
    75  	time.Sleep(delay)
    76  }
    77  
    78  func (c *Client) limitRate(r *github.Rate) {
    79  	if r.Remaining <= c.tokenReserve {
    80  		sleepDuration := time.Until(r.Reset.Time) + (time.Second * 10)
    81  		if sleepDuration > 0 {
    82  			glog.Infof("--Rate Limiting-- Tokens reached minimum reserve %d. Sleeping until reset in %v.\n", c.tokenReserve, sleepDuration)
    83  			time.Sleep(sleepDuration)
    84  		}
    85  	}
    86  }
    87  
    88  type retryAbort struct{ error }
    89  
    90  func (r *retryAbort) Error() string {
    91  	return fmt.Sprintf("aborting retry loop: %v", r.error)
    92  }
    93  
    94  // retry handles rate limiting and retry logic for a github API call.
    95  func (c *Client) retry(action string, call func() (*github.Response, error)) (*github.Response, error) {
    96  	var err error
    97  	var resp *github.Response
    98  
    99  	for retryCount := 0; retryCount <= c.retries; retryCount++ {
   100  		if resp, err = call(); err == nil {
   101  			c.limitRate(&resp.Rate)
   102  			return resp, nil
   103  		}
   104  		switch err := err.(type) {
   105  		case *github.RateLimitError:
   106  			c.limitRate(&err.Rate)
   107  		case *github.TwoFactorAuthError:
   108  			return resp, err
   109  		case *retryAbort:
   110  			return resp, err
   111  		}
   112  
   113  		if retryCount == c.retries {
   114  			return resp, err
   115  		}
   116  		glog.Errorf("error %s: %v. Will retry.\n", action, err)
   117  		c.sleepForAttempt(retryCount)
   118  	}
   119  	return resp, err
   120  }
   121  
   122  // depaginate adds depagination on top of the retry and rate limiting logic provided by retry.
   123  func (c *Client) depaginate(action string, opts *github.ListOptions, call func() ([]interface{}, *github.Response, error)) ([]interface{}, error) {
   124  	var allItems []interface{}
   125  	wrapper := func() (*github.Response, error) {
   126  		items, resp, err := call()
   127  		if err == nil {
   128  			allItems = append(allItems, items...)
   129  		}
   130  		return resp, err
   131  	}
   132  
   133  	opts.Page = 1
   134  	opts.PerPage = 100
   135  	lastPage := 1
   136  	for ; opts.Page <= lastPage; opts.Page++ {
   137  		resp, err := c.retry(action, wrapper)
   138  		if err != nil {
   139  			return allItems, fmt.Errorf("error while depaginating page %d/%d: %v", opts.Page, lastPage, err)
   140  		}
   141  		if resp.LastPage > 0 {
   142  			lastPage = resp.LastPage
   143  		}
   144  	}
   145  	return allItems, nil
   146  }