Added notification handler

Author: mxmauro

go.sum

@@ -5,6 +5,7 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

loadbalancer.go

@@ -13,8 +13,19 @@ type LoadBalancer struct {
 	primaryGroup       ServerGroup
 	backupGroup        ServerGroup
 	primaryOnlineCount int
+	eventHandlerMtx    sync.RWMutex
+	eventHandler       EventHandler
 }
 
+type EventHandler func(eventType int, server *Server)
+
+// -----------------------------------------------------------------------------
+
+const (
+	ServerUpEvent int = iota + 1
+	ServerDownEvent
+)
+
 const (
 	InvalidParamsErr = "invalid parameter"
 )
@@ -31,18 +42,30 @@ func Create() *LoadBalancer {
 		backupGroup: ServerGroup{
 			srvList: make([]Server, 0),
 		},
+		eventHandlerMtx: sync.RWMutex{},
 	}
 	return &lb
 }
 
+// SetEventHandler sets a new notification handler callback
+func (lb *LoadBalancer) SetEventHandler(handler EventHandler) {
+	lb.eventHandlerMtx.Lock()
+	lb.eventHandler = handler
+	lb.eventHandlerMtx.Unlock()
+}
+
 // Add adds a new server to the list
 func (lb *LoadBalancer) Add(opts ServerOptions, userData interface{}) error {
 	// Check options
 	if opts.Weight < 0 {
 		return errors.New(InvalidParamsErr)
 	}
 	if !opts.IsBackup {
-		if opts.MaxFails < 0 || opts.FailTimeout <= time.Duration(0) {
+		if opts.MaxFails > 0 {
+			if opts.FailTimeout <= time.Duration(0) {
+				return errors.New(InvalidParamsErr)
+			}
+		} else if opts.MaxFails < 0 {
 			return errors.New(InvalidParamsErr)
 		}
 	}
@@ -56,7 +79,7 @@ func (lb *LoadBalancer) Add(opts ServerOptions, userData interface{}) error {
 	if srv.opts.Weight == 0 {
 		srv.opts.Weight = 1
 	}
-	if opts.IsBackup {
+	if opts.IsBackup || srv.opts.MaxFails == 0 {
 		srv.opts.MaxFails = 0
 		srv.opts.FailTimeout = time.Duration(0)
 	}
@@ -89,21 +112,27 @@ func (lb *LoadBalancer) Add(opts ServerOptions, userData interface{}) error {
 
 // Next gets the next available server. It can return nil if no available server
 func (lb *LoadBalancer) Next() *Server {
+	var nextServer *Server
+
 	now := time.Now()
 
+	notifyUp := make([]*Server, 0) // NOTE: We would use defer, but they are executed LIFO
+
 	// Lock access
 	lb.mtx.Lock()
-	defer lb.mtx.Unlock()
 
 	// If all primary servers are offline, check if we can put someone up
 	if lb.primaryOnlineCount == 0 {
 		for idx := range lb.primaryGroup.srvList {
-			if now.After(lb.primaryGroup.srvList[idx].failTimestamp) {
-				// Put this server online again
-				lb.primaryGroup.srvList[idx].isDown = false
-				lb.primaryGroup.srvList[idx].failCounter = 0
+			srv := &lb.primaryGroup.srvList[idx]
 
+			if now.After(srv.failTimestamp) {
+				// Put this server online again
+				srv.isDown = false
+				srv.failCounter = 0
 				lb.primaryOnlineCount += 1
+
+				notifyUp = append(notifyUp, srv)
 			}
 		}
 	}
@@ -117,14 +146,17 @@ func (lb *LoadBalancer) Next() *Server {
 				// Set this server online again
 				srv.isDown = false
 				srv.lb.primaryOnlineCount += 1
+
+				notifyUp = append(notifyUp, srv)
 			}
 
 			if !srv.isDown && lb.primaryGroup.currServerWeight < srv.opts.Weight {
 				// Got a server!
 				lb.primaryGroup.currServerWeight += 1
 
-				// Done
-				return srv
+				// Select this server
+				nextServer = srv
+				break
 			}
 
 			// Advance to next server
@@ -137,17 +169,18 @@ func (lb *LoadBalancer) Next() *Server {
 		}
 	}
 
-	// If we reach here, there is no primary server available
-	if len(lb.backupGroup.srvList) > 0 {
+	// Look for backup servers if there is no primary available
+	if nextServer == nil && len(lb.backupGroup.srvList) > 0 {
 		for {
 			srv := &lb.backupGroup.srvList[lb.backupGroup.currServerIdx]
 
 			if lb.backupGroup.currServerWeight < srv.opts.Weight {
 				// Got a server!
 				lb.backupGroup.currServerWeight += 1
 
-				// Done
-				return srv
+				// Select this server
+				nextServer = srv
+				break
 			}
 
 			// Advance to next server
@@ -160,8 +193,16 @@ func (lb *LoadBalancer) Next() *Server {
 		}
 	}
 
-	// No available server
-	return nil
+	// Unlock access
+	lb.mtx.Unlock()
+
+	// Call event callback
+	for _, srv := range notifyUp {
+		lb.raiseEvent(ServerUpEvent, srv)
+	}
+
+	// Done
+	return nextServer
 }
 
 // WaitNext returns a channel that is fulfilled with the next available server
@@ -227,3 +268,14 @@ func (lb *LoadBalancer) WaitNext() (ch chan *Server) {
 
 	return
 }
+
+// -----------------------------------------------------------------------------
+// Private methods
+
+func (lb *LoadBalancer) raiseEvent(eventType int, server *Server) {
+	lb.eventHandlerMtx.RLock()
+	if lb.eventHandler != nil {
+		lb.eventHandler(eventType, server)
+	}
+	lb.eventHandlerMtx.RUnlock()
+}

loadbalancer_test.go

@@ -4,7 +4,7 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 // -----------------------------------------------------------------------------
@@ -31,9 +31,9 @@ func TestNoFail(t *testing.T) {
 
 		srvName, _ := srv.UserData().(string)
 		if (idx % serverTotalCount) < serverOneCount {
-			assert.Equal(t, serverOneName, srvName)
+			require.Equal(t, serverOneName, srvName)
 		} else {
-			assert.Equal(t, serverTwoName, srvName)
+			require.Equal(t, serverTwoName, srvName)
 		}
 
 		srv.SetOnline()
@@ -51,7 +51,7 @@ func TestFailAll(t *testing.T) {
 
 	// At this point next server should be none
 	srv := lb.Next()
-	assert.Equal(t, (*Server)(nil), srv)
+	require.Equal(t, (*Server)(nil), srv)
 }
 
 func TestBackup(t *testing.T) {
@@ -67,7 +67,7 @@ func TestBackup(t *testing.T) {
 	srv := lb.Next()
 
 	srvName, _ := srv.UserData().(string)
-	assert.Equal(t, backupServerName, srvName)
+	require.Equal(t, backupServerName, srvName)
 
 	srv.SetOffline() // NOTE: This call will act as a NO-OP
 }
@@ -83,15 +83,15 @@ func TestWait(t *testing.T) {
 
 	// At this point next server should be none
 	srv := lb.Next()
-	assert.Equal(t, (*Server)(nil), srv)
+	require.Equal(t, (*Server)(nil), srv)
 
 	// Wait until a server becomes available (after ~1sec)
 	ch := lb.WaitNext()
 	srv = <-ch
 
 	// At this point server 2 should be online again
 	srvName, _ := srv.UserData().(string)
-	assert.Equal(t, srvName, serverTwoName)
+	require.Equal(t, srvName, serverTwoName)
 }
 
 // -----------------------------------------------------------------------------

server.go

@@ -46,9 +46,10 @@ func (srv *Server) SetOnline() {
 		return
 	}
 
+	notifyUp := false
+
 	// Lock access
 	srv.lb.mtx.Lock()
-	defer srv.lb.mtx.Unlock()
 
 	// Reset the failure counter
 	srv.failCounter = 0
@@ -57,6 +58,16 @@ func (srv *Server) SetOnline() {
 	if srv.isDown {
 		srv.isDown = false
 		srv.lb.primaryOnlineCount += 1
+
+		notifyUp = true
+	}
+
+	// Unlock access
+	srv.lb.mtx.Unlock()
+
+	// Call event callback
+	if notifyUp {
+		srv.lb.raiseEvent(ServerUpEvent, srv)
 	}
 }
 
@@ -67,9 +78,10 @@ func (srv *Server) SetOffline() {
 		return
 	}
 
+	notifyDown := false
+
 	// Lock access
 	srv.lb.mtx.Lock()
-	defer srv.lb.mtx.Unlock()
 
 	// If server is up
 	if !srv.isDown && srv.failCounter < srv.opts.MaxFails {
@@ -94,8 +106,17 @@ func (srv *Server) SetOffline() {
 		if srv.failCounter == srv.opts.MaxFails {
 			srv.isDown = true
 			srv.failTimestamp = now.Add(srv.opts.FailTimeout)
-
 			srv.lb.primaryOnlineCount -= 1
+
+			notifyDown = true
 		}
 	}
+
+	// Unlock access
+	srv.lb.mtx.Unlock()
+
+	// Call event callback
+	if notifyDown {
+		srv.lb.raiseEvent(ServerDownEvent, srv)
+	}
 }