traefik/middlewares/retry.go
Marco Jantke f79317a435 retry only on real network errors
Now retries only happen when actual network errors occur and not only
anymore based on the HTTP status code. This is because the backend could
also send this status codes as their normal interface and in that case
we don't want to retry.
2017-06-19 20:13:46 +02:00

166 lines
4.7 KiB
Go

package middlewares
import (
"bufio"
"bytes"
"context"
"io/ioutil"
"net"
"net/http"
"github.com/containous/traefik/log"
"github.com/vulcand/oxy/utils"
)
// Compile time validation responseRecorder implements http interfaces correctly.
var (
_ Stateful = &retryResponseRecorder{}
)
// 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)
recorder := newRetryResponseRecorder()
recorder.responseWriter = rw
retry.next.ServeHTTP(recorder, r.WithContext(newCtx))
if !netErrorOccurred || attempts >= retry.attempts {
utils.CopyHeaders(rw.Header(), recorder.Header())
rw.WriteHeader(recorder.Code)
rw.Write(recorder.Body.Bytes())
break
}
attempts++
log.Debugf("New attempt %d for request: %v", attempts, r.URL)
retry.listener.Retried(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(attempt int)
}
// retryResponseRecorder is an implementation of http.ResponseWriter that
// records its mutations for later inspection.
type retryResponseRecorder struct {
Code int // the HTTP response code from WriteHeader
HeaderMap http.Header // the HTTP response headers
Body *bytes.Buffer // if non-nil, the bytes.Buffer to append written data to
responseWriter http.ResponseWriter
err error
}
// newRetryResponseRecorder returns an initialized retryResponseRecorder.
func newRetryResponseRecorder() *retryResponseRecorder {
return &retryResponseRecorder{
HeaderMap: make(http.Header),
Body: new(bytes.Buffer),
Code: 200,
}
}
// Header returns the response headers.
func (rw *retryResponseRecorder) Header() http.Header {
m := rw.HeaderMap
if m == nil {
m = make(http.Header)
rw.HeaderMap = m
}
return m
}
// Write always succeeds and writes to rw.Body, if not nil.
func (rw *retryResponseRecorder) Write(buf []byte) (int, error) {
if rw.err != nil {
return 0, rw.err
}
return rw.Body.Write(buf)
}
// WriteHeader sets rw.Code.
func (rw *retryResponseRecorder) WriteHeader(code int) {
rw.Code = code
}
// Hijack hijacks the connection
func (rw *retryResponseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
return rw.responseWriter.(http.Hijacker).Hijack()
}
// CloseNotify returns a channel that receives at most a
// single value (true) when the client connection has gone
// away.
func (rw *retryResponseRecorder) CloseNotify() <-chan bool {
return rw.responseWriter.(http.CloseNotifier).CloseNotify()
}
// Flush sends any buffered data to the client.
func (rw *retryResponseRecorder) Flush() {
_, err := rw.responseWriter.Write(rw.Body.Bytes())
if err != nil {
log.Errorf("Error writing response in retryResponseRecorder: %s", err)
rw.err = err
}
rw.Body.Reset()
flusher, ok := rw.responseWriter.(http.Flusher)
if ok {
flusher.Flush()
}
}