Fix high memory usage in retry middleware

integration/fixtures/retry/simple.toml

@@ -0,0 +1,25 @@
+defaultEntryPoints = ["http"]
+
+logLevel = "DEBUG"
+
+[entryPoints]
+  [entryPoints.http]
+  address = ":8000"
+
+[api]
+
+[retry]
+
+[file]
+[backends]
+  [backends.backend1]
+    [backends.backend1.servers.server1]
+    url = "http://{{.WhoamiEndpoint}}:8080" # not valid
+    [backends.backend1.servers.server2]
+    url = "http://{{.WhoamiEndpoint}}:80"
+
+[frontends]
+  [frontends.frontend1]
+  backend = "backend1"
+    [frontends.frontend1.routes.test_1]
+    rule = "PathPrefix:/"

integration/integration_test.go

@@ -54,6 +54,7 @@ func init() {
 		check.Suite(&MarathonSuite{})
 		check.Suite(&MesosSuite{})
 		check.Suite(&RateLimitSuite{})
+		check.Suite(&RetrySuite{})
 		check.Suite(&SimpleSuite{})
 		check.Suite(&TimeoutSuite{})
 		check.Suite(&TracingSuite{})

integration/resources/compose/retry.yml

@@ -0,0 +1,2 @@
+whoami:
+  image: emilevauge/whoami

integration/retry_test.go

@@ -0,0 +1,40 @@
+package integration
+
+import (
+	"net/http"
+	"os"
+	"time"
+
+	"github.com/containous/traefik/integration/try"
+	"github.com/go-check/check"
+	checker "github.com/vdemeester/shakers"
+)
+
+type RetrySuite struct{ BaseSuite }
+
+func (s *RetrySuite) SetUpSuite(c *check.C) {
+	s.createComposeProject(c, "retry")
+	s.composeProject.Start(c)
+}
+
+func (s *RetrySuite) TestRetry(c *check.C) {
+	whoamiEndpoint := s.composeProject.Container(c, "whoami").NetworkSettings.IPAddress
+	file := s.adaptFile(c, "fixtures/retry/simple.toml", struct {
+		WhoamiEndpoint string
+	}{whoamiEndpoint})
+	defer os.Remove(file)
+
+	cmd, display := s.traefikCmd(withConfigFile(file))
+	defer display(c)
+	err := cmd.Start()
+	c.Assert(err, checker.IsNil)
+	defer cmd.Process.Kill()
+
+	err = try.GetRequest("http://127.0.0.1:8080/api/providers", 60*time.Second, try.BodyContains("PathPrefix:/"))
+	c.Assert(err, checker.IsNil)
+
+	// This simulates a DialTimeout when connecting to the backend server.
+	response, err := http.Get("http://127.0.0.1:8000/")
+	c.Assert(err, checker.IsNil)
+	c.Assert(response.StatusCode, checker.Equals, http.StatusOK)
+}

middlewares/error_pages.go

@@ -1,6 +1,9 @@
 package middlewares
 
 import (
+	"bufio"
+	"bytes"
+	"net"
 	"net/http"
 	"strconv"
 	"strings"
@@ -11,6 +14,9 @@ import (
 	"github.com/vulcand/oxy/utils"
 )
 
+// Compile time validation that the response recorder implements http interfaces correctly.
+var _ Stateful = &errorPagesResponseRecorderWithCloseNotify{}
+
 //ErrorPagesHandler is a middleware that provides the custom error pages
 type ErrorPagesHandler struct {
 	HTTPCodeRanges     [][2]int
@@ -52,7 +58,7 @@ func NewErrorPagesHandler(errorPage *types.ErrorPage, backendURL string) (*Error
 }
 
 func (ep *ErrorPagesHandler) ServeHTTP(w http.ResponseWriter, req *http.Request, next http.HandlerFunc) {
-	recorder := newRetryResponseRecorder(w)
+	recorder := newErrorPagesResponseRecorder(w)
 
 	next.ServeHTTP(recorder, req)
 
@@ -75,3 +81,108 @@ func (ep *ErrorPagesHandler) ServeHTTP(w http.ResponseWriter, req *http.Request,
 	utils.CopyHeaders(w.Header(), recorder.Header())
 	w.Write(recorder.GetBody().Bytes())
 }
+
+type errorPagesResponseRecorder interface {
+	http.ResponseWriter
+	http.Flusher
+	GetCode() int
+	GetBody() *bytes.Buffer
+	IsStreamingResponseStarted() bool
+}
+
+// newErrorPagesResponseRecorder returns an initialized responseRecorder.
+func newErrorPagesResponseRecorder(rw http.ResponseWriter) errorPagesResponseRecorder {
+	recorder := &errorPagesResponseRecorderWithoutCloseNotify{
+		HeaderMap:      make(http.Header),
+		Body:           new(bytes.Buffer),
+		Code:           http.StatusOK,
+		responseWriter: rw,
+	}
+	if _, ok := rw.(http.CloseNotifier); ok {
+		return &errorPagesResponseRecorderWithCloseNotify{recorder}
+	}
+	return recorder
+}
+
+// errorPagesResponseRecorderWithoutCloseNotify is an implementation of http.ResponseWriter that
+// records its mutations for later inspection.
+type errorPagesResponseRecorderWithoutCloseNotify 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
+	streamingResponseStarted bool
+}
+
+type errorPagesResponseRecorderWithCloseNotify struct {
+	*errorPagesResponseRecorderWithoutCloseNotify
+}
+
+// CloseNotify returns a channel that receives at most a
+// single value (true) when the client connection has gone
+// away.
+func (rw *errorPagesResponseRecorderWithCloseNotify) CloseNotify() <-chan bool {
+	return rw.responseWriter.(http.CloseNotifier).CloseNotify()
+}
+
+// Header returns the response headers.
+func (rw *errorPagesResponseRecorderWithoutCloseNotify) Header() http.Header {
+	m := rw.HeaderMap
+	if m == nil {
+		m = make(http.Header)
+		rw.HeaderMap = m
+	}
+	return m
+}
+
+func (rw *errorPagesResponseRecorderWithoutCloseNotify) GetCode() int {
+	return rw.Code
+}
+
+func (rw *errorPagesResponseRecorderWithoutCloseNotify) GetBody() *bytes.Buffer {
+	return rw.Body
+}
+
+func (rw *errorPagesResponseRecorderWithoutCloseNotify) IsStreamingResponseStarted() bool {
+	return rw.streamingResponseStarted
+}
+
+// Write always succeeds and writes to rw.Body, if not nil.
+func (rw *errorPagesResponseRecorderWithoutCloseNotify) Write(buf []byte) (int, error) {
+	if rw.err != nil {
+		return 0, rw.err
+	}
+	return rw.Body.Write(buf)
+}
+
+// WriteHeader sets rw.Code.
+func (rw *errorPagesResponseRecorderWithoutCloseNotify) WriteHeader(code int) {
+	rw.Code = code
+}
+
+// Hijack hijacks the connection
+func (rw *errorPagesResponseRecorderWithoutCloseNotify) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return rw.responseWriter.(http.Hijacker).Hijack()
+}
+
+// Flush sends any buffered data to the client.
+func (rw *errorPagesResponseRecorderWithoutCloseNotify) Flush() {
+	if !rw.streamingResponseStarted {
+		utils.CopyHeaders(rw.responseWriter.Header(), rw.Header())
+		rw.responseWriter.WriteHeader(rw.Code)
+		rw.streamingResponseStarted = true
+	}
+
+	_, err := rw.responseWriter.Write(rw.Body.Bytes())
+	if err != nil {
+		log.Errorf("Error writing response in responseRecorder: %s", err)
+		rw.err = err
+	}
+	rw.Body.Reset()
+	flusher, ok := rw.responseWriter.(http.Flusher)
+	if ok {
+		flusher.Flush()
+	}
+}

middlewares/error_pages_test.go

@@ -152,3 +152,51 @@ func TestErrorPageSingleCode(t *testing.T) {
 	assert.Contains(t, recorder.Body.String(), "503 Test Server")
 	assert.NotContains(t, recorder.Body.String(), "oops", "Should not return the oops page")
 }
+
+func TestNewErrorPagesResponseRecorder(t *testing.T) {
+	testCases := []struct {
+		desc     string
+		rw       http.ResponseWriter
+		expected http.ResponseWriter
+	}{
+		{
+			desc:     "Without Close Notify",
+			rw:       httptest.NewRecorder(),
+			expected: &errorPagesResponseRecorderWithoutCloseNotify{},
+		},
+		{
+			desc:     "With Close Notify",
+			rw:       &mockRWCloseNotify{},
+			expected: &errorPagesResponseRecorderWithCloseNotify{},
+		},
+	}
+
+	for _, test := range testCases {
+		test := test
+		t.Run(test.desc, func(t *testing.T) {
+			t.Parallel()
+
+			rec := newErrorPagesResponseRecorder(test.rw)
+
+			assert.IsType(t, rec, test.expected)
+		})
+	}
+}
+
+type mockRWCloseNotify struct{}
+
+func (m *mockRWCloseNotify) CloseNotify() <-chan bool {
+	panic("implement me")
+}
+
+func (m *mockRWCloseNotify) Header() http.Header {
+	panic("implement me")
+}
+
+func (m *mockRWCloseNotify) Write([]byte) (int, error) {
+	panic("implement me")
+}
+
+func (m *mockRWCloseNotify) WriteHeader(int) {
+	panic("implement me")
+}