homelab/traefik
1
0
Fork 0
Code Activity

Add p2c load-balancing strategy for servers load-balancer

Browse source 
Co-authored-by: Ian Ross <ifross@gmail.com>
Co-authored-by: Kevin Pollet <pollet.kevin@gmail.com>
This commit is contained in:
Romain 2025-03-10 12:12:04 +01:00 • committed by GitHub
parent 550d96ea67
commit 9e029a84c4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
50 changed files with 1621 additions and 382 deletions
Download patch file Download diff file Expand all files Collapse all files

21
pkg/server/router/router_test.go
View file

@ -53,6 +53,7 @@ func TestRouterManager_Get(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: server.URL,
@ -76,6 +77,7 @@ func TestRouterManager_Get(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: server.URL,
@ -114,6 +116,7 @@ func TestRouterManager_Get(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: server.URL,
@ -138,6 +141,7 @@ func TestRouterManager_Get(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: server.URL,
@ -179,6 +183,7 @@ func TestRouterManager_Get(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: server.URL,
@ -219,6 +224,7 @@ func TestRouterManager_Get(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service@provider-1": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: server.URL,
@ -242,6 +248,7 @@ func TestRouterManager_Get(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service@provider-2": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: server.URL,
@ -266,6 +273,7 @@ func TestRouterManager_Get(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service@provider-1": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: server.URL,
@ -351,6 +359,7 @@ func TestRuntimeConfiguration(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: "http://127.0.0.1:8085",
@ -385,6 +394,7 @@ func TestRuntimeConfiguration(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: "http://127.0.0.1",
@ -412,6 +422,7 @@ func TestRuntimeConfiguration(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: "http://127.0.0.1",
@ -439,6 +450,7 @@ func TestRuntimeConfiguration(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: "http://127.0.0.1",
@ -482,6 +494,7 @@ func TestRuntimeConfiguration(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: "http://127.0.0.1",
@ -522,6 +535,7 @@ func TestRuntimeConfiguration(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: "http://127.0.0.1",
@ -552,6 +566,7 @@ func TestRuntimeConfiguration(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: "http://127.0.0.1",
@ -582,6 +597,7 @@ func TestRuntimeConfiguration(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: "http://127.0.0.1",
@ -608,6 +624,7 @@ func TestRuntimeConfiguration(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: "http://127.0.0.1",
@ -641,6 +658,7 @@ func TestRuntimeConfiguration(t *testing.T) {
serviceConfig: map[string]*dynamic.Service{
"foo-service": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: "http://127.0.0.1",
@ -730,7 +748,8 @@ func TestProviderOnMiddlewares(t *testing.T) {
Services: map[string]*dynamic.Service{
"test@file": {
LoadBalancer: &dynamic.ServersLoadBalancer{
Servers: []dynamic.Server{},
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{},
},
},
},

227
pkg/server/service/loadbalancer/p2c/p2c.go Normal file
View file

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

288
pkg/server/service/loadbalancer/p2c/p2c_test.go Normal file
View file

@ -0,0 +1,288 @@
package p2c
import (
"context"
"net/http"
"net/http/httptest"
"strconv"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/traefik/traefik/v3/pkg/config/dynamic"
)
func TestP2C(t *testing.T) {
testCases := []struct {
desc string
handlers []*namedHandler
rand *mockRand
expectedHandler string
}{
{
desc: "one healthy handler",
handlers: testHandlers(0),
rand: nil,
expectedHandler: "0",
},
{
desc: "two handlers zero in flight",
handlers: testHandlers(0, 0),
rand: &mockRand{vals: []int{1, 0}},
expectedHandler: "1",
},
{
desc: "chooses lower of two",
handlers: testHandlers(0, 1),
rand: &mockRand{vals: []int{1, 0}},
expectedHandler: "0",
},
{
desc: "chooses lower of three",
handlers: testHandlers(10, 90, 40),
rand: &mockRand{vals: []int{1, 1}},
expectedHandler: "2",
},
}
for _, test := range testCases {
t.Run(test.desc, func(t *testing.T) {
t.Parallel()
balancer := New(nil, false)
balancer.rand = test.rand
for _, h := range test.handlers {
balancer.handlers = append(balancer.handlers, h)
balancer.status[h.name] = struct{}{}
}
got, err := balancer.nextServer()
require.NoError(t, err)
assert.Equal(t, test.expectedHandler, got.name)
})
}
}
func TestSticky(t *testing.T) {
balancer := New(&dynamic.Sticky{
Cookie: &dynamic.Cookie{
Name: "test",
Secure: true,
HTTPOnly: true,
SameSite: "none",
MaxAge: 42,
Path: func(v string) *string { return &v }("/foo"),
},
}, false)
balancer.rand = &mockRand{vals: []int{1, 0}}
balancer.AddServer("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), dynamic.Server{})
balancer.AddServer("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "second")
rw.WriteHeader(http.StatusOK)
}), dynamic.Server{})
recorder := &responseRecorder{
ResponseRecorder: httptest.NewRecorder(),
save: map[string]int{},
cookies: make(map[string]*http.Cookie),
}
req := httptest.NewRequest(http.MethodGet, "/", nil)
for range 3 {
for _, cookie := range recorder.Result().Cookies() {
assert.NotContains(t, "first", cookie.Value)
assert.NotContains(t, "second", cookie.Value)
req.AddCookie(cookie)
}
recorder.ResponseRecorder = httptest.NewRecorder()
balancer.ServeHTTP(recorder, req)
}
assert.Equal(t, 0, recorder.save["first"])
assert.Equal(t, 3, recorder.save["second"])
assert.True(t, recorder.cookies["test"].HttpOnly)
assert.True(t, recorder.cookies["test"].Secure)
assert.Equal(t, http.SameSiteNoneMode, recorder.cookies["test"].SameSite)
assert.Equal(t, 42, recorder.cookies["test"].MaxAge)
assert.Equal(t, "/foo", recorder.cookies["test"].Path)
}
func TestSticky_Fallback(t *testing.T) {
balancer := New(&dynamic.Sticky{
Cookie: &dynamic.Cookie{Name: "test"},
}, false)
balancer.rand = &mockRand{vals: []int{1, 0}}
balancer.AddServer("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), dynamic.Server{})
balancer.AddServer("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "second")
rw.WriteHeader(http.StatusOK)
}), dynamic.Server{})
recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}, cookies: make(map[string]*http.Cookie)}
req := httptest.NewRequest(http.MethodGet, "/", nil)
req.AddCookie(&http.Cookie{Name: "test", Value: "second"})
for range 3 {
recorder.ResponseRecorder = httptest.NewRecorder()
balancer.ServeHTTP(recorder, req)
}
assert.Equal(t, 0, recorder.save["first"])
assert.Equal(t, 3, recorder.save["second"])
}
// TestSticky_Fenced checks that fenced node receive traffic if their sticky cookie matches.
func TestSticky_Fenced(t *testing.T) {
balancer := New(&dynamic.Sticky{Cookie: &dynamic.Cookie{Name: "test"}}, false)
balancer.rand = &mockRand{vals: []int{1, 0, 1, 0}}
balancer.AddServer("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), dynamic.Server{})
balancer.AddServer("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "second")
rw.WriteHeader(http.StatusOK)
}), dynamic.Server{})
balancer.AddServer("fenced", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "fenced")
rw.WriteHeader(http.StatusOK)
}), dynamic.Server{Fenced: true})
recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}, cookies: make(map[string]*http.Cookie)}
stickyReq := httptest.NewRequest(http.MethodGet, "/", nil)
stickyReq.AddCookie(&http.Cookie{Name: "test", Value: "fenced"})
req := httptest.NewRequest(http.MethodGet, "/", nil)
for range 2 {
recorder.ResponseRecorder = httptest.NewRecorder()
balancer.ServeHTTP(recorder, stickyReq)
balancer.ServeHTTP(recorder, req)
}
assert.Equal(t, 2, recorder.save["fenced"])
assert.Equal(t, 0, recorder.save["first"])
assert.Equal(t, 2, recorder.save["second"])
}
func TestBalancerPropagate(t *testing.T) {
balancer := New(nil, true)
balancer.AddServer("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), dynamic.Server{})
balancer.AddServer("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "second")
rw.WriteHeader(http.StatusOK)
}), dynamic.Server{})
var calls int
err := balancer.RegisterStatusUpdater(func(up bool) {
calls++
})
require.NoError(t, err)
recorder := httptest.NewRecorder()
balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
assert.Equal(t, http.StatusOK, recorder.Code)
// two gets downed, but balancer still up since first is still up.
balancer.SetStatus(context.Background(), "second", false)
assert.Equal(t, 0, calls)
recorder = httptest.NewRecorder()
balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
assert.Equal(t, http.StatusOK, recorder.Code)
assert.Equal(t, "first", recorder.Header().Get("server"))
// first gets downed, balancer is down.
balancer.SetStatus(context.Background(), "first", false)
assert.Equal(t, 1, calls)
recorder = httptest.NewRecorder()
balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
assert.Equal(t, http.StatusServiceUnavailable, recorder.Code)
// two gets up, balancer up.
balancer.SetStatus(context.Background(), "second", true)
assert.Equal(t, 2, calls)
recorder = httptest.NewRecorder()
balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
assert.Equal(t, http.StatusOK, recorder.Code)
assert.Equal(t, "second", recorder.Header().Get("server"))
}
func TestBalancerAllServersFenced(t *testing.T) {
balancer := New(nil, false)
balancer.AddServer("test", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), dynamic.Server{Fenced: true})
balancer.AddServer("test2", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), dynamic.Server{Fenced: true})
recorder := httptest.NewRecorder()
balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
assert.Equal(t, http.StatusServiceUnavailable, recorder.Result().StatusCode)
}
type responseRecorder struct {
*httptest.ResponseRecorder
save map[string]int
sequence []string
status []int
cookies map[string]*http.Cookie
}
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)
for _, cookie := range r.Result().Cookies() {
r.cookies[cookie.Name] = cookie
}
r.ResponseRecorder.WriteHeader(statusCode)
}
type mockRand struct {
vals []int
calls int
}
func (m *mockRand) Intn(int) int {
defer func() {
m.calls++
}()
return m.vals[m.calls]
}
func testHandlers(inflights ...int) []*namedHandler {
var out []*namedHandler
for i, inflight := range inflights {
h := &namedHandler{
name: strconv.Itoa(i),
}
h.inflight.Store(int64(inflight))
out = append(out, h)
}
return out
}

179
pkg/server/service/loadbalancer/sticky.go Normal file
View file

@ -0,0 +1,179 @@
package loadbalancer
import (
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"hash/fnv"
"net/http"
"strconv"
"sync"
"github.com/traefik/traefik/v3/pkg/config/dynamic"
)
// NamedHandler is a http.Handler with a name.
type NamedHandler struct {
http.Handler
Name string
}
// stickyCookie represents a sticky cookie.
type stickyCookie struct {
name string
secure bool
httpOnly bool
sameSite http.SameSite
maxAge int
path string
domain string
}
// Sticky ensures that client consistently interacts with the same HTTP handler by adding a sticky cookie to the response.
// This cookie allows subsequent requests from the same client to be routed to the same handler,
// enabling session persistence across multiple requests.
type Sticky struct {
// cookie is the sticky cookie configuration.
cookie *stickyCookie
// References all the handlers by name and also by the hashed value of the name.
handlersMu sync.RWMutex
hashMap map[string]string
stickyMap map[string]*NamedHandler
compatibilityStickyMap map[string]*NamedHandler
}
// NewSticky creates a new Sticky instance.
func NewSticky(cookieConfig dynamic.Cookie) *Sticky {
cookie := &stickyCookie{
name: cookieConfig.Name,
secure: cookieConfig.Secure,
httpOnly: cookieConfig.HTTPOnly,
sameSite: convertSameSite(cookieConfig.SameSite),
maxAge: cookieConfig.MaxAge,
path: "/",
domain: cookieConfig.Domain,
}
if cookieConfig.Path != nil {
cookie.path = *cookieConfig.Path
}
return &Sticky{
cookie: cookie,
hashMap: make(map[string]string),
stickyMap: make(map[string]*NamedHandler),
compatibilityStickyMap: make(map[string]*NamedHandler),
}
}
// AddHandler adds a http.Handler to the sticky pool.
func (s *Sticky) AddHandler(name string, h http.Handler) {
s.handlersMu.Lock()
defer s.handlersMu.Unlock()
sha256HashedName := sha256Hash(name)
s.hashMap[name] = sha256HashedName
handler := &NamedHandler{
Handler: h,
Name: name,
}
s.stickyMap[sha256HashedName] = handler
s.compatibilityStickyMap[name] = handler
hashedName := fnvHash(name)
s.compatibilityStickyMap[hashedName] = handler
// server.URL was fnv hashed in service.Manager
// so we can have "double" fnv hash in already existing cookies
hashedName = fnvHash(hashedName)
s.compatibilityStickyMap[hashedName] = handler
}
// StickyHandler returns the NamedHandler corresponding to the sticky cookie if one.
// It also returns a boolean which indicates if the sticky cookie has to be overwritten because it uses a deprecated hash algorithm.
func (s *Sticky) StickyHandler(req *http.Request) (*NamedHandler, bool, error) {
cookie, err := req.Cookie(s.cookie.name)
if err != nil && errors.Is(err, http.ErrNoCookie) {
return nil, false, nil
}
if err != nil {
return nil, false, fmt.Errorf("reading cookie: %w", err)
}
s.handlersMu.RLock()
handler, ok := s.stickyMap[cookie.Value]
s.handlersMu.RUnlock()
if ok && handler != nil {
return handler, false, nil
}
s.handlersMu.RLock()
handler, ok = s.compatibilityStickyMap[cookie.Value]
s.handlersMu.RUnlock()
return handler, ok, nil
}
// WriteStickyCookie writes a sticky cookie to the response to stick the client to the given handler name.
func (s *Sticky) WriteStickyCookie(rw http.ResponseWriter, name string) error {
s.handlersMu.RLock()
hash, ok := s.hashMap[name]
s.handlersMu.RUnlock()
if !ok {
return fmt.Errorf("no hash found for handler named %s", name)
}
cookie := &http.Cookie{
Name: s.cookie.name,
Value: hash,
Path: s.cookie.path,
Domain: s.cookie.domain,
HttpOnly: s.cookie.httpOnly,
Secure: s.cookie.secure,
SameSite: s.cookie.sameSite,
MaxAge: s.cookie.maxAge,
}
http.SetCookie(rw, cookie)
return nil
}
func convertSameSite(sameSite string) http.SameSite {
switch sameSite {
case "none":
return http.SameSiteNoneMode
case "lax":
return http.SameSiteLaxMode
case "strict":
return http.SameSiteStrictMode
default:
return http.SameSiteDefaultMode
}
}
// fnvHash returns the FNV-64 hash of the input string.
func fnvHash(input string) string {
hasher := fnv.New64()
// We purposely ignore the error because the implementation always returns nil.
_, _ = hasher.Write([]byte(input))
return strconv.FormatUint(hasher.Sum64(), 16)
}
// sha256 returns the SHA-256 hash, truncated to 16 characters, of the input string.
func sha256Hash(input string) string {
hash := sha256.New()
// We purposely ignore the error because the implementation always returns nil.
_, _ = hash.Write([]byte(input))
hashedInput := hex.EncodeToString(hash.Sum(nil))
if len(hashedInput) < 16 {
return hashedInput
}
return hashedInput[:16]
}

138
pkg/server/service/loadbalancer/sticky_test.go Normal file
View file

@ -0,0 +1,138 @@
package loadbalancer
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/traefik/traefik/v3/pkg/config/dynamic"
)
func pointer[T any](v T) *T { return &v }
func TestSticky_StickyHandler(t *testing.T) {
testCases := []struct {
desc string
handlers []string
cookies []*http.Cookie
wantHandler string
wantRewrite bool
}{
{
desc: "No previous cookie",
handlers: []string{"first"},
wantHandler: "",
wantRewrite: false,
},
{
desc: "Wrong previous cookie",
handlers: []string{"first"},
cookies: []*http.Cookie{
{Name: "test", Value: sha256Hash("foo")},
},
wantHandler: "",
wantRewrite: false,
},
{
desc: "Sha256 previous cookie",
handlers: []string{"first", "second"},
cookies: []*http.Cookie{
{Name: "test", Value: sha256Hash("first")},
},
wantHandler: "first",
wantRewrite: false,
},
{
desc: "Raw previous cookie",
handlers: []string{"first", "second"},
cookies: []*http.Cookie{
{Name: "test", Value: "first"},
},
wantHandler: "first",
wantRewrite: true,
},
{
desc: "Fnv previous cookie",
handlers: []string{"first", "second"},
cookies: []*http.Cookie{
{Name: "test", Value: fnvHash("first")},
},
wantHandler: "first",
wantRewrite: true,
},
{
desc: "Double fnv previous cookie",
handlers: []string{"first", "second"},
cookies: []*http.Cookie{
{Name: "test", Value: fnvHash("first")},
},
wantHandler: "first",
wantRewrite: true,
},
}
for _, test := range testCases {
t.Run(test.desc, func(t *testing.T) {
t.Parallel()
sticky := NewSticky(dynamic.Cookie{Name: "test"})
for _, handler := range test.handlers {
sticky.AddHandler(handler, http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}))
}
req := httptest.NewRequest(http.MethodGet, "/", nil)
for _, cookie := range test.cookies {
req.AddCookie(cookie)
}
h, rewrite, err := sticky.StickyHandler(req)
require.NoError(t, err)
if test.wantHandler != "" {
assert.NotNil(t, h)
assert.Equal(t, test.wantHandler, h.Name)
} else {
assert.Nil(t, h)
}
assert.Equal(t, test.wantRewrite, rewrite)
})
}
}
func TestSticky_WriteStickyCookie(t *testing.T) {
sticky := NewSticky(dynamic.Cookie{
Name: "test",
Secure: true,
HTTPOnly: true,
SameSite: "none",
MaxAge: 42,
Path: pointer("/foo"),
Domain: "foo.com",
})
// Should return an error if the handler does not exist.
res := httptest.NewRecorder()
require.Error(t, sticky.WriteStickyCookie(res, "first"))
// Should write the sticky cookie and use the sha256 hash.
sticky.AddHandler("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}))
res = httptest.NewRecorder()
require.NoError(t, sticky.WriteStickyCookie(res, "first"))
assert.Len(t, res.Result().Cookies(), 1)
cookie := res.Result().Cookies()[0]
assert.Equal(t, sha256Hash("first"), cookie.Value)
assert.Equal(t, "test", cookie.Name)
assert.True(t, cookie.Secure)
assert.True(t, cookie.HttpOnly)
assert.Equal(t, http.SameSiteNoneMode, cookie.SameSite)
assert.Equal(t, 42, cookie.MaxAge)
assert.Equal(t, "/foo", cookie.Path)
assert.Equal(t, "foo.com", cookie.Domain)
}

175
pkg/server/service/loadbalancer/wrr/wrr.go
View file

@ -3,47 +3,20 @@ package wrr
import (
"container/heap"
"context"
"crypto/sha256"
"encoding/hex"
"errors"
"hash/fnv"
"net/http"
"strconv"
"sync"
"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 string
hashedName string
weight float64
deadline float64
}
type stickyCookie struct {
name string
secure bool
httpOnly bool
sameSite string
maxAge int
path string
domain string
}
func convertSameSite(sameSite string) http.SameSite {
switch sameSite {
case "none":
return http.SameSiteNoneMode
case "lax":
return http.SameSiteLaxMode
case "strict":
return http.SameSiteStrictMode
default:
return http.SameSiteDefaultMode
}
weight float64
deadline float64
}
// Balancer is a WeightedRoundRobin load balancer based on Earliest Deadline First (EDF).
@ -52,15 +25,10 @@ func convertSameSite(sameSite string) http.SameSite {
// 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
wantsHealthCheck bool
handlersMu sync.RWMutex
// References all the handlers by name and also by the hashed value of the name.
stickyMap map[string]*namedHandler
compatibilityStickyMap map[string]*namedHandler
handlers []*namedHandler
curDeadline float64
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,
@ -71,31 +39,21 @@ type Balancer struct {
updaters []func(bool)
// fenced is the list of terminating yet still serving child services.
fenced map[string]struct{}
sticky *loadbalancer.Sticky
curDeadline float64
}
// New creates a new load balancer.
func New(sticky *dynamic.Sticky, wantHealthCheck bool) *Balancer {
func New(sticky *dynamic.Sticky, wantsHealthCheck bool) *Balancer {
balancer := &Balancer{
status: make(map[string]struct{}),
fenced: make(map[string]struct{}),
wantsHealthCheck: wantHealthCheck,
wantsHealthCheck: wantsHealthCheck,
}
if sticky != nil && sticky.Cookie != nil {
balancer.stickyCookie = &stickyCookie{
name: sticky.Cookie.Name,
secure: sticky.Cookie.Secure,
httpOnly: sticky.Cookie.HTTPOnly,
sameSite: sticky.Cookie.SameSite,
maxAge: sticky.Cookie.MaxAge,
path: "/",
domain: sticky.Cookie.Domain,
}
if sticky.Cookie.Path != nil {
balancer.stickyCookie.path = *sticky.Cookie.Path
}
balancer.stickyMap = make(map[string]*namedHandler)
balancer.compatibilityStickyMap = make(map[string]*namedHandler)
balancer.sticky = loadbalancer.NewSticky(*sticky.Cookie)
}
return balancer
@ -216,43 +174,21 @@ func (b *Balancer) nextServer() (*namedHandler, error) {
return handler, nil
}
func (b *Balancer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
if b.stickyCookie != nil {
cookie, err := req.Cookie(b.stickyCookie.name)
if err != nil && !errors.Is(err, http.ErrNoCookie) {
log.Warn().Err(err).Msg("Error while reading cookie")
}
if err == nil && cookie != nil {
b.handlersMu.RLock()
handler, ok := b.stickyMap[cookie.Value]
b.handlersMu.RUnlock()
if ok && handler != nil {
b.handlersMu.RLock()
_, isHealthy := b.status[handler.name]
b.handlersMu.RUnlock()
if isHealthy {
handler.ServeHTTP(w, req)
return
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")
}
}
}
b.handlersMu.RLock()
handler, ok = b.compatibilityStickyMap[cookie.Value]
b.handlersMu.RUnlock()
if ok && handler != nil {
b.handlersMu.RLock()
_, isHealthy := b.status[handler.name]
b.handlersMu.RUnlock()
if isHealthy {
b.writeStickyCookie(w, handler)
handler.ServeHTTP(w, req)
return
}
h.ServeHTTP(rw, req)
return
}
}
}
@ -260,32 +196,25 @@ func (b *Balancer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
server, err := b.nextServer()
if err != nil {
if errors.Is(err, errNoAvailableServer) {
http.Error(w, errNoAvailableServer.Error(), http.StatusServiceUnavailable)
http.Error(rw, errNoAvailableServer.Error(), http.StatusServiceUnavailable)
} else {
http.Error(w, err.Error(), http.StatusInternalServerError)
http.Error(rw, err.Error(), http.StatusInternalServerError)
}
return
}
if b.stickyCookie != nil {
b.writeStickyCookie(w, server)
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(w, req)
server.ServeHTTP(rw, req)
}
func (b *Balancer) writeStickyCookie(w http.ResponseWriter, handler *namedHandler) {
cookie := &http.Cookie{
Name: b.stickyCookie.name,
Value: handler.hashedName,
Path: b.stickyCookie.path,
HttpOnly: b.stickyCookie.httpOnly,
Secure: b.stickyCookie.secure,
SameSite: convertSameSite(b.stickyCookie.sameSite),
MaxAge: b.stickyCookie.maxAge,
Domain: b.stickyCookie.domain,
}
http.SetCookie(w, cookie)
// AddServer adds a handler with a server.
func (b *Balancer) AddServer(name string, handler http.Handler, server dynamic.Server) {
b.Add(name, handler, server.Weight, server.Fenced)
}
// Add adds a handler.
@ -309,41 +238,9 @@ func (b *Balancer) Add(name string, handler http.Handler, weight *int, fenced bo
if fenced {
b.fenced[name] = struct{}{}
}
if b.stickyCookie != nil {
sha256HashedName := sha256Hash(name)
h.hashedName = sha256HashedName
b.stickyMap[sha256HashedName] = h
b.compatibilityStickyMap[name] = h
hashedName := fnvHash(name)
b.compatibilityStickyMap[hashedName] = h
// server.URL was fnv hashed in service.Manager
// so we can have "double" fnv hash in already existing cookies
hashedName = fnvHash(hashedName)
b.compatibilityStickyMap[hashedName] = h
}
b.handlersMu.Unlock()
}
func fnvHash(input string) string {
hasher := fnv.New64()
// We purposely ignore the error because the implementation always returns nil.
_, _ = hasher.Write([]byte(input))
return strconv.FormatUint(hasher.Sum64(), 16)
}
func sha256Hash(input string) string {
hash := sha256.New()
// We purposely ignore the error because the implementation always returns nil.
_, _ = hash.Write([]byte(input))
hashedInput := hex.EncodeToString(hash.Sum(nil))
if len(hashedInput) < 16 {
return hashedInput
if b.sticky != nil {
b.sticky.AddHandler(name, handler)
}
return hashedInput[:16]
}

186
pkg/server/service/loadbalancer/wrr/wrr_test.go
View file

@ -10,6 +10,10 @@ import (
"github.com/traefik/traefik/v3/pkg/config/dynamic"
)
type key string
const serviceName key = "serviceName"
func pointer[T any](v T) *T { return &v }
func TestBalancer(t *testing.T) {
@ -61,10 +65,6 @@ 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, false)
@ -264,8 +264,8 @@ func TestSticky(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "/", nil)
for range 3 {
for _, cookie := range recorder.Result().Cookies() {
assert.NotContains(t, "test=first", cookie.Value)
assert.NotContains(t, "test=second", cookie.Value)
assert.NotContains(t, "first", cookie.Value)
assert.NotContains(t, "second", cookie.Value)
req.AddCookie(cookie)
}
recorder.ResponseRecorder = httptest.NewRecorder()
@ -283,7 +283,7 @@ func TestSticky(t *testing.T) {
assert.Equal(t, "/foo", recorder.cookies["test"].Path)
}
func TestSticky_FallBack(t *testing.T) {
func TestSticky_Fallback(t *testing.T) {
balancer := New(&dynamic.Sticky{
Cookie: &dynamic.Cookie{Name: "test"},
}, false)
@ -312,6 +312,44 @@ func TestSticky_FallBack(t *testing.T) {
assert.Equal(t, 3, recorder.save["second"])
}
// TestSticky_Fenced checks that fenced node receive traffic if their sticky cookie matches.
func TestSticky_Fenced(t *testing.T) {
balancer := New(&dynamic.Sticky{Cookie: &dynamic.Cookie{Name: "test"}}, false)
balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), pointer(1), false)
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "second")
rw.WriteHeader(http.StatusOK)
}), pointer(1), false)
balancer.Add("fenced", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "fenced")
rw.WriteHeader(http.StatusOK)
}), pointer(1), true)
recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}, cookies: make(map[string]*http.Cookie)}
stickyReq := httptest.NewRequest(http.MethodGet, "/", nil)
stickyReq.AddCookie(&http.Cookie{Name: "test", Value: "fenced"})
req := httptest.NewRequest(http.MethodGet, "/", nil)
for range 4 {
recorder.ResponseRecorder = httptest.NewRecorder()
balancer.ServeHTTP(recorder, stickyReq)
balancer.ServeHTTP(recorder, req)
}
assert.Equal(t, 4, recorder.save["fenced"])
assert.Equal(t, 2, recorder.save["first"])
assert.Equal(t, 2, recorder.save["second"])
}
// 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) {
@ -355,137 +393,3 @@ func (r *responseRecorder) WriteHeader(statusCode int) {
}
r.ResponseRecorder.WriteHeader(statusCode)
}
// TestSticky_Fenced checks that fenced node receive traffic if their sticky cookie matches.
func TestSticky_Fenced(t *testing.T) {
balancer := New(&dynamic.Sticky{Cookie: &dynamic.Cookie{Name: "test"}}, false)
balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), pointer(1), false)
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "second")
rw.WriteHeader(http.StatusOK)
}), pointer(1), false)
balancer.Add("fenced", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "fenced")
rw.WriteHeader(http.StatusOK)
}), pointer(1), true)
recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}, cookies: make(map[string]*http.Cookie)}
stickyReq := httptest.NewRequest(http.MethodGet, "/", nil)
stickyReq.AddCookie(&http.Cookie{Name: "test", Value: "fenced"})
req := httptest.NewRequest(http.MethodGet, "/", nil)
for range 4 {
recorder.ResponseRecorder = httptest.NewRecorder()
balancer.ServeHTTP(recorder, stickyReq)
balancer.ServeHTTP(recorder, req)
}
assert.Equal(t, 4, recorder.save["fenced"])
assert.Equal(t, 2, recorder.save["first"])
assert.Equal(t, 2, recorder.save["second"])
}
func TestStickyWithCompatibility(t *testing.T) {
testCases := []struct {
desc string
servers []string
cookies []*http.Cookie
expectedCookies []*http.Cookie
expectedServer string
}{
{
desc: "No previous cookie",
servers: []string{"first"},
expectedServer: "first",
expectedCookies: []*http.Cookie{
{Name: "test", Value: sha256Hash("first")},
},
},
{
desc: "Sha256 previous cookie",
servers: []string{"first", "second"},
cookies: []*http.Cookie{
{Name: "test", Value: sha256Hash("first")},
},
expectedServer: "first",
expectedCookies: []*http.Cookie{},
},
{
desc: "Raw previous cookie",
servers: []string{"first", "second"},
cookies: []*http.Cookie{
{Name: "test", Value: "first"},
},
expectedServer: "first",
expectedCookies: []*http.Cookie{
{Name: "test", Value: sha256Hash("first")},
},
},
{
desc: "Fnv previous cookie",
servers: []string{"first", "second"},
cookies: []*http.Cookie{
{Name: "test", Value: fnvHash("first")},
},
expectedServer: "first",
expectedCookies: []*http.Cookie{
{Name: "test", Value: sha256Hash("first")},
},
},
{
desc: "Double fnv previous cookie",
servers: []string{"first", "second"},
cookies: []*http.Cookie{
{Name: "test", Value: fnvHash(fnvHash("first"))},
},
expectedServer: "first",
expectedCookies: []*http.Cookie{
{Name: "test", Value: sha256Hash("first")},
},
},
}
for _, test := range testCases {
t.Run(test.desc, func(t *testing.T) {
t.Parallel()
balancer := New(&dynamic.Sticky{Cookie: &dynamic.Cookie{Name: "test"}}, false)
for _, server := range test.servers {
balancer.Add(server, http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.WriteHeader(http.StatusOK)
_, _ = rw.Write([]byte(server))
}), pointer(1), false)
}
// Do it twice, to be sure it's not just the luck.
for range 2 {
req := httptest.NewRequest(http.MethodGet, "/", nil)
for _, cookie := range test.cookies {
req.AddCookie(cookie)
}
recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}, cookies: make(map[string]*http.Cookie)}
balancer.ServeHTTP(recorder, req)
assert.Equal(t, test.expectedServer, recorder.Body.String())
assert.Len(t, recorder.cookies, len(test.expectedCookies))
for _, cookie := range test.expectedCookies {
assert.Equal(t, cookie.Value, recorder.cookies[cookie.Name].Value)
}
}
})
}
}

