healthcheck: add support at the load-balancers of services level

Co-authored-by: Dmitry Sharshakov <d3dx12.xx@gmail.com>
Co-authored-by: Julien Salleyron <julien.salleyron@gmail.com>
Co-authored-by: Jean-Baptiste Doumenjou <925513+jbdoumenjou@users.noreply.github.com>
Co-authored-by: Romain <rtribotte@users.noreply.github.com>
Co-authored-by: Tom Moulard <tom.moulard@traefik.io>

pkg/server/service/loadbalancer/wrr/wrr.go

@@ -2,6 +2,7 @@ package wrr
 
 import (
 	"container/heap"
+	"context"
 	"errors"
 	"fmt"
 	"net/http"
@@ -30,16 +31,28 @@ type stickyCookie struct {
 // Entries have deadlines set at currentDeadline + 1 / weight,
 // providing weighted round robin behavior with floating point weights and an O(log n) pick time.
 type Balancer struct {
-	stickyCookie *stickyCookie
+	stickyCookie     *stickyCookie
+	wantsHealthCheck bool
 
 	mutex       sync.RWMutex
 	handlers    []*namedHandler
 	curDeadline float64
+	// status is a record of which child services of the Balancer are healthy, keyed
+	// by name of child service. A service is initially added to the map when it is
+	// created via AddService, and it is later removed or added to the map as needed,
+	// through the SetStatus method.
+	status map[string]struct{}
+	// updaters is the list of hooks that are run (to update the Balancer
+	// parent(s)), whenever the Balancer status changes.
+	updaters []func(bool)
 }
 
 // New creates a new load balancer.
-func New(sticky *dynamic.Sticky) *Balancer {
-	balancer := &Balancer{}
+func New(sticky *dynamic.Sticky, hc *dynamic.HealthCheck) *Balancer {
+	balancer := &Balancer{
+		status:           make(map[string]struct{}),
+		wantsHealthCheck: hc != nil,
+	}
 	if sticky != nil && sticky.Cookie != nil {
 		balancer.stickyCookie = &stickyCookie{
 			name:     sticky.Cookie.Name,
@@ -81,6 +94,58 @@ func (b *Balancer) Pop() interface{} {
 	return h
 }
 
+// SetStatus sets on the balancer that its given child is now of the given
+// status. balancerName is only needed for logging purposes.
+func (b *Balancer) SetStatus(ctx context.Context, childName string, up bool) {
+	b.mutex.Lock()
+	defer b.mutex.Unlock()
+
+	upBefore := len(b.status) > 0
+
+	status := "DOWN"
+	if up {
+		status = "UP"
+	}
+	log.FromContext(ctx).Debugf("Setting status of %s to %v", childName, status)
+	if up {
+		b.status[childName] = struct{}{}
+	} else {
+		delete(b.status, childName)
+	}
+
+	upAfter := len(b.status) > 0
+	status = "DOWN"
+	if upAfter {
+		status = "UP"
+	}
+
+	// No Status Change
+	if upBefore == upAfter {
+		// We're still with the same status, no need to propagate
+		log.FromContext(ctx).Debugf("Still %s, no need to propagate", status)
+		return
+	}
+
+	// Status Change
+	log.FromContext(ctx).Debugf("Propagating new %s status", status)
+	for _, fn := range b.updaters {
+		fn(upAfter)
+	}
+}
+
+// RegisterStatusUpdater adds fn to the list of hooks that are run when the
+// status of the Balancer changes.
+// Not thread safe.
+func (b *Balancer) RegisterStatusUpdater(fn func(up bool)) error {
+	if !b.wantsHealthCheck {
+		return errors.New("healthCheck not enabled in config for this weighted service")
+	}
+	b.updaters = append(b.updaters, fn)
+	return nil
+}
+
+var errNoAvailableServer = errors.New("no available server")
+
 func (b *Balancer) nextServer() (*namedHandler, error) {
 	b.mutex.Lock()
 	defer b.mutex.Unlock()
@@ -88,15 +153,24 @@ func (b *Balancer) nextServer() (*namedHandler, error) {
 	if len(b.handlers) == 0 {
 		return nil, fmt.Errorf("no servers in the pool")
 	}
+	if len(b.status) == 0 {
+		return nil, errNoAvailableServer
+	}
 
-	// Pick handler with closest deadline.
-	handler := heap.Pop(b).(*namedHandler)
+	var handler *namedHandler
+	for {
+		// Pick handler with closest deadline.
+		handler = heap.Pop(b).(*namedHandler)
 
-	// curDeadline should be handler's deadline so that new added entry would have a fair competition environment with the old ones.
-	b.curDeadline = handler.deadline
-	handler.deadline += 1 / handler.weight
+		// curDeadline should be handler's deadline so that new added entry would have a fair competition environment with the old ones.
+		b.curDeadline = handler.deadline
+		handler.deadline += 1 / handler.weight
 
-	heap.Push(b, handler)
+		heap.Push(b, handler)
+		if _, ok := b.status[handler.name]; ok {
+			break
+		}
+	}
 
 	log.WithoutContext().Debugf("Service selected by WRR: %s", handler.name)
 	return handler, nil
@@ -112,17 +186,32 @@ func (b *Balancer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 
 		if err == nil && cookie != nil {
 			for _, handler := range b.handlers {
-				if handler.name == cookie.Value {
-					handler.ServeHTTP(w, req)
-					return
+				if handler.name != cookie.Value {
+					continue
 				}
+
+				b.mutex.RLock()
+				_, ok := b.status[handler.name]
+				b.mutex.RUnlock()
+				if !ok {
+					// because we already are in the only iteration that matches the cookie, so none
+					// of the following iterations are going to be a match for the cookie anyway.
+					break
+				}
+
+				handler.ServeHTTP(w, req)
+				return
 			}
 		}
 	}
 
 	server, err := b.nextServer()
 	if err != nil {
-		http.Error(w, http.StatusText(http.StatusInternalServerError)+err.Error(), http.StatusInternalServerError)
+		if errors.Is(err, errNoAvailableServer) {
+			http.Error(w, errNoAvailableServer.Error(), http.StatusServiceUnavailable)
+		} else {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+		}
 		return
 	}
 
@@ -135,7 +224,6 @@ func (b *Balancer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 }
 
 // AddService adds a handler.
-// It is not thread safe with ServeHTTP.
 // A handler with a non-positive weight is ignored.
 func (b *Balancer) AddService(name string, handler http.Handler, weight *int) {
 	w := 1
@@ -148,10 +236,9 @@ func (b *Balancer) AddService(name string, handler http.Handler, weight *int) {
 
 	h := &namedHandler{Handler: handler, name: name, weight: float64(w)}
 
-	// use RWLock to protect b.curDeadline
-	b.mutex.RLock()
+	b.mutex.Lock()
 	h.deadline = b.curDeadline + 1/h.weight
-	b.mutex.RUnlock()
-
 	heap.Push(b, h)
+	b.status[name] = struct{}{}
+	b.mutex.Unlock()
 }

pkg/server/service/loadbalancer/wrr/wrr_test.go

@@ -1,6 +1,7 @@
 package wrr
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -15,16 +16,18 @@ 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 TestBalancer(t *testing.T) {
-	balancer := New(nil)
+	balancer := New(nil, nil)
 
 	balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 		rw.Header().Set("server", "first")
@@ -46,7 +49,7 @@ func TestBalancer(t *testing.T) {
 }
 
 func TestBalancerNoService(t *testing.T) {
-	balancer := New(nil)
+	balancer := New(nil, nil)
 
 	recorder := httptest.NewRecorder()
 	balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
@@ -55,7 +58,7 @@ func TestBalancerNoService(t *testing.T) {
 }
 
 func TestBalancerOneServerZeroWeight(t *testing.T) {
-	balancer := New(nil)
+	balancer := New(nil, nil)
 
 	balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 		rw.Header().Set("server", "first")
@@ -72,8 +75,155 @@ func TestBalancerOneServerZeroWeight(t *testing.T) {
 	assert.Equal(t, 3, recorder.save["first"])
 }
 
+type key string
+
+const serviceName key = "serviceName"
+
+func TestBalancerNoServiceUp(t *testing.T) {
+	balancer := New(nil, nil)
+
+	balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.WriteHeader(http.StatusInternalServerError)
+	}), Int(1))
+
+	balancer.AddService("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.WriteHeader(http.StatusInternalServerError)
+	}), Int(1))
+
+	balancer.SetStatus(context.WithValue(context.Background(), serviceName, "parent"), "first", false)
+	balancer.SetStatus(context.WithValue(context.Background(), serviceName, "parent"), "second", false)
+
+	recorder := httptest.NewRecorder()
+	balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+
+	assert.Equal(t, http.StatusServiceUnavailable, recorder.Result().StatusCode)
+}
+
+func TestBalancerOneServerDown(t *testing.T) {
+	balancer := New(nil, nil)
+
+	balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "first")
+		rw.WriteHeader(http.StatusOK)
+	}), Int(1))
+
+	balancer.AddService("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.WriteHeader(http.StatusInternalServerError)
+	}), Int(1))
+	balancer.SetStatus(context.WithValue(context.Background(), serviceName, "parent"), "second", false)
+
+	recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	for i := 0; i < 3; i++ {
+		balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+	}
+
+	assert.Equal(t, 3, recorder.save["first"])
+}
+
+func TestBalancerDownThenUp(t *testing.T) {
+	balancer := New(nil, nil)
+
+	balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "first")
+		rw.WriteHeader(http.StatusOK)
+	}), Int(1))
+
+	balancer.AddService("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "second")
+		rw.WriteHeader(http.StatusOK)
+	}), Int(1))
+	balancer.SetStatus(context.WithValue(context.Background(), serviceName, "parent"), "second", false)
+
+	recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	for i := 0; i < 3; i++ {
+		balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+	}
+	assert.Equal(t, 3, recorder.save["first"])
+
+	balancer.SetStatus(context.WithValue(context.Background(), serviceName, "parent"), "second", true)
+	recorder = &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	for i := 0; i < 2; i++ {
+		balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+	}
+	assert.Equal(t, 1, recorder.save["first"])
+	assert.Equal(t, 1, recorder.save["second"])
+}
+
+func TestBalancerPropagate(t *testing.T) {
+	balancer1 := New(nil, &dynamic.HealthCheck{})
+
+	balancer1.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "first")
+		rw.WriteHeader(http.StatusOK)
+	}), Int(1))
+	balancer1.AddService("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "second")
+		rw.WriteHeader(http.StatusOK)
+	}), Int(1))
+
+	balancer2 := New(nil, &dynamic.HealthCheck{})
+	balancer2.AddService("third", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "third")
+		rw.WriteHeader(http.StatusOK)
+	}), Int(1))
+	balancer2.AddService("fourth", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "fourth")
+		rw.WriteHeader(http.StatusOK)
+	}), Int(1))
+
+	topBalancer := New(nil, &dynamic.HealthCheck{})
+	topBalancer.AddService("balancer1", balancer1, Int(1))
+	_ = balancer1.RegisterStatusUpdater(func(up bool) {
+		topBalancer.SetStatus(context.WithValue(context.Background(), serviceName, "top"), "balancer1", up)
+		// TODO(mpl): if test gets flaky, add channel or something here to signal that
+		// propagation is done, and wait on it before sending request.
+	})
+	topBalancer.AddService("balancer2", balancer2, Int(1))
+	_ = balancer2.RegisterStatusUpdater(func(up bool) {
+		topBalancer.SetStatus(context.WithValue(context.Background(), serviceName, "top"), "balancer2", up)
+	})
+
+	recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	for i := 0; i < 8; i++ {
+		topBalancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+	}
+	assert.Equal(t, 2, recorder.save["first"])
+	assert.Equal(t, 2, recorder.save["second"])
+	assert.Equal(t, 2, recorder.save["third"])
+	assert.Equal(t, 2, recorder.save["fourth"])
+	wantStatus := []int{200, 200, 200, 200, 200, 200, 200, 200}
+	assert.Equal(t, wantStatus, recorder.status)
+
+	// fourth gets downed, but balancer2 still up since third is still up.
+	balancer2.SetStatus(context.WithValue(context.Background(), serviceName, "top"), "fourth", false)
+	recorder = &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	for i := 0; i < 8; i++ {
+		topBalancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+	}
+	assert.Equal(t, 2, recorder.save["first"])
+	assert.Equal(t, 2, recorder.save["second"])
+	assert.Equal(t, 4, recorder.save["third"])
+	assert.Equal(t, 0, recorder.save["fourth"])
+	wantStatus = []int{200, 200, 200, 200, 200, 200, 200, 200}
+	assert.Equal(t, wantStatus, recorder.status)
+
+	// third gets downed, and the propagation triggers balancer2 to be marked as
+	// down as well for topBalancer.
+	balancer2.SetStatus(context.WithValue(context.Background(), serviceName, "top"), "third", false)
+	recorder = &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	for i := 0; i < 8; i++ {
+		topBalancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+	}
+	assert.Equal(t, 4, recorder.save["first"])
+	assert.Equal(t, 4, recorder.save["second"])
+	assert.Equal(t, 0, recorder.save["third"])
+	assert.Equal(t, 0, recorder.save["fourth"])
+	wantStatus = []int{200, 200, 200, 200, 200, 200, 200, 200}
+	assert.Equal(t, wantStatus, recorder.status)
+}
+
 func TestBalancerAllServersZeroWeight(t *testing.T) {
-	balancer := New(nil)
+	balancer := New(nil, nil)
 
 	balancer.AddService("test", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), Int(0))
 	balancer.AddService("test2", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), Int(0))
@@ -87,7 +237,7 @@ func TestBalancerAllServersZeroWeight(t *testing.T) {
 func TestSticky(t *testing.T) {
 	balancer := New(&dynamic.Sticky{
 		Cookie: &dynamic.Cookie{Name: "test"},
-	})
+	}, nil)
 
 	balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 		rw.Header().Set("server", "first")
@@ -118,7 +268,7 @@ func TestSticky(t *testing.T) {
 // TestBalancerBias makes sure that the WRR algorithm spreads elements evenly right from the start,
 // and that it does not "over-favor" the high-weighted ones with a biased start-up regime.
 func TestBalancerBias(t *testing.T) {
-	balancer := New(nil)
+	balancer := New(nil, nil)
 
 	balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 		rw.Header().Set("server", "A")