forked from holygeek/traefik
-
Notifications
You must be signed in to change notification settings - Fork 0
/
retry.go
168 lines (139 loc) · 4.86 KB
/
retry.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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
package middlewares
import (
"bufio"
"context"
"io/ioutil"
"net"
"net/http"
"github.com/containous/traefik/log"
)
// Compile time validation that the response writer implements http interfaces correctly.
var _ Stateful = &retryResponseWriterWithCloseNotify{}
// Retry is a middleware that retries requests
type Retry struct {
attempts int
next http.Handler
listener RetryListener
}
// NewRetry returns a new Retry instance
func NewRetry(attempts int, next http.Handler, listener RetryListener) *Retry {
return &Retry{
attempts: attempts,
next: next,
listener: listener,
}
}
func (retry *Retry) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
// if we might make multiple attempts, swap the body for an ioutil.NopCloser
// cf https://github.com/containous/traefik/issues/1008
if retry.attempts > 1 {
body := r.Body
defer body.Close()
r.Body = ioutil.NopCloser(body)
}
attempts := 1
for {
netErrorOccurred := false
// We pass in a pointer to netErrorOccurred so that we can set it to true on network errors
// when proxying the HTTP requests to the backends. This happens in the custom RecordingErrorHandler.
newCtx := context.WithValue(r.Context(), defaultNetErrCtxKey, &netErrorOccurred)
retryResponseWriter := newRetryResponseWriter(rw, attempts >= retry.attempts, &netErrorOccurred)
retry.next.ServeHTTP(retryResponseWriter, r.WithContext(newCtx))
if !retryResponseWriter.ShouldRetry() {
break
}
attempts++
log.Debugf("New attempt %d for request: %v", attempts, r.URL)
retry.listener.Retried(r, attempts)
}
}
// netErrorCtxKey is a custom type that is used as key for the context.
type netErrorCtxKey string
// defaultNetErrCtxKey is the actual key which value is used to record network errors.
var defaultNetErrCtxKey netErrorCtxKey = "NetErrCtxKey"
// NetErrorRecorder is an interface to record net errors.
type NetErrorRecorder interface {
// Record can be used to signal the retry middleware that an network error happened
// and therefore the request should be retried.
Record(ctx context.Context)
}
// DefaultNetErrorRecorder is the default NetErrorRecorder implementation.
type DefaultNetErrorRecorder struct{}
// Record is recording network errors by setting the context value for the defaultNetErrCtxKey to true.
func (DefaultNetErrorRecorder) Record(ctx context.Context) {
val := ctx.Value(defaultNetErrCtxKey)
if netErrorOccurred, isBoolPointer := val.(*bool); isBoolPointer {
*netErrorOccurred = true
}
}
// RetryListener is used to inform about retry attempts.
type RetryListener interface {
// Retried will be called when a retry happens, with the request attempt passed to it.
// For the first retry this will be attempt 2.
Retried(req *http.Request, attempt int)
}
// RetryListeners is a convenience type to construct a list of RetryListener and notify
// each of them about a retry attempt.
type RetryListeners []RetryListener
// Retried exists to implement the RetryListener interface. It calls Retried on each of its slice entries.
func (l RetryListeners) Retried(req *http.Request, attempt int) {
for _, retryListener := range l {
retryListener.Retried(req, attempt)
}
}
type retryResponseWriter interface {
http.ResponseWriter
http.Flusher
ShouldRetry() bool
}
func newRetryResponseWriter(rw http.ResponseWriter, attemptsExhausted bool, netErrorOccured *bool) retryResponseWriter {
responseWriter := &retryResponseWriterWithoutCloseNotify{
responseWriter: rw,
attemptsExhausted: attemptsExhausted,
netErrorOccured: netErrorOccured,
}
if _, ok := rw.(http.CloseNotifier); ok {
return &retryResponseWriterWithCloseNotify{responseWriter}
}
return responseWriter
}
type retryResponseWriterWithoutCloseNotify struct {
responseWriter http.ResponseWriter
attemptsExhausted bool
netErrorOccured *bool
}
func (rr *retryResponseWriterWithoutCloseNotify) ShouldRetry() bool {
return *rr.netErrorOccured && !rr.attemptsExhausted
}
func (rr *retryResponseWriterWithoutCloseNotify) Header() http.Header {
if rr.ShouldRetry() {
return make(http.Header)
}
return rr.responseWriter.Header()
}
func (rr *retryResponseWriterWithoutCloseNotify) Write(buf []byte) (int, error) {
if rr.ShouldRetry() {
return 0, nil
}
return rr.responseWriter.Write(buf)
}
func (rr *retryResponseWriterWithoutCloseNotify) WriteHeader(code int) {
if rr.ShouldRetry() {
return
}
rr.responseWriter.WriteHeader(code)
}
func (rr *retryResponseWriterWithoutCloseNotify) Hijack() (net.Conn, *bufio.ReadWriter, error) {
return rr.responseWriter.(http.Hijacker).Hijack()
}
func (rr *retryResponseWriterWithoutCloseNotify) Flush() {
if flusher, ok := rr.responseWriter.(http.Flusher); ok {
flusher.Flush()
}
}
type retryResponseWriterWithCloseNotify struct {
*retryResponseWriterWithoutCloseNotify
}
func (rr *retryResponseWriterWithCloseNotify) CloseNotify() <-chan bool {
return rr.responseWriter.(http.CloseNotifier).CloseNotify()
}