forked from integrations/terraform-provider-github
/
transport.go
150 lines (126 loc) · 3.91 KB
/
transport.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
package github
import (
"bytes"
"io"
"io/ioutil"
"log"
"net/http"
"sync"
"time"
"github.com/google/go-github/github"
)
const (
ctxEtag = "etag"
ctxId = "id"
writeDelay = 1 * time.Second
)
// etagTransport allows saving API quota by passing previously stored Etag
// available via context to request headers
type etagTransport struct {
transport http.RoundTripper
}
func (ett *etagTransport) RoundTrip(req *http.Request) (*http.Response, error) {
ctx := req.Context()
etag := ctx.Value(ctxEtag)
if v, ok := etag.(string); ok {
req.Header.Set("If-None-Match", v)
}
return ett.transport.RoundTrip(req)
}
func NewEtagTransport(rt http.RoundTripper) *etagTransport {
return &etagTransport{transport: rt}
}
// rateLimitTransport implements GitHub's best practices
// for avoiding rate limits
// https://developer.github.com/v3/guides/best-practices-for-integrators/#dealing-with-abuse-rate-limits
type rateLimitTransport struct {
transport http.RoundTripper
delayNextRequest bool
responseBody []byte
m sync.Mutex
}
func (rlt *rateLimitTransport) RoundTrip(req *http.Request) (*http.Response, error) {
// Make requests for a single user or client ID serially
// This is also necessary for safely saving
// and restoring bodies between retries below
rlt.lock(req)
// If you're making a large number of POST, PATCH, PUT, or DELETE requests
// for a single user or client ID, wait at least one second between each request.
if rlt.delayNextRequest {
log.Printf("[DEBUG] Sleeping %s between write operations", writeDelay)
time.Sleep(writeDelay)
}
rlt.delayNextRequest = isWriteMethod(req.Method)
resp, err := rlt.transport.RoundTrip(req)
if err != nil {
rlt.unlock(req)
return resp, err
}
// Make response body accessible for retries & debugging
// (work around bug in GitHub SDK)
// See https://github.com/google/go-github/pull/986
r1, r2, err := drainBody(resp.Body)
if err != nil {
return nil, err
}
resp.Body = r1
ghErr := github.CheckResponse(resp)
resp.Body = r2
// When you have been limited, use the Retry-After response header to slow down.
if arlErr, ok := ghErr.(*github.AbuseRateLimitError); ok {
rlt.delayNextRequest = false
retryAfter := arlErr.GetRetryAfter()
log.Printf("[DEBUG] Abuse detection mechanism triggered, sleeping for %s before retrying",
retryAfter)
time.Sleep(retryAfter)
rlt.unlock(req)
return rlt.RoundTrip(req)
}
if rlErr, ok := ghErr.(*github.RateLimitError); ok {
rlt.delayNextRequest = false
retryAfter := rlErr.Rate.Reset.Sub(time.Now())
log.Printf("[DEBUG] Rate limit %d reached, sleeping for %s (until %s) before retrying",
rlErr.Rate.Limit, retryAfter, time.Now().Add(retryAfter))
time.Sleep(retryAfter)
rlt.unlock(req)
return rlt.RoundTrip(req)
}
rlt.unlock(req)
return resp, nil
}
func (rlt *rateLimitTransport) lock(req *http.Request) {
ctx := req.Context()
log.Printf("[TRACE] Aquiring lock for GitHub API request (%q)", ctx.Value(ctxId))
rlt.m.Lock()
}
func (rlt *rateLimitTransport) unlock(req *http.Request) {
ctx := req.Context()
log.Printf("[TRACE] Releasing lock for GitHub API request (%q)", ctx.Value(ctxId))
rlt.m.Unlock()
}
func NewRateLimitTransport(rt http.RoundTripper) *rateLimitTransport {
return &rateLimitTransport{transport: rt}
}
// drainBody reads all of b to memory and then returns two equivalent
// ReadClosers yielding the same bytes.
func drainBody(b io.ReadCloser) (r1, r2 io.ReadCloser, err error) {
if b == http.NoBody {
// No copying needed. Preserve the magic sentinel meaning of NoBody.
return http.NoBody, http.NoBody, nil
}
var buf bytes.Buffer
if _, err = buf.ReadFrom(b); err != nil {
return nil, b, err
}
if err = b.Close(); err != nil {
return nil, b, err
}
return ioutil.NopCloser(&buf), ioutil.NopCloser(bytes.NewReader(buf.Bytes())), nil
}
func isWriteMethod(method string) bool {
switch method {
case "POST", "PATCH", "PUT", "DELETE":
return true
}
return false
}