1
0
Fork 0

Implementation of serving not ready endpoints

This commit is contained in:
Valéry Fouques 2024-12-11 13:54:05 +01:00 committed by GitHub
parent a4c0b1649d
commit 9588e51146
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 204 additions and 42 deletions

View file

@ -4,7 +4,6 @@ package server
import (
"context"
"net"
"testing"
"github.com/stretchr/testify/require"

View file

@ -64,12 +64,15 @@ type Balancer 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{}
}
// New creates a new load balancer.
func New(sticky *dynamic.Sticky, wantHealthCheck bool) *Balancer {
balancer := &Balancer{
status: make(map[string]struct{}),
fenced: make(map[string]struct{}),
handlerMap: make(map[string]*namedHandler),
wantsHealthCheck: wantHealthCheck,
}
@ -179,7 +182,7 @@ func (b *Balancer) nextServer() (*namedHandler, error) {
b.handlersMu.Lock()
defer b.handlersMu.Unlock()
if len(b.handlers) == 0 || len(b.status) == 0 {
if len(b.handlers) == 0 || len(b.status) == 0 || len(b.fenced) == len(b.handlers) {
return nil, errNoAvailableServer
}
@ -194,7 +197,10 @@ func (b *Balancer) nextServer() (*namedHandler, error) {
heap.Push(b, handler)
if _, ok := b.status[handler.name]; ok {
break
if _, ok := b.fenced[handler.name]; !ok {
// do not select a fenced handler.
break
}
}
}
@ -255,7 +261,7 @@ func (b *Balancer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
// Add adds a handler.
// A handler with a non-positive weight is ignored.
func (b *Balancer) Add(name string, handler http.Handler, weight *int) {
func (b *Balancer) Add(name string, handler http.Handler, weight *int, fenced bool) {
w := 1
if weight != nil {
w = *weight
@ -271,6 +277,9 @@ func (b *Balancer) Add(name string, handler http.Handler, weight *int) {
h.deadline = b.curDeadline + 1/h.weight
heap.Push(b, h)
b.status[name] = struct{}{}
if fenced {
b.fenced[name] = struct{}{}
}
b.handlerMap[name] = h
b.handlerMap[hash(name)] = h
b.handlersMu.Unlock()

View file

@ -18,12 +18,12 @@ func TestBalancer(t *testing.T) {
balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), pointer(3))
}), pointer(3), false)
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "second")
rw.WriteHeader(http.StatusOK)
}), pointer(1))
}), pointer(1), false)
recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
for range 4 {
@ -49,9 +49,9 @@ func TestBalancerOneServerZeroWeight(t *testing.T) {
balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), pointer(1))
}), pointer(1), false)
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), pointer(0))
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), pointer(0), false)
recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
for range 3 {
@ -70,11 +70,11 @@ func TestBalancerNoServiceUp(t *testing.T) {
balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.WriteHeader(http.StatusInternalServerError)
}), pointer(1))
}), pointer(1), false)
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.WriteHeader(http.StatusInternalServerError)
}), pointer(1))
}), pointer(1), false)
balancer.SetStatus(context.WithValue(context.Background(), serviceName, "parent"), "first", false)
balancer.SetStatus(context.WithValue(context.Background(), serviceName, "parent"), "second", false)
@ -91,11 +91,11 @@ func TestBalancerOneServerDown(t *testing.T) {
balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), pointer(1))
}), pointer(1), false)
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.WriteHeader(http.StatusInternalServerError)
}), pointer(1))
}), pointer(1), false)
balancer.SetStatus(context.WithValue(context.Background(), serviceName, "parent"), "second", false)
recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
@ -112,12 +112,12 @@ func TestBalancerDownThenUp(t *testing.T) {
balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), pointer(1))
}), 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))
}), pointer(1), false)
balancer.SetStatus(context.WithValue(context.Background(), serviceName, "parent"), "second", false)
recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
@ -141,30 +141,30 @@ func TestBalancerPropagate(t *testing.T) {
balancer1.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), pointer(1))
}), pointer(1), false)
balancer1.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "second")
rw.WriteHeader(http.StatusOK)
}), pointer(1))
}), pointer(1), false)
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)
}), pointer(1))
}), pointer(1), false)
balancer2.Add("fourth", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "fourth")
rw.WriteHeader(http.StatusOK)
}), pointer(1))
}), pointer(1), false)
topBalancer := New(nil, true)
topBalancer.Add("balancer1", balancer1, pointer(1))
topBalancer.Add("balancer1", balancer1, pointer(1), false)
_ = 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.Add("balancer2", balancer2, pointer(1))
topBalancer.Add("balancer2", balancer2, pointer(1), false)
_ = balancer2.RegisterStatusUpdater(func(up bool) {
topBalancer.SetStatus(context.WithValue(context.Background(), serviceName, "top"), "balancer2", up)
})
@ -211,8 +211,20 @@ func TestBalancerPropagate(t *testing.T) {
func TestBalancerAllServersZeroWeight(t *testing.T) {
balancer := New(nil, false)
balancer.Add("test", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), pointer(0))
balancer.Add("test2", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), pointer(0))
balancer.Add("test", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), pointer(0), false)
balancer.Add("test2", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), pointer(0), false)
recorder := httptest.NewRecorder()
balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
assert.Equal(t, http.StatusServiceUnavailable, recorder.Result().StatusCode)
}
func TestBalancerAllServersFenced(t *testing.T) {
balancer := New(nil, false)
balancer.Add("test", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), pointer(1), true)
balancer.Add("test2", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), pointer(1), true)
recorder := httptest.NewRecorder()
balancer.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
@ -235,12 +247,12 @@ func TestSticky(t *testing.T) {
balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), pointer(1))
}), 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(2))
}), pointer(2), false)
recorder := &responseRecorder{
ResponseRecorder: httptest.NewRecorder(),
@ -277,12 +289,12 @@ func TestSticky_FallBack(t *testing.T) {
balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "first")
rw.WriteHeader(http.StatusOK)
}), pointer(1))
}), 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(2))
}), pointer(2), false)
recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
@ -306,12 +318,12 @@ func TestBalancerBias(t *testing.T) {
balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "A")
rw.WriteHeader(http.StatusOK)
}), pointer(11))
}), pointer(11), false)
balancer.Add("second", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.Header().Set("server", "B")
rw.WriteHeader(http.StatusOK)
}), pointer(3))
}), pointer(3), false)
recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
@ -341,3 +353,41 @@ 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{}}
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"])
}

View file

@ -258,7 +258,7 @@ func (m *Manager) getWRRServiceHandler(ctx context.Context, serviceName string,
return nil, err
}
balancer.Add(service.Name, serviceHandler, service.Weight)
balancer.Add(service.Name, serviceHandler, service.Weight, false)
if config.HealthCheck == nil {
continue
@ -397,7 +397,7 @@ func (m *Manager) getLoadBalancerServiceHandler(ctx context.Context, serviceName
proxy, _ = capture.Wrap(proxy)
}
lb.Add(proxyName, proxy, server.Weight)
lb.Add(proxyName, proxy, server.Weight, server.Fenced)
// servers are considered UP by default.
info.UpdateServerStatus(target.String(), runtime.StatusUp)