Add weighted round robin load balancer on TCP

Co-authored-by: Mathieu Lonjaret <mathieu.lonjaret@gmail.com>

pkg/tcp/proxy.go

@@ -58,13 +58,14 @@ func (p Proxy) connCopy(dst, src WriteCloser, errCh chan error) {
 
 	errClose := dst.CloseWrite()
 	if errClose != nil {
-		log.WithoutContext().Errorf("Error while terminating connection: %v", errClose)
+		log.WithoutContext().Debugf("Error while terminating connection: %v", errClose)
+		return
 	}
 
 	if p.terminationDelay >= 0 {
 		err := dst.SetReadDeadline(time.Now().Add(p.terminationDelay))
 		if err != nil {
-			log.WithoutContext().Errorf("Error while setting deadline: %v", err)
+			log.WithoutContext().Debugf("Error while setting deadline: %v", err)
 		}
 	}
 }

pkg/tcp/rr_load_balancer.go

@@ -1,48 +0,0 @@
-package tcp
-
-import (
-	"sync"
-
-	"github.com/containous/traefik/v2/pkg/log"
-)
-
-// RRLoadBalancer is a naive RoundRobin load balancer for TCP services
-type RRLoadBalancer struct {
-	servers []Handler
-	lock    sync.RWMutex
-	current int
-}
-
-// NewRRLoadBalancer creates a new RRLoadBalancer
-func NewRRLoadBalancer() *RRLoadBalancer {
-	return &RRLoadBalancer{}
-}
-
-// ServeTCP forwards the connection to the right service
-func (r *RRLoadBalancer) ServeTCP(conn WriteCloser) {
-	if len(r.servers) == 0 {
-		log.WithoutContext().Error("no available server")
-		return
-	}
-
-	r.next().ServeTCP(conn)
-}
-
-// AddServer appends a server to the existing list
-func (r *RRLoadBalancer) AddServer(server Handler) {
-	r.servers = append(r.servers, server)
-}
-
-func (r *RRLoadBalancer) next() Handler {
-	r.lock.Lock()
-	defer r.lock.Unlock()
-
-	if r.current >= len(r.servers) {
-		r.current = 0
-		log.WithoutContext().Debugf("Load balancer: going back to the first available server")
-	}
-
-	handler := r.servers[r.current]
-	r.current++
-	return handler
-}

pkg/tcp/wrr_load_balancer.go

@@ -0,0 +1,122 @@
+package tcp
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/containous/traefik/v2/pkg/log"
+)
+
+type server struct {
+	Handler
+	weight int
+}
+
+// WRRLoadBalancer is a naive RoundRobin load balancer for TCP services
+type WRRLoadBalancer struct {
+	servers       []server
+	lock          sync.RWMutex
+	currentWeight int
+	index         int
+}
+
+// NewWRRLoadBalancer creates a new WRRLoadBalancer
+func NewWRRLoadBalancer() *WRRLoadBalancer {
+	return &WRRLoadBalancer{
+		index: -1,
+	}
+}
+
+// ServeTCP forwards the connection to the right service
+func (b *WRRLoadBalancer) ServeTCP(conn WriteCloser) {
+	if len(b.servers) == 0 {
+		log.WithoutContext().Error("no available server")
+		return
+	}
+
+	next, err := b.next()
+	if err != nil {
+		log.WithoutContext().Errorf("Error during load balancing: %v", err)
+		conn.Close()
+	}
+	next.ServeTCP(conn)
+}
+
+// AddServer appends a server to the existing list
+func (b *WRRLoadBalancer) AddServer(serverHandler Handler) {
+	w := 1
+	b.AddWeightServer(serverHandler, &w)
+}
+
+// AddWeightServer appends a server to the existing list with a weight
+func (b *WRRLoadBalancer) AddWeightServer(serverHandler Handler, weight *int) {
+	w := 1
+	if weight != nil {
+		w = *weight
+	}
+	b.servers = append(b.servers, server{Handler: serverHandler, weight: w})
+}
+
+func (b *WRRLoadBalancer) maxWeight() int {
+	max := -1
+	for _, s := range b.servers {
+		if s.weight > max {
+			max = s.weight
+		}
+	}
+	return max
+}
+
+func (b *WRRLoadBalancer) weightGcd() int {
+	divisor := -1
+	for _, s := range b.servers {
+		if divisor == -1 {
+			divisor = s.weight
+		} else {
+			divisor = gcd(divisor, s.weight)
+		}
+	}
+	return divisor
+}
+
+func gcd(a, b int) int {
+	for b != 0 {
+		a, b = b, a%b
+	}
+	return a
+}
+
+func (b *WRRLoadBalancer) next() (Handler, error) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	if len(b.servers) == 0 {
+		return nil, fmt.Errorf("no servers in the pool")
+	}
+
+	// The algo below may look messy, but is actually very simple
+	// it calculates the GCD  and subtracts it on every iteration, what interleaves servers
+	// and allows us not to build an iterator every time we readjust weights
+
+	// GCD across all enabled servers
+	gcd := b.weightGcd()
+	// Maximum weight across all enabled servers
+	max := b.maxWeight()
+
+	for {
+		b.index = (b.index + 1) % len(b.servers)
+		if b.index == 0 {
+			b.currentWeight -= gcd
+			if b.currentWeight <= 0 {
+				b.currentWeight = max
+				if b.currentWeight == 0 {
+					return nil, fmt.Errorf("all servers have 0 weight")
+				}
+			}
+		}
+		srv := b.servers[b.index]
+		if srv.weight >= b.currentWeight {
+			return srv, nil
+		}
+	}
+}

pkg/tcp/wrr_load_balancer_test.go

@@ -0,0 +1,131 @@
+package tcp
+
+import (
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type fakeConn struct {
+	call map[string]int
+}
+
+func (f *fakeConn) Read(b []byte) (n int, err error) {
+	panic("implement me")
+}
+
+func (f *fakeConn) Write(b []byte) (n int, err error) {
+	f.call[string(b)]++
+	return len(b), nil
+}
+
+func (f *fakeConn) Close() error {
+	panic("implement me")
+}
+
+func (f *fakeConn) LocalAddr() net.Addr {
+	panic("implement me")
+}
+
+func (f *fakeConn) RemoteAddr() net.Addr {
+	panic("implement me")
+}
+
+func (f *fakeConn) SetDeadline(t time.Time) error {
+	panic("implement me")
+}
+
+func (f *fakeConn) SetReadDeadline(t time.Time) error {
+	panic("implement me")
+}
+
+func (f *fakeConn) SetWriteDeadline(t time.Time) error {
+	panic("implement me")
+}
+
+func (f *fakeConn) CloseWrite() error {
+	panic("implement me")
+}
+
+func TestLoadBalancing(t *testing.T) {
+	testCases := []struct {
+		desc          string
+		serversWeight map[string]int
+		totalCall     int
+		expected      map[string]int
+	}{
+		{
+			desc: "RoundRobin",
+			serversWeight: map[string]int{
+				"h1": 1,
+				"h2": 1,
+			},
+			totalCall: 4,
+			expected: map[string]int{
+				"h1": 2,
+				"h2": 2,
+			},
+		},
+		{
+			desc: "WeighedRoundRobin",
+			serversWeight: map[string]int{
+				"h1": 3,
+				"h2": 1,
+			},
+			totalCall: 4,
+			expected: map[string]int{
+				"h1": 3,
+				"h2": 1,
+			},
+		},
+		{
+			desc: "WeighedRoundRobin with more call",
+			serversWeight: map[string]int{
+				"h1": 3,
+				"h2": 1,
+			},
+			totalCall: 16,
+			expected: map[string]int{
+				"h1": 12,
+				"h2": 4,
+			},
+		},
+		{
+			desc: "WeighedRoundRobin with 0 weight server",
+			serversWeight: map[string]int{
+				"h1": 3,
+				"h2": 0,
+			},
+			totalCall: 16,
+			expected: map[string]int{
+				"h1": 16,
+			},
+		},
+	}
+
+	for _, test := range testCases {
+		test := test
+		t.Run(test.desc, func(t *testing.T) {
+			t.Parallel()
+
+			balancer := NewWRRLoadBalancer()
+			for server, weight := range test.serversWeight {
+				server := server
+				balancer.AddWeightServer(HandlerFunc(func(conn WriteCloser) {
+					_, err := conn.Write([]byte(server))
+					require.NoError(t, err)
+				}), &weight)
+			}
+
+			conn := &fakeConn{call: make(map[string]int)}
+			for i := 0; i < test.totalCall; i++ {
+				balancer.ServeTCP(conn)
+			}
+
+			assert.Equal(t, test.expected, conn.call)
+		})
+	}
+}