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