Enable the retry middleware in the proxy

Co-authored-by: Romain <rtribotte@users.noreply.github.com>

pkg/middlewares/retry/retry.go

@@ -37,6 +37,48 @@ type Listener interface {
 // each of them about a retry attempt.
 type Listeners []Listener
 
+// Retried exists to implement the Listener interface. It calls Retried on each of its slice entries.
+func (l Listeners) Retried(req *http.Request, attempt int) {
+	for _, listener := range l {
+		listener.Retried(req, attempt)
+	}
+}
+
+type shouldRetryContextKey struct{}
+
+// ShouldRetry is a function allowing to enable/disable the retry middleware mechanism.
+type ShouldRetry func(shouldRetry bool)
+
+// ContextShouldRetry returns the ShouldRetry function if it has been set by the Retry middleware in the chain.
+func ContextShouldRetry(ctx context.Context) ShouldRetry {
+	f, _ := ctx.Value(shouldRetryContextKey{}).(ShouldRetry)
+	return f
+}
+
+// WrapHandler wraps a given http.Handler to inject the httptrace.ClientTrace in the request context when it is needed
+// by the retry middleware.
+func WrapHandler(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		if shouldRetry := ContextShouldRetry(req.Context()); shouldRetry != nil {
+			shouldRetry(true)
+
+			trace := &httptrace.ClientTrace{
+				WroteHeaders: func() {
+					shouldRetry(false)
+				},
+				WroteRequest: func(httptrace.WroteRequestInfo) {
+					shouldRetry(false)
+				},
+			}
+			newCtx := httptrace.WithClientTrace(req.Context(), trace)
+			next.ServeHTTP(rw, req.WithContext(newCtx))
+			return
+		}
+
+		next.ServeHTTP(rw, req)
+	})
+}
+
 // retry is a middleware that retries requests.
 type retry struct {
 	attempts        int
@@ -83,19 +125,13 @@ func (r *retry) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	attempts := 1
 
 	operation := func() error {
-		shouldRetry := attempts < r.attempts
-		retryResponseWriter := newResponseWriter(rw, shouldRetry)
+		remainAttempts := attempts < r.attempts
+		retryResponseWriter := newResponseWriter(rw)
 
-		// Disable retries when the backend already received request data
-		trace := &httptrace.ClientTrace{
-			WroteHeaders: func() {
-				retryResponseWriter.DisableRetries()
-			},
-			WroteRequest: func(httptrace.WroteRequestInfo) {
-				retryResponseWriter.DisableRetries()
-			},
+		var shouldRetry ShouldRetry = func(shouldRetry bool) {
+			retryResponseWriter.SetShouldRetry(remainAttempts && shouldRetry)
 		}
-		newCtx := httptrace.WithClientTrace(req.Context(), trace)
+		newCtx := context.WithValue(req.Context(), shouldRetryContextKey{}, shouldRetry)
 
 		r.next.ServeHTTP(retryResponseWriter, req.Clone(newCtx))
 
@@ -142,25 +178,17 @@ func (r *retry) newBackOff() backoff.BackOff {
 	return b
 }
 
-// Retried exists to implement the Listener interface. It calls Retried on each of its slice entries.
-func (l Listeners) Retried(req *http.Request, attempt int) {
-	for _, listener := range l {
-		listener.Retried(req, attempt)
-	}
-}
-
 type responseWriter interface {
 	http.ResponseWriter
 	http.Flusher
 	ShouldRetry() bool
-	DisableRetries()
+	SetShouldRetry(shouldRetry bool)
 }
 
-func newResponseWriter(rw http.ResponseWriter, shouldRetry bool) responseWriter {
+func newResponseWriter(rw http.ResponseWriter) responseWriter {
 	responseWriter := &responseWriterWithoutCloseNotify{
 		responseWriter: rw,
 		headers:        make(http.Header),
-		shouldRetry:    shouldRetry,
 	}
 	if _, ok := rw.(http.CloseNotifier); ok {
 		return &responseWriterWithCloseNotify{
@@ -181,8 +209,8 @@ func (r *responseWriterWithoutCloseNotify) ShouldRetry() bool {
 	return r.shouldRetry
 }
 
-func (r *responseWriterWithoutCloseNotify) DisableRetries() {
-	r.shouldRetry = false
+func (r *responseWriterWithoutCloseNotify) SetShouldRetry(shouldRetry bool) {
+	r.shouldRetry = shouldRetry
 }
 
 func (r *responseWriterWithoutCloseNotify) Header() http.Header {
@@ -193,7 +221,7 @@ func (r *responseWriterWithoutCloseNotify) Header() http.Header {
 }
 
 func (r *responseWriterWithoutCloseNotify) Write(buf []byte) (int, error) {
-	if r.ShouldRetry() {
+	if r.shouldRetry {
 		return len(buf), nil
 	}
 	if !r.written {
@@ -203,16 +231,7 @@ func (r *responseWriterWithoutCloseNotify) Write(buf []byte) (int, error) {
 }
 
 func (r *responseWriterWithoutCloseNotify) WriteHeader(code int) {
-	if r.ShouldRetry() && code == http.StatusServiceUnavailable {
-		// We get a 503 HTTP Status Code when there is no backend server in the pool
-		// to which the request could be sent.  Also, note that r.ShouldRetry()
-		// will never return true in case there was a connection established to
-		// the backend server and so we can be sure that the 503 was produced
-		// inside Traefik already and we don't have to retry in this cases.
-		r.DisableRetries()
-	}
-
-	if r.ShouldRetry() || r.written {
+	if r.shouldRetry || r.written {
 		return
 	}