k8s.io/test-infra@v0.0.0-20240520184403-27c6b4c223d8/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  	"net/url"
    26  	"strings"
    27  	"time"
    28  
    29  	"github.com/golang/glog"
    30  
    31  	"github.com/google/go-github/github"
    32  	"golang.org/x/oauth2"
    33  )
    34  
    35  // Client is an augmentation of the go-github client that adds retry logic, rate limiting, and pagination
    36  // handling to applicable the client functions.
    37  type Client struct {
    38  	issueService issueService
    39  	prService    pullRequestService
    40  	repoService  repositoryService
    41  	userService  usersService
    42  
    43  	retries             int
    44  	retryInitialBackoff time.Duration
    45  
    46  	tokenReserve int
    47  	dryRun       bool
    48  }
    49  
    50  // NewClient makes a new Client with the specified token and dry-run status.
    51  func NewClient(token string, dryRun bool) *Client {
    52  	return NewClientWithEndpoint("https://api.github.com", token, dryRun)
    53  }
    54  
    55  // NewClientWithEndpoint makes a new Client with the provided endpoint.
    56  func NewClientWithEndpoint(endpoint string, token string, dryRun bool) *Client {
    57  	baseURL, err := url.Parse(endpoint)
    58  	if err != nil {
    59  		glog.Fatalf("invalid github endpoint %s: %s", endpoint, err)
    60  	}
    61  	if !strings.HasSuffix(baseURL.Path, "/") {
    62  		baseURL.Path += "/"
    63  	}
    64  
    65  	httpClient := &http.Client{
    66  		Transport: &oauth2.Transport{
    67  			Base:   http.DefaultTransport,
    68  			Source: oauth2.ReuseTokenSource(nil, oauth2.StaticTokenSource(&oauth2.Token{AccessToken: token})),
    69  		},
    70  	}
    71  	client := github.NewClient(httpClient)
    72  	client.BaseURL = baseURL
    73  
    74  	return &Client{
    75  		issueService:        client.Issues,
    76  		prService:           client.PullRequests,
    77  		repoService:         client.Repositories,
    78  		userService:         client.Users,
    79  		retries:             5,
    80  		retryInitialBackoff: time.Second,
    81  		tokenReserve:        50,
    82  		dryRun:              dryRun,
    83  	}
    84  }
    85  
    86  func (c *Client) sleepForAttempt(retryCount int) {
    87  	maxDelay := 20 * time.Second
    88  	delay := c.retryInitialBackoff * time.Duration(math.Exp2(float64(retryCount)))
    89  	if delay > maxDelay {
    90  		delay = maxDelay
    91  	}
    92  	time.Sleep(delay)
    93  }
    94  
    95  func (c *Client) limitRate(r *github.Rate) {
    96  	if r.Remaining <= c.tokenReserve {
    97  		sleepDuration := time.Until(r.Reset.Time) + (time.Second * 10)
    98  		if sleepDuration > 0 {
    99  			glog.Infof("--Rate Limiting-- Tokens reached minimum reserve %d. Sleeping until reset in %v.\n", c.tokenReserve, sleepDuration)
   100  			time.Sleep(sleepDuration)
   101  		}
   102  	}
   103  }
   104  
   105  type retryAbort struct{ error }
   106  
   107  func (r *retryAbort) Error() string {
   108  	return fmt.Sprintf("aborting retry loop: %v", r.error)
   109  }
   110  
   111  // retry handles rate limiting and retry logic for a github API call.
   112  func (c *Client) retry(action string, call func() (*github.Response, error)) (*github.Response, error) {
   113  	var err error
   114  	var resp *github.Response
   115  
   116  	for retryCount := 0; retryCount <= c.retries; retryCount++ {
   117  		if resp, err = call(); err == nil {
   118  			c.limitRate(&resp.Rate)
   119  			return resp, nil
   120  		}
   121  		switch err := err.(type) {
   122  		case *github.RateLimitError:
   123  			c.limitRate(&err.Rate)
   124  		case *github.TwoFactorAuthError:
   125  			return resp, err
   126  		case *retryAbort:
   127  			return resp, err
   128  		}
   129  
   130  		if retryCount == c.retries {
   131  			return resp, err
   132  		}
   133  		glog.Errorf("error %s: %v. Will retry.\n", action, err)
   134  		c.sleepForAttempt(retryCount)
   135  	}
   136  	return resp, err
   137  }
   138  
   139  // depaginate adds depagination on top of the retry and rate limiting logic provided by retry.
   140  func (c *Client) depaginate(action string, opts *github.ListOptions, call func() ([]interface{}, *github.Response, error)) ([]interface{}, error) {
   141  	var allItems []interface{}
   142  	wrapper := func() (*github.Response, error) {
   143  		items, resp, err := call()
   144  		if err == nil {
   145  			allItems = append(allItems, items...)
   146  		}
   147  		return resp, err
   148  	}
   149  
   150  	opts.Page = 1
   151  	opts.PerPage = 100
   152  	lastPage := 1
   153  	for ; opts.Page <= lastPage; opts.Page++ {
   154  		resp, err := c.retry(action, wrapper)
   155  		if err != nil {
   156  			return allItems, fmt.Errorf("error while depaginating page %d/%d: %w", opts.Page, lastPage, err)
   157  		}
   158  		if resp.LastPage > 0 {
   159  			lastPage = resp.LastPage
   160  		}
   161  	}
   162  	return allItems, nil
   163  }