Merge branch v2.11 into v3.0

pkg/server/service/loadbalancer/wrr/wrr.go

@@ -4,6 +4,8 @@ import (
 	"container/heap"
 	"context"
 	"errors"
+	"fmt"
+	"hash/fnv"
 	"net/http"
 	"sync"
 
@@ -47,7 +49,9 @@ type Balancer struct {
 	stickyCookie     *stickyCookie
 	wantsHealthCheck bool
 
-	mutex       sync.RWMutex
+	handlersMu sync.RWMutex
+	// References all the handlers by name and also by the hashed value of the name.
+	handlerMap  map[string]*namedHandler
 	handlers    []*namedHandler
 	curDeadline float64
 	// status is a record of which child services of the Balancer are healthy, keyed
@@ -64,6 +68,7 @@ type Balancer struct {
 func New(sticky *dynamic.Sticky, wantHealthCheck bool) *Balancer {
 	balancer := &Balancer{
 		status:           make(map[string]struct{}),
+		handlerMap:       make(map[string]*namedHandler),
 		wantsHealthCheck: wantHealthCheck,
 	}
 	if sticky != nil && sticky.Cookie != nil {
@@ -74,6 +79,7 @@ func New(sticky *dynamic.Sticky, wantHealthCheck bool) *Balancer {
 			sameSite: sticky.Cookie.SameSite,
 		}
 	}
+
 	return balancer
 }
 
@@ -111,8 +117,8 @@ func (b *Balancer) Pop() interface{} {
 // 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()
+	b.handlersMu.Lock()
+	defer b.handlersMu.Unlock()
 
 	upBefore := len(b.status) > 0
 
@@ -163,8 +169,8 @@ func (b *Balancer) RegisterStatusUpdater(fn func(up bool)) error {
 var errNoAvailableServer = errors.New("no available server")
 
 func (b *Balancer) nextServer() (*namedHandler, error) {
-	b.mutex.Lock()
-	defer b.mutex.Unlock()
+	b.handlersMu.Lock()
+	defer b.handlersMu.Unlock()
 
 	if len(b.handlers) == 0 || len(b.status) == 0 {
 		return nil, errNoAvailableServer
@@ -198,22 +204,18 @@ 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 {
-					continue
-				}
+			b.handlersMu.RLock()
+			handler, ok := b.handlerMap[cookie.Value]
+			b.handlersMu.RUnlock()
 
-				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
+			if ok && handler != nil {
+				b.handlersMu.RLock()
+				_, isHealthy := b.status[handler.name]
+				b.handlersMu.RUnlock()
+				if isHealthy {
+					handler.ServeHTTP(w, req)
+					return
 				}
-
-				handler.ServeHTTP(w, req)
-				return
 			}
 		}
 	}
@@ -231,7 +233,7 @@ func (b *Balancer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	if b.stickyCookie != nil {
 		cookie := &http.Cookie{
 			Name:     b.stickyCookie.name,
-			Value:    server.name,
+			Value:    hash(server.name),
 			Path:     "/",
 			HttpOnly: b.stickyCookie.httpOnly,
 			Secure:   b.stickyCookie.secure,
@@ -257,9 +259,19 @@ func (b *Balancer) Add(name string, handler http.Handler, weight *int) {
 
 	h := &namedHandler{Handler: handler, name: name, weight: float64(w)}
 
-	b.mutex.Lock()
+	b.handlersMu.Lock()
 	h.deadline = b.curDeadline + 1/h.weight
 	heap.Push(b, h)
 	b.status[name] = struct{}{}
-	b.mutex.Unlock()
+	b.handlerMap[name] = h
+	b.handlerMap[hash(name)] = h
+	b.handlersMu.Unlock()
+}
+
+func hash(input string) string {
+	hasher := fnv.New64()
+	// We purposely ignore the error because the implementation always returns nil.
+	_, _ = hasher.Write([]byte(input))
+
+	return fmt.Sprintf("%x", hasher.Sum64())
 }

pkg/server/service/loadbalancer/wrr/wrr_test.go

@@ -247,6 +247,8 @@ func TestSticky(t *testing.T) {
 	req := httptest.NewRequest(http.MethodGet, "/", nil)
 	for i := 0; i < 3; i++ {
 		for _, cookie := range recorder.Result().Cookies() {
+			assert.NotContains(t, "test=first", cookie.Value)
+			assert.NotContains(t, "test=second", cookie.Value)
 			req.AddCookie(cookie)
 		}
 		recorder.ResponseRecorder = httptest.NewRecorder()
@@ -261,6 +263,35 @@ func TestSticky(t *testing.T) {
 	assert.Equal(t, http.SameSiteNoneMode, recorder.cookies["test"].SameSite)
 }
 
+func TestSticky_FallBack(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)
+	}), Int(1))
+
+	balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.Header().Set("server", "second")
+		rw.WriteHeader(http.StatusOK)
+	}), Int(2))
+
+	recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	req.AddCookie(&http.Cookie{Name: "test", Value: "second"})
+	for i := 0; i < 3; i++ {
+		recorder.ResponseRecorder = httptest.NewRecorder()
+
+		balancer.ServeHTTP(recorder, req)
+	}
+
+	assert.Equal(t, 0, recorder.save["first"])
+	assert.Equal(t, 3, 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) {