Add Failover service

Co-authored-by: Kevin Pollet <pollet.kevin@gmail.com>

pkg/server/service/loadbalancer/failover/failover_test.go

@@ -0,0 +1,163 @@
+package failover
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/traefik/traefik/v2/pkg/config/dynamic"
+)
+
+type responseRecorder struct {
+	*httptest.ResponseRecorder
+	save     map[string]int
+	sequence []string
+	status   []int
+}
+
+func (r *responseRecorder) WriteHeader(statusCode int) {
+	r.save[r.Header().Get("server")]++
+	r.sequence = append(r.sequence, r.Header().Get("server"))
+	r.status = append(r.status, statusCode)
+	r.ResponseRecorder.WriteHeader(statusCode)
+}
+
+func TestFailover(t *testing.T) {
+	failover := New(&dynamic.HealthCheck{})
+
+	status := true
+	require.NoError(t, failover.RegisterStatusUpdater(func(up bool) {
+		status = up
+	}))
+
+	failover.SetHandler(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "handler")
+		rw.WriteHeader(http.StatusOK)
+	}))
+
+	failover.SetFallbackHandler(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "fallback")
+		rw.WriteHeader(http.StatusOK)
+	}))
+
+	recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	failover.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+
+	assert.Equal(t, 1, recorder.save["handler"])
+	assert.Equal(t, 0, recorder.save["fallback"])
+	assert.Equal(t, []int{200}, recorder.status)
+	assert.True(t, status)
+
+	failover.SetHandlerStatus(context.Background(), false)
+
+	recorder = &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	failover.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+
+	assert.Equal(t, 0, recorder.save["handler"])
+	assert.Equal(t, 1, recorder.save["fallback"])
+	assert.Equal(t, []int{200}, recorder.status)
+	assert.True(t, status)
+
+	failover.SetFallbackHandlerStatus(context.Background(), false)
+
+	recorder = &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	failover.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+
+	assert.Equal(t, 0, recorder.save["handler"])
+	assert.Equal(t, 0, recorder.save["fallback"])
+	assert.Equal(t, []int{503}, recorder.status)
+	assert.False(t, status)
+}
+
+func TestFailoverDownThenUp(t *testing.T) {
+	failover := New(nil)
+
+	failover.SetHandler(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "handler")
+		rw.WriteHeader(http.StatusOK)
+	}))
+
+	failover.SetFallbackHandler(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "fallback")
+		rw.WriteHeader(http.StatusOK)
+	}))
+
+	recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	failover.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+
+	assert.Equal(t, 1, recorder.save["handler"])
+	assert.Equal(t, 0, recorder.save["fallback"])
+	assert.Equal(t, []int{200}, recorder.status)
+
+	failover.SetHandlerStatus(context.Background(), false)
+
+	recorder = &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	failover.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+
+	assert.Equal(t, 0, recorder.save["handler"])
+	assert.Equal(t, 1, recorder.save["fallback"])
+	assert.Equal(t, []int{200}, recorder.status)
+
+	failover.SetHandlerStatus(context.Background(), true)
+
+	recorder = &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	failover.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+
+	assert.Equal(t, 1, recorder.save["handler"])
+	assert.Equal(t, 0, recorder.save["fallback"])
+	assert.Equal(t, []int{200}, recorder.status)
+}
+
+func TestFailoverPropagate(t *testing.T) {
+	failover := New(&dynamic.HealthCheck{})
+	failover.SetHandler(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "handler")
+		rw.WriteHeader(http.StatusOK)
+	}))
+	failover.SetFallbackHandler(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "fallback")
+		rw.WriteHeader(http.StatusOK)
+	}))
+
+	topFailover := New(nil)
+	topFailover.SetHandler(failover)
+	topFailover.SetFallbackHandler(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "topFailover")
+		rw.WriteHeader(http.StatusOK)
+	}))
+	err := failover.RegisterStatusUpdater(func(up bool) {
+		topFailover.SetHandlerStatus(context.Background(), up)
+	})
+	require.NoError(t, err)
+
+	recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	topFailover.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+
+	assert.Equal(t, 1, recorder.save["handler"])
+	assert.Equal(t, 0, recorder.save["fallback"])
+	assert.Equal(t, 0, recorder.save["topFailover"])
+	assert.Equal(t, []int{200}, recorder.status)
+
+	failover.SetHandlerStatus(context.Background(), false)
+
+	recorder = &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	topFailover.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+
+	assert.Equal(t, 0, recorder.save["handler"])
+	assert.Equal(t, 1, recorder.save["fallback"])
+	assert.Equal(t, 0, recorder.save["topFailover"])
+	assert.Equal(t, []int{200}, recorder.status)
+
+	failover.SetFallbackHandlerStatus(context.Background(), false)
+
+	recorder = &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	topFailover.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+
+	assert.Equal(t, 0, recorder.save["handler"])
+	assert.Equal(t, 0, recorder.save["fallback"])
+	assert.Equal(t, 1, recorder.save["topFailover"])
+	assert.Equal(t, []int{200}, recorder.status)
+}