Fix high memory usage in retry middleware

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()
+	}
+}