Add TCP Healthcheck
This commit is contained in:
Douglas De Toni Machado
GitHub
parent
d1ab6ed489
commit
8392503df7
37 changed files with 2416 additions and 307 deletions
|
|
@ -9,50 +9,7 @@ import (
|
|||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type fakeConn struct {
|
||||
writeCall map[string]int
|
||||
closeCall 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.writeCall[string(b)]++
|
||||
return len(b), nil
|
||||
}
|
||||
|
||||
func (f *fakeConn) Close() error {
|
||||
f.closeCall++
|
||||
return nil
|
||||
}
|
||||
|
||||
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) {
|
||||
func TestWRRLoadBalancer_LoadBalancing(t *testing.T) {
|
||||
testCases := []struct {
|
||||
desc string
|
||||
serversWeight map[string]int
|
||||
|
|
@ -124,9 +81,9 @@ func TestLoadBalancing(t *testing.T) {
|
|||
t.Run(test.desc, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
balancer := NewWRRLoadBalancer()
|
||||
balancer := NewWRRLoadBalancer(false)
|
||||
for server, weight := range test.serversWeight {
|
||||
balancer.AddWeightServer(HandlerFunc(func(conn WriteCloser) {
|
||||
balancer.Add(server, HandlerFunc(func(conn WriteCloser) {
|
||||
_, err := conn.Write([]byte(server))
|
||||
require.NoError(t, err)
|
||||
}), &weight)
|
||||
|
|
@ -142,3 +99,196 @@ func TestLoadBalancing(t *testing.T) {
|
|||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWRRLoadBalancer_NoServiceUp(t *testing.T) {
|
||||
balancer := NewWRRLoadBalancer(false)
|
||||
|
||||
balancer.Add("first", HandlerFunc(func(conn WriteCloser) {
|
||||
_, err := conn.Write([]byte("first"))
|
||||
require.NoError(t, err)
|
||||
}), pointer(1))
|
||||
|
||||
balancer.Add("second", HandlerFunc(func(conn WriteCloser) {
|
||||
_, err := conn.Write([]byte("second"))
|
||||
require.NoError(t, err)
|
||||
}), pointer(1))
|
||||
|
||||
balancer.SetStatus(t.Context(), "first", false)
|
||||
balancer.SetStatus(t.Context(), "second", false)
|
||||
|
||||
conn := &fakeConn{writeCall: make(map[string]int)}
|
||||
balancer.ServeTCP(conn)
|
||||
|
||||
assert.Empty(t, conn.writeCall)
|
||||
assert.Equal(t, 1, conn.closeCall)
|
||||
}
|
||||
|
||||
func TestWRRLoadBalancer_OneServerDown(t *testing.T) {
|
||||
balancer := NewWRRLoadBalancer(false)
|
||||
|
||||
balancer.Add("first", HandlerFunc(func(conn WriteCloser) {
|
||||
_, err := conn.Write([]byte("first"))
|
||||
require.NoError(t, err)
|
||||
}), pointer(1))
|
||||
|
||||
balancer.Add("second", HandlerFunc(func(conn WriteCloser) {
|
||||
_, err := conn.Write([]byte("second"))
|
||||
require.NoError(t, err)
|
||||
}), pointer(1))
|
||||
|
||||
balancer.SetStatus(t.Context(), "second", false)
|
||||
|
||||
conn := &fakeConn{writeCall: make(map[string]int)}
|
||||
for range 3 {
|
||||
balancer.ServeTCP(conn)
|
||||
}
|
||||
assert.Equal(t, 3, conn.writeCall["first"])
|
||||
}
|
||||
|
||||
func TestWRRLoadBalancer_DownThenUp(t *testing.T) {
|
||||
balancer := NewWRRLoadBalancer(false)
|
||||
|
||||
balancer.Add("first", HandlerFunc(func(conn WriteCloser) {
|
||||
_, err := conn.Write([]byte("first"))
|
||||
require.NoError(t, err)
|
||||
}), pointer(1))
|
||||
|
||||
balancer.Add("second", HandlerFunc(func(conn WriteCloser) {
|
||||
_, err := conn.Write([]byte("second"))
|
||||
require.NoError(t, err)
|
||||
}), pointer(1))
|
||||
|
||||
balancer.SetStatus(t.Context(), "second", false)
|
||||
|
||||
conn := &fakeConn{writeCall: make(map[string]int)}
|
||||
for range 3 {
|
||||
balancer.ServeTCP(conn)
|
||||
}
|
||||
assert.Equal(t, 3, conn.writeCall["first"])
|
||||
|
||||
balancer.SetStatus(t.Context(), "second", true)
|
||||
|
||||
conn = &fakeConn{writeCall: make(map[string]int)}
|
||||
for range 2 {
|
||||
balancer.ServeTCP(conn)
|
||||
}
|
||||
assert.Equal(t, 1, conn.writeCall["first"])
|
||||
assert.Equal(t, 1, conn.writeCall["second"])
|
||||
}
|
||||
|
||||
func TestWRRLoadBalancer_Propagate(t *testing.T) {
|
||||
balancer1 := NewWRRLoadBalancer(true)
|
||||
|
||||
balancer1.Add("first", HandlerFunc(func(conn WriteCloser) {
|
||||
_, err := conn.Write([]byte("first"))
|
||||
require.NoError(t, err)
|
||||
}), pointer(1))
|
||||
|
||||
balancer1.Add("second", HandlerFunc(func(conn WriteCloser) {
|
||||
_, err := conn.Write([]byte("second"))
|
||||
require.NoError(t, err)
|
||||
}), pointer(1))
|
||||
|
||||
balancer2 := NewWRRLoadBalancer(true)
|
||||
|
||||
balancer2.Add("third", HandlerFunc(func(conn WriteCloser) {
|
||||
_, err := conn.Write([]byte("third"))
|
||||
require.NoError(t, err)
|
||||
}), pointer(1))
|
||||
|
||||
balancer2.Add("fourth", HandlerFunc(func(conn WriteCloser) {
|
||||
_, err := conn.Write([]byte("fourth"))
|
||||
require.NoError(t, err)
|
||||
}), pointer(1))
|
||||
|
||||
topBalancer := NewWRRLoadBalancer(true)
|
||||
|
||||
topBalancer.Add("balancer1", balancer1, pointer(1))
|
||||
_ = balancer1.RegisterStatusUpdater(func(up bool) {
|
||||
topBalancer.SetStatus(t.Context(), "balancer1", up)
|
||||
})
|
||||
|
||||
topBalancer.Add("balancer2", balancer2, pointer(1))
|
||||
_ = balancer2.RegisterStatusUpdater(func(up bool) {
|
||||
topBalancer.SetStatus(t.Context(), "balancer2", up)
|
||||
})
|
||||
|
||||
conn := &fakeConn{writeCall: make(map[string]int)}
|
||||
for range 8 {
|
||||
topBalancer.ServeTCP(conn)
|
||||
}
|
||||
assert.Equal(t, 2, conn.writeCall["first"])
|
||||
assert.Equal(t, 2, conn.writeCall["second"])
|
||||
assert.Equal(t, 2, conn.writeCall["third"])
|
||||
assert.Equal(t, 2, conn.writeCall["fourth"])
|
||||
|
||||
// fourth gets downed, but balancer2 still up since third is still up.
|
||||
balancer2.SetStatus(t.Context(), "fourth", false)
|
||||
|
||||
conn = &fakeConn{writeCall: make(map[string]int)}
|
||||
for range 8 {
|
||||
topBalancer.ServeTCP(conn)
|
||||
}
|
||||
assert.Equal(t, 2, conn.writeCall["first"])
|
||||
assert.Equal(t, 2, conn.writeCall["second"])
|
||||
assert.Equal(t, 4, conn.writeCall["third"])
|
||||
assert.Equal(t, 0, conn.writeCall["fourth"])
|
||||
|
||||
// third gets downed, and the propagation triggers balancer2 to be marked as
|
||||
// down as well for topBalancer.
|
||||
balancer2.SetStatus(t.Context(), "third", false)
|
||||
|
||||
conn = &fakeConn{writeCall: make(map[string]int)}
|
||||
for range 8 {
|
||||
topBalancer.ServeTCP(conn)
|
||||
}
|
||||
assert.Equal(t, 4, conn.writeCall["first"])
|
||||
assert.Equal(t, 4, conn.writeCall["second"])
|
||||
assert.Equal(t, 0, conn.writeCall["third"])
|
||||
assert.Equal(t, 0, conn.writeCall["fourth"])
|
||||
}
|
||||
|
||||
func pointer[T any](v T) *T { return &v }
|
||||
|
||||
type fakeConn struct {
|
||||
writeCall map[string]int
|
||||
closeCall 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.writeCall[string(b)]++
|
||||
return len(b), nil
|
||||
}
|
||||
|
||||
func (f *fakeConn) Close() error {
|
||||
f.closeCall++
|
||||
return nil
|
||||
}
|
||||
|
||||
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")
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue