homelab/traefik
1
0
Fork 0
Code Activity

Rework servers load-balancer to use the WRR

Browse source 
Co-authored-by: Kevin Pollet <pollet.kevin@gmail.com>
This commit is contained in:
Julien Salleyron 2022-11-16 11:38:07 +01:00 • committed by GitHub
parent 67d9c8da0b
commit fadee5e87b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
70 changed files with 2085 additions and 2211 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -4,7 +4,6 @@ import (
"container/heap"
"context"
"errors"
"fmt"
"net/http"
"sync"
@ -39,7 +38,7 @@ type Balancer struct {
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,
// 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
@ -48,10 +47,10 @@ type Balancer struct {
}
// New creates a new load balancer.
func New(sticky *dynamic.Sticky, hc *dynamic.HealthCheck) *Balancer {
func New(sticky *dynamic.Sticky, wantHealthCheck bool) *Balancer {
balancer := &Balancer{
status: make(map[string]struct{}),
wantsHealthCheck: hc != nil,
wantsHealthCheck: wantHealthCheck,
}
if sticky != nil && sticky.Cookie != nil {
balancer.stickyCookie = &stickyCookie{
@ -150,10 +149,7 @@ func (b *Balancer) nextServer() (*namedHandler, error) {
b.mutex.Lock()
defer b.mutex.Unlock()
if len(b.handlers) == 0 {
return nil, fmt.Errorf("no servers in the pool")
}
if len(b.status) == 0 {
if len(b.handlers) == 0 || len(b.status) == 0 {
return nil, errNoAvailableServer
}
@ -223,9 +219,9 @@ func (b *Balancer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
server.ServeHTTP(w, req)
}
// AddService adds a handler.
// Add adds a handler.
// A handler with a non-positive weight is ignored.
func (b *Balancer) AddService(name string, handler http.Handler, weight *int) {
func (b *Balancer) Add(name string, handler http.Handler, weight *int) {
w := 1
if weight != nil {
w = *weight

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

@ -10,31 +10,15 @@ import (
"github.com/traefik/traefik/v2/pkg/config/dynamic"
)
func Int(v int) *int { return &v }
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, nil)
balancer := New(nil, false)
balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), Int(3))
balancer.AddService("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "second")
rw.WriteHeader(http.StatusOK)
}), Int(1))
@ -49,23 +33,23 @@ func TestBalancer(t *testing.T) {
}
func TestBalancerNoService(t *testing.T) {
balancer := New(nil, nil)
balancer := New(nil, false)
recorder := httptest.NewRecorder()
balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
assert.Equal(t, http.StatusInternalServerError, recorder.Result().StatusCode)
assert.Equal(t, http.StatusServiceUnavailable, recorder.Result().StatusCode)
}
func TestBalancerOneServerZeroWeight(t *testing.T) {
balancer := New(nil, nil)
balancer := New(nil, false)
balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
balancer.Add("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) {}), Int(0))
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), Int(0))
recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
for i := 0; i < 3; i++ {
@ -80,13 +64,13 @@ type key string
const serviceName key = "serviceName"
func TestBalancerNoServiceUp(t *testing.T) {
balancer := New(nil, nil)
balancer := New(nil, false)
balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
balancer.Add("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) {
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.WriteHeader(http.StatusInternalServerError)
}), Int(1))
@ -100,14 +84,14 @@ func TestBalancerNoServiceUp(t *testing.T) {
}
func TestBalancerOneServerDown(t *testing.T) {
balancer := New(nil, nil)
balancer := New(nil, false)
balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
balancer.Add("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) {
balancer.Add("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)
@ -121,14 +105,14 @@ func TestBalancerOneServerDown(t *testing.T) {
}
func TestBalancerDownThenUp(t *testing.T) {
balancer := New(nil, nil)
balancer := New(nil, false)
balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
balancer.Add("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) {
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "second")
rw.WriteHeader(http.StatusOK)
}), Int(1))
@ -150,35 +134,35 @@ func TestBalancerDownThenUp(t *testing.T) {
}
func TestBalancerPropagate(t *testing.T) {
balancer1 := New(nil, &dynamic.HealthCheck{})
balancer1 := New(nil, true)
balancer1.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
balancer1.Add("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) {
balancer1.Add("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) {
balancer2 := New(nil, true)
balancer2.Add("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) {
balancer2.Add("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))
topBalancer := New(nil, true)
topBalancer.Add("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))
topBalancer.Add("balancer2", balancer2, Int(1))
_ = balancer2.RegisterStatusUpdater(func(up bool) {
topBalancer.SetStatus(context.WithValue(context.Background(), serviceName, "top"), "balancer2", up)
})
@ -223,28 +207,28 @@ func TestBalancerPropagate(t *testing.T) {
}
func TestBalancerAllServersZeroWeight(t *testing.T) {
balancer := New(nil, nil)
balancer := New(nil, false)
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))
balancer.Add("test", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), Int(0))
balancer.Add("test2", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), Int(0))
recorder := httptest.NewRecorder()
balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
assert.Equal(t, http.StatusInternalServerError, recorder.Result().StatusCode)
assert.Equal(t, http.StatusServiceUnavailable, recorder.Result().StatusCode)
}
func TestSticky(t *testing.T) {
balancer := New(&dynamic.Sticky{
Cookie: &dynamic.Cookie{Name: "test"},
}, nil)
}, false)
balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
balancer.Add("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) {
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "second")
rw.WriteHeader(http.StatusOK)
}), Int(2))
@ -268,14 +252,14 @@ 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, nil)
balancer := New(nil, false)
balancer.AddService("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "A")
rw.WriteHeader(http.StatusOK)
}), Int(11))
balancer.AddService("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "B")
rw.WriteHeader(http.StatusOK)
}), Int(3))
@ -290,3 +274,19 @@ func TestBalancerBias(t *testing.T) {
assert.Equal(t, wantSequence, recorder.sequence)
}
func Int(v int) *int { return &v }
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)
}