Merge branch v2.11 into v3.3

pkg/middlewares/retry/retry.go

@@ -36,6 +36,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
@@ -101,19 +143,13 @@ func (r *retry) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 			req = req.WithContext(tracingCtx)
 		}
 
-		shouldRetry := attempts < r.attempts
-		retryResponseWriter := newResponseWriter(rw, shouldRetry)
+		remainAttempts := attempts < r.attempts
+		retryResponseWriter := newResponseWriter(rw)
 
-		// Disable retries when the backend already received request data
-		clientTrace := &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(), clientTrace)
+		newCtx := context.WithValue(req.Context(), shouldRetryContextKey{}, shouldRetry)
 
 		r.next.ServeHTTP(retryResponseWriter, req.Clone(newCtx))
 
@@ -164,18 +200,10 @@ 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)
-	}
-}
-
-func newResponseWriter(rw http.ResponseWriter, shouldRetry bool) *responseWriter {
+func newResponseWriter(rw http.ResponseWriter) *responseWriter {
 	return &responseWriter{
 		responseWriter: rw,
 		headers:        make(http.Header),
-		shouldRetry:    shouldRetry,
 	}
 }
 
@@ -190,8 +218,8 @@ func (r *responseWriter) ShouldRetry() bool {
 	return r.shouldRetry
 }
 
-func (r *responseWriter) DisableRetries() {
-	r.shouldRetry = false
+func (r *responseWriter) SetShouldRetry(shouldRetry bool) {
+	r.shouldRetry = shouldRetry
 }
 
 func (r *responseWriter) Header() http.Header {
@@ -205,20 +233,14 @@ func (r *responseWriter) Write(buf []byte) (int, error) {
 	if r.ShouldRetry() {
 		return len(buf), nil
 	}
+	if !r.written {
+		r.WriteHeader(http.StatusOK)
+	}
 	return r.responseWriter.Write(buf)
 }
 
 func (r *responseWriter) 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
 	}