23
pkg/server/service/service.go
View file

@ -29,6 +29,7 @@ import (
"github.com/traefik/traefik/v3/pkg/server/provider"
"github.com/traefik/traefik/v3/pkg/server/service/loadbalancer/failover"
"github.com/traefik/traefik/v3/pkg/server/service/loadbalancer/mirror"
"github.com/traefik/traefik/v3/pkg/server/service/loadbalancer/p2c"
"github.com/traefik/traefik/v3/pkg/server/service/loadbalancer/wrr"
"google.golang.org/grpc/status"
)
@ -304,6 +305,13 @@ func (m *Manager) getServiceHandler(ctx context.Context, service dynamic.WRRServ
}
}
type serverBalancer interface {
http.Handler
healthcheck.StatusSetter
AddServer(name string, handler http.Handler, server dynamic.Server)
}
func (m *Manager) getLoadBalancerServiceHandler(ctx context.Context, serviceName string, info *runtime.ServiceInfo) (http.Handler, error) {
service := info.LoadBalancer
@ -330,7 +338,18 @@ func (m *Manager) getLoadBalancerServiceHandler(ctx context.Context, serviceName
passHostHeader = *service.PassHostHeader
}
lb := wrr.New(service.Sticky, service.HealthCheck != nil)
var lb serverBalancer
switch service.Strategy {
// Here we are handling the empty value to comply with providers that are not applying defaults (e.g. REST provider)
// TODO: remove this when all providers apply default values.
case dynamic.BalancerStrategyWRR, "":
lb = wrr.New(service.Sticky, service.HealthCheck != nil)
case dynamic.BalancerStrategyP2C:
lb = p2c.New(service.Sticky, service.HealthCheck != nil)
default:
return nil, fmt.Errorf("unsupported load-balancer strategy %q", service.Strategy)
}
healthCheckTargets := make(map[string]*url.URL)
for i, server := range shuffle(service.Servers, m.rand) {
@ -385,7 +404,7 @@ func (m *Manager) getLoadBalancerServiceHandler(ctx context.Context, serviceName
proxy, _ = capture.Wrap(proxy)
}
lb.Add(server.URL, proxy, server.Weight, server.Fenced)
lb.AddServer(server.URL, proxy, server)
// servers are considered UP by default.
info.UpdateServerStatus(target.String(), runtime.StatusUp)

40
pkg/server/service/service_test.go
View file

@ -38,6 +38,7 @@ func TestGetLoadBalancer(t *testing.T) {
desc: "Fails when provided an invalid URL",
serviceName: "test",
service: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: ":",
@ -50,7 +51,9 @@ func TestGetLoadBalancer(t *testing.T) {
{
desc: "Succeeds when there are no servers",
serviceName: "test",
service: &dynamic.ServersLoadBalancer{},
service: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
},
fwd: &forwarderMock{},
expectError: false,
},
@ -58,7 +61,8 @@ func TestGetLoadBalancer(t *testing.T) {
desc: "Succeeds when sticky.cookie is set",
serviceName: "test",
service: &dynamic.ServersLoadBalancer{
Sticky: &dynamic.Sticky{Cookie: &dynamic.Cookie{}},
Strategy: dynamic.BalancerStrategyWRR,
Sticky: &dynamic.Sticky{Cookie: &dynamic.Cookie{}},
},
fwd: &forwarderMock{},
expectError: false,
@ -140,6 +144,7 @@ func TestGetLoadBalancerServiceHandler(t *testing.T) {
desc: "Load balances between the two servers",
serviceName: "test",
service: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
PassHostHeader: boolPtr(true),
Servers: []dynamic.Server{
{
@ -165,6 +170,7 @@ func TestGetLoadBalancerServiceHandler(t *testing.T) {
desc: "StatusBadGateway when the server is not reachable",
serviceName: "test",
service: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: "http://foo",
@ -181,7 +187,8 @@ func TestGetLoadBalancerServiceHandler(t *testing.T) {
desc: "ServiceUnavailable when no servers are available",
serviceName: "test",
service: &dynamic.ServersLoadBalancer{
Servers: []dynamic.Server{},
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{},
},
expected: []ExpectedResult{
{
@ -193,7 +200,8 @@ func TestGetLoadBalancerServiceHandler(t *testing.T) {
desc: "Always call the same server when sticky.cookie is true",
serviceName: "test",
service: &dynamic.ServersLoadBalancer{
Sticky: &dynamic.Sticky{Cookie: &dynamic.Cookie{}},
Strategy: dynamic.BalancerStrategyWRR,
Sticky: &dynamic.Sticky{Cookie: &dynamic.Cookie{}},
Servers: []dynamic.Server{
{
URL: server1.URL,
@ -216,7 +224,8 @@ func TestGetLoadBalancerServiceHandler(t *testing.T) {
desc: "Sticky Cookie's options set correctly",
serviceName: "test",
service: &dynamic.ServersLoadBalancer{
Sticky: &dynamic.Sticky{Cookie: &dynamic.Cookie{HTTPOnly: true, Secure: true}},
Strategy: dynamic.BalancerStrategyWRR,
Sticky: &dynamic.Sticky{Cookie: &dynamic.Cookie{HTTPOnly: true, Secure: true}},
Servers: []dynamic.Server{
{
URL: server1.URL,
@ -236,6 +245,7 @@ func TestGetLoadBalancerServiceHandler(t *testing.T) {
desc: "PassHost passes the host instead of the IP",
serviceName: "test",
service: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Sticky: &dynamic.Sticky{Cookie: &dynamic.Cookie{}},
PassHostHeader: pointer(true),
Servers: []dynamic.Server{
@ -255,6 +265,7 @@ func TestGetLoadBalancerServiceHandler(t *testing.T) {
desc: "PassHost doesn't pass the host instead of the IP",
serviceName: "test",
service: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
PassHostHeader: pointer(false),
Sticky: &dynamic.Sticky{Cookie: &dynamic.Cookie{}},
Servers: []dynamic.Server{
@ -274,6 +285,7 @@ func TestGetLoadBalancerServiceHandler(t *testing.T) {
desc: "No user-agent",
serviceName: "test",
service: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: hasNoUserAgent.URL,
@ -291,6 +303,7 @@ func TestGetLoadBalancerServiceHandler(t *testing.T) {
serviceName: "test",
userAgent: "foobar",
service: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: hasUserAgent.URL,
@ -379,6 +392,7 @@ func Test1xxResponses(t *testing.T) {
info := &runtime.ServiceInfo{
Service: &dynamic.Service{
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
Servers: []dynamic.Server{
{
URL: backend.URL,
@ -466,7 +480,9 @@ func TestManager_ServiceBuilders(t *testing.T) {
manager := NewManager(map[string]*runtime.ServiceInfo{
"test@test": {
Service: &dynamic.Service{
LoadBalancer: &dynamic.ServersLoadBalancer{},
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
},
},
},
}, nil, nil, &TransportManager{
@ -505,7 +521,9 @@ func TestManager_Build(t *testing.T) {
configs: map[string]*runtime.ServiceInfo{
"serviceName": {
Service: &dynamic.Service{
LoadBalancer: &dynamic.ServersLoadBalancer{},
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
},
},
},
},
@ -516,7 +534,9 @@ func TestManager_Build(t *testing.T) {
configs: map[string]*runtime.ServiceInfo{
"serviceName@provider-1": {
Service: &dynamic.Service{
LoadBalancer: &dynamic.ServersLoadBalancer{},
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
},
},
},
},
@ -527,7 +547,9 @@ func TestManager_Build(t *testing.T) {
configs: map[string]*runtime.ServiceInfo{
"serviceName@provider-1": {
Service: &dynamic.Service{
LoadBalancer: &dynamic.ServersLoadBalancer{},
LoadBalancer: &dynamic.ServersLoadBalancer{
Strategy: dynamic.BalancerStrategyWRR,
},
},
},
},