Add p2c load-balancing strategy for servers load-balancer

Co-authored-by: Ian Ross <ifross@gmail.com> Co-authored-by: Kevin Pollet <pollet.kevin@gmail.com>
2025-03-10 12:12:04 +01:00 · 2025-03-10 12:12:04 +01:00 · 9e029a84c4
commit 9e029a84c4
parent 550d96ea67
50 changed files with 1621 additions and 382 deletions
--- a/pkg/server/service/loadbalancer/p2c/p2c.go
+++ b/pkg/server/service/loadbalancer/p2c/p2c.go
@ -0,0 +1,227 @@
+package p2c
+
+import (
+	"context"
+	"errors"
+	"math/rand"
+	"net/http"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/rs/zerolog/log"
+	"github.com/traefik/traefik/v3/pkg/config/dynamic"
+	"github.com/traefik/traefik/v3/pkg/server/service/loadbalancer"
+)
+
+type namedHandler struct {
+	http.Handler
+
+	// name is the handler name.
+	name string
+	// inflight is the number of inflight requests.
+	// It is used to implement the "power-of-two-random-choices" algorithm.
+	inflight atomic.Int64
+}
+
+func (h *namedHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
+	h.inflight.Add(1)
+	defer h.inflight.Add(-1)
+
+	h.Handler.ServeHTTP(rw, req)
+}
+
+type rnd interface {
+	Intn(n int) int
+}
+
+// Balancer implements the power-of-two-random-choices algorithm for load balancing.
+// The idea is to randomly select two of the available backends and choose the one with the fewest in-flight requests.
+// This algorithm balances the load more effectively than a round-robin approach, while maintaining a constant time for the selection:
+// The strategy also has more advantageous "herd" behavior than the "fewest connections" algorithm, especially when the load balancer
+// doesn't have perfect knowledge of the global number of connections to the backend, for example, when running in a distributed fashion.
+type Balancer struct {
+	wantsHealthCheck bool
+
+	handlersMu sync.RWMutex
+	handlers   []*namedHandler
+	// 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 Add, 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)
+	// fenced is the list of terminating yet still serving child services.
+	fenced map[string]struct{}
+
+	sticky *loadbalancer.Sticky
+
+	rand rnd
+}
+
+// New creates a new power-of-two-random-choices load balancer.
+func New(stickyConfig *dynamic.Sticky, wantsHealthCheck bool) *Balancer {
+	balancer := &Balancer{
+		status:           make(map[string]struct{}),
+		fenced:           make(map[string]struct{}),
+		wantsHealthCheck: wantsHealthCheck,
+		rand:             rand.New(rand.NewSource(time.Now().UnixNano())),
+	}
+	if stickyConfig != nil && stickyConfig.Cookie != nil {
+		balancer.sticky = loadbalancer.NewSticky(*stickyConfig.Cookie)
+	}
+
+	return balancer
+}
+
+// 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.handlersMu.Lock()
+	defer b.handlersMu.Unlock()
+
+	upBefore := len(b.status) > 0
+
+	status := "DOWN"
+	if up {
+		status = "UP"
+	}
+
+	log.Ctx(ctx).Debug().Msgf("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.Ctx(ctx).Debug().Msgf("Still %s, no need to propagate", status)
+		return
+	}
+
+	// Status Change
+	log.Ctx(ctx).Debug().Msgf("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) {
+	// We kept the same representation (map) as in the WRR strategy to improve maintainability.
+	// However, with the P2C strategy, we only need a slice of healthy servers.
+	b.handlersMu.RLock()
+	var healthy []*namedHandler
+	for _, h := range b.handlers {
+		if _, ok := b.status[h.name]; ok {
+			if _, fenced := b.fenced[h.name]; !fenced {
+				healthy = append(healthy, h)
+			}
+		}
+	}
+	b.handlersMu.RUnlock()
+
+	if len(healthy) == 0 {
+		return nil, errNoAvailableServer
+	}
+
+	// If there is only one healthy server, return it.
+	if len(healthy) == 1 {
+		return healthy[0], nil
+	}
+	// In order to not get the same backend twice, we make the second call to s.rand.IntN one fewer
+	// than the length of the slice. We then have to shift over the second index if it is equal or
+	// greater than the first index, wrapping round if needed.
+	n1, n2 := b.rand.Intn(len(healthy)), b.rand.Intn(len(healthy))
+	if n2 == n1 {
+		n2 = (n2 + 1) % len(healthy)
+	}
+
+	h1, h2 := healthy[n1], healthy[n2]
+	// Ensure h1 has fewer inflight requests than h2.
+	if h2.inflight.Load() < h1.inflight.Load() {
+		log.Debug().Msgf("Service selected by P2C: %s", h2.name)
+		return h2, nil
+	}
+
+	log.Debug().Msgf("Service selected by P2C: %s", h1.name)
+	return h1, nil
+}
+
+func (b *Balancer) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
+	if b.sticky != nil {
+		h, rewrite, err := b.sticky.StickyHandler(req)
+		if err != nil {
+			log.Error().Err(err).Msg("Error while getting sticky handler")
+		} else if h != nil {
+			if _, ok := b.status[h.Name]; ok {
+				if rewrite {
+					if err := b.sticky.WriteStickyCookie(rw, h.Name); err != nil {
+						log.Error().Err(err).Msg("Writing sticky cookie")
+					}
+				}
+
+				h.ServeHTTP(rw, req)
+				return
+			}
+		}
+	}
+
+	server, err := b.nextServer()
+	if err != nil {
+		if errors.Is(err, errNoAvailableServer) {
+			http.Error(rw, errNoAvailableServer.Error(), http.StatusServiceUnavailable)
+		} else {
+			http.Error(rw, err.Error(), http.StatusInternalServerError)
+		}
+		return
+	}
+
+	if b.sticky != nil {
+		if err := b.sticky.WriteStickyCookie(rw, server.name); err != nil {
+			log.Error().Err(err).Msg("Error while writing sticky cookie")
+		}
+	}
+
+	server.ServeHTTP(rw, req)
+}
+
+// AddServer adds a handler with a server.
+func (b *Balancer) AddServer(name string, handler http.Handler, server dynamic.Server) {
+	h := &namedHandler{Handler: handler, name: name}
+
+	b.handlersMu.Lock()
+	b.handlers = append(b.handlers, h)
+	b.status[name] = struct{}{}
+	if server.Fenced {
+		b.fenced[name] = struct{}{}
+	}
+	b.handlersMu.Unlock()
+
+	if b.sticky != nil {
+		b.sticky.AddHandler(name, h)
+	}
+}