github.com/shashidharatd/test-infra@v0.0.0-20171006011030-71304e1ca560/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  		} else {
   116  			glog.Errorf("error %s: %v. Will retry.\n", action, err)
   117  			c.sleepForAttempt(retryCount)
   118  		}
   119  	}
   120  	return resp, err
   121  }
   122  
   123  // depaginate adds depagination on top of the retry and rate limiting logic provided by retry.
   124  func (c *Client) depaginate(action string, opts *github.ListOptions, call func() ([]interface{}, *github.Response, error)) ([]interface{}, error) {
   125  	var allItems []interface{}
   126  	wrapper := func() (*github.Response, error) {
   127  		items, resp, err := call()
   128  		if err == nil {
   129  			allItems = append(allItems, items...)
   130  		}
   131  		return resp, err
   132  	}
   133  
   134  	opts.Page = 1
   135  	opts.PerPage = 100
   136  	lastPage := 1
   137  	for ; opts.Page <= lastPage; opts.Page++ {
   138  		resp, err := c.retry(action, wrapper)
   139  		if err != nil {
   140  			return allItems, fmt.Errorf("error while depaginating page %d/%d: %v", opts.Page, lastPage, err)
   141  		}
   142  		if resp.LastPage > 0 {
   143  			lastPage = resp.LastPage
   144  		}
   145  	}
   146  	return allItems, nil
   147  }