Move code to pkg

2019-03-15 09:42:03 +01:00 · 2019-03-15 09:42:03 +01:00 · f1b085fa36
commit f1b085fa36
parent bd4c822670
465 changed files with 656 additions and 680 deletions
--- a/pkg/middlewares/retry/retry.go
+++ b/pkg/middlewares/retry/retry.go
@ -0,0 +1,207 @@
+package retry
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"io/ioutil"
+	"net"
+	"net/http"
+	"net/http/httptrace"
+
+	"github.com/containous/traefik/pkg/config"
+	"github.com/containous/traefik/pkg/middlewares"
+	"github.com/containous/traefik/pkg/tracing"
+	"github.com/opentracing/opentracing-go/ext"
+)
+
+// Compile time validation that the response writer implements http interfaces correctly.
+var _ middlewares.Stateful = &responseWriterWithCloseNotify{}
+
+const (
+	typeName = "Retry"
+)
+
+// Listener is used to inform about retry attempts.
+type Listener 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(req *http.Request, attempt int)
+}
+
+// Listeners is a convenience type to construct a list of Listener and notify
+// each of them about a retry attempt.
+type Listeners []Listener
+
+// retry is a middleware that retries requests.
+type retry struct {
+	attempts int
+	next     http.Handler
+	listener Listener
+	name     string
+}
+
+// New returns a new retry middleware.
+func New(ctx context.Context, next http.Handler, config config.Retry, listener Listener, name string) (http.Handler, error) {
+	logger := middlewares.GetLogger(ctx, name, typeName)
+	logger.Debug("Creating middleware")
+
+	if config.Attempts <= 0 {
+		return nil, fmt.Errorf("incorrect (or empty) value for attempt (%d)", config.Attempts)
+	}
+
+	return &retry{
+		attempts: config.Attempts,
+		next:     next,
+		listener: listener,
+		name:     name,
+	}, nil
+}
+
+func (r *retry) GetTracingInformation() (string, ext.SpanKindEnum) {
+	return r.name, tracing.SpanKindNoneEnum
+}
+
+func (r *retry) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
+	// if we might make multiple attempts, swap the body for an ioutil.NopCloser
+	// cf https://github.com/containous/traefik/issues/1008
+	if r.attempts > 1 {
+		body := req.Body
+		defer body.Close()
+		req.Body = ioutil.NopCloser(body)
+	}
+
+	attempts := 1
+	for {
+		shouldRetry := attempts < r.attempts
+		retryResponseWriter := newResponseWriter(rw, shouldRetry)
+
+		// Disable retries when the backend already received request data
+		trace := &httptrace.ClientTrace{
+			WroteHeaders: func() {
+				retryResponseWriter.DisableRetries()
+			},
+			WroteRequest: func(httptrace.WroteRequestInfo) {
+				retryResponseWriter.DisableRetries()
+			},
+		}
+		newCtx := httptrace.WithClientTrace(req.Context(), trace)
+
+		r.next.ServeHTTP(retryResponseWriter, req.WithContext(newCtx))
+
+		if !retryResponseWriter.ShouldRetry() {
+			break
+		}
+
+		attempts++
+		logger := middlewares.GetLogger(req.Context(), r.name, typeName)
+		logger.Debugf("New attempt %d for request: %v", attempts, req.URL)
+		r.listener.Retried(req, attempts)
+	}
+}
+
+// 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()
+}
+
+func newResponseWriter(rw http.ResponseWriter, shouldRetry bool) responseWriter {
+	responseWriter := &responseWriterWithoutCloseNotify{
+		responseWriter: rw,
+		headers:        make(http.Header),
+		shouldRetry:    shouldRetry,
+	}
+	if _, ok := rw.(http.CloseNotifier); ok {
+		return &responseWriterWithCloseNotify{
+			responseWriterWithoutCloseNotify: responseWriter,
+		}
+	}
+	return responseWriter
+}
+
+type responseWriterWithoutCloseNotify struct {
+	responseWriter http.ResponseWriter
+	headers        http.Header
+	shouldRetry    bool
+	written        bool
+}
+
+func (r *responseWriterWithoutCloseNotify) ShouldRetry() bool {
+	return r.shouldRetry
+}
+
+func (r *responseWriterWithoutCloseNotify) DisableRetries() {
+	r.shouldRetry = false
+}
+
+func (r *responseWriterWithoutCloseNotify) Header() http.Header {
+	if r.written {
+		return r.responseWriter.Header()
+	}
+	return r.headers
+}
+
+func (r *responseWriterWithoutCloseNotify) Write(buf []byte) (int, error) {
+	if r.ShouldRetry() {
+		return len(buf), nil
+	}
+	return r.responseWriter.Write(buf)
+}
+
+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() {
+		return
+	}
+
+	// In that case retry case is set to false which means we at least managed
+	// to write headers to the backend : we are not going to perform any further retry.
+	// So it is now safe to alter current response headers with headers collected during
+	// the latest try before writing headers to client.
+	headers := r.responseWriter.Header()
+	for header, value := range r.headers {
+		headers[header] = value
+	}
+
+	r.responseWriter.WriteHeader(code)
+	r.written = true
+}
+
+func (r *responseWriterWithoutCloseNotify) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	hijacker, ok := r.responseWriter.(http.Hijacker)
+	if !ok {
+		return nil, nil, fmt.Errorf("%T is not a http.Hijacker", r.responseWriter)
+	}
+	return hijacker.Hijack()
+}
+
+func (r *responseWriterWithoutCloseNotify) Flush() {
+	if flusher, ok := r.responseWriter.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+type responseWriterWithCloseNotify struct {
+	*responseWriterWithoutCloseNotify
+}
+
+func (r *responseWriterWithCloseNotify) CloseNotify() <-chan bool {
+	return r.responseWriter.(http.CloseNotifier).CloseNotify()
+}
--- a/pkg/middlewares/retry/retry_test.go
+++ b/pkg/middlewares/retry/retry_test.go
@ -0,0 +1,312 @@
+package retry
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/http/httptrace"
+	"strings"
+	"testing"
+
+	"github.com/containous/traefik/pkg/config"
+	"github.com/containous/traefik/pkg/middlewares/emptybackendhandler"
+	"github.com/containous/traefik/pkg/testhelpers"
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/vulcand/oxy/forward"
+	"github.com/vulcand/oxy/roundrobin"
+)
+
+func TestRetry(t *testing.T) {
+	testCases := []struct {
+		desc                  string
+		config                config.Retry
+		wantRetryAttempts     int
+		wantResponseStatus    int
+		amountFaultyEndpoints int
+	}{
+		{
+			desc:                  "no retry on success",
+			config:                config.Retry{Attempts: 1},
+			wantRetryAttempts:     0,
+			wantResponseStatus:    http.StatusOK,
+			amountFaultyEndpoints: 0,
+		},
+		{
+			desc:                  "no retry when max request attempts is one",
+			config:                config.Retry{Attempts: 1},
+			wantRetryAttempts:     0,
+			wantResponseStatus:    http.StatusInternalServerError,
+			amountFaultyEndpoints: 1,
+		},
+		{
+			desc:                  "one retry when one server is faulty",
+			config:                config.Retry{Attempts: 2},
+			wantRetryAttempts:     1,
+			wantResponseStatus:    http.StatusOK,
+			amountFaultyEndpoints: 1,
+		},
+		{
+			desc:                  "two retries when two servers are faulty",
+			config:                config.Retry{Attempts: 3},
+			wantRetryAttempts:     2,
+			wantResponseStatus:    http.StatusOK,
+			amountFaultyEndpoints: 2,
+		},
+		{
+			desc:                  "max attempts exhausted delivers the 5xx response",
+			config:                config.Retry{Attempts: 3},
+			wantRetryAttempts:     2,
+			wantResponseStatus:    http.StatusInternalServerError,
+			amountFaultyEndpoints: 3,
+		},
+	}
+
+	backendServer := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.WriteHeader(http.StatusOK)
+		_, err := rw.Write([]byte("OK"))
+		if err != nil {
+			http.Error(rw, err.Error(), http.StatusInternalServerError)
+		}
+	}))
+
+	forwarder, err := forward.New()
+	require.NoError(t, err)
+
+	for _, test := range testCases {
+		test := test
+		t.Run(test.desc, func(t *testing.T) {
+			t.Parallel()
+
+			loadBalancer, err := roundrobin.New(forwarder)
+			require.NoError(t, err)
+
+			basePort := 33444
+			for i := 0; i < test.amountFaultyEndpoints; i++ {
+				// 192.0.2.0 is a non-routable IP for testing purposes.
+				// See: https://stackoverflow.com/questions/528538/non-routable-ip-address/18436928#18436928
+				// We only use the port specification here because the URL is used as identifier
+				// in the load balancer and using the exact same URL would not add a new server.
+				err = loadBalancer.UpsertServer(testhelpers.MustParseURL("http://192.0.2.0:" + string(basePort+i)))
+				require.NoError(t, err)
+			}
+
+			// add the functioning server to the end of the load balancer list
+			err = loadBalancer.UpsertServer(testhelpers.MustParseURL(backendServer.URL))
+			require.NoError(t, err)
+
+			retryListener := &countingRetryListener{}
+			retry, err := New(context.Background(), loadBalancer, test.config, retryListener, "traefikTest")
+			require.NoError(t, err)
+
+			recorder := httptest.NewRecorder()
+			req := httptest.NewRequest(http.MethodGet, "http://localhost:3000/ok", nil)
+
+			retry.ServeHTTP(recorder, req)
+
+			assert.Equal(t, test.wantResponseStatus, recorder.Code)
+			assert.Equal(t, test.wantRetryAttempts, retryListener.timesCalled)
+		})
+	}
+}
+
+func TestRetryEmptyServerList(t *testing.T) {
+	forwarder, err := forward.New()
+	require.NoError(t, err)
+
+	loadBalancer, err := roundrobin.New(forwarder)
+	require.NoError(t, err)
+
+	// The EmptyBackend middleware ensures that there is a 503
+	// response status set when there is no backend server in the pool.
+	next := emptybackendhandler.New(loadBalancer)
+
+	retryListener := &countingRetryListener{}
+	retry, err := New(context.Background(), next, config.Retry{Attempts: 3}, retryListener, "traefikTest")
+	require.NoError(t, err)
+
+	recorder := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodGet, "http://localhost:3000/ok", nil)
+
+	retry.ServeHTTP(recorder, req)
+
+	assert.Equal(t, http.StatusServiceUnavailable, recorder.Code)
+	assert.Equal(t, 0, retryListener.timesCalled)
+}
+
+func TestRetryListeners(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	retryListeners := Listeners{&countingRetryListener{}, &countingRetryListener{}}
+
+	retryListeners.Retried(req, 1)
+	retryListeners.Retried(req, 1)
+
+	for _, retryListener := range retryListeners {
+		listener := retryListener.(*countingRetryListener)
+		if listener.timesCalled != 2 {
+			t.Errorf("retry listener was called %d time(s), want %d time(s)", listener.timesCalled, 2)
+		}
+	}
+}
+
+func TestMultipleRetriesShouldNotLooseHeaders(t *testing.T) {
+	attempt := 0
+	expectedHeaderName := "X-Foo-Test-2"
+	expectedHeaderValue := "bar"
+
+	next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		headerName := fmt.Sprintf("X-Foo-Test-%d", attempt)
+		rw.Header().Add(headerName, expectedHeaderValue)
+		if attempt < 2 {
+			attempt++
+			return
+		}
+
+		// Request has been successfully written to backend
+		trace := httptrace.ContextClientTrace(req.Context())
+		trace.WroteHeaders()
+
+		// And we decide to answer to client
+		rw.WriteHeader(http.StatusNoContent)
+	})
+
+	retry, err := New(context.Background(), next, config.Retry{Attempts: 3}, &countingRetryListener{}, "traefikTest")
+	require.NoError(t, err)
+
+	responseRecorder := httptest.NewRecorder()
+	retry.ServeHTTP(responseRecorder, testhelpers.MustNewRequest(http.MethodGet, "http://test", http.NoBody))
+
+	headerValue := responseRecorder.Header().Get(expectedHeaderName)
+
+	// Validate if we have the correct header
+	if headerValue != expectedHeaderValue {
+		t.Errorf("Expected to have %s for header %s, got %s", expectedHeaderValue, expectedHeaderName, headerValue)
+	}
+
+	// Validate that we don't have headers from previous attempts
+	for i := 0; i < attempt; i++ {
+		headerName := fmt.Sprintf("X-Foo-Test-%d", i)
+		headerValue = responseRecorder.Header().Get("headerName")
+		if headerValue != "" {
+			t.Errorf("Expected no value for header %s, got %s", headerName, headerValue)
+		}
+	}
+}
+
+// countingRetryListener is a Listener implementation to count the times the Retried fn is called.
+type countingRetryListener struct {
+	timesCalled int
+}
+
+func (l *countingRetryListener) Retried(req *http.Request, attempt int) {
+	l.timesCalled++
+}
+
+func TestRetryWithFlush(t *testing.T) {
+	next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.WriteHeader(200)
+		_, err := rw.Write([]byte("FULL "))
+		if err != nil {
+			http.Error(rw, err.Error(), http.StatusInternalServerError)
+		}
+		rw.(http.Flusher).Flush()
+		_, err = rw.Write([]byte("DATA"))
+		if err != nil {
+			http.Error(rw, err.Error(), http.StatusInternalServerError)
+		}
+	})
+
+	retry, err := New(context.Background(), next, config.Retry{Attempts: 1}, &countingRetryListener{}, "traefikTest")
+	require.NoError(t, err)
+
+	responseRecorder := httptest.NewRecorder()
+
+	retry.ServeHTTP(responseRecorder, &http.Request{})
+
+	assert.Equal(t, "FULL DATA", responseRecorder.Body.String())
+}
+
+func TestRetryWebsocket(t *testing.T) {
+	testCases := []struct {
+		desc                   string
+		maxRequestAttempts     int
+		expectedRetryAttempts  int
+		expectedResponseStatus int
+		expectedError          bool
+		amountFaultyEndpoints  int
+	}{
+		{
+			desc:                   "Switching ok after 2 retries",
+			maxRequestAttempts:     3,
+			expectedRetryAttempts:  2,
+			amountFaultyEndpoints:  2,
+			expectedResponseStatus: http.StatusSwitchingProtocols,
+		},
+		{
+			desc:                   "Switching failed",
+			maxRequestAttempts:     2,
+			expectedRetryAttempts:  1,
+			amountFaultyEndpoints:  2,
+			expectedResponseStatus: http.StatusBadGateway,
+			expectedError:          true,
+		},
+	}
+
+	forwarder, err := forward.New()
+	if err != nil {
+		t.Fatalf("Error creating forwarder: %v", err)
+	}
+
+	backendServer := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		upgrader := websocket.Upgrader{}
+		_, err := upgrader.Upgrade(rw, req, nil)
+		if err != nil {
+			http.Error(rw, err.Error(), http.StatusInternalServerError)
+		}
+	}))
+
+	for _, test := range testCases {
+		test := test
+		t.Run(test.desc, func(t *testing.T) {
+			t.Parallel()
+
+			loadBalancer, err := roundrobin.New(forwarder)
+			if err != nil {
+				t.Fatalf("Error creating load balancer: %v", err)
+			}
+
+			basePort := 33444
+			for i := 0; i < test.amountFaultyEndpoints; i++ {
+				// 192.0.2.0 is a non-routable IP for testing purposes.
+				// See: https://stackoverflow.com/questions/528538/non-routable-ip-address/18436928#18436928
+				// We only use the port specification here because the URL is used as identifier
+				// in the load balancer and using the exact same URL would not add a new server.
+				_ = loadBalancer.UpsertServer(testhelpers.MustParseURL("http://192.0.2.0:" + string(basePort+i)))
+			}
+
+			// add the functioning server to the end of the load balancer list
+			err = loadBalancer.UpsertServer(testhelpers.MustParseURL(backendServer.URL))
+			if err != nil {
+				t.Fatalf("Fail to upsert server: %v", err)
+			}
+
+			retryListener := &countingRetryListener{}
+			retryH, err := New(context.Background(), loadBalancer, config.Retry{Attempts: test.maxRequestAttempts}, retryListener, "traefikTest")
+			require.NoError(t, err)
+
+			retryServer := httptest.NewServer(retryH)
+
+			url := strings.Replace(retryServer.URL, "http", "ws", 1)
+			_, response, err := websocket.DefaultDialer.Dial(url, nil)
+
+			if !test.expectedError {
+				require.NoError(t, err)
+			}
+
+			assert.Equal(t, test.expectedResponseStatus, response.StatusCode)
+			assert.Equal(t, test.expectedRetryAttempts, retryListener.timesCalled)
+		})
+	}
+}