github.com/ActiveState/cli@v0.0.0-20240508170324-6801f60cd051/scripts/internal/workflow-helpers/ghtransport.go (about)

     1  package workflow_helpers
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"net/http"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/ActiveState/cli/internal/logging"
    11  	"github.com/google/go-github/v45/github"
    12  )
    13  
    14  /*
    15  Code based on: https://github.com/hashicorp/terraform-provider-github/blob/fa73654b66e37b1fd8d886141d9c2974e24ba42f/github/transport.go#L42-L109
    16  Sourced via: https://github.com/google/go-github/issues/431
    17  
    18  MIT License
    19  
    20  Copyright (c) 2020 GitHub
    21  
    22  Permission is hereby granted, free of charge, to any person obtaining a copy
    23  of this software and associated documentation files (the "Software"), to deal
    24  in the Software without restriction, including without limitation the rights
    25  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    26  copies of the Software, and to permit persons to whom the Software is
    27  furnished to do so, subject to the following conditions:
    28  
    29  The above copyright notice and this permission notice shall be included in all
    30  copies or substantial portions of the Software.
    31  
    32  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    33  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    34  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    35  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    36  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    37  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    38  SOFTWARE.
    39  */
    40  
    41  const (
    42  	ctxId      = "id"
    43  	writeDelay = 720 * time.Millisecond // https://github.com/google/go-github/issues/431#issuecomment-248767702
    44  	// The above writeDelay is the minimum interval to avoid rate limiting as long as is only one
    45  	// job using the GitHub API at a time. Since this is likely not the case for us, we need a tunable
    46  	// parameter that can easily be adjusted based on experience.
    47  	multiplier = 1.5 // 150% of the minimum delay
    48  )
    49  
    50  // rateLimitTransport implements GitHub's best practices
    51  // for avoiding rate limits
    52  // https://developer.github.com/v3/guides/best-practices-for-integrators/#dealing-with-abuse-rate-limits
    53  type rateLimitTransport struct {
    54  	transport        http.RoundTripper
    55  	delayNextRequest bool
    56  	accessToken      string
    57  
    58  	m sync.Mutex
    59  }
    60  
    61  func (rlt *rateLimitTransport) RoundTrip(req *http.Request) (*http.Response, error) {
    62  	req.Header.Set("Authorization", "Bearer "+rlt.accessToken)
    63  	// Make requests for a single user or client ID serially
    64  	// This is also necessary for safely saving
    65  	// and restoring bodies between retries below
    66  	rlt.lock(req)
    67  
    68  	// If you're making a large number of POST, PATCH, PUT, or DELETE requests
    69  	// for a single user or client ID, wait at least one second between each request.
    70  	if rlt.delayNextRequest {
    71  		delay := time.Duration(multiplier * float64(writeDelay))
    72  		logging.Debug("Sleeping %s between write operations", delay)
    73  		time.Sleep(delay)
    74  	}
    75  
    76  	rlt.delayNextRequest = isWriteMethod(req.Method)
    77  
    78  	resp, err := rlt.transport.RoundTrip(req)
    79  	if err != nil {
    80  		rlt.unlock(req)
    81  		return resp, err
    82  	}
    83  
    84  	// Make response body accessible for retries & debugging
    85  	// (work around bug in GitHub SDK)
    86  	// See https://github.com/google/go-github/pull/986
    87  	r1, r2, err := drainBody(resp.Body)
    88  	if err != nil {
    89  		return nil, err
    90  	}
    91  	resp.Body = r1
    92  	ghErr := github.CheckResponse(resp)
    93  	resp.Body = r2
    94  
    95  	// When you have been limited, use the Retry-After response header to slow down.
    96  	if arlErr, ok := ghErr.(*github.AbuseRateLimitError); ok {
    97  		rlt.delayNextRequest = false
    98  		retryAfter := arlErr.GetRetryAfter()
    99  		logging.Debug("Abuse detection mechanism triggered, sleeping for %s before retrying",
   100  			retryAfter)
   101  		time.Sleep(retryAfter)
   102  		rlt.unlock(req)
   103  		return rlt.RoundTrip(req)
   104  	}
   105  
   106  	if rlErr, ok := ghErr.(*github.RateLimitError); ok {
   107  		rlt.delayNextRequest = false
   108  		retryAfter := time.Until(rlErr.Rate.Reset.Time)
   109  		logging.Debug("Rate limit %d reached, sleeping for %s (until %s) before retrying",
   110  			rlErr.Rate.Limit, retryAfter, time.Now().Add(retryAfter))
   111  		time.Sleep(retryAfter)
   112  		rlt.unlock(req)
   113  		return rlt.RoundTrip(req)
   114  	}
   115  
   116  	rlt.unlock(req)
   117  
   118  	return resp, nil
   119  }
   120  
   121  func (rlt *rateLimitTransport) lock(req *http.Request) {
   122  	ctx := req.Context()
   123  	logging.Debug("[TRACE] Aquiring lock for GitHub API request (%q)", ctx.Value(ctxId))
   124  	rlt.m.Lock()
   125  }
   126  
   127  func (rlt *rateLimitTransport) unlock(req *http.Request) {
   128  	ctx := req.Context()
   129  	logging.Debug("[TRACE] Releasing lock for GitHub API request (%q)", ctx.Value(ctxId))
   130  	rlt.m.Unlock()
   131  }
   132  
   133  func NewRateLimitTransport(rt http.RoundTripper, accessToken string) *rateLimitTransport {
   134  	return &rateLimitTransport{transport: rt, accessToken: accessToken}
   135  }
   136  
   137  // drainBody reads all of b to memory and then returns two equivalent
   138  // ReadClosers yielding the same bytes.
   139  func drainBody(b io.ReadCloser) (r1, r2 io.ReadCloser, err error) {
   140  	if b == http.NoBody {
   141  		// No copying needed. Preserve the magic sentinel meaning of NoBody.
   142  		return http.NoBody, http.NoBody, nil
   143  	}
   144  	var buf bytes.Buffer
   145  	if _, err = buf.ReadFrom(b); err != nil {
   146  		return nil, b, err
   147  	}
   148  	if err = b.Close(); err != nil {
   149  		return nil, b, err
   150  	}
   151  	return io.NopCloser(&buf), io.NopCloser(bytes.NewReader(buf.Bytes())), nil
   152  }
   153  
   154  func isWriteMethod(method string) bool {
   155  	switch method {
   156  	case "POST", "PATCH", "PUT", "DELETE":
   157  		return true
   158  	}
   159  	return false
   160  }